delete_model_passthrough.py

import torch, gc, psutil
import comfy.model_management as mm
from comfy.model_management import loaded_models, free_memory, get_torch_device
from comfy.comfy_types import IO
from nodes import ControlNetLoader, VAELoader, UNETLoader, CLIPTextEncode
from comfy_extras.nodes_model_patch import ModelPatchLoader

try:
    from custom_nodes.ComfyUI_GGUF.nodes import UnetLoaderGGUF
    GGUF_AVAILABLE = True
except ImportError:
    print("⚠️ ComfyUI_GGUF not available - ControlledUnetLoaderGGUF will not work")
    GGUF_AVAILABLE = False
    UnetLoaderGGUF = None

class AnyType(str):
    def __ne__(self, __value: object) -> bool: return False
any_typ = AnyType("*")


def hard_free_model(model):
    if model is None:
        return
        
    try:
        model_type = type(model).__name__
        print(f"Freeing {model_type} model...")
        
        # Handle dictionary-style models (common in some workflows)
        if isinstance(model, dict):
            for key, value in list(model.items()):
                if hasattr(value, 'parameters') or hasattr(value, 'model'):
                    hard_free_model(value)
                del model[key]
            return
        
        # Handle ModelPatcher objects (common in ComfyUI)
        if hasattr(model, 'model') and model.model is not None:
            inner_model = model.model
            if hasattr(inner_model, "parameters"):
                for p in inner_model.parameters():
                    if p is not None:
                        p.detach_()
                        del p
            model.model = None
        
        # CLIP-specific cleanup - some CLIP models have additional attributes
        if hasattr(model, "transformer") and model.transformer is not None:
            if hasattr(model.transformer, "parameters"):
                for p in model.transformer.parameters():
                    if p is not None:
                        p.detach_()
                        del p
            model.transformer = None
            
        # Handle tokenizer if present (some CLIP implementations)
        if hasattr(model, "tokenizer"):
            model.tokenizer = None
            
        # Standard parameter cleanup
        if hasattr(model, "parameters"):
            for p in model.parameters():
                if p is not None:
                    p.detach_()
                    del p
        
        # Clear buffers and tensor attributes
        if hasattr(model, "buffers"):
            for b in model.buffers():
                if b is not None:
                    b.detach_()
                    del b
        
        # Clear any tensor 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"❌ Error during model freeing: {e}")

def identify_model_type(model_obj):
    """Identify what type of model this is"""
    if model_obj is None:
        return "Unknown"
    
    if isinstance(model_obj, dict):
        model_types = []
        for key, value in model_obj.items():
            if hasattr(value, '__class__'):
                model_types.append(f"{key}:{value.__class__.__name__}")
        return f"DictContainer[{', '.join(model_types)}]"
    
    cls_name = model_obj.__class__.__name__.lower()
    
    if 'clip' in cls_name:
        return "CLIP"
    elif 'unet' in cls_name:
        return "UNet"
    elif 'vae' in cls_name:
        return "VAE"
    elif 'control' in cls_name:
        return "ControlNet"
    elif 'modelpatcher' in cls_name:
        return "ModelPatcher"
    elif 'diffusion' in cls_name:
        return "DiffusionModel"
    else:
        return f"Unknown ({model_obj.__class__.__name__})"


def print_currently_loaded():
    """Print models currently loaded in ComfyUI's management system"""
    current_models = mm.loaded_models()
    print("📋 Models in loaded_models():")
    if not current_models:
        print("   No models currently managed")
        return 0
    
    for i, model in enumerate(current_models):
        try:
            model_type = identify_model_type(model)
            # Remove the problematic formatting entirely
            if hasattr(model, 'model_memory'):
                memory_used = model.model_memory()
                print(f"   {i}: {model_type} ({memory_used})")
            else:
                print(f"   {i}: {model_type} (unknown memory)")
        except Exception as e:
            print(f"   {i}: [Error: {e}]")
    
    return len(current_models)

class DeleteModelPassthrough:
    """
    ComfyUI Custom Node: Properly deletes models using ComfyUI's memory management system
    """

    @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


        
        model_type = identify_model_type(model)
        # print(f"Target model type: {model_type}")
        
        # Print current state
        initial_count = print_currently_loaded()
        
        # Try to remove from ComfyUI's management system
        model_removed = False
        current_models = mm.loaded_models()
        
        if model in current_models:
            # print("🗑Removing model from ComfyUI management...")
            current_models.remove(model)
            model_removed = True
        else:
            # Check if it's a ModelPatcher or wrapped model
            for managed_model in current_models:
                try:
                    if (hasattr(managed_model, 'model') and 
                        (managed_model.model is model or 
                         (hasattr(model, 'model') and managed_model.model is model.model))):
                        # print("🗑Removing wrapped model from ComfyUI management...")
                        current_models.remove(managed_model)
                        model_removed = True
                        break
                except:
                    continue
        
        # Free memory using ComfyUI's proper methods
        # print("Freeing memory using ComfyUI's system...")
        mm.free_memory(1e30, mm.get_torch_device(), mm.loaded_models())
        
        # Additional forceful cleanup
        # print("Forceful cleanup...")
        hard_free_model(model)

        try:
            del model
        except:
            pass

        
        # ComfyUI's cache cleanup
        mm.soft_empty_cache(force=True)
        
        # Standard Python cleanup
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()

        # 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

        # Print final state
        final_count = print_currently_loaded()
        
        print(f"Managed models: {initial_count} → {final_count}")
        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"Final allocated: {after_vram / (1024 * 1024 * 1024):.3f} GB")
            # print(f"Final reserved: {after_reserved / (1024 * 1024 * 1024):.3f} GB")
        
        print("=" * 60)
        
        # Success determination
        if model_removed:
            print("SUCCESS: Model removed from management system!")
        elif torch.cuda.is_available() and reserved_freed > 0.1:
            print("SUCCESS: Significant memory freed (model may use custom management)")
        else:
            print("⚠PARTIAL: Model may still be referenced somewhere")
        
        print("=" * 60)

        return (data,)


def safe_hard_free_model(model):
    """Memory-safe model freeing without expensive operations"""
    if model is None:
        return
        
    try:
        # Minimal identification without memory overhead
        model_type = model.__class__.__name__
        
        # Handle dictionary-style models - but do it safely
        if isinstance(model, dict):
            # Clear dictionary without iterating through values
            model.clear()
            return
        
        # Quick and dirty cleanup - avoid iterating through parameters if possible
        if hasattr(model, 'model') and model.model is not None:
            # Just nullify the reference instead of iterating parameters
            model.model = None
        
        # Nullify common heavy attributes without iteration
        for attr_name in ['transformer', 'tokenizer', 'encoder', 'decoder']:
            if hasattr(model, attr_name):
                setattr(model, attr_name, None)
                
        # If we must detach, do it in chunks with garbage collection
        if hasattr(model, "parameters"):
            param_count = 0
            for p in model.parameters():
                if p is not None:
                    try:
                        p.detach_()
                    except:
                        pass
                    param_count += 1
                    # Periodic garbage collection to prevent memory buildup
                    if param_count % 1000 == 0:
                        gc.collect()
        
        # Clear the object itself
        try:
            del model
        except:
            pass
            
    except Exception as e:
        # Silent error handling to avoid memory overhead
        pass

def safe_print_loaded_models():
    """Memory-safe way to print loaded models"""
    try:
        current_models = mm.loaded_models()
        if not current_models:
            print("No models currently managed")
            return 0
        
        print(f"Managed models count: {len(current_models)}")
        # Just show count, avoid accessing model properties
        return len(current_models)
        
    except:
        return 0

def quick_identify_model(model_obj):
    """Quick identification without memory overhead"""
    if model_obj is None:
        return "Unknown"
    return model_obj.__class__.__name__



class DeleteModelPassthroughLight:
    @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:
            return (data,)
            
        # Get memory stats BEFORE any operations
        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
        
        # Quick identification
        model_type = quick_identify_model(model)
        
        # Safe model counting
        initial_count = safe_print_loaded_models()
        
        # Remove ONLY the target model from management
        model_removed = False
        try:
            current_models = mm.loaded_models()
            if model in current_models:
                current_models.remove(model)
                model_removed = True
                print(f"Removed target model from tracking")
        except Exception as e:
            print(f"Error removing from tracking: {e}")
    
        # Don't use free_memory() with huge values as it will delete other models
        
        # Instead, manually free only our target model
        safe_hard_free_model(model)
        
        # Optional: Use free_memory() correctly if we know how much this model used
        # But better to avoid it entirely for single-model deletion
        
        # Gentle cleanup without affecting other models
        try:
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        except:
            pass
        
        # 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
        
        # Safe final count
        final_count = safe_print_loaded_models()
        
        # Check if other models were affected
        other_models_affected = initial_count - final_count > 1 if model_removed else initial_count - final_count > 0
        
        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"Freed: {reserved_freed:.3f}GB | Models: {initial_count}→{final_count}")
            
            if other_models_affected:
                print(f"⚠️ WARNING: {initial_count - final_count} models were affected, not just the target!")
        
        return (data,)

class SmartClipDeleter:
    """
    Automatically detects if CLIP is in RAM or VRAM and handles appropriately
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"data": (any_typ,), "clip_model": ("CLIP",)}}
 
    RETURN_TYPES = (any_typ,)
    FUNCTION = "delete_smart"
    CATEGORY = "Memory Management"

    def delete_smart(self, data, clip_model):
        if clip_model is None:
            return (data,)
            
        # Detect where the model is loaded
        model_in_vram = False
        model_in_ram = False
        
        # Check if model has parameters on GPU
        if hasattr(clip_model, "parameters"):
            for p in clip_model.parameters():
                if p is not None and p.device.type != 'cpu':
                    model_in_vram = True
                    break
                elif p is not None:
                    model_in_ram = True
        
        print(f"📍 CLIP model location: {'VRAM' if model_in_vram else 'RAM'}")
        
        if model_in_vram:
            # Use VRAM cleanup approach
            try:
                if hasattr(clip_model, 'model_memory'):
                    model_size = clip_model.model_memory()
                    mm.free_memory(model_size, mm.get_torch_device())
                else:
                    mm.soft_empty_cache()
            except:
                pass
                
        else:
            # Use RAM cleanup approach
            try:
                # Clear model attributes
                for attr in ['transformer', 'tokenizer', 'text_model', 'visual', 'model']:
                    if hasattr(clip_model, attr):
                        setattr(clip_model, attr, None)
                
                # Force garbage collection
                gc.collect()
                gc.collect()
                
            except:
                pass
    
        return (data,)

class ControlledControlNetLoader:
    @classmethod
    def INPUT_TYPES(s):
        # Get the original input types and add trigger
        original_types = ControlNetLoader.INPUT_TYPES()
        if "required" in original_types:
            original_types["required"]["trigger"] = (any_typ, {"default": None})
        else:
            original_types["required"] = {"trigger": (any_typ, {"default": None})}
        return original_types

    RETURN_TYPES = ("CONTROL_NET",)
    FUNCTION = "load_controlnet"
    CATEGORY = "Memory Management"
    TITLE = "Controlled ControlNet Loader"

    def load_controlnet(self, trigger, *args, **kwargs):
        if trigger is None:
            print("⏸️  ControlNet loading paused - no trigger received")
            return (None,)
        
        print(f"🚀 Loading ControlNet...")
        # Simply call the original class method
        return ControlNetLoader.load_controlnet(self, *args, **kwargs)


class ControlledVAELoader:
    @classmethod
    def INPUT_TYPES(cls):
        # Get the original input types and add trigger
        original_types = VAELoader.INPUT_TYPES()
        original_types = original_types.copy()
        original_types["required"] = dict(original_types.get("required", {}))

        original_types["required"]["trigger"] = (any_typ, {"default": None})
        return original_types

    RETURN_TYPES = ("VAE",)
    FUNCTION = "load_vae"
    CATEGORY = "Memory Management"
    TITLE = "Controlled VAE Loader"

    def load_vae(self, vae_name, trigger=None):
        if trigger is None:
            print("VAE loading paused - no trigger received")
            return (None,)

        print("Loading VAE...")
        loader = VAELoader()
        return loader.load_vae(vae_name)



class ControlledUnetLoaderGGUF:
    @classmethod
    def INPUT_TYPES(s):
        if not GGUF_AVAILABLE:
            return {"required": {"trigger": (any_typ, {"default": None})}}
        
        # Get the original input types and add trigger
        original_types = UnetLoaderGGUF.INPUT_TYPES()
        if "required" in original_types:
            original_types["required"]["trigger"] = (any_typ, {"default": None})
        else:
            original_types["required"] = {"trigger": (any_typ, {"default": None})}
        return original_types

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load_unet"
    CATEGORY = "Memory Management"
    TITLE = "Controlled UNet Loader (GGUF)"

    def load_unet(self, trigger, *args, **kwargs):
        if not GGUF_AVAILABLE:
            print("❌ ComfyUI_GGUF not installed - cannot load UNet")
            return (None,)
            
        if trigger is None:
            print("⏸UNet loading paused - no trigger received")
            return (None,)
        
        print(f"Loading UNet...")
        # Simply call the original class method
        return UnetLoaderGGUF.load_unet(self, *args, **kwargs)


class ControlledModelPatchLoader:
    @classmethod
    def INPUT_TYPES(s):
        # Get the original input types and add trigger
        original_types = ModelPatchLoader.INPUT_TYPES()
        if "required" in original_types:
            original_types["required"]["trigger"] = (any_typ, {"default": None})
        else:
            original_types["required"] = {"trigger": (any_typ, {"default": None})}
        return original_types

    RETURN_TYPES = ("MODEL_PATCH",)
    FUNCTION = "load_model_patch"
    CATEGORY = "Memory Management"
    TITLE = "Controlled Model Patch Loader"
    EXPERIMENTAL = True

    def load_model_patch(self, trigger, *args, **kwargs):
        if trigger is None:
            print("⏸️  Model Patch loading paused - no trigger received")
            return (None,)
        
        print(f"🚀 Loading Model Patch...")
        # Simply call the original class method
        return ModelPatchLoader.load_model_patch(self, *args, **kwargs)

class ControlledUNETLoader:
    @classmethod
    def INPUT_TYPES(s):
        # Get original input types and add trigger
        original_types = UNETLoader.INPUT_TYPES()

        if "required" in original_types:
            original_types["required"]["trigger"] = (any_typ, {"default": None})
        else:
            original_types["required"] = {"trigger": (any_typ, {"default": None})}

        return original_types

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load_unet"
    CATEGORY = "Memory Management"
    TITLE = "Controlled UNet Loader"

    def load_unet(self, trigger, *args, **kwargs):
        # Pause loading if no trigger received
        if trigger is None:
            print("⏸️  UNet loading paused — no trigger received")
            return (None,)

        print("🚀 Loading UNet...")

        # Forward call to original loader
        return UNETLoader.load_unet(self, *args, **kwargs)

class ControlledCLIPTextEncode:
    @classmethod
    def INPUT_TYPES(s):
        # Get original input types and add trigger
        original_types = CLIPTextEncode.INPUT_TYPES()

        if "required" in original_types:
            original_types["required"]["trigger"] = (any_typ, {"default": None})
        else:
            original_types["required"] = {"trigger": (any_typ, {"default": None})}

        return original_types

    RETURN_TYPES = (IO.CONDITIONING,)
    FUNCTION = "encode"
    CATEGORY = "Memory Management"
    TITLE = "Controlled CLIP Text Encode"
    DESCRIPTION = "A trigger-controlled wrapper around CLIPTextEncode."

    def encode(self, trigger, *args, **kwargs):
        # Pause encoding if no trigger value given
        if trigger is None:
            print("⏸️  CLIP text encoding paused — no trigger received")
            return (None,)

        print("🚀 Encoding text with CLIP...")

        # Forward to original class method
        return CLIPTextEncode.encode(self, *args, **kwargs)


# Experimental Factory code
# def make_trigger_controlled_node(original_class, title_suffix=" (Controlled)", category="Memory Management"):
#     """
#     Creates a trigger-controlled wrapper class around any ComfyUI node class.
#     """

#     original_function_name = getattr(original_class, "FUNCTION", None)
#     if original_function_name is None:
#         raise ValueError(f"{original_class.__name__} does not define FUNCTION")

#     # --------------------------
#     # 1. Build INPUT_TYPES wrapper
#     # --------------------------
#     @classmethod
#     def INPUT_TYPES(cls):
#         original_types = original_class.INPUT_TYPES()

#         # Insert trigger field
#         if "required" in original_types:
#             original_types["required"]["trigger"] = (any_typ, {"default": None})
#         else:
#             original_types["required"] = {"trigger": (any_typ, {"default": None})}

#         return original_types

#     # --------------------------
#     # 2. Wrapped FUNCTION method
#     # --------------------------
#     def wrapped_function(self, trigger, *args, **kwargs):
#         if trigger is None:
#             print(f"⏸️  {original_class.__name__} paused — no trigger received")
#             return (None,)
        
#         print(f"🚀 Executing {original_class.__name__}...")
#         original_fn = getattr(original_class, original_function_name)
#         return original_fn(self, *args, **kwargs)

#     # --------------------------
#     # 3. Build dynamic class attrs
#     # --------------------------
#     attrs = {
#         "INPUT_TYPES": INPUT_TYPES,
#         "FUNCTION": original_function_name,
#         original_function_name: wrapped_function,
#         "RETURN_TYPES": getattr(original_class, "RETURN_TYPES", ()),
#         "CATEGORY": category,
#         "TITLE": getattr(original_class, "TITLE", original_class.__name__ + title_suffix),
#     }

#     # Copy DESCRIPTION if present
#     if hasattr(original_class, "DESCRIPTION"):
#         attrs["DESCRIPTION"] = getattr(original_class, "DESCRIPTION")

#     # --------------------------
#     # 4. Create new class
#     # --------------------------
#     new_class_name = f"Controlled_{original_class.__name__}"

#     return type(new_class_name, (original_class,), attrs)
    
# How to use 'Experimental Factory code':
# ControlledUNETLoader = make_trigger_controlled_node(UNETLoader)


NODE_CLASS_MAPPINGS = {
    "DeleteModelPassthroughLight": DeleteModelPassthroughLight,
    "DeleteModelPassthrough": DeleteModelPassthrough,
    "ControlledUNETLoader": ControlledUNETLoader,
    "ControlledUnetLoaderGGUF": ControlledUnetLoaderGGUF,
    "ControlledControlNetLoader": ControlledControlNetLoader,
    "ControlledVAELoader": ControlledVAELoader,
    "ControlledCLIPTextEncode": ControlledCLIPTextEncode,
    "ControlledModelPatchLoader": ControlledModelPatchLoader,
    "SmartClipDeleter": SmartClipDeleter
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "DeleteModelPassthroughLight": "Delete Model Light (Passthrough Any)",
    "DeleteModelPassthrough": "Delete Model (Passthrough Any)",
    "ControlledUNETLoader": "Controlled Load Diffusion Model",
    "ControlledUnetLoaderGGUF": "Controlled UNet Loader (GGUF)",
    "ControlledControlNetLoader": "Controlled ControlNet Loader",
    "ControlledVAELoader": "Controlled VAE Loader",
    "ControlledCLIPTextEncode": "Controlled CLIP Text Encode (Prompt)",
    "ControlledModelPatchLoader": "Controlled ModelPatch Loader",
    "SmartClipDeleter": "Smart CLIP Deleter (Auto-Detect)"
}