From 6f88e549ff82b615b3d7b16256f99cb2b0165fe8 Mon Sep 17 00:00:00 2001 From: Vijaysinh Date: Fri, 28 Nov 2025 17:23:32 +0530 Subject: [PATCH 1/7] Add function to unload all models from RAM and GPU This function unloads all models from RAM and GPU, freeing resources when a user cancels. It handles exceptions during the unload process and logs the outcome. --- comfy/model_management.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/comfy/model_management.py b/comfy/model_management.py index 38c506df5..2157dc34e 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -1524,3 +1524,28 @@ def throw_exception_if_processing_interrupted(): if interrupt_processing: interrupt_processing = False raise InterruptProcessingException() + +def unload_all_models_full(): + """ + Completely unloads all models from RAM and GPU when the user cancels. + Frees CPU RAM, GPU VRAM, and clears python references. + """ + global current_loaded_models + try: + # Unload every model object + for m in current_loaded_models: + try: + m.model_unload(memory_to_free=None, unpatch_weights=True) + except: + pass + current_loaded_models.clear() + + # Force Python GC + gc.collect() + + # Clear GPU memory + soft_empty_cache(force=True) + + logging.info("All models unloaded successfully (manual full unload).") + except Exception as e: + logging.warning(f"Model unload warning: {e}") From 0c85c7a92f457c409ba0814a3de676b93521b35e Mon Sep 17 00:00:00 2001 From: Vijaysinh Date: Fri, 28 Nov 2025 18:23:17 +0530 Subject: [PATCH 2/7] Create extra_config.json --- comfy/extra_config.json | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 comfy/extra_config.json diff --git a/comfy/extra_config.json b/comfy/extra_config.json new file mode 100644 index 000000000..4facaec9f --- /dev/null +++ b/comfy/extra_config.json @@ -0,0 +1,3 @@ +{ + "unload_text_encoder_after_run": true +} From 387aeed9cef2fbe6f8de606c403ce785813d01ca Mon Sep 17 00:00:00 2001 From: Vijaysinh Date: Fri, 28 Nov 2025 18:25:40 +0530 Subject: [PATCH 3/7] Update model_management.py --- comfy/model_management.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 2157dc34e..175af930c 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -15,7 +15,8 @@ You should have received a copy of the GNU General Public License along with this program. If not, see . """ - +import gc +import torch import psutil import logging from enum import Enum @@ -1524,7 +1525,7 @@ def throw_exception_if_processing_interrupted(): if interrupt_processing: interrupt_processing = False raise InterruptProcessingException() - +""" def unload_all_models_full(): """ Completely unloads all models from RAM and GPU when the user cancels. @@ -1549,3 +1550,13 @@ def unload_all_models_full(): logging.info("All models unloaded successfully (manual full unload).") except Exception as e: logging.warning(f"Model unload warning: {e}") + """ + + +def cleanup_ram(): + gc.collect() + try: + torch.cuda.empty_cache() + except: + pass + From 63419c107bc23c0e7af1842605e66c7d77c273ed Mon Sep 17 00:00:00 2001 From: Vijaysinh Date: Fri, 28 Nov 2025 18:30:33 +0530 Subject: [PATCH 4/7] Implement unload_text_encoder for memory cleanup Add unload_text_encoder function to manage encoder memory. --- comfy/model_management.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/comfy/model_management.py b/comfy/model_management.py index 175af930c..ad2ffa253 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -1550,6 +1550,7 @@ def unload_all_models_full(): logging.info("All models unloaded successfully (manual full unload).") except Exception as e: logging.warning(f"Model unload warning: {e}") + """ @@ -1559,4 +1560,14 @@ def cleanup_ram(): torch.cuda.empty_cache() except: pass +def unload_text_encoder(encoder): + if encoder is None: + return + try: + if hasattr(encoder, "model"): + del encoder.model + del encoder + except: + pass + cleanup_ram() From 74037f35db6689d5017e4484056e688e6430b4e8 Mon Sep 17 00:00:00 2001 From: Vijaysinh Date: Fri, 28 Nov 2025 18:34:15 +0530 Subject: [PATCH 5/7] Remove unnecessary blank line in model_management.py --- comfy/model_management.py | 1 - 1 file changed, 1 deletion(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index ad2ffa253..7815ab77f 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -1553,7 +1553,6 @@ def unload_all_models_full(): """ - def cleanup_ram(): gc.collect() try: From 95acb06959db5c89525b4e66b51613f46ccfd0f0 Mon Sep 17 00:00:00 2001 From: Vijaysinh Date: Fri, 28 Nov 2025 18:55:21 +0530 Subject: [PATCH 6/7] Update execution.py --- execution.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/execution.py b/execution.py index 17c77beab..0c287f26b 100644 --- a/execution.py +++ b/execution.py @@ -731,6 +731,31 @@ class PromptExecutor: "outputs": ui_outputs, "meta": meta_outputs, } + + try: + import comfy.model_management as mm + + # If ComfyUI exposes loaded text encoders (most builds do) + if hasattr(mm, "loaded_text_encoders"): + for enc in list(mm.loaded_text_encoders.values()): + try: + mm.unload_text_encoder(enc) + except: + pass + + mm.loaded_text_encoders.clear() + + # Final RAM + VRAM cleanup + try: + mm.cleanup_models_gc() + except: + pass + + print("[RAM Optimizer] Text encoders unloaded successfully after run.") + except Exception as e: + print(f"[RAM Optimizer] Failed to unload text encoders: {e}") + # --- END: Text Encoder RAM Cleanup Patch --- + self.server.last_node_id = None if comfy.model_management.DISABLE_SMART_MEMORY: comfy.model_management.unload_all_models() From 460e9e0d8a3ecf4e338a575db88205d82ae550a0 Mon Sep 17 00:00:00 2001 From: Vijay2359 Date: Mon, 8 Dec 2025 19:32:25 +0530 Subject: [PATCH 7/7] Added new feature/fix for ComfyUI --- comfy/model_management.py | 31 ++++--------------------------- nodes.py | 35 ++++++++++++++++++++++++++++------- reproduce_stream_error.py | 27 +++++++++++++++++++++++++++ 3 files changed, 59 insertions(+), 34 deletions(-) create mode 100644 reproduce_stream_error.py diff --git a/comfy/model_management.py b/comfy/model_management.py index 7815ab77f..532aac680 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -1079,11 +1079,15 @@ def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False, str if dtype is None or weight.dtype == dtype: return weight if stream is not None: + if not hasattr(stream, "__enter__"): + logging.error(f"Stream object {stream} of type {type(stream)} does not have __enter__ method") with stream: return weight.to(dtype=dtype, copy=copy) return weight.to(dtype=dtype, copy=copy) if stream is not None: + if not hasattr(stream, "__enter__"): + logging.error(f"Stream object {stream} of type {type(stream)} does not have __enter__ method") with stream: r = torch.empty_like(weight, dtype=dtype, device=device) r.copy_(weight, non_blocking=non_blocking) @@ -1525,33 +1529,6 @@ def throw_exception_if_processing_interrupted(): if interrupt_processing: interrupt_processing = False raise InterruptProcessingException() -""" -def unload_all_models_full(): - """ - Completely unloads all models from RAM and GPU when the user cancels. - Frees CPU RAM, GPU VRAM, and clears python references. - """ - global current_loaded_models - try: - # Unload every model object - for m in current_loaded_models: - try: - m.model_unload(memory_to_free=None, unpatch_weights=True) - except: - pass - current_loaded_models.clear() - - # Force Python GC - gc.collect() - - # Clear GPU memory - soft_empty_cache(force=True) - - logging.info("All models unloaded successfully (manual full unload).") - except Exception as e: - logging.warning(f"Model unload warning: {e}") - - """ def cleanup_ram(): gc.collect() diff --git a/nodes.py b/nodes.py index bf73eb90e..458a88e0f 100644 --- a/nodes.py +++ b/nodes.py @@ -43,6 +43,9 @@ import folder_paths import latent_preview import node_helpers +if getattr(args, "enable_manager", False): + import comfyui_manager + def before_node_execution(): comfy.model_management.throw_exception_if_processing_interrupted() @@ -692,8 +695,10 @@ class LoraLoaderModelOnly(LoraLoader): return (self.load_lora(model, None, lora_name, strength_model, 0)[0],) class VAELoader: + video_taes = ["taehv", "lighttaew2_2", "lighttaew2_1", "lighttaehy1_5"] + image_taes = ["taesd", "taesdxl", "taesd3", "taef1"] @staticmethod - def vae_list(): + def vae_list(s): vaes = folder_paths.get_filename_list("vae") approx_vaes = folder_paths.get_filename_list("vae_approx") sdxl_taesd_enc = False @@ -722,6 +727,11 @@ class VAELoader: f1_taesd_dec = True elif v.startswith("taef1_decoder."): f1_taesd_enc = True + else: + for tae in s.video_taes: + if v.startswith(tae): + vaes.append(v) + if sd1_taesd_dec and sd1_taesd_enc: vaes.append("taesd") if sdxl_taesd_dec and sdxl_taesd_enc: @@ -765,7 +775,7 @@ class VAELoader: @classmethod def INPUT_TYPES(s): - return {"required": { "vae_name": (s.vae_list(), )}} + return {"required": { "vae_name": (s.vae_list(s), )}} RETURN_TYPES = ("VAE",) FUNCTION = "load_vae" @@ -776,10 +786,13 @@ class VAELoader: if vae_name == "pixel_space": sd = {} sd["pixel_space_vae"] = torch.tensor(1.0) - elif vae_name in ["taesd", "taesdxl", "taesd3", "taef1"]: + elif vae_name in self.image_taes: sd = self.load_taesd(vae_name) else: - vae_path = folder_paths.get_full_path_or_raise("vae", vae_name) + if os.path.splitext(vae_name)[0] in self.video_taes: + vae_path = folder_paths.get_full_path_or_raise("vae_approx", vae_name) + else: + vae_path = folder_paths.get_full_path_or_raise("vae", vae_name) sd = comfy.utils.load_torch_file(vae_path) vae = comfy.sd.VAE(sd=sd) vae.throw_exception_if_invalid() @@ -929,7 +942,7 @@ class CLIPLoader: @classmethod def INPUT_TYPES(s): return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ), - "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2"], ), + "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis"], ), }, "optional": { "device": (["default", "cpu"], {"advanced": True}), @@ -957,7 +970,7 @@ class DualCLIPLoader: def INPUT_TYPES(s): return {"required": { "clip_name1": (folder_paths.get_filename_list("text_encoders"), ), "clip_name2": (folder_paths.get_filename_list("text_encoders"), ), - "type": (["sdxl", "sd3", "flux", "hunyuan_video", "hidream", "hunyuan_image", "hunyuan_video_15"], ), + "type": (["sdxl", "sd3", "flux", "hunyuan_video", "hidream", "hunyuan_image", "hunyuan_video_15", "kandinsky5", "kandinsky5_image"], ), }, "optional": { "device": (["default", "cpu"], {"advanced": True}), @@ -2233,6 +2246,12 @@ async def init_external_custom_nodes(): if args.disable_all_custom_nodes and possible_module not in args.whitelist_custom_nodes: logging.info(f"Skipping {possible_module} due to disable_all_custom_nodes and whitelist_custom_nodes") continue + + if getattr(args, "enable_manager", False): + if comfyui_manager.should_be_disabled(module_path): + logging.info(f"Blocked by policy: {module_path}") + continue + time_before = time.perf_counter() success = await load_custom_node(module_path, base_node_names, module_parent="custom_nodes") node_import_times.append((time.perf_counter() - time_before, module_path, success)) @@ -2336,7 +2355,9 @@ async def init_builtin_extra_nodes(): "nodes_easycache.py", "nodes_audio_encoder.py", "nodes_rope.py", + "nodes_logic.py", "nodes_nop.py", + "nodes_kandinsky5.py", ] import_failed = [] @@ -2428,4 +2449,4 @@ async def init_extra_nodes(init_custom_nodes=True, init_api_nodes=True): logging.warning("Please do a: pip install -r requirements.txt") logging.warning("") - return import_failed + return import_failed \ No newline at end of file diff --git a/reproduce_stream_error.py b/reproduce_stream_error.py new file mode 100644 index 000000000..7aef72120 --- /dev/null +++ b/reproduce_stream_error.py @@ -0,0 +1,27 @@ + +import torch +import logging + +logging.basicConfig(level=logging.INFO) + +def test_stream(): + if not torch.cuda.is_available(): + print("CUDA not available, cannot test cuda stream") + return + + device = torch.device("cuda") + stream = torch.cuda.Stream(device=device, priority=0) + + print(f"Stream type: {type(stream)}") + print(f"Has __enter__: {hasattr(stream, '__enter__')}") + + try: + with stream: + print("Stream context manager works") + except AttributeError as e: + print(f"AttributeError caught: {e}") + except Exception as e: + print(f"Other exception caught: {e}") + +if __name__ == "__main__": + test_stream()