""" Memory management module for SeedVR2 Handles VRAM usage, cache management, and memory optimization Extracted from: seedvr2.py (lines 373-405, 607-626, 1016-1044) """ import os import torch import gc import time import platform import psutil from typing import Tuple, Optional from src.common.cache import Cache from src.models.dit_v2.rope import RotaryEmbeddingBase try: from comfy import model_management as mm COMFYUI_AVAILABLE = True except: COMFYUI_AVAILABLE = False pass def get_basic_vram_info(): if platform.system() == "Darwin": mem = psutil.virtual_memory() free_memory = mem.total - mem.used total_memory = mem.total else: """🔍 Méthode basique avec PyTorch natif""" if not torch.cuda.is_available(): return {"error": "CUDA not available"} # Mémoire libre et totale (en bytes) free_memory, total_memory = torch.cuda.mem_get_info() # Conversion en GB free_gb = free_memory / (1024**3) total_gb = total_memory / (1024**3) return { "free_gb": free_gb, "total_gb": total_gb } # Initial VRAM check at module load vram_info = get_basic_vram_info() if "error" not in vram_info: print(f"📊 Initial VRAM status: {vram_info['free_gb']:.2f}GB free / {vram_info['total_gb']:.2f}GB total") else: print(f"⚠️ VRAM check: {vram_info['error']} - SeedVR2 requires an NVIDIA GPU") def get_vram_usage() -> Tuple[float, float, float]: """ Get current VRAM usage (allocated, reserved, peak) Returns: tuple: (allocated_gb, reserved_gb, max_allocated_gb) Returns (0, 0, 0) if CUDA not available """ if platform.system() == "Darwin": allocated = torch.mps.current_allocated_memory() / (1024**3) reserved = torch.mps.driver_allocated_memory() / (1024**3) max_allocated = 0 return allocated, reserved, max_allocated if torch.cuda.is_available(): allocated = torch.cuda.memory_allocated() / (1024**3) reserved = torch.cuda.memory_reserved() / (1024**3) max_allocated = torch.cuda.max_memory_allocated() / (1024**3) return allocated, reserved, max_allocated return 0, 0, 0 def clear_vram_cache(debug) -> None: """Clear VRAM cache and run garbage collection""" debug.log("Clearing VRAM cache...", category="cleanup") if platform.system() == "Darwin": torch.mps.empty_cache() if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() gc.collect() def reset_vram_peak(debug) -> None: """ Reset VRAM peak counter for new tracking """ debug.log("Resetting VRAM peak memory statistics", category="memory") if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() def preinitialize_rope_cache(runner, debug) -> None: """ 🚀 Pre-initialize RoPE cache to avoid OOM at first launch Args: runner: The model runner containing DiT and VAE models debug: Optional Debug instance """ debug.log("Pre-initializing RoPE cache to avoid OOM...", category="setup") try: # Create dummy tensors to simulate common shapes # Format: [batch, channels, frames, height, width] for vid_shape # Format: [batch, seq_len] for txt_shape common_shapes = [ # Common video resolutions (torch.tensor([[1, 3, 3]], dtype=torch.long), torch.tensor([[77]], dtype=torch.long)), # 1 frame, 77 tokens (torch.tensor([[4, 3, 3]], dtype=torch.long), torch.tensor([[77]], dtype=torch.long)), # 4 frames (torch.tensor([[5, 3, 3]], dtype=torch.long), torch.tensor([[77]], dtype=torch.long)), # 5 frames (4n+1 format) (torch.tensor([[1, 4, 4]], dtype=torch.long), torch.tensor([[77]], dtype=torch.long)), # Higher resolution ] # Create mock cache for pre-initialization temp_cache = Cache() # Access RoPE modules in DiT (recursive search) def find_rope_modules(module): rope_modules = [] for name, child in module.named_modules(): if hasattr(child, 'get_freqs') and callable(getattr(child, 'get_freqs')): rope_modules.append((name, child)) return rope_modules rope_modules = find_rope_modules(runner.dit) # Pre-calculate for each RoPE module found for name, rope_module in rope_modules: # Temporarily move module to CPU if necessary original_device = next(rope_module.parameters()).device if list(rope_module.parameters()) else torch.device('cpu') rope_module.to('cpu') try: for vid_shape, txt_shape in common_shapes: cache_key = f"720pswin_by_size_bysize_{tuple(vid_shape[0].tolist())}_sd3.mmrope_freqs_3d" def compute_freqs(): try: # Calculate with reduced dimensions to avoid OOM with torch.no_grad(): # Detect RoPE module type module_type = type(rope_module).__name__ if module_type == 'NaRotaryEmbedding3d': # NaRotaryEmbedding3d: only takes shape (vid_shape) return rope_module.get_freqs(vid_shape.cpu()) else: # Standard RoPE: takes vid_shape and txt_shape return rope_module.get_freqs(vid_shape.cpu(), txt_shape.cpu()) except Exception as e: debug.log(f"Failed for {cache_key}: {e}", level="WARNING", category="cache") # Return empty tensors as fallback time.sleep(1) clear_vram_cache(debug) return torch.zeros(1, 64) # Store in cache temp_cache(cache_key, compute_freqs) except Exception as e: debug.log(f"Error in module {name}: {e}", level="ERROR", category="cache") finally: # Restore to original device rope_module.to(original_device) # Copy temporary cache to runner cache if hasattr(runner, 'cache'): runner.cache.cache.update(temp_cache.cache) else: runner.cache = temp_cache except Exception as e: debug.log(f"Error during RoPE pre-init: {e}", level="WARNING", category="setup", force=True) debug.log("Model will work but could have OOM at first launch", level="WARNING", category="setup", force=True) def clear_rope_lru_caches(model) -> int: """Clear ALL LRU caches from RoPE modules""" cleared_count = 0 for name, module in model.named_modules(): if hasattr(module, 'get_axial_freqs') and hasattr(module.get_axial_freqs, 'cache_clear'): module.get_axial_freqs.cache_clear() cleared_count += 1 return cleared_count def fast_model_cleanup(model): """Fast model cleanup without logs""" if model is None: return # Move to CPU model.to("cpu") # Clear parameters and buffers recursively def clear_recursive(m): for child in m.children(): clear_recursive(child) for param in m.parameters(): if param is not None: param.data = param.data.cpu() param.grad = None for buffer in m.buffers(): if buffer is not None: buffer.data = buffer.data.cpu() clear_recursive(model) def fast_ram_cleanup(): """Fast RAM cleanup without excessive logging""" # Garbage collection gc.collect() # Clear MPS cache if platform.system() == "Darwin": torch.mps.empty_cache() # Clear CUDA cache if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() torch.cuda.reset_peak_memory_stats() # Clear PyTorch internal caches try: torch._C._clear_cache() except: pass def clear_all_caches(runner, debug, offload_vae=False) -> int: """ Aggressively clear all caches from runner and model. Optimized to only process what's necessary. Args: runner: The runner instance to clear caches from debug: Optional Debug instance for logging offload_vae: If True, also moves VAE to CPU and clears its caches """ if not runner: return 0 cleaned_items = 0 # Early exit if no caches to clear has_cache = hasattr(runner, 'cache') and hasattr(runner.cache, 'cache') if not has_cache and not hasattr(runner, 'dit') and not (offload_vae and hasattr(runner, 'vae')): return 0 # Clear main runner cache efficiently if has_cache and runner.cache.cache: cache_entries = len(runner.cache.cache) # Process all cache items to properly free memory for key, value in list(runner.cache.cache.items()): if torch.is_tensor(value): # Force deallocation of tensor storage if value.is_cuda or value.is_mps: value.data = value.data.cpu() value.grad = None if value.numel() > 0: value.data.set_() # Release underlying storage elif isinstance(value, (list, tuple)): for item in value: if torch.is_tensor(item): if item.is_cuda or item.is_mps: item.data = item.data.cpu() item.grad = None if item.numel() > 0: item.data.set_() # Clear the cache after processing runner.cache.cache.clear() cleaned_items += cache_entries debug.log(f"Cleared {cache_entries} cache entries", category="success") # Clear any accumulated state in blocks if hasattr(runner, 'dit'): model = runner.dit if hasattr(model, 'dit_model'): model = model.dit_model # Clear RoPE LRU caches rope_caches_cleared = clear_rope_lru_caches(model) cleaned_items += rope_caches_cleared if rope_caches_cleared > 0: debug.log(f"Cleared {rope_caches_cleared} RoPE LRU caches", category="success") # Clear block attributes if needed if hasattr(model, 'blocks'): block_attrs_cleared = 0 # Define PyTorch's essential attributes that must NOT be deleted essential_attrs = { '_modules', '_parameters', '_buffers', '_forward_hooks', '_forward_pre_hooks', '_backward_hooks', '_backward_pre_hooks', '_state_dict_hooks', '_state_dict_pre_hooks', '_load_state_dict_pre_hooks', '_load_state_dict_post_hooks', '_non_persistent_buffers_set', '_version', '_is_full_backward_hook', 'training', '_original_forward', # BlockSwap attribute '_is_io_wrapped', # BlockSwap attribute '_block_idx', # BlockSwap attribute } for idx, block in enumerate(model.blocks): # Get all attributes that look like caches attrs_to_remove = [] for attr_name in list(block.__dict__.keys()): # Only remove cache-like attributes, not essential PyTorch attributes if (attr_name not in essential_attrs and ('cache' in attr_name or 'temp' in attr_name or (attr_name.startswith('_') and not attr_name.startswith('__') and attr_name not in essential_attrs))): attrs_to_remove.append(attr_name) # Remove the identified attributes for attr_name in attrs_to_remove: try: delattr(block, attr_name) block_attrs_cleared += 1 except AttributeError: pass # Already deleted or doesn't exist if block_attrs_cleared > 0: debug.log(f"Cleared {block_attrs_cleared} temporary attributes from blocks", category="success") # Clear any temporary attributes that might accumulate temp_attrs = ['_temp_cache', '_block_cache', '_swap_cache', '_generation_cache', '_rope_cache', '_intermediate_cache', '_backward_cache'] # Check both runner and model for these attributes for obj in [runner, getattr(runner, 'dit', None)]: if obj is None: continue # Handle wrapped models if hasattr(obj, 'dit_model'): obj = obj.dit_model for attr in temp_attrs: if hasattr(obj, attr): delattr(obj, attr) cleaned_items += 1 debug.log(f"Cleared {attr} from {type(obj).__name__}", category="success") # Handle VAE offloading if requested if offload_vae and hasattr(runner, 'vae') and runner.vae is not None: debug.log("Moving VAE to CPU and clearing intermediate tensors", category="cleanup") # Clear any intermediate tensors/buffers in VAE vae_caches_cleared = 0 for module in runner.vae.modules(): # Clear module-specific caches for cache_attr in ['_temp_cache', '_intermediate_cache']: if hasattr(module, cache_attr): delattr(module, cache_attr) vae_caches_cleared += 1 # Clear any CUDA tensors in module attributes for attr_name in list(vars(module).keys()): attr = getattr(module, attr_name, None) if torch.is_tensor(attr) and attr.is_cuda: # Move tensor to CPU if it's not a parameter/buffer if attr_name not in module._parameters and attr_name not in module._buffers: setattr(module, attr_name, attr.cpu()) vae_caches_cleared += 1 if vae_caches_cleared > 0: debug.log(f"Cleared {vae_caches_cleared} VAE caches", category="success") # Move entire VAE to CPU (preserves model for reuse) runner.vae = runner.vae.to('cpu') debug.log("VAE moved to CPU, intermediate tensors cleared", category="success") cleaned_items += vae_caches_cleared # Force garbage collection gc.collect(2) # Collect all generations # Clear MPS cache if platform.system() == "Darwin": torch.mps.empty_cache() # Clear CUDA cache if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() return cleaned_items