From 6bee379ab2fe421cb8daa7ae45d3eac9e5dc269b Mon Sep 17 00:00:00 2001 From: obisin <90648621+obisin@users.noreply.github.com> Date: Sat, 30 Aug 2025 18:08:55 +0100 Subject: [PATCH] Update dgls_model_loader.py --- dgls_model_loader.py | 1774 ++++++------------------------------------ 1 file changed, 243 insertions(+), 1531 deletions(-) diff --git a/dgls_model_loader.py b/dgls_model_loader.py index ece2043..3eb4a8c 100644 --- a/dgls_model_loader.py +++ b/dgls_model_loader.py @@ -1,1595 +1,307 @@ -PROFILE = False - import comfy.sd import folder_paths -import datetime -from functools import wraps -from torch import nn -import time import torch +from safetensors.torch import load_file +import comfy.model_detection import gc -import threading -import warnings -import psutil -import types -import torch._dynamo -from functools import partial -import torch.utils.checkpoint -from collections import defaultdict -import bitsandbytes as bnb -import queue -from concurrent.futures import ThreadPoolExecutor - -if PROFILE == True: - try: - import torch.profiler - import torch.cuda.nvtx as nvtx - except ImportError: - PROFILE = False - nvtx = None - - """ -Dynamic GPU Layer Swapping for ComfyUI by obisin -====================================== -Enhanced inference with automatic memory optimization for diffusion models. +DGLS Model Loader for ComfyUI by obisin +============================ +Universal diffusion model loader with layer extraction for memory optimization. + +Extracts computation layers from UNet/DiT architectures for dynamic GPU swapping. Supports Cosmos, Flux, Wan2.1, Wan2.2, HunyuanVideo, and generic transformer blocks. -PERFORMANCE DESIGN NOTE: -This file uses module-level functions with local variable binding for optimal performance. -Hot path functions bind frequently-accessed attributes to locals at function start (LOAD_FAST) -to avoid Python's attribute lookup overhead (LOAD_ATTR). Call chains are minimized to -reduce function call overhead. - -Architecture: -- Globals: Intentionally used for performance (documented below) -- Hot functions: Bind globals to locals once per call -- Minimal call depth: Direct function calls without abstraction layers -- Factory pattern: Clean call sites with fast execution - -CRITICAL PATH: Functions marked "HOT PATH" - profile before changing! +COMPATIBILITY: Diffusion models and UNets only. Not compatible with full checkpoints yett. Author: obisin """ -class Args: - def __init__(self): - self.dynamic_swapping = True - self.auto = False - self.initial_gpu_layers = 1 # Number of initial layers to keep permanently on GPU - self.final_gpu_layers = 1 # Number of final layers to keep permanently on GPU - self.gpu_layers = None # Comma-separated list of layer indices to keep permanently on GPU - self.prefetch = 0 # Number of layers to prefetch ahead (0=off, 1=single layer, 2+=multiple layers) - self.cuda_streams = False # Enable CUDA streams for copy-compute overlap (needs more VRAM) - self.cast_target = None # Cast FROM dtype TO dtype at start-up (e.g., f32 bf16) - self.verbose = False # Enable verbose output with detailed timing and transfer information - self.cpu_threading = False # Enable CPU threading for async CPU transfers - -args = Args() -TOOLTIPS = { - "dynamic_swapping": "Smart dynamic layer swapping between GPU and CPU for optimal performance", - "initial_gpu_layers": "Number of initial layers to keep permanently on GPU. If not specified, uses reasonable defaults based on estimated VRAM", - "final_gpu_layers": "Number of final layers to keep permanently on GPU", - "gpu_layers": "Comma-separated list of layer indices to keep permanently on GPU (e.g., '0,1,2,14,18,19,20,21,22'). Overrides initial_gpu_layers and final_gpu_layers", - "prefetch": "Number of layers to prefetch ahead of single layer transfer", - "cuda_streams": "Enable CUDA streams for copy-compute overlap (needs more VRAM)", - "cpu_threading": "Enable CPU threading for async CPU transfers (May cause instability on some systems)", - "batch_move": "Use batch layer moving (experimental, may cause device issues)", - "cast_target": "Cast FROM dtype TO dtype at start-up (e.g., f32 bf16) choices=[f32, bf16, f16, f8_e4m3, f8_e5m2, nf4, fp4]", - "verbose": "Enable verbose output with detailed timing and transfer information" -} - # ============================================================================= -# GLOBAL VARIABLES -# ============================================================================= -# Global variables for transfer thread -transfer_queue = None -transfer_thread = None -transfer_active = False -original_functions = {} - -global add_smart_swapping_to_layer -swap_stats = {'to_gpu': 0, 'to_cpu': 0} -transfer_events = {} - -transfer_stats = { - 'to_gpu_times': [], 'to_cpu_times': [], 'to_gpu_speeds': [], 'to_cpu_speeds': [], - 'current_step_gpu_times': [], 'current_step_cpu_times': [], - 'current_step_gpu_speeds': [], 'current_step_cpu_speeds': []} - -overlap_stats = { - 'transfer_start_times': {}, # {layer_idx: start_time} - 'transfer_end_times': {}, # {layer_idx: end_time} - 'compute_start_times': {}, # {layer_idx: start_time} - 'compute_end_times': {}, # {layer_idx: end_time} - 'overlaps': [] # List of overlap measurements -} - -device_cache = None -packed_layers = {} -layer_sizes_mb = {} - -pending_gpu_transfers = {} # {layer_idx: event} -pending_cpu_transfers = {} # {layer_idx: event} - -layers = None - -cpu_swappable_layers = set() -gpu_resident_layers = set() - -GPU_DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' -CPU_DEVICE = 'cpu' - -print(f"Device configuration: GPU={GPU_DEVICE}, CPU={CPU_DEVICE}") -casting_handler = None - -profiler_active = False - -# ============================================================================= -# ANALYSIS FUNCTIONS by obisin +# MODEL LOADER by obisin # ============================================================================= -def analyze_layer_performance(): - """Analyze which layers should be permanent GPU residents based on performance""" - if not (hasattr(add_smart_swapping_to_layer, 'layer_transfer_times') or - hasattr(add_smart_swapping_to_layer, 'layer_compute_times')): - print("No performance data available yet") - return - - print(f"\n=== LAYER PERFORMANCE ANALYSIS (Step {add_smart_swapping_to_layer.current_step}) ===") - - # Calculate current totals - current_total_transfer = sum(sum(times) for times in add_smart_swapping_to_layer.layer_transfer_times.values()) - current_total_compute = sum(sum(times) for times in add_smart_swapping_to_layer.layer_compute_times.values()) - - # Get count for averaging - total_values = sum(len(times) for times in add_smart_swapping_to_layer.layer_transfer_times.values()) - total_values += sum(len(times) for times in add_smart_swapping_to_layer.layer_compute_times.values()) - - # OVERVIEW: Show averages - avg_total = (current_total_transfer + current_total_compute) / total_values if total_values > 0 else 0 - avg_transfer = current_total_transfer / total_values if total_values > 0 else 0 - avg_compute = current_total_compute / total_values if total_values > 0 else 0 - - gpu_mem_gb = torch.cuda.memory_allocated() / 1024 ** 3 if torch.cuda.is_available() else 0 - - current_time = time.time() - total_elapsed = current_time - add_smart_swapping_to_layer.total_start_time - steps_completed = add_smart_swapping_to_layer.current_step - seconds_per_iteration = total_elapsed / steps_completed if steps_completed > 0 else 0 - - - print(f"OVERVIEW: CPU {avg_total:.1f}ms | Transfer {avg_transfer:.1f}ms | Compute {avg_compute:.1f}ms | GPU: {gpu_mem_gb:.2f}GB") - print(f"TIMING: {seconds_per_iteration:.2f}s/it ({steps_completed} steps in {total_elapsed:.1f}s)") - print(f"STATS: GPU Swaps {swap_stats['to_gpu']} | CPU Swaps {swap_stats['to_cpu']} | Prefetch Hits {getattr(add_smart_swapping_to_layer, 'prefetch_hits', 0)} | Prefetch Calls {getattr(add_smart_swapping_to_layer, 'prefetch_misses', 0)}") - - combined_data = {} - - if hasattr(add_smart_swapping_to_layer, 'layer_transfer_times'): - for layer_idx, transfer_times in add_smart_swapping_to_layer.layer_transfer_times.items(): - if transfer_times: - avg_transfer = sum(transfer_times) / len(transfer_times) - combined_data[layer_idx] = {'event_transfer_time': avg_transfer} - - if hasattr(add_smart_swapping_to_layer, 'layer_compute_times'): - for layer_idx, compute_times in add_smart_swapping_to_layer.layer_compute_times.items(): - if compute_times: - avg_compute = sum(compute_times) / len(compute_times) - if layer_idx in combined_data: - combined_data[layer_idx]['event_compute_time'] = avg_compute - else: - combined_data[layer_idx] = {'event_compute_time': avg_compute} - - if combined_data: - # Calculate threshold from accumulated averages - all_cpu_times = [data.get('event_transfer_time', 0) + data.get('event_compute_time', 0) for data in - combined_data.values()] - avg_cpu = sum(all_cpu_times) / len(all_cpu_times) - std_cpu = (sum((x - avg_cpu) ** 2 for x in all_cpu_times) / len(all_cpu_times)) ** 0.5 - cpu_threshold = avg_cpu + std_cpu - - print(f"Threshold (avg + 1 std): {cpu_threshold:.1f}ms") - print() - - # Show problematic layers - problematic_layers = [] - for layer_idx in sorted(combined_data.keys()): - data = combined_data[layer_idx] - - event_transfer = data.get('event_transfer_time', 0) - event_compute = data.get('event_compute_time', 0) - layer_total = event_transfer + event_compute - - if layer_total > cpu_threshold: - bottleneck = "TRANSFER" if event_transfer > event_compute else "COMPUTE" - status = "RESIDENT" if layer_idx in gpu_resident_layers else "SWAPPABLE" - - print(f"Layer {layer_idx:2d}: CPU {layer_total:5.1f}ms | Transfer {event_transfer:4.1f}ms | Compute {event_compute:5.1f}ms | {bottleneck:8s} | {status}") - - problematic_layers.append(str(layer_idx)) - - if problematic_layers: - current_residents = sorted(gpu_resident_layers) - recommended = current_residents + [int(x) for x in problematic_layers] - print(f"\nRECOMMENDED LAYERS FOR GPU: \"{','.join(map(str, sorted(set(recommended))))}\"") - print("No. of Layers: ", len(recommended)) - - print("=" * 80) - -def debug_layer_locations(layers, device_cache): - """Debug exact layer locations and memory usage""" - print(f"\n=== LAYER LOCATION DEBUG ===") - - if layers is None or device_cache is None: - print("ERROR: layers/device cache is None - debug called before layer setup") - return - - total_gpu_mem = 0 - total_cpu_mem = 0 - - for i, layer in enumerate(layers): - try: - # Get actual device - actual_device = next(layer.parameters()).device - except StopIteration: - actual_device = torch.device('cpu') # No parameters - - # Get cached device - cached_device = device_cache.get_device(i) - - # Calculate layer memory - layer_mem = sum(p.numel() * p.element_size() for p in layer.parameters()) - layer_mem_mb = layer_mem / (1024 * 1024) - - # Track totals - if actual_device.type == 'cuda': - total_gpu_mem += layer_mem_mb - else: - total_cpu_mem += layer_mem_mb - - # Status indicators - status = "RESIDENT" if i in gpu_resident_layers else "SWAPPABLE" - mismatch = "MISMATCH" if actual_device != cached_device else "OK" - - print(f"Layer {i:2d}: {actual_device} | Cache: {cached_device} | " - f"{layer_mem_mb:6.1f}MB | {status} | {mismatch}") - - print(f"\nTOTAL IN GPU MEMORY: {total_gpu_mem:.1f}MB ({total_gpu_mem / 1024:.2f}GB)") - print(f"TOTAL IN CPU MEMORY: {total_cpu_mem:.1f}MB ({total_cpu_mem / 1024:.2f}GB)") - print(f"GPU ALLOCATED: {torch.cuda.memory_allocated(0) / 1024 ** 3:.3f}GB") - print("============================\n") - - -def print_memory_optimization_analysis(model, layers, args): - """ - Streamlined memory analysis for DGLS Dynamic Swapping Loader. - Provides clear, actionable recommendations without unnecessary detail. - """ - import platform - import psutil - import torch - import time - import comfy.model_management - - available_for_layers = 0 - gpu_memory_used = 0 - unused_memory = 0 - swapping_overhead = 0 - total_vram = 0 - inference_memory = 0 - - # Check if this is first run or per-step update - is_startup = not hasattr(print_memory_optimization_analysis, 'has_run') - if is_startup: - print_memory_optimization_analysis.has_run = True - print("\n" + "=" * 80) - print(" " * 20 + "DGLS MEMORY OPTIMIZATION ANALYSIS") - print("=" * 80) - else: - print("\n" + "-" * 80) - print(" " * 25 + "STEP PERFORMANCE UPDATE") - print("-" * 80) - - # ======================================================================== - # SYSTEM ANALYSIS (startup only) - # ======================================================================== - if is_startup: - print("\n--- SYSTEM ANALYSIS ---") - - if torch.cuda.is_available(): - gpu_props = torch.cuda.get_device_properties(0) - total_vram = gpu_props.total_memory - free_vram, _ = torch.cuda.mem_get_info(0) - used_vram = total_vram - free_vram - - print(f"\nGPU DEVICE:") - print(f" Name: {gpu_props.name}") - print(f" Total VRAM: {total_vram / 1024 ** 3:.2f} GB") - print(f" Free VRAM: {free_vram / 1024 ** 3:.2f} GB") - print(f" Used VRAM: {used_vram / 1024 ** 3:.2f} GB ({used_vram / total_vram * 100:.1f}%)") - else: - print("\nWARNING: No CUDA device available!") - total_vram = 0 - free_vram = 0 - - cpu_mem = psutil.virtual_memory() - print(f"\nSYSTEM MEMORY:") - print(f" Total RAM: {cpu_mem.total / 1024 ** 3:.2f} GB") - print(f" Available RAM: {cpu_mem.available / 1024 ** 3:.2f} GB") - - # ======================================================================== - # LAYER ANALYSIS - # ======================================================================== - if is_startup: - print("\n--- MODEL LAYER ANALYSIS ---") - - layer_sizes_bytes = [] - for i, layer in enumerate(layers): - size_bytes = sum(p.numel() * p.element_size() for p in layer.parameters()) - layer_sizes_bytes.append(size_bytes) - - total_model_size = sum(layer_sizes_bytes) - avg_layer_size = total_model_size / len(layers) if layers else 0 - max_layer_size = max(layer_sizes_bytes) if layer_sizes_bytes else 0 - min_layer_size = min(layer_sizes_bytes) if layer_sizes_bytes else 0 - - print(f"\nMODEL STATISTICS:") - print(f" Total layers: {len(layers)}") - print(f" Model size: {total_model_size / 1024 ** 3:.2f} GB") - print(f" Average layer: {avg_layer_size / 1024 ** 2:.1f} MB") - print(f" Max layer size: {max_layer_size / 1024 ** 2:.1f} MB") - print(f" Min layer size: {min_layer_size / 1024 ** 2:.1f} MB") - else: - # Recalculate for step updates - layer_sizes_bytes = [sum(p.numel() * p.element_size() for p in layer.parameters()) - for layer in layers] - avg_layer_size = sum(layer_sizes_bytes) / len(layers) if layers else 0 - - # ======================================================================== - # MEMORY BREAKDOWN - # ======================================================================== - if is_startup: - print("\n--- MEMORY BREAKDOWN ---") - - if torch.cuda.is_available(): - # Get actual memory state - stats = torch.cuda.memory_stats(0) - mem_free_cuda, _ = torch.cuda.mem_get_info(0) - mem_free_torch = stats['reserved_bytes.all.current'] - stats['active_bytes.all.current'] - total_free = mem_free_cuda + mem_free_torch - - device = comfy.model_management.get_torch_device() - total_free = comfy.model_management.get_free_memory(device) - - # reserved_memory = comfy.model_management.extra_reserved_memory() # OS/driver overhead - inference_memory = comfy.model_management.minimum_inference_memory() # Base inference needs - - # Calculate layer sizes - layer_sizes_bytes = [] - for i, layer in enumerate(layers): - size_bytes = sum(p.numel() * p.element_size() for p in layer.parameters()) - layer_sizes_bytes.append(size_bytes) - - # Calculate swapping overhead - gpu_layers = gpu_resident_layers.copy() - swappable_layers = cpu_swappable_layers.copy() - - swapping_overhead = 0 - prefetch_overhead = 0 - threading_overhead = 0 - max_swappable_size = 0 - if len(swappable_layers) > 0: - max_swappable_size = max(layer_sizes_bytes[i] for i in swappable_layers) - prefetch_overhead = (args.prefetch+1) * max_swappable_size * 1.25 - threading_overhead = prefetch_overhead * 2 - - if args.cpu_threading: - swapping_overhead = prefetch_overhead + threading_overhead - else: - swapping_overhead = prefetch_overhead - - # Calculate available memory step by step - after_inference = total_free - inference_memory - after_sampling = after_inference - available_for_layers = max(0, after_sampling - swapping_overhead) - - print(f"MEMORY CALCULATION:") - print(f" Total free memory: {total_free / 1024 ** 3:.2f} GB") - # print(f" Reserved (OS/driver): {reserved_memory / 1024 ** 3:.2f} GB") - print(f" Minus inference overhead: -{inference_memory / 1024 ** 3:.2f} GB = {after_inference / 1024 ** 3:.2f} GB") - print(f" Minus swapping overhead: -{swapping_overhead / 1024 ** 3:.2f} GB = {available_for_layers / 1024 ** 3:.2f} GB") - # layers_min = calculate_auto_gpu_layers(layers, args) - print(f" RECOMMENDED AMOUNT OF STARTING LAYERS: {(available_for_layers/ 1024 ** 3) / ((max_swappable_size/ 1024 ** 3) * 2.1):.1f} Layers with prefetch=0") - # print(f" RECOMMENDED AMOUNT OF STARTING LAYERS: {len(layers_min)} Layers for current settings") - - - # Current allocation - gpu_memory_used = sum(layer_sizes_bytes[i] for i in gpu_layers) if gpu_layers else 0 - unused_memory = available_for_layers - gpu_memory_used - - print(f"\nCURRENT MEMORY ALLOCATION:") - print(f" GPU layers: {len(gpu_layers)}/{len(layers)} (indices: {sorted(list(gpu_layers))})") - print(f" Memory Used for Layers: {gpu_memory_used / 1024 ** 3:.2f} GB") - print(f" Memory left for more layers: {available_for_layers / 1024 ** 3:.2f} GB") - - print() - else: - available_for_layers = 0 - gpu_memory_used = 0 - unused_memory = 0 - swapping_overhead = 0 - - else: - # ==================================================================== - # STEP UPDATES (only during steps) - # ==================================================================== - print("\n--- CURRENT GPU MEMORY ---") - - if torch.cuda.is_available(): - free_vram, _ = torch.cuda.mem_get_info(0) - allocated_memory = torch.cuda.memory_allocated(0) - reserved_memory = torch.cuda.memory_reserved(0) - - print(f" Free VRAM: {free_vram / 1024 ** 3:.2f} GB") - print(f" Allocated: {allocated_memory / 1024 ** 3:.2f} GB") - print(f" Reserved: {reserved_memory / 1024 ** 3:.2f} GB") - - print(f"\n--- SYSTEM MEMORY: ---") - cpu_mem = psutil.virtual_memory() - print(f" Total RAM: {cpu_mem.total / 1024 ** 3:.2f} GB") - print(f" Available RAM: {cpu_mem.available / 1024 ** 3:.2f} GB") - - return { - 'available_memory_gb': available_for_layers / 1024 ** 3 if available_for_layers else 0, - 'unused_memory_gb': unused_memory / 1024 ** 3 if available_for_layers else 0, - 'gpu_layers': list(gpu_layers) if available_for_layers else [], - 'swapping_overhead_gb': swapping_overhead / 1024 ** 3 if available_for_layers else 0 - } - - -# ============================================================================= -# DYNAMIC LAYER SWAPPING INFERENCE for Comfy-UI - by obisin -# ============================================================================= - - -def fix_inference_tensor_parameters(layer): - """Fix inference tensor parameters in a layer if needed""" - # Quick check if ANY parameters need fixing - needs_fixing = False - with torch.no_grad(): - for name, param in layer.named_parameters(): - if param is not None: - try: - _ = param._version - except (RuntimeError, AttributeError): - needs_fixing = True - break - - # Only fix if needed - if needs_fixing: - with torch.no_grad(): - for name, param in layer.named_parameters(): - if param is not None: +def extract_layers_from_comfy_model(model_patcher): + """Extract ComfyUI's layers from different model architectures""" + + if hasattr(model_patcher.model, 'diffusion_model'): + dm = model_patcher.model.diffusion_model + + layers = [] + + # Method 1: Standard blocks (WAN, SDXL, etc.) + if hasattr(dm, 'blocks') and len(dm.blocks) > 0: + for i, block in enumerate(dm.blocks): + layers.append(block) + print(f"Extracted {len(layers)} blocks from diffusion model") + return layers + + # Method 2: Flux DiT layers (double_blocks + single_blocks) + if hasattr(dm, 'double_blocks') or hasattr(dm, 'single_blocks'): + if hasattr(dm, 'double_blocks'): + for i, block in enumerate(dm.double_blocks): + layers.append(block) + # Debug: Check the forward signature + import inspect + sig = inspect.signature(block.forward) + print(f"Block {i} forward signature: {sig}") + print(f"Added {len(dm.double_blocks)} double_blocks") + + if hasattr(dm, 'single_blocks'): + for i, block in enumerate(dm.single_blocks): + layers.append(block) + print(f"Added {len(dm.single_blocks)} single_blocks") + + print(f"Extracted {len(layers)} Flux DiT layers") + return layers + + # Method 3: SD3 DiT layers + if hasattr(dm, 'joint_blocks'): + for i, block in enumerate(dm.joint_blocks): + layers.append(block) + print(f"Extracted {len(layers)} SD3 joint_blocks") + return layers + + # Method 4: HunyuanVideo layers + if hasattr(dm, 'blocks') and hasattr(dm, 'single_blocks'): + # HunyuanVideo has both blocks and single_blocks + for i, block in enumerate(dm.blocks): + layers.append(block) + for i, block in enumerate(dm.single_blocks): + layers.append(block) + print(f"Extracted {len(layers)} HunyuanVideo layers (blocks + single_blocks)") + return layers + + # Method 5: Transformer layers (generic DiT models) + if hasattr(dm, 'transformer_blocks'): + for i, block in enumerate(dm.transformer_blocks): + layers.append(block) + print(f"Extracted {len(layers)} transformer_blocks") + return layers + + # Method 6: Generic layer extraction by common names + layer_names = ['layers', 'encoders', 'decoders', 'blocks', 'modules'] + for attr_name in layer_names: + if hasattr(dm, attr_name): + attr = getattr(dm, attr_name) + if hasattr(attr, '__iter__') and not isinstance(attr, str): try: - _ = param._version - except (RuntimeError, AttributeError): - param.data = param.data.detach().clone() + for i, block in enumerate(attr): + if isinstance(block, torch.nn.Module): + layers.append(block) + if layers: + print(f"Extracted {len(layers)} layers from {attr_name}") + return layers + except: + continue - return needs_fixing + # Method 7: Search for block-like modules + for name, module in dm.named_children(): + if isinstance(module, torch.nn.Module): + # Look for modules that are likely to be computation blocks + if any(keyword in name.lower() for keyword in ['block', 'layer', 'transformer', 'attention', 'mlp']): + # Check if it's a container of blocks or a single block + if hasattr(module, '__iter__') and not isinstance(module, torch.nn.Parameter): + try: + for sub_module in module: + if isinstance(sub_module, torch.nn.Module): + layers.append(sub_module) + except: + layers.append(module) + else: + layers.append(module) -def get_cached_layer_size_mb(idx): - size = layer_sizes_mb.get(idx, 0) - if size == 0: - print(f"WARNING: Layer {idx} size not found in cache") - return size + if layers: + print(f"Extracted {len(layers)} layers using pattern matching") + return layers -class LayerDeviceCache: - def __init__(self, model, layers): # Add layers parameter - self.cache = {} - self.dirty = set() - # Initialize cache using the layers list - for i, layer in enumerate(layers): - self.cache[i] = self.get_layer_device(layer) + # Method 8: Last resort - extract all meaningful modules + print(" Using fallback: extracting all major modules") + skip_modules = ['input_layer', 'output_layer', 'norm', 'embedding', 'pos_embed', 'patch_embed'] + for name, module in dm.named_children(): + if isinstance(module, torch.nn.Module) and name not in skip_modules: + layers.append(module) + print(f" Added module: {name} ({type(module).__name__})") - def get_layer_device(self, layer): - """Your existing function""" - try: - return next(layer.parameters()).device - except StopIteration: - return torch.device('cpu') - - def get_device(self, layer_idx): # Simplified - no need for model parameter - """Fast cached lookup""" - if layer_idx in self.dirty: - if layers is not None and layer_idx < len(layers): - self.cache[layer_idx] = self.get_layer_device(layers[layer_idx]) - self.dirty.remove(layer_idx) - else: - return torch.device('cpu') # Safe fallback - return self.cache.get(layer_idx, torch.device('cpu')) - - def mark_moved(self, layer_idx, new_device): - """Update cache when we move a layer""" - self.cache[layer_idx] = new_device - - -def calculate_auto_gpu_layers(layers, args): - """Auto-select GPU layers allocation logic with safety mechanisms""" - - device = comfy.model_management.get_torch_device() - - # Step 1: Cleanup existing models - comfy.model_management.cleanup_models_gc() - - # Step 2: Calculate total memory required with safety multiplier - total_model_size = sum(comfy.model_management.module_size(layer) for layer in layers) - memory_required_with_safety = total_model_size * 1.1 - minimum_memory_required = comfy.model_management.minimum_inference_memory() - extra_mem = max(minimum_memory_required, - memory_required_with_safety + comfy.model_management.extra_reserved_memory()) - - # Step 3: Free memory if needed (ComfyUI's approach) - initial_free_memory = comfy.model_management.get_free_memory(device) - if initial_free_memory < memory_required_with_safety + minimum_memory_required: - comfy.model_management.free_memory(memory_required_with_safety + extra_mem, device) - print("Freed other models to make space") - - # Step 4: Secondary safety check - current_free_mem = comfy.model_management.get_free_memory(device) - if current_free_mem < minimum_memory_required: - comfy.model_management.free_memory(minimum_memory_required, device) - current_free_mem = comfy.model_management.get_free_memory(device) - - # Step 5: Platform-specific memory ratio - if comfy.model_management.is_nvidia(): - min_ratio = 0.0 - else: - min_ratio = 0.4 - - # Step 6: budget calculation - layer_memory_budget = max( - 128 * 1024 * 1024, - (current_free_mem - minimum_memory_required), - min( - current_free_mem * min_ratio, - current_free_mem - minimum_memory_required - ) - ) - - # Get currently loaded memory and subtract it - loaded_memory = sum(model.loaded_size() for model in comfy.model_management.loaded_models()) - layer_memory_budget = max(0.1, layer_memory_budget - loaded_memory) - - # Step 7: Calculate layer sizes and sort (largest first) - layer_allocation_data = [] - for i, layer in enumerate(layers): - layer_size = comfy.model_management.module_size(layer) - layer_allocation_data.append((layer_size, i)) - - layer_allocation_data.sort(reverse=True, key=lambda x: x[0]) - - # Step 8: Apply overhead from user settings - if layer_allocation_data: - max_layer_size = layer_allocation_data[0][0] - overhead = max_layer_size * 1.1 * (args.prefetch + 1) - - if args.cpu_threading: - overhead += max_layer_size * 1.25 * (args.prefetch + 1) - - if args.cuda_streams: - overhead += max_layer_size * 0.6 * (args.prefetch + 1) - - layer_memory_budget = max(128 * 1024 * 1024, layer_memory_budget - overhead) - - # Step 9: Greedy allocation - allocated_memory = 0 - selected_gpu_layers = set() - - for layer_size, layer_idx in layer_allocation_data: - if allocated_memory + layer_size <= layer_memory_budget: - allocated_memory += layer_size - selected_gpu_layers.add(layer_idx) + if layers: + print(f"Extracted {len(layers)} modules using fallback method") + return layers else: - break + print(" No layers found in diffusion model") + # Debug: show what's available + print("Available attributes:") + for name in dir(dm): + if not name.startswith('_') and hasattr(dm, name): + attr = getattr(dm, name) + if isinstance(attr, (torch.nn.Module, torch.nn.ModuleList, list)): + print(f" {name}: {type(attr)}") + return [] + else: + print(" No diffusion_model found") + return [] - # Step 10: Final validation - if allocated_memory > current_free_mem * 0.85: # Additional safety check - print("WARNING: High memory allocation detected, reducing selection") - # Remove smallest selected layers until under 85% of free memory - selected_list = [(layer_allocation_data[i][0], layer_allocation_data[i][1]) for i in - range(len(layer_allocation_data)) if layer_allocation_data[i][1] in selected_gpu_layers] - selected_list.sort() # Smallest first for removal - while allocated_memory > current_free_mem * 0.8 and selected_list: - removed_size, removed_idx = selected_list.pop(0) - selected_gpu_layers.remove(removed_idx) - allocated_memory -= removed_size +def add_to_layers_method(model_patcher): + """Add to_layers method to model so it works like your training code""" - if args.verbose: - print(f"Free VRAM: {current_free_mem / (1024 ** 2):.0f}MB") - print(f"Final budget: {layer_memory_budget / (1024 ** 2):.0f}MB") - print(f"Auto-selected GPU layers: {sorted(selected_gpu_layers)}") - print(f"Final allocated: {allocated_memory / (1024 ** 2):.0f}MB") + def to_layers(): + return extract_layers_from_comfy_model(model_patcher) - return selected_gpu_layers + # Add the method to the model_patcher so you can call model.to_layers() + model_patcher.to_layers = to_layers + return model_patcher -def calculate_needed_layers(layer_idx, prefetch): - needed = set() - needed.add(layer_idx) - - # Get sorted list of swappable layers - swappable_list = sorted(cpu_swappable_layers) - - # if not hasattr(calculate_needed_layers, 'printed_list'): - # print(f"DEBUG: swappable_list = {swappable_list}") - # calculate_needed_layers.printed_list = True - - try: - # Find current layer's position in the swappable list - current_pos = swappable_list.index(layer_idx) - - # Take next 'prefetch' layers, wrapping around if needed - for i in range(1, prefetch + 1): - next_pos = (current_pos + i) % len(swappable_list) - needed.add(swappable_list[next_pos]) - - except ValueError: - # Current layer not in swappable list - shouldn't happen but fallback to old logic - for i in range(1, prefetch + 1): - prefetch_idx = layer_idx + i - - if prefetch_idx >= len(layers): - wrap_offset = prefetch_idx - len(layers) - if wrap_offset < len(swappable_list): - prefetch_idx = swappable_list[wrap_offset] - else: - continue - - if prefetch_idx in cpu_swappable_layers: - needed.add(prefetch_idx) - - # if VERBOSE and layer_idx % 10 == 0: - # print(f"Layer {layer_idx}: needed {len(needed)} layers {sorted(needed)} (prefetch={prefetch})") - return needed - -def cleanup_excess_layers(keep_layers): - """Remove layers from GPU that are not in keep_layers set""" - def _cleanup_operation(): - layers_to_remove = [] - for idx in cpu_swappable_layers: - if (idx < len(layers) and - idx not in keep_layers and - device_cache.get_device(idx).type == 'cuda'): - layers_to_remove.append(idx) - - cleaned_count = 0 - - if layers_to_remove: - if args.cpu_threading and hasattr(add_smart_swapping_to_layer, 'cpu_thread_pool'): - def batch_cpu_transfer(): - for idx in layers_to_remove: - # Add CUDA event timing - if args.verbose: - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - start_event.record() - - layers[idx].to(CPU_DEVICE) - device_cache.mark_moved(idx, torch.device('cpu')) - swap_stats['to_cpu'] += 1 - - # Store transfer timing - if args.verbose: - end_event.record() - torch.cuda.synchronize() - transfer_time = start_event.elapsed_time(end_event) - add_smart_swapping_to_layer.layer_transfer_times[idx].append(transfer_time) - - add_smart_swapping_to_layer.cpu_thread_pool.submit(batch_cpu_transfer) - else: - for idx in layers_to_remove: - # Add CUDA event timing - if args.verbose: - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - start_event.record() - - layers[idx].to(CPU_DEVICE) - device_cache.mark_moved(idx, torch.device('cpu')) - swap_stats['to_cpu'] += 1 - - # Store transfer timing - if args.verbose: - end_event.record() - torch.cuda.synchronize() - transfer_time = start_event.elapsed_time(end_event) - add_smart_swapping_to_layer.layer_transfer_times[idx].append(transfer_time) - - cleaned_count = len(layers_to_remove) - - # Clean up transfer events - for idx in layers_to_remove: - if (hasattr(add_smart_swapping_to_layer, 'transfer_events') and - idx in add_smart_swapping_to_layer.transfer_events): - del add_smart_swapping_to_layer.transfer_events[idx] - - - return cleaned_count - - # Profile every 5th step, not every cleanup - if not hasattr(cleanup_excess_layers, 'step_counter'): - cleanup_excess_layers.step_counter = 0 - - cleanup_excess_layers.step_counter += 1 - - return _cleanup_operation() - - -def fetch_missing_layers(needed_layers, current_layer_idx=None): - """Ensure all needed layers are on GPU, but return as soon as current layer is ready""" - - def _fetch_operation(): - layers_to_fetch = [idx for idx in needed_layers - if (idx < len(layers) and - idx in cpu_swappable_layers and - device_cache.get_device(idx).type == 'cpu')] - - fetched_count = 0 - - transfer_start = time.perf_counter() - - for idx in layers_to_fetch: - overlap_stats['transfer_start_times'][idx] = transfer_start - - if layers_to_fetch: - if args.cuda_streams and hasattr(add_smart_swapping_to_layer, 'fetch_stream'): - with torch.cuda.stream(add_smart_swapping_to_layer.fetch_stream): - for idx in layers_to_fetch: - if args.verbose: - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - start_event.record() - - layers[idx].to(GPU_DEVICE, non_blocking=True) - - transfer_event = torch.cuda.Event() - transfer_event.record() - if hasattr(add_smart_swapping_to_layer, 'transfer_events'): - add_smart_swapping_to_layer.transfer_events[idx] = transfer_event - - device_cache.mark_moved(idx, torch.device('cuda')) - swap_stats['to_gpu'] += 1 - fetched_count += 1 - - # Store transfer timing - if args.verbose: - end_event.record() - torch.cuda.synchronize() - transfer_time = start_event.elapsed_time(end_event) - add_smart_swapping_to_layer.layer_transfer_times[idx].append(transfer_time) - else: - for idx in layers_to_fetch: - # Add CUDA event timing - if args.verbose: - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - start_event.record() - - layers[idx].to(GPU_DEVICE, non_blocking=True) - device_cache.mark_moved(idx, torch.device('cuda')) - swap_stats['to_gpu'] += 1 - fetched_count += 1 - - # Store transfer timing - if args.verbose: - end_event.record() - torch.cuda.synchronize() - transfer_time = start_event.elapsed_time(end_event) - add_smart_swapping_to_layer.layer_transfer_times[idx].append(transfer_time) - - # Wait only for current layer if specified - if current_layer_idx is not None and current_layer_idx in layers_to_fetch: - while True: - try: - test_param = next(layers[current_layer_idx].parameters()) - if test_param.device.type == 'cuda': - break - except StopIteration: - break - - # Record transfer end times - transfer_end = time.perf_counter() - for idx in layers_to_fetch: - overlap_stats['transfer_end_times'][idx] = transfer_end - - return fetched_count - - - if not hasattr(fetch_missing_layers, 'step_counter'): - fetch_missing_layers.step_counter = 0 - - fetch_missing_layers.step_counter += 1 - - return _fetch_operation() - - -def add_smart_swapping_to_layer(layer, layer_idx, layers_list, gpu_resident_layers, cpu_swappable_layers): - """Add swapping capability with ComfyUI coordination""" - - global layers, device_cache - layers = layers_list - original_forward = layer.forward - - @wraps(original_forward) - def swapped_forward(*fwd_args, **kwargs): - def _whole_forward_operation(): - layer_compute_start = time.perf_counter() - overlap_stats['compute_start_times'][layer_idx] = layer_compute_start - - # Detect new step when we hit layer 0 again (reset to early layers) - # In swapped_forward, after the new step detection: - if args.verbose and layer_idx == min(cpu_swappable_layers): - add_smart_swapping_to_layer.current_step += 1 - add_smart_swapping_to_layer.calls_this_step = 0 - if args.verbose: - print(f" New sampling step {add_smart_swapping_to_layer.current_step}") - - if add_smart_swapping_to_layer.current_step % 3 == 0: - analyze_layer_performance() - - if layer_idx in cpu_swappable_layers: - if PROFILE: - nvtx.range_push(f"Smart_Management_L{layer_idx}") - try: - current_device = device_cache.get_device(layer_idx) - layer_already_on_gpu = (current_device.type == 'cuda') - - if not layer_already_on_gpu: - needed_layers = calculate_needed_layers(layer_idx, args.prefetch) - cleaned = cleanup_excess_layers(needed_layers) - fetched = fetch_missing_layers(needed_layers, current_layer_idx=layer_idx) - - final_device = device_cache.get_device(layer_idx) - - if layer_already_on_gpu: - add_smart_swapping_to_layer.prefetch_hits += 1 - else: - add_smart_swapping_to_layer.prefetch_misses += 1 - - device = device_cache.get_device(layer_idx) - gpu_success = device.type == 'cuda' - - finally: - if PROFILE: - nvtx.range_pop() - - device = device_cache.get_device(layer_idx) - gpu_success = device.type == 'cuda' - else: - # Layer not in cpu_swappable_layers (permanent resident) - device = device_cache.get_device(layer_idx) - gpu_success = device.type == 'cuda' - - # Detect Flux calling patterns - layer_type = type(layer).__name__ - is_flux_call = 'DoubleStream' in layer_type or 'SingleStream' in layer_type - is_keyword_only_call = any(pattern in layer_type for pattern in ['QwenImage', - # 'OmniGen', - 'HiDream']) - - - if is_flux_call: - # Flux patterns - preserve exact argument structure - if fwd_args: - # SingleStreamBlock: block(img, vec=vec, pe=pe) - x = fwd_args[0] # img - new_args = fwd_args[1:] # should be empty - new_kwargs = kwargs # vec=vec, pe=pe - else: - # DoubleStreamBlock: block(img=img, txt=txt, vec=vec, pe=pe) - x = kwargs.get('img') - new_args = () - new_kwargs = kwargs - elif is_keyword_only_call: - x = None - new_args = () - new_kwargs = kwargs - else: - # WAN/SDXL pattern - x = fwd_args[0] if fwd_args else None - new_args = fwd_args[1:] - new_kwargs = kwargs - - add_smart_swapping_to_layer.total_forward_calls += 1 - add_smart_swapping_to_layer.calls_this_step += 1 - - if args.verbose and add_smart_swapping_to_layer.total_forward_calls % 100 == 0: - global current_model - if 'current_model' in globals(): - print_memory_optimization_analysis(current_model, layers, args) - # debug_layer_locations(layers, device_cache) - - if PROFILE: - nvtx.range_push(f"Layer_{layer_idx}_{'Forward'}") - try: - if not add_smart_swapping_to_layer.current_forward_logged: - current_step = getattr(add_smart_swapping_to_layer, 'current_step', 0) - if args.verbose: - print( - f" Forward pass step {current_step}: layers {min(cpu_swappable_layers)}-{max(cpu_swappable_layers)}") - add_smart_swapping_to_layer.current_forward_logged = True - - # Handle GPU failure case - if not gpu_success: - print(f" Layer {layer_idx} failed to be on GPU, forcing aggressive cleanup...") - for cleanup_idx in cpu_swappable_layers: - if cleanup_idx != layer_idx and cleanup_idx < len(layers): - try: - layers[cleanup_idx].to('cpu') - except: - pass - - gc.collect() - torch.cuda.empty_cache() - - # Try again after cleanup - try: - layer.to(GPU_DEVICE) - device = torch.device(GPU_DEVICE) - gpu_success = True - print(f" Layer {layer_idx} moved to GPU after aggressive cleanup") - except RuntimeError: - print(f" CRITICAL: Layer {layer_idx} cannot fit on GPU, skipping computation!") - return x # Pass input through unchanged - - def move_to_device(tensor, target_device): - if PROFILE: - nvtx.range_push(f"Move_Tensor_To_{target_device.type}") - try: - # Only move device, don't touch dtype at all - if hasattr(tensor, 'device') and tensor.device != target_device: - tensor = tensor.to(target_device) - - # Fix inference tensors (without dtype changes) - if isinstance(tensor, torch.Tensor): - try: - _ = tensor._version - return tensor - except (RuntimeError, AttributeError): - return tensor.detach().clone() - - return tensor - finally: - if PROFILE: - nvtx.range_pop() - - move_to_device.current_layer = layer_idx - - # Move tensors to device (respect the calling pattern you already set up) - if is_flux_call: - if fwd_args: - # SingleStreamBlock - move fwd_args and kwargs - new_args = tuple(move_to_device(arg, device) for arg in fwd_args) - new_kwargs = {k: move_to_device(v, device) for k, v in kwargs.items()} - else: - # DoubleStreamBlock - move only kwargs - new_args = () - new_kwargs = {k: move_to_device(v, device) for k, v in kwargs.items()} - elif is_keyword_only_call: - # move only kwargs, no positional args - x = None - new_args = () - new_kwargs = {k: move_to_device(v, device) for k, v in kwargs.items()} - else: - x = move_to_device(x, device) - new_kwargs = {k: move_to_device(v, device) for k, v in kwargs.items()} - new_args = tuple(move_to_device(arg, device) for arg in new_args) - - # Wait for current layer's transfer to complete before computation - if args.cuda_streams and layer_idx in cpu_swappable_layers: - # Check if we have a transfer event for this specific layer - if (hasattr(add_smart_swapping_to_layer, 'transfer_events') and - layer_idx in add_smart_swapping_to_layer.transfer_events): - # Make compute stream wait for this layer's transfer event - event = add_smart_swapping_to_layer.transfer_events[layer_idx] - add_smart_swapping_to_layer.compute_stream.wait_event(event) - - # 4. COMPUTATION with profiling - def _compute_operation(): - if args.verbose: - compute_start = time.perf_counter() - - if args.cuda_streams and hasattr(add_smart_swapping_to_layer, 'compute_stream'): - with torch.cuda.stream(add_smart_swapping_to_layer.compute_stream): - fix_inference_tensor_parameters(layer) - if is_flux_call: - if fwd_args: - result = original_forward(*new_args, **new_kwargs) - else: - result = original_forward(**new_kwargs) - - elif is_keyword_only_call: - result = original_forward(**new_kwargs) - else: - result = original_forward(x, *new_args, **new_kwargs) - - - else: - fix_inference_tensor_parameters(layer) - if is_flux_call: - if fwd_args: - result = original_forward(*new_args, **new_kwargs) - else: - result = original_forward(**new_kwargs) - elif is_keyword_only_call: - result = original_forward(**new_kwargs) - else: - result = original_forward(x, *new_args, **new_kwargs) - - # Store compute timing - if args.verbose: - compute_end = time.perf_counter() - compute_time = (compute_end - compute_start) * 1000 # Convert to ms - add_smart_swapping_to_layer.layer_compute_times[layer_idx].append(compute_time) - - layer_compute_end = time.perf_counter() - overlap_stats['compute_end_times'][layer_idx] = layer_compute_end - - # Calculate overlap for this layer - if (layer_idx in overlap_stats['transfer_start_times'] and - layer_idx in overlap_stats['transfer_end_times']): - - t_start = overlap_stats['transfer_start_times'][layer_idx] - t_end = overlap_stats['transfer_end_times'][layer_idx] - c_start = overlap_stats['compute_start_times'][layer_idx] - c_end = overlap_stats['compute_end_times'][layer_idx] - - # Calculate overlap duration - overlap_start = max(t_start, c_start) - overlap_end = min(t_end, c_end) - overlap_duration = max(0, overlap_end - overlap_start) - - transfer_duration = t_end - t_start - compute_duration = c_end - c_start - - if transfer_duration > 0 and compute_duration > 0: - overlap_ratio = overlap_duration / min(transfer_duration, compute_duration) - overlap_stats['overlaps'].append({ - 'layer': layer_idx, - 'overlap_ms': overlap_duration * 1000, - 'transfer_ms': transfer_duration * 1000, - 'compute_ms': compute_duration * 1000, - 'overlap_ratio': overlap_ratio - }) - - return result - - # Profile compute every 5th step, not every layer - current_step = getattr(add_smart_swapping_to_layer, 'current_step', 0) - should_profile_compute = (args.verbose) - - if PROFILE: - nvtx.range_push(f"Compute_L{layer_idx}") - try: - return _compute_operation() - finally: - if PROFILE: - nvtx.range_pop() - finally: - if PROFILE: - nvtx.range_pop() - - result = _whole_forward_operation() - - return result - - # Replace forward method - layer.forward = swapped_forward - -# ============================================================================= -# COMFYUI NODE CLASS - by obisin -# ============================================================================= - -class DynamicSwappingLoader: +class DGLSModelLoader: @classmethod def INPUT_TYPES(cls): + model_list = folder_paths.get_filename_list("unet") + folder_paths.get_filename_list("diffusion_models")#+ folder_paths.get_filename_list("checkpoints") return { "required": { - "model": ("MODEL",), - "layers": ("LAYERS",), - "initial_gpu_layers": ("INT",{"default": 1, "min": 0, "max": 100, "tooltip": TOOLTIPS["initial_gpu_layers"]}), - "final_gpu_layers": ("INT",{"default": 1, "min": 0, "max": 100, "tooltip": TOOLTIPS["final_gpu_layers"]}), - "auto": ("BOOLEAN", {"default": False, "tooltip": "Automatically determine layer placement based on available VRAM. Overrides any other layer placement."}), - "prefetch": ("INT", {"default": 1, "min": 0, "max": 100, "tooltip": TOOLTIPS["prefetch"]}), - "cpu_threading": ("BOOLEAN", {"default": False, "tooltip": TOOLTIPS["cpu_threading"]}), - "cuda_streams": ("BOOLEAN", {"default": False, "tooltip": TOOLTIPS["cuda_streams"]}), - "verbose": ("BOOLEAN", {"default": True, "tooltip": TOOLTIPS["verbose"]}), + "model_name": (model_list, {"default": model_list[0] if model_list else ""}), + "model_type": (["default", "hunyuan", "unet"], {"default": "default", "tooltip": "Unet is just for Debugging, its a legacy loader- Use default"}), + "cast_dtype": (["default", "fp32", "fp16", "bf16", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],{"default": "default"}), + "clear_model_cache": ("BOOLEAN", {"default": False, "tooltip": "Force reload model from disk, ignoring ComfyUI's model cache"}), + "verbose": ("BOOLEAN", {"default": False}), }, - "optional": { - "gpu_layer_indices": ("STRING", {"default": "", "tooltip": TOOLTIPS["gpu_layers"]}), - "cast_target": ("STRING", {"default": "", "tooltip": TOOLTIPS["cast_target"]}), + "optional": {"nuke_all_caches": ("BOOLEAN", {"default": False, "tooltip": "CAUTION: Aggressive Clear all ComfyUI caches"}), + "custom_ckpt_path": ("STRING", {"default": ""}), } } - RETURN_TYPES = ("MODEL",) - RETURN_NAMES = ("model",) - - FUNCTION = "load_model_with_swapping" + RETURN_TYPES = ("MODEL", "LAYERS") + RETURN_NAMES = ("model", "layers") + FUNCTION = "load_dgls_model" CATEGORY = "loaders" - TITLE = "Dynamic Swapping Loader" + TITLE = "DGLS Model Loader" + + def load_dgls_model(self, model_name, model_type, cast_dtype, verbose, clear_model_cache=False, nuke_all_caches=False, custom_ckpt_path="", override_model_type="auto"): - def load_model_with_swapping(self, model, layers, initial_gpu_layers, final_gpu_layers, auto, - prefetch, verbose, gpu_layer_indices="", - cast_target="", cuda_streams=False, cpu_threading=False): + self.verbose = verbose - # Reset all stats for new generation - if hasattr(add_smart_swapping_to_layer, 'total_start_time'): - del add_smart_swapping_to_layer.total_start_time - if hasattr(add_smart_swapping_to_layer, 'current_step'): - add_smart_swapping_to_layer.current_step = 0 - if hasattr(add_smart_swapping_to_layer, 'total_forward_calls'): - add_smart_swapping_to_layer.total_forward_calls = 0 - if hasattr(add_smart_swapping_to_layer, 'calls_this_step'): - add_smart_swapping_to_layer.calls_this_step = 0 - if hasattr(add_smart_swapping_to_layer, 'current_forward_logged'): - add_smart_swapping_to_layer.current_forward_logged = False - if hasattr(add_smart_swapping_to_layer, 'prefetch_hits'): - add_smart_swapping_to_layer.prefetch_hits = 0 - if hasattr(add_smart_swapping_to_layer, 'prefetch_misses'): - add_smart_swapping_to_layer.prefetch_misses = 0 - if hasattr(add_smart_swapping_to_layer, 'layer_transfer_times'): - add_smart_swapping_to_layer.layer_transfer_times.clear() - if hasattr(add_smart_swapping_to_layer, 'layer_compute_times'): - add_smart_swapping_to_layer.layer_compute_times.clear() + # Cache clearing logic + if nuke_all_caches: + self.nuke_all_caches() + elif clear_model_cache: + self.clear_model_cache_only() - global swap_stats, transfer_stats - swap_stats = {'to_gpu': 0, 'to_cpu': 0} - transfer_stats = { - 'to_gpu_times': [], 'to_cpu_times': [], 'to_gpu_speeds': [], 'to_cpu_speeds': [], - 'current_step_gpu_times': [], 'current_step_cpu_times': [], - 'current_step_gpu_speeds': [], 'current_step_cpu_speeds': [] - } - - if not hasattr(add_smart_swapping_to_layer, 'total_start_time'): - add_smart_swapping_to_layer.total_start_time = time.time() - - # Reset layer-specific stats - if hasattr(add_smart_swapping_to_layer, 'prefetch_hits'): - add_smart_swapping_to_layer.prefetch_hits = 0 - add_smart_swapping_to_layer.prefetch_misses = 0 - - # Update ALL args with node parameters - args.dynamic_swapping = True - args.initial_gpu_layers = initial_gpu_layers - args.final_gpu_layers = final_gpu_layers - args.prefetch = prefetch - args.cuda_streams = cuda_streams - args.cpu_threading = cpu_threading - args.verbose = verbose - args.auto = auto - - args.gpu_layers = gpu_layer_indices.strip() if gpu_layer_indices.strip() else None - args.cast_target = cast_target.strip() if cast_target.strip() else None - - add_smart_swapping_to_layer.layer_transfer_times = defaultdict(list) - add_smart_swapping_to_layer.layer_compute_times = defaultdict(list) - - if not hasattr(add_smart_swapping_to_layer, 'stats_initialized'): - add_smart_swapping_to_layer.stats_initialized = True - add_smart_swapping_to_layer.prefetch_hits = 0 - add_smart_swapping_to_layer.prefetch_misses = 0 - - if not hasattr(add_smart_swapping_to_layer, 'prefetch_hits'): - add_smart_swapping_to_layer.prefetch_hits = 0 - add_smart_swapping_to_layer.prefetch_misses = 0 - - if not hasattr(add_smart_swapping_to_layer, 'layer_compute_times'): - add_smart_swapping_to_layer.layer_compute_times = [] - - # Initialize coordination state (once) - this should match _patch_model_patcher - if not hasattr(add_smart_swapping_to_layer, 'current_sampling_step'): - add_smart_swapping_to_layer.current_sampling_step = 0 - add_smart_swapping_to_layer.last_processed_step = -1 - add_smart_swapping_to_layer.apply_model_call_count = 0 - add_smart_swapping_to_layer.current_forward_logged = False - - # Initialize coordination state (once) - if not hasattr(add_smart_swapping_to_layer, 'current_sampling_step'): - add_smart_swapping_to_layer.current_sampling_step = 0 - add_smart_swapping_to_layer.last_processed_step = -1 - - if not hasattr(add_smart_swapping_to_layer, 'step_timing_initialized'): - add_smart_swapping_to_layer.step_timing_initialized = True - add_smart_swapping_to_layer.step_start_time = time.time() - add_smart_swapping_to_layer.current_step = 0 - - if not hasattr(add_smart_swapping_to_layer, 'total_forward_calls'): - add_smart_swapping_to_layer.total_forward_calls = 0 - add_smart_swapping_to_layer.current_step = 0 - add_smart_swapping_to_layer.calls_this_step = 0 - - # CUDA STREAMS: Initialize streams (once) - only if enabled - if args.cuda_streams and not hasattr(add_smart_swapping_to_layer, 'streams_initialized'): + if custom_ckpt_path.strip(): + model_path = custom_ckpt_path.strip() + else: try: - add_smart_swapping_to_layer.streams_initialized = True - add_smart_swapping_to_layer.fetch_stream = torch.cuda.Stream() # All transfers - add_smart_swapping_to_layer.compute_stream = torch.cuda.Stream() # Only computation - add_smart_swapping_to_layer.transfer_events = {} # Track transfer completion - # torch.cuda.synchronize() # Ensure streams are ready - print("✓ CUDA Streams enabled for copy-compute overlap") - except Exception as e: - print(f"✗ CUDA Streams failed to initialize: {e}") - args.cuda_streams = False - print("✓ Falling back to default stream") + model_path = folder_paths.get_full_path('unet', model_name) + except: + model_path = folder_paths.get_full_path('diffusion_models', model_name) - # CPU THREADING: Initialize thread pool (once) - only if enabled - if args.cpu_threading and not hasattr(add_smart_swapping_to_layer, 'cpu_thread_pool'): - try: - add_smart_swapping_to_layer.cpu_thread_pool = ThreadPoolExecutor(max_workers=4, - thread_name_prefix="cpu_transfer") - print("✓ CPU Threading enabled") - except Exception as e: - print(f"✗ CPU Threading failed to initialize: {e}") - args.cpu_threading = False - print("✓ Falling back to default CPU transfers") - - if args.verbose: - print(f"\n=== LAYER DTYPE INSPECTION ===") - print(f"Total layers: {len(layers)}") - - for i, layer in enumerate(layers): - layer_type = type(layer).__name__ - - # Get parameter dtypes for this layer - param_dtypes = set() - param_count = 0 - for name, param in layer.named_parameters(): - param_dtypes.add(param.dtype) - param_count += 1 - - # Get buffer dtypes - buffer_dtypes = set() - buffer_count = 0 - for name, buffer in layer.named_buffers(): - buffer_dtypes.add(buffer.dtype) - buffer_count += 1 - - # Format dtype info - dtype_info = [] - if param_dtypes: - dtype_info.append(f"params: {list(param_dtypes)}") - if buffer_dtypes: - dtype_info.append(f"buffers: {list(buffer_dtypes)}") - - dtype_str = ", ".join(dtype_info) if dtype_info else "no tensors" - - print(f" {i:2d}: {layer_type:<20} | {dtype_str}") - - print("===============================\n") + if verbose: + print(f"Loading model using DGLS loader...") + print(f"Model path: {model_path}") + print(f"Weight dtype: {cast_dtype}") + print(f"Override model type: {override_model_type}") - if cast_target: - # Parse the cast_target string first - try: - cast_from, cast_to = cast_target.strip().split() - parsed_cast_target = (cast_from, cast_to) + # Set up model options for dtype control + model_options = {} + if cast_dtype == "fp32": + model_options["dtype"] = torch.float32 + elif cast_dtype == "fp16": + model_options["dtype"] = torch.float16 + elif cast_dtype == "bf16": + model_options["dtype"] = torch.bfloat16 + elif cast_dtype == "fp8_e4m3fn": + model_options["dtype"] = torch.float8_e4m3fn + elif cast_dtype == "fp8_e4m3fn_fast": + model_options["dtype"] = torch.float8_e4m3fn + model_options["fp8_optimizations"] = True + elif cast_dtype == "fp8_e5m2": + model_options["dtype"] = torch.float8_e5m2 + elif cast_dtype == "default": + model_options = {} - casting_handler = CastingHandler() - casting_handler.precast_all_layers(cpu_swappable_layers, layers, - parsed_cast_target) - except ValueError: - print(f" Invalid cast_target format: '{cast_target}'. Expected: 'from_dtype to_dtype'") + if verbose and model_options: + print(f"Model options: {model_options}") - # #FORCE first and last layer onto card regardles sof settings for stability - # if gpu_layer_indices and gpu_layer_indices.strip(): - # specified_gpu_layers = set(map(int, gpu_layer_indices.split(','))) - # - # specified_gpu_layers.add(0) # First layer - # specified_gpu_layers.add(len(layers) - 1) # Last layer - # if args.verbose: - # print(f" Automatically added layers 0 and {len(layers) - 1} for stability") - # else: - # # Use initial/final logic, but ensure at least 1 initial and 1 final - # initial_gpu_layers = max(1, args.initial_gpu_layers) # At least 1 - # final_gpu_layers = max(1, args.final_gpu_layers) # At least 1 - - if args.dynamic_swapping: - if args.verbose: - # print("Applying ENHANCED smart GPU allocation + dynamic swapping...") - print(f"Setting up enhanced swapping for {len(layers)} layers...") - - if args.auto: - specified_gpu_layers = calculate_auto_gpu_layers(layers, args) - elif args.gpu_layers and not args.auto: - try: - specified_gpu_layers = set(map(int, args.gpu_layers.split(','))) - if args.verbose: - print(f" Using specified GPU layers: {sorted(specified_gpu_layers)}") - - # Validate layer indices - invalid_layers = [idx for idx in specified_gpu_layers if idx >= len(layers) or idx < 0] - if invalid_layers: - raise ValueError(f"Invalid layer indices: {invalid_layers}. Must be 0-{len(layers) - 1}") - - except ValueError as e: - print(f" Error parsing --gpu_layers: {e}") - print("Example usage: --gpu_layers 0,1,2,14,18,19,20,21,22") - raise e + if verbose: + if model_path.endswith('.safetensors'): + from safetensors.torch import load_file + sd = load_file(model_path) else: - specified_gpu_layers = None - if args.verbose: - print(f"Using initial/final layer allocation: {initial_gpu_layers}/{final_gpu_layers}") + sd = torch.load(model_path, map_location='cpu') + + print(f"Sample state dict keys: {list(sd.keys())[:5]}") + detected_config = comfy.model_detection.model_config_from_unet(sd, "") + detected_type = type(detected_config).__name__ + print(f"ComfyUI auto-detected: {detected_type}") - global current_model - current_model = model + if model_type == "hunyuan": + # Use state dict method for HunyuanVideo (avoids disconnection) + if model_path.endswith('.safetensors'): + from safetensors.torch import load_file + sd = load_file(model_path) + else: + sd = torch.load(model_path, map_location='cpu') - # Store original values before any modifications - original_initial_gpu_layers = initial_gpu_layers - original_final_gpu_layers = final_gpu_layers + model_patcher = comfy.sd.load_diffusion_model_state_dict(sd, model_options=model_options) - # Clear sets to prevent double-addition from previous runs - gpu_resident_layers.clear() - cpu_swappable_layers.clear() + elif model_type == "unet": + # import comfy.model_management + unet_path = model_path + model_patcher = comfy.sd.load_unet(unet_path) + else: + model_patcher = comfy.sd.load_diffusion_model(model_path, model_options=model_options) - # PHASE 1: Determine which layers go where based on specified layers or initial/final counts - for i, layer in enumerate(layers): - if hasattr(layer, 'to'): - if specified_gpu_layers is not None: - # Use specified layer indices for GPU placement - if i in specified_gpu_layers: - try: - layer.to(GPU_DEVICE) - gpu_resident_layers.add(i) # This is correct - after successful move - if args.verbose: - print(f"Layer {i} (specified) -> GPU permanent") + model_patcher = add_to_layers_method(model_patcher) - # ADD: Check memory after each GPU layer - # if torch.cuda.is_available(): - # current_memory = torch.cuda.memory_allocated(0) - # layer_memory = current_memory - baseline_memory - # if args.verbose: - # print(f"Layer {i} -> GPU: +{layer_memory / 1024 ** 2:.1f}MB (total: {current_memory / 1024 ** 3:.3f}GB)") + # Extract layers + layers = model_patcher.to_layers() - except RuntimeError as e: - print(f"CRITICAL: Cannot fit specified layer {i} on GPU!") - print(f"GPU memory may be insufficient. Consider removing layer {i} from --gpu_layers") - raise e - else: - # Not in specified list, make swappable - layer.to(CPU_DEVICE) - cpu_swappable_layers.add(i) - if args.verbose: - print(f"Layer {i} (not specified) -> CPU swappable") - else: - # Use original initial/final logic as fallback - if i < initial_gpu_layers: - # Initial layers on GPU permanently - try: - layer.to(GPU_DEVICE) - gpu_resident_layers.add(i) # MOVED: Only add if successful - if args.verbose: - print(f"Layer {i} (initial) -> GPU permanent") - except RuntimeError as e: - print(f"GPU exhausted at layer {i}, moving to CPU with swapping") - layer.to(CPU_DEVICE) - cpu_swappable_layers.add(i) - elif i >= (len(layers) - final_gpu_layers): - # Final layers on GPU permanently - try: - layer.to(GPU_DEVICE) - gpu_resident_layers.add(i) # MOVED: Only add if successful - if args.verbose: - print(f"Layer {i} (final) -> GPU permanent") - except RuntimeError as e: - print(f"CRITICAL: Cannot fit final layer {i} on GPU!") - raise e - else: - # Middle layers on CPU with swapping capability - layer.to(CPU_DEVICE) - cpu_swappable_layers.add(i) - if args.verbose: - print(f"Layer {i} (middle) -> CPU swappable") + if verbose: + print(f" Model loaded successfully") + print(f"Total layers extracted: {len(layers)}") - print(f"✓ {len(gpu_resident_layers)} layers permanently on GPU: {sorted(gpu_resident_layers)}") - print(f"✓ {len(cpu_swappable_layers)} layers on CPU with smart swapping: {sorted(cpu_swappable_layers)}") + return (model_patcher, layers) + def clear_model_cache_only(self): + """Safely clear only model-related caches, preserve node execution cache""" + import comfy.model_management as model_management + if hasattr(self, 'verbose') and self.verbose: + print("Clearing model cache only...") - # if args.verbose: - # if torch.cuda.is_available(): - # final_memory = torch.cuda.memory_allocated(0) - # dgls_memory = final_memory - baseline_memory - # print( - # f" DGLS TOTAL MEMORY: {dgls_memory / 1024 ** 3:.3f} GB for {len(gpu_resident_layers)} GPU layers") - # print(f" Average per GPU layer: {dgls_memory / len(gpu_resident_layers) / 1024 ** 2:.1f}MB") + model_management.unload_all_models() + comfy.model_management.cleanup_models_gc() + model_management.soft_empty_cache(True) - # print("\n=== Module Size Debug ===") - # total_module_mem = 0 - # for name, module in model.model.named_modules(): - # if hasattr(module, 'weight') or len(list(module.parameters())) > 0: - # module_mem = sum(p.numel() * p.element_size() for p in module.parameters()) - # if module_mem > 0: - # print(f"Module {name}: {module_mem / 1024 ** 2:.1f} MB") - # total_module_mem += module_mem - # print(f"Total module memory: {total_module_mem / 1024 ** 3:.2f} GB") + # GPU cleanup + torch.cuda.empty_cache() + gc.collect() + torch.cuda.synchronize() - global device_cache - device_cache = LayerDeviceCache(model, layers) + if hasattr(self, 'verbose') and self.verbose: + print("✓ Model cache cleared safely") - if args.verbose: - analysis_result = print_memory_optimization_analysis(model, layers, args) - # debug_layer_locations(layers, device_cache) + def nuke_all_caches(self): + """Clear ALL caches - may break downstream nodes""" + if hasattr(self, 'verbose') and self.verbose: + print("Clearing ALL ComfyUI caches...") - print("Calculating layer sizes...") - global layer_sizes_mb - for i, layer in enumerate(layers): - layer_sizes_mb[i] = sum(p.numel() * p.element_size() for p in layer.parameters()) / (1024 * 1024) - if args.verbose: - print(f" Layer {i}: {layer_sizes_mb[i]:.1f}MB") - print() + # Clear models first + self.clear_model_cache_only() - - if cast_target: - print(f"\n=== LAYER CASTED DTYPES ===") - print(f"Total layers: {len(layers)}") - - for i, layer in enumerate(layers): - layer_type = type(layer).__name__ - - # Get parameter dtypes for this layer - param_dtypes = set() - param_count = 0 - for name, param in layer.named_parameters(): - param_dtypes.add(param.dtype) - param_count += 1 - - # Get buffer dtypes - buffer_dtypes = set() - buffer_count = 0 - for name, buffer in layer.named_buffers(): - buffer_dtypes.add(buffer.dtype) - buffer_count += 1 - - # Format dtype info - dtype_info = [] - if param_dtypes: - dtype_info.append(f"params: {list(param_dtypes)}") - if buffer_dtypes: - dtype_info.append(f"buffers: {list(buffer_dtypes)}") - - dtype_str = ", ".join(dtype_info) if dtype_info else "no tensors" - - print(f" {i:2d}: {layer_type:<20} | {dtype_str}") - - print("Calculating layer sizes for Casted...") - print(f"\n=== LAYER CASTED SIZES ===") - # global layer_sizes_mb - for i, layer in enumerate(layers): - layer_sizes_mb[i] = sum(p.numel() * p.element_size() for p in layer.parameters()) / (1024 * 1024) - if args.verbose: - print(f" Layer {i}: {layer_sizes_mb[i]:.1f}MB") - - print("===============================\n") - print() - - # # Add this after casting - # for i in [15, 20, 25]: # Sample swappable layers - # for name, param in layers[i].named_parameters(): - # print( - # f"CHECK: Layer {i} param {name}: {param.dtype}, size: {param.numel() * param.element_size() / 1024 ** 2:.1f}MB") - - for layer_idx in cpu_swappable_layers: - add_smart_swapping_to_layer( - layers[layer_idx], - layer_idx, - layers, - gpu_resident_layers, - cpu_swappable_layers) - - for layer_idx in gpu_resident_layers: - layer = layers[layer_idx] - original_forward = layer.forward - - def create_resident_forward(layer_idx, original_forward): - current_args = args - @wraps(original_forward) - def resident_forward( *args_tuple, **kwargs): - result = original_forward( *args_tuple, **kwargs) - return result - return resident_forward - - layer.forward = create_resident_forward(layer_idx, original_forward) - - if args.verbose: - print("✓ Dynamic swapping successfully integrated with ComfyUI") - - return (model,) - -# ============================================================================= -# MIXED PRECISION FRAMEWORK CODE by obisin -# ============================================================================= - - -class CastingHandler: - def __init__(self): - self.original_clip_grad_norm = None - self.dtype_conversion_map = { - torch.float8_e4m3fn: torch.bfloat16, - torch.float8_e5m2: torch.bfloat16, - } - self.safe_dtypes = {torch.float32, torch.float16, torch.bfloat16} - self.is_patched = False - - - def cast_to_dtype(self, param, target_device, cast_target=None): - """Cast specific source dtype to target dtype""" - source_dtype = param.dtype - - if cast_target: - cast_from, cast_to = cast_target - - # Map string to torch dtypes - dtype_map = { - 'f32': torch.float32, - 'bf16': torch.bfloat16, - 'f16': torch.float16, - 'f8_e4m3': torch.float8_e4m3fn, - 'f8_e5m2': torch.float8_e5m2, - 'nf4': 'nf4', # Special handling needed NOT TESTED - 'fp4': 'fp4' # Special handling needed - } - - # Only cast if source matches the FROM dtype - if source_dtype == dtype_map[cast_from]: - # Handle 4-bit quantization - if cast_to in ['nf4', 'fp4']: - try: - import bitsandbytes as bnb - quantized, quant_state = bnb.functional.quantize_4bit( - param.data, blocksize=64, quant_type=cast_to - ) - # Return dequantized version on target device - return bnb.functional.dequantize_4bit( - quantized.to(target_device), quant_state.to(target_device) - ) - except ImportError: - raise RuntimeError( - " bitsandbytes required for 4-bit quantization. Install with: pip install bitsandbytes") - except Exception as e: - raise RuntimeError(f" 4-bit quantization failed: {e}") - else: - # Regular dtype casting - return param.to(target_device, dtype=dtype_map[cast_to], non_blocking=True) - - # No casting needed - return param.to(target_device, non_blocking=True) - - def precast_all_layers(self, cpu_swappable_layers, layers, cast_target=None): - """Cast all layer dtypes once at startup using existing cast_to_dtype""" - if not cast_target: - print(" No cast_target specified, skipping precasting") - return - - print(" Pre-casting layer dtypes (one-time setup)...") - casted_layers = 0 - - for i, layer in enumerate(layers): - if i in cpu_swappable_layers: # Only cast swappable layers - layer_casted = False - - for name, param in layer.named_parameters(): - # Use your existing function but stay on current device - new_param = self.cast_to_dtype(param, param.device, cast_target) # Keep on same device - if new_param.dtype != param.dtype: - param.data = new_param.data - layer_casted = True - - if layer_casted: - casted_layers += 1 - - print(f"✓ Pre-casted {casted_layers} layers using existing cast_to_dtype function") + # Clear execution cache + try: + import execution + # This is the risky part - affects ALL nodes + if hasattr(execution, 'PromptServer'): + server = execution.PromptServer.instance + if hasattr(server, 'cache'): + server.cache.outputs.clear() + server.cache.ui.clear() + server.cache.objects.clear() + print(" Execution cache cleared - other nodes may reload") + except Exception as e: + print(f" Could not clear execution cache: {e}")