From 82b4ba9c911f3fa5e249bfb50edc474559af1a5f Mon Sep 17 00:00:00 2001 From: Isi Date: Fri, 29 Aug 2025 09:17:10 +0100 Subject: [PATCH] Update delete_model_passthrough.py --- delete_model_passthrough.py | 213 +++++++++++++++++++++++++++++++++++- 1 file changed, 208 insertions(+), 5 deletions(-) diff --git a/delete_model_passthrough.py b/delete_model_passthrough.py index e7db62f..a65ffee 100644 --- a/delete_model_passthrough.py +++ b/delete_model_passthrough.py @@ -1,3 +1,153 @@ +# import torch, gc, psutil +# import comfy.model_management as mm +# from comfy.model_management import current_loaded_models, LoadedModel + +# class AnyType(str): +# def __ne__(self, __value: object) -> bool: return False +# any_typ = AnyType("*") + + +# def hard_free(model): +# """ +# Forcefully free model memory by detaching parameters and clearing caches. +# """ +# try: +# # Clear parameters and buffers +# if hasattr(model, "parameters"): +# for p in model.parameters(): +# if p is not None: +# p.detach_() +# if hasattr(p, 'data'): +# del p.data +# del p + +# if hasattr(model, "buffers"): +# for b in model.buffers(): +# if b is not None: +# b.detach_() +# if hasattr(b, 'data'): +# del b.data +# del b + +# # Clear any cached attributes +# for attr_name in list(vars(model).keys()): +# attr = getattr(model, attr_name) +# if isinstance(attr, torch.Tensor): +# attr.detach_() +# delattr(model, attr_name) + +# except Exception as e: +# print(f"Hard free error: {e}") + +# try: +# del model +# except: +# pass + +# gc.collect() +# mm.soft_empty_cache() + + +# def purge_specific_from_comfy(model_obj): +# """ +# Remove a specific model from ComfyUI's current_loaded_models list. +# Returns the model type if found and removed. +# """ +# deleted_type = "Unknown" + +# try: +# # Check if it's a LoadedModel wrapper +# if isinstance(model_obj, LoadedModel): +# if model_obj in current_loaded_models: +# current_loaded_models.remove(model_obj) +# # Also detach the underlying model +# if hasattr(model_obj, 'model_unload'): +# model_obj.model_unload() +# deleted_type = f"LoadedModel ({type(model_obj.model).__name__})" + +# # Check if it's a raw model that might be wrapped in LoadedModel +# else: +# for i, loaded_model in enumerate(current_loaded_models): +# if loaded_model.model is model_obj: +# current_loaded_models.pop(i) +# if hasattr(loaded_model, 'model_unload'): +# loaded_model.model_unload() +# deleted_type = f"RawModel ({type(model_obj).__name__})" +# break + +# # Additional cleanup for common model types +# if hasattr(model_obj, '__class__'): +# cls_name = model_obj.__class__.__name__.lower() +# if 'unet' in cls_name: +# deleted_type = "UNet" +# elif 'clip' in cls_name: +# deleted_type = "CLIP" +# elif 'vae' in cls_name: +# deleted_type = "VAE" + +# except Exception as e: +# print(f"Cache purge error: {e}") + +# return deleted_type + + +# class DeleteModelPassthrough: +# """ +# ComfyUI Custom Node: +# - Accepts any input + a model (UNet/CLIP/VAE/latent). +# - Deletes ONLY that model from VRAM, RAM, and ComfyUI cache. +# - Auto-detects the model type and logs it. +# - Passes the data unchanged. +# """ + +# @classmethod +# def INPUT_TYPES(cls): +# return {"required": {"data": (any_typ,), "model": (any_typ,)}} + +# RETURN_TYPES = (any_typ,) +# RETURN_NAMES = ("output",) +# FUNCTION = "run" +# CATEGORY = "Memory Management" + +# def run(self, data, model): +# if model is None: +# print("โš ๏ธ No model provided to delete") +# return (data,) + +# # Get memory stats before deletion +# before_vram = torch.cuda.memory_allocated() if torch.cuda.is_available() else 0 +# before_reserved = torch.cuda.memory_reserved() if torch.cuda.is_available() else 0 +# before_ram = psutil.virtual_memory().percent + +# # Remove from ComfyUI management and free memory +# deleted_type = purge_specific_from_comfy(model) +# hard_free(model) + +# # Get memory stats after deletion +# after_vram = torch.cuda.memory_allocated() if torch.cuda.is_available() else 0 +# after_reserved = torch.cuda.memory_reserved() if torch.cuda.is_available() else 0 +# after_ram = psutil.virtual_memory().percent + +# # Log results +# print(f"๐Ÿ—‘๏ธ Deleted model type: {deleted_type}") +# print(f"System RAM change: {before_ram - after_ram:.2f}%") + +# if torch.cuda.is_available(): +# vram_freed = (before_vram - after_vram) / (1024 * 1024 * 1024) +# reserved_freed = (before_reserved - after_reserved) / (1024 * 1024 * 1024) +# print(f"GPU allocated freed: {vram_freed:.3f} GB") +# print(f"GPU reserved freed: {reserved_freed:.3f} GB") + +# return (data,) + + +# # Node mappings +# NODE_CLASS_MAPPINGS = {"DeleteModelPassthrough": DeleteModelPassthrough} +# NODE_DISPLAY_NAME_MAPPINGS = {"DeleteModelPassthrough": "Delete Model (Passthrough Any)"} + + +# Testing Deepseek's more rigorous version + import torch, gc, psutil import comfy.model_management as mm from comfy.model_management import current_loaded_models, LoadedModel @@ -48,6 +198,23 @@ def hard_free(model): mm.soft_empty_cache() +def print_tracked_models(): + """Print all models currently tracked by the model manager""" + print("๐Ÿ“‹ Currently tracked models:") + if not current_loaded_models: + print(" No models currently tracked") + return + + for i, loaded_model in enumerate(current_loaded_models): + try: + model_obj = loaded_model.model + model_type = type(model_obj).__name__ + model_id = id(model_obj) + print(f" {i}: {model_type} (id: {model_id})") + except Exception as e: + print(f" {i}: [Error accessing model: {e}]") + + def purge_specific_from_comfy(model_obj): """ Remove a specific model from ComfyUI's current_loaded_models list. @@ -56,9 +223,14 @@ def purge_specific_from_comfy(model_obj): deleted_type = "Unknown" try: + # Print initial state + print("๐Ÿ” Before removal:") + print_tracked_models() + # Check if it's a LoadedModel wrapper if isinstance(model_obj, LoadedModel): if model_obj in current_loaded_models: + print(f"๐Ÿ—‘๏ธ Removing LoadedModel wrapper: {type(model_obj.model).__name__}") current_loaded_models.remove(model_obj) # Also detach the underlying model if hasattr(model_obj, 'model_unload'): @@ -67,13 +239,30 @@ def purge_specific_from_comfy(model_obj): # Check if it's a raw model that might be wrapped in LoadedModel else: + model_removed = False for i, loaded_model in enumerate(current_loaded_models): if loaded_model.model is model_obj: + print(f"๐Ÿ—‘๏ธ Removing raw model: {type(model_obj).__name__}") current_loaded_models.pop(i) if hasattr(loaded_model, 'model_unload'): loaded_model.model_unload() deleted_type = f"RawModel ({type(model_obj).__name__})" + model_removed = True break + + if not model_removed: + print("โš ๏ธ Model not found in current_loaded_models") + # Check if it's a different type of model reference + model_id = id(model_obj) + for i, loaded_model in enumerate(current_loaded_models): + loaded_model_id = id(loaded_model.model) if hasattr(loaded_model, 'model') else None + if loaded_model_id == model_id: + print(f"๐Ÿ—‘๏ธ Removing model by ID match: {type(model_obj).__name__}") + current_loaded_models.pop(i) + if hasattr(loaded_model, 'model_unload'): + loaded_model.model_unload() + deleted_type = f"RawModel ({type(model_obj).__name__})" + break # Additional cleanup for common model types if hasattr(model_obj, '__class__'): @@ -85,8 +274,12 @@ def purge_specific_from_comfy(model_obj): elif 'vae' in cls_name: deleted_type = "VAE" + # Print final state + print("๐Ÿ” After removal:") + print_tracked_models() + except Exception as e: - print(f"Cache purge error: {e}") + print(f"โŒ Cache purge error: {e}") return deleted_type @@ -119,6 +312,10 @@ class DeleteModelPassthrough: before_reserved = torch.cuda.memory_reserved() if torch.cuda.is_available() else 0 before_ram = psutil.virtual_memory().percent + print("=" * 50) + print("๐Ÿงน Starting model deletion process") + print("=" * 50) + # Remove from ComfyUI management and free memory deleted_type = purge_specific_from_comfy(model) hard_free(model) @@ -129,18 +326,24 @@ class DeleteModelPassthrough: after_ram = psutil.virtual_memory().percent # Log results + print("=" * 50) + print("๐Ÿ“Š Deletion Results:") + print("=" * 50) print(f"๐Ÿ—‘๏ธ Deleted model type: {deleted_type}") - print(f"System RAM change: {before_ram - after_ram:.2f}%") + print(f"๐Ÿ’พ System RAM change: {before_ram - after_ram:.2f}%") if torch.cuda.is_available(): vram_freed = (before_vram - after_vram) / (1024 * 1024 * 1024) reserved_freed = (before_reserved - after_reserved) / (1024 * 1024 * 1024) - print(f"GPU allocated freed: {vram_freed:.3f} GB") - print(f"GPU reserved freed: {reserved_freed:.3f} GB") + print(f"๐ŸŽฎ GPU allocated freed: {vram_freed:.3f} GB") + print(f"๐ŸŽฎ GPU reserved freed: {reserved_freed:.3f} GB") + + print("=" * 50) + print("โœ… Deletion process completed") + print("=" * 50) return (data,) -# Node mappings NODE_CLASS_MAPPINGS = {"DeleteModelPassthrough": DeleteModelPassthrough} NODE_DISPLAY_NAME_MAPPINGS = {"DeleteModelPassthrough": "Delete Model (Passthrough Any)"}