diff --git a/__init__.py b/__init__.py index f10917e..90c4c6e 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,3 @@ -from .src.interfaces.comfyui_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .src.interfaces import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/src/core/generation.py b/src/core/generation.py index 4522c7a..1aa1be4 100644 --- a/src/core/generation.py +++ b/src/core/generation.py @@ -1190,7 +1190,7 @@ def postprocess_all_batches(ctx: Optional[Dict[str, Any]] = None, # Format channel info for readability channels_str = "RGBA" if Cf == 4 else "RGB" if Cf == 3 else f"{Cf}-channel" - debug.log(f"Final video assembled: Total frames: {total_frames}, Resolution: {Wf}x{Hf}px, Channels: {channels_str}", category="video", force=True) + debug.log(f"Final video assembled: Total frames: {total_frames}, Resolution: {Wf}x{Hf}px, Channels: {channels_str}", category="generation", force=True) else: ctx['final_video'] = torch.empty((0, 0, 0, 0), dtype=torch.float16) debug.log("No frame to assemble", level="WARNING", category="video", force=True) diff --git a/src/core/model_manager.py b/src/core/model_manager.py index 5b311fa..2e7375d 100644 --- a/src/core/model_manager.py +++ b/src/core/model_manager.py @@ -93,13 +93,16 @@ def _describe_blockswap_config(config: Optional[Dict[str, Any]]) -> str: if config is None: return "disabled" - blocks = config.get("blocks_to_swap", 0) - if blocks <= 0: + blocks_to_swap = config.get("blocks_to_swap", 0) + swap_io_components = config.get("swap_io_components", False) + + # Early return only if both block swap and I/O swap are disabled + if blocks_to_swap <= 0 and not swap_io_components: return "disabled" - block_text = "block" if blocks == 1 else "blocks" - parts = [f"{blocks} {block_text}"] - if config.get("swap_io_components", False): + block_text = "block" if blocks_to_swap <= 1 else "blocks" + parts = [f"{blocks_to_swap} {block_text}"] + if swap_io_components: parts.append("I/O offload") return f"enabled ({', '.join(parts)})" @@ -338,11 +341,17 @@ def _update_dit_config_inplace( if blockswap_changed: # Determine change type from config comparison old_blocks = cached_blockswap_config.get("blocks_to_swap", 0) if cached_blockswap_config else 0 + old_swap_io = cached_blockswap_config.get("swap_io_components", False) if cached_blockswap_config else False new_blocks = block_swap_config.get("blocks_to_swap", 0) if block_swap_config else 0 + new_swap_io = block_swap_config.get("swap_io_components", False) if block_swap_config else False - # If old config had BlockSwap, clean it up first - if old_blocks > 0: - if new_blocks == 0: + # Check if old config had ANY BlockSwap features (blocks or I/O components) + had_blockswap = old_blocks > 0 or old_swap_io + has_blockswap = new_blocks > 0 or new_swap_io + + # If old config had BlockSwap features, clean them up first + if had_blockswap: + if not has_blockswap: # Disabling BlockSwap completely debug.log("Disabling BlockSwap completely", category="blockswap") cleanup_blockswap(runner=cached_runner, keep_state_for_cache=False) @@ -354,7 +363,8 @@ def _update_dit_config_inplace( cached_runner._blockswap_active = False # Apply new BlockSwap if configured (before torch.compile) - if block_swap_config and block_swap_config.get("blocks_to_swap", 0) > 0: + if block_swap_config and (block_swap_config.get("blocks_to_swap", 0) > 0 or + block_swap_config.get("swap_io_components", False)): cached_runner.dit = model # Temporarily set for BlockSwap apply_block_swap_to_dit(cached_runner, block_swap_config, debug) model = cached_runner.dit # Get potentially wrapped model @@ -735,7 +745,7 @@ def _load_gguf_state(checkpoint_path: str, device: str, debug: Optional[Debug], # Progress reporting if (i + 1) % 100 == 0: - debug.log(f" Loaded {i+1}/{total_tensors} tensors...", category="dit") + debug.log(f" Loaded {i+1}/{total_tensors} tensors...", category="dit") debug.log(f"Successfully loaded {len(state_dict)} tensors to {device}", category="success") @@ -842,9 +852,9 @@ class GGUFTensor(torch.Tensor): return final_result except Exception as e: self.debug.log(f"Numpy fallback also failed: {e}", level="WARNING", category="dit", force=True) - self.debug.log(f" Tensor type: {self.tensor_type}", level="WARNING", category="dit", force=True) - self.debug.log(f" Shape: {self.shape}", level="WARNING", category="dit", force=True) - self.debug.log(f" Target shape: {self.tensor_shape}", level="WARNING", category="dit", force=True) + self.debug.log(f" Tensor type: {self.tensor_type}", level="WARNING", category="dit", force=True) + self.debug.log(f" Shape: {self.shape}", level="WARNING", category="dit", force=True) + self.debug.log(f" Target shape: {self.tensor_shape}", level="WARNING", category="dit", force=True) traceback.print_exc() # Return regular tensor as last resort @@ -1405,14 +1415,14 @@ def _report_parameter_mismatches(state: Dict[str, torch.Tensor], if unmatched: debug.log(f"Warning: {len(unmatched)} parameters from GGUF not found in model", level="WARNING", category="dit", force=True) - debug.log(f" First few unmatched: {unmatched[:5]}", level="WARNING", category="dit", force=True) + debug.log(f" First few unmatched: {unmatched[:5]}", level="WARNING", category="dit", force=True) # Check for missing model parameters missing = [name for name in model_state if name not in loaded_names] if missing: debug.log(f"Warning: {len(missing)} model parameters not loaded from GGUF", level="WARNING", category="dit", force=True) - debug.log(f" First few missing: {missing[:5]}", level="WARNING", category="dit", force=True) + debug.log(f" First few missing: {missing[:5]}", level="WARNING", category="dit", force=True) def _initialize_meta_buffers_wrapped(model: torch.nn.Module, target_device: str, debug: Debug) -> None: diff --git a/src/interfaces/__init__.py b/src/interfaces/__init__.py index e69de29..05c1674 100644 --- a/src/interfaces/__init__.py +++ b/src/interfaces/__init__.py @@ -0,0 +1,38 @@ +""" +SeedVR2 ComfyUI Nodes +Central registry for all SeedVR2 nodes +""" + +from .video_upscaler import SeedVR2VideoUpscaler +from .dit_model_loader import SeedVR2LoadDiTModel +from .vae_model_loader import SeedVR2LoadVAEModel +from .torch_compile_settings import SeedVR2TorchCompileSettings + +# ComfyUI Node Mappings - +NODE_CLASS_MAPPINGS = { + "SeedVR2VideoUpscaler": SeedVR2VideoUpscaler, + "SeedVR2LoadDiTModel": SeedVR2LoadDiTModel, + "SeedVR2LoadVAEModel": SeedVR2LoadVAEModel, + "SeedVR2TorchCompileSettings": SeedVR2TorchCompileSettings, +} + +# Human-readable node display names - unchanged +NODE_DISPLAY_NAME_MAPPINGS = { + "SeedVR2VideoUpscaler": "SeedVR2 Video Upscaler", + "SeedVR2LoadDiTModel": "SeedVR2 (Down)Load DiT Model", + "SeedVR2LoadVAEModel": "SeedVR2 (Down)Load VAE Model", + "SeedVR2TorchCompileSettings": "SeedVR2 Torch Compile Settings", +} + +__version__ = "2.0.0" +__author__ = "numz, adrientoupet" +__description__ = "ComfyUI integration of ByteDance-Seed's SeedVR2: One-step diffusion-based video/image upscaling with memory-efficient inference" + +__all__ = [ + 'SeedVR2VideoUpscaler', + 'SeedVR2LoadDiTModel', + 'SeedVR2LoadVAEModel', + 'SeedVR2TorchCompileSettings', + 'NODE_CLASS_MAPPINGS', + 'NODE_DISPLAY_NAME_MAPPINGS' +] \ No newline at end of file diff --git a/src/interfaces/dit_model_loader.py b/src/interfaces/dit_model_loader.py new file mode 100644 index 0000000..f29b0fd --- /dev/null +++ b/src/interfaces/dit_model_loader.py @@ -0,0 +1,95 @@ +""" +SeedVR2 DiT Model Loader Node +Configure DiT (Diffusion Transformer) model with memory optimization +""" + +from typing import Dict, Any, Tuple + +from ..utils.model_registry import get_available_dit_models, DEFAULT_DIT +from ..optimization.memory_manager import get_device_list + + +class SeedVR2LoadDiTModel: + """ + Configure DiT (Diffusion Transformer) model loader with memory optimization + + Provides configuration for: + - Model selection and device placement + - BlockSwap memory optimization for limited VRAM + - Model caching between runs + - Optional torch.compile integration + + Returns: + SEEDVR2_DIT configuration dictionary for main upscaler node + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]: + devices = get_device_list() + return { + "required": { + "model": (get_available_dit_models(), { + "default": DEFAULT_DIT, + "tooltip": "DiT model for upscaling. Models will automatically download on first use. Additional models can be added to the ComfyUI models folder." + }), + "device": (devices, { + "default": devices[0], + "tooltip": "Device to use for DiT processing" + }), + }, + "optional": { + "blocks_to_swap": ("INT", { + "default": 0, + "min": 0, + "max": 36, + "step": 1, + "tooltip": "Number of transformer blocks to swap to CPU. 0=disabled (fastest, most VRAM). Higher values save VRAM but are slower. 3B model: 0-32 blocks, 7B model: 0-36 blocks." + }), + "swap_io_components": ("BOOLEAN", { + "default": False, + "tooltip": "Offload input/output embeddings and norm layers to CPU for additional VRAM savings (slower)" + }), + "cache_in_ram": ("BOOLEAN", { + "default": False, + "tooltip": "Keep model in RAM between runs for faster reuse. Useful for batch processing." + }), + "torch_compile_args": ("TORCH_COMPILE_ARGS", { + "tooltip": "Optional torch.compile settings from SeedVR2 Torch Compile Settings node for speedup" + }), + } + } + + RETURN_TYPES = ("SEEDVR2_DIT",) + FUNCTION = "create_config" + CATEGORY = "SEEDVR2" + DESCRIPTION = ( + "Configure DiT model for SeedVR2 upscaling. Supports BlockSwap for limited VRAM, " + "model caching, and torch.compile optimization. Connect output to SeedVR2 Video Upscaler." + ) + + def create_config(self, model: str, device: str, blocks_to_swap: int = 0, + swap_io_components: bool = False, cache_in_ram: bool = False, + torch_compile_args: Dict[str, Any] = None) -> Tuple[Dict[str, Any]]: + """ + Create DiT model configuration for SeedVR2 main node + + Args: + model: Model filename to load + device: Target device for model execution + blocks_to_swap: Number of transformer blocks to swap to CPU (0=disabled) + swap_io_components: Whether to offload I/O components to CPU + cache_in_ram: Whether to keep model in RAM between runs + torch_compile_args: Optional torch.compile configuration from settings node + + Returns: + Tuple containing configuration dictionary for SeedVR2 main node + """ + config = { + "model": model, + "device": device, + "blocks_to_swap": blocks_to_swap, + "swap_io_components": swap_io_components, + "cache_in_ram": cache_in_ram, + "torch_compile_args": torch_compile_args, + } + return (config,) \ No newline at end of file diff --git a/src/interfaces/torch_compile_settings.py b/src/interfaces/torch_compile_settings.py new file mode 100644 index 0000000..0c67cd7 --- /dev/null +++ b/src/interfaces/torch_compile_settings.py @@ -0,0 +1,91 @@ +""" +SeedVR2 Torch Compile Settings Node +Configure torch.compile optimization for DiT and VAE models +""" + +from typing import Dict, Any, Tuple + + +class SeedVR2TorchCompileSettings: + """Configure torch.compile optimization for DiT and VAE models""" + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]: + return { + "required": { + "backend": (["inductor", "cudagraphs"], { + "default": "inductor", + "tooltip": ( + "Compilation backend:\n" + "• inductor: Full optimization with Triton kernel generation and fusion\n" + "• cudagraphs: Lightweight, only wraps model with CUDA graphs, no kernel optimization" + ) + }), + "mode": (["default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"], { + "default": "default", + "tooltip": ( + "Optimization level (compilation time vs runtime speed):\n" + "• default: Fast compilation, good speedup\n" + "• reduce-overhead: Lower overhead, better for smaller models\n" + "• max-autotune: Slowest compilation, best runtime (recommended for production)\n" + "• max-autotune-no-cudagraphs: Like max-autotune but without cudagraphs" + ) + }), + "fullgraph": ("BOOLEAN", { + "default": False, + "tooltip": "Compile entire model as single graph (faster but less flexible). May fail with dynamic shapes." + }), + "dynamic": ("BOOLEAN", { + "default": False, + "tooltip": "Handle varying input shapes without recompilation. Useful for different resolutions/batch sizes." + }), + "dynamo_cache_size_limit": ("INT", { + "default": 64, + "min": 0, + "max": 1024, + "step": 1, + "tooltip": "Maximum cached compiled versions per function. Increase if using many different input shapes." + }), + "dynamo_recompile_limit": ("INT", { + "default": 128, + "min": 0, + "max": 1024, + "step": 1, + "tooltip": "Maximum recompilation attempts before fallback to eager mode. Only increase if you see recompile_limit warnings" + }), + } + } + + RETURN_TYPES = ("TORCH_COMPILE_ARGS",) + FUNCTION = "create_args" + CATEGORY = "SEEDVR2" + DESCRIPTION = ( + "Configure torch.compile optimization for DiT and VAE speedup. " + "Connect to DiT and/or VAE model loader. Requires PyTorch 2.0+ and Triton." + ) + + def create_args(self, backend: str, mode: str, fullgraph: bool, dynamic: bool, + dynamo_cache_size_limit: int, dynamo_recompile_limit: int = 128) -> Tuple[Dict[str, Any]]: + """ + Create torch.compile configuration for model optimization + + Args: + backend: Compilation backend ("inductor" or "cudagraphs") + mode: Optimization mode ("default", "reduce-overhead", "max-autotune", etc.) + fullgraph: Whether to compile entire model as single graph + dynamic: Whether to handle varying input shapes without recompilation + dynamo_cache_size_limit: Maximum cached compiled versions per function + dynamo_recompile_limit: Maximum recompilation attempts before fallback + + Returns: + Tuple containing torch.compile configuration dictionary + """ + compile_args = { + "backend": backend, + "mode": mode, + "fullgraph": fullgraph, + "dynamic": dynamic, + "dynamo_cache_size_limit": dynamo_cache_size_limit, + "dynamo_recompile_limit": dynamo_recompile_limit, + } + return (compile_args,) \ No newline at end of file diff --git a/src/interfaces/vae_model_loader.py b/src/interfaces/vae_model_loader.py new file mode 100644 index 0000000..1e60f84 --- /dev/null +++ b/src/interfaces/vae_model_loader.py @@ -0,0 +1,127 @@ +""" +SeedVR2 VAE Model Loader Node +Configure VAE (Variational Autoencoder) model with tiling support +""" + +from typing import Dict, Any, Tuple + +from ..utils.model_registry import get_available_vae_models, DEFAULT_VAE +from ..optimization.memory_manager import get_device_list + + +class SeedVR2LoadVAEModel: + """ + Configure VAE (Variational Autoencoder) model loader with tiling support + + Provides configuration for: + - Model selection and device placement + - Tiled encoding/decoding for VRAM reduction + - Tile size and overlap control + - Model caching between runs + - Optional torch.compile integration + + Returns: + SEEDVR2_VAE configuration dictionary for main upscaler node + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]: + devices = get_device_list() + return { + "required": { + "model": (get_available_vae_models(), { + "default": DEFAULT_VAE, + "tooltip": "VAE model for encoding/decoding. Models will automatically download on first use. Additional models can be added to the ComfyUI models folder." + }), + "device": (devices, { + "default": devices[0], + "tooltip": "Device to use for VAE processing" + }), + }, + "optional": { + "encode_tiled": ("BOOLEAN", { + "default": False, + "tooltip": "Enable tiled encoding to reduce VRAM during encoding" + }), + "encode_tile_size": ("INT", { + "default": 512, + "min": 64, + "step": 32, + "tooltip": "Size of encoding tiles in pixels. Smaller = less VRAM but more seams/artifacts and slower. Larger = more VRAM but better quality and faster." + }), + "encode_tile_overlap": ("INT", { + "default": 64, + "min": 0, + "step": 32, + "tooltip": "Pixel overlap between encoding tiles to reduce visible seams. Higher = better blending but slower processing." + }), + "decode_tiled": ("BOOLEAN", { + "default": False, + "tooltip": "Enable tiled decoding to reduce VRAM during decoding" + }), + "decode_tile_size": ("INT", { + "default": 512, + "min": 64, + "step": 32, + "tooltip": "Size of decoding tiles in pixels. Smaller = less VRAM but more seams/artifacts and slower. Larger = more VRAM but better quality and faster." + }), + "decode_tile_overlap": ("INT", { + "default": 64, + "min": 0, + "step": 32, + "tooltip": "Pixel overlap between decoding tiles to reduce visible seams. Higher = better blending but slower processing." + }), + "cache_in_ram": ("BOOLEAN", { + "default": False, + "tooltip": "Keep model in RAM between runs for faster reuse. Useful for batch processing." + }), + "torch_compile_args": ("TORCH_COMPILE_ARGS", { + "tooltip": "Optional torch.compile settings from SeedVR2 Torch Compile Settings node for speedup" + }), + } + } + + RETURN_TYPES = ("SEEDVR2_VAE",) + FUNCTION = "create_config" + CATEGORY = "SEEDVR2" + DESCRIPTION = ( + "Configure VAE model for SeedVR2 encoding/decoding. Supports tiled processing for VRAM reduction, " + "model caching, and torch.compile optimization. Connect output to SeedVR2 Video Upscaler." + ) + + def create_config(self, model: str, device: str, encode_tiled: bool = False, + encode_tile_size: int = 512, encode_tile_overlap: int = 64, + decode_tiled: bool = False, decode_tile_size: int = 512, + decode_tile_overlap: int = 64, cache_in_ram: bool = False, + torch_compile_args: Dict[str, Any] = None) -> Tuple[Dict[str, Any]]: + """ + Create VAE model configuration for SeedVR2 main node + + Args: + model: Model filename to load + device: Target device for model execution + encode_tiled: Enable tiled encoding + encode_tile_size: Tile size for encoding + encode_tile_overlap: Tile overlap for encoding + decode_tiled: Enable tiled decoding + decode_tile_size: Tile size for decoding + decode_tile_overlap: Tile overlap for decoding + cache_in_ram: Whether to keep model in RAM between runs + torch_compile_args: Optional torch.compile configuration from settings node + + Returns: + Tuple containing configuration dictionary for SeedVR2 main node + """ + config = { + "model": model, + "device": device, + "encode_tiled": encode_tiled, + "encode_tile_size": encode_tile_size, + "encode_tile_overlap": encode_tile_overlap, + "decode_tiled": decode_tiled, + "decode_tile_size": decode_tile_size, + "decode_tile_overlap": decode_tile_overlap, + "cache_in_ram": cache_in_ram, + "torch_compile_args": torch_compile_args, + } + return (config,) \ No newline at end of file diff --git a/src/interfaces/comfyui_node.py b/src/interfaces/video_upscaler.py similarity index 50% rename from src/interfaces/comfyui_node.py rename to src/interfaces/video_upscaler.py index dda7b00..a2de75d 100644 --- a/src/interfaces/comfyui_node.py +++ b/src/interfaces/video_upscaler.py @@ -1,22 +1,14 @@ -# ComfyUI Node Interface -# Clean interface for SeedVR2 VideoUpscaler integration with ComfyUI -# Extracted from original seedvr2.py lines 1731-1812 +""" +SeedVR2 Video Upscaler Node +Main ComfyUI node for high-quality video upscaling using diffusion models +""" -import os -import time import torch from typing import Tuple, Dict, Any, Optional -from ..utils.constants import get_base_cache_dir, get_script_directory +from ..utils.constants import get_base_cache_dir from ..utils.downloads import download_weight -from ..utils.model_registry import ( - get_available_vae_models, - get_available_dit_models, - DEFAULT_DIT, - DEFAULT_VAE -) from ..utils.debug import Debug -from ..core.model_manager import configure_runner from ..core.generation import ( setup_device_environment, prepare_generation_context, @@ -29,8 +21,7 @@ from ..core.generation import ( ) from ..optimization.memory_manager import ( cleanup_text_embeddings, - complete_cleanup, - get_device_list + complete_cleanup ) # Import ComfyUI progress reporting @@ -39,9 +30,8 @@ try: except ImportError: ProgressBar = None -script_directory = get_script_directory() -class SeedVR2: +class SeedVR2VideoUpscaler: """ SeedVR2 Video Upscaler ComfyUI Node @@ -54,14 +44,13 @@ class SeedVR2: """ def __init__(self): - """Initialize SeedVR2 node""" + """Initialize SeedVR2VideoUpscaler node""" self.runner = None self.ctx = None self.debug = None self._dit_model_name = "" self._vae_model_name = "" - @classmethod def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]: """ @@ -133,7 +122,6 @@ class SeedVR2: } } - # Define return types for ComfyUI RETURN_TYPES = ("IMAGE",) FUNCTION = "execute" CATEGORY = "SEEDVR2" @@ -152,24 +140,8 @@ class SeedVR2: Args: pixels: Input video frames as tensor (N, H, W, C) in [0, 1] range - dit: DiT model configuration from SeedVR2LoadDiTModel node containing: - - model: Model filename - - device: Target device - - blocks_to_swap: BlockSwap configuration - - swap_io_components: I/O component swap flag - - cache_in_ram: Model caching flag - - torch_compile_args: Optional compilation settings - vae: VAE model configuration from SeedVR2LoadVAEModel node containing: - - model: Model filename - - device: Target device - - encode_tiled: Enable tiled encoding - - encode_tile_size: Encoding tile size - - encode_tile_overlap: Encoding tile overlap - - decode_tiled: Enable tiled decoding - - decode_tile_size: Decoding tile size - - decode_tile_overlap: Decoding tile overlap - - cache_in_ram: Model caching flag - - torch_compile_args: Optional compilation settings + dit: DiT model configuration from SeedVR2LoadDiTModel node + vae: VAE model configuration from SeedVR2LoadVAEModel node seed: Random seed for reproducible generation new_resolution: Target resolution for shortest edge (maintains aspect ratio) batch_size: Frames per batch (minimum 5 for temporal consistency) @@ -184,70 +156,63 @@ class SeedVR2: Raises: ValueError: If model files cannot be downloaded or configuration is invalid - RuntimeError: If generation pipeline fails - - Note: - - Models are automatically downloaded on first use - - Minimum batch_size of 5 recommended for temporal consistency - - Higher batch_size improves quality but requires more VRAM + RuntimeError: If generation fails """ - - # Unpack DiT configuration - dit_model = dit["model"] - dit_device = dit["device"] - blocks_to_swap = dit.get("blocks_to_swap", 0) - swap_io_components = dit.get("swap_io_components", False) - cache_model_dit = dit.get("cache_in_ram", False) - dit_torch_compile_args = dit.get("torch_compile_args", None) - - # Unpack VAE configuration - vae_model = vae["model"] - vae_device = vae["device"] - encode_tiled = vae.get("encode_tiled", False) - encode_tile_size = vae.get("encode_tile_size", 512) - encode_tile_overlap = vae.get("encode_tile_overlap", 64) - decode_tiled = vae.get("decode_tiled", False) - decode_tile_size = vae.get("decode_tile_size", 512) - decode_tile_overlap = vae.get("decode_tile_overlap", 64) - cache_model_vae = vae.get("cache_in_ram", False) - vae_torch_compile_args = vae.get("torch_compile_args", None) - - # Create block_swap_config if blocks_to_swap > 0 - block_swap_config = None - if blocks_to_swap > 0: - block_swap_config = { - "blocks_to_swap": blocks_to_swap, - "swap_io_components": swap_io_components, - } - - # Fixed parameters (could be exposed in future) - temporal_overlap = 0 - - # Initialize debug + # Initialize debug instance self.debug = Debug(enabled=enable_debug) debug = self.debug + + # Extract configuration from dict inputs + dit_model = dit["model"] + vae_model = vae["model"] + dit_device = dit["device"] + vae_device = vae["device"] + cache_model_dit = dit["cache_in_ram"] + cache_model_vae = vae["cache_in_ram"] + dit_torch_compile_args = dit.get("torch_compile_args") + vae_torch_compile_args = vae.get("torch_compile_args") + + # Extract VAE tiling configuration + encode_tiled = vae["encode_tiled"] + encode_tile_size = vae["encode_tile_size"] + encode_tile_overlap = vae["encode_tile_overlap"] + decode_tiled = vae["decode_tiled"] + decode_tile_size = vae["decode_tile_size"] + decode_tile_overlap = vae["decode_tile_overlap"] + + # Extract DiT BlockSwap configuration + block_swap_config = None + if dit.get("blocks_to_swap", 0) > 0 or dit.get("swap_io_components", False): + block_swap_config = { + "blocks_to_swap": dit["blocks_to_swap"], + "swap_io_components": dit["swap_io_components"], + } + + # Fixed parameters + temporal_overlap = 0 + + # Intro logo + debug.log("", category="none", force=True) + debug.log(" ╔══════════════════════════════════════════════════════════╗", category="none", force=True) + debug.log(" ║ ███████ ███████ ███████ ██████ ██ ██ ██████ ███████ ║", category="none", force=True) + debug.log(" ║ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ║", category="none", force=True) + debug.log(" ║ ███████ █████ █████ ██ ██ ██ ██ ██████ █████ ║", category="none", force=True) + debug.log(" ║ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ║", category="none", force=True) + debug.log(" ║ ███████ ███████ ███████ ██████ ████ ██ ██ ███████ ║", category="none", force=True) + debug.log(" ║ © ByteDance Seed · NumZ · AInVFX ║", category="none", force=True) + debug.log(" ╚══════════════════════════════════════════════════════════╝", category="none", force=True) + debug.log("", category="none", force=True) - self.debug.log("", category="none", force=True) - self.debug.log(" ╔══════════════════════════════════════════════════════════╗", category="none", force=True) - self.debug.log(" ║ ███████ ███████ ███████ ██████ ██ ██ ██████ ███████ ║", category="none", force=True) - self.debug.log(" ║ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ║", category="none", force=True) - self.debug.log(" ║ ███████ █████ █████ ██ ██ ██ ██ ██████ █████ ║", category="none", force=True) - self.debug.log(" ║ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ║", category="none", force=True) - self.debug.log(" ║ ███████ ███████ ███████ ██████ ████ ██ ██ ███████ ║", category="none", force=True) - self.debug.log(" ║ © ByteDance Seed · NumZ · AInVFX ║", category="none", force=True) - self.debug.log(" ╚══════════════════════════════════════════════════════════╝", category="none", force=True) - self.debug.log("", category="none", force=True) + debug.start_timer("total_execution", force=True) - self.debug.start_timer("total_execution", force=True) - - self.debug.log("━━━━━━━━━ Model Preparation ━━━━━━━━━", category="none") + debug.log("━━━━━━━━━ Model Preparation ━━━━━━━━━", category="none") # Initial memory state - self.debug.log_memory_state("Before model preparation", show_tensors=False, detailed_tensors=False) - self.debug.start_timer("model_preparation") + debug.log_memory_state("Before model preparation", show_tensors=False, detailed_tensors=False) + debug.start_timer("model_preparation") # Check if download succeeded - if not download_weight(dit_model=dit_model, vae_model=vae_model, debug=self.debug): + if not download_weight(dit_model=dit_model, vae_model=vae_model, debug=debug): raise RuntimeError( f"Failed to download required model files. " f"DiT model: {dit_model}, VAE model: {vae_model}. " @@ -261,7 +226,7 @@ class SeedVR2: # RGBA detected: ComfyUI inverts alpha when loading images # (ComfyUI mask convention: 1=hidden, 0=visible vs standard alpha: 1=opaque, 0=transparent) # Invert alpha in-place: multiply by -1 then add 1 (equivalent to 1.0 - alpha) - self.debug.log("RGBA input detected - inverting alpha channel to match mask convention", category="info") + debug.log("RGBA input detected - inverting alpha channel to match mask convention", category="info") pixels[..., 3].mul_(-1).add_(1) return self._internal_execute(pixels, dit_model, vae_model, seed, new_resolution, cfg_scale, @@ -272,45 +237,8 @@ class SeedVR2: cache_model_dit, cache_model_vae, dit_device, vae_device, dit_torch_compile_args, vae_torch_compile_args, block_swap_config) except Exception as e: - self.cleanup(cache_model_dit=cache_model_dit, cache_model_vae=cache_model_vae, debug=self.debug) + self.cleanup(cache_model_dit=cache_model_dit, cache_model_vae=cache_model_vae, debug=debug) raise e - - - def cleanup(self, cache_model_dit: bool = False, cache_model_vae: bool = False, - debug: Optional['Debug'] = None) -> None: - """ - Cleanup runner and free memory - - Args: - cache_model_dit: If True, keep DiT model in RAM for future runs - cache_model_vae: If True, keep VAE model in RAM for future runs - debug: Optional debug instance for logging cleanup operations - """ - # Clear progress bar if it exists - if hasattr(self, '_pbar') and self._pbar is not None: - self._pbar.update_absolute(0, 100) - self._pbar = None - - # Get debug from runner if not provided - if debug is None and self.runner and hasattr(self.runner, 'debug'): - debug = self.runner.debug - - # Use complete_cleanup for all cleanup operations - if self.runner: - complete_cleanup(runner=self.runner, debug=debug, keep_dit_in_ram=cache_model_dit, keep_vae_in_ram=cache_model_vae) - - # Delete runner only if neither model is cached - if not (cache_model_dit or cache_model_vae): - del self.runner - self.runner = None - self._dit_model_name = "" - self._vae_model_name = "" - - # Clean up context text embeddings if they exist - if self.ctx: - cleanup_text_embeddings(self.ctx, debug) - self.ctx = None - def _internal_execute( self, @@ -392,7 +320,6 @@ class SeedVR2: This method manages the complete lifecycle including model loading, inference, and cleanup based on caching preferences. """ - debug = self.debug # Initialize progress bar @@ -453,11 +380,11 @@ class SeedVR2: # Transform outputs CTHW format, get dimensions from last two axes output_h, output_w = transformed_sample.shape[-2:] # Get actual output size - debug.log(f" Total frames: {total_frames}, Input: {input_w}x{input_h}px → Output: {output_w}x{output_h}px, Batch size: {batch_size}, Channels: {channels_info}", category="generation", force=True) + debug.log(f" Total frames: {total_frames}, Input: {input_w}x{input_h}px → Output: {output_w}x{output_h}px, Batch size: {batch_size}, Channels: {channels_info}", category="generation", force=True) # Store transform in context for reuse during encoding (avoids recreating it) ctx['video_transform'] = video_transform - + # Phase 1: Encode all batches ctx = encode_all_batches( self.runner, @@ -512,7 +439,7 @@ class SeedVR2: sample = sample.cpu() debug.log("", category="none", force=True) - debug.log("Video upscaling completed successfully!", category="generation", force=True) + debug.log("Video upscaling completed successfully!", category="success", force=True) debug.end_timer("generation", "Video generation") @@ -548,8 +475,13 @@ class SeedVR2: if total_execution_time > 0: fps = total_frames / total_execution_time debug.log(f"Average FPS: {fps:.2f} frames/sec", category="timing", force=True) - - debug.log("━━━━━━━━━━━━━━━━━━", category="none") + + # Footer with support info + debug.log("", category="none", force=True) + debug.log("━━━━━━━━━━━━━━━━━━", category="none", force=True) + debug.log("Questions? Updates? Watch the videos, star the repo & join us!", category="dialogue", force=True) + debug.log("https://www.youtube.com/@AInVFX", category="generation", force=True) + debug.log("https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler", category="star", force=True) # Clear history for next run (do this last, after all logging) debug.clear_history() @@ -561,31 +493,38 @@ class SeedVR2: return (sample,) def _progress_callback(self, current_step: int, total_steps: int, - current_batch_frames: int, phase_name: str = "") -> None: + current_frames: int, phase_name: str) -> None: """ - Progress callback for generation phases + Update progress bar based on pipeline phase Args: - current_step: Current step number within the phase - total_steps: Total steps in the current phase - current_batch_frames: Number of frames in current batch - phase_name: Name of the current phase (e.g., "Phase 1: Encoding") + current_step: Current step within phase + total_steps: Total steps in phase + current_frames: Number of frames being processed + phase_name: Name of current phase """ - - if not hasattr(self, '_pbar') or self._pbar is None: + if self._pbar is None: return - # Calculate overall progress across all phases - # Phase weights: Encode=20%, Upscale=60%, Decode=20% - phase_weights = {"Encoding": 0.2, "Upscaling": 0.6, "Decoding": 0.2} - phase_offset = {"Encoding": 0.0, "Upscaling": 0.2, "Decoding": 0.8} + # Define phase weights and offsets for overall progress + phase_weights = { + "Phase 1: Encoding": 0.25, + "Phase 2: Upscaling": 0.50, + "Phase 3: Decoding": 0.20, + "Phase 4: Post-processing": 0.05 + } - # Extract the phase type from "Phase X: Type" format - if ":" in phase_name: - phase_key = phase_name.split(":")[1].strip() - else: - phase_key = "Upscaling" + phase_offset = { + "Phase 1: Encoding": 0.0, + "Phase 2: Upscaling": 0.25, + "Phase 3: Decoding": 0.75, + "Phase 4: Post-processing": 0.95 + } + # Extract phase key from phase_name + phase_key = phase_name.split(" (")[0] if " (" in phase_name else phase_name + + # Get weight and offset weight = phase_weights.get(phase_key, 1.0) offset = phase_offset.get(phase_key, 0.0) @@ -597,6 +536,41 @@ class SeedVR2: progress_value = int(overall_progress * 100) self._pbar.update_absolute(progress_value, 100) + def cleanup(self, cache_model_dit: bool = False, cache_model_vae: bool = False, + debug: Optional[Debug] = None) -> None: + """ + Cleanup runner and free memory + + Args: + cache_model_dit: If True, keep DiT model in RAM for future runs + cache_model_vae: If True, keep VAE model in RAM for future runs + debug: Optional debug instance for logging cleanup operations + """ + # Clear progress bar if it exists + if hasattr(self, '_pbar') and self._pbar is not None: + self._pbar.update_absolute(0, 100) + self._pbar = None + + # Get debug from runner if not provided + if debug is None and self.runner and hasattr(self.runner, 'debug'): + debug = self.runner.debug + + # Use complete_cleanup for all cleanup operations + if self.runner: + complete_cleanup(runner=self.runner, debug=debug, keep_dit_in_ram=cache_model_dit, keep_vae_in_ram=cache_model_vae) + + # Delete runner only if neither model is cached + if not (cache_model_dit or cache_model_vae): + del self.runner + self.runner = None + self._dit_model_name = "" + self._vae_model_name = "" + + # Clean up context text embeddings if they exist + if self.ctx: + cleanup_text_embeddings(self.ctx, debug) + self.ctx = None + def __del__(self): """Destructor""" try: @@ -612,341 +586,4 @@ class SeedVR2: if hasattr(self, attr): delattr(self, attr) except: - pass - - -class SeedVR2LoadDiTModel: - """ - Configure DiT (Diffusion Transformer) model loader with memory optimization - - Provides configuration for: - - Model selection and device placement - - BlockSwap memory optimization for limited VRAM - - Model caching between runs - - Optional torch.compile integration - - Returns: - SEEDVR2_DIT configuration dictionary for main upscaler node - """ - - @classmethod - def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]: - devices = get_device_list() - return { - "required": { - "model": (get_available_dit_models(), { - "default": DEFAULT_DIT, - "tooltip": "DiT model for upscaling. Models will automatically download on first use. Additional models can be added to the ComfyUI models folder." - }), - "device": (devices, { - "default": devices[0], - "tooltip": "Device to use for DiT upscaling" - }), - }, - "optional": { - "blocks_to_swap": ("INT", { - "default": 0, - "min": 0, - "max": 36, - "step": 1, - "tooltip": "BlockSwap: Number of transformer blocks to offload (0=disabled, 16=balanced, 32=max savings for 3b model, 36=max savings for 7b model)" - }), - "swap_io_components": ("BOOLEAN", { - "default": False, - "tooltip": "Offload embeddings/IO layers to CPU for additional VRAM savings" - }), - "cache_in_ram": ("BOOLEAN", { - "default": False, - "tooltip": "Keep DiT model in RAM between runs for faster batch processing" - }), - "torch_compile_args": ("TORCH_COMPILE_ARGS", { - "tooltip": "Optional torch.compile optimization settings from SeedVR2 Torch Compile Settings node" - }), - } - } - - # Define return types for ComfyUI - RETURN_TYPES = ("SEEDVR2_DIT",) - FUNCTION = "create_config" - CATEGORY = "SEEDVR2" - DESCRIPTION = "Configure DiT model loading and memory optimization settings" - - def create_config(self, model: str, device: str, blocks_to_swap: int = 0, - swap_io_components: bool = False, cache_in_ram: bool = False, - torch_compile_args: Optional[Dict[str, Any]] = None) -> Tuple[Dict[str, Any]]: - """ - Create DiT model configuration dictionary - - Args: - model: DiT model filename (e.g., "seedvr2_ema_3b_fp16.safetensors") - device: Target device for DiT model (e.g., "cuda:0", "cpu") - blocks_to_swap: Number of transformer blocks to offload for BlockSwap (0=disabled) - swap_io_components: Whether to offload input/output layers to CPU - cache_in_ram: Whether to keep model in RAM between runs - torch_compile_args: Optional torch.compile configuration from settings node - - Returns: - Tuple containing configuration dictionary for SeedVR2 main node - """ - config = { - "model": model, - "device": device, - "blocks_to_swap": blocks_to_swap, - "swap_io_components": swap_io_components, - "cache_in_ram": cache_in_ram, - "torch_compile_args": torch_compile_args, - } - return (config,) - - -class SeedVR2LoadVAEModel: - """ - Configure VAE (Variational Autoencoder) model loader with tiling support - - Provides configuration for: - - Model selection and device placement - - Tiled encoding/decoding for VRAM reduction - - Tile size and overlap control - - Model caching between runs - - Optional torch.compile integration - - Returns: - SEEDVR2_VAE configuration dictionary for main upscaler node - """ - - @classmethod - def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]: - devices = get_device_list() - return { - "required": { - "model": (get_available_vae_models(), { - "default": DEFAULT_VAE, - "tooltip": "VAE model for encoding/decoding. Models will automatically download on first use. Additional models can be added to the ComfyUI models folder." - }), - "device": (devices, { - "default": devices[0], - "tooltip": "Device to use for VAE processing" - }), - }, - "optional": { - "encode_tiled": ("BOOLEAN", { - "default": False, - "tooltip": "Enable tiled encoding to reduce VRAM during encoding" - }), - "encode_tile_size": ("INT", { - "default": 512, - "min": 64, - "step": 32, - "tooltip": "Size of encoding tiles in pixels. Smaller = less VRAM but more seams/artifacts and slower. Larger = more VRAM but better quality and faster." - }), - "encode_tile_overlap": ("INT", { - "default": 64, - "min": 0, - "step": 32, - "tooltip": "Pixel overlap between encoding tiles to reduce visible seams. Higher = better blending but slower processing." - }), - "decode_tiled": ("BOOLEAN", { - "default": False, - "tooltip": "Enable tiled decoding to reduce VRAM during decoding" - }), - "decode_tile_size": ("INT", { - "default": 512, - "min": 64, - "step": 32, - "tooltip": "Size of decoding tiles in pixels. Smaller = less VRAM but more seams/artifacts and slower. Larger = more VRAM but better quality and faster." - }), - "decode_tile_overlap": ("INT", { - "default": 64, - "min": 0, - "step": 32, - "tooltip": "Pixel overlap between decoding tiles to reduce visible seams. Higher = better blending but slower processing." - }), - "cache_in_ram": ("BOOLEAN", { - "default": False, - "tooltip": "Keep VAE model in RAM between runs for faster batch processing" - }), - "torch_compile_args": ("TORCH_COMPILE_ARGS", { - "tooltip": "Optional torch.compile optimization settings from SeedVR2 Torch Compile Settings node" - }), - } - } - - # Define return types for ComfyUI - RETURN_TYPES = ("SEEDVR2_VAE",) - FUNCTION = "create_config" - CATEGORY = "SEEDVR2" - DESCRIPTION = "Configure VAE settings and tiling options for memory optimization" - - def create_config(self, model: str, device: str, encode_tiled: bool = False, - encode_tile_size: int = 512, encode_tile_overlap: int = 64, - decode_tiled: bool = False, decode_tile_size: int = 512, - decode_tile_overlap: int = 64, cache_in_ram: bool = False, - torch_compile_args: Optional[Dict[str, Any]] = None) -> Tuple[Dict[str, Any]]: - """ - Create VAE model configuration dictionary - - Args: - model: VAE model filename (e.g., "ema_vae_fp16.safetensors") - device: Target device for VAE model (e.g., "cuda:0", "cpu") - encode_tiled: Enable tiled encoding to reduce VRAM usage - encode_tile_size: Size of encoding tiles in pixels - encode_tile_overlap: Overlap between encoding tiles to reduce seams - decode_tiled: Enable tiled decoding to reduce VRAM usage - decode_tile_size: Size of decoding tiles in pixels - decode_tile_overlap: Overlap between decoding tiles to reduce seams - cache_in_ram: Whether to keep model in RAM between runs - torch_compile_args: Optional torch.compile configuration from settings node - - Returns: - Tuple containing configuration dictionary for SeedVR2 main node - """ - config = { - "model": model, - "device": device, - "encode_tiled": encode_tiled, - "encode_tile_size": encode_tile_size, - "encode_tile_overlap": encode_tile_overlap, - "decode_tiled": decode_tiled, - "decode_tile_size": decode_tile_size, - "decode_tile_overlap": decode_tile_overlap, - "cache_in_ram": cache_in_ram, - "torch_compile_args": torch_compile_args, - } - return (config,) - -class SeedVR2TorchCompileSettings: - """Configure torch.compile optimization for DiT and VAE models""" - -class SeedVR2TorchCompileSettings: - """Configure torch.compile optimization for DiT and VAE models""" - - @classmethod - def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]: - return { - "required": { - "backend": (["inductor", "cudagraphs"], { - "default": "inductor", - "tooltip": ( - "Compilation backend:\n" - "• inductor: Full optimization with Triton kernel generation and fusion\n" - "• cudagraphs: Lightweight, only wraps model with CUDA graphs, no kernel optimization" - ) - }), - "mode": (["default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"], { - "default": "default", - "tooltip": ( - "Optimization level (compilation time vs runtime speed):\n" - "• default: Fast compile, good speedup, lowest memory\n" - "• reduce-overhead: Fast compile, better speedup with CUDA graphs, +10-20% memory\n" - "• max-autotune: Slow compile, best speedup, highest memory\n" - "• max-autotune-no-cudagraphs: Like max-autotune but without CUDA graphs (more compatible)" - ) - }), - "fullgraph": ("BOOLEAN", { - "default": False, - "tooltip": ( - "Compile entire model as single graph:\n" - "• False: Allows graph breaks, compiles what it can, more compatible\n" - "• True: Enforces no graph breaks, errors if any found, maximum optimization but fragile" - ) - }), - "dynamic": ("BOOLEAN", { - "default": False, - "tooltip": ( - "Handle varying input shapes without recompilation:\n" - "• False: Specialized for exact shapes, recompiles if shape changes\n" - "• True: Creates dynamic kernels upfront, slower but handles shape variations" - ) - }), - "dynamo_cache_size_limit": ("INT", { - "default": 64, - "min": 0, - "max": 1024, - "step": 1, - "tooltip": ( - "Max cached compiled versions per function (prevents memory bloat):\n" - "Higher value = more variations cached = more memory used\n" - "Lower value = more recompilation if inputs vary\n" - "Default 64 is good for most cases. Increase if you see cache_size_limit warnings" - ) - }), - }, - "optional": { - "dynamo_recompile_limit": ("INT", { - "default": 128, - "min": 0, - "max": 1024, - "step": 1, - "tooltip": ( - "Max recompilation attempts before giving up (falls back to no compilation):\n" - "Safety limit to prevent infinite recompilation loops\n" - "Default 128 is sufficient. Only increase if you see recompile_limit warnings" - ) - }), - } - } - - # Define return types for ComfyUI - RETURN_TYPES = ("TORCH_COMPILE_ARGS",) - FUNCTION = "create_args" - CATEGORY = "SEEDVR2" - DESCRIPTION = ( - "Configure torch.compile optimization for DiT and VAE speedup. " - "Connect to DiT and/or VAE model loader. Requires PyTorch 2.0+ and Triton." - ) - - def create_args(self, backend: str, mode: str, fullgraph: bool, dynamic: bool, - dynamo_cache_size_limit: int, dynamo_recompile_limit: int = 128) -> Tuple[Dict[str, Any]]: - """ - Create torch.compile configuration for model optimization - - Args: - backend: Compilation backend ("inductor" or "cudagraphs") - mode: Optimization mode ("default", "reduce-overhead", "max-autotune", etc.) - fullgraph: Whether to compile entire model as single graph - dynamic: Whether to handle varying input shapes without recompilation - dynamo_cache_size_limit: Maximum cached compiled versions per function - dynamo_recompile_limit: Maximum recompilation attempts before fallback - - Returns: - Tuple containing torch.compile configuration dictionary - """ - compile_args = { - "backend": backend, - "mode": mode, - "fullgraph": fullgraph, - "dynamic": dynamic, - "dynamo_cache_size_limit": dynamo_cache_size_limit, - "dynamo_recompile_limit": dynamo_recompile_limit, - } - return (compile_args,) - -# ComfyUI Node Mappings -NODE_CLASS_MAPPINGS = { - "SeedVR2": SeedVR2, - "SeedVR2LoadDiTModel": SeedVR2LoadDiTModel, - "SeedVR2LoadVAEModel": SeedVR2LoadVAEModel, - "SeedVR2TorchCompileSettings": SeedVR2TorchCompileSettings, -} - -# Human-readable node display names -NODE_DISPLAY_NAME_MAPPINGS = { - "SeedVR2": "SeedVR2 Video Upscaler", - "SeedVR2LoadDiTModel": "SeedVR2 (Down)Load DiT Model", - "SeedVR2LoadVAEModel": "SeedVR2 (Down)Load VAE Model", - "SeedVR2TorchCompileSettings": "SeedVR2 Torch Compile Settings", -} - -# Export version and metadata -__version__ = "2.0.0-modular" -__author__ = "SeedVR2 Team" -__description__ = "High-quality video upscaling using advanced diffusion models" - -# Additional exports for introspection -__all__ = [ - 'SeedVR2', - 'SeedVR2LoadDiTModel', - 'SeedVR2LoadVAEModel', - 'NODE_CLASS_MAPPINGS', - 'NODE_DISPLAY_NAME_MAPPINGS' -] + pass \ No newline at end of file diff --git a/src/optimization/blockswap.py b/src/optimization/blockswap.py index b5484b2..65dd45a 100644 --- a/src/optimization/blockswap.py +++ b/src/optimization/blockswap.py @@ -79,7 +79,10 @@ def apply_block_swap_to_dit(runner, block_swap_config: Dict[str, Any], debug) -> return blocks_to_swap = block_swap_config.get("blocks_to_swap", 0) - if blocks_to_swap <= 0: + swap_io_components = block_swap_config.get("swap_io_components", False) + + # Early return only if both block swap and I/O swap are disabled + if blocks_to_swap <= 0 and not swap_io_components: return if debug is None: @@ -100,48 +103,50 @@ def apply_block_swap_to_dit(runner, block_swap_config: Dict[str, Any], debug) -> offload_device = "cpu" configs = [] - blocks_to_swap = block_swap_config.get("blocks_to_swap", 0) if blocks_to_swap > 0: - block_text = "block" if blocks_to_swap == 1 else "blocks" + block_text = "block" if blocks_to_swap <= 1 else "blocks" configs.append(f"{blocks_to_swap} {block_text}") - if block_swap_config.get("swap_io_components", False): + if swap_io_components: configs.append("I/O components") debug.log(f"BlockSwap configured: {', '.join(configs)}", category="blockswap", force=True) - debug.log("BlockSwap will swap blocks to GPU during inference", category="info", force=True) - - # Validate model structure + + # Validate model structure for block operations if not hasattr(model, "blocks"): debug.log("Model doesn't have 'blocks' attribute for BlockSwap", level="ERROR", category="blockswap") return total_blocks = len(model.blocks) debug.log(f"Model has {total_blocks} blocks total", category="blockswap") - blocks_to_swap = min(blocks_to_swap, total_blocks) - + # Configure model with blockswap attributes - model.blocks_to_swap = blocks_to_swap - 1 # Convert to 0-indexed + if blocks_to_swap > 0: + blocks_to_swap = min(blocks_to_swap, total_blocks) + model.blocks_to_swap = blocks_to_swap - 1 # Convert to 0-indexed + debug.log(f"Configuring: {blocks_to_swap}/{total_blocks} blocks for swapping", category="blockswap") + debug.log("BlockSwap will swap blocks to GPU during inference", category="info", force=True) + else: + # No block swapping, set to -1 so no blocks match the swap condition + model.blocks_to_swap = -1 + debug.log("Block swapping disabled (blocks_to_swap=0)", category="blockswap") + model.main_device = device model.offload_device = offload_device - debug.log(f"Configuring: {blocks_to_swap}/{total_blocks} blocks for swapping", category="blockswap") - # Configure I/O components - swap_io_components = block_swap_config.get("swap_io_components", False) io_components_offloaded = _configure_io_components(model, device, offload_device, swap_io_components, debug) - - # Configure block placement and memory tracking memory_stats = _configure_blocks(model, device, offload_device, debug) memory_stats['io_components'] = io_components_offloaded - # Log memory summary + # Log memory summary _log_memory_summary(memory_stats, offload_device, device, swap_io_components, debug) - # Wrap block forward methods for dynamic swapping - for b, block in enumerate(model.blocks): - if b <= model.blocks_to_swap: - _wrap_block_forward(block, b, model, debug) + # Wrap block forward methods for dynamic swapping (only if blocks_to_swap > 0) + if blocks_to_swap > 0: + for b, block in enumerate(model.blocks): + if b <= model.blocks_to_swap: + _wrap_block_forward(block, b, model, debug) # Patch RoPE modules for robust error handling _patch_rope_for_blockswap(model, debug) diff --git a/src/utils/debug.py b/src/utils/debug.py index 8fd4cb2..3732a4a 100644 --- a/src/utils/debug.py +++ b/src/utils/debug.py @@ -31,20 +31,20 @@ class Debug: # Icon mapping for different categories CATEGORY_ICONS = { "general": "🔄", # General operations/processing - "timing": "⚡", # Performance timing + "timing": "⚡", # Performance timing "memory": "📊", # Memory usage tracking "cache": "💾", # Cache operations "cleanup": "🧹", # Cleanup operations "setup": "🔧", # Configuration/setup "generation": "🎬", # Generation process - "dit": "🚀", # Model loading/operations + "dit": "🚀", # Model loading/operations "blockswap": "🔀", # BlockSwap operations "download": "📥", # Download operations "success": "✅", # Successful completion "warning": "⚠️", # Warnings "error": "❌", # Errors "info": "ℹ️", # Statistics/info - "tip" :"💡", # Tip/suggestion + "tip" :"💡", # Tip/suggestion "video": "📹", # Video/sequence info "reuse": "♻️", # Reusing/recycling "runner": "🏃", # Runner operations @@ -53,6 +53,8 @@ class Debug: "precision": "🎯", # Precision "device": "🖥️", # Device info "file": "📂", # File operations + "star": "⭐", # Star + "dialogue": "💬", # Dialogue "none" : "", } @@ -220,7 +222,7 @@ class Debug: grandchild_duration = self.timer_durations.get(grandchild, 0) if grandchild_duration >= 0.01: # Only show if >= 10ms grandchild_message = self.timer_messages.get(grandchild, grandchild) - self.log(f" └─ {grandchild_message}: {grandchild_duration:.2f}s", category="timing", force=force) + self.log(f" └─ {grandchild_message}: {grandchild_duration:.2f}s", category="timing", force=force) if unaccounted > 0.01: # Show if more than 10ms unaccounted self.log(f" └─ (other operations): {unaccounted:.2f}s", category="timing", force=force) @@ -411,18 +413,18 @@ class Debug: # Show top 5 largest largest = sorted(details['gpu_tensors'], key=lambda x: x['size_mb'], reverse=True)[:5] for t in largest: - self.log(f" {t['shape']}: {t['size_mb']:.2f}MB, {t['dtype']}", category="memory", force=force) + self.log(f" {t['shape']}: {t['size_mb']:.2f}MB, {t['dtype']}", category="memory", force=force) # Large CPU tensors if details['large_cpu_tensors']: cpu_large_gb = sum(t['size_mb'] for t in details['large_cpu_tensors']) / 1024 self.log(f" Large CPU tensors (>10MB):", category="memory", force=force) - self.log(f" {len(details['large_cpu_tensors'])} using {cpu_large_gb:.2f}GB", category="memory", force=force) + self.log(f" {len(details['large_cpu_tensors'])} using {cpu_large_gb:.2f}GB", category="memory", force=force) # Show top 3 largest largest = sorted(details['large_cpu_tensors'], key=lambda x: x['size_mb'], reverse=True)[:3] for t in largest: - self.log(f" {t['shape']}: {t['size_mb']:.2f}MB, {t['dtype']}", category="memory", force=force) + self.log(f" {t['shape']}: {t['size_mb']:.2f}MB, {t['dtype']}", category="memory", force=force) # Common shape patterns if details['shape_patterns']: @@ -432,7 +434,7 @@ class Debug: self.log(" Common tensor shapes:", category="memory", force=force) for shape, count in common_shapes: if count > 1: - self.log(f" {shape}: {count} instances", category="memory", force=force) + self.log(f" {shape}: {count} instances", category="memory", force=force) # Module instances if details['module_types']: @@ -440,7 +442,7 @@ class Debug: if multi_instance: self.log(" Multiple module instances:", category="memory", force=force) for mtype, count in sorted(multi_instance, key=lambda x: x[1], reverse=True)[:5]: - self.log(f" {mtype}: {count} instances", category="memory", force=force) + self.log(f" {mtype}: {count} instances", category="memory", force=force) def _log_memory_diff(self, current_metrics: Dict[str, Any], force: bool = False) -> None: