From 7ba6f9dae8ba8b9cbb8c88072b41e1fd17ddb177 Mon Sep 17 00:00:00 2001 From: Isi Date: Thu, 28 Aug 2025 19:42:20 +0100 Subject: [PATCH] Update delete_model_passthrough.py --- delete_model_passthrough.py | 116 +++++++++++++++++++++++------------- 1 file changed, 73 insertions(+), 43 deletions(-) diff --git a/delete_model_passthrough.py b/delete_model_passthrough.py index 1715dcd..e7db62f 100644 --- a/delete_model_passthrough.py +++ b/delete_model_passthrough.py @@ -1,5 +1,6 @@ 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 @@ -7,23 +8,34 @@ any_typ = AnyType("*") def hard_free(model): + """ + Forcefully free model memory by detaching parameters and clearing caches. + """ try: - # Explicitly nuke parameters and buffers + # 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 state_dict if exists - if hasattr(model, "state_dict"): - sd = model.state_dict() - for k in list(sd.keys()): - del sd[k] - del sd + + # 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}") @@ -33,44 +45,49 @@ def hard_free(model): pass gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.ipc_collect() + mm.soft_empty_cache() -def purge_specific_from_comfy(model): +def purge_specific_from_comfy(model_obj): """ - Surgically remove THIS model object from ComfyUI's caches - without touching unrelated ones. Returns a type string. + Remove a specific model from ComfyUI's current_loaded_models list. + Returns the model type if found and removed. """ deleted_type = "Unknown" - + try: - if mm.current_model == model: - mm.current_model = None - deleted_type = "Diffusion/UNet" - - if hasattr(mm, "loaded_models"): - for k, v in list(mm.loaded_models.items()): - if v is model: - del mm.loaded_models[k] - deleted_type = "Diffusion/UNet" - - if hasattr(mm, "loaded_clip"): - for k, v in list(mm.loaded_clip.items()): - if v is model: - del mm.loaded_clip[k] - deleted_type = "CLIP" - - if hasattr(mm, "loaded_vae"): - for k, v in list(mm.loaded_vae.items()): - if v is model: - del mm.loaded_vae[k] - deleted_type = "VAE" - + # 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 @@ -93,24 +110,37 @@ class DeleteModelPassthrough: 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_res = torch.cuda.memory_reserved() if torch.cuda.is_available() else 0 - before_ram = psutil.virtual_memory().percent + 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_res = torch.cuda.memory_reserved() if torch.cuda.is_available() else 0 - after_ram = psutil.virtual_memory().percent + 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}%") - print(f"GPU allocated freed: {(before_vram - after_vram)/1e9:.2f} GB") - print(f"GPU reserved freed: {(before_res - after_res)/1e9:.2f} GB") + + 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)"}