diff --git a/delete_model_passthrough.py b/delete_model_passthrough.py index 38562e1..1715dcd 100644 --- a/delete_model_passthrough.py +++ b/delete_model_passthrough.py @@ -1,32 +1,91 @@ -import torch -import gc -import psutil -import os -import ctypes +import torch, gc, psutil +import comfy.model_management as mm class AnyType(str): - def __ne__(self, __value: object) -> bool: - return False - + def __ne__(self, __value: object) -> bool: return False any_typ = AnyType("*") +def hard_free(model): + try: + # Explicitly nuke parameters and buffers + if hasattr(model, "parameters"): + for p in model.parameters(): + if p is not None: + p.detach_() + del p + if hasattr(model, "buffers"): + for b in model.buffers(): + if b is not None: + 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 + except Exception as e: + print(f"Hard free error: {e}") + + try: + del model + except: + pass + + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + + +def purge_specific_from_comfy(model): + """ + Surgically remove THIS model object from ComfyUI's caches + without touching unrelated ones. Returns a type string. + """ + 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" + + except Exception as e: + print(f"Cache purge error: {e}") + + return deleted_type + + class DeleteModelPassthrough: """ - ComfyUI Custom Node that: - - Accepts any input (data) and a MODEL (`model`). - - Deletes the MODEL completely (VRAM + RAM). - - Passes through the data unchanged. + 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 {"required": {"data": (any_typ,), "model": (any_typ,)}} RETURN_TYPES = (any_typ,) RETURN_NAMES = ("output",) @@ -34,38 +93,24 @@ class DeleteModelPassthrough: CATEGORY = "Memory Management" def run(self, data, model): + 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 = psutil.virtual_memory().percent - collected = gc.collect() - after = psutil.virtual_memory().percent - freedr = before - after - print(f"System RAM freed: {freedr:.2f}%, GC collected {collected} objects.") + deleted_type = purge_specific_from_comfy(model) + hard_free(model) - if not torch.cuda.is_available(): - return "CUDA not available." - initial = torch.cuda.memory_allocated() - print(f"Initial GPU memory: {initial}") - - try: - del model - torch.cuda.empty_cache() - gc.collect() - print("Deleted MODEL completely from VRAM.") - except Exception as e: - print(f"Failed to delete model: {e}") + 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 - final = torch.cuda.memory_allocated() - print(f"Final GPU memory: {final}") - freed = max(0, initial - final) - print(f"GPU VRAM freed: {freed/1e9:.2f} GB") + 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") return (data,) -NODE_CLASS_MAPPINGS = { - "DeleteModelPassthrough": DeleteModelPassthrough, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "DeleteModelPassthrough": "Delete Model (Passthrough Any)", -} +NODE_CLASS_MAPPINGS = {"DeleteModelPassthrough": DeleteModelPassthrough} +NODE_DISPLAY_NAME_MAPPINGS = {"DeleteModelPassthrough": "Delete Model (Passthrough Any)"}