diff --git a/__init__.py b/__init__.py index 3cfd08e..44dfc29 100644 --- a/__init__.py +++ b/__init__.py @@ -313,6 +313,7 @@ from .nodes import ( HyVideoModelLoader, HyVideoVAELoader, DownloadAndLoadHyVideoTextEncoder, + UNetLoaderLP, FullCleanupMultiGPU, ) @@ -376,13 +377,28 @@ def soft_empty_cache_distorch2_patched(force=False): is_distorch_active = False # Detect DisTorch2-managed models - for lm in mm.current_loaded_models: + logger.mgpu_mm_log(f"[DETECT_DEBUG] Checking DisTorch2 active status - loaded models: {len(mm.current_loaded_models)}, store entries: {len(safetensor_allocation_store)}") + + for i, lm in enumerate(mm.current_loaded_models): mp = lm.model # weakref call to ModelPatcher if mp is not None: - model_hash = create_safetensor_model_hash(mp, "cache_patch_check") - if model_hash in safetensor_allocation_store and safetensor_allocation_store.get(model_hash): - is_distorch_active = True - break + try: + model_hash = create_safetensor_model_hash(mp, "cache_patch_check") + in_store = model_hash in safetensor_allocation_store + alloc_value = safetensor_allocation_store.get(model_hash, "") + model_name = type(getattr(mp, 'model', mp)).__name__ + keep_loaded = getattr(getattr(mp, 'model', None), '_mgpu_keep_loaded', False) + + logger.mgpu_mm_log(f"[DETECT_DEBUG] Model {i}: {model_name}, hash={model_hash[:8]}, in_store={in_store}, alloc_value='{alloc_value}', keep_loaded={keep_loaded}") + + if in_store and alloc_value: + is_distorch_active = True + logger.mgpu_mm_log(f"[DETECT_DEBUG] DisTorch2 ACTIVE detected on model: {model_name}") + break + except Exception as e: + logger.mgpu_mm_log(f"[DETECT_DEBUG] Model {i}: Error during detection - {e}") + + logger.mgpu_mm_log(f"[DETECT_DEBUG] Final DisTorch2 active status: {is_distorch_active}") # Phase 2: adaptive CPU memory management check_cpu_memory_threshold() @@ -645,19 +661,6 @@ if hasattr(mm, 'load_models_gpu') and not hasattr(mm.load_models_gpu, "_distorch result = original_load_models_gpu(models, memory_required, force_patch_weights, minimum_memory_required, force_full_load) multigpu_memory_log("patched_load_models_gpu", "post-original-call") - # Cleanup policy triggers (flags-only, Manager semantics) - if MGPU_CLEANUP_POLICY in ("threshold", "every_load+threshold", "threshold+every_load"): - try: - check_cpu_memory_threshold(threshold_percent=MGPU_CPU_RESET_THRESHOLD * 100.0) - except Exception: - pass - if MGPU_CLEANUP_POLICY in ("every_load", "every_load+threshold", "threshold+every_load"): - try: - # flags-only; prompt worker performs unload/reset/gc - force_full_system_cleanup(reason="policy_every_load", force=False) - except Exception: - pass - return result # Mark and apply the patch @@ -681,6 +684,7 @@ NODE_CLASS_MAPPINGS = { "HunyuanVideoEmbeddingsAdapter": HunyuanVideoEmbeddingsAdapter, "CheckpointLoaderAdvancedMultiGPU": CheckpointLoaderAdvancedMultiGPU, "CheckpointLoaderAdvancedDisTorch2MultiGPU": CheckpointLoaderAdvancedDisTorch2MultiGPU, + "UNetLoaderLP": UNetLoaderLP, } # Standard MultiGPU nodes diff --git a/distorch_2.py b/distorch_2.py index cd66bef..f7d6043 100644 --- a/distorch_2.py +++ b/distorch_2.py @@ -67,8 +67,11 @@ def register_patched_safetensor_modelpatcher(): multigpu_memory_log(f"safetensor:{debug_hash[:8]}", "pre-load") allocations = safetensor_allocation_store.get(debug_hash) - if not hasattr(self.model, '_distorch_high_precision_loras') or not allocations: + # Set default precision flag before checking + if not hasattr(self.model, '_distorch_high_precision_loras'): + self.model._distorch_high_precision_loras = True + if not allocations: result = original_partially_load(self, device_to, extra_memory, force_patch_weights) multigpu_memory_log(f"safetensor:{debug_hash[:8]}", "post-load") if hasattr(self, '_distorch_block_assignments'): @@ -103,7 +106,7 @@ def register_patched_safetensor_modelpatcher(): device_assignments = analyze_safetensor_loading(self, allocations) model_original_dtype = comfy.utils.weight_dtype(self.model.state_dict()) - high_precision_loras = self.model._distorch_high_precision_loras + high_precision_loras = getattr(self.model, "_distorch_high_precision_loras", True) loading = self._load_list() loading.sort(reverse=True) for module_size, module_name, module_object, params in loading: @@ -813,7 +816,7 @@ def override_class_with_distorch_safetensor_v2(cls): inputs["optional"]["virtual_vram_gb"] = ("FLOAT", {"default": 4.0, "min": 0.0, "max": 128.0, "step": 0.1}) inputs["optional"]["donor_device"] = (devices, {"default": "cpu"}) inputs["optional"]["expert_mode_allocations"] = ("STRING", {"multiline": False, "default": ""}) - inputs["optional"]["high_precision_loras"] = ("BOOLEAN", {"default": True}) + inputs["optional"]["keep_loaded"] = ("BOOLEAN", {"default": True}) return inputs CATEGORY = "multigpu/distorch_2" @@ -822,13 +825,13 @@ def override_class_with_distorch_safetensor_v2(cls): @classmethod def IS_CHANGED(s, *args, compute_device=None, virtual_vram_gb=4.0, - donor_device="cpu", expert_mode_allocations="", high_precision_loras=True, **kwargs): + donor_device="cpu", expert_mode_allocations="", keep_loaded=True, **kwargs): # Create a hash of our specific settings - settings_str = f"{compute_device}{virtual_vram_gb}{donor_device}{expert_mode_allocations}{high_precision_loras}" + settings_str = f"{compute_device}{virtual_vram_gb}{donor_device}{expert_mode_allocations}{keep_loaded}" return hashlib.sha256(settings_str.encode()).hexdigest() def override(self, *args, compute_device=None, virtual_vram_gb=4.0, - donor_device="cpu", expert_mode_allocations="", high_precision_loras=True, **kwargs): + donor_device="cpu", expert_mode_allocations="", keep_loaded=True, **kwargs): from . import set_current_device if compute_device is not None: @@ -837,39 +840,7 @@ def override_class_with_distorch_safetensor_v2(cls): # Register our patched ModelPatcher register_patched_safetensor_modelpatcher() - # Call original function - fn = getattr(super(), cls.FUNCTION) - - # --- Check if we need to unload the model due to settings change --- - # This logic is a bit redundant with IS_CHANGED, but provides clear logging - settings_str = f"{compute_device}{virtual_vram_gb}{donor_device}{expert_mode_allocations}" - settings_hash = hashlib.sha256(settings_str.encode()).hexdigest() - - # Temporarily load to get hash without applying our patch - temp_out = fn(*args, **kwargs) - model_to_check = None - if hasattr(temp_out[0], 'model'): - model_to_check = temp_out[0] - elif hasattr(temp_out[0], 'patcher') and hasattr(temp_out[0].patcher, 'model'): - model_to_check = temp_out[0].patcher - - if model_to_check: - model_hash = create_safetensor_model_hash(model_to_check, "override_check") - last_settings_hash = safetensor_settings_store.get(model_hash) - - if last_settings_hash != settings_hash: - logger.mgpu_mm_log(f"Settings changed for model {model_hash[:8]}. Previous settings hash: {last_settings_hash}, New settings hash: {settings_hash}. Forcing reload.") - else: - logger.mgpu_mm_log(f"Settings unchanged for model {model_hash[:8]}. Using cached model.") - - out = fn(*args, **kwargs) - - # Store high_precision_loras in the model for later retrieval - if hasattr(out[0], 'model'): - out[0].model._distorch_high_precision_loras = high_precision_loras - elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'): - out[0].patcher.model._distorch_high_precision_loras = high_precision_loras - + # Build allocation string vram_string = "" if virtual_vram_gb > 0: vram_string = f"{compute_device};{virtual_vram_gb};{donor_device}" @@ -877,17 +848,35 @@ def override_class_with_distorch_safetensor_v2(cls): vram_string = compute_device full_allocation = f"{expert_mode_allocations}#{vram_string}" if expert_mode_allocations or vram_string else "" + + fn = getattr(super(), cls.FUNCTION) + + # Load the model and get hash, then store allocation for future runs + out = fn(*args, **kwargs) + + model_to_check = None + if hasattr(out[0], 'model'): + model_to_check = out[0] + elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'): + model_to_check = out[0].patcher + + if model_to_check: + model_hash = create_safetensor_model_hash(model_to_check, "override_store") + settings_str = f"{compute_device}{virtual_vram_gb}{donor_device}{expert_mode_allocations}" + settings_hash = hashlib.sha256(settings_str.encode()).hexdigest() + + # Store allocation for next run - this enables DisTorch for subsequent loads + safetensor_allocation_store[model_hash] = full_allocation + safetensor_settings_store[model_hash] = settings_hash + logger.debug(f"[MultiGPU DisTorch V2] Stored allocation for model {model_hash[:8]}: {full_allocation}") logger.info(f"[MultiGPU DisTorch V2] Full allocation string: {full_allocation}") + # Store keep_loaded in the model for later retrieval by unload_all_models patch if hasattr(out[0], 'model'): - model_hash = create_safetensor_model_hash(out[0], "override") - safetensor_allocation_store[model_hash] = full_allocation - safetensor_settings_store[model_hash] = settings_hash + out[0].model._mgpu_keep_loaded = keep_loaded elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'): - model_hash = create_safetensor_model_hash(out[0].patcher, "override") - safetensor_allocation_store[model_hash] = full_allocation - safetensor_settings_store[model_hash] = settings_hash + out[0].patcher.model._mgpu_keep_loaded = keep_loaded return out @@ -909,7 +898,7 @@ def override_class_with_distorch_safetensor_v2_clip(cls): inputs["optional"]["virtual_vram_gb"] = ("FLOAT", {"default": 4.0, "min": 0.0, "max": 128.0, "step": 0.1}) inputs["optional"]["donor_device"] = (devices, {"default": "cpu"}) inputs["optional"]["expert_mode_allocations"] = ("STRING", {"multiline": False, "default": ""}) - inputs["optional"]["high_precision_loras"] = ("BOOLEAN", {"default": True}) + inputs["optional"]["keep_loaded"] = ("BOOLEAN", {"default": True}) return inputs CATEGORY = "multigpu/distorch_2" @@ -918,13 +907,13 @@ def override_class_with_distorch_safetensor_v2_clip(cls): @classmethod def IS_CHANGED(s, *args, device=None, virtual_vram_gb=4.0, # Changed from compute_device - donor_device="cpu", expert_mode_allocations="", high_precision_loras=True, **kwargs): + donor_device="cpu", expert_mode_allocations="", keep_loaded=True, **kwargs): # Create a hash of our specific settings - settings_str = f"{device}{virtual_vram_gb}{donor_device}{expert_mode_allocations}{high_precision_loras}" # Changed from compute_device + settings_str = f"{device}{virtual_vram_gb}{donor_device}{expert_mode_allocations}{keep_loaded}" # Changed from compute_device return hashlib.sha256(settings_str.encode()).hexdigest() def override(self, *args, device=None, virtual_vram_gb=4.0, # Changed from compute_device - donor_device="cpu", expert_mode_allocations="", high_precision_loras=True, **kwargs): + donor_device="cpu", expert_mode_allocations="", keep_loaded=True, **kwargs): from . import set_current_text_encoder_device # Use text encoder device setter if device is not None: @@ -938,35 +927,14 @@ def override_class_with_distorch_safetensor_v2_clip(cls): # Call original function fn = getattr(super(), cls.FUNCTION) - # --- Check if we need to unload the model due to settings change --- - # This logic is a bit redundant with IS_CHANGED, but provides clear logging - settings_str = f"{device}{virtual_vram_gb}{donor_device}{expert_mode_allocations}" # Changed from compute_device - settings_hash = hashlib.sha256(settings_str.encode()).hexdigest() - - # Temporarily load to get hash without applying our patch - temp_out = fn(*args, **kwargs) - model_to_check = None - if hasattr(temp_out[0], 'model'): - model_to_check = temp_out[0] - elif hasattr(temp_out[0], 'patcher') and hasattr(temp_out[0].patcher, 'model'): - model_to_check = temp_out[0].patcher - - if model_to_check: - model_hash = create_safetensor_model_hash(model_to_check, "override_check") - last_settings_hash = safetensor_settings_store.get(model_hash) - - if last_settings_hash != settings_hash: - logger.mgpu_mm_log(f"Settings changed for model {model_hash[:8]}. Previous settings hash: {last_settings_hash}, New settings hash: {settings_hash}. Forcing reload.") - else: - logger.mgpu_mm_log(f"Settings unchanged for model {model_hash[:8]}. Using cached model.") - + # Call the main function once out = fn(*args, **kwargs) - # Store high_precision_loras in the model for later retrieval + # Store keep_loaded in the model for later retrieval by unload_all_models patch if hasattr(out[0], 'model'): - out[0].model._distorch_high_precision_loras = high_precision_loras + out[0].model._mgpu_keep_loaded = keep_loaded elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'): - out[0].patcher.model._distorch_high_precision_loras = high_precision_loras + out[0].patcher.model._mgpu_keep_loaded = keep_loaded vram_string = "" if virtual_vram_gb > 0: @@ -978,12 +946,19 @@ def override_class_with_distorch_safetensor_v2_clip(cls): logger.info(f"[MultiGPU DisTorch V2] Full allocation string: {full_allocation}") + # Store allocation AFTER loading for next time + model_to_check = None if hasattr(out[0], 'model'): - model_hash = create_safetensor_model_hash(out[0], "override") - safetensor_allocation_store[model_hash] = full_allocation - safetensor_settings_store[model_hash] = settings_hash + model_to_check = out[0] elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'): - model_hash = create_safetensor_model_hash(out[0].patcher, "override") + model_to_check = out[0].patcher + + if model_to_check: + model_hash = create_safetensor_model_hash(model_to_check, "override_store") + settings_str = f"{device}{virtual_vram_gb}{donor_device}{expert_mode_allocations}" + settings_hash = hashlib.sha256(settings_str.encode()).hexdigest() + + # Store allocation for next time safetensor_allocation_store[model_hash] = full_allocation safetensor_settings_store[model_hash] = settings_hash @@ -1006,7 +981,7 @@ def override_class_with_distorch_safetensor_v2_clip_no_device(cls): inputs["optional"]["virtual_vram_gb"] = ("FLOAT", {"default": 4.0, "min": 0.0, "max": 128.0, "step": 0.1}) inputs["optional"]["donor_device"] = (devices, {"default": "cpu"}) inputs["optional"]["expert_mode_allocations"] = ("STRING", {"multiline": False, "default": ""}) - inputs["optional"]["high_precision_loras"] = ("BOOLEAN", {"default": True}) + inputs["optional"]["keep_loaded"] = ("BOOLEAN", {"default": True}) return inputs CATEGORY = "multigpu/distorch_2" @@ -1015,13 +990,13 @@ def override_class_with_distorch_safetensor_v2_clip_no_device(cls): @classmethod def IS_CHANGED(s, *args, device=None, virtual_vram_gb=4.0, # Changed from compute_device - donor_device="cpu", expert_mode_allocations="", high_precision_loras=True, **kwargs): + donor_device="cpu", expert_mode_allocations="", keep_loaded=True, **kwargs): # Create a hash of our specific settings - settings_str = f"{device}{virtual_vram_gb}{donor_device}{expert_mode_allocations}{high_precision_loras}" # Changed from compute_device + settings_str = f"{device}{virtual_vram_gb}{donor_device}{expert_mode_allocations}{keep_loaded}" # Changed from compute_device return hashlib.sha256(settings_str.encode()).hexdigest() def override(self, *args, device=None, virtual_vram_gb=4.0, # Changed from compute_device - donor_device="cpu", expert_mode_allocations="", high_precision_loras=True, **kwargs): + donor_device="cpu", expert_mode_allocations="", keep_loaded=True, **kwargs): from . import set_current_text_encoder_device # Use text encoder device setter if device is not None: @@ -1033,35 +1008,14 @@ def override_class_with_distorch_safetensor_v2_clip_no_device(cls): # Call original function fn = getattr(super(), cls.FUNCTION) - # --- Check if we need to unload the model due to settings change --- - # This logic is a bit redundant with IS_CHANGED, but provides clear logging - settings_str = f"{device}{virtual_vram_gb}{donor_device}{expert_mode_allocations}" # Changed from compute_device - settings_hash = hashlib.sha256(settings_str.encode()).hexdigest() - - # Temporarily load to get hash without applying our patch - temp_out = fn(*args, **kwargs) - model_to_check = None - if hasattr(temp_out[0], 'model'): - model_to_check = temp_out[0] - elif hasattr(temp_out[0], 'patcher') and hasattr(temp_out[0].patcher, 'model'): - model_to_check = temp_out[0].patcher - - if model_to_check: - model_hash = create_safetensor_model_hash(model_to_check, "override_check") - last_settings_hash = safetensor_settings_store.get(model_hash) - - if last_settings_hash != settings_hash: - logger.info(f"[MultiGPU DisTorch V2] Settings changed for model {model_hash[:8]}. Previous settings hash: {last_settings_hash}, New settings hash: {settings_hash}. Forcing reload.") - else: - logger.info(f"[MultiGPU DisTorch V2] Settings unchanged for model {model_hash[:8]}. Using cached model.") - + # Call the main function once out = fn(*args, **kwargs) - # Store high_precision_loras in the model for later retrieval + # Store keep_loaded in the model for later retrieval by unload_all_models patch if hasattr(out[0], 'model'): - out[0].model._distorch_high_precision_loras = high_precision_loras + out[0].model._mgpu_keep_loaded = keep_loaded elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'): - out[0].patcher.model._distorch_high_precision_loras = high_precision_loras + out[0].patcher.model._mgpu_keep_loaded = keep_loaded vram_string = "" if virtual_vram_gb > 0: @@ -1073,12 +1027,19 @@ def override_class_with_distorch_safetensor_v2_clip_no_device(cls): logger.info(f"[MultiGPU DisTorch V2] Full allocation string: {full_allocation}") + # Store allocation AFTER loading for next time + model_to_check = None if hasattr(out[0], 'model'): - model_hash = create_safetensor_model_hash(out[0], "override") - safetensor_allocation_store[model_hash] = full_allocation - safetensor_settings_store[model_hash] = settings_hash + model_to_check = out[0] elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'): - model_hash = create_safetensor_model_hash(out[0].patcher, "override") + model_to_check = out[0].patcher + + if model_to_check: + model_hash = create_safetensor_model_hash(model_to_check, "override_store") + settings_str = f"{device}{virtual_vram_gb}{donor_device}{expert_mode_allocations}" + settings_hash = hashlib.sha256(settings_str.encode()).hexdigest() + + # Store allocation for next time safetensor_allocation_store[model_hash] = full_allocation safetensor_settings_store[model_hash] = settings_hash diff --git a/memory-bank/code-references.md b/memory-bank/code-references.md new file mode 100644 index 0000000..3061f98 --- /dev/null +++ b/memory-bank/code-references.md @@ -0,0 +1,146 @@ +# Code References (Definitive): ComfyUI Manager “Free model and node cache” + +Purpose +- Provide an end-to-end, fully verified lineage of the ComfyUI Manager “Free model and node cache” button through to the exact consumption of flags in ComfyUI core, with exact file paths and code excerpts captured from the current snapshot in this workspace. + +End‑to‑End Flow (Current Snapshot) +1) UI Button (Manager) → 2) JS helper free_models(...) → 3) POST /free (Comfy core) → 4) main.py prompt_worker thread polls flags and performs: + - unload_models: comfy.model_management.unload_all_models() + - free_memory: PromptExecutor.reset() + - Additionally triggers GC and comfy.model_management.soft_empty_cache() + +A) Frontend UI trigger (ComfyUI Manager) +- File: ../ComfyUI-Manager/js/comfyui-manager.js +- Location: app.registerExtension({ name: "Comfy.ManagerMenu", ... }) → setup() → ComfyButtonGroup +```js +new(await import("../../scripts/ui/components/button.js")).ComfyButton({ + icon: "vacuum-outline", + action: () => { + free_models(); + }, + tooltip: "Unload Models" +}).element, +new(await import("../../scripts/ui/components/button.js")).ComfyButton({ + icon: "vacuum", + action: () => { + free_models(true); + }, + tooltip: "Free model and node cache" +}).element, +``` +Semantics: +- “Unload Models” → free_models() (models only) +- “Free model and node cache” → free_models(true) (models + execution cache) + +B) Frontend request construction (ComfyUI Manager) +- File: ../ComfyUI-Manager/js/common.js +- Function: export async function free_models(free_execution_cache) +```js +export async function free_models(free_execution_cache) { + try { + let mode = ""; + if (free_execution_cache) { + mode = '{"unload_models": true, "free_memory": true}'; + } else { + mode = '{"unload_models": true}'; + } + + console.log(`[ManagerFreePath] POST /free payload: ${mode}`); + let res = await api.fetchApi(`/free`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: mode + }); + console.log(`[ManagerFreePath] /free status: ${res.status}`); + + if (res.status == 200) { + if (free_execution_cache) { + showToast("'Models' and 'Execution Cache' have been cleared.", 3000); + } else { + showToast("Models' have been unloaded.", 3000); + } + } else { + showToast('Unloading of models failed. Installed ComfyUI may be an outdated version.', 5000); + } + } catch (error) { + console.error('[ManagerFreePath] /free error:', error); + showToast('An error occurred while trying to unload models.', 5000); + } +} +``` +Semantics: +- free_models(true) → POST /free with {"unload_models": true, "free_memory": true} +- free_models() → POST /free with {"unload_models": true} + +C) Core server endpoint (flags are set on the queue) +- File: ../../server.py +- Route: @routes.post("/free") +```py +@routes.post("/free") +async def post_free(request): + json_data = await request.json() + unload_models = json_data.get("unload_models", False) + free_memory = json_data.get("free_memory", False) + if unload_models: + self.prompt_queue.set_flag("unload_models", unload_models) + if free_memory: + self.prompt_queue.set_flag("free_memory", free_memory) + return web.Response(status=200) +``` +Semantics: +- The HTTP endpoint itself does not unload/reset; instead it sets flags on PromptServer.prompt_queue for the background worker to consume. + +D) Flag consumption and execution (definitive mechanism) +- File: ../../main.py +- Function: prompt_worker(q, server_instance) +- Excerpt (poll and handle flags, then clean up): +```py + flags = q.get_flags() + free_memory = flags.get("free_memory", False) + + if flags.get("unload_models", free_memory): + comfy.model_management.unload_all_models() + need_gc = True + last_gc_collect = 0 + + if free_memory: + e.reset() + need_gc = True + last_gc_collect = 0 + + if need_gc: + current_time = time.perf_counter() + if (current_time - last_gc_collect) > gc_collect_interval: + gc.collect() + comfy.model_management.soft_empty_cache() + last_gc_collect = current_time + need_gc = False + hook_breaker_ac10a0.restore_functions() +``` +Context: +- e is a PromptExecutor (created earlier in prompt_worker): `e = execution.PromptExecutor(server_instance, ...)` +- The worker thread is started in start_comfyui(): +```py +threading.Thread(target=prompt_worker, daemon=True, args=(prompt_server.prompt_queue, prompt_server,)).start() +``` + +Interpretation (What the Manager button actually does) +- “Free model and node cache” sets unload_models: true and free_memory: true via POST /free. +- The background prompt_worker then: + - Calls comfy.model_management.unload_all_models() + - Calls e.reset() on the PromptExecutor to drop execution caches + - Performs gc.collect() and comfy.model_management.soft_empty_cache() +- This matches the “benchmark button” behavior required for CPU memory reclamation (models fully unloaded + executor reset + allocator/cache cleanup). + +Implications for MultiGPU P1 (force_full_system_cleanup) +- To 100% replicate the benchmark button behavior from within MultiGPU code paths: + - Call comfy.model_management.unload_all_models() + - Trigger PromptExecutor.reset() on the active executor + - Follow up with gc.collect() and comfy.model_management.soft_empty_cache() +- Or, trigger the core behavior indirectly by POST /free with both flags set, relying on ComfyUI’s running prompt worker. + +Verification Status +- All file paths and snippets above were extracted from this workspace: + - Manager JS files under ../ComfyUI-Manager/js/ + - ComfyUI server and main under ../../server.py and ../../main.py +- Consumption site conclusively identified in ../../main.py prompt_worker via q.get_flags → unload_all_models + PromptExecutor.reset diff --git a/memory-bank/comfy_core.py b/memory-bank/comfy_core.py new file mode 100644 index 0000000..d20add8 --- /dev/null +++ b/memory-bank/comfy_core.py @@ -0,0 +1,4042 @@ +################## Comfy Core model_management.py #################### + +""" + This file is part of ComfyUI. + Copyright (C) 2024 Comfy + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License + along with this program. If not, see . +""" + +import psutil +import logging +from enum import Enum +from comfy.cli_args import args, PerformanceFeature +import torch +import sys +import importlib +import platform +import weakref +import gc + +class VRAMState(Enum): + DISABLED = 0 #No vram present: no need to move models to vram + NO_VRAM = 1 #Very low vram: enable all the options to save vram + LOW_VRAM = 2 + NORMAL_VRAM = 3 + HIGH_VRAM = 4 + SHARED = 5 #No dedicated vram: memory shared between CPU and GPU but models still need to be moved between both. + +class CPUState(Enum): + GPU = 0 + CPU = 1 + MPS = 2 + +# Determine VRAM State +vram_state = VRAMState.NORMAL_VRAM +set_vram_to = VRAMState.NORMAL_VRAM +cpu_state = CPUState.GPU + +total_vram = 0 + +def get_supported_float8_types(): + float8_types = [] + try: + float8_types.append(torch.float8_e4m3fn) + except: + pass + try: + float8_types.append(torch.float8_e4m3fnuz) + except: + pass + try: + float8_types.append(torch.float8_e5m2) + except: + pass + try: + float8_types.append(torch.float8_e5m2fnuz) + except: + pass + try: + float8_types.append(torch.float8_e8m0fnu) + except: + pass + return float8_types + +FLOAT8_TYPES = get_supported_float8_types() + +xpu_available = False +torch_version = "" +try: + torch_version = torch.version.__version__ + temp = torch_version.split(".") + torch_version_numeric = (int(temp[0]), int(temp[1])) +except: + pass + +lowvram_available = True +if args.deterministic: + logging.info("Using deterministic algorithms for pytorch") + torch.use_deterministic_algorithms(True, warn_only=True) + +directml_enabled = False +if args.directml is not None: + import torch_directml + directml_enabled = True + device_index = args.directml + if device_index < 0: + directml_device = torch_directml.device() + else: + directml_device = torch_directml.device(device_index) + logging.info("Using directml with device: {}".format(torch_directml.device_name(device_index))) + # torch_directml.disable_tiled_resources(True) + lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default. + +try: + import intel_extension_for_pytorch as ipex # noqa: F401 +except: + pass + +try: + _ = torch.xpu.device_count() + xpu_available = torch.xpu.is_available() +except: + xpu_available = False + +try: + if torch.backends.mps.is_available(): + cpu_state = CPUState.MPS + import torch.mps +except: + pass + +try: + import torch_npu # noqa: F401 + _ = torch.npu.device_count() + npu_available = torch.npu.is_available() +except: + npu_available = False + +try: + import torch_mlu # noqa: F401 + _ = torch.mlu.device_count() + mlu_available = torch.mlu.is_available() +except: + mlu_available = False + +try: + ixuca_available = hasattr(torch, "corex") +except: + ixuca_available = False + +if args.cpu: + cpu_state = CPUState.CPU + +def is_intel_xpu(): + global cpu_state + global xpu_available + if cpu_state == CPUState.GPU: + if xpu_available: + return True + return False + +def is_ascend_npu(): + global npu_available + if npu_available: + return True + return False + +def is_mlu(): + global mlu_available + if mlu_available: + return True + return False + +def is_ixuca(): + global ixuca_available + if ixuca_available: + return True + return False + +def get_torch_device(): + global directml_enabled + global cpu_state + if directml_enabled: + global directml_device + return directml_device + if cpu_state == CPUState.MPS: + return torch.device("mps") + if cpu_state == CPUState.CPU: + return torch.device("cpu") + else: + if is_intel_xpu(): + return torch.device("xpu", torch.xpu.current_device()) + elif is_ascend_npu(): + return torch.device("npu", torch.npu.current_device()) + elif is_mlu(): + return torch.device("mlu", torch.mlu.current_device()) + else: + return torch.device(torch.cuda.current_device()) + +def get_total_memory(dev=None, torch_total_too=False): + global directml_enabled + if dev is None: + dev = get_torch_device() + + if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'): + mem_total = psutil.virtual_memory().total + mem_total_torch = mem_total + else: + if directml_enabled: + mem_total = 1024 * 1024 * 1024 #TODO + mem_total_torch = mem_total + elif is_intel_xpu(): + stats = torch.xpu.memory_stats(dev) + mem_reserved = stats['reserved_bytes.all.current'] + mem_total_xpu = torch.xpu.get_device_properties(dev).total_memory + mem_total_torch = mem_reserved + mem_total = mem_total_xpu + elif is_ascend_npu(): + stats = torch.npu.memory_stats(dev) + mem_reserved = stats['reserved_bytes.all.current'] + _, mem_total_npu = torch.npu.mem_get_info(dev) + mem_total_torch = mem_reserved + mem_total = mem_total_npu + elif is_mlu(): + stats = torch.mlu.memory_stats(dev) + mem_reserved = stats['reserved_bytes.all.current'] + _, mem_total_mlu = torch.mlu.mem_get_info(dev) + mem_total_torch = mem_reserved + mem_total = mem_total_mlu + else: + stats = torch.cuda.memory_stats(dev) + mem_reserved = stats['reserved_bytes.all.current'] + _, mem_total_cuda = torch.cuda.mem_get_info(dev) + mem_total_torch = mem_reserved + mem_total = mem_total_cuda + + if torch_total_too: + return (mem_total, mem_total_torch) + else: + return mem_total + +def mac_version(): + try: + return tuple(int(n) for n in platform.mac_ver()[0].split(".")) + except: + return None + +total_vram = get_total_memory(get_torch_device()) / (1024 * 1024) +total_ram = psutil.virtual_memory().total / (1024 * 1024) +logging.info("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram)) + +try: + logging.info("pytorch version: {}".format(torch_version)) + mac_ver = mac_version() + if mac_ver is not None: + logging.info("Mac Version {}".format(mac_ver)) +except: + pass + +try: + OOM_EXCEPTION = torch.cuda.OutOfMemoryError +except: + OOM_EXCEPTION = Exception + +XFORMERS_VERSION = "" +XFORMERS_ENABLED_VAE = True +if args.disable_xformers: + XFORMERS_IS_AVAILABLE = False +else: + try: + import xformers + import xformers.ops + XFORMERS_IS_AVAILABLE = True + try: + XFORMERS_IS_AVAILABLE = xformers._has_cpp_library + except: + pass + try: + XFORMERS_VERSION = xformers.version.__version__ + logging.info("xformers version: {}".format(XFORMERS_VERSION)) + if XFORMERS_VERSION.startswith("0.0.18"): + logging.warning("\nWARNING: This version of xformers has a major bug where you will get black images when generating high resolution images.") + logging.warning("Please downgrade or upgrade xformers to a different version.\n") + XFORMERS_ENABLED_VAE = False + except: + pass + except: + XFORMERS_IS_AVAILABLE = False + +def is_nvidia(): + global cpu_state + if cpu_state == CPUState.GPU: + if torch.version.cuda: + return True + return False + +def is_amd(): + global cpu_state + if cpu_state == CPUState.GPU: + if torch.version.hip: + return True + return False + +def amd_min_version(device=None, min_rdna_version=0): + if not is_amd(): + return False + + if is_device_cpu(device): + return False + + arch = torch.cuda.get_device_properties(device).gcnArchName + if arch.startswith('gfx') and len(arch) == 7: + try: + cmp_rdna_version = int(arch[4]) + 2 + except: + cmp_rdna_version = 0 + if cmp_rdna_version >= min_rdna_version: + return True + + return False + +MIN_WEIGHT_MEMORY_RATIO = 0.4 +if is_nvidia(): + MIN_WEIGHT_MEMORY_RATIO = 0.0 + +ENABLE_PYTORCH_ATTENTION = False +if args.use_pytorch_cross_attention: + ENABLE_PYTORCH_ATTENTION = True + XFORMERS_IS_AVAILABLE = False + +try: + if is_nvidia(): + if torch_version_numeric[0] >= 2: + if ENABLE_PYTORCH_ATTENTION == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False: + ENABLE_PYTORCH_ATTENTION = True + if is_intel_xpu() or is_ascend_npu() or is_mlu() or is_ixuca(): + if args.use_split_cross_attention == False and args.use_quad_cross_attention == False: + ENABLE_PYTORCH_ATTENTION = True +except: + pass + + +SUPPORT_FP8_OPS = args.supports_fp8_compute +try: + if is_amd(): + try: + rocm_version = tuple(map(int, str(torch.version.hip).split(".")[:2])) + except: + rocm_version = (6, -1) + arch = torch.cuda.get_device_properties(get_torch_device()).gcnArchName + logging.info("AMD arch: {}".format(arch)) + logging.info("ROCm version: {}".format(rocm_version)) + if args.use_split_cross_attention == False and args.use_quad_cross_attention == False: + if importlib.util.find_spec('triton') is not None: # AMD efficient attention implementation depends on triton. TODO: better way of detecting if it's compiled in or not. + if torch_version_numeric >= (2, 7): # works on 2.6 but doesn't actually seem to improve much + if any((a in arch) for a in ["gfx90a", "gfx942", "gfx1100", "gfx1101", "gfx1151"]): # TODO: more arches, TODO: gfx950 + ENABLE_PYTORCH_ATTENTION = True +# if torch_version_numeric >= (2, 8): +# if any((a in arch) for a in ["gfx1201"]): +# ENABLE_PYTORCH_ATTENTION = True + if torch_version_numeric >= (2, 7) and rocm_version >= (6, 4): + if any((a in arch) for a in ["gfx1200", "gfx1201", "gfx942", "gfx950"]): # TODO: more arches + SUPPORT_FP8_OPS = True + +except: + pass + + +if ENABLE_PYTORCH_ATTENTION: + torch.backends.cuda.enable_math_sdp(True) + torch.backends.cuda.enable_flash_sdp(True) + torch.backends.cuda.enable_mem_efficient_sdp(True) + + +PRIORITIZE_FP16 = False # TODO: remove and replace with something that shows exactly which dtype is faster than the other +try: + if (is_nvidia() or is_amd()) and PerformanceFeature.Fp16Accumulation in args.fast: + torch.backends.cuda.matmul.allow_fp16_accumulation = True + PRIORITIZE_FP16 = True # TODO: limit to cards where it actually boosts performance + logging.info("Enabled fp16 accumulation.") +except: + pass + +try: + if torch_version_numeric >= (2, 5): + torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(True) +except: + logging.warning("Warning, could not set allow_fp16_bf16_reduction_math_sdp") + +if args.lowvram: + set_vram_to = VRAMState.LOW_VRAM + lowvram_available = True +elif args.novram: + set_vram_to = VRAMState.NO_VRAM +elif args.highvram or args.gpu_only: + vram_state = VRAMState.HIGH_VRAM + +FORCE_FP32 = False +if args.force_fp32: + logging.info("Forcing FP32, if this improves things please report it.") + FORCE_FP32 = True + +if lowvram_available: + if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM): + vram_state = set_vram_to + + +if cpu_state != CPUState.GPU: + vram_state = VRAMState.DISABLED + +if cpu_state == CPUState.MPS: + vram_state = VRAMState.SHARED + +logging.info(f"Set vram state to: {vram_state.name}") + +DISABLE_SMART_MEMORY = args.disable_smart_memory + +if DISABLE_SMART_MEMORY: + logging.info("Disabling smart memory management") + +def get_torch_device_name(device): + if hasattr(device, 'type'): + if device.type == "cuda": + try: + allocator_backend = torch.cuda.get_allocator_backend() + except: + allocator_backend = "" + return "{} {} : {}".format(device, torch.cuda.get_device_name(device), allocator_backend) + elif device.type == "xpu": + return "{} {}".format(device, torch.xpu.get_device_name(device)) + else: + return "{}".format(device.type) + elif is_intel_xpu(): + return "{} {}".format(device, torch.xpu.get_device_name(device)) + elif is_ascend_npu(): + return "{} {}".format(device, torch.npu.get_device_name(device)) + elif is_mlu(): + return "{} {}".format(device, torch.mlu.get_device_name(device)) + else: + return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device)) + +try: + logging.info("Device: {}".format(get_torch_device_name(get_torch_device()))) +except: + logging.warning("Could not pick default device.") + + +current_loaded_models = [] + +def module_size(module): + module_mem = 0 + sd = module.state_dict() + for k in sd: + t = sd[k] + module_mem += t.nelement() * t.element_size() + return module_mem + +class LoadedModel: + def __init__(self, model): + self._set_model(model) + self.device = model.load_device + self.real_model = None + self.currently_used = True + self.model_finalizer = None + self._patcher_finalizer = None + + def _set_model(self, model): + self._model = weakref.ref(model) + if model.parent is not None: + self._parent_model = weakref.ref(model.parent) + self._patcher_finalizer = weakref.finalize(model, self._switch_parent) + + def _switch_parent(self): + model = self._parent_model() + if model is not None: + self._set_model(model) + + @property + def model(self): + return self._model() + + def model_memory(self): + return self.model.model_size() + + def model_loaded_memory(self): + return self.model.loaded_size() + + def model_offloaded_memory(self): + return self.model.model_size() - self.model.loaded_size() + + def model_memory_required(self, device): + if device == self.model.current_loaded_device(): + return self.model_offloaded_memory() + else: + return self.model_memory() + + def model_load(self, lowvram_model_memory=0, force_patch_weights=False): + self.model.model_patches_to(self.device) + self.model.model_patches_to(self.model.model_dtype()) + + # if self.model.loaded_size() > 0: + use_more_vram = lowvram_model_memory + if use_more_vram == 0: + use_more_vram = 1e32 + self.model_use_more_vram(use_more_vram, force_patch_weights=force_patch_weights) + real_model = self.model.model + + if is_intel_xpu() and not args.disable_ipex_optimize and 'ipex' in globals() and real_model is not None: + with torch.no_grad(): + real_model = ipex.optimize(real_model.eval(), inplace=True, graph_mode=True, concat_linear=True) + + self.real_model = weakref.ref(real_model) + self.model_finalizer = weakref.finalize(real_model, cleanup_models) + return real_model + + def should_reload_model(self, force_patch_weights=False): + if force_patch_weights and self.model.lowvram_patch_counter() > 0: + return True + return False + + def model_unload(self, memory_to_free=None, unpatch_weights=True): + if memory_to_free is not None: + if memory_to_free < self.model.loaded_size(): + freed = self.model.partially_unload(self.model.offload_device, memory_to_free) + if freed >= memory_to_free: + return False + self.model.detach(unpatch_weights) + self.model_finalizer.detach() + self.model_finalizer = None + self.real_model = None + return True + + def model_use_more_vram(self, extra_memory, force_patch_weights=False): + return self.model.partially_load(self.device, extra_memory, force_patch_weights=force_patch_weights) + + def __eq__(self, other): + return self.model is other.model + + def __del__(self): + if self._patcher_finalizer is not None: + self._patcher_finalizer.detach() + + def is_dead(self): + return self.real_model() is not None and self.model is None + + +def use_more_memory(extra_memory, loaded_models, device): + for m in loaded_models: + if m.device == device: + extra_memory -= m.model_use_more_vram(extra_memory) + if extra_memory <= 0: + break + +def offloaded_memory(loaded_models, device): + offloaded_mem = 0 + for m in loaded_models: + if m.device == device: + offloaded_mem += m.model_offloaded_memory() + return offloaded_mem + +WINDOWS = any(platform.win32_ver()) + +EXTRA_RESERVED_VRAM = 400 * 1024 * 1024 +if WINDOWS: + EXTRA_RESERVED_VRAM = 600 * 1024 * 1024 #Windows is higher because of the shared vram issue + if total_vram > (15 * 1024): # more extra reserved vram on 16GB+ cards + EXTRA_RESERVED_VRAM += 100 * 1024 * 1024 + +if args.reserve_vram is not None: + EXTRA_RESERVED_VRAM = args.reserve_vram * 1024 * 1024 * 1024 + logging.debug("Reserving {}MB vram for other applications.".format(EXTRA_RESERVED_VRAM / (1024 * 1024))) + +def extra_reserved_memory(): + return EXTRA_RESERVED_VRAM + +def minimum_inference_memory(): + return (1024 * 1024 * 1024) * 0.8 + extra_reserved_memory() + +def free_memory(memory_required, device, keep_loaded=[]): + cleanup_models_gc() + unloaded_model = [] + can_unload = [] + unloaded_models = [] + + for i in range(len(current_loaded_models) -1, -1, -1): + shift_model = current_loaded_models[i] + if shift_model.device == device: + if shift_model not in keep_loaded and not shift_model.is_dead(): + can_unload.append((-shift_model.model_offloaded_memory(), sys.getrefcount(shift_model.model), shift_model.model_memory(), i)) + shift_model.currently_used = False + + for x in sorted(can_unload): + i = x[-1] + memory_to_free = None + if not DISABLE_SMART_MEMORY: + free_mem = get_free_memory(device) + if free_mem > memory_required: + break + memory_to_free = memory_required - free_mem + logging.debug(f"Unloading {current_loaded_models[i].model.model.__class__.__name__}") + if current_loaded_models[i].model_unload(memory_to_free): + unloaded_model.append(i) + + for i in sorted(unloaded_model, reverse=True): + unloaded_models.append(current_loaded_models.pop(i)) + + if len(unloaded_model) > 0: + soft_empty_cache() + else: + if vram_state != VRAMState.HIGH_VRAM: + mem_free_total, mem_free_torch = get_free_memory(device, torch_free_too=True) + if mem_free_torch > mem_free_total * 0.25: + soft_empty_cache() + return unloaded_models + +def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None, force_full_load=False): + cleanup_models_gc() + global vram_state + + inference_memory = minimum_inference_memory() + extra_mem = max(inference_memory, memory_required + extra_reserved_memory()) + if minimum_memory_required is None: + minimum_memory_required = extra_mem + else: + minimum_memory_required = max(inference_memory, minimum_memory_required + extra_reserved_memory()) + + models_temp = set() + for m in models: + models_temp.add(m) + for mm in m.model_patches_models(): + models_temp.add(mm) + + models = models_temp + + models_to_load = [] + + for x in models: + loaded_model = LoadedModel(x) + try: + loaded_model_index = current_loaded_models.index(loaded_model) + except: + loaded_model_index = None + + if loaded_model_index is not None: + loaded = current_loaded_models[loaded_model_index] + loaded.currently_used = True + models_to_load.append(loaded) + else: + if hasattr(x, "model"): + logging.info(f"Requested to load {x.model.__class__.__name__}") + models_to_load.append(loaded_model) + + for loaded_model in models_to_load: + to_unload = [] + for i in range(len(current_loaded_models)): + if loaded_model.model.is_clone(current_loaded_models[i].model): + to_unload = [i] + to_unload + for i in to_unload: + model_to_unload = current_loaded_models.pop(i) + model_to_unload.model.detach(unpatch_all=False) + model_to_unload.model_finalizer.detach() + + total_memory_required = {} + for loaded_model in models_to_load: + total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device) + + for device in total_memory_required: + if device != torch.device("cpu"): + free_memory(total_memory_required[device] * 1.1 + extra_mem, device) + + for device in total_memory_required: + if device != torch.device("cpu"): + free_mem = get_free_memory(device) + if free_mem < minimum_memory_required: + models_l = free_memory(minimum_memory_required, device) + logging.info("{} models unloaded.".format(len(models_l))) + + for loaded_model in models_to_load: + model = loaded_model.model + torch_dev = model.load_device + if is_device_cpu(torch_dev): + vram_set_state = VRAMState.DISABLED + else: + vram_set_state = vram_state + lowvram_model_memory = 0 + if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM) and not force_full_load: + loaded_memory = loaded_model.model_loaded_memory() + current_free_mem = get_free_memory(torch_dev) + loaded_memory + + lowvram_model_memory = max(128 * 1024 * 1024, (current_free_mem - minimum_memory_required), min(current_free_mem * MIN_WEIGHT_MEMORY_RATIO, current_free_mem - minimum_inference_memory())) + lowvram_model_memory = max(0.1, lowvram_model_memory - loaded_memory) + + if vram_set_state == VRAMState.NO_VRAM: + lowvram_model_memory = 0.1 + + loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights) + current_loaded_models.insert(0, loaded_model) + return + +def load_model_gpu(model): + return load_models_gpu([model]) + +def loaded_models(only_currently_used=False): + output = [] + for m in current_loaded_models: + if only_currently_used: + if not m.currently_used: + continue + + output.append(m.model) + return output + + +def cleanup_models_gc(): + do_gc = False + for i in range(len(current_loaded_models)): + cur = current_loaded_models[i] + if cur.is_dead(): + logging.info("Potential memory leak detected with model {}, doing a full garbage collect, for maximum performance avoid circular references in the model code.".format(cur.real_model().__class__.__name__)) + do_gc = True + break + + if do_gc: + gc.collect() + soft_empty_cache() + + for i in range(len(current_loaded_models)): + cur = current_loaded_models[i] + if cur.is_dead(): + logging.warning("WARNING, memory leak with model {}. Please make sure it is not being referenced from somewhere.".format(cur.real_model().__class__.__name__)) + + + +def cleanup_models(): + to_delete = [] + for i in range(len(current_loaded_models)): + if current_loaded_models[i].real_model() is None: + to_delete = [i] + to_delete + + for i in to_delete: + x = current_loaded_models.pop(i) + del x + +def dtype_size(dtype): + dtype_size = 4 + if dtype == torch.float16 or dtype == torch.bfloat16: + dtype_size = 2 + elif dtype == torch.float32: + dtype_size = 4 + else: + try: + dtype_size = dtype.itemsize + except: #Old pytorch doesn't have .itemsize + pass + return dtype_size + +def unet_offload_device(): + if vram_state == VRAMState.HIGH_VRAM: + return get_torch_device() + else: + return torch.device("cpu") + +def unet_inital_load_device(parameters, dtype): + torch_dev = get_torch_device() + if vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.SHARED: + return torch_dev + + cpu_dev = torch.device("cpu") + if DISABLE_SMART_MEMORY or vram_state == VRAMState.NO_VRAM: + return cpu_dev + + model_size = dtype_size(dtype) * parameters + + mem_dev = get_free_memory(torch_dev) + mem_cpu = get_free_memory(cpu_dev) + if mem_dev > mem_cpu and model_size < mem_dev: + return torch_dev + else: + return cpu_dev + +def maximum_vram_for_weights(device=None): + return (get_total_memory(device) * 0.88 - minimum_inference_memory()) + +def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32], weight_dtype=None): + if model_params < 0: + model_params = 1000000000000000000000 + if args.fp32_unet: + return torch.float32 + if args.fp64_unet: + return torch.float64 + if args.bf16_unet: + return torch.bfloat16 + if args.fp16_unet: + return torch.float16 + if args.fp8_e4m3fn_unet: + return torch.float8_e4m3fn + if args.fp8_e5m2_unet: + return torch.float8_e5m2 + if args.fp8_e8m0fnu_unet: + return torch.float8_e8m0fnu + + fp8_dtype = None + if weight_dtype in FLOAT8_TYPES: + fp8_dtype = weight_dtype + + if fp8_dtype is not None: + if supports_fp8_compute(device): #if fp8 compute is supported the casting is most likely not expensive + return fp8_dtype + + free_model_memory = maximum_vram_for_weights(device) + if model_params * 2 > free_model_memory: + return fp8_dtype + + if PRIORITIZE_FP16 or weight_dtype == torch.float16: + if torch.float16 in supported_dtypes and should_use_fp16(device=device, model_params=model_params): + return torch.float16 + + for dt in supported_dtypes: + if dt == torch.float16 and should_use_fp16(device=device, model_params=model_params): + if torch.float16 in supported_dtypes: + return torch.float16 + if dt == torch.bfloat16 and should_use_bf16(device, model_params=model_params): + if torch.bfloat16 in supported_dtypes: + return torch.bfloat16 + + for dt in supported_dtypes: + if dt == torch.float16 and should_use_fp16(device=device, model_params=model_params, manual_cast=True): + if torch.float16 in supported_dtypes: + return torch.float16 + if dt == torch.bfloat16 and should_use_bf16(device, model_params=model_params, manual_cast=True): + if torch.bfloat16 in supported_dtypes: + return torch.bfloat16 + + return torch.float32 + +# None means no manual cast +def unet_manual_cast(weight_dtype, inference_device, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]): + if weight_dtype == torch.float32 or weight_dtype == torch.float64: + return None + + fp16_supported = should_use_fp16(inference_device, prioritize_performance=False) + if fp16_supported and weight_dtype == torch.float16: + return None + + bf16_supported = should_use_bf16(inference_device) + if bf16_supported and weight_dtype == torch.bfloat16: + return None + + fp16_supported = should_use_fp16(inference_device, prioritize_performance=True) + if PRIORITIZE_FP16 and fp16_supported and torch.float16 in supported_dtypes: + return torch.float16 + + for dt in supported_dtypes: + if dt == torch.float16 and fp16_supported: + return torch.float16 + if dt == torch.bfloat16 and bf16_supported: + return torch.bfloat16 + + return torch.float32 + +def text_encoder_offload_device(): + if args.gpu_only: + return get_torch_device() + else: + return torch.device("cpu") + +def text_encoder_device(): + if args.gpu_only: + return get_torch_device() + elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM: + if should_use_fp16(prioritize_performance=False): + return get_torch_device() + else: + return torch.device("cpu") + else: + return torch.device("cpu") + +def text_encoder_initial_device(load_device, offload_device, model_size=0): + if load_device == offload_device or model_size <= 1024 * 1024 * 1024: + return offload_device + + if is_device_mps(load_device): + return load_device + + mem_l = get_free_memory(load_device) + mem_o = get_free_memory(offload_device) + if mem_l > (mem_o * 0.5) and model_size * 1.2 < mem_l: + return load_device + else: + return offload_device + +def text_encoder_dtype(device=None): + if args.fp8_e4m3fn_text_enc: + return torch.float8_e4m3fn + elif args.fp8_e5m2_text_enc: + return torch.float8_e5m2 + elif args.fp16_text_enc: + return torch.float16 + elif args.bf16_text_enc: + return torch.bfloat16 + elif args.fp32_text_enc: + return torch.float32 + + if is_device_cpu(device): + return torch.float16 + + return torch.float16 + + +def intermediate_device(): + if args.gpu_only: + return get_torch_device() + else: + return torch.device("cpu") + +def vae_device(): + if args.cpu_vae: + return torch.device("cpu") + return get_torch_device() + +def vae_offload_device(): + if args.gpu_only: + return get_torch_device() + else: + return torch.device("cpu") + +def vae_dtype(device=None, allowed_dtypes=[]): + if args.fp16_vae: + return torch.float16 + elif args.bf16_vae: + return torch.bfloat16 + elif args.fp32_vae: + return torch.float32 + + for d in allowed_dtypes: + if d == torch.float16 and should_use_fp16(device): + return d + + # NOTE: bfloat16 seems to work on AMD for the VAE but is extremely slow in some cases compared to fp32 + # slowness still a problem on pytorch nightly 2.9.0.dev20250720+rocm6.4 tested on RDNA3 + # also a problem on RDNA4 except fp32 is also slow there. + # This is due to large bf16 convolutions being extremely slow. + if d == torch.bfloat16 and ((not is_amd()) or amd_min_version(device, min_rdna_version=4)) and should_use_bf16(device): + return d + + return torch.float32 + +def get_autocast_device(dev): + if hasattr(dev, 'type'): + return dev.type + return "cuda" + +def supports_dtype(device, dtype): #TODO + if dtype == torch.float32: + return True + if is_device_cpu(device): + return False + if dtype == torch.float16: + return True + if dtype == torch.bfloat16: + return True + return False + +def supports_cast(device, dtype): #TODO + if dtype == torch.float32: + return True + if dtype == torch.float16: + return True + if directml_enabled: #TODO: test this + return False + if dtype == torch.bfloat16: + return True + if is_device_mps(device): + return False + if dtype == torch.float8_e4m3fn: + return True + if dtype == torch.float8_e5m2: + return True + return False + +def pick_weight_dtype(dtype, fallback_dtype, device=None): + if dtype is None: + dtype = fallback_dtype + elif dtype_size(dtype) > dtype_size(fallback_dtype): + dtype = fallback_dtype + + if not supports_cast(device, dtype): + dtype = fallback_dtype + + return dtype + +def device_supports_non_blocking(device): + if args.force_non_blocking: + return True + if is_device_mps(device): + return False #pytorch bug? mps doesn't support non blocking + if is_intel_xpu(): #xpu does support non blocking but it is slower on iGPUs for some reason so disable by default until situation changes + return False + if args.deterministic: #TODO: figure out why deterministic breaks non blocking from gpu to cpu (previews) + return False + if directml_enabled: + return False + return True + +def device_should_use_non_blocking(device): + if not device_supports_non_blocking(device): + return False + return False + # return True #TODO: figure out why this causes memory issues on Nvidia and possibly others + +def force_channels_last(): + if args.force_channels_last: + return True + + #TODO + return False + + +STREAMS = {} +NUM_STREAMS = 1 +if args.async_offload: + NUM_STREAMS = 2 + logging.info("Using async weight offloading with {} streams".format(NUM_STREAMS)) + +stream_counters = {} +def get_offload_stream(device): + stream_counter = stream_counters.get(device, 0) + if NUM_STREAMS <= 1: + return None + + if device in STREAMS: + ss = STREAMS[device] + s = ss[stream_counter] + stream_counter = (stream_counter + 1) % len(ss) + if is_device_cuda(device): + ss[stream_counter].wait_stream(torch.cuda.current_stream()) + elif is_device_xpu(device): + ss[stream_counter].wait_stream(torch.xpu.current_stream()) + stream_counters[device] = stream_counter + return s + elif is_device_cuda(device): + ss = [] + for k in range(NUM_STREAMS): + ss.append(torch.cuda.Stream(device=device, priority=0)) + STREAMS[device] = ss + s = ss[stream_counter] + stream_counter = (stream_counter + 1) % len(ss) + stream_counters[device] = stream_counter + return s + elif is_device_xpu(device): + ss = [] + for k in range(NUM_STREAMS): + ss.append(torch.xpu.Stream(device=device, priority=0)) + STREAMS[device] = ss + s = ss[stream_counter] + stream_counter = (stream_counter + 1) % len(ss) + stream_counters[device] = stream_counter + return s + return None + +def sync_stream(device, stream): + if stream is None: + return + if is_device_cuda(device): + torch.cuda.current_stream().wait_stream(stream) + elif is_device_xpu(device): + torch.xpu.current_stream().wait_stream(stream) + +def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False, stream=None): + if device is None or weight.device == device: + if not copy: + if dtype is None or weight.dtype == dtype: + return weight + if stream is not None: + with stream: + return weight.to(dtype=dtype, copy=copy) + return weight.to(dtype=dtype, copy=copy) + + if stream is not None: + with stream: + r = torch.empty_like(weight, dtype=dtype, device=device) + r.copy_(weight, non_blocking=non_blocking) + else: + r = torch.empty_like(weight, dtype=dtype, device=device) + r.copy_(weight, non_blocking=non_blocking) + return r + +def cast_to_device(tensor, device, dtype, copy=False): + non_blocking = device_supports_non_blocking(device) + return cast_to(tensor, dtype=dtype, device=device, non_blocking=non_blocking, copy=copy) + +def sage_attention_enabled(): + return args.use_sage_attention + +def flash_attention_enabled(): + return args.use_flash_attention + +def xformers_enabled(): + global directml_enabled + global cpu_state + if cpu_state != CPUState.GPU: + return False + if is_intel_xpu(): + return False + if is_ascend_npu(): + return False + if is_mlu(): + return False + if is_ixuca(): + return False + if directml_enabled: + return False + return XFORMERS_IS_AVAILABLE + + +def xformers_enabled_vae(): + enabled = xformers_enabled() + if not enabled: + return False + + return XFORMERS_ENABLED_VAE + +def pytorch_attention_enabled(): + global ENABLE_PYTORCH_ATTENTION + return ENABLE_PYTORCH_ATTENTION + +def pytorch_attention_enabled_vae(): + if is_amd(): + return False # enabling pytorch attention on AMD currently causes crash when doing high res + return pytorch_attention_enabled() + +def pytorch_attention_flash_attention(): + global ENABLE_PYTORCH_ATTENTION + if ENABLE_PYTORCH_ATTENTION: + #TODO: more reliable way of checking for flash attention? + if is_nvidia(): + return True + if is_intel_xpu(): + return True + if is_ascend_npu(): + return True + if is_mlu(): + return True + if is_amd(): + return True #if you have pytorch attention enabled on AMD it probably supports at least mem efficient attention + if is_ixuca(): + return True + return False + +def force_upcast_attention_dtype(): + upcast = args.force_upcast_attention + + macos_version = mac_version() + if macos_version is not None and ((14, 5) <= macos_version): # black image bug on recent versions of macOS, I don't think it's ever getting fixed + upcast = True + + if upcast: + return {torch.float16: torch.float32} + else: + return None + +def get_free_memory(dev=None, torch_free_too=False): + global directml_enabled + if dev is None: + dev = get_torch_device() + + if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'): + mem_free_total = psutil.virtual_memory().available + mem_free_torch = mem_free_total + else: + if directml_enabled: + mem_free_total = 1024 * 1024 * 1024 #TODO + mem_free_torch = mem_free_total + elif is_intel_xpu(): + stats = torch.xpu.memory_stats(dev) + mem_active = stats['active_bytes.all.current'] + mem_reserved = stats['reserved_bytes.all.current'] + mem_free_xpu = torch.xpu.get_device_properties(dev).total_memory - mem_reserved + mem_free_torch = mem_reserved - mem_active + mem_free_total = mem_free_xpu + mem_free_torch + elif is_ascend_npu(): + stats = torch.npu.memory_stats(dev) + mem_active = stats['active_bytes.all.current'] + mem_reserved = stats['reserved_bytes.all.current'] + mem_free_npu, _ = torch.npu.mem_get_info(dev) + mem_free_torch = mem_reserved - mem_active + mem_free_total = mem_free_npu + mem_free_torch + elif is_mlu(): + stats = torch.mlu.memory_stats(dev) + mem_active = stats['active_bytes.all.current'] + mem_reserved = stats['reserved_bytes.all.current'] + mem_free_mlu, _ = torch.mlu.mem_get_info(dev) + mem_free_torch = mem_reserved - mem_active + mem_free_total = mem_free_mlu + mem_free_torch + else: + stats = torch.cuda.memory_stats(dev) + mem_active = stats['active_bytes.all.current'] + mem_reserved = stats['reserved_bytes.all.current'] + mem_free_cuda, _ = torch.cuda.mem_get_info(dev) + mem_free_torch = mem_reserved - mem_active + mem_free_total = mem_free_cuda + mem_free_torch + + if torch_free_too: + return (mem_free_total, mem_free_torch) + else: + return mem_free_total + +def cpu_mode(): + global cpu_state + return cpu_state == CPUState.CPU + +def mps_mode(): + global cpu_state + return cpu_state == CPUState.MPS + +def is_device_type(device, type): + if hasattr(device, 'type'): + if (device.type == type): + return True + return False + +def is_device_cpu(device): + return is_device_type(device, 'cpu') + +def is_device_mps(device): + return is_device_type(device, 'mps') + +def is_device_xpu(device): + return is_device_type(device, 'xpu') + +def is_device_cuda(device): + return is_device_type(device, 'cuda') + +def is_directml_enabled(): + global directml_enabled + if directml_enabled: + return True + + return False + +def should_use_fp16(device=None, model_params=0, prioritize_performance=True, manual_cast=False): + if device is not None: + if is_device_cpu(device): + return False + + if args.force_fp16: + return True + + if FORCE_FP32: + return False + + if is_directml_enabled(): + return True + + if (device is not None and is_device_mps(device)) or mps_mode(): + return True + + if cpu_mode(): + return False + + if is_intel_xpu(): + if torch_version_numeric < (2, 3): + return True + else: + return torch.xpu.get_device_properties(device).has_fp16 + + if is_ascend_npu(): + return True + + if is_mlu(): + return True + + if is_ixuca(): + return True + + if torch.version.hip: + return True + + props = torch.cuda.get_device_properties(device) + if props.major >= 8: + return True + + if props.major < 6: + return False + + #FP16 is confirmed working on a 1080 (GP104) and on latest pytorch actually seems faster than fp32 + nvidia_10_series = ["1080", "1070", "titan x", "p3000", "p3200", "p4000", "p4200", "p5000", "p5200", "p6000", "1060", "1050", "p40", "p100", "p6", "p4"] + for x in nvidia_10_series: + if x in props.name.lower(): + if WINDOWS or manual_cast: + return True + else: + return False #weird linux behavior where fp32 is faster + + if manual_cast: + free_model_memory = maximum_vram_for_weights(device) + if (not prioritize_performance) or model_params * 4 > free_model_memory: + return True + + if props.major < 7: + return False + + #FP16 is just broken on these cards + nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600", "MX550", "MX450", "CMP 30HX", "T2000", "T1000", "T1200"] + for x in nvidia_16_series: + if x in props.name: + return False + + return True + +def should_use_bf16(device=None, model_params=0, prioritize_performance=True, manual_cast=False): + if device is not None: + if is_device_cpu(device): #TODO ? bf16 works on CPU but is extremely slow + return False + + if FORCE_FP32: + return False + + if directml_enabled: + return False + + if (device is not None and is_device_mps(device)) or mps_mode(): + if mac_version() < (14,): + return False + return True + + if cpu_mode(): + return False + + if is_intel_xpu(): + if torch_version_numeric < (2, 3): + return True + else: + return torch.xpu.is_bf16_supported() + + if is_ascend_npu(): + return True + + if is_ixuca(): + return True + + if is_amd(): + arch = torch.cuda.get_device_properties(device).gcnArchName + if any((a in arch) for a in ["gfx1030", "gfx1031", "gfx1010", "gfx1011", "gfx1012", "gfx906", "gfx900", "gfx803"]): # RDNA2 and older don't support bf16 + if manual_cast: + return True + return False + + props = torch.cuda.get_device_properties(device) + + if is_mlu(): + if props.major > 3: + return True + + if props.major >= 8: + return True + + bf16_works = torch.cuda.is_bf16_supported() + + if bf16_works and manual_cast: + free_model_memory = maximum_vram_for_weights(device) + if (not prioritize_performance) or model_params * 4 > free_model_memory: + return True + + return False + +def supports_fp8_compute(device=None): + if SUPPORT_FP8_OPS: + return True + + if not is_nvidia(): + return False + + props = torch.cuda.get_device_properties(device) + if props.major >= 9: + return True + if props.major < 8: + return False + if props.minor < 9: + return False + + if torch_version_numeric < (2, 3): + return False + + if WINDOWS: + if torch_version_numeric < (2, 4): + return False + + return True + +def extended_fp16_support(): + # TODO: check why some models work with fp16 on newer torch versions but not on older + if torch_version_numeric < (2, 7): + return False + + return True + +def soft_empty_cache(force=False): + global cpu_state + if cpu_state == CPUState.MPS: + torch.mps.empty_cache() + elif is_intel_xpu(): + torch.xpu.empty_cache() + elif is_ascend_npu(): + torch.npu.empty_cache() + elif is_mlu(): + torch.mlu.empty_cache() + elif torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + +def unload_all_models(): + free_memory(1e30, get_torch_device()) + + +#TODO: might be cleaner to put this somewhere else +import threading + +class InterruptProcessingException(Exception): + pass + +interrupt_processing_mutex = threading.RLock() + +interrupt_processing = False +def interrupt_current_processing(value=True): + global interrupt_processing + global interrupt_processing_mutex + with interrupt_processing_mutex: + interrupt_processing = value + +def processing_interrupted(): + global interrupt_processing + global interrupt_processing_mutex + with interrupt_processing_mutex: + return interrupt_processing + +def throw_exception_if_processing_interrupted(): + global interrupt_processing + global interrupt_processing_mutex + with interrupt_processing_mutex: + if interrupt_processing: + interrupt_processing = False + raise InterruptProcessingException() + + +################### ComfyUI Core server.py #################### +import os +import sys +import asyncio +import traceback + +import nodes +import folder_paths +import execution +import uuid +import urllib +import json +import glob +import struct +import ssl +import socket +import ipaddress +from PIL import Image, ImageOps +from PIL.PngImagePlugin import PngInfo +from io import BytesIO + +import aiohttp +from aiohttp import web +import logging + +import mimetypes +from comfy.cli_args import args +import comfy.utils +import comfy.model_management +from comfy_api import feature_flags +import node_helpers +from comfyui_version import __version__ +from app.frontend_management import FrontendManager +from comfy_api.internal import _ComfyNodeInternal + +from app.user_manager import UserManager +from app.model_manager import ModelFileManager +from app.custom_node_manager import CustomNodeManager +from typing import Optional, Union +from api_server.routes.internal.internal_routes import InternalRoutes +from protocol import BinaryEventTypes + +# Import cache control middleware +from middleware.cache_middleware import cache_control + +async def send_socket_catch_exception(function, message): + try: + await function(message) + except (aiohttp.ClientError, aiohttp.ClientPayloadError, ConnectionResetError, BrokenPipeError, ConnectionError) as err: + logging.warning("send error: {}".format(err)) + +@web.middleware +async def compress_body(request: web.Request, handler): + accept_encoding = request.headers.get("Accept-Encoding", "") + response: web.Response = await handler(request) + if not isinstance(response, web.Response): + return response + if response.content_type not in ["application/json", "text/plain"]: + return response + if response.body and "gzip" in accept_encoding: + response.enable_compression() + return response + + +def create_cors_middleware(allowed_origin: str): + @web.middleware + async def cors_middleware(request: web.Request, handler): + if request.method == "OPTIONS": + # Pre-flight request. Reply successfully: + response = web.Response() + else: + response = await handler(request) + + response.headers['Access-Control-Allow-Origin'] = allowed_origin + response.headers['Access-Control-Allow-Methods'] = 'POST, GET, DELETE, PUT, OPTIONS' + response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization' + response.headers['Access-Control-Allow-Credentials'] = 'true' + return response + + return cors_middleware + +def is_loopback(host): + if host is None: + return False + try: + if ipaddress.ip_address(host).is_loopback: + return True + else: + return False + except: + pass + + loopback = False + for family in (socket.AF_INET, socket.AF_INET6): + try: + r = socket.getaddrinfo(host, None, family, socket.SOCK_STREAM) + for family, _, _, _, sockaddr in r: + if not ipaddress.ip_address(sockaddr[0]).is_loopback: + return loopback + else: + loopback = True + except socket.gaierror: + pass + + return loopback + + +def create_origin_only_middleware(): + @web.middleware + async def origin_only_middleware(request: web.Request, handler): + #this code is used to prevent the case where a random website can queue comfy workflows by making a POST to 127.0.0.1 which browsers don't prevent for some dumb reason. + #in that case the Host and Origin hostnames won't match + #I know the proper fix would be to add a cookie but this should take care of the problem in the meantime + if 'Host' in request.headers and 'Origin' in request.headers: + host = request.headers['Host'] + origin = request.headers['Origin'] + host_domain = host.lower() + parsed = urllib.parse.urlparse(origin) + origin_domain = parsed.netloc.lower() + host_domain_parsed = urllib.parse.urlsplit('//' + host_domain) + + #limit the check to when the host domain is localhost, this makes it slightly less safe but should still prevent the exploit + loopback = is_loopback(host_domain_parsed.hostname) + + if parsed.port is None: #if origin doesn't have a port strip it from the host to handle weird browsers, same for host + host_domain = host_domain_parsed.hostname + if host_domain_parsed.port is None: + origin_domain = parsed.hostname + + if loopback and host_domain is not None and origin_domain is not None and len(host_domain) > 0 and len(origin_domain) > 0: + if host_domain != origin_domain: + logging.warning("WARNING: request with non matching host and origin {} != {}, returning 403".format(host_domain, origin_domain)) + return web.Response(status=403) + + if request.method == "OPTIONS": + response = web.Response() + else: + response = await handler(request) + + return response + + return origin_only_middleware + +class PromptServer(): + def __init__(self, loop): + PromptServer.instance = self + + mimetypes.init() + mimetypes.add_type('application/javascript; charset=utf-8', '.js') + mimetypes.add_type('image/webp', '.webp') + + self.user_manager = UserManager() + self.model_file_manager = ModelFileManager() + self.custom_node_manager = CustomNodeManager() + self.internal_routes = InternalRoutes(self) + self.supports = ["custom_nodes_from_web"] + self.prompt_queue = execution.PromptQueue(self) + self.loop = loop + self.messages = asyncio.Queue() + self.client_session:Optional[aiohttp.ClientSession] = None + self.number = 0 + + middlewares = [cache_control] + if args.enable_compress_response_body: + middlewares.append(compress_body) + + if args.enable_cors_header: + middlewares.append(create_cors_middleware(args.enable_cors_header)) + else: + middlewares.append(create_origin_only_middleware()) + + max_upload_size = round(args.max_upload_size * 1024 * 1024) + self.app = web.Application(client_max_size=max_upload_size, middlewares=middlewares) + self.sockets = dict() + self.sockets_metadata = dict() + self.web_root = ( + FrontendManager.init_frontend(args.front_end_version) + if args.front_end_root is None + else args.front_end_root + ) + logging.info(f"[Prompt Server] web root: {self.web_root}") + routes = web.RouteTableDef() + self.routes = routes + self.last_node_id = None + self.client_id = None + + self.on_prompt_handlers = [] + + @routes.get('/ws') + async def websocket_handler(request): + ws = web.WebSocketResponse() + await ws.prepare(request) + sid = request.rel_url.query.get('clientId', '') + if sid: + # Reusing existing session, remove old + self.sockets.pop(sid, None) + else: + sid = uuid.uuid4().hex + + # Store WebSocket for backward compatibility + self.sockets[sid] = ws + # Store metadata separately + self.sockets_metadata[sid] = {"feature_flags": {}} + + try: + # Send initial state to the new client + await self.send("status", {"status": self.get_queue_info(), "sid": sid}, sid) + # On reconnect if we are the currently executing client send the current node + if self.client_id == sid and self.last_node_id is not None: + await self.send("executing", { "node": self.last_node_id }, sid) + + # Flag to track if we've received the first message + first_message = True + + async for msg in ws: + if msg.type == aiohttp.WSMsgType.ERROR: + logging.warning('ws connection closed with exception %s' % ws.exception()) + elif msg.type == aiohttp.WSMsgType.TEXT: + try: + data = json.loads(msg.data) + # Check if first message is feature flags + if first_message and data.get("type") == "feature_flags": + # Store client feature flags + client_flags = data.get("data", {}) + self.sockets_metadata[sid]["feature_flags"] = client_flags + + # Send server feature flags in response + await self.send( + "feature_flags", + feature_flags.get_server_features(), + sid, + ) + + logging.debug( + f"Feature flags negotiated for client {sid}: {client_flags}" + ) + first_message = False + except json.JSONDecodeError: + logging.warning( + f"Invalid JSON received from client {sid}: {msg.data}" + ) + except Exception as e: + logging.error(f"Error processing WebSocket message: {e}") + finally: + self.sockets.pop(sid, None) + self.sockets_metadata.pop(sid, None) + return ws + + @routes.get("/") + async def get_root(request): + response = web.FileResponse(os.path.join(self.web_root, "index.html")) + response.headers['Cache-Control'] = 'no-cache' + response.headers["Pragma"] = "no-cache" + response.headers["Expires"] = "0" + return response + + @routes.get("/embeddings") + def get_embeddings(request): + embeddings = folder_paths.get_filename_list("embeddings") + return web.json_response(list(map(lambda a: os.path.splitext(a)[0], embeddings))) + + @routes.get("/models") + def list_model_types(request): + model_types = list(folder_paths.folder_names_and_paths.keys()) + + return web.json_response(model_types) + + @routes.get("/models/{folder}") + async def get_models(request): + folder = request.match_info.get("folder", None) + if not folder in folder_paths.folder_names_and_paths: + return web.Response(status=404) + files = folder_paths.get_filename_list(folder) + return web.json_response(files) + + @routes.get("/extensions") + async def get_extensions(request): + files = glob.glob(os.path.join( + glob.escape(self.web_root), 'extensions/**/*.js'), recursive=True) + + extensions = list(map(lambda f: "/" + os.path.relpath(f, self.web_root).replace("\\", "/"), files)) + + for name, dir in nodes.EXTENSION_WEB_DIRS.items(): + files = glob.glob(os.path.join(glob.escape(dir), '**/*.js'), recursive=True) + extensions.extend(list(map(lambda f: "/extensions/" + urllib.parse.quote( + name) + "/" + os.path.relpath(f, dir).replace("\\", "/"), files))) + + return web.json_response(extensions) + + def get_dir_by_type(dir_type): + if dir_type is None: + dir_type = "input" + + if dir_type == "input": + type_dir = folder_paths.get_input_directory() + elif dir_type == "temp": + type_dir = folder_paths.get_temp_directory() + elif dir_type == "output": + type_dir = folder_paths.get_output_directory() + + return type_dir, dir_type + + def compare_image_hash(filepath, image): + hasher = node_helpers.hasher() + + # function to compare hashes of two images to see if it already exists, fix to #3465 + if os.path.exists(filepath): + a = hasher() + b = hasher() + with open(filepath, "rb") as f: + a.update(f.read()) + b.update(image.file.read()) + image.file.seek(0) + return a.hexdigest() == b.hexdigest() + return False + + def image_upload(post, image_save_function=None): + image = post.get("image") + overwrite = post.get("overwrite") + image_is_duplicate = False + + image_upload_type = post.get("type") + upload_dir, image_upload_type = get_dir_by_type(image_upload_type) + + if image and image.file: + filename = image.filename + if not filename: + return web.Response(status=400) + + subfolder = post.get("subfolder", "") + full_output_folder = os.path.join(upload_dir, os.path.normpath(subfolder)) + filepath = os.path.abspath(os.path.join(full_output_folder, filename)) + + if os.path.commonpath((upload_dir, filepath)) != upload_dir: + return web.Response(status=400) + + if not os.path.exists(full_output_folder): + os.makedirs(full_output_folder) + + split = os.path.splitext(filename) + + if overwrite is not None and (overwrite == "true" or overwrite == "1"): + pass + else: + i = 1 + while os.path.exists(filepath): + if compare_image_hash(filepath, image): #compare hash to prevent saving of duplicates with same name, fix for #3465 + image_is_duplicate = True + break + filename = f"{split[0]} ({i}){split[1]}" + filepath = os.path.join(full_output_folder, filename) + i += 1 + + if not image_is_duplicate: + if image_save_function is not None: + image_save_function(image, post, filepath) + else: + with open(filepath, "wb") as f: + f.write(image.file.read()) + + return web.json_response({"name" : filename, "subfolder": subfolder, "type": image_upload_type}) + else: + return web.Response(status=400) + + @routes.post("/upload/image") + async def upload_image(request): + post = await request.post() + return image_upload(post) + + + @routes.post("/upload/mask") + async def upload_mask(request): + post = await request.post() + + def image_save_function(image, post, filepath): + original_ref = json.loads(post.get("original_ref")) + filename, output_dir = folder_paths.annotated_filepath(original_ref['filename']) + + if not filename: + return web.Response(status=400) + + # validation for security: prevent accessing arbitrary path + if filename[0] == '/' or '..' in filename: + return web.Response(status=400) + + if output_dir is None: + type = original_ref.get("type", "output") + output_dir = folder_paths.get_directory_by_type(type) + + if output_dir is None: + return web.Response(status=400) + + if original_ref.get("subfolder", "") != "": + full_output_dir = os.path.join(output_dir, original_ref["subfolder"]) + if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir: + return web.Response(status=403) + output_dir = full_output_dir + + file = os.path.join(output_dir, filename) + + if os.path.isfile(file): + with Image.open(file) as original_pil: + metadata = PngInfo() + if hasattr(original_pil,'text'): + for key in original_pil.text: + metadata.add_text(key, original_pil.text[key]) + original_pil = original_pil.convert('RGBA') + mask_pil = Image.open(image.file).convert('RGBA') + + # alpha copy + new_alpha = mask_pil.getchannel('A') + original_pil.putalpha(new_alpha) + original_pil.save(filepath, compress_level=4, pnginfo=metadata) + + return image_upload(post, image_save_function) + + @routes.get("/view") + async def view_image(request): + if "filename" in request.rel_url.query: + filename = request.rel_url.query["filename"] + filename, output_dir = folder_paths.annotated_filepath(filename) + + if not filename: + return web.Response(status=400) + + # validation for security: prevent accessing arbitrary path + if filename[0] == '/' or '..' in filename: + return web.Response(status=400) + + if output_dir is None: + type = request.rel_url.query.get("type", "output") + output_dir = folder_paths.get_directory_by_type(type) + + if output_dir is None: + return web.Response(status=400) + + if "subfolder" in request.rel_url.query: + full_output_dir = os.path.join(output_dir, request.rel_url.query["subfolder"]) + if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir: + return web.Response(status=403) + output_dir = full_output_dir + + filename = os.path.basename(filename) + file = os.path.join(output_dir, filename) + + if os.path.isfile(file): + if 'preview' in request.rel_url.query: + with Image.open(file) as img: + preview_info = request.rel_url.query['preview'].split(';') + image_format = preview_info[0] + if image_format not in ['webp', 'jpeg'] or 'a' in request.rel_url.query.get('channel', ''): + image_format = 'webp' + + quality = 90 + if preview_info[-1].isdigit(): + quality = int(preview_info[-1]) + + buffer = BytesIO() + if image_format in ['jpeg'] or request.rel_url.query.get('channel', '') == 'rgb': + img = img.convert("RGB") + img.save(buffer, format=image_format, quality=quality) + buffer.seek(0) + + return web.Response(body=buffer.read(), content_type=f'image/{image_format}', + headers={"Content-Disposition": f"filename=\"{filename}\""}) + + if 'channel' not in request.rel_url.query: + channel = 'rgba' + else: + channel = request.rel_url.query["channel"] + + if channel == 'rgb': + with Image.open(file) as img: + if img.mode == "RGBA": + r, g, b, a = img.split() + new_img = Image.merge('RGB', (r, g, b)) + else: + new_img = img.convert("RGB") + + buffer = BytesIO() + new_img.save(buffer, format='PNG') + buffer.seek(0) + + return web.Response(body=buffer.read(), content_type='image/png', + headers={"Content-Disposition": f"filename=\"{filename}\""}) + + elif channel == 'a': + with Image.open(file) as img: + if img.mode == "RGBA": + _, _, _, a = img.split() + else: + a = Image.new('L', img.size, 255) + + # alpha img + alpha_img = Image.new('RGBA', img.size) + alpha_img.putalpha(a) + alpha_buffer = BytesIO() + alpha_img.save(alpha_buffer, format='PNG') + alpha_buffer.seek(0) + + return web.Response(body=alpha_buffer.read(), content_type='image/png', + headers={"Content-Disposition": f"filename=\"{filename}\""}) + else: + # Get content type from mimetype, defaulting to 'application/octet-stream' + content_type = mimetypes.guess_type(filename)[0] or 'application/octet-stream' + + # For security, force certain mimetypes to download instead of display + if content_type in {'text/html', 'text/html-sandboxed', 'application/xhtml+xml', 'text/javascript', 'text/css'}: + content_type = 'application/octet-stream' # Forces download + + return web.FileResponse( + file, + headers={ + "Content-Disposition": f"filename=\"{filename}\"", + "Content-Type": content_type + } + ) + + return web.Response(status=404) + + @routes.get("/view_metadata/{folder_name}") + async def view_metadata(request): + folder_name = request.match_info.get("folder_name", None) + if folder_name is None: + return web.Response(status=404) + if not "filename" in request.rel_url.query: + return web.Response(status=404) + + filename = request.rel_url.query["filename"] + if not filename.endswith(".safetensors"): + return web.Response(status=404) + + safetensors_path = folder_paths.get_full_path(folder_name, filename) + if safetensors_path is None: + return web.Response(status=404) + out = comfy.utils.safetensors_header(safetensors_path, max_size=1024*1024) + if out is None: + return web.Response(status=404) + dt = json.loads(out) + if not "__metadata__" in dt: + return web.Response(status=404) + return web.json_response(dt["__metadata__"]) + + @routes.get("/system_stats") + async def system_stats(request): + device = comfy.model_management.get_torch_device() + device_name = comfy.model_management.get_torch_device_name(device) + cpu_device = comfy.model_management.torch.device("cpu") + ram_total = comfy.model_management.get_total_memory(cpu_device) + ram_free = comfy.model_management.get_free_memory(cpu_device) + vram_total, torch_vram_total = comfy.model_management.get_total_memory(device, torch_total_too=True) + vram_free, torch_vram_free = comfy.model_management.get_free_memory(device, torch_free_too=True) + required_frontend_version = FrontendManager.get_required_frontend_version() + + system_stats = { + "system": { + "os": os.name, + "ram_total": ram_total, + "ram_free": ram_free, + "comfyui_version": __version__, + "required_frontend_version": required_frontend_version, + "python_version": sys.version, + "pytorch_version": comfy.model_management.torch_version, + "embedded_python": os.path.split(os.path.split(sys.executable)[0])[1] == "python_embeded", + "argv": sys.argv + }, + "devices": [ + { + "name": device_name, + "type": device.type, + "index": device.index, + "vram_total": vram_total, + "vram_free": vram_free, + "torch_vram_total": torch_vram_total, + "torch_vram_free": torch_vram_free, + } + ] + } + return web.json_response(system_stats) + + @routes.get("/features") + async def get_features(request): + return web.json_response(feature_flags.get_server_features()) + + @routes.get("/prompt") + async def get_prompt(request): + return web.json_response(self.get_queue_info()) + + def node_info(node_class): + obj_class = nodes.NODE_CLASS_MAPPINGS[node_class] + if issubclass(obj_class, _ComfyNodeInternal): + return obj_class.GET_NODE_INFO_V1() + info = {} + info['input'] = obj_class.INPUT_TYPES() + info['input_order'] = {key: list(value.keys()) for (key, value) in obj_class.INPUT_TYPES().items()} + info['output'] = obj_class.RETURN_TYPES + info['output_is_list'] = obj_class.OUTPUT_IS_LIST if hasattr(obj_class, 'OUTPUT_IS_LIST') else [False] * len(obj_class.RETURN_TYPES) + info['output_name'] = obj_class.RETURN_NAMES if hasattr(obj_class, 'RETURN_NAMES') else info['output'] + info['name'] = node_class + info['display_name'] = nodes.NODE_DISPLAY_NAME_MAPPINGS[node_class] if node_class in nodes.NODE_DISPLAY_NAME_MAPPINGS.keys() else node_class + info['description'] = obj_class.DESCRIPTION if hasattr(obj_class,'DESCRIPTION') else '' + info['python_module'] = getattr(obj_class, "RELATIVE_PYTHON_MODULE", "nodes") + info['category'] = 'sd' + if hasattr(obj_class, 'OUTPUT_NODE') and obj_class.OUTPUT_NODE == True: + info['output_node'] = True + else: + info['output_node'] = False + + if hasattr(obj_class, 'CATEGORY'): + info['category'] = obj_class.CATEGORY + + if hasattr(obj_class, 'OUTPUT_TOOLTIPS'): + info['output_tooltips'] = obj_class.OUTPUT_TOOLTIPS + + if getattr(obj_class, "DEPRECATED", False): + info['deprecated'] = True + if getattr(obj_class, "EXPERIMENTAL", False): + info['experimental'] = True + + if hasattr(obj_class, 'API_NODE'): + info['api_node'] = obj_class.API_NODE + return info + + @routes.get("/object_info") + async def get_object_info(request): + with folder_paths.cache_helper: + out = {} + for x in nodes.NODE_CLASS_MAPPINGS: + try: + out[x] = node_info(x) + except Exception: + logging.error(f"[ERROR] An error occurred while retrieving information for the '{x}' node.") + logging.error(traceback.format_exc()) + return web.json_response(out) + + @routes.get("/object_info/{node_class}") + async def get_object_info_node(request): + node_class = request.match_info.get("node_class", None) + out = {} + if (node_class is not None) and (node_class in nodes.NODE_CLASS_MAPPINGS): + out[node_class] = node_info(node_class) + return web.json_response(out) + + @routes.get("/history") + async def get_history(request): + max_items = request.rel_url.query.get("max_items", None) + if max_items is not None: + max_items = int(max_items) + + offset = request.rel_url.query.get("offset", None) + if offset is not None: + offset = int(offset) + else: + offset = -1 + + return web.json_response(self.prompt_queue.get_history(max_items=max_items, offset=offset)) + + @routes.get("/history/{prompt_id}") + async def get_history_prompt_id(request): + prompt_id = request.match_info.get("prompt_id", None) + return web.json_response(self.prompt_queue.get_history(prompt_id=prompt_id)) + + @routes.get("/queue") + async def get_queue(request): + queue_info = {} + current_queue = self.prompt_queue.get_current_queue_volatile() + queue_info['queue_running'] = current_queue[0] + queue_info['queue_pending'] = current_queue[1] + return web.json_response(queue_info) + + @routes.post("/prompt") + async def post_prompt(request): + logging.info("got prompt") + json_data = await request.json() + json_data = self.trigger_on_prompt(json_data) + + if "number" in json_data: + number = float(json_data['number']) + else: + number = self.number + if "front" in json_data: + if json_data['front']: + number = -number + + self.number += 1 + + if "prompt" in json_data: + prompt = json_data["prompt"] + prompt_id = str(json_data.get("prompt_id", uuid.uuid4())) + + partial_execution_targets = None + if "partial_execution_targets" in json_data: + partial_execution_targets = json_data["partial_execution_targets"] + + valid = await execution.validate_prompt(prompt_id, prompt, partial_execution_targets) + extra_data = {} + if "extra_data" in json_data: + extra_data = json_data["extra_data"] + + if "client_id" in json_data: + extra_data["client_id"] = json_data["client_id"] + if valid[0]: + outputs_to_execute = valid[2] + self.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute)) + response = {"prompt_id": prompt_id, "number": number, "node_errors": valid[3]} + return web.json_response(response) + else: + logging.warning("invalid prompt: {}".format(valid[1])) + return web.json_response({"error": valid[1], "node_errors": valid[3]}, status=400) + else: + error = { + "type": "no_prompt", + "message": "No prompt provided", + "details": "No prompt provided", + "extra_info": {} + } + return web.json_response({"error": error, "node_errors": {}}, status=400) + + @routes.post("/queue") + async def post_queue(request): + json_data = await request.json() + if "clear" in json_data: + if json_data["clear"]: + self.prompt_queue.wipe_queue() + if "delete" in json_data: + to_delete = json_data['delete'] + for id_to_delete in to_delete: + delete_func = lambda a: a[1] == id_to_delete + self.prompt_queue.delete_queue_item(delete_func) + + return web.Response(status=200) + + @routes.post("/interrupt") + async def post_interrupt(request): + try: + json_data = await request.json() + except json.JSONDecodeError: + json_data = {} + + # Check if a specific prompt_id was provided for targeted interruption + prompt_id = json_data.get('prompt_id') + if prompt_id: + currently_running, _ = self.prompt_queue.get_current_queue() + + # Check if the prompt_id matches any currently running prompt + should_interrupt = False + for item in currently_running: + # item structure: (number, prompt_id, prompt, extra_data, outputs_to_execute) + if item[1] == prompt_id: + logging.info(f"Interrupting prompt {prompt_id}") + should_interrupt = True + break + + if should_interrupt: + nodes.interrupt_processing() + else: + logging.info(f"Prompt {prompt_id} is not currently running, skipping interrupt") + else: + # No prompt_id provided, do a global interrupt + logging.info("Global interrupt (no prompt_id specified)") + nodes.interrupt_processing() + + return web.Response(status=200) + + @routes.post("/free") + async def post_free(request): + json_data = await request.json() + unload_models = json_data.get("unload_models", False) + free_memory = json_data.get("free_memory", False) + if unload_models: + self.prompt_queue.set_flag("unload_models", unload_models) + if free_memory: + self.prompt_queue.set_flag("free_memory", free_memory) + return web.Response(status=200) + + @routes.post("/history") + async def post_history(request): + json_data = await request.json() + if "clear" in json_data: + if json_data["clear"]: + self.prompt_queue.wipe_history() + if "delete" in json_data: + to_delete = json_data['delete'] + for id_to_delete in to_delete: + self.prompt_queue.delete_history_item(id_to_delete) + + return web.Response(status=200) + + async def setup(self): + timeout = aiohttp.ClientTimeout(total=None) # no timeout + self.client_session = aiohttp.ClientSession(timeout=timeout) + + def add_routes(self): + self.user_manager.add_routes(self.routes) + self.model_file_manager.add_routes(self.routes) + self.custom_node_manager.add_routes(self.routes, self.app, nodes.LOADED_MODULE_DIRS.items()) + self.app.add_subapp('/internal', self.internal_routes.get_app()) + + # Prefix every route with /api for easier matching for delegation. + # This is very useful for frontend dev server, which need to forward + # everything except serving of static files. + # Currently both the old endpoints without prefix and new endpoints with + # prefix are supported. + api_routes = web.RouteTableDef() + for route in self.routes: + # Custom nodes might add extra static routes. Only process non-static + # routes to add /api prefix. + if isinstance(route, web.RouteDef): + api_routes.route(route.method, "/api" + route.path)(route.handler, **route.kwargs) + self.app.add_routes(api_routes) + self.app.add_routes(self.routes) + + # Add routes from web extensions. + for name, dir in nodes.EXTENSION_WEB_DIRS.items(): + self.app.add_routes([web.static('/extensions/' + name, dir)]) + + workflow_templates_path = FrontendManager.templates_path() + if workflow_templates_path: + self.app.add_routes([ + web.static('/templates', workflow_templates_path) + ]) + + # Serve embedded documentation from the package + embedded_docs_path = FrontendManager.embedded_docs_path() + if embedded_docs_path: + self.app.add_routes([ + web.static('/docs', embedded_docs_path) + ]) + + self.app.add_routes([ + web.static('/', self.web_root), + ]) + + def get_queue_info(self): + prompt_info = {} + exec_info = {} + exec_info['queue_remaining'] = self.prompt_queue.get_tasks_remaining() + prompt_info['exec_info'] = exec_info + return prompt_info + + async def send(self, event, data, sid=None): + if event == BinaryEventTypes.UNENCODED_PREVIEW_IMAGE: + await self.send_image(data, sid=sid) + elif event == BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA: + # data is (preview_image, metadata) + preview_image, metadata = data + await self.send_image_with_metadata(preview_image, metadata, sid=sid) + elif isinstance(data, (bytes, bytearray)): + await self.send_bytes(event, data, sid) + else: + await self.send_json(event, data, sid) + + def encode_bytes(self, event, data): + if not isinstance(event, int): + raise RuntimeError(f"Binary event types must be integers, got {event}") + + packed = struct.pack(">I", event) + message = bytearray(packed) + message.extend(data) + return message + + async def send_image(self, image_data, sid=None): + image_type = image_data[0] + image = image_data[1] + max_size = image_data[2] + if max_size is not None: + if hasattr(Image, 'Resampling'): + resampling = Image.Resampling.BILINEAR + else: + resampling = Image.Resampling.LANCZOS + + image = ImageOps.contain(image, (max_size, max_size), resampling) + type_num = 1 + if image_type == "JPEG": + type_num = 1 + elif image_type == "PNG": + type_num = 2 + + bytesIO = BytesIO() + header = struct.pack(">I", type_num) + bytesIO.write(header) + image.save(bytesIO, format=image_type, quality=95, compress_level=1) + preview_bytes = bytesIO.getvalue() + await self.send_bytes(BinaryEventTypes.PREVIEW_IMAGE, preview_bytes, sid=sid) + + async def send_image_with_metadata(self, image_data, metadata=None, sid=None): + image_type = image_data[0] + image = image_data[1] + max_size = image_data[2] + if max_size is not None: + if hasattr(Image, 'Resampling'): + resampling = Image.Resampling.BILINEAR + else: + resampling = Image.Resampling.LANCZOS + + image = ImageOps.contain(image, (max_size, max_size), resampling) + + mimetype = "image/png" if image_type == "PNG" else "image/jpeg" + + # Prepare metadata + if metadata is None: + metadata = {} + metadata["image_type"] = mimetype + + # Serialize metadata as JSON + import json + metadata_json = json.dumps(metadata).encode('utf-8') + metadata_length = len(metadata_json) + + # Prepare image data + bytesIO = BytesIO() + image.save(bytesIO, format=image_type, quality=95, compress_level=1) + image_bytes = bytesIO.getvalue() + + # Combine metadata and image + combined_data = bytearray() + combined_data.extend(struct.pack(">I", metadata_length)) + combined_data.extend(metadata_json) + combined_data.extend(image_bytes) + + await self.send_bytes(BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA, combined_data, sid=sid) + + async def send_bytes(self, event, data, sid=None): + message = self.encode_bytes(event, data) + + if sid is None: + sockets = list(self.sockets.values()) + for ws in sockets: + await send_socket_catch_exception(ws.send_bytes, message) + elif sid in self.sockets: + await send_socket_catch_exception(self.sockets[sid].send_bytes, message) + + async def send_json(self, event, data, sid=None): + message = {"type": event, "data": data} + + if sid is None: + sockets = list(self.sockets.values()) + for ws in sockets: + await send_socket_catch_exception(ws.send_json, message) + elif sid in self.sockets: + await send_socket_catch_exception(self.sockets[sid].send_json, message) + + def send_sync(self, event, data, sid=None): + self.loop.call_soon_threadsafe( + self.messages.put_nowait, (event, data, sid)) + + def queue_updated(self): + self.send_sync("status", { "status": self.get_queue_info() }) + + async def publish_loop(self): + while True: + msg = await self.messages.get() + await self.send(*msg) + + async def start(self, address, port, verbose=True, call_on_start=None): + await self.start_multi_address([(address, port)], call_on_start=call_on_start) + + async def start_multi_address(self, addresses, call_on_start=None, verbose=True): + runner = web.AppRunner(self.app, access_log=None) + await runner.setup() + ssl_ctx = None + scheme = "http" + if args.tls_keyfile and args.tls_certfile: + ssl_ctx = ssl.SSLContext(protocol=ssl.PROTOCOL_TLS_SERVER, verify_mode=ssl.CERT_NONE) + ssl_ctx.load_cert_chain(certfile=args.tls_certfile, + keyfile=args.tls_keyfile) + scheme = "https" + + if verbose: + logging.info("Starting server\n") + for addr in addresses: + address = addr[0] + port = addr[1] + site = web.TCPSite(runner, address, port, ssl_context=ssl_ctx) + await site.start() + + if not hasattr(self, 'address'): + self.address = address #TODO: remove this + self.port = port + + if ':' in address: + address_print = "[{}]".format(address) + else: + address_print = address + + if verbose: + logging.info("To see the GUI go to: {}://{}:{}".format(scheme, address_print, port)) + + if call_on_start is not None: + call_on_start(scheme, self.address, self.port) + + def add_on_prompt_handler(self, handler): + self.on_prompt_handlers.append(handler) + + def trigger_on_prompt(self, json_data): + for handler in self.on_prompt_handlers: + try: + json_data = handler(json_data) + except Exception: + logging.warning("[ERROR] An error occurred during the on_prompt_handler processing") + logging.warning(traceback.format_exc()) + + return json_data + + def send_progress_text( + self, text: Union[bytes, bytearray, str], node_id: str, sid=None + ): + if isinstance(text, str): + text = text.encode("utf-8") + node_id_bytes = str(node_id).encode("utf-8") + + # Pack the node_id length as a 4-byte unsigned integer, followed by the node_id bytes + message = struct.pack(">I", len(node_id_bytes)) + node_id_bytes + text + + self.send_sync(BinaryEventTypes.TEXT, message, sid) +########################## Comfy Core main.py ########################## +import comfy.options +comfy.options.enable_args_parsing() + +import os +import importlib.util +import folder_paths +import time +from comfy.cli_args import args +from app.logger import setup_logger +import itertools +import utils.extra_config +import logging +import sys +from comfy_execution.progress import get_progress_state +from comfy_execution.utils import get_executing_context +from comfy_api import feature_flags + +if __name__ == "__main__": + #NOTE: These do not do anything on core ComfyUI, they are for custom nodes. + os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1' + os.environ['DO_NOT_TRACK'] = '1' + +setup_logger(log_level=args.verbose, use_stdout=args.log_stdout) + +def apply_custom_paths(): + # extra model paths + extra_model_paths_config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml") + if os.path.isfile(extra_model_paths_config_path): + utils.extra_config.load_extra_path_config(extra_model_paths_config_path) + + if args.extra_model_paths_config: + for config_path in itertools.chain(*args.extra_model_paths_config): + utils.extra_config.load_extra_path_config(config_path) + + # --output-directory, --input-directory, --user-directory + if args.output_directory: + output_dir = os.path.abspath(args.output_directory) + logging.info(f"Setting output directory to: {output_dir}") + folder_paths.set_output_directory(output_dir) + + # These are the default folders that checkpoints, clip and vae models will be saved to when using CheckpointSave, etc.. nodes + folder_paths.add_model_folder_path("checkpoints", os.path.join(folder_paths.get_output_directory(), "checkpoints")) + folder_paths.add_model_folder_path("clip", os.path.join(folder_paths.get_output_directory(), "clip")) + folder_paths.add_model_folder_path("vae", os.path.join(folder_paths.get_output_directory(), "vae")) + folder_paths.add_model_folder_path("diffusion_models", + os.path.join(folder_paths.get_output_directory(), "diffusion_models")) + folder_paths.add_model_folder_path("loras", os.path.join(folder_paths.get_output_directory(), "loras")) + + if args.input_directory: + input_dir = os.path.abspath(args.input_directory) + logging.info(f"Setting input directory to: {input_dir}") + folder_paths.set_input_directory(input_dir) + + if args.user_directory: + user_dir = os.path.abspath(args.user_directory) + logging.info(f"Setting user directory to: {user_dir}") + folder_paths.set_user_directory(user_dir) + + +def execute_prestartup_script(): + if args.disable_all_custom_nodes and len(args.whitelist_custom_nodes) == 0: + return + + def execute_script(script_path): + module_name = os.path.splitext(script_path)[0] + try: + spec = importlib.util.spec_from_file_location(module_name, script_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return True + except Exception as e: + logging.error(f"Failed to execute startup-script: {script_path} / {e}") + return False + + node_paths = folder_paths.get_folder_paths("custom_nodes") + for custom_node_path in node_paths: + possible_modules = os.listdir(custom_node_path) + node_prestartup_times = [] + + for possible_module in possible_modules: + module_path = os.path.join(custom_node_path, possible_module) + if os.path.isfile(module_path) or module_path.endswith(".disabled") or module_path == "__pycache__": + continue + + script_path = os.path.join(module_path, "prestartup_script.py") + if os.path.exists(script_path): + if args.disable_all_custom_nodes and possible_module not in args.whitelist_custom_nodes: + logging.info(f"Prestartup Skipping {possible_module} due to disable_all_custom_nodes and whitelist_custom_nodes") + continue + time_before = time.perf_counter() + success = execute_script(script_path) + node_prestartup_times.append((time.perf_counter() - time_before, module_path, success)) + if len(node_prestartup_times) > 0: + logging.info("\nPrestartup times for custom nodes:") + for n in sorted(node_prestartup_times): + if n[2]: + import_message = "" + else: + import_message = " (PRESTARTUP FAILED)" + logging.info("{:6.1f} seconds{}: {}".format(n[0], import_message, n[1])) + logging.info("") + +apply_custom_paths() +execute_prestartup_script() + + +# Main code +import asyncio +import shutil +import threading +import gc + + +if os.name == "nt": + os.environ['MIMALLOC_PURGE_DELAY'] = '0' + +if __name__ == "__main__": + if args.default_device is not None: + default_dev = args.default_device + devices = list(range(32)) + devices.remove(default_dev) + devices.insert(0, default_dev) + devices = ','.join(map(str, devices)) + os.environ['CUDA_VISIBLE_DEVICES'] = str(devices) + os.environ['HIP_VISIBLE_DEVICES'] = str(devices) + + if args.cuda_device is not None: + os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device) + os.environ['HIP_VISIBLE_DEVICES'] = str(args.cuda_device) + logging.info("Set cuda device to: {}".format(args.cuda_device)) + + if args.oneapi_device_selector is not None: + os.environ['ONEAPI_DEVICE_SELECTOR'] = args.oneapi_device_selector + logging.info("Set oneapi device selector to: {}".format(args.oneapi_device_selector)) + + if args.deterministic: + if 'CUBLAS_WORKSPACE_CONFIG' not in os.environ: + os.environ['CUBLAS_WORKSPACE_CONFIG'] = ":4096:8" + + import cuda_malloc + +if 'torch' in sys.modules: + logging.warning("WARNING: Potential Error in code: Torch already imported, torch should never be imported before this point.") + +import comfy.utils + +import execution +import server +from protocol import BinaryEventTypes +import nodes +import comfy.model_management +import comfyui_version +import app.logger +import hook_breaker_ac10a0 + +def cuda_malloc_warning(): + device = comfy.model_management.get_torch_device() + device_name = comfy.model_management.get_torch_device_name(device) + cuda_malloc_warning = False + if "cudaMallocAsync" in device_name: + for b in cuda_malloc.blacklist: + if b in device_name: + cuda_malloc_warning = True + if cuda_malloc_warning: + logging.warning("\nWARNING: this card most likely does not support cuda-malloc, if you get \"CUDA error\" please run ComfyUI with: --disable-cuda-malloc\n") + + +def prompt_worker(q, server_instance): + current_time: float = 0.0 + cache_type = execution.CacheType.CLASSIC + if args.cache_lru > 0: + cache_type = execution.CacheType.LRU + elif args.cache_none: + cache_type = execution.CacheType.DEPENDENCY_AWARE + + e = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_size=args.cache_lru) + last_gc_collect = 0 + need_gc = False + gc_collect_interval = 10.0 + + while True: + timeout = 1000.0 + if need_gc: + timeout = max(gc_collect_interval - (current_time - last_gc_collect), 0.0) + + queue_item = q.get(timeout=timeout) + if queue_item is not None: + item, item_id = queue_item + execution_start_time = time.perf_counter() + prompt_id = item[1] + server_instance.last_prompt_id = prompt_id + + e.execute(item[2], prompt_id, item[3], item[4]) + need_gc = True + q.task_done(item_id, + e.history_result, + status=execution.PromptQueue.ExecutionStatus( + status_str='success' if e.success else 'error', + completed=e.success, + messages=e.status_messages)) + if server_instance.client_id is not None: + server_instance.send_sync("executing", {"node": None, "prompt_id": prompt_id}, server_instance.client_id) + + current_time = time.perf_counter() + execution_time = current_time - execution_start_time + + # Log Time in a more readable way after 10 minutes + if execution_time > 600: + execution_time = time.strftime("%H:%M:%S", time.gmtime(execution_time)) + logging.info(f"Prompt executed in {execution_time}") + else: + logging.info("Prompt executed in {:.2f} seconds".format(execution_time)) + + flags = q.get_flags() + free_memory = flags.get("free_memory", False) + + if flags.get("unload_models", free_memory): + comfy.model_management.unload_all_models() + need_gc = True + last_gc_collect = 0 + + if free_memory: + e.reset() + need_gc = True + last_gc_collect = 0 + + if need_gc: + current_time = time.perf_counter() + if (current_time - last_gc_collect) > gc_collect_interval: + gc.collect() + comfy.model_management.soft_empty_cache() + last_gc_collect = current_time + need_gc = False + hook_breaker_ac10a0.restore_functions() + + +async def run(server_instance, address='', port=8188, verbose=True, call_on_start=None): + addresses = [] + for addr in address.split(","): + addresses.append((addr, port)) + await asyncio.gather( + server_instance.start_multi_address(addresses, call_on_start, verbose), server_instance.publish_loop() + ) + +def hijack_progress(server_instance): + def hook(value, total, preview_image, prompt_id=None, node_id=None): + executing_context = get_executing_context() + if prompt_id is None and executing_context is not None: + prompt_id = executing_context.prompt_id + if node_id is None and executing_context is not None: + node_id = executing_context.node_id + comfy.model_management.throw_exception_if_processing_interrupted() + if prompt_id is None: + prompt_id = server_instance.last_prompt_id + if node_id is None: + node_id = server_instance.last_node_id + progress = {"value": value, "max": total, "prompt_id": prompt_id, "node": node_id} + get_progress_state().update_progress(node_id, value, total, preview_image) + + server_instance.send_sync("progress", progress, server_instance.client_id) + if preview_image is not None: + # Only send old method if client doesn't support preview metadata + if not feature_flags.supports_feature( + server_instance.sockets_metadata, + server_instance.client_id, + "supports_preview_metadata", + ): + server_instance.send_sync( + BinaryEventTypes.UNENCODED_PREVIEW_IMAGE, + preview_image, + server_instance.client_id, + ) + + comfy.utils.set_progress_bar_global_hook(hook) + + +def cleanup_temp(): + temp_dir = folder_paths.get_temp_directory() + if os.path.exists(temp_dir): + shutil.rmtree(temp_dir, ignore_errors=True) + + +def setup_database(): + try: + from app.database.db import init_db, dependencies_available + if dependencies_available(): + init_db() + except Exception as e: + logging.error(f"Failed to initialize database. Please ensure you have installed the latest requirements. If the error persists, please report this as in future the database will be required: {e}") + + +def start_comfyui(asyncio_loop=None): + """ + Starts the ComfyUI server using the provided asyncio event loop or creates a new one. + Returns the event loop, server instance, and a function to start the server asynchronously. + """ + if args.temp_directory: + temp_dir = os.path.join(os.path.abspath(args.temp_directory), "temp") + logging.info(f"Setting temp directory to: {temp_dir}") + folder_paths.set_temp_directory(temp_dir) + cleanup_temp() + + if args.windows_standalone_build: + try: + import new_updater + new_updater.update_windows_updater() + except: + pass + + if not asyncio_loop: + asyncio_loop = asyncio.new_event_loop() + asyncio.set_event_loop(asyncio_loop) + prompt_server = server.PromptServer(asyncio_loop) + + hook_breaker_ac10a0.save_functions() + asyncio_loop.run_until_complete(nodes.init_extra_nodes( + init_custom_nodes=(not args.disable_all_custom_nodes) or len(args.whitelist_custom_nodes) > 0, + init_api_nodes=not args.disable_api_nodes + )) + hook_breaker_ac10a0.restore_functions() + + cuda_malloc_warning() + setup_database() + + prompt_server.add_routes() + hijack_progress(prompt_server) + + threading.Thread(target=prompt_worker, daemon=True, args=(prompt_server.prompt_queue, prompt_server,)).start() + + if args.quick_test_for_ci: + exit(0) + + os.makedirs(folder_paths.get_temp_directory(), exist_ok=True) + call_on_start = None + if args.auto_launch: + def startup_server(scheme, address, port): + import webbrowser + if os.name == 'nt' and address == '0.0.0.0': + address = '127.0.0.1' + if ':' in address: + address = "[{}]".format(address) + webbrowser.open(f"{scheme}://{address}:{port}") + call_on_start = startup_server + + async def start_all(): + await prompt_server.setup() + await run(prompt_server, address=args.listen, port=args.port, verbose=not args.dont_print_server, call_on_start=call_on_start) + + # Returning these so that other code can integrate with the ComfyUI loop and server + return asyncio_loop, prompt_server, start_all + + +if __name__ == "__main__": + # Running directly, just start ComfyUI. + logging.info("Python version: {}".format(sys.version)) + logging.info("ComfyUI version: {}".format(comfyui_version.__version__)) + + if sys.version_info.major == 3 and sys.version_info.minor < 10: + logging.warning("WARNING: You are using a python version older than 3.10, please upgrade to a newer one. 3.12 and above is recommended.") + + event_loop, _, start_all_func = start_comfyui() + try: + x = start_all_func() + app.logger.print_startup_warnings() + event_loop.run_until_complete(x) + except KeyboardInterrupt: + logging.info("\nStopped server") + + cleanup_temp() +###################### Comfy Core execution.py ###################### +import copy +import heapq +import inspect +import logging +import sys +import threading +import time +import traceback +from enum import Enum +from typing import List, Literal, NamedTuple, Optional, Union +import asyncio + +import torch + +import comfy.model_management +import nodes +from comfy_execution.caching import ( + BasicCache, + CacheKeySetID, + CacheKeySetInputSignature, + DependencyAwareCache, + HierarchicalCache, + LRUCache, +) +from comfy_execution.graph import ( + DynamicPrompt, + ExecutionBlocker, + ExecutionList, + get_input_info, +) +from comfy_execution.graph_utils import GraphBuilder, is_link +from comfy_execution.validation import validate_node_input +from comfy_execution.progress import get_progress_state, reset_progress_state, add_progress_handler, WebUIProgressHandler +from comfy_execution.utils import CurrentNodeContext +from comfy_api.internal import _ComfyNodeInternal, _NodeOutputInternal, first_real_override, is_class, make_locked_method_func +from comfy_api.latest import io + + +class ExecutionResult(Enum): + SUCCESS = 0 + FAILURE = 1 + PENDING = 2 + +class DuplicateNodeError(Exception): + pass + +class IsChangedCache: + def __init__(self, prompt_id: str, dynprompt: DynamicPrompt, outputs_cache: BasicCache): + self.prompt_id = prompt_id + self.dynprompt = dynprompt + self.outputs_cache = outputs_cache + self.is_changed = {} + + async def get(self, node_id): + if node_id in self.is_changed: + return self.is_changed[node_id] + + node = self.dynprompt.get_node(node_id) + class_type = node["class_type"] + class_def = nodes.NODE_CLASS_MAPPINGS[class_type] + has_is_changed = False + is_changed_name = None + if issubclass(class_def, _ComfyNodeInternal) and first_real_override(class_def, "fingerprint_inputs") is not None: + has_is_changed = True + is_changed_name = "fingerprint_inputs" + elif hasattr(class_def, "IS_CHANGED"): + has_is_changed = True + is_changed_name = "IS_CHANGED" + if not has_is_changed: + self.is_changed[node_id] = False + return self.is_changed[node_id] + + if "is_changed" in node: + self.is_changed[node_id] = node["is_changed"] + return self.is_changed[node_id] + + # Intentionally do not use cached outputs here. We only want constants in IS_CHANGED + input_data_all, _, hidden_inputs = get_input_data(node["inputs"], class_def, node_id, None) + try: + is_changed = await _async_map_node_over_list(self.prompt_id, node_id, class_def, input_data_all, is_changed_name) + is_changed = await resolve_map_node_over_list_results(is_changed) + node["is_changed"] = [None if isinstance(x, ExecutionBlocker) else x for x in is_changed] + except Exception as e: + logging.warning("WARNING: {}".format(e)) + node["is_changed"] = float("NaN") + finally: + self.is_changed[node_id] = node["is_changed"] + return self.is_changed[node_id] + + +class CacheType(Enum): + CLASSIC = 0 + LRU = 1 + DEPENDENCY_AWARE = 2 + + +class CacheSet: + def __init__(self, cache_type=None, cache_size=None): + if cache_type == CacheType.DEPENDENCY_AWARE: + self.init_dependency_aware_cache() + logging.info("Disabling intermediate node cache.") + elif cache_type == CacheType.LRU: + if cache_size is None: + cache_size = 0 + self.init_lru_cache(cache_size) + logging.info("Using LRU cache") + else: + self.init_classic_cache() + + self.all = [self.outputs, self.ui, self.objects] + + # Performs like the old cache -- dump data ASAP + def init_classic_cache(self): + self.outputs = HierarchicalCache(CacheKeySetInputSignature) + self.ui = HierarchicalCache(CacheKeySetInputSignature) + self.objects = HierarchicalCache(CacheKeySetID) + + def init_lru_cache(self, cache_size): + self.outputs = LRUCache(CacheKeySetInputSignature, max_size=cache_size) + self.ui = LRUCache(CacheKeySetInputSignature, max_size=cache_size) + self.objects = HierarchicalCache(CacheKeySetID) + + # only hold cached items while the decendents have not executed + def init_dependency_aware_cache(self): + self.outputs = DependencyAwareCache(CacheKeySetInputSignature) + self.ui = DependencyAwareCache(CacheKeySetInputSignature) + self.objects = DependencyAwareCache(CacheKeySetID) + + def recursive_debug_dump(self): + result = { + "outputs": self.outputs.recursive_debug_dump(), + "ui": self.ui.recursive_debug_dump(), + } + return result + +SENSITIVE_EXTRA_DATA_KEYS = ("auth_token_comfy_org", "api_key_comfy_org") + +def get_input_data(inputs, class_def, unique_id, outputs=None, dynprompt=None, extra_data={}): + is_v3 = issubclass(class_def, _ComfyNodeInternal) + if is_v3: + valid_inputs, schema = class_def.INPUT_TYPES(include_hidden=False, return_schema=True) + else: + valid_inputs = class_def.INPUT_TYPES() + input_data_all = {} + missing_keys = {} + hidden_inputs_v3 = {} + for x in inputs: + input_data = inputs[x] + _, input_category, input_info = get_input_info(class_def, x, valid_inputs) + def mark_missing(): + missing_keys[x] = True + input_data_all[x] = (None,) + if is_link(input_data) and (not input_info or not input_info.get("rawLink", False)): + input_unique_id = input_data[0] + output_index = input_data[1] + if outputs is None: + mark_missing() + continue # This might be a lazily-evaluated input + cached_output = outputs.get(input_unique_id) + if cached_output is None: + mark_missing() + continue + if output_index >= len(cached_output): + mark_missing() + continue + obj = cached_output[output_index] + input_data_all[x] = obj + elif input_category is not None: + input_data_all[x] = [input_data] + + if is_v3: + if schema.hidden: + if io.Hidden.prompt in schema.hidden: + hidden_inputs_v3[io.Hidden.prompt] = dynprompt.get_original_prompt() if dynprompt is not None else {} + if io.Hidden.dynprompt in schema.hidden: + hidden_inputs_v3[io.Hidden.dynprompt] = dynprompt + if io.Hidden.extra_pnginfo in schema.hidden: + hidden_inputs_v3[io.Hidden.extra_pnginfo] = extra_data.get('extra_pnginfo', None) + if io.Hidden.unique_id in schema.hidden: + hidden_inputs_v3[io.Hidden.unique_id] = unique_id + if io.Hidden.auth_token_comfy_org in schema.hidden: + hidden_inputs_v3[io.Hidden.auth_token_comfy_org] = extra_data.get("auth_token_comfy_org", None) + if io.Hidden.api_key_comfy_org in schema.hidden: + hidden_inputs_v3[io.Hidden.api_key_comfy_org] = extra_data.get("api_key_comfy_org", None) + else: + if "hidden" in valid_inputs: + h = valid_inputs["hidden"] + for x in h: + if h[x] == "PROMPT": + input_data_all[x] = [dynprompt.get_original_prompt() if dynprompt is not None else {}] + if h[x] == "DYNPROMPT": + input_data_all[x] = [dynprompt] + if h[x] == "EXTRA_PNGINFO": + input_data_all[x] = [extra_data.get('extra_pnginfo', None)] + if h[x] == "UNIQUE_ID": + input_data_all[x] = [unique_id] + if h[x] == "AUTH_TOKEN_COMFY_ORG": + input_data_all[x] = [extra_data.get("auth_token_comfy_org", None)] + if h[x] == "API_KEY_COMFY_ORG": + input_data_all[x] = [extra_data.get("api_key_comfy_org", None)] + return input_data_all, missing_keys, hidden_inputs_v3 + +map_node_over_list = None #Don't hook this please + +async def resolve_map_node_over_list_results(results): + remaining = [x for x in results if isinstance(x, asyncio.Task) and not x.done()] + if len(remaining) == 0: + return [x.result() if isinstance(x, asyncio.Task) else x for x in results] + else: + done, pending = await asyncio.wait(remaining) + for task in done: + exc = task.exception() + if exc is not None: + raise exc + return [x.result() if isinstance(x, asyncio.Task) else x for x in results] + +async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, func, allow_interrupt=False, execution_block_cb=None, pre_execute_cb=None, hidden_inputs=None): + # check if node wants the lists + input_is_list = getattr(obj, "INPUT_IS_LIST", False) + + if len(input_data_all) == 0: + max_len_input = 0 + else: + max_len_input = max(len(x) for x in input_data_all.values()) + + # get a slice of inputs, repeat last input when list isn't long enough + def slice_dict(d, i): + return {k: v[i if len(v) > i else -1] for k, v in d.items()} + + results = [] + async def process_inputs(inputs, index=None, input_is_list=False): + if allow_interrupt: + nodes.before_node_execution() + execution_block = None + for k, v in inputs.items(): + if input_is_list: + for e in v: + if isinstance(e, ExecutionBlocker): + v = e + break + if isinstance(v, ExecutionBlocker): + execution_block = execution_block_cb(v) if execution_block_cb else v + break + if execution_block is None: + if pre_execute_cb is not None and index is not None: + pre_execute_cb(index) + # V3 + if isinstance(obj, _ComfyNodeInternal) or (is_class(obj) and issubclass(obj, _ComfyNodeInternal)): + # if is just a class, then assign no resources or state, just create clone + if is_class(obj): + type_obj = obj + obj.VALIDATE_CLASS() + class_clone = obj.PREPARE_CLASS_CLONE(hidden_inputs) + # otherwise, use class instance to populate/reuse some fields + else: + type_obj = type(obj) + type_obj.VALIDATE_CLASS() + class_clone = type_obj.PREPARE_CLASS_CLONE(hidden_inputs) + f = make_locked_method_func(type_obj, func, class_clone) + # V1 + else: + f = getattr(obj, func) + if inspect.iscoroutinefunction(f): + async def async_wrapper(f, prompt_id, unique_id, list_index, args): + with CurrentNodeContext(prompt_id, unique_id, list_index): + return await f(**args) + task = asyncio.create_task(async_wrapper(f, prompt_id, unique_id, index, args=inputs)) + # Give the task a chance to execute without yielding + await asyncio.sleep(0) + if task.done(): + result = task.result() + results.append(result) + else: + results.append(task) + else: + with CurrentNodeContext(prompt_id, unique_id, index): + result = f(**inputs) + results.append(result) + else: + results.append(execution_block) + + if input_is_list: + await process_inputs(input_data_all, 0, input_is_list=input_is_list) + elif max_len_input == 0: + await process_inputs({}) + else: + for i in range(max_len_input): + input_dict = slice_dict(input_data_all, i) + await process_inputs(input_dict, i) + return results + + +def merge_result_data(results, obj): + # check which outputs need concatenating + output = [] + output_is_list = [False] * len(results[0]) + if hasattr(obj, "OUTPUT_IS_LIST"): + output_is_list = obj.OUTPUT_IS_LIST + + # merge node execution results + for i, is_list in zip(range(len(results[0])), output_is_list): + if is_list: + value = [] + for o in results: + if isinstance(o[i], ExecutionBlocker): + value.append(o[i]) + else: + value.extend(o[i]) + output.append(value) + else: + output.append([o[i] for o in results]) + return output + +async def get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=None, pre_execute_cb=None, hidden_inputs=None): + return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, hidden_inputs=hidden_inputs) + has_pending_task = any(isinstance(r, asyncio.Task) and not r.done() for r in return_values) + if has_pending_task: + return return_values, {}, False, has_pending_task + output, ui, has_subgraph = get_output_from_returns(return_values, obj) + return output, ui, has_subgraph, False + +def get_output_from_returns(return_values, obj): + results = [] + uis = [] + subgraph_results = [] + has_subgraph = False + for i in range(len(return_values)): + r = return_values[i] + if isinstance(r, dict): + if 'ui' in r: + uis.append(r['ui']) + if 'expand' in r: + # Perform an expansion, but do not append results + has_subgraph = True + new_graph = r['expand'] + result = r.get("result", None) + if isinstance(result, ExecutionBlocker): + result = tuple([result] * len(obj.RETURN_TYPES)) + subgraph_results.append((new_graph, result)) + elif 'result' in r: + result = r.get("result", None) + if isinstance(result, ExecutionBlocker): + result = tuple([result] * len(obj.RETURN_TYPES)) + results.append(result) + subgraph_results.append((None, result)) + elif isinstance(r, _NodeOutputInternal): + # V3 + if r.ui is not None: + if isinstance(r.ui, dict): + uis.append(r.ui) + else: + uis.append(r.ui.as_dict()) + if r.expand is not None: + has_subgraph = True + new_graph = r.expand + result = r.result + if r.block_execution is not None: + result = tuple([ExecutionBlocker(r.block_execution)] * len(obj.RETURN_TYPES)) + subgraph_results.append((new_graph, result)) + elif r.result is not None: + result = r.result + if r.block_execution is not None: + result = tuple([ExecutionBlocker(r.block_execution)] * len(obj.RETURN_TYPES)) + results.append(result) + subgraph_results.append((None, result)) + else: + if isinstance(r, ExecutionBlocker): + r = tuple([r] * len(obj.RETURN_TYPES)) + results.append(r) + subgraph_results.append((None, r)) + + if has_subgraph: + output = subgraph_results + elif len(results) > 0: + output = merge_result_data(results, obj) + else: + output = [] + ui = dict() + # TODO: Think there's an existing bug here + # If we're performing a subgraph expansion, we probably shouldn't be returning UI values yet. + # They'll get cached without the completed subgraphs. It's an edge case and I'm not aware of + # any nodes that use both subgraph expansion and custom UI outputs, but might be a problem in the future. + if len(uis) > 0: + ui = {k: [y for x in uis for y in x[k]] for k in uis[0].keys()} + return output, ui, has_subgraph + +def format_value(x): + if x is None: + return None + elif isinstance(x, (int, float, bool, str)): + return x + else: + return str(x) + +async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes): + unique_id = current_item + real_node_id = dynprompt.get_real_node_id(unique_id) + display_node_id = dynprompt.get_display_node_id(unique_id) + parent_node_id = dynprompt.get_parent_node_id(unique_id) + inputs = dynprompt.get_node(unique_id)['inputs'] + class_type = dynprompt.get_node(unique_id)['class_type'] + class_def = nodes.NODE_CLASS_MAPPINGS[class_type] + if caches.outputs.get(unique_id) is not None: + if server.client_id is not None: + cached_output = caches.ui.get(unique_id) or {} + server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": cached_output.get("output",None), "prompt_id": prompt_id }, server.client_id) + get_progress_state().finish_progress(unique_id) + return (ExecutionResult.SUCCESS, None, None) + + input_data_all = None + try: + if unique_id in pending_async_nodes: + results = [] + for r in pending_async_nodes[unique_id]: + if isinstance(r, asyncio.Task): + try: + results.append(r.result()) + except Exception as ex: + # An async task failed - propagate the exception up + del pending_async_nodes[unique_id] + raise ex + else: + results.append(r) + del pending_async_nodes[unique_id] + output_data, output_ui, has_subgraph = get_output_from_returns(results, class_def) + elif unique_id in pending_subgraph_results: + cached_results = pending_subgraph_results[unique_id] + resolved_outputs = [] + for is_subgraph, result in cached_results: + if not is_subgraph: + resolved_outputs.append(result) + else: + resolved_output = [] + for r in result: + if is_link(r): + source_node, source_output = r[0], r[1] + node_output = caches.outputs.get(source_node)[source_output] + for o in node_output: + resolved_output.append(o) + + else: + resolved_output.append(r) + resolved_outputs.append(tuple(resolved_output)) + output_data = merge_result_data(resolved_outputs, class_def) + output_ui = [] + has_subgraph = False + else: + get_progress_state().start_progress(unique_id) + input_data_all, missing_keys, hidden_inputs = get_input_data(inputs, class_def, unique_id, caches.outputs, dynprompt, extra_data) + if server.client_id is not None: + server.last_node_id = display_node_id + server.send_sync("executing", { "node": unique_id, "display_node": display_node_id, "prompt_id": prompt_id }, server.client_id) + + obj = caches.objects.get(unique_id) + if obj is None: + obj = class_def() + caches.objects.set(unique_id, obj) + + if issubclass(class_def, _ComfyNodeInternal): + lazy_status_present = first_real_override(class_def, "check_lazy_status") is not None + else: + lazy_status_present = getattr(obj, "check_lazy_status", None) is not None + if lazy_status_present: + required_inputs = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, "check_lazy_status", allow_interrupt=True, hidden_inputs=hidden_inputs) + required_inputs = await resolve_map_node_over_list_results(required_inputs) + required_inputs = set(sum([r for r in required_inputs if isinstance(r,list)], [])) + required_inputs = [x for x in required_inputs if isinstance(x,str) and ( + x not in input_data_all or x in missing_keys + )] + if len(required_inputs) > 0: + for i in required_inputs: + execution_list.make_input_strong_link(unique_id, i) + return (ExecutionResult.PENDING, None, None) + + def execution_block_cb(block): + if block.message is not None: + mes = { + "prompt_id": prompt_id, + "node_id": unique_id, + "node_type": class_type, + "executed": list(executed), + + "exception_message": f"Execution Blocked: {block.message}", + "exception_type": "ExecutionBlocked", + "traceback": [], + "current_inputs": [], + "current_outputs": [], + } + server.send_sync("execution_error", mes, server.client_id) + return ExecutionBlocker(None) + else: + return block + def pre_execute_cb(call_index): + # TODO - How to handle this with async functions without contextvars (which requires Python 3.12)? + GraphBuilder.set_default_prefix(unique_id, call_index, 0) + output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, hidden_inputs=hidden_inputs) + if has_pending_tasks: + pending_async_nodes[unique_id] = output_data + unblock = execution_list.add_external_block(unique_id) + async def await_completion(): + tasks = [x for x in output_data if isinstance(x, asyncio.Task)] + await asyncio.gather(*tasks, return_exceptions=True) + unblock() + asyncio.create_task(await_completion()) + return (ExecutionResult.PENDING, None, None) + if len(output_ui) > 0: + caches.ui.set(unique_id, { + "meta": { + "node_id": unique_id, + "display_node": display_node_id, + "parent_node": parent_node_id, + "real_node_id": real_node_id, + }, + "output": output_ui + }) + if server.client_id is not None: + server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": output_ui, "prompt_id": prompt_id }, server.client_id) + if has_subgraph: + cached_outputs = [] + new_node_ids = [] + new_output_ids = [] + new_output_links = [] + for i in range(len(output_data)): + new_graph, node_outputs = output_data[i] + if new_graph is None: + cached_outputs.append((False, node_outputs)) + else: + # Check for conflicts + for node_id in new_graph.keys(): + if dynprompt.has_node(node_id): + raise DuplicateNodeError(f"Attempt to add duplicate node {node_id}. Ensure node ids are unique and deterministic or use graph_utils.GraphBuilder.") + for node_id, node_info in new_graph.items(): + new_node_ids.append(node_id) + display_id = node_info.get("override_display_id", unique_id) + dynprompt.add_ephemeral_node(node_id, node_info, unique_id, display_id) + # Figure out if the newly created node is an output node + class_type = node_info["class_type"] + class_def = nodes.NODE_CLASS_MAPPINGS[class_type] + if hasattr(class_def, 'OUTPUT_NODE') and class_def.OUTPUT_NODE == True: + new_output_ids.append(node_id) + for i in range(len(node_outputs)): + if is_link(node_outputs[i]): + from_node_id, from_socket = node_outputs[i][0], node_outputs[i][1] + new_output_links.append((from_node_id, from_socket)) + cached_outputs.append((True, node_outputs)) + new_node_ids = set(new_node_ids) + for cache in caches.all: + subcache = await cache.ensure_subcache_for(unique_id, new_node_ids) + subcache.clean_unused() + for node_id in new_output_ids: + execution_list.add_node(node_id) + for link in new_output_links: + execution_list.add_strong_link(link[0], link[1], unique_id) + pending_subgraph_results[unique_id] = cached_outputs + return (ExecutionResult.PENDING, None, None) + caches.outputs.set(unique_id, output_data) + except comfy.model_management.InterruptProcessingException as iex: + logging.info("Processing interrupted") + + # skip formatting inputs/outputs + error_details = { + "node_id": real_node_id, + } + + return (ExecutionResult.FAILURE, error_details, iex) + except Exception as ex: + typ, _, tb = sys.exc_info() + exception_type = full_type_name(typ) + input_data_formatted = {} + if input_data_all is not None: + input_data_formatted = {} + for name, inputs in input_data_all.items(): + input_data_formatted[name] = [format_value(x) for x in inputs] + + logging.error(f"!!! Exception during processing !!! {ex}") + logging.error(traceback.format_exc()) + tips = "" + + if isinstance(ex, comfy.model_management.OOM_EXCEPTION): + tips = "This error means you ran out of memory on your GPU.\n\nTIPS: If the workflow worked before you might have accidentally set the batch_size to a large number." + logging.error("Got an OOM, unloading all loaded models.") + comfy.model_management.unload_all_models() + + error_details = { + "node_id": real_node_id, + "exception_message": "{}\n{}".format(ex, tips), + "exception_type": exception_type, + "traceback": traceback.format_tb(tb), + "current_inputs": input_data_formatted + } + + return (ExecutionResult.FAILURE, error_details, ex) + + get_progress_state().finish_progress(unique_id) + executed.add(unique_id) + + return (ExecutionResult.SUCCESS, None, None) + +class PromptExecutor: + def __init__(self, server, cache_type=False, cache_size=None): + self.cache_size = cache_size + self.cache_type = cache_type + self.server = server + self.reset() + + def reset(self): + self.caches = CacheSet(cache_type=self.cache_type, cache_size=self.cache_size) + self.status_messages = [] + self.success = True + + def add_message(self, event, data: dict, broadcast: bool): + data = { + **data, + "timestamp": int(time.time() * 1000), + } + self.status_messages.append((event, data)) + if self.server.client_id is not None or broadcast: + self.server.send_sync(event, data, self.server.client_id) + + def handle_execution_error(self, prompt_id, prompt, current_outputs, executed, error, ex): + node_id = error["node_id"] + class_type = prompt[node_id]["class_type"] + + # First, send back the status to the frontend depending + # on the exception type + if isinstance(ex, comfy.model_management.InterruptProcessingException): + mes = { + "prompt_id": prompt_id, + "node_id": node_id, + "node_type": class_type, + "executed": list(executed), + } + self.add_message("execution_interrupted", mes, broadcast=True) + else: + mes = { + "prompt_id": prompt_id, + "node_id": node_id, + "node_type": class_type, + "executed": list(executed), + "exception_message": error["exception_message"], + "exception_type": error["exception_type"], + "traceback": error["traceback"], + "current_inputs": error["current_inputs"], + "current_outputs": list(current_outputs), + } + self.add_message("execution_error", mes, broadcast=False) + + def execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]): + asyncio.run(self.execute_async(prompt, prompt_id, extra_data, execute_outputs)) + + async def execute_async(self, prompt, prompt_id, extra_data={}, execute_outputs=[]): + nodes.interrupt_processing(False) + + if "client_id" in extra_data: + self.server.client_id = extra_data["client_id"] + else: + self.server.client_id = None + + self.status_messages = [] + self.add_message("execution_start", { "prompt_id": prompt_id}, broadcast=False) + + with torch.inference_mode(): + dynamic_prompt = DynamicPrompt(prompt) + reset_progress_state(prompt_id, dynamic_prompt) + add_progress_handler(WebUIProgressHandler(self.server)) + is_changed_cache = IsChangedCache(prompt_id, dynamic_prompt, self.caches.outputs) + for cache in self.caches.all: + await cache.set_prompt(dynamic_prompt, prompt.keys(), is_changed_cache) + cache.clean_unused() + + cached_nodes = [] + for node_id in prompt: + if self.caches.outputs.get(node_id) is not None: + cached_nodes.append(node_id) + + comfy.model_management.cleanup_models_gc() + self.add_message("execution_cached", + { "nodes": cached_nodes, "prompt_id": prompt_id}, + broadcast=False) + pending_subgraph_results = {} + pending_async_nodes = {} # TODO - Unify this with pending_subgraph_results + executed = set() + execution_list = ExecutionList(dynamic_prompt, self.caches.outputs) + current_outputs = self.caches.outputs.all_node_ids() + for node_id in list(execute_outputs): + execution_list.add_node(node_id) + + while not execution_list.is_empty(): + node_id, error, ex = await execution_list.stage_node_execution() + if error is not None: + self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex) + break + + assert node_id is not None, "Node ID should not be None at this point" + result, error, ex = await execute(self.server, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes) + self.success = result != ExecutionResult.FAILURE + if result == ExecutionResult.FAILURE: + self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex) + break + elif result == ExecutionResult.PENDING: + execution_list.unstage_node_execution() + else: # result == ExecutionResult.SUCCESS: + execution_list.complete_node_execution() + else: + # Only execute when the while-loop ends without break + self.add_message("execution_success", { "prompt_id": prompt_id }, broadcast=False) + + ui_outputs = {} + meta_outputs = {} + all_node_ids = self.caches.ui.all_node_ids() + for node_id in all_node_ids: + ui_info = self.caches.ui.get(node_id) + if ui_info is not None: + ui_outputs[node_id] = ui_info["output"] + meta_outputs[node_id] = ui_info["meta"] + self.history_result = { + "outputs": ui_outputs, + "meta": meta_outputs, + } + self.server.last_node_id = None + if comfy.model_management.DISABLE_SMART_MEMORY: + comfy.model_management.unload_all_models() + + +async def validate_inputs(prompt_id, prompt, item, validated): + unique_id = item + if unique_id in validated: + return validated[unique_id] + + inputs = prompt[unique_id]['inputs'] + class_type = prompt[unique_id]['class_type'] + obj_class = nodes.NODE_CLASS_MAPPINGS[class_type] + + class_inputs = obj_class.INPUT_TYPES() + valid_inputs = set(class_inputs.get('required',{})).union(set(class_inputs.get('optional',{}))) + + errors = [] + valid = True + + validate_function_inputs = [] + validate_has_kwargs = False + if issubclass(obj_class, _ComfyNodeInternal): + validate_function_name = "validate_inputs" + validate_function = first_real_override(obj_class, validate_function_name) + else: + validate_function_name = "VALIDATE_INPUTS" + validate_function = getattr(obj_class, validate_function_name, None) + if validate_function is not None: + argspec = inspect.getfullargspec(validate_function) + validate_function_inputs = argspec.args + validate_has_kwargs = argspec.varkw is not None + received_types = {} + + for x in valid_inputs: + input_type, input_category, extra_info = get_input_info(obj_class, x, class_inputs) + assert extra_info is not None + if x not in inputs: + if input_category == "required": + error = { + "type": "required_input_missing", + "message": "Required input is missing", + "details": f"{x}", + "extra_info": { + "input_name": x + } + } + errors.append(error) + continue + + val = inputs[x] + info = (input_type, extra_info) + if isinstance(val, list): + if len(val) != 2: + error = { + "type": "bad_linked_input", + "message": "Bad linked input, must be a length-2 list of [node_id, slot_index]", + "details": f"{x}", + "extra_info": { + "input_name": x, + "input_config": info, + "received_value": val + } + } + errors.append(error) + continue + + o_id = val[0] + o_class_type = prompt[o_id]['class_type'] + r = nodes.NODE_CLASS_MAPPINGS[o_class_type].RETURN_TYPES + received_type = r[val[1]] + received_types[x] = received_type + if 'input_types' not in validate_function_inputs and not validate_node_input(received_type, input_type): + details = f"{x}, received_type({received_type}) mismatch input_type({input_type})" + error = { + "type": "return_type_mismatch", + "message": "Return type mismatch between linked nodes", + "details": details, + "extra_info": { + "input_name": x, + "input_config": info, + "received_type": received_type, + "linked_node": val + } + } + errors.append(error) + continue + try: + r = await validate_inputs(prompt_id, prompt, o_id, validated) + if r[0] is False: + # `r` will be set in `validated[o_id]` already + valid = False + continue + except Exception as ex: + typ, _, tb = sys.exc_info() + valid = False + exception_type = full_type_name(typ) + reasons = [{ + "type": "exception_during_inner_validation", + "message": "Exception when validating inner node", + "details": str(ex), + "extra_info": { + "input_name": x, + "input_config": info, + "exception_message": str(ex), + "exception_type": exception_type, + "traceback": traceback.format_tb(tb), + "linked_node": val + } + }] + validated[o_id] = (False, reasons, o_id) + continue + else: + try: + # Unwraps values wrapped in __value__ key. This is used to pass + # list widget value to execution, as by default list value is + # reserved to represent the connection between nodes. + if isinstance(val, dict) and "__value__" in val: + val = val["__value__"] + inputs[x] = val + + if input_type == "INT": + val = int(val) + inputs[x] = val + if input_type == "FLOAT": + val = float(val) + inputs[x] = val + if input_type == "STRING": + val = str(val) + inputs[x] = val + if input_type == "BOOLEAN": + val = bool(val) + inputs[x] = val + except Exception as ex: + error = { + "type": "invalid_input_type", + "message": f"Failed to convert an input value to a {input_type} value", + "details": f"{x}, {val}, {ex}", + "extra_info": { + "input_name": x, + "input_config": info, + "received_value": val, + "exception_message": str(ex) + } + } + errors.append(error) + continue + + if x not in validate_function_inputs and not validate_has_kwargs: + if "min" in extra_info and val < extra_info["min"]: + error = { + "type": "value_smaller_than_min", + "message": "Value {} smaller than min of {}".format(val, extra_info["min"]), + "details": f"{x}", + "extra_info": { + "input_name": x, + "input_config": info, + "received_value": val, + } + } + errors.append(error) + continue + if "max" in extra_info and val > extra_info["max"]: + error = { + "type": "value_bigger_than_max", + "message": "Value {} bigger than max of {}".format(val, extra_info["max"]), + "details": f"{x}", + "extra_info": { + "input_name": x, + "input_config": info, + "received_value": val, + } + } + errors.append(error) + continue + + if isinstance(input_type, list): + combo_options = input_type + if val not in combo_options: + input_config = info + list_info = "" + + # Don't send back gigantic lists like if they're lots of + # scanned model filepaths + if len(combo_options) > 20: + list_info = f"(list of length {len(combo_options)})" + input_config = None + else: + list_info = str(combo_options) + + error = { + "type": "value_not_in_list", + "message": "Value not in list", + "details": f"{x}: '{val}' not in {list_info}", + "extra_info": { + "input_name": x, + "input_config": input_config, + "received_value": val, + } + } + errors.append(error) + continue + + if len(validate_function_inputs) > 0 or validate_has_kwargs: + input_data_all, _, hidden_inputs = get_input_data(inputs, obj_class, unique_id) + input_filtered = {} + for x in input_data_all: + if x in validate_function_inputs or validate_has_kwargs: + input_filtered[x] = input_data_all[x] + if 'input_types' in validate_function_inputs: + input_filtered['input_types'] = [received_types] + + ret = await _async_map_node_over_list(prompt_id, unique_id, obj_class, input_filtered, validate_function_name, hidden_inputs=hidden_inputs) + ret = await resolve_map_node_over_list_results(ret) + for x in input_filtered: + for i, r in enumerate(ret): + if r is not True and not isinstance(r, ExecutionBlocker): + details = f"{x}" + if r is not False: + details += f" - {str(r)}" + + error = { + "type": "custom_validation_failed", + "message": "Custom validation failed for node", + "details": details, + "extra_info": { + "input_name": x, + } + } + errors.append(error) + continue + + if len(errors) > 0 or valid is not True: + ret = (False, errors, unique_id) + else: + ret = (True, [], unique_id) + + validated[unique_id] = ret + return ret + +def full_type_name(klass): + module = klass.__module__ + if module == 'builtins': + return klass.__qualname__ + return module + '.' + klass.__qualname__ + +async def validate_prompt(prompt_id, prompt, partial_execution_list: Union[list[str], None]): + outputs = set() + for x in prompt: + if 'class_type' not in prompt[x]: + error = { + "type": "invalid_prompt", + "message": "Cannot execute because a node is missing the class_type property.", + "details": f"Node ID '#{x}'", + "extra_info": {} + } + return (False, error, [], {}) + + class_type = prompt[x]['class_type'] + class_ = nodes.NODE_CLASS_MAPPINGS.get(class_type, None) + if class_ is None: + error = { + "type": "invalid_prompt", + "message": f"Cannot execute because node {class_type} does not exist.", + "details": f"Node ID '#{x}'", + "extra_info": {} + } + return (False, error, [], {}) + + if hasattr(class_, 'OUTPUT_NODE') and class_.OUTPUT_NODE is True: + if partial_execution_list is None or x in partial_execution_list: + outputs.add(x) + + if len(outputs) == 0: + error = { + "type": "prompt_no_outputs", + "message": "Prompt has no outputs", + "details": "", + "extra_info": {} + } + return (False, error, [], {}) + + good_outputs = set() + errors = [] + node_errors = {} + validated = {} + for o in outputs: + valid = False + reasons = [] + try: + m = await validate_inputs(prompt_id, prompt, o, validated) + valid = m[0] + reasons = m[1] + except Exception as ex: + typ, _, tb = sys.exc_info() + valid = False + exception_type = full_type_name(typ) + reasons = [{ + "type": "exception_during_validation", + "message": "Exception when validating node", + "details": str(ex), + "extra_info": { + "exception_type": exception_type, + "traceback": traceback.format_tb(tb) + } + }] + validated[o] = (False, reasons, o) + + if valid is True: + good_outputs.add(o) + else: + logging.error(f"Failed to validate prompt for output {o}:") + if len(reasons) > 0: + logging.error("* (prompt):") + for reason in reasons: + logging.error(f" - {reason['message']}: {reason['details']}") + errors += [(o, reasons)] + for node_id, result in validated.items(): + valid = result[0] + reasons = result[1] + # If a node upstream has errors, the nodes downstream will also + # be reported as invalid, but there will be no errors attached. + # So don't return those nodes as having errors in the response. + if valid is not True and len(reasons) > 0: + if node_id not in node_errors: + class_type = prompt[node_id]['class_type'] + node_errors[node_id] = { + "errors": reasons, + "dependent_outputs": [], + "class_type": class_type + } + logging.error(f"* {class_type} {node_id}:") + for reason in reasons: + logging.error(f" - {reason['message']}: {reason['details']}") + node_errors[node_id]["dependent_outputs"].append(o) + logging.error("Output will be ignored") + + if len(good_outputs) == 0: + errors_list = [] + for o, errors in errors: + for error in errors: + errors_list.append(f"{error['message']}: {error['details']}") + errors_list = "\n".join(errors_list) + + error = { + "type": "prompt_outputs_failed_validation", + "message": "Prompt outputs failed validation", + "details": errors_list, + "extra_info": {} + } + + return (False, error, list(good_outputs), node_errors) + + return (True, None, list(good_outputs), node_errors) + +MAXIMUM_HISTORY_SIZE = 10000 + +class PromptQueue: + def __init__(self, server): + self.server = server + self.mutex = threading.RLock() + self.not_empty = threading.Condition(self.mutex) + self.task_counter = 0 + self.queue = [] + self.currently_running = {} + self.history = {} + self.flags = {} + + def put(self, item): + with self.mutex: + heapq.heappush(self.queue, item) + self.server.queue_updated() + self.not_empty.notify() + + def get(self, timeout=None): + with self.not_empty: + while len(self.queue) == 0: + self.not_empty.wait(timeout=timeout) + if timeout is not None and len(self.queue) == 0: + return None + item = heapq.heappop(self.queue) + i = self.task_counter + self.currently_running[i] = copy.deepcopy(item) + self.task_counter += 1 + self.server.queue_updated() + return (item, i) + + class ExecutionStatus(NamedTuple): + status_str: Literal['success', 'error'] + completed: bool + messages: List[str] + + def task_done(self, item_id, history_result, + status: Optional['PromptQueue.ExecutionStatus']): + with self.mutex: + prompt = self.currently_running.pop(item_id) + if len(self.history) > MAXIMUM_HISTORY_SIZE: + self.history.pop(next(iter(self.history))) + + status_dict: Optional[dict] = None + if status is not None: + status_dict = copy.deepcopy(status._asdict()) + + # Remove sensitive data from extra_data before storing in history + for sensitive_val in SENSITIVE_EXTRA_DATA_KEYS: + if sensitive_val in prompt[3]: + prompt[3].pop(sensitive_val) + + self.history[prompt[1]] = { + "prompt": prompt, + "outputs": {}, + 'status': status_dict, + } + self.history[prompt[1]].update(history_result) + self.server.queue_updated() + + # Note: slow + def get_current_queue(self): + with self.mutex: + out = [] + for x in self.currently_running.values(): + out += [x] + return (out, copy.deepcopy(self.queue)) + + # read-safe as long as queue items are immutable + def get_current_queue_volatile(self): + with self.mutex: + running = [x for x in self.currently_running.values()] + queued = copy.copy(self.queue) + return (running, queued) + + def get_tasks_remaining(self): + with self.mutex: + return len(self.queue) + len(self.currently_running) + + def wipe_queue(self): + with self.mutex: + self.queue = [] + self.server.queue_updated() + + def delete_queue_item(self, function): + with self.mutex: + for x in range(len(self.queue)): + if function(self.queue[x]): + if len(self.queue) == 1: + self.wipe_queue() + else: + self.queue.pop(x) + heapq.heapify(self.queue) + self.server.queue_updated() + return True + return False + + def get_history(self, prompt_id=None, max_items=None, offset=-1, map_function=None): + with self.mutex: + if prompt_id is None: + out = {} + i = 0 + if offset < 0 and max_items is not None: + offset = len(self.history) - max_items + for k in self.history: + if i >= offset: + p = self.history[k] + if map_function is not None: + p = map_function(p) + out[k] = p + if max_items is not None and len(out) >= max_items: + break + i += 1 + return out + elif prompt_id in self.history: + p = self.history[prompt_id] + if map_function is None: + p = copy.deepcopy(p) + else: + p = map_function(p) + return {prompt_id: p} + else: + return {} + + def wipe_history(self): + with self.mutex: + self.history = {} + + def delete_history_item(self, id_to_delete): + with self.mutex: + self.history.pop(id_to_delete, None) + + def set_flag(self, name, data): + with self.mutex: + self.flags[name] = data + self.not_empty.notify() + + def get_flags(self, reset=True): + with self.mutex: + if reset: + ret = self.flags + self.flags = {} + return ret + else: + return self.flags.copy() diff --git a/model_management_mgpu.py b/model_management_mgpu.py index 7a58aeb..64bf315 100644 --- a/model_management_mgpu.py +++ b/model_management_mgpu.py @@ -19,6 +19,34 @@ from collections import defaultdict logger = logging.getLogger("MultiGPU") +# ========================================================================================== +# GC Anchor System for Model Retention Testing +# ========================================================================================== + +# Global anchor set to prevent GC of models with keep_loaded=True +_MGPU_RETENTION_ANCHORS = set() + +def add_retention_anchor(model_patcher, reason="keep_loaded"): + """Add a model patcher to the GC anchor set to prevent premature garbage collection""" + if model_patcher is not None: + _MGPU_RETENTION_ANCHORS.add(model_patcher) + model_name = type(getattr(model_patcher, 'model', model_patcher)).__name__ + logger.mgpu_mm_log(f"[GC_ANCHOR] Added retention anchor for {model_name}, reason: {reason}, total anchors: {len(_MGPU_RETENTION_ANCHORS)}") + +def remove_retention_anchor(model_patcher, reason="cleanup"): + """Remove a model patcher from the GC anchor set""" + if model_patcher is not None and model_patcher in _MGPU_RETENTION_ANCHORS: + _MGPU_RETENTION_ANCHORS.discard(model_patcher) + model_name = type(getattr(model_patcher, 'model', model_patcher)).__name__ + logger.mgpu_mm_log(f"[GC_ANCHOR] Removed retention anchor for {model_name}, reason: {reason}, total anchors: {len(_MGPU_RETENTION_ANCHORS)}") + +def clear_all_retention_anchors(reason="manual_clear"): + """Clear all retention anchors""" + count = len(_MGPU_RETENTION_ANCHORS) + _MGPU_RETENTION_ANCHORS.clear() + logger.mgpu_mm_log(f"[GC_ANCHOR] Cleared all {count} retention anchors, reason: {reason}") + + # ========================================================================================== # Model Analysis and Store Management (DisTorch V1 & V2) # ========================================================================================== @@ -63,27 +91,44 @@ def prune_distorch_stores(): active_hashes_v2 = set() active_hashes_v1 = set() - for lm in mm.current_loaded_models: + logger.mgpu_mm_log(f"[PRUNE_DEBUG] Starting prune - current_loaded_models count: {len(mm.current_loaded_models)}") + + for i, lm in enumerate(mm.current_loaded_models): mp = lm.model if mp is not None: - active_hashes_v2.add(create_safetensor_model_hash(mp, "prune_check_v2")) - active_hashes_v1.add(create_model_hash(mp, "prune_check_v1")) + try: + hash_v2 = create_safetensor_model_hash(mp, "prune_check_v2") + hash_v1 = create_model_hash(mp, "prune_check_v1") + active_hashes_v2.add(hash_v2) + active_hashes_v1.add(hash_v1) + + model_name = type(getattr(mp, 'model', mp)).__name__ + keep_loaded = getattr(getattr(mp, 'model', None), '_mgpu_keep_loaded', False) + has_v2_alloc = hash_v2 in safetensor_allocation_store + logger.mgpu_mm_log(f"[PRUNE_DEBUG] Model {i}: {model_name}, keep_loaded={keep_loaded}, hash={hash_v2[:8]}, has_v2_allocation={has_v2_alloc}") + except Exception as e: + logger.mgpu_mm_log(f"[PRUNE_DEBUG] Model {i}: Error getting hash - {e}") + + logger.mgpu_mm_log(f"[PRUNE_DEBUG] Active hashes V2: {len(active_hashes_v2)}, Store has: {len(safetensor_allocation_store)}") # V1 pruning stale_v1 = set(model_allocation_store.keys()) - active_hashes_v1 if stale_v1: - logger.info(f"[MultiGPU_Memory_Management] Pruning {len(stale_v1)} DisTorch V1 entries") + logger.mgpu_mm_log(f"[MultiGPU_Memory_Management] Pruning {len(stale_v1)} DisTorch V1 entries") for k in stale_v1: del model_allocation_store[k] - # V2 pruning + # V2 pruning with diagnostics for store, name in ((safetensor_allocation_store, "allocation"), (safetensor_settings_store, "settings")): stale_v2 = set(store.keys()) - active_hashes_v2 if stale_v2: - logger.info(f"[MultiGPU_Memory_Management] Pruning {len(stale_v2)} V2 {name} entries") + logger.mgpu_mm_log(f"[MultiGPU_Memory_Management] Would prune {len(stale_v2)} V2 {name} entries: {[h[:8] for h in list(stale_v2)[:5]]}") for k in stale_v2: del store[k] + else: + logger.mgpu_mm_log(f"[PRUNE_DEBUG] No stale {name} entries to prune") + logger.mgpu_mm_log(f"[PRUNE_DEBUG] After pruning - V2 allocation store has: {len(safetensor_allocation_store)} entries") multigpu_memory_log("distorch_prune", "end") # ========================================================================================== @@ -117,7 +162,7 @@ def _capture_memory_snapshot(): return snapshot def multigpu_memory_log(identifier, tag): - """Record timestamped memory snapshot with delta logging""" + """Record timestamped memory snapshot with clean aligned logging""" if identifier == "print_summary": for id_key in sorted(_MEM_SNAPSHOT_SERIES.keys()): series = _MEM_SNAPSHOT_SERIES[id_key] @@ -125,12 +170,13 @@ def multigpu_memory_log(identifier, tag): for ts, tag_name, snap in series: parts = [] cpu_used, cpu_total = snap.get("cpu", (0, 0)) - parts.append(f"cpu={cpu_used/(1024**3):.2f}/{cpu_total/(1024**3):.2f}") + parts.append(f"cpu|{cpu_used/(1024**3):.2f}") for dev in sorted([k for k in snap.keys() if k != "cpu"]): used, total = snap[dev] - parts.append(f"{dev}={used/(1024**3):.2f}/{total/(1024**3):.2f}") + parts.append(f"{dev}|{used/(1024**3):.2f}") ts_str = ts.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z" - logger.mgpu_mm_log(f"{ts_str} {id_key} {tag_name} | " + " | ".join(parts)) + tag_padded = f"{id_key}_{tag_name}".ljust(35) + logger.mgpu_mm_log(f"{ts_str} {tag_padded} {' '.join(parts)}") return ts = datetime.now(timezone.utc) @@ -141,28 +187,20 @@ def multigpu_memory_log(identifier, tag): _MEM_SNAPSHOT_SERIES[identifier] = [] _MEM_SNAPSHOT_SERIES[identifier].append((ts, tag, curr)) - # Compute delta - if identifier in _MEM_SNAPSHOT_LAST: - prev_tag, prev = _MEM_SNAPSHOT_LAST[identifier] - keys = set(prev.keys()) | set(curr.keys()) - ordered = ["cpu"] + sorted([k for k in keys if k != "cpu"]) - parts = [] - for k in ordered: - p_used, _ = prev.get(k, (0, 0)) - c_used, _ = curr.get(k, (0, 0)) - delta = c_used - p_used - sign = "+" if delta >= 0 else "-" - parts.append(f"{k}={sign}{abs(delta)/(1024**3):.2f}") - logger.mgpu_mm_log(f"{identifier} {tag} - {prev_tag}: " + " | ".join(parts)) - else: - # Baseline - ordered = ["cpu"] + sorted([k for k in curr.keys() if k != "cpu"]) - parts = [] - for k in ordered: - c_used, _ = curr.get(k, (0, 0)) - parts.append(f"{k}=+{c_used/(1024**3):.2f}") - logger.mgpu_mm_log(f"{identifier} {tag} - : " + " | ".join(parts)) - + # Clean aligned format: timestamp + padded tag + memory values + ts_str = ts.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z" + tag_padded = f"{identifier}_{tag}".ljust(35) + + parts = [] + cpu_used, _ = curr.get("cpu", (0, 0)) + parts.append(f"cpu|{cpu_used/(1024**3):.2f}") + + for dev in sorted([k for k in curr.keys() if k != "cpu"]): + used, _ = curr[dev] + parts.append(f"{dev}|{used/(1024**3):.2f}") + + logger.mgpu_mm_log(f"{ts_str} {tag_padded} {' '.join(parts)}") + _MEM_SNAPSHOT_LAST[identifier] = (tag, curr) def clear_memory_snapshot_history(): @@ -331,3 +369,73 @@ def force_full_system_cleanup(reason="manual", force=True): summary = f"[ManagerMatch] Cleanup requested (reason={reason}) | models {pre_models}->{post_models}, cpu_delta_mb={delta_cpu_mb:.2f}" logger.mgpu_mm_log(summary) return summary + + +# ========================================================================================== +# Core Patching: unload_all_models with keep_loaded retention +# ========================================================================================== + +if hasattr(mm, 'unload_all_models') and not hasattr(mm.unload_all_models, '_mgpu_keep_loaded_patched'): + logger.info("[MultiGPU Core Patching] Patching mm.unload_all_models to respect keep_loaded flag for DisTorch models") + + _mgpu_original_unload_all_models = mm.unload_all_models + + def _mgpu_patched_unload_all_models(): + """ + Patched mm.unload_all_models that preserves DisTorch models with _mgpu_keep_loaded=True. + All other models (including DisTorch models without the flag) unload normally. + """ + logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Patched unload_all_models called - initial model count: {len(mm.current_loaded_models)}") + + # Direct approach: iterate through loaded models and selectively unload + models_to_unload = [] + kept_models = [] + + for i, lm in enumerate(mm.current_loaded_models): + mp = lm.model # weakref call to ModelPatcher + if mp is not None and hasattr(mp, 'model'): + # Check if this is a DisTorch model with keep_loaded flag + keep_loaded = getattr(mp.model, '_mgpu_keep_loaded', False) + model_name = type(getattr(mp, 'model', mp)).__name__ + logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Model {i}: {model_name}, keep_loaded={keep_loaded}") + + if keep_loaded: + kept_models.append(lm) + logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Adding to kept_models: {model_name}") + # GC ANCHOR TEST: Prevent premature GC of clone patchers + add_retention_anchor(mp, "keep_loaded_test") + else: + models_to_unload.append(lm) + else: + logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Model {i}: ModelPatcher is None or missing model attribute") + models_to_unload.append(lm) + + logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Final counts - kept_models: {len(kept_models)}, models_to_unload: {len(models_to_unload)}") + + if kept_models: + logger.mgpu_mm_log(f"Found {len(kept_models)} model(s) to retain, unloading {len(models_to_unload)} model(s)") + + # Unload models that don't have keep_loaded flag + for lm in models_to_unload: + try: + lm.model_unload(unpatch_weights=True) + logger.debug(f"Unloaded model: {type(lm.model.model).__name__ if lm.model else 'Unknown'}") + except Exception as e: + logger.warning(f"Error unloading model: {e}") + + # Remove unloaded models from current_loaded_models + mm.current_loaded_models = kept_models + logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Updated mm.current_loaded_models, new count: {len(mm.current_loaded_models)}") + logger.mgpu_mm_log(f"Successfully retained {len(kept_models)} model(s) during unload") + else: + logger.mgpu_mm_log("No models with keep_loaded=True found - delegating to original unload_all_models") + _mgpu_original_unload_all_models() + + mm.unload_all_models = _mgpu_patched_unload_all_models + mm.unload_all_models._mgpu_keep_loaded_patched = True + logger.info("[MultiGPU Core Patching] mm.unload_all_models patched successfully") +else: + if not hasattr(mm, 'unload_all_models'): + logger.warning("[MultiGPU Core Patching] mm.unload_all_models not found - cannot patch keep_loaded retention") + else: + logger.debug("[MultiGPU Core Patching] mm.unload_all_models already patched for keep_loaded - skipping") diff --git a/nodes.py b/nodes.py index a6f6a7d..447a02c 100644 --- a/nodes.py +++ b/nodes.py @@ -530,6 +530,30 @@ class DownloadAndLoadHyVideoTextEncoder: return original_loader.loadmodel(llm_model, clip_model, precision, apply_final_norm, hidden_state_skip_layer, quantization) +class UNetLoaderLP: + """UNet Loader (Low Precision) - sets LoRA precision to False for CPU storage optimization""" + @classmethod + def INPUT_TYPES(s): + return {"required": { "unet_name": (folder_paths.get_filename_list("unet"), ), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "load_unet" + CATEGORY = "loaders" + TITLE = "UNet Loader (LP)" + + def load_unet(self, unet_name): + original_loader = NODE_CLASS_MAPPINGS["UNETLoader"]() + out = original_loader.load_unet(unet_name) + + # Set the low-precision LoRA flag on the loaded model + if hasattr(out[0], 'model'): + out[0].model._distorch_high_precision_loras = False + elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'): + out[0].patcher.model._distorch_high_precision_loras = False + + return out + + class FullCleanupMultiGPU: @classmethod def INPUT_TYPES(s):