diff --git a/__init__.py b/__init__.py index f0d6552..a9d3e71 100644 --- a/__init__.py +++ b/__init__.py @@ -69,7 +69,9 @@ async def start_download(request): ) # Get the first folder path for this model type - output_dir = folder_paths.folder_names_and_paths[save_path][0][0] + # Get all eligible folder paths and find the first writable one + available_paths = folder_paths.folder_names_and_paths[save_path][0] + output_dir = next((p for p in available_paths if os.access(p, os.W_OK)), available_paths[0]) output_path = os.path.join(output_dir, safe_filename) # Final security check: ensure the resolved path is within the intended directory @@ -258,8 +260,11 @@ async def download_file(url, output_path, download_id): logging.info(f"File size for {download_id}: {total_size} bytes, supports range: {supports_range}") + # Atomic Download: Use a temporary file ending in .part + part_path = output_path + ".part" + # Create file with full size - with open(output_path, 'wb') as f: + with open(part_path, 'wb') as f: f.seek(total_size - 1) f.write(b'\0') @@ -279,7 +284,7 @@ async def download_file(url, output_path, download_id): end = start + chunk_size - 1 if i < NUM_CONNECTIONS - 1 else total_size - 1 tasks.append(download_chunk_with_progress( - session, url, start, end, output_path, i, download_id, total_size + session, url, start, end, part_path, i, download_id, total_size )) # Download all chunks in parallel @@ -293,13 +298,17 @@ async def download_file(url, output_path, download_id): else: # Fallback to single connection download logging.info(f"Using single connection for {download_id}") - await download_single_connection(session, url, output_path, download_id, total_size) + await download_single_connection(session, url, part_path, download_id, total_size) # Check if cancelled if download_control[download_id]["cancelled"]: - os.remove(output_path) + if os.path.exists(part_path): + os.remove(part_path) return + # Atomic Move: Rename .part to actual filename on success + os.rename(part_path, output_path) + # Mark as complete active_downloads[download_id]["status"] = "completed" active_downloads[download_id]["progress"] = 100 @@ -337,6 +346,7 @@ async def download_chunk_with_progress(session, url, start, end, output_path, ch chunk_size = end - start + 1 downloaded = 0 last_report_time = 0 + last_log_time = 0 try: async with session.get(url, headers=headers) as response: @@ -367,6 +377,12 @@ async def download_chunk_with_progress(session, url, start, end, output_path, ch # Send progress updates every 100ms to avoid spam (only from chunk 0) import time current_time = time.time() + + if chunk_index == 0 and (current_time - last_log_time) >= 60.0: + progress = (total_downloaded / total_size) * 100 + print(f"[AutoModelDownloader] {download_id}: {progress:.1f}%", flush=True) + last_log_time = current_time + if chunk_index == 0 and (current_time - last_report_time) >= 0.1: progress = (total_downloaded / total_size) * 100 active_downloads[download_id]["progress"] = progress @@ -389,6 +405,8 @@ async def download_chunk_with_progress(session, url, start, end, output_path, ch async def download_single_connection(session, url, output_path, download_id, total_size): """Fallback single connection download""" downloaded_size = 0 + last_log_time = 0 + import time async with session.get(url) as response: if response.status != 200: @@ -412,6 +430,11 @@ async def download_single_connection(session, url, output_path, download_id, tot active_downloads[download_id]["progress"] = progress active_downloads[download_id]["downloaded"] = downloaded_size + current_time = time.time() + if (current_time - last_log_time) >= 60.0: + print(f"[AutoModelDownloader] {download_id}: {progress:.1f}%", flush=True) + last_log_time = current_time + await PromptServer.instance.send("server_download_progress", { "download_id": download_id, "progress": progress, @@ -565,4 +588,4 @@ NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] -print("[AutoModelDownloader] Linked and Loaded!", flush=True) +print("[AutoModelDownloader] Linked and Loaded!", flush=True) diff --git a/icon.png b/icon.png new file mode 100644 index 0000000..5a38c68 Binary files /dev/null and b/icon.png differ diff --git a/pyproject.toml b/pyproject.toml index cf820fc..a7e62b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,4 +31,5 @@ Documentation = "https://github.com/MadiatorLabs/ComfyUI-AutoModelDownloader#rea [tool.comfy] PublisherId = "madiator2011" DisplayName = "Madiator2011" -includes = [] \ No newline at end of file +includes = [] +Icon = "https://raw.githubusercontent.com/FNGarvin/ComfyUI-AutoModelDownloader/main/icon.png" diff --git a/web/serverDownload.js b/web/serverDownload.js index 446530d..742e1a4 100644 --- a/web/serverDownload.js +++ b/web/serverDownload.js @@ -626,7 +626,7 @@ function injectServerDownloadButtons() { // Create button content const btnLabel = document.createElement('span'); btnLabel.className = 'p-button-label'; - btnLabel.textContent = 'Download to Host'; + btnLabel.textContent = 'Download to Comfy'; serverDownloadBtn.appendChild(btnLabel); // Status indicator (icon)