""" ComfyUI RunpodDirect - Direct Model Downloads for RunPod Download models directly to your RunPod instance with multi-connection support """ import os import json import logging import asyncio import hashlib import time import threading import folder_paths import comfy.model_management as comfy_mm from aiohttp import web from server import PromptServer # Track active downloads active_downloads = {} # Download control (for pause/resume) download_control = {} # Download queue management download_queue = [] current_download_task = None # Only one download at a time # Configuration optimized for datacenter connections (Runpod) CHUNK_SIZE = 32 * 1024 * 1024 # 32MB chunks - balanced for 500MB to 30GB+ files NUM_CONNECTIONS = 8 # 8 parallel connections - optimal for DC bandwidth CHUNK_MAX_RETRIES = 4 PROGRESS_EVENT_INTERVAL = 0.1 # seconds SOCK_CONNECT_TIMEOUT = 30 SOCK_READ_TIMEOUT = 120 HASH_READ_CHUNK_SIZE = 8 * 1024 * 1024 WS_KEEPALIVE_INTERVAL_SECONDS = 45 SETTINGS_FILENAME = "runpoddirect_settings.json" PREQUEUE_CHECK_MAX_MODELS = 512 _ws_keepalive_task = None _ws_keepalive_enabled = True _cgroup_ram_patch_applied = False _cgroup_ram_patch_enabled = True _original_virtual_memory_fn = None _settings_file_lock = threading.Lock() async def _ws_keepalive_loop(): """Send lightweight periodic websocket traffic to reduce proxy idle disconnects.""" while True: try: server = getattr(PromptServer, "instance", None) sockets = getattr(server, "sockets", None) if server is not None else None if _ws_keepalive_enabled and server is not None and sockets: stale_ids = [] for sid, ws in list(sockets.items()): try: if ws is None or ws.closed: stale_ids.append(sid) continue # Use websocket ping control frame so frontend APIs stay silent. await ws.ping() except Exception: stale_ids.append(sid) for sid in stale_ids: sockets.pop(sid, None) metadata = getattr(server, "sockets_metadata", None) if metadata is not None: metadata.pop(sid, None) except Exception: # Silent by design: keepalive should never create user-facing log noise. pass await asyncio.sleep(WS_KEEPALIVE_INTERVAL_SECONDS) def _ensure_ws_keepalive_task(): global _ws_keepalive_task if _ws_keepalive_task is not None and not _ws_keepalive_task.done(): return try: server = getattr(PromptServer, "instance", None) loop = getattr(server, "loop", None) if server is not None else None if loop is None: return _ws_keepalive_task = loop.create_task(_ws_keepalive_loop()) except Exception: _ws_keepalive_task = None def _parse_bool(value, default=True): if isinstance(value, bool): return value if isinstance(value, (int, float)): return bool(value) if isinstance(value, str): lowered = value.strip().lower() if lowered in ("1", "true", "yes", "on"): return True if lowered in ("0", "false", "no", "off"): return False return default def _read_int_file(path): try: with open(path, "r", encoding="utf-8") as f: raw = f.read().strip() if not raw or raw == "max": return None return int(raw) except Exception: return None def _read_cgroup_memory_limit_and_usage(): # cgroup v2 v2_limit = _read_int_file("/sys/fs/cgroup/memory.max") v2_usage = _read_int_file("/sys/fs/cgroup/memory.current") if v2_limit is not None or v2_usage is not None: limit = v2_limit usage = v2_usage if v2_usage is not None else 0 else: # cgroup v1 limit = _read_int_file("/sys/fs/cgroup/memory/memory.limit_in_bytes") usage = _read_int_file("/sys/fs/cgroup/memory/memory.usage_in_bytes") usage = usage if usage is not None else 0 # Very large limits in cgroup v1 commonly mean "unlimited". if limit is not None and limit >= (1 << 60): limit = None if usage < 0: usage = 0 return limit, usage def _patch_comfy_ram_detection_for_cgroups(): """Monkey-patch ComfyUI RAM detection to respect cgroup memory limits.""" global _cgroup_ram_patch_applied global _cgroup_ram_patch_enabled global _original_virtual_memory_fn if _cgroup_ram_patch_applied: return env_disabled = _parse_bool(os.environ.get("RPD_DISABLE_CGROUP_RAM_PATCH"), default=False) if env_disabled: _cgroup_ram_patch_enabled = False logging.info("[RunpodDirect] Cgroup RAM patch disabled via RPD_DISABLE_CGROUP_RAM_PATCH") return try: original_virtual_memory = comfy_mm.psutil.virtual_memory _original_virtual_memory_fn = original_virtual_memory sample = original_virtual_memory() host_total = int(getattr(sample, "total", 0) or 0) host_available = int(getattr(sample, "available", 0) or 0) limit, usage = _read_cgroup_memory_limit_and_usage() if not limit or host_total <= 0 or limit >= host_total: return def _virtual_memory_cgroup_aware(): vm = original_virtual_memory() if not _cgroup_ram_patch_enabled: return vm vm_total = int(getattr(vm, "total", 0) or 0) vm_available = int(getattr(vm, "available", 0) or 0) c_limit, c_usage = _read_cgroup_memory_limit_and_usage() if not c_limit or vm_total <= 0 or c_limit >= vm_total: return vm effective_total = min(vm_total, c_limit) cgroup_free = max(c_limit - (c_usage or 0), 0) effective_available = min(vm_available, cgroup_free, effective_total) try: return vm._replace(total=effective_total, available=effective_available) except Exception: return vm comfy_mm.psutil.virtual_memory = _virtual_memory_cgroup_aware _cgroup_ram_patch_applied = True _cgroup_ram_patch_enabled = True limit_gib = limit / (1024 ** 3) host_gib = host_total / (1024 ** 3) avail_gib = min(host_available, max(limit - usage, 0)) / (1024 ** 3) logging.info( f"[RunpodDirect] Cgroup RAM patch enabled: host={host_gib:.2f} GiB, " f"limit={limit_gib:.2f} GiB, available={avail_gib:.2f} GiB" ) except Exception as e: logging.warning(f"[RunpodDirect] Failed to apply cgroup RAM patch: {e}") def _settings_file_path(): return os.path.join(os.path.dirname(__file__), SETTINGS_FILENAME) def _build_runtime_settings_payload(): return { "keepalive_enabled": bool(_ws_keepalive_enabled), "cgroup_ram_patch_enabled": bool(_cgroup_ram_patch_enabled), } def _persist_runtime_settings(): settings_path = _settings_file_path() temp_path = f"{settings_path}.tmp" payload = _build_runtime_settings_payload() try: with _settings_file_lock: with open(temp_path, "w", encoding="utf-8") as f: json.dump(payload, f, indent=2, sort_keys=True) f.write("\n") os.replace(temp_path, settings_path) return True except Exception as e: logging.warning(f"[RunpodDirect] Failed to persist settings: {e}") try: if os.path.exists(temp_path): os.remove(temp_path) except Exception: pass return False def _load_persisted_settings(): global _ws_keepalive_enabled global _cgroup_ram_patch_enabled settings_path = _settings_file_path() if not os.path.exists(settings_path): return try: with open(settings_path, "r", encoding="utf-8") as f: data = json.load(f) if not isinstance(data, dict): return if "keepalive_enabled" in data: _ws_keepalive_enabled = _parse_bool(data.get("keepalive_enabled"), default=True) cgroup_locked_by_env = _parse_bool( os.environ.get("RPD_DISABLE_CGROUP_RAM_PATCH"), default=False ) if cgroup_locked_by_env: _cgroup_ram_patch_enabled = False elif "cgroup_ram_patch_enabled" in data: _cgroup_ram_patch_enabled = _parse_bool( data.get("cgroup_ram_patch_enabled"), default=_cgroup_ram_patch_enabled, ) logging.info( "[RunpodDirect] Loaded settings: " f"keepalive={_ws_keepalive_enabled}, " f"cgroup_ram_patch={_cgroup_ram_patch_enabled}" ) except Exception as e: logging.warning(f"[RunpodDirect] Failed to load settings: {e}") def _normalize_expected_hash(expected_hash): if not expected_hash or not isinstance(expected_hash, str): return None value = expected_hash.strip().lower() if value.startswith("sha256:"): value = value.split(":", 1)[1] if not value: return None if len(value) == 64 and all(c in "0123456789abcdef" for c in value): return value return None def _resolve_hash_type(expected_hash, expected_hash_type): hash_type = (expected_hash_type or "").strip().lower() if hash_type in ("", None): if expected_hash and len(expected_hash) == 64: return "sha256" return None if hash_type in ("sha256", "sha-256"): return "sha256" return None def _compute_file_hash(path, hash_type): if hash_type != "sha256": raise ValueError(f"Unsupported hash type: {hash_type}") digest = hashlib.sha256() with open(path, "rb") as f: while True: chunk = f.read(HASH_READ_CHUNK_SIZE) if not chunk: break digest.update(chunk) return digest.hexdigest() def _cleanup_partial_file(path): try: if path and os.path.exists(path): os.remove(path) except Exception as e: logging.warning(f"[RunpodDirect] Failed to remove partial file {path}: {e}") def _sanitize_simple_filename(filename): if not filename or not isinstance(filename, str): raise ValueError("Missing filename") if "/" in filename or "\\" in filename or os.path.sep in filename: raise ValueError("Invalid filename: must not contain path separators") if ".." in filename or filename.startswith("/") or filename.startswith("~"): raise ValueError("Invalid filename: path traversal patterns detected") safe_filename = os.path.basename(filename) if safe_filename != filename: raise ValueError("Invalid filename: must be a simple filename without path components") return safe_filename def _resolve_model_path(directory, safe_filename): if directory not in folder_paths.folder_names_and_paths: raise ValueError(f"Invalid directory: {directory}") output_dir = os.path.abspath(folder_paths.folder_names_and_paths[directory][0][0]) model_path = os.path.abspath(os.path.join(output_dir, safe_filename)) if not model_path.startswith(output_dir + os.sep): raise ValueError("Security error: attempted directory escape") return output_dir, model_path def _model_exists_in_directory(directory, safe_filename): try: _, model_path = _resolve_model_path(directory, safe_filename) except Exception: return False return os.path.exists(model_path) def _find_model_directory_anywhere(safe_filename, preferred_directories=None): ordered_directories = [] seen = set() for directory in preferred_directories or []: if directory in folder_paths.folder_names_and_paths and directory not in seen: ordered_directories.append(directory) seen.add(directory) for directory in folder_paths.folder_names_and_paths.keys(): if directory not in seen: ordered_directories.append(directory) seen.add(directory) for directory in ordered_directories: if _model_exists_in_directory(directory, safe_filename): return directory return None async def _emit_download_progress(download_id, total_size, force=False): ctrl = download_control.get(download_id) if not ctrl: return now = time.monotonic() total_downloaded = 0 should_emit = force async with ctrl["lock"]: total_downloaded = ctrl.get("total_downloaded", 0) if not should_emit: last_emit = ctrl.get("last_progress_emit", 0.0) should_emit = (now - last_emit) >= PROGRESS_EVENT_INTERVAL if should_emit: ctrl["last_progress_emit"] = now progress = 0.0 if total_size <= 0 else min((total_downloaded / total_size) * 100.0, 100.0) if download_id in active_downloads: active_downloads[download_id]["progress"] = progress active_downloads[download_id]["downloaded"] = total_downloaded if should_emit: progress = 0.0 if total_size <= 0 else min((total_downloaded / total_size) * 100.0, 100.0) await PromptServer.instance.send("server_download_progress", { "download_id": download_id, "progress": progress, "downloaded": total_downloaded, "total": total_size }) def _verify_download_integrity(output_path, total_size, expected_hash=None, expected_hash_type=None): if not os.path.exists(output_path): raise Exception("Downloaded file is missing on disk") actual_size = os.path.getsize(output_path) if actual_size != total_size: raise Exception(f"Size mismatch: expected {total_size} bytes, got {actual_size} bytes") normalized_hash = _normalize_expected_hash(expected_hash) if not normalized_hash: return {"size_verified": True, "hash_verified": False, "hash_type": None} hash_type = _resolve_hash_type(normalized_hash, expected_hash_type) if not hash_type: raise Exception(f"Unsupported hash_type: {expected_hash_type}. Supported: sha256") actual_hash = _compute_file_hash(output_path, hash_type) if actual_hash.lower() != normalized_hash: raise Exception(f"{hash_type} mismatch for downloaded file") return {"size_verified": True, "hash_verified": True, "hash_type": hash_type} @PromptServer.instance.routes.post("/server_download/start") async def start_download(request): """Start downloading a model file to the server""" try: json_data = await request.json() url = json_data.get("url") save_path = json_data.get("save_path") # e.g., "checkpoints" filename = json_data.get("filename") # e.g., "model.safetensors" token = json_data.get("token") # Optional HF token for gated models expected_hash_input = json_data.get("hash") # Optional integrity hash expected_hash_type = json_data.get("hash_type") # Optional hash type, e.g. sha256 if not url or not save_path or not filename: return web.json_response( {"error": "Missing required parameters: url, save_path, filename"}, status=400 ) expected_hash = _normalize_expected_hash(expected_hash_input) if expected_hash_input and not expected_hash: return web.json_response( {"error": "Invalid hash format. Expected SHA-256 hex digest."}, status=400 ) if expected_hash: resolved_hash_type = _resolve_hash_type(expected_hash, expected_hash_type) if not resolved_hash_type: return web.json_response( {"error": f"Unsupported hash_type: {expected_hash_type}. Supported: sha256"}, status=400 ) expected_hash_type = resolved_hash_type else: expected_hash_type = None # Validate save_path if save_path not in folder_paths.folder_names_and_paths: return web.json_response( {"error": f"Invalid save_path: {save_path}. Must be one of: {list(folder_paths.folder_names_and_paths.keys())}"}, status=400 ) # Security: Validate filename to prevent path traversal attacks # Check for any directory separators (both Unix and Windows style) if "/" in filename or "\\" in filename or os.path.sep in filename: return web.json_response( {"error": "Invalid filename: must not contain path separators"}, status=400 ) # Additional check for various path traversal patterns if ".." in filename or filename.startswith("/") or filename.startswith("~"): return web.json_response( {"error": "Invalid filename: path traversal patterns detected"}, status=400 ) # Normalize the filename to remove any potential tricks safe_filename = os.path.basename(filename) if safe_filename != filename: return web.json_response( {"error": "Invalid filename: must be a simple filename without path components"}, status=400 ) # Get the first folder path for this model type output_dir = folder_paths.folder_names_and_paths[save_path][0][0] output_path = os.path.join(output_dir, safe_filename) # Final security check: ensure the resolved path is within the intended directory output_path = os.path.abspath(output_path) output_dir = os.path.abspath(output_dir) if not output_path.startswith(output_dir + os.sep): return web.json_response( {"error": "Security error: attempted directory escape"}, status=400 ) temp_output_path = f"{output_path}.runpoddirect.part" # If a previous session crashed mid-download, remove stale temp file and start clean. if os.path.exists(temp_output_path): _cleanup_partial_file(temp_output_path) # Check if final file already exists. # If hash is available and mismatched, treat it as corrupted and redownload. if os.path.exists(output_path): if expected_hash: hash_type = _resolve_hash_type(expected_hash, expected_hash_type) actual_hash = _compute_file_hash(output_path, hash_type) if actual_hash.lower() != expected_hash.lower(): logging.warning( f"[RunpodDirect] Existing file hash mismatch for {output_path}; removing and redownloading" ) _cleanup_partial_file(output_path) else: return web.json_response( {"error": f"File already exists and hash matches: {output_path}"}, status=400 ) else: return web.json_response( {"error": f"File already exists: {output_path}"}, status=400 ) # Create directory if it doesn't exist os.makedirs(os.path.dirname(output_path), exist_ok=True) # Mark as queued download_id = f"{save_path}/{safe_filename}" active_downloads[download_id] = { "url": url, "filename": safe_filename, "save_path": save_path, "output_path": output_path, "temp_output_path": temp_output_path, "progress": 0, "status": "queued", "priority": None, "expected_hash": expected_hash, "expected_hash_type": expected_hash_type, } # Resolve token: prefer explicit token, fall back to HF_TOKEN env var resolved_token = token or os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") # Add to queue download_queue.append({ "download_id": download_id, "url": url, "output_path": output_path, "temp_output_path": temp_output_path, "token": resolved_token, "expected_hash": expected_hash, "expected_hash_type": expected_hash_type, }) # Process queue (will start download if slot available) asyncio.create_task(process_download_queue()) return web.json_response({ "success": True, "download_id": download_id, "message": "Download queued" }) except Exception as e: logging.error(f"Error starting download: {e}") return web.json_response( {"error": str(e)}, status=500 ) async def process_download_queue(): """Process the download queue - one download at a time""" global download_queue, current_download_task # Check if already downloading if current_download_task is not None and not current_download_task.done(): logging.info("[RunpodDirect] Download already in progress, waiting...") return # Already downloading if len(download_queue) == 0: logging.info("[RunpodDirect] Queue is empty") return # Nothing to process # Get next download from queue download_item = download_queue.pop(0) download_id = download_item["download_id"] url = download_item["url"] output_path = download_item["output_path"] temp_output_path = download_item.get("temp_output_path") token = download_item.get("token") expected_hash = download_item.get("expected_hash") expected_hash_type = download_item.get("expected_hash_type") # Set status to downloading active_downloads[download_id]["status"] = "downloading" active_downloads[download_id]["progress"] = 0 active_downloads[download_id]["downloaded"] = 0 logging.info(f"[RunpodDirect] Starting download {download_id} with {NUM_CONNECTIONS} connections (full speed)") # Notify frontend that download is starting await PromptServer.instance.send("server_download_progress", { "download_id": download_id, "progress": 0, "downloaded": 0, "total": 0 }) # Start download task current_download_task = asyncio.create_task( download_file( url, output_path, temp_output_path, download_id, token=token, expected_hash=expected_hash, expected_hash_type=expected_hash_type ) ) # Add completion callback to process next in queue current_download_task.add_done_callback(lambda t: on_download_complete(download_id)) def on_download_complete(download_id): """Advance the queue after a download ends, preserving its actual outcome.""" global current_download_task current_download_task = None download = active_downloads.get(download_id, {}) status = download.get("status", "unknown") if status == "completed": logging.info(f"[RunpodDirect] Download completed: {download_id}, processing next in queue...") elif status == "error": logging.error( f"[RunpodDirect] Download failed: {download_id}: {download.get('error', 'Unknown error')}, " "processing next in queue..." ) else: logging.info(f"[RunpodDirect] Download ended ({status}): {download_id}, processing next in queue...") # Process next in queue asyncio.create_task(process_download_queue()) async def download_chunk(session, url, start, end, output_path, chunk_index, download_id): """Download a specific chunk of the file""" headers = {'Range': f'bytes={start}-{end}'} try: async with session.get(url, headers=headers) as response: if response.status not in [200, 206]: return None chunk_data = await response.read() # Write chunk to file at specific position with open(output_path, 'r+b') as f: f.seek(start) f.write(chunk_data) return len(chunk_data) except Exception as e: logging.error(f"Error downloading chunk {chunk_index} for {download_id}: {e}") return None class HuggingFaceAccessError(Exception): """Authentication/access failures that must not be retried as network errors.""" def _check_hf_access(status_code, url): if 'huggingface.co' not in url or status_code not in (401, 403, 451): return repo_url = url.split('/resolve/', 1)[0] if status_code == 401: message = ( "Authentication required. Provide a valid Hugging Face token using HF_TOKEN " "or the HF token input. For gated models, first accept the repository's " f"access terms with the same account. Repo: {repo_url}" ) elif status_code == 403: message = ( "Access denied. Check that your Hugging Face token has read access and " f"that its account has accepted the model's terms at {repo_url}" ) else: message = f"Model access is restricted. Check repository access requirements at {repo_url}" raise HuggingFaceAccessError(message) async def download_file(url, output_path, temp_output_path, download_id, token=None, expected_hash=None, expected_hash_type=None): """Download file with multi-connection support and progress tracking""" import aiohttp logging.info(f"[RunpodDirect] Download {download_id} using {NUM_CONNECTIONS} connections (full speed)") working_output_path = temp_output_path or output_path try: # Initialize control for this download download_control[download_id] = { "paused": False, "cancelled": False, "total_downloaded": 0, # Shared counter for all chunks "lock": asyncio.Lock(), # Lock for thread-safe updates "last_progress_emit": 0.0, } # Build auth headers for gated HF models auth_headers = {} if token and 'huggingface.co' in url: auth_headers['Authorization'] = f'Bearer {token}' logging.info(f"[RunpodDirect] Using HF token for {download_id}") timeout = aiohttp.ClientTimeout( total=None, sock_connect=SOCK_CONNECT_TIMEOUT, sock_read=SOCK_READ_TIMEOUT, ) async with aiohttp.ClientSession(timeout=timeout, headers=auth_headers) as session: # Get file size - try HEAD first, then fall back to GET with Range total_size = 0 supports_range = False try: # Try HEAD request first async with session.head(url, allow_redirects=True) as response: _check_hf_access(response.status, url) if response.status == 200: total_size = int(response.headers.get('content-length', 0)) supports_range = response.headers.get('accept-ranges') == 'bytes' except HuggingFaceAccessError: raise except Exception as e: logging.warning(f"HEAD request failed for {download_id}: {e}") # If HEAD didn't give us the size, try GET with Range header if total_size == 0: logging.info(f"HEAD request didn't return size, trying GET with Range for {download_id}") try: headers = {'Range': 'bytes=0-0'} async with session.get(url, headers=headers, allow_redirects=True) as response: _check_hf_access(response.status, url) if response.status in [200, 206]: # Try to get size from Content-Range header first content_range = response.headers.get('content-range', '') if content_range: # Format: "bytes 0-0/12345" where 12345 is total size parts = content_range.split('/') if len(parts) == 2: total_size = int(parts[1]) supports_range = True # Fallback to Content-Length if total_size == 0: total_size = int(response.headers.get('content-length', 0)) except HuggingFaceAccessError: raise except Exception as e: logging.warning(f"GET with Range failed for {download_id}: {e}") if total_size == 0: raise Exception("Could not determine file size from server") logging.info(f"File size for {download_id}: {total_size} bytes, supports range: {supports_range}") # Create file with full size with open(working_output_path, 'wb') as f: f.seek(total_size - 1) f.write(b'\0') active_downloads[download_id]["total"] = total_size active_downloads[download_id]["downloaded"] = 0 await _emit_download_progress(download_id, total_size, force=True) # Use multi-connection download if server supports range requests if supports_range and total_size > CHUNK_SIZE: logging.info(f"Using {NUM_CONNECTIONS} connections for {download_id}") # Calculate chunk ranges chunk_size = total_size // NUM_CONNECTIONS tasks = [] for i in range(NUM_CONNECTIONS): start = i * chunk_size end = start + chunk_size - 1 if i < NUM_CONNECTIONS - 1 else total_size - 1 tasks.append(download_chunk_with_progress( session, url, start, end, working_output_path, i, download_id, total_size )) # Download all chunks in parallel results = await asyncio.gather(*tasks, return_exceptions=True) # Check for errors for result in results: if isinstance(result, Exception): raise result else: # Fallback to single connection download logging.info(f"Using single connection for {download_id}") await download_single_connection(session, url, working_output_path, download_id, total_size) # Check if cancelled if download_control[download_id]["cancelled"]: _cleanup_partial_file(working_output_path) return async with download_control[download_id]["lock"]: total_downloaded = download_control[download_id]["total_downloaded"] if total_downloaded != total_size: raise Exception(f"Incomplete download: received {total_downloaded} of {total_size} bytes") integrity_result = _verify_download_integrity( working_output_path, total_size, expected_hash=expected_hash, expected_hash_type=expected_hash_type, ) # Atomic finalize: only move into final model path after integrity checks pass. os.replace(working_output_path, output_path) # Mark as complete active_downloads[download_id]["status"] = "completed" active_downloads[download_id]["progress"] = 100 active_downloads[download_id]["downloaded"] = total_size active_downloads[download_id]["size_verified"] = integrity_result["size_verified"] active_downloads[download_id]["hash_verified"] = integrity_result["hash_verified"] active_downloads[download_id]["hash_type"] = integrity_result["hash_type"] await _emit_download_progress(download_id, total_size, force=True) # Send completion message await PromptServer.instance.send("server_download_complete", { "download_id": download_id, "path": output_path, "size": total_size, "size_verified": integrity_result["size_verified"], "hash_verified": integrity_result["hash_verified"], "hash_type": integrity_result["hash_type"], }) logging.info(f"Successfully downloaded {download_id} to {output_path}") # Cleanup del download_control[download_id] except Exception as e: logging.error(f"Error downloading {download_id}: {e}") active_downloads[download_id]["status"] = "error" active_downloads[download_id]["error"] = str(e) _cleanup_partial_file(working_output_path) await PromptServer.instance.send("server_download_error", { "download_id": download_id, "error": str(e) }) # Cleanup if download_id in download_control: del download_control[download_id] async def download_chunk_with_progress(session, url, start, end, output_path, chunk_index, download_id, total_size): """Download chunk with progress tracking""" chunk_size = end - start + 1 chunk_downloaded = 0 retries = 0 try: while chunk_downloaded < chunk_size: # Check if paused/cancelled before opening a new request while download_control.get(download_id, {}).get("paused", False): await asyncio.sleep(0.5) if download_control.get(download_id, {}).get("cancelled", False): return range_start = start + chunk_downloaded headers = {'Range': f'bytes={range_start}-{end}'} try: async with session.get(url, headers=headers) as response: _check_hf_access(response.status, url) if response.status not in [200, 206]: raise Exception(f"HTTP {response.status} for chunk {chunk_index}") # If server ignores ranged request in the middle, fail fast to avoid corruption. if range_start > 0 and response.status == 200: raise Exception(f"Server ignored range request for chunk {chunk_index} retry") with open(output_path, 'r+b') as f: f.seek(range_start) async for chunk in response.content.iter_chunked(CHUNK_SIZE): # Check if paused while download_control.get(download_id, {}).get("paused", False): await asyncio.sleep(0.5) # Check if cancelled if download_control.get(download_id, {}).get("cancelled", False): return remaining = chunk_size - chunk_downloaded if remaining <= 0: break if len(chunk) > remaining: chunk = chunk[:remaining] f.write(chunk) chunk_len = len(chunk) if chunk_len <= 0: continue chunk_downloaded += chunk_len # Update shared progress counter with lock async with download_control[download_id]["lock"]: download_control[download_id]["total_downloaded"] += chunk_len await _emit_download_progress(download_id, total_size) # Connection finished. If this chunk isn't complete yet, retry from last offset. if chunk_downloaded < chunk_size: retries += 1 if retries > CHUNK_MAX_RETRIES: raise Exception( f"Chunk {chunk_index} stalled after {CHUNK_MAX_RETRIES} retries " f"({chunk_downloaded}/{chunk_size} bytes)" ) logging.warning( f"[RunpodDirect] Chunk {chunk_index} incomplete, retry {retries}/{CHUNK_MAX_RETRIES} " f"at byte {start + chunk_downloaded}" ) await asyncio.sleep(1) else: break except (asyncio.CancelledError, HuggingFaceAccessError): raise except Exception as e: retries += 1 if retries > CHUNK_MAX_RETRIES: raise Exception(f"Chunk {chunk_index} failed after {CHUNK_MAX_RETRIES} retries: {e}") logging.warning( f"[RunpodDirect] Chunk {chunk_index} error retry {retries}/{CHUNK_MAX_RETRIES}: {e}" ) await asyncio.sleep(1) if chunk_downloaded != chunk_size: raise Exception( f"Chunk {chunk_index} size mismatch: expected {chunk_size}, got {chunk_downloaded}" ) except Exception as e: logging.error(f"Error in chunk {chunk_index} for {download_id}: {e}") raise async def download_single_connection(session, url, output_path, download_id, total_size): """Fallback single connection download""" downloaded_size = 0 async with session.get(url) as response: _check_hf_access(response.status, url) if response.status != 200: raise Exception(f"HTTP {response.status}") with open(output_path, 'wb') as f: async for chunk in response.content.iter_chunked(CHUNK_SIZE): # Check if paused while download_control.get(download_id, {}).get("paused", False): await asyncio.sleep(0.5) # Check if cancelled if download_control.get(download_id, {}).get("cancelled", False): return f.write(chunk) chunk_len = len(chunk) if chunk_len <= 0: continue downloaded_size += chunk_len async with download_control[download_id]["lock"]: download_control[download_id]["total_downloaded"] += chunk_len await _emit_download_progress(download_id, total_size) if downloaded_size != total_size: raise Exception(f"Single-connection size mismatch: expected {total_size}, got {downloaded_size}") @PromptServer.instance.routes.get("/server_download/status") async def get_download_status(request): """Get status of all downloads""" return web.json_response(active_downloads) @PromptServer.instance.routes.get("/server_download/status/{download_id:.*}") async def get_single_download_status(request): """Get status of a specific download""" download_id = request.match_info.get("download_id", "") if download_id in active_downloads: return web.json_response(active_downloads[download_id]) else: return web.json_response( {"error": "Download not found"}, status=404 ) @PromptServer.instance.routes.post("/server_download/pause") async def pause_download(request): """Pause an active download""" try: json_data = await request.json() download_id = json_data.get("download_id") if not download_id: return web.json_response( {"error": "Missing download_id"}, status=400 ) if download_id not in download_control: return web.json_response( {"error": "Download not found or already completed"}, status=404 ) download_control[download_id]["paused"] = True active_downloads[download_id]["status"] = "paused" await PromptServer.instance.send("server_download_paused", { "download_id": download_id }) return web.json_response({"success": True, "message": "Download paused"}) except Exception as e: return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.post("/server_download/resume") async def resume_download(request): """Resume a paused download""" try: json_data = await request.json() download_id = json_data.get("download_id") if not download_id: return web.json_response( {"error": "Missing download_id"}, status=400 ) if download_id not in download_control: return web.json_response( {"error": "Download not found or already completed"}, status=404 ) download_control[download_id]["paused"] = False active_downloads[download_id]["status"] = "downloading" await PromptServer.instance.send("server_download_resumed", { "download_id": download_id }) return web.json_response({"success": True, "message": "Download resumed"}) except Exception as e: return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.post("/server_download/cancel") async def cancel_download(request): """Cancel an active download""" global download_queue, current_download_task try: json_data = await request.json() download_id = json_data.get("download_id") if not download_id: return web.json_response( {"error": "Missing download_id"}, status=400 ) # Check if download is queued (not started yet) download_queue[:] = [d for d in download_queue if d["download_id"] != download_id] # Check if download is active if download_id in download_control: download_control[download_id]["cancelled"] = True # Update status if download_id in active_downloads: active_downloads[download_id]["status"] = "cancelled" await PromptServer.instance.send("server_download_cancelled", { "download_id": download_id }) return web.json_response({"success": True, "message": "Download cancelled"}) except Exception as e: return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.get("/server_download/hf_token_status") async def hf_token_status(request): """Check if HF_TOKEN environment variable is set (without exposing the value)""" has_token = bool(os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")) return web.json_response({"has_token": has_token}) @PromptServer.instance.routes.get("/server_download/keepalive_status") async def keepalive_status(request): """Get current websocket keepalive setting.""" return web.json_response({ "enabled": bool(_ws_keepalive_enabled), "interval_seconds": WS_KEEPALIVE_INTERVAL_SECONDS, "settings_file": _settings_file_path(), }) @PromptServer.instance.routes.post("/server_download/keepalive") async def set_keepalive_status(request): """Enable or disable websocket keepalive loop.""" global _ws_keepalive_enabled try: json_data = await request.json() previous = bool(_ws_keepalive_enabled) enabled = _parse_bool(json_data.get("enabled"), default=True) _ws_keepalive_enabled = bool(enabled) if not _persist_runtime_settings(): _ws_keepalive_enabled = previous return web.json_response({ "success": False, "error": "Failed to persist settings", "enabled": previous, }, status=500) return web.json_response({ "success": True, "enabled": _ws_keepalive_enabled, "interval_seconds": WS_KEEPALIVE_INTERVAL_SECONDS, "settings_file": _settings_file_path(), }) except Exception as e: return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.get("/server_download/cgroup_ram_patch_status") async def cgroup_ram_patch_status(request): """Get current cgroup RAM patch setting.""" limit, usage = _read_cgroup_memory_limit_and_usage() return web.json_response({ "enabled": bool(_cgroup_ram_patch_enabled), "applied": bool(_cgroup_ram_patch_applied), "cgroup_limit_bytes": limit, "cgroup_used_bytes": usage, "settings_file": _settings_file_path(), }) @PromptServer.instance.routes.post("/server_download/cgroup_ram_patch") async def set_cgroup_ram_patch_status(request): """Enable or disable cgroup-aware RAM override.""" global _cgroup_ram_patch_enabled try: if _parse_bool(os.environ.get("RPD_DISABLE_CGROUP_RAM_PATCH"), default=False): return web.json_response({ "success": False, "enabled": False, "locked_by_env": True, "settings_file": _settings_file_path(), }, status=409) json_data = await request.json() previous = bool(_cgroup_ram_patch_enabled) enabled = _parse_bool(json_data.get("enabled"), default=True) _cgroup_ram_patch_enabled = bool(enabled) if not _persist_runtime_settings(): _cgroup_ram_patch_enabled = previous return web.json_response({ "success": False, "error": "Failed to persist settings", "enabled": previous, "applied": bool(_cgroup_ram_patch_applied), "locked_by_env": False, }, status=500) return web.json_response({ "success": True, "enabled": _cgroup_ram_patch_enabled, "applied": bool(_cgroup_ram_patch_applied), "locked_by_env": False, "settings_file": _settings_file_path(), }) except Exception as e: return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.get("/server_download/folder_paths") async def server_download_folder_paths(request): """Return canonical ComfyUI folder path keys from backend runtime.""" try: payload = {} for key, entries in folder_paths.folder_names_and_paths.items(): normalized_key = str(key) normalized_entries = [] if isinstance(entries, (list, tuple)): for entry in entries: if isinstance(entry, (list, tuple)) and len(entry) > 0 and isinstance(entry[0], str): normalized_entries.append(entry[0]) elif isinstance(entry, str): normalized_entries.append(entry) payload[normalized_key] = normalized_entries return web.json_response(payload) except Exception as e: logging.error(f"Error getting folder paths: {e}") return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.post("/server_download/check_missing_models") async def check_missing_models(request): """Batch-check workflow models without blocking queueing on many frontend round-trips.""" try: json_data = await request.json() models = json_data.get("models") verify_hashes = _parse_bool(json_data.get("verify_hashes"), default=False) if not isinstance(models, list): return web.json_response( {"error": "Missing required parameter: models"}, status=400 ) if len(models) > PREQUEUE_CHECK_MAX_MODELS: return web.json_response( {"error": f"Too many models requested. Max: {PREQUEUE_CHECK_MAX_MODELS}"}, status=400 ) missing = [] unresolved = [] for raw_model in models: if not isinstance(raw_model, dict): continue report_model = dict(raw_model) filename = raw_model.get("filename") or raw_model.get("name") try: safe_filename = _sanitize_simple_filename(filename) except ValueError as e: report_model["filename"] = str(filename or "") report_model["reason"] = "invalid_filename" report_model["error"] = str(e) unresolved.append(report_model) continue report_model["filename"] = safe_filename directory = raw_model.get("directory") directory = str(directory) if isinstance(directory, str) and directory else None if directory and directory not in folder_paths.folder_names_and_paths: directory = None expected_hash = _normalize_expected_hash(raw_model.get("hash")) expected_hash_type = raw_model.get("hash_type") found_directory = None if directory: if _model_exists_in_directory(directory, safe_filename): found_directory = directory else: found_directory = _find_model_directory_anywhere( safe_filename, preferred_directories=[directory], ) if not found_directory: report_model["directory"] = directory report_model["reason"] = "missing" missing.append(report_model) continue else: found_directory = _find_model_directory_anywhere(safe_filename) if not found_directory: report_model["reason"] = "directory_unresolved" unresolved.append(report_model) continue if verify_hashes and expected_hash: hash_type = _resolve_hash_type(expected_hash, expected_hash_type) if not hash_type: report_model["directory"] = found_directory report_model["reason"] = "unsupported_hash_type" unresolved.append(report_model) continue _, model_path = _resolve_model_path(found_directory, safe_filename) actual_hash = _compute_file_hash(model_path, hash_type) if actual_hash.lower() != expected_hash.lower(): report_model["directory"] = found_directory report_model["corrupted"] = True report_model["reason"] = "hash_mismatch" missing.append(report_model) return web.json_response({ "success": True, "missing": missing, "unresolved": unresolved, "checked": len(models), "verify_hashes": verify_hashes, }) except Exception as e: logging.error(f"Error checking missing models: {e}") return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.post("/server_download/verify_model_integrity") async def verify_model_integrity(request): """Verify whether a model file exists and, when hash is provided, whether it matches.""" try: json_data = await request.json() directory = json_data.get("directory") or json_data.get("save_path") filename = json_data.get("filename") expected_hash_input = json_data.get("hash") expected_hash_type = json_data.get("hash_type") if not directory or not filename: return web.json_response( {"error": "Missing required parameters: directory, filename"}, status=400 ) if directory not in folder_paths.folder_names_and_paths: # Keep response successful so frontend can treat unknown folders as # non-existent instead of surfacing noisy request errors in console. return web.json_response({ "exists": False, "valid": False, "reason": "invalid_directory", }) if "/" in filename or "\\" in filename or os.path.sep in filename: return web.json_response( {"error": "Invalid filename: must not contain path separators"}, status=400 ) safe_filename = os.path.basename(filename) if safe_filename != filename: return web.json_response( {"error": "Invalid filename: must be a simple filename without path components"}, status=400 ) output_dir = os.path.abspath(folder_paths.folder_names_and_paths[directory][0][0]) model_path = os.path.abspath(os.path.join(output_dir, safe_filename)) if not model_path.startswith(output_dir + os.sep): return web.json_response( {"error": "Security error: attempted directory escape"}, status=400 ) if not os.path.exists(model_path): return web.json_response({ "exists": False, "valid": False, "reason": "missing", }) normalized_hash = _normalize_expected_hash(expected_hash_input) if expected_hash_input and not normalized_hash: return web.json_response( {"error": "Invalid hash format. Expected SHA-256 hex digest."}, status=400 ) # No hash available: we can only confirm presence. if not normalized_hash: return web.json_response({ "exists": True, "valid": True, "hash_verified": False, "reason": "no_hash", }) hash_type = _resolve_hash_type(normalized_hash, expected_hash_type) if not hash_type: return web.json_response( {"error": f"Unsupported hash_type: {expected_hash_type}. Supported: sha256"}, status=400 ) actual_hash = _compute_file_hash(model_path, hash_type) hash_ok = actual_hash.lower() == normalized_hash.lower() return web.json_response({ "exists": True, "valid": hash_ok, "hash_verified": hash_ok, "hash_type": hash_type, "reason": "ok" if hash_ok else "hash_mismatch", }) except Exception as e: logging.error(f"Error verifying model integrity: {e}") return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.post("/server_download/validate_hf_token") async def validate_hf_token(request): """Validate a Hugging Face token and optionally check access to specific model URLs""" import aiohttp try: json_data = await request.json() token = json_data.get("token", "") urls = json_data.get("urls", []) # Optional: list of model URLs to check access # If token is the sentinel '__env__', use the environment variable if token == "__env__": token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or "" if not token: return web.json_response({"valid": False, "error": "No token provided"}, status=400) # Validate token against HF API timeout = aiohttp.ClientTimeout(total=10) async with aiohttp.ClientSession(timeout=timeout) as session: # Check token validity via whoami endpoint async with session.get( "https://huggingface.co/api/whoami-v2", headers={"Authorization": f"Bearer {token}"} ) as response: if response.status != 200: return web.json_response({"valid": False, "error": "Invalid token"}) user_data = await response.json() username = user_data.get("name", "unknown") # If URLs provided, check access to each url_access = {} for url in urls[:10]: # Limit to 10 URLs try: async with session.head( url, headers={"Authorization": f"Bearer {token}"}, allow_redirects=True ) as resp: if resp.status == 200: url_access[url] = {"accessible": True} elif resp.status in (401, 403, 451): repo_url = url[:url.find('/resolve/')] if '/resolve/' in url else url url_access[url] = { "accessible": False, "reason": "terms_not_accepted", "repo_url": repo_url } else: url_access[url] = {"accessible": False, "reason": f"HTTP {resp.status}"} except Exception: url_access[url] = {"accessible": False, "reason": "request_failed"} return web.json_response({ "valid": True, "username": username, "url_access": url_access }) except Exception as e: logging.error(f"Error validating HF token: {e}") return web.json_response({"valid": False, "error": str(e)}, status=500) @PromptServer.instance.routes.get("/extensions/ComfyUI-RunpodDirect/serverDownload.js") async def serve_js_with_version(request): """Serve JS file with cache-busting headers""" js_path = os.path.join(os.path.dirname(__file__), "web", "serverDownload.js") response = web.FileResponse(js_path) # Add cache control headers to force revalidation response.headers['Cache-Control'] = 'no-cache, must-revalidate' response.headers['Pragma'] = 'no-cache' response.headers['Expires'] = '0' response.headers['X-Version'] = __version__ return response # Set the web directory for frontend files WEB_DIRECTORY = "./web" # Version for cache busting - increment this when you update the JS __version__ = "1.0.12" # Apply cgroup-aware RAM patch, then load persisted settings, then start keepalive. _patch_comfy_ram_detection_for_cgroups() _load_persisted_settings() _ensure_ws_keepalive_task() NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]