diff --git a/delete_model_passthrough.py b/delete_model_passthrough.py index a65ffee..b20df5b 100644 --- a/delete_model_passthrough.py +++ b/delete_model_passthrough.py @@ -160,26 +160,53 @@ any_typ = AnyType("*") def hard_free(model): """ Forcefully free model memory by detaching parameters and clearing caches. + Special handling for different model types. """ + if model is None: + return + try: - # Clear parameters and buffers + model_type = type(model).__name__ + print(f"๐Ÿ”ง Freeing {model_type} model...") + + # Handle ModelPatcher objects (common for CLIP, UNet, VAE) + if model_type == 'ModelPatcher' and hasattr(model, 'model'): + print(" Detaching ModelPatcher parameters...") + real_model = model.model + if hasattr(real_model, "parameters"): + for p in real_model.parameters(): + if p is not None: + p.detach_() + del p + model.model = None + + # Handle ControlNet models (they have different structure) + elif 'Control' in model_type or 'control' in model_type.lower(): + print(" Detaching ControlNet parameters...") + if hasattr(model, "parameters"): + for p in model.parameters(): + if p is not None: + p.detach_() + del p + if hasattr(model, "weights"): + model.weights = None + if hasattr(model, "model"): + model.model = None + + # Standard cleanup for any model 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 + # Clear any cached tensor attributes for attr_name in list(vars(model).keys()): attr = getattr(model, attr_name) if isinstance(attr, torch.Tensor): @@ -187,7 +214,7 @@ def hard_free(model): delattr(model, attr_name) except Exception as e: - print(f"Hard free error: {e}") + print(f"โŒ Hard free error: {e}") try: del model @@ -199,98 +226,117 @@ def hard_free(model): def print_tracked_models(): - """Print all models currently tracked by the model manager""" - print("๐Ÿ“‹ Currently tracked models:") + """Print all models currently tracked by ComfyUI's main system""" + print("๐Ÿ“‹ Models in current_loaded_models:") if not current_loaded_models: - print(" No models currently tracked") - return + print(" No models tracked in main system") + return 0 + count = 0 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})") + if hasattr(loaded_model, 'model'): + model_obj = loaded_model.model + model_type = type(model_obj).__name__ + model_id = id(model_obj) + print(f" {i}: {model_type} (id: {model_id})") + count += 1 except Exception as e: - print(f" {i}: [Error accessing model: {e}]") + print(f" {i}: [Error: {e}]") + + return count + + +def identify_model_type(model_obj): + """Identify what type of model this is""" + if model_obj is None: + return "Unknown" + + 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: + # Check what's inside the ModelPatcher + if hasattr(model_obj, 'model') and model_obj.model is not None: + inner_type = model_obj.model.__class__.__name__.lower() + if 'clip' in inner_type: + return "CLIP (ModelPatcher)" + elif 'unet' in inner_type: + return "UNet (ModelPatcher)" + elif 'vae' in inner_type: + return "VAE (ModelPatcher)" + return "ModelPatcher" + elif 'diffusion' in cls_name: + return "DiffusionModel" + else: + return f"Unknown ({model_obj.__class__.__name__})" 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. + Try to remove model from ComfyUI's tracking systems. + Returns the identified model type. """ - deleted_type = "Unknown" + if model_obj is None: + return "None" + + model_type = identify_model_type(model_obj) + model_removed = False + + print(f"๐Ÿ” Target model type: {model_type}") + print(f"๐Ÿ” Model object type: {type(model_obj).__name__}") + print(f"๐Ÿ” Model ID: {id(model_obj)}") + + # Print current tracking state + initial_count = print_tracked_models() 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'): - 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: - model_removed = False + # Try to remove from main tracking system (for CLIP/UNet/VAE) + if current_loaded_models: 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__}") + try: + # Check if this LoadedModel contains our target model + if (hasattr(loaded_model, 'model') and + (loaded_model.model is model_obj or + (hasattr(model_obj, 'model') and loaded_model.model is model_obj.model))): + + print(f"๐Ÿ—‘๏ธ Removing from tracked models: {type(loaded_model.model).__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 + except: + continue - # 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" - - # Print final state - print("๐Ÿ” After removal:") - print_tracked_models() - + if not model_removed: + print("โš ๏ธ Model not in main tracking system - may use different management") + + # For ControlNet and other specially managed models, we can't remove them from tracking + # but we can still free their memory + if 'control' in model_type.lower(): + print("๐Ÿ’ก ControlNet models often use custom management systems") + except Exception as e: - print(f"โŒ Cache purge error: {e}") + print(f"โŒ Error during purge attempt: {e}") - return deleted_type + # Print final tracking state + final_count = print_tracked_models() + print(f"๐Ÿ“Š Tracking change: {initial_count} โ†’ {final_count} models") + + return model_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. + ComfyUI Custom Node: Deletes models from memory regardless of tracking system. """ @classmethod @@ -312,11 +358,11 @@ class DeleteModelPassthrough: before_reserved = torch.cuda.memory_reserved() if torch.cuda.is_available() else 0 before_ram = psutil.virtual_memory().percent - print("=" * 50) + print("=" * 60) print("๐Ÿงน Starting model deletion process") - print("=" * 50) + print("=" * 60) - # Remove from ComfyUI management and free memory + # Try to remove from tracking and free memory deleted_type = purge_specific_from_comfy(model) hard_free(model) @@ -326,24 +372,35 @@ class DeleteModelPassthrough: after_ram = psutil.virtual_memory().percent # Log results - print("=" * 50) - print("๐Ÿ“Š Deletion Results:") - print("=" * 50) + print("=" * 60) + print("๐Ÿ“Š Memory Free Results:") + print("=" * 60) 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"๐ŸŽฎ Final allocated: {after_vram / (1024 * 1024 * 1024):.3f} GB") + print(f"๐ŸŽฎ Final reserved: {after_reserved / (1024 * 1024 * 1024):.3f} GB") - print("=" * 50) - print("โœ… Deletion process completed") - print("=" * 50) + print("=" * 60) + + # Success determination based on memory freed + if torch.cuda.is_available() and reserved_freed > 0.1: # Freed at least 100MB + print("โœ… SUCCESS: Significant memory freed!") + elif torch.cuda.is_available() and reserved_freed > 0: + print("โš ๏ธ PARTIAL: Some memory freed, but model may still be referenced") + else: + print("โŒ FAILED: No significant memory freed - model may be already freed or invalid") + + print("=" * 60) return (data,) +# Node mappings NODE_CLASS_MAPPINGS = {"DeleteModelPassthrough": DeleteModelPassthrough} NODE_DISPLAY_NAME_MAPPINGS = {"DeleteModelPassthrough": "Delete Model (Passthrough Any)"}