| """
|
| SeedVR2 DiT Model Loader Node
|
| Configure DiT (Diffusion Transformer) model with memory optimization
|
| """
|
|
|
| from comfy_api.latest import io
|
| from comfy_execution.utils import get_executing_context
|
| 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(io.ComfyNode):
|
| """
|
| 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 define_schema(cls) -> io.Schema:
|
| devices = get_device_list()
|
| dit_models = get_available_dit_models()
|
|
|
| return io.Schema(
|
| node_id="SeedVR2LoadDiTModel",
|
| display_name="SeedVR2 (Down)Load DiT Model",
|
| category="SEEDVR2",
|
| description=(
|
| "Load and configure SeedVR2 DiT (Diffusion Transformer) model for video upscaling. "
|
| "Supports BlockSwap memory optimization for low VRAM systems, model caching for batch processing, "
|
| "multi-GPU offloading, and torch.compile acceleration. \n\n"
|
| "Connect to Video Upscaler node."
|
| ),
|
| inputs=[
|
| io.Combo.Input("model",
|
| options=dit_models,
|
| default=DEFAULT_DIT,
|
| tooltip=(
|
| "DiT (Diffusion Transformer) model for video upscaling.\n"
|
| "Models automatically download on first use.\n"
|
| "Additional models can be added to the ComfyUI models folder."
|
| )
|
| ),
|
| io.Combo.Input("device",
|
| options=devices,
|
| default=devices[0],
|
| tooltip="GPU device for DiT model inference (upscaling phase)"
|
| ),
|
| io.Int.Input("blocks_to_swap",
|
| default=0,
|
| min=0,
|
| max=36,
|
| step=1,
|
| optional=True,
|
| tooltip=(
|
| "Number of transformer blocks to swap between devices for VRAM optimization.\n"
|
| "• 0: Disabled (default)\n"
|
| "• 3B model: 0-32 blocks\n"
|
| "• 7B model: 0-36 blocks\n"
|
| "\n"
|
| "Requires offload_device to be set and different from device.\n"
|
| "Not available on macOS (unified memory architecture)."
|
| )
|
| ),
|
| io.Boolean.Input("swap_io_components",
|
| default=False,
|
| optional=True,
|
| tooltip=(
|
| "Offload input/output embeddings and normalization layers to reduce VRAM.\n"
|
| "Requires offload_device to be set and different from device.\n"
|
| "Not available on macOS (unified memory architecture)."
|
| )
|
| ),
|
| io.Combo.Input("offload_device",
|
| options=get_device_list(include_none=True, include_cpu=True),
|
| default="none",
|
| optional=True,
|
| tooltip=(
|
| "Device to offload DiT model when not actively processing.\n"
|
| "• 'none': Keep model on inference device (default, fastest)\n"
|
| "• 'cpu': Offload to system RAM (reduces VRAM usage)\n"
|
| "• 'cuda:X': Offload to another GPU (good balance if available)\n"
|
| "\n"
|
| "Required for BlockSwap (blocks_to_swap or swap_io_components)."
|
| )
|
| ),
|
| io.Boolean.Input("cache_model",
|
| default=False,
|
| optional=True,
|
| tooltip=(
|
| "Keep DiT model loaded on offload_device between workflow runs.\n"
|
| "Useful for batch processing to avoid repeated loading.\n"
|
| "Requires offload_device to be set."
|
| )
|
| ),
|
| io.Combo.Input("attention_mode",
|
| options=["sdpa", "flash_attn_2", "flash_attn_3", "sageattn_2", "sageattn_3"],
|
| default="sdpa",
|
| optional=True,
|
| tooltip=(
|
| "Attention computation backend:\n"
|
| "• sdpa: PyTorch scaled_dot_product_attention (default, stable, always available)\n"
|
| "• flash_attn_2: Flash Attention 2 (Ampere+, requires flash-attn package)\n"
|
| "• flash_attn_3: Flash Attention 3 (Hopper+, requires flash-attn with FA3 support)\n"
|
| "• sageattn_2: SageAttention 2 (requires sageattention package)\n"
|
| "• sageattn_3: SageAttention 3 (Blackwell/RTX 50xx only, requires sageattn3 package)\n"
|
| "\n"
|
| "SDPA is recommended - stable and works everywhere.\n"
|
| "Flash Attention and SageAttention provide speedup through optimized CUDA kernels on compatible GPUs."
|
| )
|
| ),
|
| io.Custom("TORCH_COMPILE_ARGS").Input("torch_compile_args",
|
| optional=True,
|
| tooltip=(
|
| "Optional torch.compile optimization settings from SeedVR2 Torch Compile Settings node.\n"
|
| "Provides 20-40% speedup with compatible PyTorch 2.0+ and Triton installation."
|
| )
|
| ),
|
| ],
|
| outputs=[
|
| io.Custom("SEEDVR2_DIT").Output(
|
| tooltip="DiT model configuration containing model path, device settings, BlockSwap parameters, and compilation options. Connect to Video Upscaler node."
|
| )
|
| ]
|
| )
|
|
|
| @classmethod
|
| def execute(cls, model: str, device: str, offload_device: str = "none",
|
| cache_model: bool = False, blocks_to_swap: int = 0,
|
| swap_io_components: bool = False, attention_mode: str = "sdpa",
|
| torch_compile_args: Dict[str, Any] = None) -> io.NodeOutput:
|
| """
|
| Create DiT model configuration for SeedVR2 main node
|
|
|
| Args:
|
| model: Model filename to load
|
| device: Target device for model execution
|
| offload_device: Device to offload model to when not in use
|
| cache_model: Whether to keep model loaded between runs
|
| blocks_to_swap: Number of transformer blocks to swap (requires offload_device != device)
|
| swap_io_components: Whether to offload I/O components (requires offload_device != device)
|
| attention_mode: Attention computation backend ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
|
| torch_compile_args: Optional torch.compile configuration from settings node
|
|
|
| Returns:
|
| NodeOutput containing configuration dictionary for SeedVR2 main node
|
|
|
| Raises:
|
| ValueError: If cache_model is enabled but offload_device is not set
|
| """
|
|
|
| if cache_model and offload_device == "none":
|
| raise ValueError(
|
| "Model caching (cache_model=True) requires offload_device to be set. "
|
| f"Current: offload_device='{offload_device}'. "
|
| "Please set offload_device to specify where the cached DiT model should be stored "
|
| "(e.g., 'cpu' or another device). Set cache_model=False if you don't want to cache the model."
|
| )
|
|
|
| config = {
|
| "model": model,
|
| "device": device,
|
| "offload_device": offload_device,
|
| "cache_model": cache_model,
|
| "blocks_to_swap": blocks_to_swap,
|
| "swap_io_components": swap_io_components,
|
| "attention_mode": attention_mode,
|
| "torch_compile_args": torch_compile_args,
|
| "node_id": get_executing_context().node_id,
|
| }
|
|
|
| return io.NodeOutput(config) |
