alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-03-21 21:22:11 -03:00
Code Issues Packages Projects Releases Wiki Activity

feat: Implement download progress WebSocket and enhance download manager with unique IDs

Browse Source
This commit is contained in:
Will Miao
2025-07-02 23:48:35 +08:00
parent 49c4a4068b
commit 7a4b5a4667
7 changed files with 160 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
static/js/managers/CheckpointDownloadManager.js
View File

@@ -301,13 +301,24 @@ export class CheckpointDownloadManager {
const updateProgress = this.loadingManager.showDownloadProgress(1);
updateProgress(0, 0, this.currentVersion.name);
// Setup WebSocket for progress updates using checkpoint-specific endpoint
// Generate a unique ID for this download
const downloadId = Date.now().toString();
// Setup WebSocket for progress updates using download-specific endpoint
const wsProtocol = window.location.protocol === 'https:' ? 'wss://' : 'ws://';
const ws = new WebSocket(`${wsProtocol}${window.location.host}/ws/fetch-progress`);
const ws = new WebSocket(`${wsProtocol}${window.location.host}/ws/download-progress?id=${downloadId}`);
ws.onmessage = (event) => {
const data = JSON.parse(event.data);
if (data.status === 'progress') {
// Handle download ID confirmation
if (data.type === 'download_id') {
console.log(`Connected to checkpoint download progress with ID: ${data.download_id}`);
return;
}
// Only process progress updates for our download
if (data.status === 'progress' && data.download_id === downloadId) {
// Update progress display with current progress
updateProgress(data.progress, 0, this.currentVersion.name);
@@ -329,7 +340,7 @@ export class CheckpointDownloadManager {
// Continue with download even if WebSocket fails
};
// Start download using checkpoint download endpoint
// Start download using checkpoint download endpoint with download ID
const response = await fetch('/api/download-model', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
@@ -337,7 +348,8 @@ export class CheckpointDownloadManager {
model_id: this.modelId,
model_version_id: this.currentVersion.id,
model_root: checkpointRoot,
relative_path: targetFolder
relative_path: targetFolder,
download_id: downloadId
})
});

25
static/js/managers/DownloadManager.js
View File

@@ -311,13 +311,24 @@ export class DownloadManager {
const updateProgress = this.loadingManager.showDownloadProgress(1);
updateProgress(0, 0, this.currentVersion.name);
// Setup WebSocket for progress updates
// Generate a unique ID for this download
const downloadId = Date.now().toString();
// Setup WebSocket for progress updates - use download-specific endpoint
const wsProtocol = window.location.protocol === 'https:' ? 'wss://' : 'ws://';
const ws = new WebSocket(`${wsProtocol}${window.location.host}/ws/fetch-progress`);
const ws = new WebSocket(`${wsProtocol}${window.location.host}/ws/download-progress?id=${downloadId}`);
ws.onmessage = (event) => {
const data = JSON.parse(event.data);
if (data.status === 'progress') {
// Handle download ID confirmation
if (data.type === 'download_id') {
console.log(`Connected to download progress with ID: ${data.download_id}`);
return;
}
// Only process progress updates for our download
if (data.status === 'progress' && data.download_id === downloadId) {
// Update progress display with current progress
updateProgress(data.progress, 0, this.currentVersion.name);
@@ -339,7 +350,7 @@ export class DownloadManager {
// Continue with download even if WebSocket fails
};
// Start download
// Start download with our download ID
const response = await fetch('/api/download-model', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
@@ -347,7 +358,8 @@ export class DownloadManager {
model_id: this.modelId,
model_version_id: this.currentVersion.id,
model_root: loraRoot,
relative_path: targetFolder
relative_path: targetFolder,
download_id: downloadId
})
});
@@ -358,6 +370,9 @@ export class DownloadManager {
showToast('Download completed successfully', 'success');
modalManager.closeModal('downloadModal');
// Close WebSocket after download completes
ws.close();
// Update state and trigger reload with folder update
state.activeFolder = targetFolder;
await resetAndReload(true); // Pass true to update folders

25
static/js/managers/import/DownloadManager.js
View File

@@ -128,9 +128,12 @@ export class DownloadManager {
targetPath += '/' + newFolder;
}
// Generate a unique ID for this batch download
const batchDownloadId = Date.now().toString();
// Set up WebSocket for progress updates
const wsProtocol = window.location.protocol === 'https:' ? 'wss://' : 'ws://';
const ws = new WebSocket(`${wsProtocol}${window.location.host}/ws/fetch-progress`);
const ws = new WebSocket(`${wsProtocol}${window.location.host}/ws/download-progress?id=${batchDownloadId}`);
// Show enhanced loading with progress details for multiple items
const updateProgress = this.importManager.loadingManager.showDownloadProgress(
@@ -145,7 +148,15 @@ export class DownloadManager {
// Set up progress tracking for current download
ws.onmessage = (event) => {
const data = JSON.parse(event.data);
if (data.status === 'progress') {
// Handle download ID confirmation
if (data.type === 'download_id') {
console.log(`Connected to batch download progress with ID: ${data.download_id}`);
return;
}
// Process progress updates for our current active download
if (data.status === 'progress' && data.download_id && data.download_id.startsWith(batchDownloadId)) {
// Update current LoRA progress
currentLoraProgress = data.progress;
@@ -188,16 +199,16 @@ export class DownloadManager {
updateProgress(0, completedDownloads, lora.name);
try {
// Download the LoRA
// Download the LoRA with download ID
const response = await fetch('/api/download-model', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
download_url: lora.downloadUrl,
model_version_id: lora.modelVersionId,
model_hash: lora.hash,
model_id: lora.modelId,
model_version_id: lora.id,
model_root: loraRoot,
relative_path: targetPath.replace(loraRoot + '/', '')
relative_path: targetPath.replace(loraRoot + '/', ''),
download_id: batchDownloadId
})
});