prototype seedvr

Signed-off-by: Vladimir Mandic <mandic00@live.com>

.pylintrc

@@ -29,6 +29,7 @@ ignore-paths=/usr/lib/.*$,
              modules/postprocess/aurasr_arch.py,
              modules/prompt_parser_xhinker.py,
              modules/ras,
+             modules/seedvr,
              modules/rife,
              modules/schedulers,
              modules/taesd,

.ruff.toml

@@ -12,6 +12,7 @@ exclude = [
     "modules/pag",
     "modules/schedulers",
     "modules/teacache",
+    "modules/seedvr",
 
     "modules/control/proc",
     "modules/control/units",

CHANGELOG.md

@@ -42,12 +42,13 @@
   - **ROCm** for Windows  
     support for both official torch preview release of `torch-rocm` for windows and **TheRock** unoffical `torch-rocm` builds for windows  
     note that rocm for windows is still in preview and has limited gpu support, please check rocm docs for details  
-  - **DirectML** warn as end-of-life  
+  - **DirectML** warn as *end-of-life*  
     `torch-directml` received no updates in over 1 year and its currently superceded by `rocm` or `zluda`  
   - command line params `--use-zluda` and `--use-rocm` will attempt desired operation or fail if not possible  
     previously sdnext was performing a fallback to `torch-cpu` which is not desired  
-  - if `--use-cuda` or `--use-rocm` are specified and `torch-cpu` is installed, installer will attempt to reinstall correct torch package  
-  - **installer**: warn if cuda or rocm are available and `torch-cpu` is installed  
+  - **installer** if `--use-cuda` or `--use-rocm` are specified and `torch-cpu` is installed, installer will attempt to reinstall correct torch package  
+  - **installer** warn if *cuda* or *rocm* are available and `torch-cpu` is installed  
+  - support for `torch==2.10-nightly` with `cuda==13.0`  
 - **Extensions**  
   - [Agent-Scheduler](https://github.com/SipherAGI/sd-webui-agent-scheduler)  
     was a high-value built-in extension, but it has not been maintained for 1.5 years  

installer.py

@@ -1353,6 +1353,7 @@ def set_environment():
         allocator += ',backend:cudaMallocAsync'
     if opts.get("torch_expandable_segments", False):
         allocator += ',expandable_segments:True'
+    os.environ.setdefault('PYTORCH_ALLOC_CONF', allocator)
     os.environ.setdefault('PYTORCH_CUDA_ALLOC_CONF', allocator)
     os.environ.setdefault('PYTORCH_HIP_ALLOC_CONF', allocator)
     log.debug(f'Torch allocator: "{allocator}"')

modules/postprocess/seedvr_model.py

@@ -0,0 +1,76 @@
+import time
+import numpy as np
+import torch
+from PIL import Image
+from torchvision.transforms import ToPILImage
+from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn
+from modules import devices
+from modules.shared import opts, log
+from modules.upscaler import Upscaler, UpscalerData
+
+
+MODELS_MAP = {
+    "SeedVR2 3B": "seedvr2_ema_3b_fp16.safetensors",
+    "SeedVR2 7B": "seedvr2_ema_7b_fp16.safetensors",
+    "SeedVR2 7B Sharp": "seedvr2_ema_7b_sharp_fp16.safetensors",
+}
+to_pil = ToPILImage()
+
+
+class UpscalerSeedVR(Upscaler):
+    def __init__(self, dirname=None):
+        self.name = "SeedVR"
+        super().__init__()
+        self.scalers = [
+            UpscalerData(name="SeedVR2 3B", path=None, upscaler=self, model=None, scale=1),
+            UpscalerData(name="SeedVR2 7B", path=None, upscaler=self, model=None, scale=1),
+            UpscalerData(name="SeedVR2 7B Sharp", path=None, upscaler=self, model=None, scale=1),
+        ]
+        self.model = None
+        self.model_loaded = None
+
+    def load_model(self, path: str):
+        model_name = MODELS_MAP.get(path, None)
+        if (self.model is None) or (self.model_loaded != model_name):
+            log.debug(f'Upscaler load: name="{self.name}" model="{model_name}"')
+            from modules.seedvr.src.core.model_manager import configure_runner
+            self.model = configure_runner(
+                model_name=model_name,
+                cache_dir=opts.hfcache_dir,
+                device=devices.device,
+                dtype=devices.dtype,
+            )
+
+    def do_upscale(self, img: Image.Image, selected_file):
+        devices.torch_gc()
+        self.load_model(selected_file)
+        if self.model is None:
+            return img
+
+        from modules.seedvr.src.core.generation import generation_loop
+
+        width = int(self.scale * img.width) // 8 * 8
+        image_tensor = np.array(img)
+        image_tensor = torch.from_numpy(image_tensor).to(device=devices.device, dtype=devices.dtype).unsqueeze(0) / 255.0
+
+        t0 = time.time()
+        result_tensor = generation_loop(
+            runner=self.model,
+            images=image_tensor,
+            cfg_scale=1.0,
+            seed=42,
+            res_w=width,
+            batch_size=1,
+            temporal_overlap=0,
+            device=devices.device,
+        )
+        t1 = time.time()
+        log.info(f'Upscaler: type="{self.name}" model="{selected_file}" scale={self.scale} time={t1 - t0:.2f}')
+        img = to_pil(result_tensor.squeeze().permute((2, 0, 1)))
+        devices.torch_gc()
+
+        if opts.upscaler_unload:
+            self.model = None
+            log.debug(f'Upscaler unload: type="{self.name}" model="{selected_file}"')
+            devices.torch_gc(force=True)
+        return img

modules/seedvr/config_3b.yaml

@@ -0,0 +1,97 @@
+__object__:
+  path: projects.video_diffusion_sr.train
+  name: VideoDiffusionTrainer
+
+dit:
+  model:
+    __object__:
+      path:
+        - "custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.dit_v2.nadit"
+        - "ComfyUI.custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.dit_v2.nadit"
+        - "src.models.dit_v2.nadit"
+      name: "NaDiT"
+      args: "as_params"
+    vid_in_channels: 33
+    vid_out_channels: 16
+    vid_dim: 2560
+    vid_out_norm: fusedrms
+    txt_in_dim: 5120
+    txt_in_norm: fusedln
+    txt_dim: ${.vid_dim}
+    emb_dim: ${eval:'6 * ${.vid_dim}'}
+    heads: 20
+    head_dim: 128 # llm-like
+    expand_ratio: 4
+    norm: fusedrms
+    norm_eps: 1.0e-05
+    ada: single
+    qk_bias: False
+    qk_norm: fusedrms
+    patch_size: [1, 2, 2]
+    num_layers: 32 # llm-like
+    mm_layers: 10
+    mlp_type: swiglu
+    msa_type: None
+    block_type: ${eval:'${.num_layers} * ["mmdit_sr"]'} # space-full
+    window: ${eval:'${.num_layers} * [(4,3,3)]'} # space-full
+    window_method: ${eval:'${.num_layers} // 2 * ["720pwin_by_size_bysize","720pswin_by_size_bysize"]'} # space-full
+    rope_type: mmrope3d
+    rope_dim: 128
+  compile: False
+  gradient_checkpoint: True
+  fsdp:
+    sharding_strategy: _HYBRID_SHARD_ZERO2
+
+ema:
+  decay: 0.9998
+
+vae:
+  model:
+    __object__:
+      path:
+        - "custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.video_vae_v3.modules.attn_video_vae"
+        - "ComfyUI.custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.video_vae_v3.modules.attn_video_vae"
+        - "src.models.video_vae_v3.modules.attn_video_vae"
+      name: "VideoAutoencoderKLWrapper"
+      args: "as_params"
+    freeze_encoder: False
+    gradient_checkpoint: True # Disabled to prevent VRAM leaks in inference
+  slicing:
+    split_size: 4
+    memory_device: same
+  memory_limit:
+    conv_max_mem: 0.5
+    norm_max_mem: 0.5
+  checkpoint: ema_vae_fp16.safetensors
+  scaling_factor: 0.9152
+  compile: False
+  grouping: False
+  dtype: float16
+
+diffusion:
+  schedule:
+    type: lerp
+    T: 1000.0
+  sampler:
+    type: euler
+    prediction_type: v_lerp
+  timesteps:
+    training:
+      type: logitnormal
+      loc: 0.0
+      scale: 1.0
+    sampling:
+      type: uniform_trailing
+      steps: 50
+    transform: True
+  loss:
+    type: v_lerp
+  cfg:
+    scale: 7.5
+    rescale: 0
+
+condition:
+  i2v: 0.0
+  v2v: 0.0
+  sr: 1.0
+  noise_scale: 0.25

modules/seedvr/config_7b.yaml

@@ -0,0 +1,94 @@
+__object__:
+  path: projects.video_diffusion_sr.train
+  name: VideoDiffusionTrainer
+
+dit:
+  model:
+    __object__:
+      path:
+        - "custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.dit.nadit"
+        - "ComfyUI.custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.dit.nadit"
+        - "src.models.dit.nadit"
+      name: "NaDiT"
+      args: "as_params"
+    vid_in_channels: 33
+    vid_out_channels: 16
+    vid_dim: 3072
+    txt_in_dim: 5120
+    txt_dim: ${.vid_dim}
+    emb_dim: ${eval:'6 * ${.vid_dim}'}
+    heads: 24
+    head_dim: 128 # llm-like
+    expand_ratio: 4
+    norm: fusedrms
+    norm_eps: 1e-5
+    ada: single
+    qk_bias: False
+    qk_rope: True
+    qk_norm: fusedrms
+    patch_size: [1, 2, 2]
+    num_layers: 36 # llm-like
+    shared_mlp: False
+    shared_qkv: False
+    mlp_type: normal
+    block_type: ${eval:'${.num_layers} * ["mmdit_sr"]'} # space-full
+    window: ${eval:'${.num_layers} * [(4,3,3)]'} # space-full
+    window_method: ${eval:'${.num_layers} // 2 * ["720pwin_by_size_bysize","720pswin_by_size_bysize"]'} # space-full
+  compile: False
+  gradient_checkpoint: True
+  fsdp:
+    sharding_strategy: _HYBRID_SHARD_ZERO2
+
+ema:
+  decay: 0.9998
+
+vae:
+  model:
+    __object__:
+      path:
+        - "custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.video_vae_v3.modules.attn_video_vae"
+        - "ComfyUI.custom_nodes.ComfyUI-SeedVR2_VideoUpscaler.src.models.video_vae_v3.modules.attn_video_vae"
+        - "src.models.video_vae_v3.modules.attn_video_vae"
+      name: "VideoAutoencoderKLWrapper"
+      args: "as_params"
+    freeze_encoder: False
+    # gradient_checkpoint: True
+  slicing:
+    split_size: 4
+    memory_device: same
+  memory_limit:
+    conv_max_mem: 0.5
+    norm_max_mem: 0.5
+  checkpoint: ema_vae_fp16.safetensors
+  scaling_factor: 0.9152
+  compile: False
+  grouping: False
+  dtype: float16
+
+diffusion:
+  schedule:
+    type: lerp
+    T: 1000.0
+  sampler:
+    type: euler
+    prediction_type: v_lerp
+  timesteps:
+    training:
+      type: logitnormal
+      loc: 0.0
+      scale: 1.0
+    sampling:
+      type: uniform_trailing
+      steps: 50
+    transform: True
+  loss:
+    type: v_lerp
+  cfg:
+    scale: 7.5
+    rescale: 0
+
+condition:
+  i2v: 0.0
+  v2v: 0.0
+  sr: 1.0
+  noise_scale: 0.25

modules/seedvr/src/__init__.py

@@ -0,0 +1,37 @@
+"""
+SeedVR2 Video Upscaler - Modular Architecture
+Refactored from monolithic seedvr2.py for better maintainability
+
+Author: Refactored codebase
+Version: 2.0.0 - Modular
+
+Available Modules:
+- utils: Download and path utilities
+- optimization: Memory, performance, and compatibility optimizations  
+- core: Model management and generation pipeline (NEW)
+- processing: Video and tensor processing (coming next)
+- interfaces: ComfyUI integration
+"""
+'''
+# Track which modules are available for progressive migration
+MODULES_AVAILABLE = {
+    'downloads': True,          # ✅ Module 1 - Downloads and model management
+    'memory_manager': True,     # ✅ Module 2 - Memory optimization
+    'performance': True,        # ✅ Module 3 - Performance optimizations  
+    'compatibility': True,      # ✅ Module 4 - FP8/FP16 compatibility
+    'model_manager': True,      # ✅ Module 5 - Model configuration and loading
+    'generation': True,         # ✅ Module 6 - Generation loop and inference
+    'video_transforms': True,   # ✅ Module 7 - Video processing and transforms
+    'comfyui_node': True,       # ✅ Module 8 - ComfyUI node interface (COMPLETE!)
+    'infer': True,              # ✅ Module 9 - Infer
+}
+'''
+# Core imports (always available)
+import os
+import sys
+
+# Add current directory to path for fallback imports
+current_dir = os.path.dirname(os.path.abspath(__file__))
+parent_dir = os.path.dirname(current_dir)
+if parent_dir not in sys.path:
+    sys.path.insert(0, parent_dir)

modules/seedvr/src/common/cache.py

@@ -0,0 +1,33 @@
+from typing import Callable
+
+
+class Cache:
+    """Caching reusable args for faster inference"""
+
+    def __init__(self, disable=False, prefix="", cache=None):
+        self.cache = cache if cache is not None else {}
+        self.disable = disable
+        self.prefix = prefix
+
+    def __call__(self, key: str, fn: Callable):
+        if self.disable:
+            return fn()
+
+        key = self.prefix + key
+        try:
+            result = self.cache[key]
+        except KeyError:
+            result = fn()
+            self.cache[key] = result
+        return result
+
+    def namespace(self, namespace: str):
+        return Cache(
+            disable=self.disable,
+            prefix=self.prefix + namespace + ".",
+            cache=self.cache,
+        )
+
+    def get(self, key: str):
+        key = self.prefix + key
+        return self.cache[key]

modules/seedvr/src/common/config.py

@@ -0,0 +1,128 @@
+import importlib
+from typing import Any, Callable, List, Union
+from omegaconf import DictConfig, ListConfig, OmegaConf
+
+try:
+    OmegaConf.register_new_resolver("eval", eval)
+except Exception as e:
+    if "already registered" not in str(e):
+        raise
+
+
+
+def load_config(path: str, argv: List[str] = None) -> Union[DictConfig, ListConfig]:
+    """
+    Load a configuration. Will resolve inheritance.
+    """
+
+    config = OmegaConf.load(path)
+    if argv is not None:
+        config_argv = OmegaConf.from_dotlist(argv)
+        config = OmegaConf.merge(config, config_argv)
+    config = resolve_recursive(config, resolve_inheritance)
+    return config
+
+
+def resolve_recursive(
+    config: Any,
+    resolver: Callable[[Union[DictConfig, ListConfig]], Union[DictConfig, ListConfig]],
+) -> Any:
+    config = resolver(config)
+    if isinstance(config, DictConfig):
+        for k in config.keys():
+            v = config.get(k)
+            if isinstance(v, (DictConfig, ListConfig)):
+                config[k] = resolve_recursive(v, resolver)
+    if isinstance(config, ListConfig):
+        for i in range(len(config)):
+            v = config.get(i)
+            if isinstance(v, (DictConfig, ListConfig)):
+                config[i] = resolve_recursive(v, resolver)
+    return config
+
+
+def resolve_inheritance(config: Union[DictConfig, ListConfig]) -> Any:
+    """
+    Recursively resolve inheritance if the config contains:
+    __inherit__: path/to/parent.yaml or a ListConfig of such paths.
+    """
+    if isinstance(config, DictConfig):
+        inherit = config.pop("__inherit__", None)
+
+        if inherit:
+            inherit_list = inherit if isinstance(inherit, ListConfig) else [inherit]
+
+            parent_config = None
+            for parent_path in inherit_list:
+                assert isinstance(parent_path, str)
+                parent_config = (
+                    load_config(parent_path)
+                    if parent_config is None
+                    else OmegaConf.merge(parent_config, load_config(parent_path))
+                )
+
+            if len(config.keys()) > 0:
+                config = OmegaConf.merge(parent_config, config)
+            else:
+                config = parent_config
+    return config
+
+
+def import_item(path: Union[str, List[str]], name: str) -> Any:
+    """
+    Import a python item with fallback support.
+    
+    Args:
+        path: Single path string or list of paths to try (fallback order)
+        name: Class/function name to import
+        
+    Returns:
+        Imported object
+        
+    Example: 
+        import_item("path.to.file", "MyClass") -> MyClass
+        import_item(["path1.to.file", "path2.to.file"], "MyClass") -> MyClass (first working path)
+    """
+    if isinstance(path, str):
+        # Single path - original behavior
+        return getattr(importlib.import_module(path), name)
+
+    elif isinstance(path, (list, ListConfig)):
+        # Multiple paths - try each until one works
+        last_error = None
+        for single_path in path:
+            try:
+                return getattr(importlib.import_module(single_path), name)
+            except ImportError as e:
+                last_error = e
+                continue
+
+        # If we get here, none of the paths worked
+        raise ImportError(f"Could not import '{name}' from any of the paths: {path}. Last error: {last_error}")
+
+    else:
+        raise ValueError(f"Path must be string or list of strings, got: {type(path)}")
+
+
+def create_object(config: DictConfig) -> Any:
+    """
+    Create an object from config.
+    The config is expected to contains the following:
+    __object__:
+      path: path.to.module
+      name: MyClass
+      args: as_config | as_params (default to as_config)
+    """
+
+    item = import_item(
+        path=config.__object__.path,
+        name=config.__object__.name,
+    )
+    args = config.__object__.get("args", "as_config")
+    if args == "as_config":
+        return item(config)
+    if args == "as_params":
+        config = OmegaConf.to_object(config)
+        config.pop("__object__")
+        return item(**config)
+    raise NotImplementedError(f"Unknown args type: {args}")

modules/seedvr/src/common/decorators.py

@@ -0,0 +1,126 @@
+import functools
+import threading
+from typing import Callable
+import torch
+
+from .distributed import barrier_if_distributed, get_global_rank, get_local_rank
+from .logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def log_on_entry(func: Callable) -> Callable:
+    """
+    Functions with this decorator will log the function name at entry.
+    When using multiple decorators, this must be applied innermost to properly capture the name.
+    """
+
+    def log_on_entry_wrapper(*args, **kwargs):
+        logger.info(f"Entering {func.__name__}")
+        return func(*args, **kwargs)
+
+    return log_on_entry_wrapper
+
+
+def barrier_on_entry(func: Callable) -> Callable:
+    """
+    Functions with this decorator will start executing when all ranks are ready to enter.
+    """
+
+    def barrier_on_entry_wrapper(*args, **kwargs):
+        barrier_if_distributed()
+        return func(*args, **kwargs)
+
+    return barrier_on_entry_wrapper
+
+
+def _conditional_execute_wrapper_factory(execute: bool, func: Callable) -> Callable:
+    """
+    Helper function for local_rank_zero_only and global_rank_zero_only.
+    """
+
+    def conditional_execute_wrapper(*args, **kwargs):
+        # Only execute if needed.
+        result = func(*args, **kwargs) if execute else None
+        # All GPUs must wait.
+        barrier_if_distributed()
+        # Return results.
+        return result
+
+    return conditional_execute_wrapper
+
+
+def _asserted_wrapper_factory(condition: bool, func: Callable, err_msg: str = "") -> Callable:
+    """
+    Helper function for some functions with special constraints,
+    especially functions called by other global_rank_zero_only / local_rank_zero_only ones,
+    in case they are wrongly invoked in other scenarios.
+    """
+
+    def asserted_execute_wrapper(*args, **kwargs):
+        assert condition, err_msg
+        result = func(*args, **kwargs)
+        return result
+
+    return asserted_execute_wrapper
+
+
+def local_rank_zero_only(func: Callable) -> Callable:
+    """
+    Functions with this decorator will only execute on local rank zero.
+    """
+    return _conditional_execute_wrapper_factory(get_local_rank() == 0, func)
+
+
+def global_rank_zero_only(func: Callable) -> Callable:
+    """
+    Functions with this decorator will only execute on global rank zero.
+    """
+    return _conditional_execute_wrapper_factory(get_global_rank() == 0, func)
+
+
+def assert_only_global_rank_zero(func: Callable) -> Callable:
+    """
+    Functions with this decorator are only accessible to processes with global rank zero.
+    """
+    return _asserted_wrapper_factory(
+        get_global_rank() == 0, func, err_msg="Not accessible to processes with global_rank != 0"
+    )
+
+
+def assert_only_local_rank_zero(func: Callable) -> Callable:
+    """
+    Functions with this decorator are only accessible to processes with local rank zero.
+    """
+    return _asserted_wrapper_factory(
+        get_local_rank() == 0, func, err_msg="Not accessible to processes with local_rank != 0"
+    )
+
+
+def new_thread(func: Callable) -> Callable:
+    """
+    Functions with this decorator will run in a new thread.
+    The function will return the thread, which can be joined to wait for completion.
+    """
+
+    def new_thread_wrapper(*args, **kwargs):
+        thread = threading.Thread(target=func, args=args, kwargs=kwargs)
+        thread.start()
+        return thread
+
+    return new_thread_wrapper
+
+
+def log_runtime(func: Callable) -> Callable:
+    """
+    Functions with this decorator will logging the runtime.
+    """
+
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+        torch.distributed.barrier()
+        result = func(*args, **kwargs)
+        torch.distributed.barrier()
+        return result
+
+    return wrapped

modules/seedvr/src/common/diffusion/__init__.py

@@ -0,0 +1,56 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Diffusion package.
+"""
+
+from .config import (
+    create_sampler_from_config,
+    create_sampling_timesteps_from_config,
+    create_schedule_from_config,
+)
+from .samplers.base import Sampler
+from .samplers.euler import EulerSampler
+from .schedules.base import Schedule
+from .schedules.lerp import LinearInterpolationSchedule
+from .timesteps.base import SamplingTimesteps, Timesteps
+from .timesteps.sampling.trailing import UniformTrailingSamplingTimesteps
+from .types import PredictionType, SamplingDirection
+from .utils import classifier_free_guidance, classifier_free_guidance_dispatcher, expand_dims
+
+__all__ = [
+    # Configs
+    "create_sampler_from_config",
+    "create_sampling_timesteps_from_config",
+    "create_schedule_from_config",
+    # Schedules
+    "Schedule",
+    "DiscreteVariancePreservingSchedule",
+    "LinearInterpolationSchedule",
+    # Samplers
+    "Sampler",
+    "EulerSampler",
+    # Timesteps
+    "Timesteps",
+    "SamplingTimesteps",
+    # Types
+    "PredictionType",
+    "SamplingDirection",
+    "UniformTrailingSamplingTimesteps",
+    # Utils
+    "classifier_free_guidance",
+    "classifier_free_guidance_dispatcher",
+    "expand_dims",
+]

modules/seedvr/src/common/diffusion/config.py

@@ -0,0 +1,71 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Utility functions for creating schedules and samplers from config.
+"""
+
+import torch
+from omegaconf import DictConfig
+
+from .samplers.base import Sampler
+from .samplers.euler import EulerSampler
+from .schedules.base import Schedule
+from .schedules.lerp import LinearInterpolationSchedule
+from .timesteps.base import SamplingTimesteps
+from .timesteps.sampling.trailing import UniformTrailingSamplingTimesteps
+
+
+def create_schedule_from_config(
+    config: DictConfig,
+) -> Schedule:
+    """
+    Create a schedule from configuration.
+    """
+    if config.type == "lerp":
+        return LinearInterpolationSchedule(T=config.get("T", 1.0))
+
+    raise NotImplementedError
+
+
+def create_sampler_from_config(
+    config: DictConfig,
+    schedule: Schedule,
+    timesteps: SamplingTimesteps,
+) -> Sampler:
+    """
+    Create a sampler from configuration.
+    """
+    if config.type == "euler":
+        return EulerSampler(
+            schedule=schedule,
+            timesteps=timesteps,
+            prediction_type=config.prediction_type,
+        )
+    raise NotImplementedError
+
+
+def create_sampling_timesteps_from_config(
+    config: DictConfig,
+    schedule: Schedule,
+    device: torch.device,
+) -> SamplingTimesteps:
+    if config.type == "uniform_trailing":
+        return UniformTrailingSamplingTimesteps(
+            T=schedule.T,
+            steps=config.steps,
+            shift=config.get("shift", 1.0),
+            device=device,
+        )
+    raise NotImplementedError

modules/seedvr/src/common/diffusion/samplers/base.py

@@ -0,0 +1,108 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Sampler base class.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Callable
+import torch
+from tqdm import tqdm
+
+from ..schedules.base import Schedule
+from ..timesteps.base import SamplingTimesteps
+from ..types import PredictionType, SamplingDirection
+from ..utils import assert_schedule_timesteps_compatible
+
+
+@dataclass
+class SamplerModelArgs:
+    x_t: torch.Tensor
+    t: torch.Tensor
+    i: int
+
+
+class Sampler(ABC):
+    """
+    Samplers are ODE/SDE solvers.
+    """
+
+    def __init__(
+        self,
+        schedule: Schedule,
+        timesteps: SamplingTimesteps,
+        prediction_type: PredictionType,
+        return_endpoint: bool = True,
+    ):
+        assert_schedule_timesteps_compatible(
+            schedule=schedule,
+            timesteps=timesteps,
+        )
+        self.schedule = schedule
+        self.timesteps = timesteps
+        self.prediction_type = prediction_type
+        self.return_endpoint = return_endpoint
+
+    @abstractmethod
+    def sample(
+        self,
+        x: torch.Tensor,
+        f: Callable[[SamplerModelArgs], torch.Tensor],
+    ) -> torch.Tensor:
+        """
+        Generate a new sample given the the intial sample x and score function f.
+        """
+
+    def get_next_timestep(
+        self,
+        t: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Get the next sample timestep.
+        Support multiple different timesteps t in a batch.
+        If no more steps, return out of bound value -1 or T+1.
+        """
+        T = self.timesteps.T
+        steps = len(self.timesteps)
+        curr_idx = self.timesteps.index(t)
+        next_idx = curr_idx + 1
+        bound = -1 if self.timesteps.direction == SamplingDirection.backward else T + 1
+
+        s = self.timesteps[next_idx.clamp_max(steps - 1)]
+        s = s.where(next_idx < steps, bound)
+        return s
+
+    def get_endpoint(
+        self,
+        pred: torch.Tensor,
+        x_t: torch.Tensor,
+        t: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Get to the endpoint of the probability flow.
+        """
+        x_0, x_T = self.schedule.convert_from_pred(pred, self.prediction_type, x_t, t)
+        return x_0 if self.timesteps.direction == SamplingDirection.backward else x_T
+
+    def get_progress_bar(self):
+        """
+        Get progress bar for sampling.
+        """
+        return tqdm(
+            iterable=range(len(self.timesteps) - (0 if self.return_endpoint else 1)),
+            dynamic_ncols=True,
+            desc=self.__class__.__name__,
+        )

modules/seedvr/src/common/diffusion/samplers/euler.py

@@ -0,0 +1,107 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+
+"""
+Euler ODE solver.
+"""
+
+from typing import Callable
+import torch
+from einops import rearrange
+from torch.nn import functional as F
+
+#from ....models.dit_v2 import na
+
+from ..types import PredictionType
+from ..utils import expand_dims
+from .base import Sampler, SamplerModelArgs
+
+
+class EulerSampler(Sampler):
+    """
+    The Euler method is the simplest ODE solver.
+    <https://en.wikipedia.org/wiki/Euler_method>
+    """
+
+    def sample(
+        self,
+        x: torch.Tensor,
+        f: Callable[[SamplerModelArgs], torch.Tensor],
+    ) -> torch.Tensor:
+        timesteps = self.timesteps.timesteps
+        #progress = self.get_progress_bar()
+        i = 0
+
+        # Optimisations VRAM
+        original_dtype = x.dtype
+        device = x.device
+
+        for t, s in zip(timesteps[:-1], timesteps[1:]):
+            # Appel du modèle avec monitoring
+            pred = f(SamplerModelArgs(x, t, i))
+
+            # Étape suivante
+            x = self.step_to(pred, x, t, s)
+
+            # Nettoyer les tenseurs temporaires
+            del pred
+
+            i += 1
+            #progress.update()
+
+        if self.return_endpoint:
+            t = timesteps[-1]
+            pred = f(SamplerModelArgs(x, t, i))
+            x = self.get_endpoint(pred, x, t)
+            del pred
+            #progress.update()
+
+        # Restaurer le dtype original si nécessaire
+        if original_dtype != torch.float16:
+            x = x.to(original_dtype)
+
+        return x
+
+    def step(
+        self,
+        pred: torch.Tensor,
+        x_t: torch.Tensor,
+        t: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Step to the next timestep.
+        """
+        return self.step_to(pred, x_t, t, self.get_next_timestep(t))
+
+    def step_to(
+        self,
+        pred: torch.Tensor,
+        x_t: torch.Tensor,
+        t: torch.Tensor,
+        s: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Steps from x_t at timestep t to x_s at timestep s. Returns x_s.
+        """
+        t = expand_dims(t, x_t.ndim)
+        s = expand_dims(s, x_t.ndim)
+        T = self.schedule.T
+        # Step from x_t to x_s.
+        pred_x_0, pred_x_T = self.schedule.convert_from_pred(pred, self.prediction_type, x_t, t)
+        pred_x_s = self.schedule.forward(pred_x_0, pred_x_T, s.clamp(0, T))
+        # Clamp x_s to x_0 and x_T if s is out of bound.
+        pred_x_s = pred_x_s.where(s >= 0, pred_x_0)
+        pred_x_s = pred_x_s.where(s <= T, pred_x_T)
+        return pred_x_s

modules/seedvr/src/common/diffusion/schedules/base.py

@@ -0,0 +1,131 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Schedule base class.
+"""
+
+from abc import ABC, abstractmethod, abstractproperty
+from typing import Tuple, Union
+import torch
+
+from ..types import PredictionType
+from ..utils import expand_dims
+
+
+class Schedule(ABC):
+    """
+    Diffusion schedules are uniquely defined by T, A, B:
+
+        x_t = A(t) * x_0 + B(t) * x_T, where t in [0, T]
+
+    Schedules can be continuous or discrete.
+    """
+
+    @abstractproperty
+    def T(self) -> Union[int, float]:
+        """
+        Maximum timestep inclusive.
+        Schedule is continuous if float, discrete if int.
+        """
+
+    @abstractmethod
+    def A(self, t: torch.Tensor) -> torch.Tensor:
+        """
+        Interpolation coefficient A.
+        Returns tensor with the same shape as t.
+        """
+
+    @abstractmethod
+    def B(self, t: torch.Tensor) -> torch.Tensor:
+        """
+        Interpolation coefficient B.
+        Returns tensor with the same shape as t.
+        """
+
+    # ----------------------------------------------------
+
+    def snr(self, t: torch.Tensor) -> torch.Tensor:
+        """
+        Signal to noise ratio.
+        Returns tensor with the same shape as t.
+        """
+        return (self.A(t) ** 2) / (self.B(t) ** 2)
+
+    def isnr(self, snr: torch.Tensor) -> torch.Tensor:
+        """
+        Inverse signal to noise ratio.
+        Returns tensor with the same shape as snr.
+        Subclass may implement.
+        """
+        raise NotImplementedError
+
+    # ----------------------------------------------------
+
+    def is_continuous(self) -> bool:
+        """
+        Whether the schedule is continuous.
+        """
+        return isinstance(self.T, float)
+
+    def forward(self, x_0: torch.Tensor, x_T: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        """
+        Diffusion forward function.
+        """
+        t = expand_dims(t, x_0.ndim)
+        return self.A(t) * x_0 + self.B(t) * x_T
+
+    def convert_from_pred(
+        self, pred: torch.Tensor, pred_type: PredictionType, x_t: torch.Tensor, t: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Convert from prediction. Return predicted x_0 and x_T.
+        """
+        t = expand_dims(t, x_t.ndim)
+        A_t = self.A(t)
+        B_t = self.B(t)
+
+        if pred_type == PredictionType.x_T:
+            pred_x_T = pred
+            pred_x_0 = (x_t - B_t * pred_x_T) / A_t
+        elif pred_type == PredictionType.x_0:
+            pred_x_0 = pred
+            pred_x_T = (x_t - A_t * pred_x_0) / B_t
+        elif pred_type == PredictionType.v_cos:
+            pred_x_0 = A_t * x_t - B_t * pred
+            pred_x_T = A_t * pred + B_t * x_t
+        elif pred_type == PredictionType.v_lerp:
+            pred_x_0 = (x_t - B_t * pred) / (A_t + B_t)
+            pred_x_T = (x_t + A_t * pred) / (A_t + B_t)
+        else:
+            raise NotImplementedError
+
+        return pred_x_0, pred_x_T
+
+    def convert_to_pred(
+        self, x_0: torch.Tensor, x_T: torch.Tensor, t: torch.Tensor, pred_type: PredictionType
+    ) -> torch.FloatTensor:
+        """
+        Convert to prediction target given x_0 and x_T.
+        """
+        if pred_type == PredictionType.x_T:
+            return x_T
+        if pred_type == PredictionType.x_0:
+            return x_0
+        if pred_type == PredictionType.v_cos:
+            t = expand_dims(t, x_0.ndim)
+            return self.A(t) * x_T - self.B(t) * x_0
+        if pred_type == PredictionType.v_lerp:
+            return x_T - x_0
+        raise NotImplementedError

modules/seedvr/src/common/diffusion/schedules/lerp.py

@@ -0,0 +1,55 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Linear interpolation schedule (lerp).
+"""
+
+from typing import Union
+import torch
+
+from .base import Schedule
+
+
+class LinearInterpolationSchedule(Schedule):
+    """
+    Linear interpolation schedule (lerp) is proposed by flow matching and rectified flow.
+    It leads to straighter probability flow theoretically. It is also used by Stable Diffusion 3.
+    <https://arxiv.org/abs/2209.03003>
+    <https://arxiv.org/abs/2210.02747>
+
+        x_t = (1 - t) * x_0 + t * x_T
+
+    Can be either continuous or discrete.
+    """
+
+    def __init__(self, T: Union[int, float] = 1.0):
+        self._T = T
+
+    @property
+    def T(self) -> Union[int, float]:
+        return self._T
+
+    def A(self, t: torch.Tensor) -> torch.Tensor:
+        return 1 - (t / self.T)
+
+    def B(self, t: torch.Tensor) -> torch.Tensor:
+        return t / self.T
+
+    # ----------------------------------------------------
+
+    def isnr(self, snr: torch.Tensor) -> torch.Tensor:
+        t = self.T / (1 + snr**0.5)
+        t = t if self.is_continuous() else t.round().int()
+        return t

modules/seedvr/src/common/diffusion/timesteps/base.py

@@ -0,0 +1,72 @@
+from abc import ABC, abstractmethod
+from typing import Sequence, Union
+import torch
+
+from ..types import SamplingDirection
+
+
+class Timesteps(ABC):
+    """
+    Timesteps base class.
+    """
+
+    def __init__(self, T: Union[int, float]):
+        assert T > 0
+        self._T = T
+
+    @property
+    def T(self) -> Union[int, float]:
+        """
+        Maximum timestep inclusive.
+        int if discrete, float if continuous.
+        """
+        return self._T
+
+    def is_continuous(self) -> bool:
+        """
+        Whether the schedule is continuous.
+        """
+        return isinstance(self.T, float)
+
+
+class SamplingTimesteps(Timesteps):
+    """
+    Sampling timesteps.
+    It defines the discretization of sampling steps.
+    """
+
+    def __init__(
+        self,
+        T: Union[int, float],
+        timesteps: torch.Tensor,
+        direction: SamplingDirection,
+    ):
+        assert timesteps.ndim == 1
+        super().__init__(T)
+        self.timesteps = timesteps
+        self.direction = direction
+
+    def __len__(self) -> int:
+        """
+        Number of sampling steps.
+        """
+        return len(self.timesteps)
+
+    def __getitem__(self, idx: Union[int, torch.IntTensor]) -> torch.Tensor:
+        """
+        The timestep at the sampling step.
+        Returns a scalar tensor if idx is int,
+        or tensor of the same size if idx is a tensor.
+        """
+        return self.timesteps[idx]
+
+    def index(self, t: torch.Tensor) -> torch.Tensor:
+        """
+        Find index by t.
+        Return index of the same shape as t.
+        Index is -1 if t not found in timesteps.
+        """
+        i, j = t.reshape(-1, 1).eq(self.timesteps).nonzero(as_tuple=True)
+        idx = torch.full_like(t, fill_value=-1, dtype=torch.int)
+        idx.view(-1)[i] = j.int()
+        return idx

modules/seedvr/src/common/diffusion/timesteps/sampling/trailing.py

@@ -0,0 +1,49 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+import torch
+
+from ...types import SamplingDirection
+from ..base import SamplingTimesteps
+
+
+class UniformTrailingSamplingTimesteps(SamplingTimesteps):
+    """
+    Uniform trailing sampling timesteps.
+    Defined in (https://arxiv.org/abs/2305.08891)
+
+    Shift is proposed in SD3 for RF schedule.
+    Defined in (https://arxiv.org/pdf/2403.03206) eq.23
+    """
+
+    def __init__(
+        self,
+        T: int,
+        steps: int,
+        shift: float = 1.0,
+        device: torch.device = "cpu",
+    ):
+        # Create trailing timesteps.
+        timesteps = torch.arange(1.0, 0.0, -1.0 / steps, device=device)
+
+        # Shift timesteps.
+        timesteps = shift * timesteps / (1 + (shift - 1) * timesteps)
+
+        # Scale to T range.
+        if isinstance(T, float):
+            timesteps = timesteps * T
+        else:
+            timesteps = timesteps.mul(T + 1).sub(1).round().int()
+
+        super().__init__(T=T, timesteps=timesteps, direction=SamplingDirection.backward)

modules/seedvr/src/common/diffusion/types.py

@@ -0,0 +1,59 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Type definitions.
+"""
+
+from enum import Enum
+
+
+class PredictionType(str, Enum):
+    """
+    x_0:
+        Predict data sample.
+    x_T:
+        Predict noise sample.
+        Proposed by DDPM (https://arxiv.org/abs/2006.11239)
+        Proved problematic by zsnr paper (https://arxiv.org/abs/2305.08891)
+    v_cos:
+        Predict velocity dx/dt based on the cosine schedule (A_t * x_T - B_t * x_0).
+        Proposed by progressive distillation (https://arxiv.org/abs/2202.00512)
+    v_lerp:
+        Predict velocity dx/dt based on the lerp schedule (x_T - x_0).
+        Proposed by rectified flow (https://arxiv.org/abs/2209.03003)
+    """
+
+    x_0 = "x_0"
+    x_T = "x_T"
+    v_cos = "v_cos"
+    v_lerp = "v_lerp"
+
+
+class SamplingDirection(str, Enum):
+    """
+    backward: Sample from x_T to x_0 for data generation.
+    forward:  Sample from x_0 to x_T for noise inversion.
+    """
+
+    backward = "backward"
+    forward = "forward"
+
+    @staticmethod
+    def reverse(direction):
+        if direction == SamplingDirection.backward:
+            return SamplingDirection.forward
+        if direction == SamplingDirection.forward:
+            return SamplingDirection.backward
+        raise NotImplementedError

modules/seedvr/src/common/diffusion/utils.py

@@ -0,0 +1,84 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Utility functions.
+"""
+
+from typing import Callable
+import torch
+
+
+def expand_dims(tensor: torch.Tensor, ndim: int):
+    """
+    Expand tensor to target ndim. New dims are added to the right.
+    For example, if the tensor shape was (8,), target ndim is 4, return (8, 1, 1, 1).
+    """
+    shape = tensor.shape + (1,) * (ndim - tensor.ndim)
+    return tensor.reshape(shape)
+
+
+def assert_schedule_timesteps_compatible(schedule, timesteps):
+    """
+    Check if schedule and timesteps are compatible.
+    """
+    if schedule.T != timesteps.T:
+        raise ValueError("Schedule and timesteps must have the same T.")
+    if schedule.is_continuous() != timesteps.is_continuous():
+        raise ValueError("Schedule and timesteps must have the same continuity.")
+
+
+def classifier_free_guidance(
+    pos: torch.Tensor,
+    neg: torch.Tensor,
+    scale: float,
+    rescale: float = 0.0,
+):
+    """
+    Apply classifier-free guidance.
+    """
+    # Classifier-free guidance (https://arxiv.org/abs/2207.12598)
+    cfg = neg + scale * (pos - neg)
+
+    # Classifier-free guidance rescale (https://arxiv.org/pdf/2305.08891.pdf)
+    if rescale != 0.0:
+        pos_std = pos.std(dim=list(range(1, pos.ndim)), keepdim=True)
+        cfg_std = cfg.std(dim=list(range(1, cfg.ndim)), keepdim=True)
+        factor = pos_std / cfg_std
+        factor = rescale * factor + (1 - rescale)
+        cfg *= factor
+
+    return cfg
+
+
+def classifier_free_guidance_dispatcher(
+    pos: Callable,
+    neg: Callable,
+    scale: float,
+    rescale: float = 0.0,
+):
+    """
+    Optionally execute models depending on classifer-free guidance scale.
+    """
+    # If scale is 1, no need to execute neg model.
+    if scale == 1.0:
+        return pos()
+
+    # Otherwise, execute both pos nad neg models and apply cfg.
+    return classifier_free_guidance(
+        pos=pos(),
+        neg=neg(),
+        scale=scale,
+        rescale=rescale,
+    )

modules/seedvr/src/common/distributed/__init__.py

@@ -0,0 +1,35 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Distributed package.
+"""
+
+from .basic import (
+    barrier_if_distributed,
+    convert_to_ddp,
+    get_device,
+    get_global_rank,
+    get_local_rank,
+    get_world_size,
+)
+
+__all__ = [
+    "barrier_if_distributed",
+    "convert_to_ddp",
+    "get_device",
+    "get_global_rank",
+    "get_local_rank",
+    "get_world_size",
+]

modules/seedvr/src/common/distributed/advanced.py

@@ -0,0 +1,208 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Advanced distributed functions for sequence parallel.
+"""
+
+from typing import Optional, List
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
+from torch.distributed.fsdp import ShardingStrategy
+
+from .basic import get_global_rank, get_world_size
+
+
+_DATA_PARALLEL_GROUP = None
+_SEQUENCE_PARALLEL_GROUP = None
+_SEQUENCE_PARALLEL_CPU_GROUP = None
+_MODEL_SHARD_CPU_INTER_GROUP = None
+_MODEL_SHARD_CPU_INTRA_GROUP = None
+_MODEL_SHARD_INTER_GROUP = None
+_MODEL_SHARD_INTRA_GROUP = None
+_SEQUENCE_PARALLEL_GLOBAL_RANKS = None
+
+
+def get_data_parallel_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get data parallel process group.
+    """
+    return _DATA_PARALLEL_GROUP
+
+
+def get_sequence_parallel_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get sequence parallel process group.
+    """
+    return _SEQUENCE_PARALLEL_GROUP
+
+
+def get_sequence_parallel_cpu_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get sequence parallel CPU process group.
+    """
+    return _SEQUENCE_PARALLEL_CPU_GROUP
+
+
+def get_data_parallel_rank() -> int:
+    """
+    Get data parallel rank.
+    """
+    group = get_data_parallel_group()
+    return dist.get_rank(group) if group else get_global_rank()
+
+
+def get_data_parallel_world_size() -> int:
+    """
+    Get data parallel world size.
+    """
+    group = get_data_parallel_group()
+    return dist.get_world_size(group) if group else get_world_size()
+
+
+def get_sequence_parallel_rank() -> int:
+    """
+    Get sequence parallel rank.
+    """
+    group = get_sequence_parallel_group()
+    return dist.get_rank(group) if group else 0
+
+
+def get_sequence_parallel_world_size() -> int:
+    """
+    Get sequence parallel world size.
+    """
+    group = get_sequence_parallel_group()
+    return dist.get_world_size(group) if group else 1
+
+
+def get_model_shard_cpu_intra_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get the CPU intra process group of model sharding.
+    """
+    return _MODEL_SHARD_CPU_INTRA_GROUP
+
+
+def get_model_shard_cpu_inter_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get the CPU inter process group of model sharding.
+    """
+    return _MODEL_SHARD_CPU_INTER_GROUP
+
+
+def get_model_shard_intra_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get the GPU intra process group of model sharding.
+    """
+    return _MODEL_SHARD_INTRA_GROUP
+
+
+def get_model_shard_inter_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get the GPU inter process group of model sharding.
+    """
+    return _MODEL_SHARD_INTER_GROUP
+
+
+def init_sequence_parallel(sequence_parallel_size: int):
+    """
+    Initialize sequence parallel.
+    """
+    global _DATA_PARALLEL_GROUP
+    global _SEQUENCE_PARALLEL_GROUP
+    global _SEQUENCE_PARALLEL_CPU_GROUP
+    global _SEQUENCE_PARALLEL_GLOBAL_RANKS
+    assert dist.is_initialized()
+    world_size = dist.get_world_size()
+    rank = dist.get_rank()
+    data_parallel_size = world_size // sequence_parallel_size
+    for i in range(data_parallel_size):
+        start_rank = i * sequence_parallel_size
+        end_rank = (i + 1) * sequence_parallel_size
+        ranks = range(start_rank, end_rank)
+        group = dist.new_group(ranks)
+        cpu_group = dist.new_group(ranks, backend="gloo")
+        if rank in ranks:
+            _SEQUENCE_PARALLEL_GROUP = group
+            _SEQUENCE_PARALLEL_CPU_GROUP = cpu_group
+            _SEQUENCE_PARALLEL_GLOBAL_RANKS = list(ranks)
+
+
+def init_model_shard_group(
+    *,
+    sharding_strategy: ShardingStrategy,
+    device_mesh: Optional[DeviceMesh] = None,
+):
+    """
+    Initialize process group of model sharding.
+    """
+    global _MODEL_SHARD_INTER_GROUP
+    global _MODEL_SHARD_INTRA_GROUP
+    global _MODEL_SHARD_CPU_INTER_GROUP
+    global _MODEL_SHARD_CPU_INTRA_GROUP
+    assert dist.is_initialized()
+    world_size = dist.get_world_size()
+    if device_mesh is not None:
+        num_shards_per_group = device_mesh.shape[1]
+    elif sharding_strategy == ShardingStrategy.NO_SHARD:
+        num_shards_per_group = 1
+    elif sharding_strategy in [
+        ShardingStrategy.HYBRID_SHARD,
+        ShardingStrategy._HYBRID_SHARD_ZERO2,
+    ]:
+        num_shards_per_group = torch.cuda.device_count()
+    else:
+        num_shards_per_group = world_size
+    num_groups = world_size // num_shards_per_group
+    device_mesh = (num_groups, num_shards_per_group)
+
+    gpu_mesh_2d = init_device_mesh("cuda", device_mesh, mesh_dim_names=("inter", "intra"))
+    cpu_mesh_2d = init_device_mesh("cpu", device_mesh, mesh_dim_names=("inter", "intra"))
+
+    _MODEL_SHARD_INTER_GROUP = gpu_mesh_2d.get_group("inter")
+    _MODEL_SHARD_INTRA_GROUP = gpu_mesh_2d.get_group("intra")
+    _MODEL_SHARD_CPU_INTER_GROUP = cpu_mesh_2d.get_group("inter")
+    _MODEL_SHARD_CPU_INTRA_GROUP = cpu_mesh_2d.get_group("intra")
+
+def get_sequence_parallel_global_ranks() -> List[int]:
+    """
+    Get all global ranks of the sequence parallel process group
+    that the caller rank belongs to.
+    """
+    if _SEQUENCE_PARALLEL_GLOBAL_RANKS is None:
+        return [dist.get_rank()]
+    return _SEQUENCE_PARALLEL_GLOBAL_RANKS
+
+
+def get_next_sequence_parallel_rank() -> int:
+    """
+    Get the next global rank of the sequence parallel process group
+    that the caller rank belongs to.
+    """
+    sp_global_ranks = get_sequence_parallel_global_ranks()
+    sp_rank = get_sequence_parallel_rank()
+    sp_size = get_sequence_parallel_world_size()
+    return sp_global_ranks[(sp_rank + 1) % sp_size]
+
+
+def get_prev_sequence_parallel_rank() -> int:
+    """
+    Get the previous global rank of the sequence parallel process group
+    that the caller rank belongs to.
+    """
+    sp_global_ranks = get_sequence_parallel_global_ranks()
+    sp_rank = get_sequence_parallel_rank()
+    sp_size = get_sequence_parallel_world_size()
+    return sp_global_ranks[(sp_rank + sp_size - 1) % sp_size]

modules/seedvr/src/common/distributed/basic.py

@@ -0,0 +1,67 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Distributed basic functions.
+"""
+
+import os
+import torch
+
+
+def get_global_rank() -> int:
+    """
+    Get the global rank, the global index of the GPU.
+    """
+    return int(os.environ.get("RANK", "0"))
+
+
+def get_local_rank() -> int:
+    """
+    Get the local rank, the local index of the GPU.
+    """
+    return int(os.environ.get("LOCAL_RANK", "0"))
+
+
+def get_world_size() -> int:
+    """
+    Get the world size, the total amount of GPUs.
+    """
+    return int(os.environ.get("WORLD_SIZE", "1"))
+
+
+def get_device() -> torch.device:
+    """
+    Get current rank device.
+    """
+    return torch.device("cuda", get_local_rank())
+
+
+def barrier_if_distributed(*args, **kwargs):
+    """
+    Synchronizes all processes if under distributed context.
+    """
+    import torch.distributed as dist
+    if dist.is_initialized():
+        return dist.barrier(*args, **kwargs)
+
+
+def convert_to_ddp(module: torch.nn.Module, **kwargs):
+    from torch.nn.parallel import DistributedDataParallel
+    return DistributedDataParallel(
+        module=module,
+        device_ids=[get_local_rank()],
+        output_device=get_local_rank(),
+        **kwargs,
+    )

modules/seedvr/src/common/distributed/meta_init_utils.py

@@ -0,0 +1,41 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+import torch
+from rotary_embedding_torch import RotaryEmbedding
+from torch import nn
+from torch.distributed.fsdp._common_utils import _is_fsdp_flattened
+
+__all__ = ["meta_non_persistent_buffer_init_fn"]
+
+
+def meta_non_persistent_buffer_init_fn(module: nn.Module) -> nn.Module:
+    """
+    Used for materializing `non-persistent tensor buffers` while model resuming.
+
+    Since non-persistent tensor buffers are not saved in state_dict,
+    when initializing model with meta device, user should materialize those buffers manually.
+
+    Currently, only `rope.dummy` is this special case.
+    """
+    with torch.no_grad():
+        for submodule in module.modules():
+            if not isinstance(submodule, RotaryEmbedding):
+                continue
+            for buffer_name, buffer in submodule.named_buffers(recurse=False):
+                if buffer.is_meta and "dummy" in buffer_name:
+                    materialized_buffer = torch.zeros_like(buffer, device="cpu")
+                    setattr(submodule, buffer_name, materialized_buffer)
+    assert not any(b.is_meta for n, b in module.named_buffers())
+    return module

modules/seedvr/src/common/distributed/ops.py

@@ -0,0 +1,494 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Distributed ops for supporting sequence parallel.
+"""
+
+from collections import defaultdict
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+import torch
+import torch.distributed as dist
+from torch import Tensor
+
+from ..cache import Cache
+from .advanced import (
+    get_sequence_parallel_group,
+    get_sequence_parallel_rank,
+    get_sequence_parallel_world_size,
+)
+
+from .basic import get_device
+
+_SEQ_DATA_BUF = defaultdict(lambda: [None, None, None])
+_SEQ_DATA_META_SHAPES = defaultdict()
+_SEQ_DATA_META_DTYPES = defaultdict()
+_SEQ_DATA_ASYNC_COMMS = defaultdict(list)
+_SYNC_BUFFER = defaultdict(dict)
+
+
+def single_all_to_all(
+    local_input: Tensor,
+    scatter_dim: int,
+    gather_dim: int,
+    group: dist.ProcessGroup,
+    async_op: bool = False,
+):
+    """
+    A function to do all-to-all on a tensor
+    """
+    seq_world_size = dist.get_world_size(group)
+    prev_scatter_dim = scatter_dim
+    if scatter_dim != 0:
+        local_input = local_input.transpose(0, scatter_dim)
+        if gather_dim == 0:
+            gather_dim = scatter_dim
+        scatter_dim = 0
+
+    inp_shape = list(local_input.shape)
+    inp_shape[scatter_dim] = inp_shape[scatter_dim] // seq_world_size
+    input_t = local_input.reshape(
+        [seq_world_size, inp_shape[scatter_dim]] + inp_shape[scatter_dim + 1 :]
+    ).contiguous()
+    output = torch.empty_like(input_t)
+    comm = dist.all_to_all_single(output, input_t, group=group, async_op=async_op)
+    if async_op:
+        # let user's code transpose & reshape
+        return output, comm, prev_scatter_dim
+
+    # first dim is seq_world_size, so we can split it directly
+    output = torch.cat(output.split(1), dim=gather_dim + 1).squeeze(0)
+    if prev_scatter_dim:
+        output = output.transpose(0, prev_scatter_dim).contiguous()
+    return output
+
+
+def _all_to_all(
+    local_input: Tensor,
+    scatter_dim: int,
+    gather_dim: int,
+    group: dist.ProcessGroup,
+):
+    seq_world_size = dist.get_world_size(group)
+    input_list = [
+        t.contiguous() for t in torch.tensor_split(local_input, seq_world_size, scatter_dim)
+    ]
+    output_list = [torch.empty_like(input_list[0]) for _ in range(seq_world_size)]
+    dist.all_to_all(output_list, input_list, group=group)
+    return torch.cat(output_list, dim=gather_dim).contiguous()
+
+
+class SeqAllToAll(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: Any,
+        group: dist.ProcessGroup,
+        local_input: Tensor,
+        scatter_dim: int,
+        gather_dim: int,
+        async_op: bool,
+    ) -> Tensor:
+        ctx.group = group
+        ctx.scatter_dim = scatter_dim
+        ctx.gather_dim = gather_dim
+        ctx.async_op = async_op
+        if async_op:
+            output, comm, prev_scatter_dim = single_all_to_all(
+                local_input, scatter_dim, gather_dim, group, async_op=async_op
+            )
+            ctx.prev_scatter_dim = prev_scatter_dim
+            return output, comm
+
+        return _all_to_all(local_input, scatter_dim, gather_dim, group)
+
+    @staticmethod
+    def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]:
+        if ctx.async_op:
+            input_t = torch.cat(grad_output[0].split(1), dim=ctx.gather_dim + 1).squeeze(0)
+            if ctx.prev_scatter_dim:
+                input_t = input_t.transpose(0, ctx.prev_scatter_dim)
+        else:
+            input_t = grad_output[0]
+        return (
+            None,
+            _all_to_all(input_t, ctx.gather_dim, ctx.scatter_dim, ctx.group),
+            None,
+            None,
+            None,
+        )
+
+
+class Slice(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx: Any, group: dist.ProcessGroup, local_input: Tensor, dim: int) -> Tensor:
+        ctx.group = group
+        ctx.rank = dist.get_rank(group)
+        seq_world_size = dist.get_world_size(group)
+        ctx.seq_world_size = seq_world_size
+        ctx.dim = dim
+        dim_size = local_input.shape[dim]
+        return local_input.split(dim_size // seq_world_size, dim=dim)[ctx.rank].contiguous()
+
+    @staticmethod
+    def backward(ctx: Any, grad_output: Tensor) -> Tuple[None, Tensor, None]:
+        dim_size = list(grad_output.size())
+        split_size = dim_size[0]
+        dim_size[0] = dim_size[0] * ctx.seq_world_size
+        output = torch.empty(dim_size, dtype=grad_output.dtype, device=torch.cuda.current_device())
+        dist._all_gather_base(output, grad_output, group=ctx.group)
+        return (None, torch.cat(output.split(split_size), dim=ctx.dim), None)
+
+
+class Gather(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: Any,
+        group: dist.ProcessGroup,
+        local_input: Tensor,
+        dim: int,
+        grad_scale: Optional[bool] = False,
+    ) -> Tensor:
+        ctx.group = group
+        ctx.rank = dist.get_rank(group)
+        ctx.dim = dim
+        ctx.grad_scale = grad_scale
+        seq_world_size = dist.get_world_size(group)
+        ctx.seq_world_size = seq_world_size
+        dim_size = list(local_input.size())
+        split_size = dim_size[0]
+        ctx.part_size = dim_size[dim]
+        dim_size[0] = dim_size[0] * seq_world_size
+        output = torch.empty(dim_size, dtype=local_input.dtype, device=torch.cuda.current_device())
+        dist._all_gather_base(output, local_input.contiguous(), group=ctx.group)
+        return torch.cat(output.split(split_size), dim=dim)
+
+    @staticmethod
+    def backward(ctx: Any, grad_output: Tensor) -> Tuple[None, Tensor]:
+        if ctx.grad_scale:
+            grad_output = grad_output * ctx.seq_world_size
+        return (
+            None,
+            grad_output.split(ctx.part_size, dim=ctx.dim)[ctx.rank].contiguous(),
+            None,
+            None,
+        )
+
+
+def gather_seq_scatter_heads_qkv(
+    qkv_tensor: Tensor,
+    *,
+    seq_dim: int,
+    qkv_shape: Optional[Tensor] = None,
+    cache: Cache = Cache(disable=True),
+    restore_shape: bool = True,
+):
+    """
+    A func to sync splited qkv tensor
+    qkv_tensor: the tensor we want to do alltoall with. The last dim must
+        be the projection_idx, which we will split into 3 part. After
+        spliting, the gather idx will be projecttion_idx + 1
+    seq_dim: gather_dim for all2all comm
+    restore_shape: if True, output will has the same shape length as input
+    """
+    group = get_sequence_parallel_group()
+    if not group:
+        return qkv_tensor
+    world = get_sequence_parallel_world_size()
+    orig_shape = qkv_tensor.shape
+    scatter_dim = qkv_tensor.dim()
+    bef_all2all_shape = list(orig_shape)
+    qkv_proj_dim = bef_all2all_shape[-1]
+    bef_all2all_shape = bef_all2all_shape[:-1] + [3, qkv_proj_dim // 3]
+    qkv_tensor = qkv_tensor.view(bef_all2all_shape)
+    qkv_tensor = SeqAllToAll.apply(group, qkv_tensor, scatter_dim, seq_dim, False)
+    if restore_shape:
+        out_shape = list(orig_shape)
+        out_shape[seq_dim] *= world
+        out_shape[-1] = qkv_proj_dim // world
+        qkv_tensor = qkv_tensor.view(out_shape)
+
+    # remove padding
+    if qkv_shape is not None:
+        unpad_dim_size = cache(
+            "unpad_dim_size", lambda: torch.sum(torch.prod(qkv_shape, dim=-1)).item()
+        )
+        if unpad_dim_size % world != 0:
+            padding_size = qkv_tensor.size(seq_dim) - unpad_dim_size
+            qkv_tensor = _unpad_tensor(qkv_tensor, seq_dim, padding_size)
+    return qkv_tensor
+
+
+def slice_inputs(x: Tensor, dim: int, padding: bool = True):
+    """
+    A func to slice the input sequence in sequence parallel
+    """
+    group = get_sequence_parallel_group()
+    if group is None:
+        return x
+    sp_rank = get_sequence_parallel_rank()
+    sp_world = get_sequence_parallel_world_size()
+    dim_size = x.shape[dim]
+    unit = (dim_size + sp_world - 1) // sp_world
+    if padding and dim_size % sp_world:
+        padding_size = sp_world - (dim_size % sp_world)
+        x = _pad_tensor(x, dim, padding_size)
+    slc = [slice(None)] * len(x.shape)
+    slc[dim] = slice(unit * sp_rank, unit * (sp_rank + 1))
+    return x[slc]
+
+
+def remove_seqeunce_parallel_padding(x: Tensor, dim: int, unpad_dim_size: int):
+    """
+    A func to remove the padding part of the tensor based on its original shape
+    """
+    group = get_sequence_parallel_group()
+    if group is None:
+        return x
+    sp_world = get_sequence_parallel_world_size()
+    if unpad_dim_size % sp_world == 0:
+        return x
+    padding_size = sp_world - (unpad_dim_size % sp_world)
+    assert (padding_size + unpad_dim_size) % sp_world == 0
+    return _unpad_tensor(x, dim=dim, padding_size=padding_size)
+
+
+def gather_heads_scatter_seq(x: Tensor, head_dim: int, seq_dim: int) -> Tensor:
+    """
+    A func to sync attention result with alltoall in sequence parallel
+    """
+    group = get_sequence_parallel_group()
+    if not group:
+        return x
+    dim_size = x.size(seq_dim)
+    sp_world = get_sequence_parallel_world_size()
+    if dim_size % sp_world != 0:
+        padding_size = sp_world - (dim_size % sp_world)
+        x = _pad_tensor(x, seq_dim, padding_size)
+    return SeqAllToAll.apply(group, x, seq_dim, head_dim, False)
+
+
+def gather_seq_scatter_heads(x: Tensor, seq_dim: int, head_dim: int) -> Tensor:
+    """
+    A func to sync embedding input with alltoall in sequence parallel
+    """
+    group = get_sequence_parallel_group()
+    if not group:
+        return x
+    return SeqAllToAll.apply(group, x, head_dim, seq_dim, False)
+
+
+def scatter_heads(x: Tensor, dim: int) -> Tensor:
+    """
+    A func to split heads before attention in sequence parallel
+    """
+    group = get_sequence_parallel_group()
+    if not group:
+        return x
+    return Slice.apply(group, x, dim)
+
+
+def gather_heads(x: Tensor, dim: int, grad_scale: Optional[bool] = False) -> Tensor:
+    """
+    A func to gather heads for the attention result in sequence parallel
+    """
+    group = get_sequence_parallel_group()
+    if not group:
+        return x
+    return Gather.apply(group, x, dim, grad_scale)
+
+
+def gather_outputs(
+    x: Tensor,
+    *,
+    gather_dim: int,
+    padding_dim: Optional[int] = None,
+    unpad_shape: Optional[Tensor] = None,
+    cache: Cache = Cache(disable=True),
+    scale_grad=True,
+):
+    """
+    A func to gather the outputs for the model result in sequence parallel
+    """
+    group = get_sequence_parallel_group()
+    if not group:
+        return x
+    x = Gather.apply(group, x, gather_dim, scale_grad)
+    if padding_dim is not None:
+        unpad_dim_size = cache(
+            "unpad_dim_size", lambda: torch.sum(torch.prod(unpad_shape, dim=1)).item()
+        )
+        x = remove_seqeunce_parallel_padding(x, padding_dim, unpad_dim_size)
+    return x
+
+
+def _pad_tensor(x: Tensor, dim: int, padding_size: int):
+    shape = list(x.shape)
+    shape[dim] = padding_size
+    pad = torch.zeros(shape, dtype=x.dtype, device=x.device)
+    return torch.cat([x, pad], dim=dim)
+
+
+def _unpad_tensor(x: Tensor, dim: int, padding_size):
+    slc = [slice(None)] * len(x.shape)
+    slc[dim] = slice(0, -padding_size)
+    return x[slc]
+
+
+def _broadcast_data(data, shape, dtype, src, group, async_op):
+    comms = []
+    if isinstance(data, (list, tuple)):
+        for i, sub_shape in enumerate(shape):
+            comms += _broadcast_data(data[i], sub_shape, dtype[i], src, group, async_op)
+    elif isinstance(data, dict):
+        for key, sub_data in data.items():
+            comms += _broadcast_data(sub_data, shape[key], dtype[key], src, group, async_op)
+    elif isinstance(data, Tensor):
+        comms.append(dist.broadcast(data, src=src, group=group, async_op=async_op))
+    return comms
+
+
+def _traverse(data: Any, op: Callable) -> Union[None, List, Dict, Any]:
+    if isinstance(data, (list, tuple)):
+        return [_traverse(sub_data, op) for sub_data in data]
+    elif isinstance(data, dict):
+        return {key: _traverse(sub_data, op) for key, sub_data in data.items()}
+    elif isinstance(data, Tensor):
+        return op(data)
+    else:
+        return None
+
+
+def _get_shapes(data):
+    return _traverse(data, op=lambda x: x.shape)
+
+
+def _get_dtypes(data):
+    return _traverse(data, op=lambda x: x.dtype)
+
+
+def _construct_broadcast_buffer(shapes, dtypes, device):
+    if isinstance(shapes, torch.Size):
+        return torch.empty(shapes, dtype=dtypes, device=device)
+
+    if isinstance(shapes, (list, tuple)):
+        buffer = []
+        for i, sub_shape in enumerate(shapes):
+            buffer.append(_construct_broadcast_buffer(sub_shape, dtypes[i], device))
+    elif isinstance(shapes, dict):
+        buffer = {}
+        for key, sub_shape in shapes.items():
+            buffer[key] = _construct_broadcast_buffer(sub_shape, dtypes[key], device)
+    else:
+        return None
+    return buffer
+
+
+class SPDistForward:
+    """A forward tool to sync different result across sp group
+
+    Args:
+        module: a function or module to process users input
+        sp_step: current training step to judge which rank to broadcast its result to all
+        name: a distinct str to save meta and async comm
+        comm_shape: if different ranks have different shape, mark this arg to True
+        device: the device for current rank, can be empty
+    """
+
+    def __init__(
+        self,
+        name: str,
+        comm_shape: bool,
+        device: torch.device = None,
+    ):
+        self.name = name
+        self.comm_shape = comm_shape
+        if device:
+            self.device = device
+        else:
+            self.device = get_device()
+
+    def __call__(self, inputs) -> Any:
+        group = get_sequence_parallel_group()
+        if not group:
+            yield inputs
+        else:
+            device = self.device
+            sp_world = get_sequence_parallel_world_size()
+            sp_rank = get_sequence_parallel_rank()
+            for local_step in range(sp_world):
+                src_rank = dist.get_global_rank(group, local_step)
+                is_src = sp_rank == local_step
+                local_shapes = []
+                local_dtypes = []
+                if local_step == 0:
+                    local_result = inputs
+                    _SEQ_DATA_BUF[self.name][-1] = local_result
+                    local_shapes = _get_shapes(local_result)
+                    local_dtypes = _get_dtypes(local_result)
+                    if self.comm_shape:
+                        group_shapes_lists = [None] * sp_world
+                        dist.all_gather_object(group_shapes_lists, local_shapes, group=group)
+                        _SEQ_DATA_META_SHAPES[self.name] = group_shapes_lists
+                    else:
+                        _SEQ_DATA_META_SHAPES[self.name] = [local_shapes] * sp_world
+                    _SEQ_DATA_META_DTYPES[self.name] = local_dtypes
+                shapes = _SEQ_DATA_META_SHAPES[self.name][local_step]
+                dtypes = _SEQ_DATA_META_DTYPES[self.name]
+                buf_id = local_step % 2
+                if local_step == 0:
+                    sync_data = (
+                        local_result
+                        if is_src
+                        else _construct_broadcast_buffer(shapes, dtypes, device)
+                    )
+                    _broadcast_data(sync_data, shapes, dtypes, src_rank, group, False)
+                    _SEQ_DATA_BUF[self.name][buf_id] = sync_data
+
+                # wait for async comm ops
+                if _SEQ_DATA_ASYNC_COMMS[self.name]:
+                    for comm in _SEQ_DATA_ASYNC_COMMS[self.name]:
+                        comm.wait()
+                # before return the sync result, do async broadcast for next batch
+                if local_step < sp_world - 1:
+                    next_buf_id = 1 - buf_id
+                    shapes = _SEQ_DATA_META_SHAPES[self.name][local_step + 1]
+                    src_rank = dist.get_global_rank(group, local_step + 1)
+                    is_src = sp_rank == local_step + 1
+                    next_sync_data = (
+                        _SEQ_DATA_BUF[self.name][-1]
+                        if is_src
+                        else _construct_broadcast_buffer(shapes, dtypes, device)
+                    )
+                    _SEQ_DATA_ASYNC_COMMS[self.name] = _broadcast_data(
+                        next_sync_data, shapes, dtypes, src_rank, group, True
+                    )
+                    _SEQ_DATA_BUF[self.name][next_buf_id] = next_sync_data
+                yield _SEQ_DATA_BUF[self.name][buf_id]
+
+
+sync_inputs = SPDistForward(name="bef_fwd", comm_shape=True)
+
+
+def sync_data(data, sp_idx, name="tmp"):
+    group = get_sequence_parallel_group()
+    if group is None:
+        return data
+    # if sp_idx in _SYNC_BUFFER[name]:
+    #     return _SYNC_BUFFER[name][sp_idx]
+    sp_rank = get_sequence_parallel_rank()
+    src_rank = dist.get_global_rank(group, sp_idx)
+    objects = [data] if sp_rank == sp_idx else [None]
+    dist.broadcast_object_list(objects, src=src_rank, group=group)
+    # _SYNC_BUFFER[name] = {sp_idx: objects[0]}
+    return objects[0]

modules/seedvr/src/common/half_precision_fixes.py

@@ -0,0 +1,61 @@
+import torch.nn.functional as F
+
+
+def safe_pad_operation(x, padding, mode='constant', value=0.0):
+    """Safe padding operation that handles Half precision only for problematic modes"""
+    # Modes qui nécessitent le fix Half precision
+    problematic_modes = ['replicate', 'reflect', 'circular']
+
+    if mode in problematic_modes:
+        try:
+            return F.pad(x, padding, mode=mode, value=value)
+        except RuntimeError as e:
+            if "not implemented for 'Half'" in str(e):
+                original_dtype = x.dtype
+                return F.pad(x.float(), padding, mode=mode, value=value).to(original_dtype)
+            else:
+                raise e
+    else:
+        # Pour 'constant' et autres modes compatibles, pas de fix nécessaire
+        return F.pad(x, padding, mode=mode, value=value)
+
+
+def safe_interpolate_operation(x, size=None, scale_factor=None, mode='nearest', align_corners=None, recompute_scale_factor=None):
+    """Safe interpolate operation that handles Half precision for problematic modes"""
+    # Modes qui peuvent causer des problèmes avec Half precision
+    problematic_modes = ['bilinear', 'bicubic', 'trilinear']
+
+    if mode in problematic_modes:
+        try:
+            return F.interpolate(
+                x,
+                size=size,
+                scale_factor=scale_factor,
+                mode=mode,
+                align_corners=align_corners,
+                recompute_scale_factor=recompute_scale_factor
+            )
+        except RuntimeError as e:
+            if ("not implemented for 'Half'" in str(e) or
+                "compute_indices_weights" in str(e)):
+                original_dtype = x.dtype
+                return F.interpolate(
+                    x.float(),
+                    size=size,
+                    scale_factor=scale_factor,
+                    mode=mode,
+                    align_corners=align_corners,
+                    recompute_scale_factor=recompute_scale_factor
+                ).to(original_dtype)
+            else:
+                raise e
+    else:
+        # Pour 'nearest' et autres modes compatibles, pas de fix nécessaire
+        return F.interpolate(
+            x,
+            size=size,
+            scale_factor=scale_factor,
+            mode=mode,
+            align_corners=align_corners,
+            recompute_scale_factor=recompute_scale_factor
+        )

modules/seedvr/src/common/logger.py

@@ -0,0 +1,44 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Logging utility functions.
+"""
+
+import logging
+import sys
+from typing import Optional
+
+from .distributed import get_global_rank, get_local_rank, get_world_size
+
+_default_handler = logging.StreamHandler(sys.stdout)
+_default_handler.setFormatter(
+    logging.Formatter(
+        "%(asctime)s "
+        + (f"[Rank:{get_global_rank()}]" if get_world_size() > 1 else "")
+        + (f"[LocalRank:{get_local_rank()}]" if get_world_size() > 1 else "")
+        + "[%(threadName).12s][%(name)s][%(levelname).5s] "
+        + "%(message)s"
+    )
+)
+
+
+def get_logger(name: Optional[str] = None) -> logging.Logger:
+    """
+    Get a logger.
+    """
+    logger = logging.getLogger(name)
+    logger.addHandler(_default_handler)
+    logger.setLevel(logging.INFO)
+    return logger

modules/seedvr/src/common/partition.py

@@ -0,0 +1,59 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Partition utility functions.
+"""
+
+from typing import Any, List
+
+
+def partition_by_size(data: List[Any], size: int) -> List[List[Any]]:
+    """
+    Partition a list by size.
+    When indivisible, the last group contains fewer items than the target size.
+
+    Examples:
+        - data: [1,2,3,4,5]
+        - size: 2
+        - return: [[1,2], [3,4], [5]]
+    """
+    assert size > 0
+    return [data[i : (i + size)] for i in range(0, len(data), size)]
+
+
+def partition_by_groups(data: List[Any], groups: int) -> List[List[Any]]:
+    """
+    Partition a list by groups.
+    When indivisible, some groups may have more items than others.
+
+    Examples:
+        - data: [1,2,3,4,5]
+        - groups: 2
+        - return: [[1,3,5], [2,4]]
+    """
+    assert groups > 0
+    return [data[i::groups] for i in range(groups)]
+
+
+def shift_list(data: List[Any], n: int) -> List[Any]:
+    """
+    Rotate a list by n elements.
+
+    Examples:
+        - data: [1,2,3,4,5]
+        - n: 3
+        - return: [4,5,1,2,3]
+    """
+    return data[(n % len(data)) :] + data[: (n % len(data))]

modules/seedvr/src/common/seed.py

@@ -0,0 +1,30 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+import random
+from typing import Optional
+import numpy as np
+import torch
+
+from .distributed import get_global_rank
+
+
+def set_seed(seed: Optional[int], same_across_ranks: bool = False):
+    """Function that sets the seed for pseudo-random number generators."""
+    if seed is not None:
+        seed += get_global_rank() if not same_across_ranks else 0
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+

modules/seedvr/src/core/__init__.py

@@ -0,0 +1,47 @@
+"""
+Core Module for SeedVR2
+
+Contains the main business logic and model management functionality:
+- Model configuration and loading
+- Architecture detection and memory estimation
+- Runner creation and management
+- Generation pipeline and logic
+"""
+'''
+from .model_manager import (
+    configure_runner,
+    load_quantized_state_dict,
+    configure_dit_model_inference,
+    configure_vae_model_inference,
+)
+
+from .generation import (
+    generation_step,
+    generation_loop,
+    cut_videos,
+    prepare_video_transforms,
+    load_text_embeddings,
+    calculate_optimal_batch_params
+)
+
+from .infer import VideoDiffusionInfer
+
+__all__ = [
+    # Model management
+    'configure_runner',
+    'load_quantized_state_dict', 
+    'configure_dit_model_inference',
+    'configure_vae_model_inference',
+    
+    # Generation logic
+    'generation_step',
+    'generation_loop', 
+    'cut_videos',
+    'prepare_video_transforms',
+    'load_text_embeddings',
+    'calculate_optimal_batch_params',
+
+    # Infer
+    'VideoDiffusionInfer'
+] 
+'''

modules/seedvr/src/core/generation.py

@@ -0,0 +1,365 @@
+import torch
+from torchvision.transforms import Compose, Lambda, Normalize
+
+from src.optimization.performance import optimized_video_rearrange, optimized_single_video_rearrange, optimized_sample_to_image_format
+from src.common.seed import set_seed
+from src.data.image.transforms.divisible_crop import DivisibleCrop
+from src.data.image.transforms.na_resize import NaResize
+from src.utils.color_fix import wavelet_reconstruction
+
+
+
+def generation_step(runner, text_embeds_dict, cond_latents, temporal_overlap, device):
+    """
+    Execute a single generation step with adaptive dtype handling
+    
+    Args:
+        runner: VideoDiffusionInfer instance
+        text_embeds_dict (dict): Text embeddings for positive and negative prompts
+        cond_latents (list): Conditional latents for generation
+        temporal_overlap (int): Number of frames for temporal overlap
+        
+    Returns:
+        tuple: (samples, last_latents) for potential temporal continuation
+        
+    Features:
+        - Adaptive dtype detection (FP8/FP16/BFloat16)
+        - Optimal autocast configuration for each model type
+        - Memory-efficient noise generation and reuse
+        - Automatic device placement with dtype preservation
+        - Advanced inference optimization
+    """
+    # Adaptive dtype detection for optimal performance
+    model_dtype = next(runner.dit.parameters()).dtype
+
+    # Configure dtypes according to model architecture
+    if model_dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+        # FP8 native: use BFloat16 for intermediate calculations (optimal compatibility)
+        dtype = torch.bfloat16
+    elif model_dtype == torch.float16:
+        dtype = torch.float16
+    else:
+        dtype = torch.bfloat16
+
+    def _move_to_cuda(x):
+        """Move tensors to CUDA with adaptive optimal dtype"""
+        return [i.to(device, dtype=dtype) for i in x]
+
+    # Memory optimization: Generate noise once and reuse to save VRAM
+    with torch.cuda.device(device):
+        base_noise = torch.randn_like(cond_latents[0], dtype=dtype)
+        noises = [base_noise]
+        aug_noises = [base_noise * 0.1 + torch.randn_like(base_noise) * 0.05]
+
+    # Move tensors with adaptive dtype (optimized for FP8/FP16/BFloat16)
+    noises, aug_noises, cond_latents = _move_to_cuda(noises), _move_to_cuda(aug_noises), _move_to_cuda(cond_latents)
+
+    cond_noise_scale = 0.0
+
+    def _add_noise(x, aug_noise):
+        # Use adaptive optimal dtype
+        t = (
+            torch.tensor([1000.0], device=device, dtype=dtype)
+            * cond_noise_scale
+        )
+        shape = torch.tensor(x.shape[1:], device=device)[None]
+        t = runner.timestep_transform(t, shape)
+        x = runner.schedule.forward(x, aug_noise, t)
+        return x
+
+    # Generate conditions with memory optimization
+    runner.dit.to(device=device)
+    condition = runner.get_condition(
+        noises[0],
+        task="sr",
+        latent_blur=_add_noise(cond_latents[0], aug_noises[0]),
+    )
+    conditions = [condition]
+
+    # Use adaptive autocast for optimal performance
+    with torch.no_grad():
+        video_tensors = runner.inference(
+            noises=noises,
+            conditions=conditions,
+            temporal_overlap=temporal_overlap,
+            **text_embeds_dict,
+        )
+
+    # Process samples with advanced optimization
+    samples = optimized_video_rearrange(video_tensors)
+    noises = noises[0].to("cpu")
+    aug_noises = aug_noises[0].to("cpu")
+    cond_latents = cond_latents[0].to("cpu")
+    conditions = conditions[0].to("cpu")
+    condition = condition.to("cpu")
+
+    return samples #, last_latents
+
+
+def cut_videos(videos):
+    t = videos.size(1)
+
+    if t % 4 == 1:
+        return videos
+
+    padding_needed = (4 - (t % 4)) % 4 + 1
+    last_frame = videos[:, -1:].expand(-1, padding_needed, -1, -1).contiguous()
+    result = torch.cat([videos, last_frame], dim=1)
+    return result
+
+
+def generation_loop(runner, images, cfg_scale=1.0, seed=666, res_w=720, batch_size=90, temporal_overlap=0, progress_callback=None, device:str='cpu'):
+    """
+    Main generation loop with context-aware temporal processing
+    
+    Args:
+        runner: VideoDiffusionInfer instance
+        images (torch.Tensor): Input images for upscaling
+        cfg_scale (float): Classifier-free guidance scale
+        seed (int): Random seed for reproducibility
+        res_w (int): Target resolution width
+        batch_size (int): Batch size for processing
+        temporal_overlap (int): Frames for temporal continuity
+        progress_callback (callable): Optional callback for progress reporting
+        
+    Returns:
+        torch.Tensor: Generated video frames
+        
+    Features:
+        - Context-aware generation with temporal overlap
+        - Adaptive dtype pipeline (FP8/FP16/BFloat16)
+        - Memory-optimized batch processing
+        - Advanced video transformation pipeline
+        - Intelligent VRAM management throughout process
+        - Real-time progress reporting
+    """
+
+    model_dtype = None
+    model_dtype = next(runner.dit.parameters()).dtype
+    compute_dtype = model_dtype
+
+    # Configure classifier-free guidance
+    runner.config.diffusion.cfg.scale = cfg_scale
+    runner.config.diffusion.cfg.rescale = 0.0
+    # Configure sampling steps
+    runner.config.diffusion.timesteps.sampling.steps = 1
+    runner.configure_diffusion()
+
+    # Set random seed
+    set_seed(seed)
+
+    # Advanced video transformation pipeline
+    video_transform = Compose([
+        NaResize(
+            resolution=(res_w),
+            mode="side",
+            downsample_only=False,
+        ),
+        Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),
+        DivisibleCrop((16, 16)),
+        Normalize(0.5, 0.5),
+        Lambda(lambda x: x.permute(1, 0, 2, 3)),  # t c h w -> c t h w (faster than Rearrange)
+    ])
+
+    # Initialize generation state
+    batch_samples = []
+
+    # Load text embeddings with adaptive dtype
+    text_embeds = {"texts_pos": [runner.text_pos_embeds], "texts_neg": [runner.text_neg_embeds]}
+
+    # Calculate processing parameters
+    step = batch_size - temporal_overlap
+    if step <= 0:
+        step = batch_size
+        temporal_overlap = 0
+
+    # Calculate total batches for progress reporting
+    total_batches = len(range(0, len(images), step))
+
+    # Main processing loop with context awareness
+    for batch_count, batch_idx in enumerate(range(0, len(images), step)):
+        # Calculate batch indices with overlap
+        if batch_idx == 0:
+            # First batch: no overlap
+            start_idx = 0
+            end_idx = min(batch_size, len(images))
+            effective_batch_size = end_idx - start_idx
+        else:
+            # Subsequent batches: temporal overlap
+            start_idx = batch_idx
+            end_idx = min(start_idx + batch_size, len(images))
+            effective_batch_size = end_idx - start_idx
+            if effective_batch_size <= temporal_overlap:
+                break  # Not enough new frames, stop
+
+        current_frames = end_idx - start_idx
+
+        # Process current batch
+        video = images[start_idx:end_idx]
+        # Use adaptive computation dtype
+        video = video.permute(0, 3, 1, 2).to(device, dtype=compute_dtype)
+
+        # Apply video transformations with memory optimization
+        transformed_video = video_transform(video)
+        del video
+        #video = video.to("cpu")
+        #del video
+        ori_lengths = [transformed_video.size(1)]
+
+        # Handle correct format: frames % 4 == 1
+        t = transformed_video.size(1)
+
+        if len(images) >= 5 and t % 4 != 1:
+            transformed_video = cut_videos(transformed_video)
+
+        # Context-aware temporal strategy
+        # First batch: standard complete diffusion
+        cond_latents = runner.vae_encode([transformed_video])
+
+        # Normal generation
+        samples = generation_step(runner, text_embeds, cond_latents=cond_latents, temporal_overlap=temporal_overlap, device=device)
+        #del cond_latents
+        del cond_latents
+
+        # Post-process samples
+        sample = samples[0]
+        del samples
+        #del samples
+        if ori_lengths[0] < sample.shape[0]:
+            sample = sample[:ori_lengths[0]]
+
+        # Apply color correction if available
+        transformed_video = transformed_video.to(device)
+
+        input_video = [optimized_single_video_rearrange(transformed_video)]
+        del transformed_video
+        sample = wavelet_reconstruction(sample, input_video[0][:sample.size(0)])
+        del input_video
+
+        # Convert to final image format
+        sample = optimized_sample_to_image_format(sample)
+        sample = sample.clip(-1, 1).mul_(0.5).add_(0.5)
+        sample_cpu = sample.to(torch.float16).to("cpu")
+        del sample
+        batch_samples.append(sample_cpu)
+        #del sample
+
+        # Aggressive cleanup after each batch
+        # Progress callback - batch start
+        if progress_callback:
+            progress_callback(batch_count+1, total_batches, current_frames, "Processing batch...")
+
+    runner.vae.to(device="cpu")
+    runner.dit.to(device="cpu")
+    # OPTIMISATION ULTIME : Pré-allocation et copie directe (évite les torch.cat multiples)
+
+    # 1. Calculer la taille totale finale
+    total_frames = sum(batch.shape[0] for batch in batch_samples)
+    if len(batch_samples) > 0:
+        sample_shape = batch_samples[0].shape
+        H, W, C = sample_shape[1], sample_shape[2], sample_shape[3]
+
+        # 2. Pré-allouer le tensor final directement sur CPU (évite concatenations)
+        final_video_images = torch.empty((total_frames, H, W, C), dtype=torch.float16)
+
+        # 3. Copier par blocs directement dans le tensor final
+        block_size = 500
+        current_idx = 0
+
+        for block_start in range(0, len(batch_samples), block_size):
+            block_end = min(block_start + block_size, len(batch_samples))
+
+            # Charger le bloc en VRAM
+            current_block = []
+            for i in range(block_start, block_end):
+                current_block.append(batch_samples[i].to(device))
+
+            # Concatener en VRAM (rapide)
+            block_result = torch.cat(current_block, dim=0)
+
+            # Convertir en Float16 sur GPU
+            #if block_result.dtype != torch.float16:
+            #    block_result = block_result.to(torch.float16)
+
+            # Copier directement dans le tensor final (pas de concatenation!)
+            block_frames = block_result.shape[0]
+            final_video_images[current_idx:current_idx + block_frames] = block_result.to("cpu")
+            current_idx += block_frames
+
+            # Nettoyage immédiat VRAM
+            del current_block, block_result
+    else:
+        print("SeedVR2: No batch_samples to process")
+        final_video_images = torch.empty((0, 0, 0, 0), dtype=torch.float16)
+
+    # Cleanup batch_samples
+    #del batch_samples
+    return final_video_images
+
+
+def prepare_video_transforms(res_w):
+    """
+    Prepare optimized video transformation pipeline
+    
+    Args:
+        res_w (int): Target resolution width
+        
+    Returns:
+        Compose: Configured transformation pipeline
+        
+    Features:
+        - Resolution-aware upscaling (no downsampling)
+        - Proper normalization for model compatibility
+        - Memory-efficient tensor operations
+    """
+    return Compose([
+        NaResize(
+            resolution=(res_w),
+            mode="side",
+            downsample_only=False,  # Model trained for high resolution
+        ),
+        Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),
+        DivisibleCrop((16, 16)),
+        Normalize(0.5, 0.5),
+        Lambda(lambda x: x.permute(1, 0, 2, 3)),  # t c h w -> c t h w
+    ])
+
+
+def calculate_optimal_batch_params(total_frames, batch_size, temporal_overlap):
+    """
+    Calculate optimal batch processing parameters
+    
+    Args:
+        total_frames (int): Total number of frames
+        batch_size (int): Desired batch size
+        temporal_overlap (int): Temporal overlap frames
+        
+    Returns:
+        dict: Optimized parameters and recommendations
+        
+    Features:
+        - 4n+1 constraint optimization
+        - Padding waste calculation
+        - Performance recommendations
+    """
+    step = batch_size - temporal_overlap
+    if step <= 0:
+        step = batch_size
+        temporal_overlap = 0
+
+    # Find optimal batch sizes (4n+1 constraint)
+    optimal_batches = [x for x in [i for i in range(1, 200) if i % 4 == 1] if x <= total_frames]
+    best_batch = max(optimal_batches) if optimal_batches else 1
+
+    # Calculate potential padding waste
+    padding_waste = 0
+    if batch_size not in optimal_batches:
+        padding_waste = sum(((i // 4) + 1) * 4 + 1 - i for i in range(batch_size, total_frames, batch_size))
+
+    return {
+        'step': step,
+        'temporal_overlap': temporal_overlap,
+        'best_batch': best_batch,
+        'padding_waste': padding_waste,
+        'is_optimal': batch_size in optimal_batches
+    }

modules/seedvr/src/core/infer.py

@@ -0,0 +1,357 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import List, Optional, Tuple, Union
+import torch
+from einops import rearrange
+from omegaconf import DictConfig, ListConfig
+from torch import Tensor
+from src.common.diffusion import (
+    classifier_free_guidance_dispatcher,
+    create_sampler_from_config,
+    create_sampling_timesteps_from_config,
+    create_schedule_from_config,
+)
+from src.common.distributed import (
+    get_device,
+)
+
+# from common.fs import download
+
+from src.models.dit_v2 import na
+
+
+def optimized_channels_to_last(tensor):
+    """🚀 Optimized replacement for rearrange(tensor, 'b c ... -> b ... c')
+    Moves channels from position 1 to last position using PyTorch native operations.
+    """
+    if tensor.ndim == 3:  # [batch, channels, spatial]
+        return tensor.permute(0, 2, 1)
+    elif tensor.ndim == 4:  # [batch, channels, height, width]
+        return tensor.permute(0, 2, 3, 1)
+    elif tensor.ndim == 5:  # [batch, channels, depth, height, width]
+        return tensor.permute(0, 2, 3, 4, 1)
+    else:
+        # Fallback for other dimensions - move channel (dim=1) to last
+        dims = list(range(tensor.ndim))
+        dims = [dims[0]] + dims[2:] + [dims[1]]  # [0, 2, 3, ..., 1]
+        return tensor.permute(*dims)
+
+def optimized_channels_to_second(tensor):
+    """🚀 Optimized replacement for rearrange(tensor, 'b ... c -> b c ...')
+    Moves channels from last position to position 1 using PyTorch native operations.
+    """
+    if tensor.ndim == 3:  # [batch, spatial, channels]
+        return tensor.permute(0, 2, 1)
+    elif tensor.ndim == 4:  # [batch, height, width, channels]
+        return tensor.permute(0, 3, 1, 2)
+    elif tensor.ndim == 5:  # [batch, depth, height, width, channels]
+        return tensor.permute(0, 4, 1, 2, 3)
+    else:
+        # Fallback for other dimensions - move last dim to position 1
+        dims = list(range(tensor.ndim))
+        dims = [dims[0], dims[-1]] + dims[1:-1]  # [0, -1, 1, 2, ..., -2]
+        return tensor.permute(*dims)
+
+
+class VideoDiffusionInfer():
+    def __init__(self, config: DictConfig, device: str, dtype: torch.dtype):
+        self.config = config
+        self.device = device
+        self.dtype = dtype
+        self.vae = None
+        self.dit = None
+        self.sampler = None
+        self.schedule = None
+    def get_condition(self, latent: Tensor, latent_blur: Tensor, task: str) -> Tensor:
+        t, h, w, c = latent.shape
+        cond = torch.zeros([t, h, w, c + 1], device=latent.device, dtype=latent.dtype)
+        if task == "t2v" or t == 1:
+            # t2i or t2v generation.
+            if task == "sr":
+                cond[:, ..., :-1] = latent_blur[:]
+                cond[:, ..., -1:] = 1.0
+            return cond
+        if task == "i2v":
+            # i2v generation.
+            cond[:1, ..., :-1] = latent[:1]
+            cond[:1, ..., -1:] = 1.0
+            return cond
+        if task == "v2v":
+            # v2v frame extension.
+            cond[:2, ..., :-1] = latent[:2]
+            cond[:2, ..., -1:] = 1.0
+            return cond
+        if task == "sr":
+            # sr generation.
+            cond[:, ..., :-1] = latent_blur[:]
+            cond[:, ..., -1:] = 1.0
+            return cond
+        raise NotImplementedError
+
+    def configure_diffusion(self):
+        self.schedule = create_schedule_from_config(
+            config=self.config.diffusion.schedule,
+        )
+        self.sampling_timesteps = create_sampling_timesteps_from_config( # pylint: disable=attribute-defined-outside-init
+            config=self.config.diffusion.timesteps.sampling,
+            schedule=self.schedule,
+            device=self.device,
+        )
+        self.sampler = create_sampler_from_config(
+            config=self.config.diffusion.sampler,
+            schedule=self.schedule,
+            timesteps=self.sampling_timesteps,
+        )
+
+    # -------------------------------- Helper ------------------------------- #
+
+    @torch.no_grad()
+    def vae_encode(self, samples: List[Tensor]) -> List[Tensor]:
+        self.dit.to(device="cpu")
+        self.vae.to(device=self.device)
+
+        use_sample = self.config.vae.get("use_sample", True)
+        latents = []
+        if len(samples) > 0:
+            device = get_device()
+            dtype = self.vae.dtype
+            scale = self.config.vae.scaling_factor
+            shift = self.config.vae.get("shifting_factor", 0.0)
+
+            if isinstance(scale, ListConfig):
+                scale = torch.tensor(scale, device=device, dtype=dtype)
+            if isinstance(shift, ListConfig):
+                shift = torch.tensor(shift, device=device, dtype=dtype)
+
+            # Group samples of the same shape to batches if enabled.
+            if self.config.vae.grouping:
+                batches, indices = na.pack(samples)
+            else:
+                batches = [sample.unsqueeze(0) for sample in samples]
+
+            # Vae process by each group.
+            for sample in batches:
+                sample = sample.to(device, dtype)
+                if hasattr(self.vae, "preprocess"):
+                    sample = self.vae.preprocess(sample)
+                if use_sample:
+                    latent = self.vae.encode(sample).latent
+                else:
+                    # Deterministic vae encode, only used for i2v inference (optionally)
+                    latent = self.vae.encode(sample).posterior.mode().squeeze(2)
+                latent = latent.unsqueeze(2) if latent.ndim == 4 else latent
+                latent = rearrange(latent, "b c ... -> b ... c")
+                #latent = optimized_channels_to_last(latent)
+                latent = (latent - shift) * scale
+                latents.append(latent)
+
+            # Ungroup back to individual latent with the original order.
+            if self.config.vae.grouping:
+                latents = na.unpack(latents, indices)
+            else:
+                latents = [latent.squeeze(0) for latent in latents]
+
+        self.vae.to(device="cpu")
+        return latents
+
+
+    @torch.no_grad()
+    def vae_decode(self, latents: List[Tensor], target_dtype: torch.dtype = None) -> List[Tensor]:
+        """🚀 VAE decode optimisé - décodage direct sans chunking, compatible avec autocast externe"""
+        self.dit.to(device="cpu")
+        self.vae.to(device=self.device)
+        samples = []
+        if len(latents) > 0:
+            device = get_device()
+            dtype = self.vae.dtype
+            scale = self.config.vae.scaling_factor
+            shift = self.config.vae.get("shifting_factor", 0.0)
+
+            if isinstance(scale, ListConfig):
+                scale = torch.tensor(scale, device=device, dtype=dtype)
+            if isinstance(shift, ListConfig):
+                shift = torch.tensor(shift, device=device, dtype=dtype)
+
+
+            # 🚀 OPTIMISATION 1: Group latents intelligemment pour batch processing
+            if self.config.vae.grouping:
+                latents, indices = na.pack(latents)
+            else:
+                latents = [latent.unsqueeze(0) for latent in latents]
+
+            # 🚀 OPTIMISATION 2: Traitement batch optimisé avec dtype adaptatif
+            for _i, latent in enumerate(latents):
+                # Préparation optimisée du latent
+                # Utiliser target_dtype si fourni (évite double autocast)
+                effective_dtype = target_dtype if target_dtype is not None else dtype
+                latent = latent.to(device, effective_dtype, non_blocking=True)
+                latent = latent / scale + shift
+                latent = rearrange(latent, "b ... c -> b c ...")
+                #latent = optimized_channels_to_second(latent)
+                latent = latent.squeeze(2)
+
+                # 🚀 OPTIMISATION 3: Décodage direct SANS autocast (utilise l'autocast externe)
+                sample = self.vae.decode(latent).sample
+                #sample = self.vae.decode(latent).sample
+                #sample = self.vae.decode(latent).sample
+
+                # 🚀 OPTIMISATION 4: Post-processing conditionnel
+                if hasattr(self.vae, "postprocess"):
+                    sample = self.vae.postprocess(sample)
+
+                samples.append(sample)
+
+            # Ungroup back to individual sample with the original order.
+            if self.config.vae.grouping:
+                samples = na.unpack(samples, indices)
+            else:
+                samples = [sample.squeeze(0) for sample in samples]
+        self.vae.to(device="cpu")
+        return samples
+
+    def timestep_transform(self, timesteps: Tensor, latents_shapes: Tensor):
+        # Skip if not needed.
+        if not self.config.diffusion.timesteps.get("transform", False):
+            return timesteps
+
+        # Compute resolution.
+        vt = self.config.vae.model.get("temporal_downsample_factor", 4)
+        vs = self.config.vae.model.get("spatial_downsample_factor", 8)
+        frames = (latents_shapes[:, 0] - 1) * vt + 1
+        heights = latents_shapes[:, 1] * vs
+        widths = latents_shapes[:, 2] * vs
+
+        # Compute shift factor.
+        def get_lin_function(x1, y1, x2, y2):
+            m = (y2 - y1) / (x2 - x1)
+            b = y1 - m * x1
+            return lambda x: m * x + b
+
+        img_shift_fn = get_lin_function(x1=256 * 256, y1=1.0, x2=1024 * 1024, y2=3.2)
+        vid_shift_fn = get_lin_function(x1=256 * 256 * 37, y1=1.0, x2=1280 * 720 * 145, y2=5.0)
+        shift = torch.where(
+            frames > 1,
+            vid_shift_fn(heights * widths * frames),
+            img_shift_fn(heights * widths),
+        )
+
+        # Shift timesteps.
+        timesteps = timesteps / self.schedule.T
+        timesteps = shift * timesteps / (1 + (shift - 1) * timesteps)
+        timesteps = timesteps * self.schedule.T
+        return timesteps
+
+    @torch.no_grad()
+    def inference(
+        self,
+        noises: List[Tensor],
+        conditions: List[Tensor],
+        texts_pos: Union[List[str], List[Tensor], List[Tuple[Tensor]]],
+        texts_neg: Union[List[str], List[Tensor], List[Tuple[Tensor]]],
+        cfg_scale: Optional[float] = None,
+        temporal_overlap: int = 0, # pylint: disable=unused-argument
+    ) -> List[Tensor]:
+        assert len(noises) == len(conditions) == len(texts_pos) == len(texts_neg)
+        batch_size = len(noises)
+
+        # Return if empty.
+        if batch_size == 0:
+            return []
+
+        # Set cfg scale
+        if cfg_scale is None:
+            cfg_scale = self.config.diffusion.cfg.scale
+
+        # 🚀 OPTIMISATION: Détecter le dtype du modèle pour performance optimale
+        model_dtype = next(self.dit.parameters()).dtype
+        # Adapter les dtypes selon le modèle
+        if model_dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+            target_dtype = torch.float16
+        elif model_dtype == torch.float16:
+            target_dtype = torch.float16
+        else:
+            target_dtype = torch.bfloat16
+        # Text embeddings.
+        assert type(texts_pos[0]) is type(texts_neg[0])
+        if isinstance(texts_pos[0], str):
+            text_pos_embeds, text_pos_shapes = self.text_encode(texts_pos) # pylint: disable=no-member
+            text_neg_embeds, text_neg_shapes = self.text_encode(texts_neg) # pylint: disable=no-member
+        elif isinstance(texts_pos[0], tuple):
+            text_pos_embeds, text_pos_shapes = [], []
+            text_neg_embeds, text_neg_shapes = [], []
+            for pos in zip(*texts_pos):
+                emb, shape = na.flatten(pos)
+                text_pos_embeds.append(emb)
+                text_pos_shapes.append(shape)
+            for neg in zip(*texts_neg):
+                emb, shape = na.flatten(neg)
+                text_neg_embeds.append(emb)
+                text_neg_shapes.append(shape)
+        else:
+            text_pos_embeds, text_pos_shapes = na.flatten(texts_pos)
+            text_neg_embeds, text_neg_shapes = na.flatten(texts_neg)
+
+        # Adapter les embeddings texte au dtype cible (compatible avec FP8)
+        if isinstance(text_pos_embeds, torch.Tensor):
+            text_pos_embeds = text_pos_embeds.to(target_dtype)
+        if isinstance(text_neg_embeds, torch.Tensor):
+            text_neg_embeds = text_neg_embeds.to(target_dtype)
+
+        # Flatten.
+        latents, latents_shapes = na.flatten(noises)
+        latents_cond, _ = na.flatten(conditions)
+
+        # Adapter les latents au dtype cible (compatible avec FP8)
+        latents = latents.to(target_dtype) if latents.dtype != target_dtype else latents
+        latents_cond = latents_cond.to(target_dtype) if latents_cond.dtype != target_dtype else latents_cond
+
+        latents = self.sampler.sample(
+            x=latents,
+            f=lambda args: classifier_free_guidance_dispatcher(
+                pos=lambda: self.dit(
+                    vid=torch.cat([args.x_t, latents_cond], dim=-1),
+                    txt=text_pos_embeds,
+                    vid_shape=latents_shapes,
+                    txt_shape=text_pos_shapes,
+                    timestep=args.t.repeat(batch_size),
+                ).vid_sample,
+                neg=lambda: self.dit(
+                    vid=torch.cat([args.x_t, latents_cond], dim=-1),
+                    txt=text_neg_embeds,
+                    vid_shape=latents_shapes,
+                    txt_shape=text_neg_shapes,
+                    timestep=args.t.repeat(batch_size),
+                ).vid_sample,
+                scale=(
+                    cfg_scale
+                    if (args.i + 1) / len(self.sampler.timesteps)
+                    <= self.config.diffusion.cfg.get("partial", 1)
+                    else 1.0
+                ),
+                rescale=self.config.diffusion.cfg.rescale,
+            ),
+        )
+
+        latents = na.unflatten(latents, latents_shapes)
+
+        # 🎯 Pré-calcul des dtypes (une seule fois)
+        vae_dtype = self.vae.dtype
+        decode_dtype = torch.float16 if (vae_dtype == torch.float16 or target_dtype == torch.float16) else vae_dtype
+        samples = self.vae_decode(latents, target_dtype=decode_dtype)
+
+        if samples and len(samples) > 0 and samples[0].dtype != torch.float16:
+            samples = [sample.to(torch.float16, non_blocking=True) for sample in samples]
+
+        return samples

modules/seedvr/src/core/model_manager.py

@@ -0,0 +1,61 @@
+import os
+import torch
+from omegaconf import OmegaConf
+from safetensors.torch import load_file as load_safetensors_file
+from huggingface_hub import hf_hub_download
+
+from src.optimization.memory_manager import preinitialize_rope_cache
+from src.common.config import load_config, create_object
+from src.core.infer import VideoDiffusionInfer
+
+
+def configure_runner(model_name, cache_dir, device:str='cpu', dtype:torch.dtype=None):
+    repo_id = "vladmandic/SeedVR2"
+    script_directory = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    config_path = os.path.join(script_directory, './config_7b.yaml') if "7b" in model_name else os.path.join(script_directory, './config_3b.yaml')
+    config = load_config(config_path)
+
+    vae_config_path = os.path.join(script_directory, 'src/models/video_vae_v3/s8_c16_t4_inflation_sd3.yaml')
+    vae_config = OmegaConf.load(vae_config_path)
+    vae_config.spatial_downsample_factor = vae_config.get('spatial_downsample_factor', 8)
+    vae_config.temporal_downsample_factor = vae_config.get('temporal_downsample_factor', 4)
+    config.vae.model = OmegaConf.merge(config.vae.model, vae_config)
+
+    runner = VideoDiffusionInfer(config, device=device, dtype=dtype)
+    OmegaConf.set_readonly(runner.config, False)
+
+    # load dit
+    with torch.device("meta"):
+        runner.dit = create_object(config.dit.model)
+        runner.dit.eval().to(dtype)
+        runner.dit.to_empty(device="cpu")
+
+    model_file = hf_hub_download(repo_id=repo_id, filename=model_name, cache_dir=cache_dir)
+    state_dict = load_safetensors_file(model_file)
+    runner.dit.load_state_dict(state_dict, assign=True)
+    del state_dict
+    runner.dit = runner.dit.to(device=device, dtype=dtype)
+
+    # load vae
+    vae_file = hf_hub_download(repo_id=repo_id, filename=config.vae.checkpoint, cache_dir=cache_dir)
+    with torch.device("meta"):
+        runner.vae = create_object(config.vae.model)
+        runner.vae.requires_grad_(False).eval()
+        runner.vae.to_empty(device="cpu")
+
+    state_dict = load_safetensors_file(vae_file)
+    runner.vae.load_state_dict(state_dict)
+    del state_dict
+    runner.vae = runner.vae.to(device=device, dtype=dtype)
+    runner.config.vae.dtype = str(dtype)
+    runner.vae.set_causal_slicing(**config.vae.slicing)
+    runner.vae.set_memory_limit(**runner.config.vae.memory_limit)
+
+    # load embeds
+    pos_embeds_file = hf_hub_download(repo_id=repo_id, filename='pos_emb.pt', cache_dir=cache_dir)
+    neg_embeds_file = hf_hub_download(repo_id=repo_id, filename='neg_emb.pt', cache_dir=cache_dir)
+    runner.text_pos_embeds = torch.load(pos_embeds_file).to(device=device, dtype=dtype)
+    runner.text_neg_embeds = torch.load(neg_embeds_file).to(device=device, dtype=dtype)
+
+    preinitialize_rope_cache(runner)
+    return runner

modules/seedvr/src/data/image/transforms/area_resize.py

@@ -0,0 +1,131 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+import math
+import random
+from typing import Union
+import torch
+from PIL import Image
+from torchvision.transforms import functional as TVF
+from torchvision.transforms.functional import InterpolationMode
+
+
+class AreaResize:
+    def __init__(
+        self,
+        max_area: float,
+        downsample_only: bool = False,
+        interpolation: InterpolationMode = InterpolationMode.BICUBIC,
+    ):
+        self.max_area = max_area
+        self.downsample_only = downsample_only
+        self.interpolation = interpolation
+
+    def __call__(self, image: Union[torch.Tensor, Image.Image]):
+
+        if isinstance(image, torch.Tensor):
+            height, width = image.shape[-2:]
+        elif isinstance(image, Image.Image):
+            width, height = image.size
+        else:
+            raise NotImplementedError
+
+        scale = math.sqrt(self.max_area / (height * width))
+
+        # keep original height and width for small pictures.
+        scale = 1 if scale >= 1 and self.downsample_only else scale
+
+        resized_height, resized_width = round(height * scale), round(width * scale)
+
+        return TVF.resize(
+            image,
+            size=(resized_height, resized_width),
+            interpolation=self.interpolation,
+        )
+
+
+class AreaRandomCrop:
+    def __init__(
+        self,
+        max_area: float,
+    ):
+        self.max_area = max_area
+
+    def get_params(self, input_size, output_size):
+        """Get parameters for ``crop`` for a random crop.
+
+        Args:
+            img (PIL Image): Image to be cropped.
+            output_size (tuple): Expected output size of the crop.
+
+        Returns:
+            tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
+        """
+        # w, h = _get_image_size(img)
+        h, w = input_size
+        th, tw = output_size
+        if w <= tw and h <= th:
+            return 0, 0, h, w
+
+        i = random.randint(0, h - th)
+        j = random.randint(0, w - tw)
+        return i, j, th, tw
+
+    def __call__(self, image: Union[torch.Tensor, Image.Image]):
+        if isinstance(image, torch.Tensor):
+            height, width = image.shape[-2:]
+        elif isinstance(image, Image.Image):
+            width, height = image.size
+        else:
+            raise NotImplementedError
+
+        resized_height = math.sqrt(self.max_area / (width / height))
+        resized_width = (width / height) * resized_height
+
+        resized_height, resized_width = round(resized_height), round(resized_width)
+        i, j, h, w = self.get_params((height, width), (resized_height, resized_width))
+        image = TVF.crop(image, i, j, h, w)
+        return image
+
+class ScaleResize:
+    def __init__(
+        self,
+        scale: float,
+    ):
+        self.scale = scale
+
+    def __call__(self, image: Union[torch.Tensor, Image.Image]):
+        if isinstance(image, torch.Tensor):
+            height, width = image.shape[-2:]
+            interpolation_mode = InterpolationMode.BILINEAR
+            antialias = True if image.ndim == 4 else "warn"
+        elif isinstance(image, Image.Image):
+            width, height = image.size
+            interpolation_mode = InterpolationMode.LANCZOS
+            antialias = "warn"
+        else:
+            raise NotImplementedError
+
+        scale = self.scale
+
+        # keep original height and width for small pictures
+
+        resized_height, resized_width = round(height * scale), round(width * scale)
+        image = TVF.resize(
+            image,
+            size=(resized_height, resized_width),
+            interpolation=interpolation_mode,
+            antialias=antialias,
+        )
+        return image

modules/seedvr/src/data/image/transforms/divisible_crop.py

@@ -0,0 +1,40 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Union
+import torch
+from PIL import Image
+from torchvision.transforms import functional as TVF
+
+
+class DivisibleCrop:
+    def __init__(self, factor):
+        if not isinstance(factor, tuple):
+            factor = (factor, factor)
+
+        self.height_factor, self.width_factor = factor[0], factor[1]
+
+    def __call__(self, image: Union[torch.Tensor, Image.Image]):
+        if isinstance(image, torch.Tensor):
+            height, width = image.shape[-2:]
+        elif isinstance(image, Image.Image):
+            width, height = image.size
+        else:
+            raise NotImplementedError
+
+        cropped_height = height - (height % self.height_factor)
+        cropped_width = width - (width % self.width_factor)
+
+        image = TVF.center_crop(img=image, output_size=(cropped_height, cropped_width))
+        return image

modules/seedvr/src/data/image/transforms/na_resize.py

@@ -0,0 +1,50 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Literal
+from torchvision.transforms import CenterCrop, Compose, InterpolationMode, Resize
+
+from .area_resize import AreaResize
+from .side_resize import SideResize
+
+
+def NaResize(
+    resolution: int,
+    mode: Literal["area", "side"],
+    downsample_only: bool,
+    interpolation: InterpolationMode = InterpolationMode.BICUBIC,
+):
+    if mode == "area":
+        return AreaResize(
+            max_area=resolution**2,
+            downsample_only=downsample_only,
+            interpolation=interpolation,
+        )
+    if mode == "side":
+        return SideResize(
+            size=resolution,
+            downsample_only=downsample_only,
+            interpolation=interpolation,
+        )
+    if mode == "square":
+        return Compose(
+            [
+                Resize(
+                    size=resolution,
+                    interpolation=interpolation,
+                ),
+                CenterCrop(resolution),
+            ]
+        )
+    raise ValueError(f"Unknown resize mode: {mode}")

modules/seedvr/src/data/image/transforms/side_resize.py

@@ -0,0 +1,54 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Union
+import torch
+from PIL import Image
+from torchvision.transforms import InterpolationMode
+from torchvision.transforms import functional as TVF
+
+
+class SideResize:
+    def __init__(
+        self,
+        size: int,
+        downsample_only: bool = False,
+        interpolation: InterpolationMode = InterpolationMode.BICUBIC,
+    ):
+        self.size = size
+        self.downsample_only = downsample_only
+        self.interpolation = interpolation
+
+    def __call__(self, image: Union[torch.Tensor, Image.Image]):
+        """
+        Args:
+            image (PIL Image or Tensor): Image to be scaled.
+
+        Returns:
+            PIL Image or Tensor: Rescaled image.
+        """
+        if isinstance(image, torch.Tensor):
+            height, width = image.shape[-2:]
+        elif isinstance(image, Image.Image):
+            width, height = image.size
+        else:
+            raise NotImplementedError
+
+        if self.downsample_only and min(width, height) < self.size:
+            # keep original height and width for small pictures.
+            size = min(width, height)
+        else:
+            size = self.size
+
+        return TVF.resize(image, size, self.interpolation)

modules/seedvr/src/models/dit/attention.py

@@ -0,0 +1,94 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+import torch
+import torch.nn.functional as F
+
+#from flash_attn import flash_attn_varlen_func
+
+from torch import nn
+
+
+def pytorch_varlen_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p=0.0, softmax_scale=None, causal=False, deterministic=False):
+    """
+    A PyTorch-based implementation of variable-length attention to replace flash_attn_varlen_func.
+    It processes each sequence in the batch individually.
+    """
+    # Create an empty tensor to store the output.
+    output = torch.empty_like(q)
+
+    # Iterate over each sequence in the batch. The batch size is the number of sequences.
+    for i in range(len(cu_seqlens_q) - 1):
+        # Determine the start and end indices for the current sequence.
+        start_q, end_q = cu_seqlens_q[i], cu_seqlens_q[i+1]
+        start_k, end_k = cu_seqlens_k[i], cu_seqlens_k[i+1]
+
+        # Slice the q, k, and v tensors to get the data for the current sequence.
+        # The shape is (seq_len, heads, head_dim).
+        q_i = q[start_q:end_q]
+        k_i = k[start_k:end_k]
+        v_i = v[start_k:end_k]
+
+        # Reshape for torch's scaled_dot_product_attention which expects (batch, heads, seq, dim).
+        # Here, we treat each sequence as a batch of 1.
+        q_i = q_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_q, head_dim)
+        k_i = k_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_k, head_dim)
+        v_i = v_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_k, head_dim)
+
+        # Use PyTorch's built-in scaled dot-product attention.
+        output_i = F.scaled_dot_product_attention(
+            q_i, k_i, v_i,
+            dropout_p=dropout_p if not deterministic else 0.0,
+            is_causal=causal
+        )
+
+        # Reshape the output back to the original format (seq_len, heads, head_dim)
+        output_i = output_i.squeeze(0).permute(1, 0, 2)
+
+        # Place the result for the current sequence into the main output tensor.
+        output[start_q:end_q] = output_i
+
+    return output
+
+
+class TorchAttention(nn.Module):
+    def tflops(self, args, kwargs, output) -> float:
+        assert len(args) == 0 or len(args) > 2, "query, key should both provided by args / kwargs"
+        q = kwargs.get("query") or args[0]
+        k = kwargs.get("key") or args[1]
+        b, h, sq, d = q.shape
+        b, h, sk, d = k.shape
+        return b * h * (4 * d * (sq / 1e6) * (sk / 1e6))
+
+    def forward(self, *args, **kwargs):
+        #return pytorch_varlen_attention(*args, **kwargs)
+        return F.scaled_dot_product_attention(*args, **kwargs)
+
+
+class FlashAttentionVarlen(nn.Module):
+    def tflops(self, args, kwargs, output) -> float:
+        cu_seqlens_q = kwargs["cu_seqlens_q"]
+        cu_seqlens_k = kwargs["cu_seqlens_k"]
+        _, h, d = output.shape
+        seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]) / 1e6
+        seqlens_k = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]) / 1e6
+        return h * (4 * d * (seqlens_q * seqlens_k).sum())
+
+    def forward(self, *args, **kwargs):
+        kwargs["deterministic"] = torch.are_deterministic_algorithms_enabled()
+        try:
+            from flash_attn import flash_attn_varlen_func
+            return flash_attn_varlen_func(*args, **kwargs)
+        except ImportError:
+            return pytorch_varlen_attention(*args, **kwargs)

modules/seedvr/src/models/dit/blocks/__init__.py

@@ -0,0 +1,25 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from .mmdit_window_block import MMWindowTransformerBlock
+
+dit_blocks = {
+    "mmdit_window": MMWindowTransformerBlock,
+}
+
+
+def get_block(block_type: str):
+    if block_type in dit_blocks:
+        return dit_blocks[block_type]
+    raise NotImplementedError(f"{block_type} is not supported")

modules/seedvr/src/models/dit/blocks/mmdit_window_block.py

@@ -0,0 +1,233 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Tuple, Union
+import torch
+from einops import rearrange
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.modules.utils import _triple
+from ....common.half_precision_fixes import safe_pad_operation
+from ....common.distributed.ops import (
+    gather_heads,
+    gather_heads_scatter_seq,
+    gather_seq_scatter_heads_qkv,
+    scatter_heads,
+)
+
+from ..attention import TorchAttention
+from ..mlp import get_mlp
+from ..mm import MMArg, MMModule
+from ..modulation import ada_layer_type
+from ..normalization import norm_layer_type
+from ..rope import RotaryEmbedding3d
+
+
+class MMWindowAttention(nn.Module):
+    def __init__(
+        self,
+        vid_dim: int,
+        txt_dim: int,
+        heads: int,
+        head_dim: int,
+        qk_bias: bool,
+        qk_rope: bool,
+        qk_norm: norm_layer_type,
+        qk_norm_eps: float,
+        window: Union[int, Tuple[int, int, int]],
+        window_method: str,
+        shared_qkv: bool,
+    ):
+        super().__init__()
+        dim = MMArg(vid_dim, txt_dim)
+        inner_dim = heads * head_dim
+        qkv_dim = inner_dim * 3
+
+        self.window = _triple(window)
+        self.window_method = window_method
+        assert all(map(lambda v: isinstance(v, int) and v >= 0, self.window))
+
+        self.head_dim = head_dim
+        self.proj_qkv = MMModule(nn.Linear, dim, qkv_dim, bias=qk_bias, shared_weights=shared_qkv)
+        self.proj_out = MMModule(nn.Linear, inner_dim, dim, shared_weights=shared_qkv)
+        self.norm_q = MMModule(qk_norm, dim=head_dim, eps=qk_norm_eps, elementwise_affine=True)
+        self.norm_k = MMModule(qk_norm, dim=head_dim, eps=qk_norm_eps, elementwise_affine=True)
+        self.rope = RotaryEmbedding3d(dim=head_dim // 2) if qk_rope else None
+        self.attn = TorchAttention()
+
+    def forward(
+        self,
+        vid: torch.FloatTensor,  # b T H W c
+        txt: torch.FloatTensor,  # b L c
+        txt_mask: torch.BoolTensor,  # b L
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.FloatTensor,
+    ]:
+        # Project q, k, v.
+        vid_qkv, txt_qkv = self.proj_qkv(vid, txt)
+        vid_qkv = gather_seq_scatter_heads_qkv(vid_qkv, seq_dim=2)
+        _, T, H, W, _ = vid_qkv.shape
+        _, L, _ = txt.shape
+
+        if self.window_method == "win":
+            nt, nh, nw = self.window
+            tt, hh, ww = T // nt, H // nh, W // nw
+        elif self.window_method == "win_by_size":
+            tt, hh, ww = self.window
+            tt, hh, ww = (
+                tt if tt > 0 else T,
+                hh if hh > 0 else H,
+                ww if ww > 0 else W,
+            )
+            nt, nh, nw = T // tt, H // hh, W // ww
+        else:
+            raise NotImplementedError
+
+        vid_qkv = rearrange(vid_qkv, "b T H W (o h d) -> o b h (T H W) d", o=3, d=self.head_dim)
+        txt_qkv = rearrange(txt_qkv, "b L (o h d) -> o b h L d", o=3, d=self.head_dim)
+        txt_qkv = scatter_heads(txt_qkv, dim=2)
+
+        vid_q, vid_k, vid_v = vid_qkv.unbind()
+        txt_q, txt_k, txt_v = txt_qkv.unbind()
+
+        vid_q, txt_q = self.norm_q(vid_q, txt_q)
+        vid_k, txt_k = self.norm_k(vid_k, txt_k)
+
+        if self.rope:
+            vid_q, vid_k = self.rope(vid_q, vid_k, (T, H, W))
+
+        def vid_window(v):
+            return rearrange(
+                v,
+                "b h (nt tt nh hh nw ww) d -> b h (nt nh nw) (tt hh ww) d",
+                hh=hh,
+                ww=ww,
+                tt=tt,
+                nh=nh,
+                nw=nw,
+                nt=nt,
+            )
+
+        def txt_window(t):
+            return rearrange(t, "b h L d -> b h 1 L d").expand(-1, -1, nt * nh * nw, -1, -1)
+
+        # Process video attention.
+        vid_msk = safe_pad_operation(txt_mask, (tt * hh * ww, 0), value=True)
+        vid_msk = rearrange(vid_msk, "b l -> b 1 1 1 l").expand(-1, 1, 1, tt * hh * ww, -1)
+        vid_out = self.attn(
+            vid_window(vid_q),
+            torch.cat([vid_window(vid_k), txt_window(txt_k)], dim=-2),
+            torch.cat([vid_window(vid_v), txt_window(txt_v)], dim=-2),
+            vid_msk,
+        )
+        vid_out = rearrange(
+            vid_out,
+            "b h (nt nh nw) (tt hh ww) d -> b (nt tt) (nh hh) (nw ww) (h d)",
+            hh=hh,
+            ww=ww,
+            tt=tt,
+            nh=nh,
+            nw=nw,
+        )
+        vid_out = gather_heads_scatter_seq(vid_out, head_dim=4, seq_dim=2)
+
+        # Process text attention.
+        txt_msk = safe_pad_operation(txt_mask, (T * H * W, 0), value=True)
+        txt_msk = rearrange(txt_msk, "b l -> b 1 1 l").expand(-1, 1, L, -1)
+        txt_out = self.attn(
+            txt_q,
+            torch.cat([vid_k, txt_k], dim=-2),
+            torch.cat([vid_v, txt_v], dim=-2),
+            txt_msk,
+        )
+        txt_out = rearrange(txt_out, "b h L d -> b L (h d)")
+        txt_out = gather_heads(txt_out, dim=2)
+
+        # Project output.
+        vid_out, txt_out = self.proj_out(vid_out, txt_out)
+        return vid_out, txt_out
+
+
+class MMWindowTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        vid_dim: int,
+        txt_dim: int,
+        emb_dim: int,
+        heads: int,
+        head_dim: int,
+        expand_ratio: int,
+        norm: norm_layer_type,
+        norm_eps: float,
+        ada: ada_layer_type,
+        qk_bias: bool,
+        qk_rope: bool,
+        qk_norm: norm_layer_type,
+        window: Union[int, Tuple[int, int, int]],
+        window_method: str,
+        shared_qkv: bool,
+        shared_mlp: bool,
+        mlp_type: str,
+        **kwargs,
+    ):
+        super().__init__()
+        dim = MMArg(vid_dim, txt_dim)
+        self.attn_norm = MMModule(norm, dim=dim, eps=norm_eps, elementwise_affine=False)
+        self.attn = MMWindowAttention(
+            vid_dim=vid_dim,
+            txt_dim=txt_dim,
+            heads=heads,
+            head_dim=head_dim,
+            qk_bias=qk_bias,
+            qk_rope=qk_rope,
+            qk_norm=qk_norm,
+            qk_norm_eps=norm_eps,
+            window=window,
+            window_method=window_method,
+            shared_qkv=shared_qkv,
+        )
+        self.mlp_norm = MMModule(norm, dim=dim, eps=norm_eps, elementwise_affine=False)
+        self.mlp = MMModule(
+            get_mlp(mlp_type),
+            dim=dim,
+            expand_ratio=expand_ratio,
+            shared_weights=shared_mlp,
+        )
+        self.ada = MMModule(ada, dim=dim, emb_dim=emb_dim, layers=["attn", "mlp"])
+
+    def forward(
+        self,
+        vid: torch.FloatTensor,
+        txt: torch.FloatTensor,
+        txt_mask: torch.BoolTensor,
+        emb: torch.FloatTensor,
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.FloatTensor,
+    ]:
+        vid_attn, txt_attn = self.attn_norm(vid, txt)
+        vid_attn, txt_attn = self.ada(vid_attn, txt_attn, emb=emb, layer="attn", mode="in")
+        vid_attn, txt_attn = self.attn(vid_attn, txt_attn, txt_mask=txt_mask)
+        vid_attn, txt_attn = self.ada(vid_attn, txt_attn, emb=emb, layer="attn", mode="out")
+        vid_attn, txt_attn = (vid_attn + vid), (txt_attn + txt)
+
+        vid_mlp, txt_mlp = self.mlp_norm(vid_attn, txt_attn)
+        vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, emb=emb, layer="mlp", mode="in")
+        vid_mlp, txt_mlp = self.mlp(vid_mlp, txt_mlp)
+        vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, emb=emb, layer="mlp", mode="out")
+        vid_mlp, txt_mlp = (vid_mlp + vid_attn), (txt_mlp + txt_attn)
+
+        return vid_mlp, txt_mlp

modules/seedvr/src/models/dit/embedding.py

@@ -0,0 +1,62 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Optional, Union
+import torch
+from diffusers.models.embeddings import get_timestep_embedding
+from torch import nn
+
+
+def emb_add(emb1: torch.Tensor, emb2: Optional[torch.Tensor]):
+    return emb1 if emb2 is None else emb1 + emb2
+
+
+class TimeEmbedding(nn.Module):
+    def __init__(
+        self,
+        sinusoidal_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+    ):
+        super().__init__()
+        self.sinusoidal_dim = sinusoidal_dim
+        self.proj_in = nn.Linear(sinusoidal_dim, hidden_dim)
+        self.proj_hid = nn.Linear(hidden_dim, hidden_dim)
+        self.proj_out = nn.Linear(hidden_dim, output_dim)
+        self.act = nn.SiLU()
+
+    def forward(
+        self,
+        timestep: Union[int, float, torch.IntTensor, torch.FloatTensor],
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.FloatTensor:
+        if not torch.is_tensor(timestep):
+            timestep = torch.tensor([timestep], device=device, dtype=dtype)
+        if timestep.ndim == 0:
+            timestep = timestep[None]
+
+        emb = get_timestep_embedding(
+            timesteps=timestep,
+            embedding_dim=self.sinusoidal_dim,
+            flip_sin_to_cos=False,
+            downscale_freq_shift=0,
+        )
+        emb = emb.to(dtype)
+        emb = self.proj_in(emb)
+        emb = self.act(emb)
+        emb = self.proj_hid(emb)
+        emb = self.act(emb)
+        emb = self.proj_out(emb)
+        return emb

modules/seedvr/src/models/dit/mlp.py

@@ -0,0 +1,62 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Optional
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+def get_mlp(mlp_type: Optional[str] = "normal"):
+    if mlp_type == "normal":
+        return MLP
+    elif mlp_type == "swiglu":
+        return SwiGLUMLP
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        expand_ratio: int,
+    ):
+        super().__init__()
+        self.proj_in = nn.Linear(dim, dim * expand_ratio)
+        self.act = nn.GELU("tanh")
+        self.proj_out = nn.Linear(dim * expand_ratio, dim)
+
+    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
+        x = self.proj_in(x)
+        x = self.act(x)
+        x = self.proj_out(x)
+        return x
+
+
+class SwiGLUMLP(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        expand_ratio: int,
+        multiple_of: int = 256,
+    ):
+        super().__init__()
+        hidden_dim = int(2 * dim * expand_ratio / 3)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+        self.proj_in_gate = nn.Linear(dim, hidden_dim, bias=False)
+        self.proj_out = nn.Linear(hidden_dim, dim, bias=False)
+        self.proj_in = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
+        x = self.proj_out(F.silu(self.proj_in_gate(x)) * self.proj_in(x))
+        return x

modules/seedvr/src/models/dit/mm.py

@@ -0,0 +1,67 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Tuple
+import torch
+from torch import nn
+
+
+@dataclass
+class MMArg:
+    vid: Any
+    txt: Any
+
+
+def get_args(key: str, args: List[Any]) -> List[Any]:
+    return [getattr(v, key) if isinstance(v, MMArg) else v for v in args]
+
+
+def get_kwargs(key: str, kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    return {k: getattr(v, key) if isinstance(v, MMArg) else v for k, v in kwargs.items()}
+
+
+class MMModule(nn.Module):
+    def __init__(
+        self,
+        module: Callable[..., nn.Module],
+        *args,
+        shared_weights: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.shared_weights = shared_weights
+        if self.shared_weights:
+            assert get_args("vid", args) == get_args("txt", args)
+            assert get_kwargs("vid", kwargs) == get_kwargs("txt", kwargs)
+            self.all = module(*get_args("vid", args), **get_kwargs("vid", kwargs))
+        else:
+            self.vid = module(*get_args("vid", args), **get_kwargs("vid", kwargs))
+            self.txt = module(*get_args("txt", args), **get_kwargs("txt", kwargs))
+
+    def forward(
+        self,
+        vid: torch.FloatTensor,
+        txt: torch.FloatTensor,
+        *args,
+        **kwargs,
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.FloatTensor,
+    ]:
+        vid_module = self.vid if not self.shared_weights else self.all
+        txt_module = self.txt if not self.shared_weights else self.all
+        vid = vid_module(vid, *get_args("vid", args), **get_kwargs("vid", kwargs))
+        txt = txt_module(txt, *get_args("txt", args), **get_kwargs("txt", kwargs))
+        return vid, txt

modules/seedvr/src/models/dit/modulation.py

@@ -0,0 +1,97 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Callable, List, Optional
+import torch
+from einops import rearrange
+from torch import nn
+
+from ...common.cache import Cache
+from ...common.distributed.ops import slice_inputs
+
+# (dim: int, emb_dim: int)
+ada_layer_type = Callable[[int, int], nn.Module]
+
+
+def get_ada_layer(ada_layer: str) -> ada_layer_type:
+    if ada_layer == "single":
+        return AdaSingle
+    raise NotImplementedError(f"{ada_layer} is not supported")
+
+
+def expand_dims(x: torch.Tensor, dim: int, ndim: int):
+    """
+    Expand tensor "x" to "ndim" by adding empty dims at "dim".
+    Example: x is (b d), target ndim is 5, add dim at 1, return (b 1 1 1 d).
+    """
+    shape = x.shape
+    shape = shape[:dim] + (1,) * (ndim - len(shape)) + shape[dim:]
+    return x.reshape(shape)
+
+
+class AdaSingle(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        emb_dim: int,
+        layers: List[str],
+    ):
+        assert emb_dim == 6 * dim, "AdaSingle requires emb_dim == 6 * dim"
+        super().__init__()
+        self.dim = dim
+        self.emb_dim = emb_dim
+        self.layers = layers
+        for l in layers:
+            self.register_parameter(f"{l}_shift", nn.Parameter(torch.randn(dim) / dim**0.5))
+            self.register_parameter(f"{l}_scale", nn.Parameter(torch.randn(dim) / dim**0.5 + 1))
+            self.register_parameter(f"{l}_gate", nn.Parameter(torch.randn(dim) / dim**0.5))
+
+    def forward(
+        self,
+        hid: torch.FloatTensor,  # b ... c
+        emb: torch.FloatTensor,  # b d
+        layer: str,
+        mode: str,
+        cache: Cache = Cache(disable=True),
+        branch_tag: str = "",
+        hid_len: Optional[torch.LongTensor] = None,  # b
+    ) -> torch.FloatTensor:
+        idx = self.layers.index(layer)
+        emb = rearrange(emb, "b (d l g) -> b d l g", l=len(self.layers), g=3)[..., idx, :]
+        emb = expand_dims(emb, 1, hid.ndim + 1)
+
+        if hid_len is not None:
+            emb = cache(
+                f"emb_repeat_{idx}_{branch_tag}",
+                lambda: slice_inputs(
+                    torch.cat([e.repeat(l, *([1] * e.ndim)) for e, l in zip(emb, hid_len)]),
+                    dim=0,
+                ),
+            )
+
+        shiftA, scaleA, gateA = emb.unbind(-1)
+        shiftB, scaleB, gateB = (
+            getattr(self, f"{layer}_shift"),
+            getattr(self, f"{layer}_scale"),
+            getattr(self, f"{layer}_gate"),
+        )
+
+        if mode == "in":
+            return hid.mul_(scaleA + scaleB).add_(shiftA + shiftB)
+        if mode == "out":
+            return hid.mul_(gateA + gateB)
+        raise NotImplementedError
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, emb_dim={self.emb_dim}, layers={self.layers}"

modules/seedvr/src/models/dit/na.py

@@ -0,0 +1,241 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from itertools import chain
+from typing import Callable, Dict, List, Tuple
+import einops
+import torch
+
+
+def flatten(
+    hid: List[torch.FloatTensor],  # List of (*** c)
+) -> Tuple[
+    torch.FloatTensor,  # (L c)
+    torch.LongTensor,  # (b n)
+]:
+    assert len(hid) > 0
+    shape = torch.stack([torch.tensor(x.shape[:-1], device=hid[0].device) for x in hid])
+    hid = torch.cat([x.flatten(0, -2) for x in hid])
+    return hid, shape
+
+
+def unflatten(
+    hid: torch.FloatTensor,  # (L c) or (L ... c)
+    hid_shape: torch.LongTensor,  # (b n)
+) -> List[torch.Tensor]:  # List of (*** c) or (*** ... c)
+    hid_len = hid_shape.prod(-1)
+    hid = hid.split(hid_len.tolist())
+    hid = [x.unflatten(0, s.tolist()) for x, s in zip(hid, hid_shape)]
+    return hid
+
+
+def concat(
+    vid: torch.FloatTensor,  # (VL ... c)
+    txt: torch.FloatTensor,  # (TL ... c)
+    vid_len: torch.LongTensor,  # (b)
+    txt_len: torch.LongTensor,  # (b)
+) -> torch.FloatTensor:  # (L ... c)
+    vid = torch.split(vid, vid_len.tolist())
+    txt = torch.split(txt, txt_len.tolist())
+    return torch.cat(list(chain(*zip(vid, txt))))
+
+
+def concat_idx(
+    vid_len: torch.LongTensor,  # (b)
+    txt_len: torch.LongTensor,  # (b)
+) -> Tuple[
+    Callable,
+    Callable,
+]:
+    device = vid_len.device
+    vid_idx = torch.arange(vid_len.sum(), device=device)
+    txt_idx = torch.arange(len(vid_idx), len(vid_idx) + txt_len.sum(), device=device)
+    tgt_idx = concat(vid_idx, txt_idx, vid_len, txt_len)
+    src_idx = torch.argsort(tgt_idx)
+    return (
+        lambda vid, txt: torch.index_select(torch.cat([vid, txt]), 0, tgt_idx),
+        lambda all: torch.index_select(all, 0, src_idx).split([len(vid_idx), len(txt_idx)]),
+    )
+
+
+def unconcat(
+    all: torch.FloatTensor,  # (L ... c)
+    vid_len: torch.LongTensor,  # (b)
+    txt_len: torch.LongTensor,  # (b)
+) -> Tuple[
+    torch.FloatTensor,  # (VL ... c)
+    torch.FloatTensor,  # (TL ... c)
+]:
+    interleave_len = list(chain(*zip(vid_len.tolist(), txt_len.tolist())))
+    all = all.split(interleave_len)
+    vid = torch.cat(all[0::2])
+    txt = torch.cat(all[1::2])
+    return vid, txt
+
+
+def repeat_concat(
+    vid: torch.FloatTensor,  # (VL ... c)
+    txt: torch.FloatTensor,  # (TL ... c)
+    vid_len: torch.LongTensor,  # (n*b)
+    txt_len: torch.LongTensor,  # (b)
+    txt_repeat: List,  # (n)
+) -> torch.FloatTensor:  # (L ... c)
+    vid = torch.split(vid, vid_len.tolist())
+    txt = torch.split(txt, txt_len.tolist())
+    txt = [[x] * n for x, n in zip(txt, txt_repeat)]
+    txt = list(chain(*txt))
+    return torch.cat(list(chain(*zip(vid, txt))))
+
+
+def repeat_concat_idx(
+    vid_len: torch.LongTensor,  # (n*b)
+    txt_len: torch.LongTensor,  # (b)
+    txt_repeat: torch.LongTensor,  # (n)
+) -> Tuple[
+    Callable,
+    Callable,
+]:
+    device = vid_len.device
+    vid_idx = torch.arange(vid_len.sum(), device=device)
+    txt_idx = torch.arange(len(vid_idx), len(vid_idx) + txt_len.sum(), device=device)
+    txt_repeat_list = txt_repeat.tolist()
+    tgt_idx = repeat_concat(vid_idx, txt_idx, vid_len, txt_len, txt_repeat)
+    src_idx = torch.argsort(tgt_idx)
+    txt_idx_len = len(tgt_idx) - len(vid_idx)
+    repeat_txt_len = (txt_len * txt_repeat).tolist()
+
+    def unconcat_coalesce(all):
+        """
+        Un-concat vid & txt, and coalesce the repeated txt.
+        e.g. vid [0 1 2 3 4 5 6 7 8] -> 3 splits -> [0 1 2] [3 4 5] [6 7 8]
+             txt [9 10]
+             repeat_concat ==> [0 1 2 9 10 3 4 5 9 10 6 7 8 9 10]
+             1. argsort re-index ==> [0 1 2 3 4 5 6 7 8 9 9 9 10 10 10]
+                           split ==> vid_out [0 1 2 3 4 5 6 7 8] txt_out [9 9 9 10 10 10]
+             2. reshape & mean for each sample to coalesce the repeated txt.
+        """
+        vid_out, txt_out = all[src_idx].split([len(vid_idx), txt_idx_len])
+        txt_out_coalesced = []
+        for txt, repeat_time in zip(txt_out.split(repeat_txt_len), txt_repeat_list):
+            txt = txt.reshape(-1, repeat_time, *txt.shape[1:]).mean(1)
+            txt_out_coalesced.append(txt)
+        return vid_out, torch.cat(txt_out_coalesced)
+
+    # Note: Backward of torch.index_select is non-deterministic when existing repeated index,
+    # the difference may cumulative like torch.repeat_interleave, so we use vanilla index here.
+    return (
+        lambda vid, txt: torch.cat([vid, txt])[tgt_idx],
+        lambda all: unconcat_coalesce(all),
+    )
+
+
+def rearrange(
+    hid: torch.FloatTensor,  # (L c)
+    hid_shape: torch.LongTensor,  # (b n)
+    pattern: str,
+    **kwargs: Dict[str, int],
+) -> Tuple[
+    torch.FloatTensor,
+    torch.LongTensor,
+]:
+    return flatten([einops.rearrange(h, pattern, **kwargs) for h in unflatten(hid, hid_shape)])
+
+
+def rearrange_idx(
+    hid_shape: torch.LongTensor,  # (b n)
+    pattern: str,
+    **kwargs: Dict[str, int],
+) -> Tuple[Callable, Callable, torch.LongTensor]:
+    hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(-1)
+    tgt_idx, tgt_shape = rearrange(hid_idx, hid_shape, pattern, **kwargs)
+    tgt_idx = tgt_idx.squeeze(-1)
+    src_idx = torch.argsort(tgt_idx)
+    return (
+        lambda hid: torch.index_select(hid, 0, tgt_idx),
+        lambda hid: torch.index_select(hid, 0, src_idx),
+        tgt_shape,
+    )
+
+
+def repeat(
+    hid: torch.FloatTensor,  # (L c)
+    hid_shape: torch.LongTensor,  # (b n)
+    pattern: str,
+    **kwargs: Dict[str, torch.LongTensor],  # (b)
+) -> Tuple[
+    torch.FloatTensor,
+    torch.LongTensor,
+]:
+    hid = unflatten(hid, hid_shape)
+    kwargs = [{k: v[i].item() for k, v in kwargs.items()} for i in range(len(hid))]
+    return flatten([einops.repeat(h, pattern, **a) for h, a in zip(hid, kwargs)])
+
+
+def pack(
+    samples: List[torch.Tensor],  # List of (h w c).
+) -> Tuple[
+    List[torch.Tensor],  # groups [(b1 h1 w1 c1), (b2 h2 w2 c2)]
+    List[List[int]],  # reversal indices.
+]:
+    batches = {}
+    indices = {}
+    for i, sample in enumerate(samples):
+        shape = sample.shape
+        batches[shape] = batches.get(shape, [])
+        indices[shape] = indices.get(shape, [])
+        batches[shape].append(sample)
+        indices[shape].append(i)
+
+    batches = list(map(torch.stack, batches.values()))
+    indices = list(indices.values())
+    return batches, indices
+
+
+def unpack(
+    batches: List[torch.Tensor],
+    indices: List[List[int]],
+) -> List[torch.Tensor]:
+    samples = [None] * (max(chain(*indices)) + 1)
+    for batch, index in zip(batches, indices):
+        for sample, i in zip(batch.unbind(), index):
+            samples[i] = sample
+    return samples
+
+
+def window(
+    hid: torch.FloatTensor,  # (L c)
+    hid_shape: torch.LongTensor,  # (b n)
+    window_fn: Callable[[torch.Tensor], List[torch.Tensor]],
+):
+    hid = unflatten(hid, hid_shape)
+    hid = list(map(window_fn, hid))
+    hid_windows = torch.tensor(list(map(len, hid)), device=hid_shape.device)
+    hid, hid_shape = flatten(list(chain(*hid)))
+    return hid, hid_shape, hid_windows
+
+
+def window_idx(
+    hid_shape: torch.LongTensor,  # (b n)
+    window_fn: Callable[[torch.Tensor], List[torch.Tensor]],
+):
+    hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(-1)
+    tgt_idx, tgt_shape, tgt_windows = window(hid_idx, hid_shape, window_fn)
+    tgt_idx = tgt_idx.squeeze(-1)
+    src_idx = torch.argsort(tgt_idx)
+    return (
+        lambda hid: torch.index_select(hid, 0, tgt_idx),
+        lambda hid: torch.index_select(hid, 0, src_idx),
+        tgt_shape,
+        tgt_windows,
+    )

modules/seedvr/src/models/dit/nablocks/__init__.py

@@ -0,0 +1,25 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from .mmsr_block import NaMMSRTransformerBlock
+
+nadit_blocks = {
+    "mmdit_sr": NaMMSRTransformerBlock,
+}
+
+
+def get_nablock(block_type: str):
+    if block_type in nadit_blocks:
+        return nadit_blocks[block_type]
+    raise NotImplementedError(f"{block_type} is not supported")

modules/seedvr/src/models/dit/nablocks/mmsr_block.py

@@ -0,0 +1,248 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Tuple, Union
+import torch
+from einops import rearrange
+from torch.nn import functional as F
+
+# from ..cache import Cache
+from ....common.cache import Cache
+from ....common.distributed.ops import gather_heads_scatter_seq, gather_seq_scatter_heads_qkv
+
+from .. import na
+from ..attention import FlashAttentionVarlen
+from ..blocks.mmdit_window_block import MMWindowAttention, MMWindowTransformerBlock
+from ..mm import MMArg
+from ..modulation import ada_layer_type
+from ..normalization import norm_layer_type
+from ..rope import NaRotaryEmbedding3d
+from ..window import get_window_op
+from ....common.half_precision_fixes import safe_pad_operation
+
+class NaSwinAttention(MMWindowAttention):
+    def __init__(
+        self,
+        vid_dim: int,
+        txt_dim: int,
+        heads: int,
+        head_dim: int,
+        qk_bias: bool,
+        qk_rope: bool,
+        qk_norm: norm_layer_type,
+        qk_norm_eps: float,
+        window: Union[int, Tuple[int, int, int]],
+        window_method: str,
+        shared_qkv: bool,
+        **kwargs,
+    ):
+        super().__init__(
+            vid_dim=vid_dim,
+            txt_dim=txt_dim,
+            heads=heads,
+            head_dim=head_dim,
+            qk_bias=qk_bias,
+            qk_rope=qk_rope,
+            qk_norm=qk_norm,
+            qk_norm_eps=qk_norm_eps,
+            window=window,
+            window_method=window_method,
+            shared_qkv=shared_qkv,
+        )
+        self.rope = NaRotaryEmbedding3d(dim=head_dim // 2) if qk_rope else None
+        self.attn = FlashAttentionVarlen()
+        self.window_op = get_window_op(window_method)
+
+    def forward(
+        self,
+        vid: torch.FloatTensor,  # l c
+        txt: torch.FloatTensor,  # l c
+        vid_shape: torch.LongTensor,  # b 3
+        txt_shape: torch.LongTensor,  # b 1
+        cache: Cache,
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.FloatTensor,
+    ]:
+
+        vid_qkv, txt_qkv = self.proj_qkv(vid, txt)
+        vid_qkv = gather_seq_scatter_heads_qkv(
+            vid_qkv,
+            seq_dim=0,
+            qkv_shape=vid_shape,
+            cache=cache.namespace("vid"),
+        )
+        txt_qkv = gather_seq_scatter_heads_qkv(
+            txt_qkv,
+            seq_dim=0,
+            qkv_shape=txt_shape,
+            cache=cache.namespace("txt"),
+        )
+
+        # re-org the input seq for window attn
+        cache_win = cache.namespace(f"{self.window_method}_{self.window}_sd3")
+
+        def make_window(x: torch.Tensor):
+            t, h, w, _ = x.shape
+            window_slices = self.window_op((t, h, w), self.window)
+            return [x[st, sh, sw] for (st, sh, sw) in window_slices]
+
+        window_partition, window_reverse, window_shape, window_count = cache_win(
+            "win_transform",
+            lambda: na.window_idx(vid_shape, make_window),
+        )
+        vid_qkv_win = window_partition(vid_qkv)
+
+        vid_qkv_win = rearrange(vid_qkv_win, "l (o h d) -> l o h d", o=3, d=self.head_dim)
+        txt_qkv = rearrange(txt_qkv, "l (o h d) -> l o h d", o=3, d=self.head_dim)
+
+        vid_q, vid_k, vid_v = vid_qkv_win.unbind(1)
+        txt_q, txt_k, txt_v = txt_qkv.unbind(1)
+
+        vid_q, txt_q = self.norm_q(vid_q, txt_q)
+        vid_k, txt_k = self.norm_k(vid_k, txt_k)
+
+        txt_len = cache("txt_len", lambda: txt_shape.prod(-1))
+
+        vid_len_win = cache_win("vid_len", lambda: window_shape.prod(-1))
+        txt_len_win = cache_win("txt_len", lambda: txt_len.repeat_interleave(window_count))
+        all_len_win = cache_win("all_len", lambda: vid_len_win + txt_len_win)
+        concat_win, unconcat_win = cache_win(
+            "mm_pnp", lambda: na.repeat_concat_idx(vid_len_win, txt_len, window_count)
+        )
+
+        # window rope
+        if self.rope:
+            vid_q, vid_k = self.rope(vid_q, vid_k, window_shape, cache_win)
+
+        out = self.attn(
+            q=concat_win(vid_q, txt_q).bfloat16(),
+            k=concat_win(vid_k, txt_k).bfloat16(),
+            v=concat_win(vid_v, txt_v).bfloat16(),
+            cu_seqlens_q=cache_win(
+                "vid_seqlens_q", lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int()
+            ),
+            cu_seqlens_k=cache_win(
+                "vid_seqlens_k", lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int()
+            ),
+            max_seqlen_q=cache_win("vid_max_seqlen_q", lambda: all_len_win.max().item()),
+            max_seqlen_k=cache_win("vid_max_seqlen_k", lambda: all_len_win.max().item()),
+        ).type_as(vid_q)
+
+        # text pooling
+        vid_out, txt_out = unconcat_win(out)
+
+        vid_out = rearrange(vid_out, "l h d -> l (h d)")
+        txt_out = rearrange(txt_out, "l h d -> l (h d)")
+        vid_out = window_reverse(vid_out)
+
+        vid_out = gather_heads_scatter_seq(vid_out, head_dim=1, seq_dim=0)
+        txt_out = gather_heads_scatter_seq(txt_out, head_dim=1, seq_dim=0)
+
+        vid_out, txt_out = self.proj_out(vid_out, txt_out)
+
+        return vid_out, txt_out
+
+
+class NaMMSRTransformerBlock(MMWindowTransformerBlock):
+    def __init__(
+        self,
+        *,
+        vid_dim: int,
+        txt_dim: int,
+        emb_dim: int,
+        heads: int,
+        head_dim: int,
+        expand_ratio: int,
+        norm: norm_layer_type,
+        norm_eps: float,
+        ada: ada_layer_type,
+        qk_bias: bool,
+        qk_rope: bool,
+        qk_norm: norm_layer_type,
+        shared_qkv: bool,
+        shared_mlp: bool,
+        mlp_type: str,
+        **kwargs,
+    ):
+        super().__init__(
+            vid_dim=vid_dim,
+            txt_dim=txt_dim,
+            emb_dim=emb_dim,
+            heads=heads,
+            head_dim=head_dim,
+            expand_ratio=expand_ratio,
+            norm=norm,
+            norm_eps=norm_eps,
+            ada=ada,
+            qk_bias=qk_bias,
+            qk_rope=qk_rope,
+            qk_norm=qk_norm,
+            shared_qkv=shared_qkv,
+            shared_mlp=shared_mlp,
+            mlp_type=mlp_type,
+            **kwargs,
+        )
+
+        self.attn = NaSwinAttention(
+            vid_dim=vid_dim,
+            txt_dim=txt_dim,
+            heads=heads,
+            head_dim=head_dim,
+            qk_bias=qk_bias,
+            qk_rope=qk_rope,
+            qk_norm=qk_norm,
+            qk_norm_eps=norm_eps,
+            shared_qkv=shared_qkv,
+            **kwargs,
+        )
+
+    def forward(
+        self,
+        vid: torch.FloatTensor,  # l c
+        txt: torch.FloatTensor,  # l c
+        vid_shape: torch.LongTensor,  # b 3
+        txt_shape: torch.LongTensor,  # b 1
+        emb: torch.FloatTensor,
+        cache: Cache,
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.FloatTensor,
+        torch.LongTensor,
+        torch.LongTensor,
+    ]:
+        hid_len = MMArg(
+            cache("vid_len", lambda: vid_shape.prod(-1)),
+            cache("txt_len", lambda: txt_shape.prod(-1)),
+        )
+        ada_kwargs = {
+            "emb": emb,
+            "hid_len": hid_len,
+            "cache": cache,
+            "branch_tag": MMArg("vid", "txt"),
+        }
+
+        vid_attn, txt_attn = self.attn_norm(vid, txt)
+        vid_attn, txt_attn = self.ada(vid_attn, txt_attn, layer="attn", mode="in", **ada_kwargs)
+        vid_attn, txt_attn = self.attn(vid_attn, txt_attn, vid_shape, txt_shape, cache)
+        vid_attn, txt_attn = self.ada(vid_attn, txt_attn, layer="attn", mode="out", **ada_kwargs)
+        vid_attn, txt_attn = (vid_attn + vid), (txt_attn + txt)
+
+        vid_mlp, txt_mlp = self.mlp_norm(vid_attn, txt_attn)
+        vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, layer="mlp", mode="in", **ada_kwargs)
+        vid_mlp, txt_mlp = self.mlp(vid_mlp, txt_mlp)
+        vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, layer="mlp", mode="out", **ada_kwargs)
+        vid_mlp, txt_mlp = (vid_mlp + vid_attn), (txt_mlp + txt_attn)
+
+        return vid_mlp, txt_mlp, vid_shape, txt_shape

modules/seedvr/src/models/dit/nadit.py

@@ -0,0 +1,350 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union, Callable
+import torch
+from torch import nn
+
+from ...common.cache import Cache
+from ...common.distributed.ops import slice_inputs
+
+from . import na
+from .embedding import TimeEmbedding
+from .modulation import get_ada_layer
+from .nablocks import get_nablock
+from .normalization import get_norm_layer
+from .patch import NaPatchIn, NaPatchOut
+
+# Fake func, no checkpointing is required for inference
+def gradient_checkpointing(module: Union[Callable, nn.Module], *args, enabled: bool, **kwargs):
+    return module(*args, **kwargs)
+
+@dataclass
+class NaDiTOutput:
+    vid_sample: torch.Tensor
+
+
+class NaDiT(nn.Module):
+    """
+    Native Resolution Diffusion Transformer (NaDiT)
+    """
+
+    gradient_checkpointing = False
+
+    def __init__(
+        self,
+        vid_in_channels: int,
+        vid_out_channels: int,
+        vid_dim: int,
+        txt_in_dim: Optional[int],
+        txt_dim: Optional[int],
+        emb_dim: int,
+        heads: int,
+        head_dim: int,
+        expand_ratio: int,
+        norm: Optional[str],
+        norm_eps: float,
+        ada: str,
+        qk_bias: bool,
+        qk_rope: bool,
+        qk_norm: Optional[str],
+        patch_size: Union[int, Tuple[int, int, int]],
+        num_layers: int,
+        block_type: Union[str, Tuple[str]],
+        shared_qkv: bool = False,
+        shared_mlp: bool = False,
+        mlp_type: str = "normal",
+        window: Optional[Tuple] = None,
+        window_method: Optional[Tuple[str]] = None,
+        temporal_window_size: int = None,
+        temporal_shifted: bool = False,
+        **kwargs,
+    ):
+        ada = get_ada_layer(ada)
+        norm = get_norm_layer(norm)
+        qk_norm = get_norm_layer(qk_norm)
+        if isinstance(block_type, str):
+            block_type = [block_type] * num_layers
+        elif len(block_type) != num_layers:
+            raise ValueError("The ``block_type`` list should equal to ``num_layers``.")
+        super().__init__()
+        self.vid_in = NaPatchIn(
+            in_channels=vid_in_channels,
+            patch_size=patch_size,
+            dim=vid_dim,
+        )
+        self.txt_in = (
+            nn.Linear(txt_in_dim, txt_dim)
+            if txt_in_dim and txt_in_dim != txt_dim
+            else nn.Identity()
+        )
+        self.emb_in = TimeEmbedding(
+            sinusoidal_dim=256,
+            hidden_dim=max(vid_dim, txt_dim),
+            output_dim=emb_dim,
+        )
+
+        if window is None or isinstance(window[0], int):
+            window = [window] * num_layers
+        if window_method is None or isinstance(window_method, str):
+            window_method = [window_method] * num_layers
+        if temporal_window_size is None or isinstance(temporal_window_size, int):
+            temporal_window_size = [temporal_window_size] * num_layers
+        if temporal_shifted is None or isinstance(temporal_shifted, bool):
+            temporal_shifted = [temporal_shifted] * num_layers
+
+        self.blocks = nn.ModuleList(
+            [
+                get_nablock(block_type[i])(
+                    vid_dim=vid_dim,
+                    txt_dim=txt_dim,
+                    emb_dim=emb_dim,
+                    heads=heads,
+                    head_dim=head_dim,
+                    expand_ratio=expand_ratio,
+                    norm=norm,
+                    norm_eps=norm_eps,
+                    ada=ada,
+                    qk_bias=qk_bias,
+                    qk_rope=qk_rope,
+                    qk_norm=qk_norm,
+                    shared_qkv=shared_qkv,
+                    shared_mlp=shared_mlp,
+                    mlp_type=mlp_type,
+                    window=window[i],
+                    window_method=window_method[i],
+                    temporal_window_size=temporal_window_size[i],
+                    temporal_shifted=temporal_shifted[i],
+                    **kwargs,
+                )
+                for i in range(num_layers)
+            ]
+        )
+        self.vid_out = NaPatchOut(
+            out_channels=vid_out_channels,
+            patch_size=patch_size,
+            dim=vid_dim,
+        )
+
+        self.need_txt_repeat = block_type[0] in [
+            "mmdit_stwin",
+            "mmdit_stwin_spatial",
+            "mmdit_stwin_3d_spatial",
+        ]
+
+    def set_gradient_checkpointing(self, enable: bool):
+        self.gradient_checkpointing = enable
+
+    def forward(
+        self,
+        vid: torch.FloatTensor,  # l c
+        txt: torch.FloatTensor,  # l c
+        vid_shape: torch.LongTensor,  # b 3
+        txt_shape: torch.LongTensor,  # b 1
+        timestep: Union[int, float, torch.IntTensor, torch.FloatTensor],  # b
+        disable_cache: bool = True,  # for test
+    ):
+        # Text input.
+        if txt_shape.size(-1) == 1 and self.need_txt_repeat:
+            txt, txt_shape = na.repeat(txt, txt_shape, "l c -> t l c", t=vid_shape[:, 0])
+        # slice vid after patching in when using sequence parallelism
+        txt = slice_inputs(txt, dim=0)
+        txt = self.txt_in(txt)
+
+        # Video input.
+        # Sequence parallel slicing is done inside patching class.
+        vid, vid_shape = self.vid_in(vid, vid_shape)
+
+        # Embedding input.
+        emb = self.emb_in(timestep, device=vid.device, dtype=vid.dtype)
+
+        # Body
+        cache = Cache(disable=disable_cache)
+        for i, block in enumerate(self.blocks):
+            vid, txt, vid_shape, txt_shape = gradient_checkpointing(
+                enabled=(self.gradient_checkpointing and self.training),
+                module=block,
+                vid=vid,
+                txt=txt,
+                vid_shape=vid_shape,
+                txt_shape=txt_shape,
+                emb=emb,
+                cache=cache,
+            )
+
+        vid, vid_shape = self.vid_out(vid, vid_shape, cache)
+        return NaDiTOutput(vid_sample=vid)
+
+
+class NaDiTUpscaler(nn.Module):
+    """
+    Native Resolution Diffusion Transformer (NaDiT)
+    """
+
+    gradient_checkpointing = False
+
+    def __init__(
+        self,
+        vid_in_channels: int,
+        vid_out_channels: int,
+        vid_dim: int,
+        txt_in_dim: Optional[int],
+        txt_dim: Optional[int],
+        emb_dim: int,
+        heads: int,
+        head_dim: int,
+        expand_ratio: int,
+        norm: Optional[str],
+        norm_eps: float,
+        ada: str,
+        qk_bias: bool,
+        qk_rope: bool,
+        qk_norm: Optional[str],
+        patch_size: Union[int, Tuple[int, int, int]],
+        num_layers: int,
+        block_type: Union[str, Tuple[str]],
+        shared_qkv: bool = False,
+        shared_mlp: bool = False,
+        mlp_type: str = "normal",
+        window: Optional[Tuple] = None,
+        window_method: Optional[Tuple[str]] = None,
+        temporal_window_size: int = None,
+        temporal_shifted: bool = False,
+        **kwargs,
+    ):
+        ada = get_ada_layer(ada)
+        norm = get_norm_layer(norm)
+        qk_norm = get_norm_layer(qk_norm)
+        if isinstance(block_type, str):
+            block_type = [block_type] * num_layers
+        elif len(block_type) != num_layers:
+            raise ValueError("The ``block_type`` list should equal to ``num_layers``.")
+        super().__init__()
+        self.vid_in = NaPatchIn(
+            in_channels=vid_in_channels,
+            patch_size=patch_size,
+            dim=vid_dim,
+        )
+        self.txt_in = (
+            nn.Linear(txt_in_dim, txt_dim)
+            if txt_in_dim and txt_in_dim != txt_dim
+            else nn.Identity()
+        )
+        self.emb_in = TimeEmbedding(
+            sinusoidal_dim=256,
+            hidden_dim=max(vid_dim, txt_dim),
+            output_dim=emb_dim,
+        )
+
+        self.emb_scale = TimeEmbedding(
+            sinusoidal_dim=256,
+            hidden_dim=max(vid_dim, txt_dim),
+            output_dim=emb_dim,
+        )
+
+        if window is None or isinstance(window[0], int):
+            window = [window] * num_layers
+        if window_method is None or isinstance(window_method, str):
+            window_method = [window_method] * num_layers
+        if temporal_window_size is None or isinstance(temporal_window_size, int):
+            temporal_window_size = [temporal_window_size] * num_layers
+        if temporal_shifted is None or isinstance(temporal_shifted, bool):
+            temporal_shifted = [temporal_shifted] * num_layers
+
+        self.blocks = nn.ModuleList(
+            [
+                get_nablock(block_type[i])(
+                    vid_dim=vid_dim,
+                    txt_dim=txt_dim,
+                    emb_dim=emb_dim,
+                    heads=heads,
+                    head_dim=head_dim,
+                    expand_ratio=expand_ratio,
+                    norm=norm,
+                    norm_eps=norm_eps,
+                    ada=ada,
+                    qk_bias=qk_bias,
+                    qk_rope=qk_rope,
+                    qk_norm=qk_norm,
+                    shared_qkv=shared_qkv,
+                    shared_mlp=shared_mlp,
+                    mlp_type=mlp_type,
+                    window=window[i],
+                    window_method=window_method[i],
+                    temporal_window_size=temporal_window_size[i],
+                    temporal_shifted=temporal_shifted[i],
+                    **kwargs,
+                )
+                for i in range(num_layers)
+            ]
+        )
+        self.vid_out = NaPatchOut(
+            out_channels=vid_out_channels,
+            patch_size=patch_size,
+            dim=vid_dim,
+        )
+
+        self.need_txt_repeat = block_type[0] in [
+            "mmdit_stwin",
+            "mmdit_stwin_spatial",
+            "mmdit_stwin_3d_spatial",
+        ]
+
+    def set_gradient_checkpointing(self, enable: bool):
+        self.gradient_checkpointing = enable
+
+    def forward(
+        self,
+        vid: torch.FloatTensor,  # l c
+        txt: torch.FloatTensor,  # l c
+        vid_shape: torch.LongTensor,  # b 3
+        txt_shape: torch.LongTensor,  # b 1
+        timestep: Union[int, float, torch.IntTensor, torch.FloatTensor],  # b
+        downscale: Union[int, float, torch.IntTensor, torch.FloatTensor],  # b
+        disable_cache: bool = False,  # for test
+    ):
+
+        # Text input.
+        if txt_shape.size(-1) == 1 and self.need_txt_repeat:
+            txt, txt_shape = na.repeat(txt, txt_shape, "l c -> t l c", t=vid_shape[:, 0])
+        # slice vid after patching in when using sequence parallelism
+        txt = slice_inputs(txt, dim=0)
+        txt = self.txt_in(txt)
+
+        # Video input.
+        # Sequence parallel slicing is done inside patching class.
+        vid, vid_shape = self.vid_in(vid, vid_shape)
+
+        # Embedding input.
+        emb = self.emb_in(timestep, device=vid.device, dtype=vid.dtype)
+        emb_scale = self.emb_scale(downscale, device=vid.device, dtype=vid.dtype)
+        emb = emb + emb_scale
+
+        # Body
+        cache = Cache(disable=disable_cache)
+        for i, block in enumerate(self.blocks):
+            vid, txt, vid_shape, txt_shape = gradient_checkpointing(
+                enabled=(self.gradient_checkpointing and self.training),
+                module=block,
+                vid=vid,
+                txt=txt,
+                vid_shape=vid_shape,
+                txt_shape=txt_shape,
+                emb=emb,
+                cache=cache,
+            )
+
+        vid, vid_shape = self.vid_out(vid, vid_shape, cache)
+        return NaDiTOutput(vid_sample=vid)

modules/seedvr/src/models/dit/normalization.py

@@ -0,0 +1,131 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Callable, Optional
+from diffusers.models.normalization import RMSNorm
+from torch import nn
+import torch
+import torch.nn.functional as F
+import numbers
+from torch.nn.parameter import Parameter
+from torch.nn import init
+
+# (dim: int, eps: float, elementwise_affine: bool)
+norm_layer_type = Callable[[int, float, bool], nn.Module]
+
+
+class CustomLayerNorm(nn.Module):
+    """
+    Custom LayerNorm implementation to replace Apex FusedLayerNorm
+    """
+    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
+        super(CustomLayerNorm, self).__init__()
+
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        self.normalized_shape = torch.Size(normalized_shape)
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+
+        if self.elementwise_affine:
+            self.weight = Parameter(torch.Tensor(*normalized_shape))
+            self.bias = Parameter(torch.Tensor(*normalized_shape))
+        else:
+            self.register_parameter('weight', None)
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        if self.elementwise_affine:
+            init.ones_(self.weight)
+            init.zeros_(self.bias)
+
+    def forward(self, input):
+        return F.layer_norm(
+            input, self.normalized_shape, self.weight, self.bias, self.eps)
+
+
+class CustomRMSNorm(nn.Module):
+    """
+    Custom RMSNorm implementation to replace Apex FusedRMSNorm
+    """
+    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
+        super(CustomRMSNorm, self).__init__()
+
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        self.normalized_shape = torch.Size(normalized_shape)
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+
+        if self.elementwise_affine:
+            self.weight = Parameter(torch.ones(*normalized_shape))
+        else:
+            self.register_parameter('weight', None)
+
+    def forward(self, input):
+        # RMS normalization: x / sqrt(mean(x^2) + eps) * weight
+        dims = tuple(range(-len(self.normalized_shape), 0))
+
+        # Calculate RMS: sqrt(mean(x^2))
+        variance = input.pow(2).mean(dim=dims, keepdim=True)
+        rms = torch.sqrt(variance + self.eps)
+
+        # Normalize
+        normalized = input / rms
+
+        if self.elementwise_affine:
+            return normalized * self.weight
+        return normalized
+
+
+def get_norm_layer(norm_type: Optional[str]) -> norm_layer_type:
+
+    def _norm_layer(dim: int, eps: float, elementwise_affine: bool):
+        if norm_type is None:
+            return nn.Identity()
+
+        if norm_type == "layer":
+            return nn.LayerNorm(
+                normalized_shape=dim,
+                eps=eps,
+                elementwise_affine=elementwise_affine,
+            )
+
+        if norm_type == "rms":
+            return RMSNorm(
+                dim=dim,
+                eps=eps,
+                elementwise_affine=elementwise_affine,
+            )
+
+        if norm_type == "fusedln":
+            # Use custom LayerNorm instead of Apex FusedLayerNorm
+            return CustomLayerNorm(
+                normalized_shape=dim,
+                elementwise_affine=elementwise_affine,
+                eps=eps,
+            )
+
+        if norm_type == "fusedrms":
+            # Use custom RMSNorm instead of Apex FusedRMSNorm
+            return CustomRMSNorm(
+                normalized_shape=dim,
+                elementwise_affine=elementwise_affine,
+                eps=eps,
+            )
+
+        raise NotImplementedError(f"{norm_type} is not supported")
+
+    return _norm_layer

modules/seedvr/src/models/dit/patch.py

@@ -0,0 +1,112 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Tuple, Union
+import torch
+from einops import rearrange
+from torch import nn
+from torch.nn.modules.utils import _triple
+
+from ...common.cache import Cache
+from ...common.distributed.ops import gather_outputs, slice_inputs
+
+from . import na
+
+
+class PatchIn(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        patch_size: Union[int, Tuple[int, int, int]],
+        dim: int,
+    ):
+        super().__init__()
+        t, h, w = _triple(patch_size)
+        self.patch_size = t, h, w
+        self.proj = nn.Linear(in_channels * t * h * w, dim)
+
+    def forward(
+        self,
+        vid: torch.Tensor,
+    ) -> torch.Tensor:
+        t, h, w = self.patch_size
+        vid = rearrange(vid, "b c (T t) (H h) (W w) -> b T H W (t h w c)", t=t, h=h, w=w)
+        vid = self.proj(vid)
+        return vid
+
+
+class PatchOut(nn.Module):
+    def __init__(
+        self,
+        out_channels: int,
+        patch_size: Union[int, Tuple[int, int, int]],
+        dim: int,
+    ):
+        super().__init__()
+        t, h, w = _triple(patch_size)
+        self.patch_size = t, h, w
+        self.proj = nn.Linear(dim, out_channels * t * h * w)
+
+    def forward(
+        self,
+        vid: torch.Tensor,
+    ) -> torch.Tensor:
+        t, h, w = self.patch_size
+        vid = self.proj(vid)
+        vid = rearrange(vid, "b T H W (t h w c) -> b c (T t) (H h) (W w)", t=t, h=h, w=w)
+        return vid
+
+
+class NaPatchIn(PatchIn):
+    def forward(
+        self,
+        vid: torch.Tensor,  # l c
+        vid_shape: torch.LongTensor,
+    ) -> torch.Tensor:
+        t, h, w = self.patch_size
+        if not (t == h == w == 1):
+            vid, vid_shape = na.rearrange(
+                vid, vid_shape, "(T t) (H h) (W w) c -> T H W (t h w c)", t=t, h=h, w=w
+            )
+        # slice vid after patching in when using sequence parallelism
+        vid = slice_inputs(vid, dim=0)
+        vid = self.proj(vid)
+        return vid, vid_shape
+
+
+class NaPatchOut(PatchOut):
+    def forward(
+        self,
+        vid: torch.FloatTensor,  # l c
+        vid_shape: torch.LongTensor,
+        cache: Cache = Cache(disable=True),
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.LongTensor,
+    ]:
+        t, h, w = self.patch_size
+        vid = self.proj(vid)
+        # gather vid before patching out when enabling sequence parallelism
+        vid = gather_outputs(
+            vid,
+            gather_dim=0,
+            padding_dim=0,
+            unpad_shape=vid_shape,
+            cache=cache.namespace("vid"),
+        )
+        if not (t == h == w == 1):
+            vid, vid_shape = na.rearrange(
+                vid, vid_shape, "T H W (t h w c) -> (T t) (H h) (W w) c", t=t, h=h, w=w
+            )
+        return vid, vid_shape

modules/seedvr/src/models/dit/rope.py

@@ -0,0 +1,102 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from functools import lru_cache
+from typing import Tuple
+import torch
+from einops import rearrange
+from rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb
+from torch import nn
+
+from ...common.cache import Cache
+
+
+class RotaryEmbeddingBase(nn.Module):
+    def __init__(self, dim: int, rope_dim: int):
+        super().__init__()
+        self.rope = RotaryEmbedding(
+            dim=dim // rope_dim,
+            freqs_for="pixel",
+            max_freq=256,
+        )
+        # 1. Set model.requires_grad_(True) after model creation will make
+        #    the `requires_grad=False` for rope freqs no longer hold.
+        # 2. Even if we don't set requires_grad_(True) explicitly,
+        #    FSDP is not memory efficient when handling fsdp_wrap
+        #    with mixed requires_grad=True/False.
+        # With above consideration, it is easier just remove the freqs
+        # out of nn.Parameters when `learned_freq=False`
+        freqs = self.rope.freqs
+        del self.rope.freqs
+        self.rope.register_buffer("freqs", freqs.data)
+
+    @lru_cache(maxsize=128)
+    def get_axial_freqs(self, *dims):
+        return self.rope.get_axial_freqs(*dims)
+
+
+class RotaryEmbedding3d(RotaryEmbeddingBase):
+    def __init__(self, dim: int):
+        super().__init__(dim, rope_dim=3)
+
+    def forward(
+        self,
+        q: torch.FloatTensor,  # b h l d
+        k: torch.FloatTensor,  # b h l d
+        size: Tuple[int, int, int],
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.FloatTensor,
+    ]:
+        T, H, W = size
+        freqs = self.get_axial_freqs(T, H, W)
+        q = rearrange(q, "b h (T H W) d -> b h T H W d", T=T, H=H, W=W)
+        k = rearrange(k, "b h (T H W) d -> b h T H W d", T=T, H=H, W=W)
+        q = apply_rotary_emb(freqs, q)
+        k = apply_rotary_emb(freqs, k)
+        q = rearrange(q, "b h T H W d -> b h (T H W) d")
+        k = rearrange(k, "b h T H W d -> b h (T H W) d")
+        return q, k
+
+
+class NaRotaryEmbedding3d(RotaryEmbedding3d):
+    def forward(
+        self,
+        q: torch.FloatTensor,  # L h d
+        k: torch.FloatTensor,  # L h d
+        shape: torch.LongTensor,
+        cache: Cache,
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.FloatTensor,
+    ]:
+        freqs = cache("rope_freqs_3d", lambda: self.get_freqs(shape))
+        freqs = freqs.to(device=q.device, dtype=q.dtype)
+        q = rearrange(q, "L h d -> h L d")
+        k = rearrange(k, "L h d -> h L d")
+        q = apply_rotary_emb(freqs, q.float()).to(q.dtype)
+        k = apply_rotary_emb(freqs, k.float()).to(k.dtype)
+        q = rearrange(q, "h L d -> L h d")
+        k = rearrange(k, "h L d -> L h d")
+        return q, k
+
+    def get_freqs(
+        self,
+        shape: torch.LongTensor,
+    ) -> torch.Tensor:
+        freq_list = []
+        for f, h, w in shape.tolist():
+            freqs = self.get_axial_freqs(f, h, w)
+            freq_list.append(freqs.view(-1, freqs.size(-1)))
+        return torch.cat(freq_list, dim=0)

modules/seedvr/src/models/dit/window.py

@@ -0,0 +1,83 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from math import ceil
+from typing import Tuple
+import math
+
+def get_window_op(name: str):
+    if name == "720pwin_by_size_bysize":
+        return make_720Pwindows_bysize
+    if name == "720pswin_by_size_bysize":
+        return make_shifted_720Pwindows_bysize
+    raise ValueError(f"Unknown windowing method: {name}")
+
+
+# -------------------------------- Windowing -------------------------------- #
+def make_720Pwindows_bysize(size: Tuple[int, int, int], num_windows: Tuple[int, int, int]):
+    t, h, w = size
+    resized_nt, resized_nh, resized_nw = num_windows
+    #cal windows under 720p
+    scale = math.sqrt((45 * 80) / (h * w))
+    resized_h, resized_w = round(h * scale), round(w * scale)
+    wh, ww = ceil(resized_h / resized_nh), ceil(resized_w / resized_nw)  # window size.
+    wt = ceil(min(t, 30) / resized_nt)  # window size.
+    nt, nh, nw = ceil(t / wt), ceil(h / wh), ceil(w / ww)  # window size.
+    return [
+        (
+            slice(it * wt, min((it + 1) * wt, t)),
+            slice(ih * wh, min((ih + 1) * wh, h)),
+            slice(iw * ww, min((iw + 1) * ww, w)),
+        )
+        for iw in range(nw)
+        if min((iw + 1) * ww, w) > iw * ww
+        for ih in range(nh)
+        if min((ih + 1) * wh, h) > ih * wh
+        for it in range(nt)
+        if min((it + 1) * wt, t) > it * wt
+    ]
+
+def make_shifted_720Pwindows_bysize(size: Tuple[int, int, int], num_windows: Tuple[int, int, int]):
+    t, h, w = size
+    resized_nt, resized_nh, resized_nw = num_windows
+    #cal windows under 720p
+    scale = math.sqrt((45 * 80) / (h * w))
+    resized_h, resized_w = round(h * scale), round(w * scale)
+    wh, ww = ceil(resized_h / resized_nh), ceil(resized_w / resized_nw)  # window size.
+    wt = ceil(min(t, 30) / resized_nt)  # window size.
+
+    st, sh, sw = (  # shift size.
+        0.5 if wt < t else 0,
+        0.5 if wh < h else 0,
+        0.5 if ww < w else 0,
+    )
+    nt, nh, nw = ceil((t - st) / wt), ceil((h - sh) / wh), ceil((w - sw) / ww)  # window size.
+    nt, nh, nw = (  # number of window.
+        nt + 1 if st > 0 else 1,
+        nh + 1 if sh > 0 else 1,
+        nw + 1 if sw > 0 else 1,
+    )
+    return [
+        (
+            slice(max(int((it - st) * wt), 0), min(int((it - st + 1) * wt), t)),
+            slice(max(int((ih - sh) * wh), 0), min(int((ih - sh + 1) * wh), h)),
+            slice(max(int((iw - sw) * ww), 0), min(int((iw - sw + 1) * ww), w)),
+        )
+        for iw in range(nw)
+        if min(int((iw - sw + 1) * ww), w) > max(int((iw - sw) * ww), 0)
+        for ih in range(nh)
+        if min(int((ih - sh + 1) * wh), h) > max(int((ih - sh) * wh), 0)
+        for it in range(nt)
+        if min(int((it - st + 1) * wt), t) > max(int((it - st) * wt), 0)
+    ]

modules/seedvr/src/models/dit_v2/attention.py

@@ -0,0 +1,93 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+import torch
+import torch.nn.functional as F
+
+#from flash_attn import flash_attn_varlen_func
+
+from torch import nn
+
+
+def pytorch_varlen_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p=0.0, softmax_scale=None, causal=False, deterministic=False):
+    """
+    A PyTorch-based implementation of variable-length attention to replace flash_attn_varlen_func.
+    It processes each sequence in the batch individually.
+    """
+    # Create an empty tensor to store the output.
+    output = torch.empty_like(q)
+
+    # Iterate over each sequence in the batch. The batch size is the number of sequences.
+    for i in range(len(cu_seqlens_q) - 1):
+        # Determine the start and end indices for the current sequence.
+        start_q, end_q = cu_seqlens_q[i], cu_seqlens_q[i+1]
+        start_k, end_k = cu_seqlens_k[i], cu_seqlens_k[i+1]
+
+        # Slice the q, k, and v tensors to get the data for the current sequence.
+        # The shape is (seq_len, heads, head_dim).
+        q_i = q[start_q:end_q]
+        k_i = k[start_k:end_k]
+        v_i = v[start_k:end_k]
+
+        # Reshape for torch's scaled_dot_product_attention which expects (batch, heads, seq, dim).
+        # Here, we treat each sequence as a batch of 1.
+        q_i = q_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_q, head_dim)
+        k_i = k_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_k, head_dim)
+        v_i = v_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_k, head_dim)
+
+        # Use PyTorch's built-in scaled dot-product attention.
+        output_i = F.scaled_dot_product_attention(
+            q_i, k_i, v_i,
+            dropout_p=dropout_p if not deterministic else 0.0,
+            is_causal=causal
+        )
+
+        # Reshape the output back to the original format (seq_len, heads, head_dim)
+        output_i = output_i.squeeze(0).permute(1, 0, 2)
+
+        # Place the result for the current sequence into the main output tensor.
+        output[start_q:end_q] = output_i
+
+    return output
+
+class TorchAttention(nn.Module):
+    def tflops(self, args, kwargs, output) -> float:
+        assert len(args) == 0 or len(args) > 2, "query, key should both provided by args / kwargs"
+        q = kwargs.get("query") or args[0]
+        k = kwargs.get("key") or args[1]
+        b, h, sq, d = q.shape
+        b, h, sk, d = k.shape
+        return b * h * (4 * d * (sq / 1e6) * (sk / 1e6))
+
+    def forward(self, *args, **kwargs):
+        return F.scaled_dot_product_attention(*args, **kwargs)
+
+
+class FlashAttentionVarlen(nn.Module):
+    def tflops(self, args, kwargs, output) -> float:
+        cu_seqlens_q = kwargs["cu_seqlens_q"]
+        cu_seqlens_k = kwargs["cu_seqlens_k"]
+        _, h, d = output.shape
+        seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]) / 1e6
+        seqlens_k = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]) / 1e6
+        return h * (4 * d * (seqlens_q * seqlens_k).sum())
+
+    def forward(self, *args, **kwargs):
+        kwargs["deterministic"] = torch.are_deterministic_algorithms_enabled()
+        try:
+            from flash_attn import flash_attn_varlen_func
+            return flash_attn_varlen_func(*args, **kwargs)
+        except ImportError:
+            return pytorch_varlen_attention(*args, **kwargs)
+

modules/seedvr/src/models/dit_v2/embedding.py

@@ -0,0 +1,62 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Optional, Union
+import torch
+from diffusers.models.embeddings import get_timestep_embedding
+from torch import nn
+
+
+def emb_add(emb1: torch.Tensor, emb2: Optional[torch.Tensor]):
+    return emb1 if emb2 is None else emb1 + emb2
+
+
+class TimeEmbedding(nn.Module):
+    def __init__(
+        self,
+        sinusoidal_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+    ):
+        super().__init__()
+        self.sinusoidal_dim = sinusoidal_dim
+        self.proj_in = nn.Linear(sinusoidal_dim, hidden_dim)
+        self.proj_hid = nn.Linear(hidden_dim, hidden_dim)
+        self.proj_out = nn.Linear(hidden_dim, output_dim)
+        self.act = nn.SiLU()
+
+    def forward(
+        self,
+        timestep: Union[int, float, torch.IntTensor, torch.FloatTensor],
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.FloatTensor:
+        if not torch.is_tensor(timestep):
+            timestep = torch.tensor([timestep], device=device, dtype=dtype)
+        if timestep.ndim == 0:
+            timestep = timestep[None]
+
+        emb = get_timestep_embedding(
+            timesteps=timestep,
+            embedding_dim=self.sinusoidal_dim,
+            flip_sin_to_cos=False,
+            downscale_freq_shift=0,
+        )
+        emb = emb.to(dtype)
+        emb = self.proj_in(emb)
+        emb = self.act(emb)
+        emb = self.proj_hid(emb)
+        emb = self.act(emb)
+        emb = self.proj_out(emb)
+        return emb

modules/seedvr/src/models/dit_v2/mlp.py

@@ -0,0 +1,62 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Optional
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+def get_mlp(mlp_type: Optional[str] = "normal"):
+    if mlp_type == "normal":
+        return MLP
+    elif mlp_type == "swiglu":
+        return SwiGLUMLP
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        expand_ratio: int,
+    ):
+        super().__init__()
+        self.proj_in = nn.Linear(dim, dim * expand_ratio)
+        self.act = nn.GELU("tanh")
+        self.proj_out = nn.Linear(dim * expand_ratio, dim)
+
+    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
+        x = self.proj_in(x)
+        x = self.act(x)
+        x = self.proj_out(x)
+        return x
+
+
+class SwiGLUMLP(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        expand_ratio: int,
+        multiple_of: int = 256,
+    ):
+        super().__init__()
+        hidden_dim = int(2 * dim * expand_ratio / 3)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+        self.proj_in_gate = nn.Linear(dim, hidden_dim, bias=False)
+        self.proj_out = nn.Linear(hidden_dim, dim, bias=False)
+        self.proj_in = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
+        x = self.proj_out(F.silu(self.proj_in_gate(x)) * self.proj_in(x))
+        return x

modules/seedvr/src/models/dit_v2/mm.py

@@ -0,0 +1,74 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Tuple
+import torch
+from torch import nn
+
+
+@dataclass
+class MMArg:
+    vid: Any
+    txt: Any
+
+
+def get_args(key: str, args: List[Any]) -> List[Any]:
+    return [getattr(v, key) if isinstance(v, MMArg) else v for v in args]
+
+
+def get_kwargs(key: str, kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    return {k: getattr(v, key) if isinstance(v, MMArg) else v for k, v in kwargs.items()}
+
+
+class MMModule(nn.Module):
+    def __init__(
+        self,
+        module: Callable[..., nn.Module],
+        *args,
+        shared_weights: bool = False,
+        vid_only: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.shared_weights = shared_weights
+        self.vid_only = vid_only
+        if self.shared_weights:
+            assert get_args("vid", args) == get_args("txt", args)
+            assert get_kwargs("vid", kwargs) == get_kwargs("txt", kwargs)
+            self.all = module(*get_args("vid", args), **get_kwargs("vid", kwargs))
+        else:
+            self.vid = module(*get_args("vid", args), **get_kwargs("vid", kwargs))
+            self.txt = (
+                module(*get_args("txt", args), **get_kwargs("txt", kwargs))
+                if not vid_only
+                else None
+            )
+
+    def forward(
+        self,
+        vid: torch.FloatTensor,
+        txt: torch.FloatTensor,
+        *args,
+        **kwargs,
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.FloatTensor,
+    ]:
+        vid_module = self.vid if not self.shared_weights else self.all
+        vid = vid_module(vid, *get_args("vid", args), **get_kwargs("vid", kwargs))
+        if not self.vid_only:
+            txt_module = self.txt if not self.shared_weights else self.all
+            txt = txt_module(txt, *get_args("txt", args), **get_kwargs("txt", kwargs))
+        return vid, txt

modules/seedvr/src/models/dit_v2/modulation.py

@@ -0,0 +1,118 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Callable, List, Optional
+import torch
+from einops import rearrange
+from torch import nn
+
+from ...common.cache import Cache
+from ...common.distributed.ops import slice_inputs
+
+# (dim: int, emb_dim: int)
+ada_layer_type = Callable[[int, int], nn.Module]
+
+
+def get_ada_layer(ada_layer: str) -> ada_layer_type:
+    if ada_layer == "single":
+        return AdaSingle
+    raise NotImplementedError(f"{ada_layer} is not supported")
+
+
+def expand_dims(x: torch.Tensor, dim: int, ndim: int):
+    """
+    Expand tensor "x" to "ndim" by adding empty dims at "dim".
+    Example: x is (b d), target ndim is 5, add dim at 1, return (b 1 1 1 d).
+    """
+    shape = x.shape
+    shape = shape[:dim] + (1,) * (ndim - len(shape)) + shape[dim:]
+    return x.reshape(shape)
+
+
+class AdaSingle(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        emb_dim: int,
+        layers: List[str],
+        modes: List[str] = ["in", "out"],
+    ):
+        assert emb_dim == 6 * dim, "AdaSingle requires emb_dim == 6 * dim"
+        super().__init__()
+        self.dim = dim
+        self.emb_dim = emb_dim
+        self.layers = layers
+        for l in layers:
+            if "in" in modes:
+                self.register_parameter(f"{l}_shift", nn.Parameter(torch.randn(dim) / dim**0.5))
+                self.register_parameter(
+                    f"{l}_scale", nn.Parameter(torch.randn(dim) / dim**0.5 + 1)
+                )
+            if "out" in modes:
+                self.register_parameter(f"{l}_gate", nn.Parameter(torch.randn(dim) / dim**0.5))
+
+    def forward(
+        self,
+        hid: torch.FloatTensor,  # b ... c
+        emb: torch.FloatTensor,  # b d
+        layer: str,
+        mode: str,
+        cache: Cache = Cache(disable=True),
+        branch_tag: str = "",
+        hid_len: Optional[torch.LongTensor] = None,  # b
+    ) -> torch.FloatTensor:
+        idx = self.layers.index(layer)
+        emb = rearrange(emb, "b (d l g) -> b d l g", l=len(self.layers), g=3)[..., idx, :]
+        emb = expand_dims(emb, 1, hid.ndim + 1)
+
+        if hid_len is not None:
+            emb = cache(
+                f"emb_repeat_{idx}_{branch_tag}",
+                lambda: slice_inputs(
+                    torch.cat([e.repeat(l, *([1] * e.ndim)) for e, l in zip(emb, hid_len)]),
+                    dim=0,
+                ),
+            )
+
+        shiftA, scaleA, gateA = emb.unbind(-1)
+        shiftB, scaleB, gateB = (
+            getattr(self, f"{layer}_shift", None),
+            getattr(self, f"{layer}_scale", None),
+            getattr(self, f"{layer}_gate", None),
+        )
+
+        # 🚀 FP8 COMPATIBILITY: Convert parameters to match embedding dtype
+        # This prevents "Promotion for Float8 Types is not supported" errors
+        target_dtype = shiftA.dtype
+
+        if mode == "in":
+            # Convert parameters to match embedding dtype for FP8 compatibility
+            if scaleB is not None and scaleB.dtype != target_dtype:
+                scaleB = scaleB.to(target_dtype)
+            if shiftB is not None and shiftB.dtype != target_dtype:
+                shiftB = shiftB.to(target_dtype)
+
+            return hid.mul_(scaleA + scaleB).add_(shiftA + shiftB)
+
+        if mode == "out":
+            # Convert gate parameter to match embedding dtype for FP8 compatibility
+            if gateB is not None and gateB.dtype != target_dtype:
+                gateB = gateB.to(target_dtype)
+
+            return hid.mul_(gateA + gateB)
+
+        raise NotImplementedError
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, emb_dim={self.emb_dim}, layers={self.layers}"

modules/seedvr/src/models/dit_v2/na.py

@@ -0,0 +1,241 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from itertools import chain
+from typing import Callable, Dict, List, Tuple
+import einops
+import torch
+
+
+def flatten(
+    hid: List[torch.FloatTensor],  # List of (*** c)
+) -> Tuple[
+    torch.FloatTensor,  # (L c)
+    torch.LongTensor,  # (b n)
+]:
+    assert len(hid) > 0
+    shape = torch.stack([torch.tensor(x.shape[:-1], device=hid[0].device) for x in hid])
+    hid = torch.cat([x.flatten(0, -2) for x in hid])
+    return hid, shape
+
+
+def unflatten(
+    hid: torch.FloatTensor,  # (L c) or (L ... c)
+    hid_shape: torch.LongTensor,  # (b n)
+) -> List[torch.Tensor]:  # List of (*** c) or (*** ... c)
+    hid_len = hid_shape.prod(-1)
+    hid = hid.split(hid_len.tolist())
+    hid = [x.unflatten(0, s.tolist()) for x, s in zip(hid, hid_shape)]
+    return hid
+
+
+def concat(
+    vid: torch.FloatTensor,  # (VL ... c)
+    txt: torch.FloatTensor,  # (TL ... c)
+    vid_len: torch.LongTensor,  # (b)
+    txt_len: torch.LongTensor,  # (b)
+) -> torch.FloatTensor:  # (L ... c)
+    vid = torch.split(vid, vid_len.tolist())
+    txt = torch.split(txt, txt_len.tolist())
+    return torch.cat(list(chain(*zip(vid, txt))))
+
+
+def concat_idx(
+    vid_len: torch.LongTensor,  # (b)
+    txt_len: torch.LongTensor,  # (b)
+) -> Tuple[
+    Callable,
+    Callable,
+]:
+    device = vid_len.device
+    vid_idx = torch.arange(vid_len.sum(), device=device)
+    txt_idx = torch.arange(len(vid_idx), len(vid_idx) + txt_len.sum(), device=device)
+    tgt_idx = concat(vid_idx, txt_idx, vid_len, txt_len)
+    src_idx = torch.argsort(tgt_idx)
+    return (
+        lambda vid, txt: torch.index_select(torch.cat([vid, txt]), 0, tgt_idx),
+        lambda all: torch.index_select(all, 0, src_idx).split([len(vid_idx), len(txt_idx)]),
+    )
+
+
+def unconcat(
+    all: torch.FloatTensor,  # (L ... c)
+    vid_len: torch.LongTensor,  # (b)
+    txt_len: torch.LongTensor,  # (b)
+) -> Tuple[
+    torch.FloatTensor,  # (VL ... c)
+    torch.FloatTensor,  # (TL ... c)
+]:
+    interleave_len = list(chain(*zip(vid_len.tolist(), txt_len.tolist())))
+    all = all.split(interleave_len)
+    vid = torch.cat(all[0::2])
+    txt = torch.cat(all[1::2])
+    return vid, txt
+
+
+def repeat_concat(
+    vid: torch.FloatTensor,  # (VL ... c)
+    txt: torch.FloatTensor,  # (TL ... c)
+    vid_len: torch.LongTensor,  # (n*b)
+    txt_len: torch.LongTensor,  # (b)
+    txt_repeat: List,  # (n)
+) -> torch.FloatTensor:  # (L ... c)
+    vid = torch.split(vid, vid_len.tolist())
+    txt = torch.split(txt, txt_len.tolist())
+    txt = [[x] * n for x, n in zip(txt, txt_repeat)]
+    txt = list(chain(*txt))
+    return torch.cat(list(chain(*zip(vid, txt))))
+
+
+def repeat_concat_idx(
+    vid_len: torch.LongTensor,  # (n*b)
+    txt_len: torch.LongTensor,  # (b)
+    txt_repeat: torch.LongTensor,  # (n)
+) -> Tuple[
+    Callable,
+    Callable,
+]:
+    device = vid_len.device
+    vid_idx = torch.arange(vid_len.sum(), device=device)
+    txt_idx = torch.arange(len(vid_idx), len(vid_idx) + txt_len.sum(), device=device)
+    txt_repeat_list = txt_repeat.tolist()
+    tgt_idx = repeat_concat(vid_idx, txt_idx, vid_len, txt_len, txt_repeat)
+    src_idx = torch.argsort(tgt_idx)
+    txt_idx_len = len(tgt_idx) - len(vid_idx)
+    repeat_txt_len = (txt_len * txt_repeat).tolist()
+
+    def unconcat_coalesce(all):
+        """
+        Un-concat vid & txt, and coalesce the repeated txt.
+        e.g. vid [0 1 2 3 4 5 6 7 8] -> 3 splits -> [0 1 2] [3 4 5] [6 7 8]
+             txt [9 10]
+             repeat_concat ==> [0 1 2 9 10 3 4 5 9 10 6 7 8 9 10]
+             1. argsort re-index ==> [0 1 2 3 4 5 6 7 8 9 9 9 10 10 10]
+                           split ==> vid_out [0 1 2 3 4 5 6 7 8] txt_out [9 9 9 10 10 10]
+             2. reshape & mean for each sample to coalesce the repeated txt.
+        """
+        vid_out, txt_out = all[src_idx].split([len(vid_idx), txt_idx_len])
+        txt_out_coalesced = []
+        for txt, repeat_time in zip(txt_out.split(repeat_txt_len), txt_repeat_list):
+            txt = txt.reshape(-1, repeat_time, *txt.shape[1:]).mean(1)
+            txt_out_coalesced.append(txt)
+        return vid_out, torch.cat(txt_out_coalesced)
+
+    # Note: Backward of torch.index_select is non-deterministic when existing repeated index,
+    # the difference may cumulative like torch.repeat_interleave, so we use vanilla index here.
+    return (
+        lambda vid, txt: torch.cat([vid, txt])[tgt_idx],
+        lambda all: unconcat_coalesce(all),
+    )
+
+
+def rearrange(
+    hid: torch.FloatTensor,  # (L c)
+    hid_shape: torch.LongTensor,  # (b n)
+    pattern: str,
+    **kwargs: Dict[str, int],
+) -> Tuple[
+    torch.FloatTensor,
+    torch.LongTensor,
+]:
+    return flatten([einops.rearrange(h, pattern, **kwargs) for h in unflatten(hid, hid_shape)])
+
+
+def rearrange_idx(
+    hid_shape: torch.LongTensor,  # (b n)
+    pattern: str,
+    **kwargs: Dict[str, int],
+) -> Tuple[Callable, Callable, torch.LongTensor]:
+    hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(-1)
+    tgt_idx, tgt_shape = rearrange(hid_idx, hid_shape, pattern, **kwargs)
+    tgt_idx = tgt_idx.squeeze(-1)
+    src_idx = torch.argsort(tgt_idx)
+    return (
+        lambda hid: torch.index_select(hid, 0, tgt_idx),
+        lambda hid: torch.index_select(hid, 0, src_idx),
+        tgt_shape,
+    )
+
+
+def repeat(
+    hid: torch.FloatTensor,  # (L c)
+    hid_shape: torch.LongTensor,  # (b n)
+    pattern: str,
+    **kwargs: Dict[str, torch.LongTensor],  # (b)
+) -> Tuple[
+    torch.FloatTensor,
+    torch.LongTensor,
+]:
+    hid = unflatten(hid, hid_shape)
+    kwargs = [{k: v[i].item() for k, v in kwargs.items()} for i in range(len(hid))]
+    return flatten([einops.repeat(h, pattern, **a) for h, a in zip(hid, kwargs)])
+
+
+def pack(
+    samples: List[torch.Tensor],  # List of (h w c).
+) -> Tuple[
+    List[torch.Tensor],  # groups [(b1 h1 w1 c1), (b2 h2 w2 c2)]
+    List[List[int]],  # reversal indices.
+]:
+    batches = {}
+    indices = {}
+    for i, sample in enumerate(samples):
+        shape = sample.shape
+        batches[shape] = batches.get(shape, [])
+        indices[shape] = indices.get(shape, [])
+        batches[shape].append(sample)
+        indices[shape].append(i)
+
+    batches = list(map(torch.stack, batches.values()))
+    indices = list(indices.values())
+    return batches, indices
+
+
+def unpack(
+    batches: List[torch.Tensor],
+    indices: List[List[int]],
+) -> List[torch.Tensor]:
+    samples = [None] * (max(chain(*indices)) + 1)
+    for batch, index in zip(batches, indices):
+        for sample, i in zip(batch.unbind(), index):
+            samples[i] = sample
+    return samples
+
+
+def window(
+    hid: torch.FloatTensor,  # (L c)
+    hid_shape: torch.LongTensor,  # (b n)
+    window_fn: Callable[[torch.Tensor], List[torch.Tensor]],
+):
+    hid = unflatten(hid, hid_shape)
+    hid = list(map(window_fn, hid))
+    hid_windows = torch.tensor(list(map(len, hid)), device=hid_shape.device)
+    hid, hid_shape = flatten(list(chain(*hid)))
+    return hid, hid_shape, hid_windows
+
+
+def window_idx(
+    hid_shape: torch.LongTensor,  # (b n)
+    window_fn: Callable[[torch.Tensor], List[torch.Tensor]],
+):
+    hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(-1)
+    tgt_idx, tgt_shape, tgt_windows = window(hid_idx, hid_shape, window_fn)
+    tgt_idx = tgt_idx.squeeze(-1)
+    src_idx = torch.argsort(tgt_idx)
+    return (
+        lambda hid: torch.index_select(hid, 0, tgt_idx),
+        lambda hid: torch.index_select(hid, 0, src_idx),
+        tgt_shape,
+        tgt_windows,
+    )

modules/seedvr/src/models/dit_v2/nablocks/__init__.py

@@ -0,0 +1,26 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from .mmsr_block import NaMMSRTransformerBlock
+
+
+nadit_blocks = {
+    "mmdit_sr": NaMMSRTransformerBlock,
+}
+
+
+def get_nablock(block_type: str):
+    if block_type in nadit_blocks:
+        return nadit_blocks[block_type]
+    raise NotImplementedError(f"{block_type} is not supported")

modules/seedvr/src/models/dit_v2/nablocks/attention/__init__.py

@@ -0,0 +1,25 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from .mmattn import NaMMAttention
+
+attns = {
+    "mm_full": NaMMAttention,
+}
+
+
+def get_attn(attn_type: str):
+    if attn_type in attns:
+        return attns[attn_type]
+    raise NotImplementedError(f"{attn_type} is not supported")

modules/seedvr/src/models/dit_v2/nablocks/attention/mmattn.py

@@ -0,0 +1,267 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Optional, Tuple, Union
+import torch
+from einops import rearrange
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.modules.utils import _triple
+
+from .....common.cache import Cache
+from .....common.distributed.ops import gather_heads_scatter_seq, gather_seq_scatter_heads_qkv
+from .....common.half_precision_fixes import safe_pad_operation
+
+from ... import na
+from ...attention import FlashAttentionVarlen
+from ...mm import MMArg, MMModule
+from ...normalization import norm_layer_type
+from ...rope import get_na_rope
+from ...window import get_window_op
+from itertools import chain
+
+
+class NaMMAttention(nn.Module):
+    def __init__(
+        self,
+        vid_dim: int,
+        txt_dim: int,
+        heads: int,
+        head_dim: int,
+        qk_bias: bool,
+        qk_norm: norm_layer_type,
+        qk_norm_eps: float,
+        rope_type: Optional[str],
+        rope_dim: int,
+        shared_weights: bool,
+        **kwargs,
+    ):
+        super().__init__()
+        dim = MMArg(vid_dim, txt_dim)
+        inner_dim = heads * head_dim
+        qkv_dim = inner_dim * 3
+        self.head_dim = head_dim
+        self.proj_qkv = MMModule(
+            nn.Linear, dim, qkv_dim, bias=qk_bias, shared_weights=shared_weights
+        )
+        self.proj_out = MMModule(nn.Linear, inner_dim, dim, shared_weights=shared_weights)
+        self.norm_q = MMModule(
+            qk_norm,
+            dim=head_dim,
+            eps=qk_norm_eps,
+            elementwise_affine=True,
+            shared_weights=shared_weights,
+        )
+        self.norm_k = MMModule(
+            qk_norm,
+            dim=head_dim,
+            eps=qk_norm_eps,
+            elementwise_affine=True,
+            shared_weights=shared_weights,
+        )
+
+        self.rope = get_na_rope(rope_type=rope_type, dim=rope_dim)
+        self.attn = FlashAttentionVarlen()
+
+    def forward(
+        self,
+        vid: torch.FloatTensor,  # l c
+        txt: torch.FloatTensor,  # l c
+        vid_shape: torch.LongTensor,  # b 3
+        txt_shape: torch.LongTensor,  # b 1
+        cache: Cache,
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.FloatTensor,
+    ]:
+        vid_qkv, txt_qkv = self.proj_qkv(vid, txt)
+        vid_qkv = gather_seq_scatter_heads_qkv(
+            vid_qkv,
+            seq_dim=0,
+            qkv_shape=vid_shape,
+            cache=cache.namespace("vid"),
+        )
+        txt_qkv = gather_seq_scatter_heads_qkv(
+            txt_qkv,
+            seq_dim=0,
+            qkv_shape=txt_shape,
+            cache=cache.namespace("txt"),
+        )
+        vid_qkv = rearrange(vid_qkv, "l (o h d) -> l o h d", o=3, d=self.head_dim)
+        txt_qkv = rearrange(txt_qkv, "l (o h d) -> l o h d", o=3, d=self.head_dim)
+
+        vid_q, vid_k, vid_v = vid_qkv.unbind(1)
+        txt_q, txt_k, txt_v = txt_qkv.unbind(1)
+
+        vid_q, txt_q = self.norm_q(vid_q, txt_q)
+        vid_k, txt_k = self.norm_k(vid_k, txt_k)
+
+        if self.rope:
+            if self.rope.mm:
+                vid_q, vid_k, txt_q, txt_k = self.rope(
+                    vid_q, vid_k, vid_shape, txt_q, txt_k, txt_shape, cache
+                )
+            else:
+                vid_q, vid_k = self.rope(vid_q, vid_k, vid_shape, cache)
+
+        vid_len = cache("vid_len", lambda: vid_shape.prod(-1))
+        txt_len = cache("txt_len", lambda: txt_shape.prod(-1))
+        all_len = cache("all_len", lambda: vid_len + txt_len)
+
+        concat, unconcat = cache("mm_pnp", lambda: na.concat_idx(vid_len, txt_len))
+
+        attn = self.attn(
+            q=concat(vid_q, txt_q).bfloat16(),
+            k=concat(vid_k, txt_k).bfloat16(),
+            v=concat(vid_v, txt_v).bfloat16(),
+            cu_seqlens_q=cache("mm_seqlens", lambda: safe_pad_operation(all_len.cumsum(0), (1, 0)).int()),
+            cu_seqlens_k=cache("mm_seqlens", lambda: safe_pad_operation(all_len.cumsum(0), (1, 0)).int()),
+            max_seqlen_q=cache("mm_maxlen", lambda: all_len.max().item()),
+            max_seqlen_k=cache("mm_maxlen", lambda: all_len.max().item()),
+        ).type_as(vid_q)
+
+        attn = rearrange(attn, "l h d -> l (h d)")
+        vid_out, txt_out = unconcat(attn)
+        vid_out = gather_heads_scatter_seq(vid_out, head_dim=1, seq_dim=0)
+        txt_out = gather_heads_scatter_seq(txt_out, head_dim=1, seq_dim=0)
+
+        vid_out, txt_out = self.proj_out(vid_out, txt_out)
+        return vid_out, txt_out
+
+
+class NaSwinAttention(NaMMAttention):
+    def __init__(
+        self,
+        *args,
+        window: Union[int, Tuple[int, int, int]],
+        window_method: str,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.window = _triple(window)
+        self.window_method = window_method
+        assert all(map(lambda v: isinstance(v, int) and v >= 0, self.window))
+
+        self.window_op = get_window_op(window_method)
+
+    def forward(
+        self,
+        vid: torch.FloatTensor,  # l c
+        txt: torch.FloatTensor,  # l c
+        vid_shape: torch.LongTensor,  # b 3
+        txt_shape: torch.LongTensor,  # b 1
+        cache: Cache,
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.FloatTensor,
+    ]:
+
+        vid_qkv, txt_qkv = self.proj_qkv(vid, txt)
+        vid_qkv = gather_seq_scatter_heads_qkv(
+            vid_qkv,
+            seq_dim=0,
+            qkv_shape=vid_shape,
+            cache=cache.namespace("vid"),
+        )
+        txt_qkv = gather_seq_scatter_heads_qkv(
+            txt_qkv,
+            seq_dim=0,
+            qkv_shape=txt_shape,
+            cache=cache.namespace("txt"),
+        )
+
+        # re-org the input seq for window attn
+        cache_win = cache.namespace(f"{self.window_method}_{self.window}_sd3")
+
+        def make_window(x: torch.Tensor):
+            t, h, w, _ = x.shape
+            window_slices = self.window_op((t, h, w), self.window)
+            return [x[st, sh, sw] for (st, sh, sw) in window_slices]
+
+        window_partition, window_reverse, window_shape, window_count = cache_win(
+            "win_transform",
+            lambda: na.window_idx(vid_shape, make_window),
+        )
+        vid_qkv_win = window_partition(vid_qkv)
+
+        vid_qkv_win = rearrange(vid_qkv_win, "l (o h d) -> l o h d", o=3, d=self.head_dim)
+        txt_qkv = rearrange(txt_qkv, "l (o h d) -> l o h d", o=3, d=self.head_dim)
+
+        vid_q, vid_k, vid_v = vid_qkv_win.unbind(1)
+        txt_q, txt_k, txt_v = txt_qkv.unbind(1)
+
+        vid_q, txt_q = self.norm_q(vid_q, txt_q)
+        vid_k, txt_k = self.norm_k(vid_k, txt_k)
+
+        txt_len = cache("txt_len", lambda: txt_shape.prod(-1))
+
+        vid_len_win = cache_win("vid_len", lambda: window_shape.prod(-1))
+        txt_len_win = cache_win("txt_len", lambda: txt_len.repeat_interleave(window_count))
+        all_len_win = cache_win("all_len", lambda: vid_len_win + txt_len_win)
+        concat_win, unconcat_win = cache_win(
+            "mm_pnp", lambda: na.repeat_concat_idx(vid_len_win, txt_len, window_count)
+        )
+
+        # window rope
+        if self.rope:
+            if self.rope.mm:
+                # repeat text q and k for window mmrope
+                _, num_h, _ = txt_q.shape
+                txt_q_repeat = rearrange(txt_q, "l h d -> l (h d)")
+                txt_q_repeat = na.unflatten(txt_q_repeat, txt_shape)
+                txt_q_repeat = [[x] * n for x, n in zip(txt_q_repeat, window_count)]
+                txt_q_repeat = list(chain(*txt_q_repeat))
+                txt_q_repeat, txt_shape_repeat = na.flatten(txt_q_repeat)
+                txt_q_repeat = rearrange(txt_q_repeat, "l (h d) -> l h d", h=num_h)
+
+                txt_k_repeat = rearrange(txt_k, "l h d -> l (h d)")
+                txt_k_repeat = na.unflatten(txt_k_repeat, txt_shape)
+                txt_k_repeat = [[x] * n for x, n in zip(txt_k_repeat, window_count)]
+                txt_k_repeat = list(chain(*txt_k_repeat))
+                txt_k_repeat, _ = na.flatten(txt_k_repeat)
+                txt_k_repeat = rearrange(txt_k_repeat, "l (h d) -> l h d", h=num_h)
+
+                vid_q, vid_k, txt_q, txt_k = self.rope(
+                    vid_q, vid_k, window_shape, txt_q_repeat, txt_k_repeat, txt_shape_repeat, cache_win
+                )
+            else:
+                vid_q, vid_k = self.rope(vid_q, vid_k, window_shape, cache_win)
+
+        out = self.attn(
+            q=concat_win(vid_q, txt_q).bfloat16(),
+            k=concat_win(vid_k, txt_k).bfloat16(),
+            v=concat_win(vid_v, txt_v).bfloat16(),
+            cu_seqlens_q=cache_win(
+                "vid_seqlens_q", lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int()
+            ),
+            cu_seqlens_k=cache_win(
+                "vid_seqlens_k", lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int()
+            ),
+            max_seqlen_q=cache_win("vid_max_seqlen_q", lambda: all_len_win.max().item()),
+            max_seqlen_k=cache_win("vid_max_seqlen_k", lambda: all_len_win.max().item()),
+        ).type_as(vid_q)
+
+        # text pooling
+        vid_out, txt_out = unconcat_win(out)
+
+        vid_out = rearrange(vid_out, "l h d -> l (h d)")
+        txt_out = rearrange(txt_out, "l h d -> l (h d)")
+        vid_out = window_reverse(vid_out)
+
+        vid_out = gather_heads_scatter_seq(vid_out, head_dim=1, seq_dim=0)
+        txt_out = gather_heads_scatter_seq(txt_out, head_dim=1, seq_dim=0)
+
+        vid_out, txt_out = self.proj_out(vid_out, txt_out)
+
+        return vid_out, txt_out

modules/seedvr/src/models/dit_v2/nablocks/mmsr_block.py

@@ -0,0 +1,126 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Tuple
+import torch
+import torch.nn as nn
+
+# from ..cache import Cache
+from ....common.cache import Cache
+
+from .attention.mmattn import NaSwinAttention
+from ..mm import MMArg
+from ..modulation import ada_layer_type
+from ..normalization import norm_layer_type
+from ..mm import MMModule
+from ..mlp import get_mlp
+
+
+class NaMMSRTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        vid_dim: int,
+        txt_dim: int,
+        emb_dim: int,
+        heads: int,
+        head_dim: int,
+        expand_ratio: int,
+        norm: norm_layer_type,
+        norm_eps: float,
+        ada: ada_layer_type,
+        qk_bias: bool,
+        qk_norm: norm_layer_type,
+        mlp_type: str,
+        shared_weights: bool,
+        rope_type: str,
+        rope_dim: int,
+        is_last_layer: bool,
+        **kwargs,
+    ):
+        super().__init__()
+        dim = MMArg(vid_dim, txt_dim)
+        self.attn_norm = MMModule(norm, dim=dim, eps=norm_eps, elementwise_affine=False, shared_weights=shared_weights,)
+
+        self.attn = NaSwinAttention(
+            vid_dim=vid_dim,
+            txt_dim=txt_dim,
+            heads=heads,
+            head_dim=head_dim,
+            qk_bias=qk_bias,
+            qk_norm=qk_norm,
+            qk_norm_eps=norm_eps,
+            rope_type=rope_type,
+            rope_dim=rope_dim,
+            shared_weights=shared_weights,
+            window=kwargs.pop("window", None),
+            window_method=kwargs.pop("window_method", None),
+        )
+
+        self.mlp_norm = MMModule(norm, dim=dim, eps=norm_eps, elementwise_affine=False, shared_weights=shared_weights, vid_only=is_last_layer)
+        self.mlp = MMModule(
+            get_mlp(mlp_type),
+            dim=dim,
+            expand_ratio=expand_ratio,
+            shared_weights=shared_weights,
+            vid_only=is_last_layer
+        )
+        self.ada = MMModule(ada, dim=dim, emb_dim=emb_dim, layers=["attn", "mlp"], shared_weights=shared_weights, vid_only=is_last_layer)
+        self.is_last_layer = is_last_layer
+
+    def forward(
+        self,
+        vid: torch.FloatTensor,  # l c
+        txt: torch.FloatTensor,  # l c
+        vid_shape: torch.LongTensor,  # b 3
+        txt_shape: torch.LongTensor,  # b 1
+        emb: torch.FloatTensor,
+        cache: Cache,
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.FloatTensor,
+        torch.LongTensor,
+        torch.LongTensor,
+    ]:
+        hid_len = MMArg(
+            cache("vid_len", lambda: vid_shape.prod(-1)),
+            cache("txt_len", lambda: txt_shape.prod(-1)),
+        )
+        ada_kwargs = {
+            "emb": emb,
+            "hid_len": hid_len,
+            "cache": cache,
+            "branch_tag": MMArg("vid", "txt"),
+        }
+
+        vid_attn, txt_attn = self.attn_norm(vid, txt)
+
+        vid_attn, txt_attn = self.ada(vid_attn, txt_attn, layer="attn", mode="in", **ada_kwargs)
+        vid_attn, txt_attn = self.attn(vid_attn, txt_attn, vid_shape, txt_shape, cache)
+        vid_attn, txt_attn = self.ada(vid_attn, txt_attn, layer="attn", mode="out", **ada_kwargs)
+        vid_attn, txt_attn = (vid_attn + vid), (txt_attn + txt)
+
+        vid_mlp, txt_mlp = self.mlp_norm(vid_attn, txt_attn)
+        # ADD BY NUMZ
+        if vid_mlp.dtype != vid_attn.dtype:
+            vid_mlp = vid_mlp.to(vid_attn.dtype)
+        if txt_mlp.dtype != txt_attn.dtype:
+            txt_mlp = txt_mlp.to(txt_attn.dtype)
+        # END BY NUMZ
+        vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, layer="mlp", mode="in", **ada_kwargs)
+        vid_mlp, txt_mlp = self.mlp(vid_mlp, txt_mlp)
+        vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, layer="mlp", mode="out", **ada_kwargs)
+        vid_mlp, txt_mlp = (vid_mlp + vid_attn), (txt_mlp + txt_attn)
+
+        return vid_mlp, txt_mlp, vid_shape, txt_shape

modules/seedvr/src/models/dit_v2/nadit.py

@@ -0,0 +1,246 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union, Callable
+import torch
+from torch import nn
+
+from ...common.cache import Cache
+from ...common.distributed.ops import slice_inputs
+
+from . import na
+from .embedding import TimeEmbedding
+from .modulation import get_ada_layer
+from .nablocks import get_nablock
+from .normalization import get_norm_layer
+from .patch import get_na_patch_layers
+
+# Fake func, no checkpointing is required for inference
+def gradient_checkpointing(module: Union[Callable, nn.Module], *args, enabled: bool, **kwargs):
+    return module(*args, **kwargs)
+
+@dataclass
+class NaDiTOutput:
+    vid_sample: torch.Tensor
+
+
+class NaDiT(nn.Module):
+    """
+    Native Resolution Diffusion Transformer (NaDiT)
+    """
+
+    gradient_checkpointing = False
+
+    def __init__(
+        self,
+        vid_in_channels: int,
+        vid_out_channels: int,
+        vid_dim: int,
+        txt_in_dim: Union[int, List[int]],
+        txt_dim: Optional[int],
+        emb_dim: int,
+        heads: int,
+        head_dim: int,
+        expand_ratio: int,
+        norm: Optional[str],
+        norm_eps: float,
+        ada: str,
+        qk_bias: bool,
+        qk_norm: Optional[str],
+        patch_size: Union[int, Tuple[int, int, int]],
+        num_layers: int,
+        block_type: Union[str, Tuple[str]],
+        mm_layers: Union[int, Tuple[bool]],
+        mlp_type: str = "normal",
+        patch_type: str = "v1",
+        rope_type: Optional[str] = "rope3d",
+        rope_dim: Optional[int] = None,
+        window: Optional[Tuple] = None,
+        window_method: Optional[Tuple[str]] = None,
+        msa_type: Optional[Tuple[str]] = None,
+        mca_type: Optional[Tuple[str]] = None,
+        txt_in_norm: Optional[str] = None,
+        txt_in_norm_scale_factor: int = 0.01,
+        txt_proj_type: Optional[str] = "linear",
+        vid_out_norm: Optional[str] = None,
+        **kwargs,
+    ):
+        ada = get_ada_layer(ada)
+        norm = get_norm_layer(norm)
+        qk_norm = get_norm_layer(qk_norm)
+        rope_dim = rope_dim if rope_dim is not None else head_dim // 2
+        if isinstance(block_type, str):
+            block_type = [block_type] * num_layers
+        elif len(block_type) != num_layers:
+            raise ValueError("The ``block_type`` list should equal to ``num_layers``.")
+        super().__init__()
+        NaPatchIn, NaPatchOut = get_na_patch_layers(patch_type)
+        self.vid_in = NaPatchIn(
+            in_channels=vid_in_channels,
+            patch_size=patch_size,
+            dim=vid_dim,
+        )
+        if not isinstance(txt_in_dim, int):
+            self.txt_in = nn.ModuleList([])
+            for in_dim in txt_in_dim:
+                txt_norm_layer = get_norm_layer(txt_in_norm)(txt_dim, norm_eps, True)
+                if txt_proj_type == "linear":
+                    txt_proj_layer = nn.Linear(in_dim, txt_dim)
+                else:
+                    txt_proj_layer = nn.Sequential(
+                        nn.Linear(in_dim, in_dim), nn.GELU("tanh"), nn.Linear(in_dim, txt_dim)
+                    )
+                torch.nn.init.constant_(txt_norm_layer.weight, txt_in_norm_scale_factor)
+                self.txt_in.append(
+                    nn.Sequential(
+                        txt_proj_layer,
+                        txt_norm_layer,
+                    )
+                )
+        else:
+            self.txt_in = (
+                nn.Linear(txt_in_dim, txt_dim)
+                if txt_in_dim and txt_in_dim != txt_dim
+                else nn.Identity()
+            )
+        self.emb_in = TimeEmbedding(
+            sinusoidal_dim=256,
+            hidden_dim=max(vid_dim, txt_dim),
+            output_dim=emb_dim,
+        )
+
+        if window is None or isinstance(window[0], int):
+            window = [window] * num_layers
+        if window_method is None or isinstance(window_method, str):
+            window_method = [window_method] * num_layers
+
+        if msa_type is None or isinstance(msa_type, str):
+            msa_type = [msa_type] * num_layers
+        if mca_type is None or isinstance(mca_type, str):
+            mca_type = [mca_type] * num_layers
+
+        self.blocks = nn.ModuleList(
+            [
+                get_nablock(block_type[i])(
+                    vid_dim=vid_dim,
+                    txt_dim=txt_dim,
+                    emb_dim=emb_dim,
+                    heads=heads,
+                    head_dim=head_dim,
+                    expand_ratio=expand_ratio,
+                    norm=norm,
+                    norm_eps=norm_eps,
+                    ada=ada,
+                    qk_bias=qk_bias,
+                    qk_norm=qk_norm,
+                    shared_weights=not (
+                        (i < mm_layers) if isinstance(mm_layers, int) else mm_layers[i]
+                    ),
+                    mlp_type=mlp_type,
+                    window=window[i],
+                    window_method=window_method[i],
+                    msa_type=msa_type[i],
+                    mca_type=mca_type[i],
+                    rope_type=rope_type,
+                    rope_dim=rope_dim,
+                    is_last_layer=(i == num_layers - 1),
+                    **kwargs,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        self.vid_out_norm = None
+        if vid_out_norm is not None:
+            self.vid_out_norm = get_norm_layer(vid_out_norm)(
+                dim=vid_dim,
+                eps=norm_eps,
+                elementwise_affine=True,
+            )
+            self.vid_out_ada = ada(
+                dim=vid_dim,
+                emb_dim=emb_dim,
+                layers=["out"],
+                modes=["in"],
+            )
+
+        self.vid_out = NaPatchOut(
+            out_channels=vid_out_channels,
+            patch_size=patch_size,
+            dim=vid_dim,
+        )
+
+    def set_gradient_checkpointing(self, enable: bool):
+        self.gradient_checkpointing = enable
+
+    def forward(
+        self,
+        vid: torch.FloatTensor,  # l c
+        txt: Union[torch.FloatTensor, List[torch.FloatTensor]],  # l c
+        vid_shape: torch.LongTensor,  # b 3
+        txt_shape: Union[torch.LongTensor, List[torch.LongTensor]],  # b 1
+        timestep: Union[int, float, torch.IntTensor, torch.FloatTensor],  # b
+        disable_cache: bool = False,  # for test
+    ):
+        cache = Cache(disable=disable_cache)
+
+        # slice vid after patching in when using sequence parallelism
+        if isinstance(txt, list):
+            assert isinstance(self.txt_in, nn.ModuleList)
+            txt = [
+                na.unflatten(fc(i), s) for fc, i, s in zip(self.txt_in, txt, txt_shape)
+            ]  # B L D
+            txt, txt_shape = na.flatten([torch.cat(t, dim=0) for t in zip(*txt)])
+            txt = slice_inputs(txt, dim=0)
+        else:
+            txt = slice_inputs(txt, dim=0)
+            txt = self.txt_in(txt)
+
+        # Video input.
+        # Sequence parallel slicing is done inside patching class.
+        vid, vid_shape = self.vid_in(vid, vid_shape, cache)
+
+        # Embedding input.
+        emb = self.emb_in(timestep, device=vid.device, dtype=vid.dtype)
+
+        # Body
+        for i, block in enumerate(self.blocks):
+            vid, txt, vid_shape, txt_shape = gradient_checkpointing(
+                enabled=(self.gradient_checkpointing and self.training),
+                module=block,
+                vid=vid,
+                txt=txt,
+                vid_shape=vid_shape,
+                txt_shape=txt_shape,
+                emb=emb,
+                cache=cache,
+            )
+
+        # Video output norm.
+        if self.vid_out_norm:
+            vid = self.vid_out_norm(vid)
+            vid = self.vid_out_ada(
+                vid,
+                emb=emb,
+                layer="out",
+                mode="in",
+                hid_len=cache("vid_len", lambda: vid_shape.prod(-1)),
+                cache=cache,
+                branch_tag="vid",
+            )
+
+        # Video output.
+        vid, vid_shape = self.vid_out(vid, vid_shape, cache)
+        return NaDiTOutput(vid_sample=vid)

modules/seedvr/src/models/dit_v2/normalization.py

@@ -0,0 +1,147 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Callable, Optional
+from diffusers.models.normalization import RMSNorm
+from torch import nn
+import torch
+import torch.nn.functional as F
+import numbers
+from torch.nn.parameter import Parameter
+from torch.nn import init
+
+# (dim: int, eps: float, elementwise_affine: bool)
+norm_layer_type = Callable[[int, float, bool], nn.Module]
+
+
+class CustomLayerNorm(nn.Module):
+    """
+    Custom LayerNorm implementation to replace Apex FusedLayerNorm
+    """
+    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
+        super(CustomLayerNorm, self).__init__()
+
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        self.normalized_shape = torch.Size(normalized_shape)
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+
+        if self.elementwise_affine:
+            self.weight = Parameter(torch.Tensor(*normalized_shape))
+            self.bias = Parameter(torch.Tensor(*normalized_shape))
+        else:
+            self.register_parameter('weight', None)
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        if self.elementwise_affine:
+            init.ones_(self.weight)
+            init.zeros_(self.bias)
+
+    def forward(self, input):
+        # 🚀 FP8 COMPATIBILITY: Convert parameters to match input dtype
+        # This prevents "Promotion for Float8 Types is not supported" errors
+        weight = self.weight
+        bias = self.bias
+
+        if self.elementwise_affine and weight is not None:
+            if weight.dtype != input.dtype:
+                weight = weight.to(input.dtype)
+            if bias is not None and bias.dtype != input.dtype:
+                bias = bias.to(input.dtype)
+
+        return F.layer_norm(
+            input, self.normalized_shape, weight, bias, self.eps)
+
+
+class CustomRMSNorm(nn.Module):
+    """
+    Custom RMSNorm implementation to replace Apex FusedRMSNorm
+    """
+    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
+        super(CustomRMSNorm, self).__init__()
+
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        self.normalized_shape = torch.Size(normalized_shape)
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+
+        if self.elementwise_affine:
+            self.weight = Parameter(torch.ones(*normalized_shape))
+        else:
+            self.register_parameter('weight', None)
+
+    def forward(self, input):
+        # RMS normalization: x / sqrt(mean(x^2) + eps) * weight
+        dims = tuple(range(-len(self.normalized_shape), 0))
+
+        # Calculate RMS: sqrt(mean(x^2))
+        variance = input.pow(2).mean(dim=dims, keepdim=True)
+        rms = torch.sqrt(variance + self.eps)
+
+        # Normalize
+        normalized = input / rms
+
+        if self.elementwise_affine:
+            # 🚀 FP8 COMPATIBILITY: Convert weight to match normalized dtype
+            # This prevents "Promotion for Float8 Types is not supported" errors
+            weight = self.weight
+            if weight.dtype != normalized.dtype:
+                weight = weight.to(normalized.dtype)
+            return normalized * weight
+        return normalized
+
+
+def get_norm_layer(norm_type: Optional[str]) -> norm_layer_type:
+
+    def _norm_layer(dim: int, eps: float, elementwise_affine: bool):
+        if norm_type is None:
+            return nn.Identity()
+
+        if norm_type == "layer":
+            return nn.LayerNorm(
+                normalized_shape=dim,
+                eps=eps,
+                elementwise_affine=elementwise_affine,
+            )
+
+        if norm_type == "rms":
+            return RMSNorm(
+                dim=dim,
+                eps=eps,
+                elementwise_affine=elementwise_affine,
+            )
+
+        if norm_type == "fusedln":
+            # Use custom LayerNorm instead of Apex FusedLayerNorm
+            return CustomLayerNorm(
+                normalized_shape=dim,
+                elementwise_affine=elementwise_affine,
+                eps=eps,
+            )
+
+        if norm_type == "fusedrms":
+            # Use custom RMSNorm instead of Apex FusedRMSNorm
+            return CustomRMSNorm(
+                normalized_shape=dim,
+                elementwise_affine=elementwise_affine,
+                eps=eps,
+            )
+
+        raise NotImplementedError(f"{norm_type} is not supported")
+
+    return _norm_layer

modules/seedvr/src/models/dit_v2/patch/__init__.py

@@ -0,0 +1,19 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+def get_na_patch_layers(patch_type="v1"):
+    assert patch_type in ["v1"]
+    if patch_type == "v1":
+        from .patch_v1 import NaPatchIn, NaPatchOut
+    return NaPatchIn, NaPatchOut

modules/seedvr/src/models/dit_v2/patch/patch_v1.py

@@ -0,0 +1,127 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Tuple, Union
+import torch
+from einops import rearrange
+from torch import nn
+from torch.nn.modules.utils import _triple
+
+from ....common.cache import Cache
+from ....common.distributed.ops import gather_outputs, slice_inputs
+
+from .. import na
+
+
+class PatchIn(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        patch_size: Union[int, Tuple[int, int, int]],
+        dim: int,
+    ):
+        super().__init__()
+        t, h, w = _triple(patch_size)
+        self.patch_size = t, h, w
+        self.proj = nn.Linear(in_channels * t * h * w, dim)
+
+    def forward(
+        self,
+        vid: torch.Tensor,
+    ) -> torch.Tensor:
+        t, h, w = self.patch_size
+        if t > 1:
+            assert vid.size(2) % t == 1
+            vid = torch.cat([vid[:, :, :1]] * (t - 1) + [vid], dim=2)
+        vid = rearrange(vid, "b c (T t) (H h) (W w) -> b T H W (t h w c)", t=t, h=h, w=w)
+        vid = self.proj(vid)
+        return vid
+
+
+class PatchOut(nn.Module):
+    def __init__(
+        self,
+        out_channels: int,
+        patch_size: Union[int, Tuple[int, int, int]],
+        dim: int,
+    ):
+        super().__init__()
+        t, h, w = _triple(patch_size)
+        self.patch_size = t, h, w
+        self.proj = nn.Linear(dim, out_channels * t * h * w)
+
+    def forward(
+        self,
+        vid: torch.Tensor,
+    ) -> torch.Tensor:
+        t, h, w = self.patch_size
+        vid = self.proj(vid)
+        vid = rearrange(vid, "b T H W (t h w c) -> b c (T t) (H h) (W w)", t=t, h=h, w=w)
+        if t > 1:
+            vid = vid[:, :, (t - 1) :]
+        return vid
+
+
+class NaPatchIn(PatchIn):
+    def forward(
+        self,
+        vid: torch.Tensor,  # l c
+        vid_shape: torch.LongTensor,
+        cache: Cache = Cache(disable=True),  # for test
+    ) -> torch.Tensor:
+        cache = cache.namespace("patch")
+        vid_shape_before_patchify = cache("vid_shape_before_patchify", lambda: vid_shape)
+        t, h, w = self.patch_size
+        if not (t == h == w == 1):
+            vid = na.unflatten(vid, vid_shape)
+            for i in range(len(vid)):
+                if t > 1 and vid_shape_before_patchify[i, 0] % t != 0:
+                    vid[i] = torch.cat([vid[i][:1]] * (t - vid[i].size(0) % t) + [vid[i]], dim=0)
+                vid[i] = rearrange(vid[i], "(T t) (H h) (W w) c -> T H W (t h w c)", t=t, h=h, w=w)
+            vid, vid_shape = na.flatten(vid)
+
+        # slice vid after patching in when using sequence parallelism
+        vid = slice_inputs(vid, dim=0)
+        vid = self.proj(vid)
+        return vid, vid_shape
+
+
+class NaPatchOut(PatchOut):
+    def forward(
+        self,
+        vid: torch.FloatTensor,  # l c
+        vid_shape: torch.LongTensor,
+        cache: Cache = Cache(disable=True),  # for test
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.LongTensor,
+    ]:
+        cache = cache.namespace("patch")
+        vid_shape_before_patchify = cache.get("vid_shape_before_patchify")
+
+        t, h, w = self.patch_size
+        vid = self.proj(vid)
+        # gather vid before patching out when enabling sequence parallelism
+        vid = gather_outputs(
+            vid, gather_dim=0, padding_dim=0, unpad_shape=vid_shape, cache=cache.namespace("vid")
+        )
+        if not (t == h == w == 1):
+            vid = na.unflatten(vid, vid_shape)
+            for i in range(len(vid)):
+                vid[i] = rearrange(vid[i], "T H W (t h w c) -> (T t) (H h) (W w) c", t=t, h=h, w=w)
+                if t > 1 and vid_shape_before_patchify[i, 0] % t != 0:
+                    vid[i] = vid[i][(t - vid_shape_before_patchify[i, 0] % t) :]
+            vid, vid_shape = na.flatten(vid)
+
+        return vid, vid_shape

modules/seedvr/src/models/dit_v2/rope.py

@@ -0,0 +1,155 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from functools import lru_cache
+from typing import Optional, Tuple
+import torch
+from einops import rearrange
+from rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb
+from torch import nn
+
+from src.common.cache import Cache
+
+
+class RotaryEmbeddingBase(nn.Module):
+    def __init__(self, dim: int, rope_dim: int):
+        super().__init__()
+        self.rope = RotaryEmbedding(
+            dim=dim // rope_dim,
+            freqs_for="pixel",
+            max_freq=256,
+        )
+        # 1. Set model.requires_grad_(True) after model creation will make
+        #    the `requires_grad=False` for rope freqs no longer hold.
+        # 2. Even if we don't set requires_grad_(True) explicitly,
+        #    FSDP is not memory efficient when handling fsdp_wrap
+        #    with mixed requires_grad=True/False.
+        # With above consideration, it is easier just remove the freqs
+        # out of nn.Parameters when `learned_freq=False`
+        freqs = self.rope.freqs
+        del self.rope.freqs
+        self.rope.register_buffer("freqs", freqs.data)
+
+    @lru_cache(maxsize=128)
+    def get_axial_freqs(self, *dims):
+        return self.rope.get_axial_freqs(*dims)
+
+
+class RotaryEmbedding3d(RotaryEmbeddingBase):
+    def __init__(self, dim: int):
+        super().__init__(dim, rope_dim=3)
+        self.mm = False
+
+    def forward(
+        self,
+        q: torch.FloatTensor,  # b h l d
+        k: torch.FloatTensor,  # b h l d
+        size: Tuple[int, int, int],
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.FloatTensor,
+    ]:
+        T, H, W = size
+        freqs = self.get_axial_freqs(T, H, W)
+        q = rearrange(q, "b h (T H W) d -> b h T H W d", T=T, H=H, W=W)
+        k = rearrange(k, "b h (T H W) d -> b h T H W d", T=T, H=H, W=W)
+        q = apply_rotary_emb(freqs, q.float()).to(q.dtype)
+        k = apply_rotary_emb(freqs, k.float()).to(k.dtype)
+        q = rearrange(q, "b h T H W d -> b h (T H W) d")
+        k = rearrange(k, "b h T H W d -> b h (T H W) d")
+        return q, k
+
+
+class MMRotaryEmbeddingBase(RotaryEmbeddingBase):
+    def __init__(self, dim: int, rope_dim: int):
+        super().__init__(dim, rope_dim)
+        self.rope = RotaryEmbedding(
+            dim=dim // rope_dim,
+            freqs_for="lang",
+            theta=10000,
+        )
+        freqs = self.rope.freqs
+        del self.rope.freqs
+        self.rope.register_buffer("freqs", freqs.data)
+        self.mm = True
+
+
+class NaMMRotaryEmbedding3d(MMRotaryEmbeddingBase):
+    def __init__(self, dim: int):
+        super().__init__(dim, rope_dim=3)
+
+    def forward(
+        self,
+        vid_q: torch.FloatTensor,  # L h d
+        vid_k: torch.FloatTensor,  # L h d
+        vid_shape: torch.LongTensor,  # B 3
+        txt_q: torch.FloatTensor,  # L h d
+        txt_k: torch.FloatTensor,  # L h d
+        txt_shape: torch.LongTensor,  # B 1
+        cache: Cache,
+    ) -> Tuple[
+        torch.FloatTensor,
+        torch.FloatTensor,
+        torch.FloatTensor,
+        torch.FloatTensor,
+    ]:
+        vid_freqs, txt_freqs = cache(
+            "mmrope_freqs_3d",
+            lambda: self.get_freqs(vid_shape, txt_shape),
+        )
+        target_device = vid_q.device
+        if vid_freqs.device != target_device:
+            vid_freqs = vid_freqs.to(target_device)
+        if txt_freqs.device != target_device:
+            txt_freqs = txt_freqs.to(target_device)
+        vid_q = rearrange(vid_q, "L h d -> h L d")
+        vid_k = rearrange(vid_k, "L h d -> h L d")
+        vid_q = apply_rotary_emb(vid_freqs, vid_q.float()).to(vid_q.dtype)
+        vid_k = apply_rotary_emb(vid_freqs, vid_k.float()).to(vid_k.dtype)
+        vid_q = rearrange(vid_q, "h L d -> L h d")
+        vid_k = rearrange(vid_k, "h L d -> L h d")
+
+        txt_q = rearrange(txt_q, "L h d -> h L d")
+        txt_k = rearrange(txt_k, "L h d -> h L d")
+        txt_q = apply_rotary_emb(txt_freqs, txt_q.float()).to(txt_q.dtype)
+        txt_k = apply_rotary_emb(txt_freqs, txt_k.float()).to(txt_k.dtype)
+        txt_q = rearrange(txt_q, "h L d -> L h d")
+        txt_k = rearrange(txt_k, "h L d -> L h d")
+        return vid_q, vid_k, txt_q, txt_k
+
+    def get_freqs(
+        self,
+        vid_shape: torch.LongTensor,
+        txt_shape: torch.LongTensor,
+    ) -> Tuple[
+        torch.Tensor,
+        torch.Tensor,
+    ]:
+        vid_freqs = self.get_axial_freqs(1024, 128, 128)
+        txt_freqs = self.get_axial_freqs(1024)
+        vid_freq_list, txt_freq_list = [], []
+        for (f, h, w), l in zip(vid_shape.tolist(), txt_shape[:, 0].tolist()):
+            vid_freq = vid_freqs[l : l + f, :h, :w].reshape(-1, vid_freqs.size(-1))
+            txt_freq = txt_freqs[:l].repeat(1, 3).reshape(-1, vid_freqs.size(-1))
+            vid_freq_list.append(vid_freq)
+            txt_freq_list.append(txt_freq)
+        return torch.cat(vid_freq_list, dim=0), torch.cat(txt_freq_list, dim=0)
+
+
+def get_na_rope(rope_type: Optional[str], dim: int):
+    if rope_type is None:
+        return None
+    if rope_type == "mmrope3d":
+        return NaMMRotaryEmbedding3d(dim=dim)
+    raise NotImplementedError(f"{rope_type} is not supported.")

modules/seedvr/src/models/dit_v2/window.py

@@ -0,0 +1,83 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from math import ceil
+from typing import Tuple
+import math
+
+def get_window_op(name: str):
+    if name == "720pwin_by_size_bysize":
+        return make_720Pwindows_bysize
+    if name == "720pswin_by_size_bysize":
+        return make_shifted_720Pwindows_bysize
+    raise ValueError(f"Unknown windowing method: {name}")
+
+
+# -------------------------------- Windowing -------------------------------- #
+def make_720Pwindows_bysize(size: Tuple[int, int, int], num_windows: Tuple[int, int, int]):
+    t, h, w = size
+    resized_nt, resized_nh, resized_nw = num_windows
+    #cal windows under 720p
+    scale = math.sqrt((45 * 80) / (h * w))
+    resized_h, resized_w = round(h * scale), round(w * scale)
+    wh, ww = ceil(resized_h / resized_nh), ceil(resized_w / resized_nw)  # window size.
+    wt = ceil(min(t, 30) / resized_nt)  # window size.
+    nt, nh, nw = ceil(t / wt), ceil(h / wh), ceil(w / ww)  # window size.
+    return [
+        (
+            slice(it * wt, min((it + 1) * wt, t)),
+            slice(ih * wh, min((ih + 1) * wh, h)),
+            slice(iw * ww, min((iw + 1) * ww, w)),
+        )
+        for iw in range(nw)
+        if min((iw + 1) * ww, w) > iw * ww
+        for ih in range(nh)
+        if min((ih + 1) * wh, h) > ih * wh
+        for it in range(nt)
+        if min((it + 1) * wt, t) > it * wt
+    ]
+
+def make_shifted_720Pwindows_bysize(size: Tuple[int, int, int], num_windows: Tuple[int, int, int]):
+    t, h, w = size
+    resized_nt, resized_nh, resized_nw = num_windows
+    #cal windows under 720p
+    scale = math.sqrt((45 * 80) / (h * w))
+    resized_h, resized_w = round(h * scale), round(w * scale)
+    wh, ww = ceil(resized_h / resized_nh), ceil(resized_w / resized_nw)  # window size.
+    wt = ceil(min(t, 30) / resized_nt)  # window size.
+
+    st, sh, sw = (  # shift size.
+        0.5 if wt < t else 0,
+        0.5 if wh < h else 0,
+        0.5 if ww < w else 0,
+    )
+    nt, nh, nw = ceil((t - st) / wt), ceil((h - sh) / wh), ceil((w - sw) / ww)  # window size.
+    nt, nh, nw = (  # number of window.
+        nt + 1 if st > 0 else 1,
+        nh + 1 if sh > 0 else 1,
+        nw + 1 if sw > 0 else 1,
+    )
+    return [
+        (
+            slice(max(int((it - st) * wt), 0), min(int((it - st + 1) * wt), t)),
+            slice(max(int((ih - sh) * wh), 0), min(int((ih - sh + 1) * wh), h)),
+            slice(max(int((iw - sw) * ww), 0), min(int((iw - sw + 1) * ww), w)),
+        )
+        for iw in range(nw)
+        if min(int((iw - sw + 1) * ww), w) > max(int((iw - sw) * ww), 0)
+        for ih in range(nh)
+        if min(int((ih - sh + 1) * wh), h) > max(int((ih - sh) * wh), 0)
+        for it in range(nt)
+        if min(int((it - st + 1) * wt), t) > max(int((it - st) * wt), 0)
+    ]

modules/seedvr/src/models/video_vae_v3/modules/causal_inflation_lib.py

@@ -0,0 +1,437 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+import math
+from contextlib import contextmanager
+from typing import List, Optional, Union
+import torch
+import torch.nn.functional as F
+from diffusers.models.normalization import RMSNorm
+from einops import rearrange
+from torch import Tensor, nn
+from torch.nn import Conv3d
+
+from .context_parallel_lib import cache_send_recv, get_cache_size
+from .global_config import get_norm_limit
+from .types import MemoryState, _inflation_mode_t, _memory_device_t
+from ....common.half_precision_fixes import safe_pad_operation
+
+# Single GPU inference - no distributed processing needed
+
+# Mock distributed functions for single GPU inference
+def get_sequence_parallel_group():
+    return None
+
+def get_sequence_parallel_rank():
+    return 0
+
+def get_sequence_parallel_world_size():
+    return 1
+
+def get_next_sequence_parallel_rank():
+    return 0
+
+def get_prev_sequence_parallel_rank():
+    return 0
+
+
+@contextmanager
+def ignore_padding(model):
+    orig_padding = model.padding
+    model.padding = (0, 0, 0)
+    try:
+        yield
+    finally:
+        model.padding = orig_padding
+
+
+class InflatedCausalConv3d(Conv3d):
+    def __init__(
+        self,
+        *args,
+        inflation_mode: _inflation_mode_t,
+        memory_device: _memory_device_t = "same",
+        **kwargs,
+    ):
+        self.inflation_mode = inflation_mode
+        self.memory = None
+        super().__init__(*args, **kwargs)
+        self.temporal_padding = self.padding[0]
+        self.memory_device = memory_device
+        self.padding = (0, *self.padding[1:])  # Remove temporal pad to keep causal.
+        self.memory_limit = float("inf")
+
+    def set_memory_limit(self, value: float):
+        self.memory_limit = value
+
+    def set_memory_device(self, memory_device: _memory_device_t):
+        self.memory_device = memory_device
+
+    def memory_limit_conv(
+        self,
+        x,
+        *,
+        split_dim=3,
+        padding=(0, 0, 0, 0, 0, 0),
+        prev_cache=None,
+    ):
+        # Compatible with no limit.
+        if math.isinf(self.memory_limit):
+            if prev_cache is not None:
+                x = torch.cat([prev_cache, x], dim=split_dim - 1)
+            return super().forward(x)
+
+        # Compute tensor shape after concat & padding.
+        shape = torch.tensor(x.size())
+        if prev_cache is not None:
+            shape[split_dim - 1] += prev_cache.size(split_dim - 1)
+        shape[-3:] += torch.tensor(padding).view(3, 2).sum(-1).flip(0)
+        memory_occupy = shape.prod() * x.element_size() / 1024**3  # GiB
+        if memory_occupy < self.memory_limit or split_dim == x.ndim:
+            if prev_cache is not None:
+                x = torch.cat([prev_cache, x], dim=split_dim - 1)
+            x = safe_pad_operation(x, padding, mode='constant', value=0.0)
+            with ignore_padding(self):
+                return super().forward(x)
+
+        # Exceed memory limit, splitting tensor
+
+        # Split input (& prev_cache).
+        num_splits = math.ceil(memory_occupy / self.memory_limit)
+        size_per_split = x.size(split_dim) // num_splits
+        split_sizes = [size_per_split] * (num_splits - 1)
+        split_sizes += [x.size(split_dim) - sum(split_sizes)]
+
+        x = list(x.split(split_sizes, dim=split_dim))
+        if prev_cache is not None:
+            prev_cache = list(prev_cache.split(split_sizes, dim=split_dim))
+        # Loop Fwd.
+        cache = None
+        for idx in range(len(x)):
+            # Concat prev cache from last dim
+            if prev_cache is not None:
+                x[idx] = torch.cat([prev_cache[idx], x[idx]], dim=split_dim - 1)
+
+            # Get padding pattern.
+            lpad_dim = (x[idx].ndim - split_dim - 1) * 2
+            rpad_dim = lpad_dim + 1
+            padding = list(padding)
+            padding[lpad_dim] = self.padding[split_dim - 2] if idx == 0 else 0
+            padding[rpad_dim] = self.padding[split_dim - 2] if idx == len(x) - 1 else 0
+            pad_len = padding[lpad_dim] + padding[rpad_dim]
+            padding = tuple(padding)
+
+            # Prepare cache for next slice (this dim).
+            next_cache = None
+            cache_len = cache.size(split_dim) if cache is not None else 0
+            next_catch_size = get_cache_size(
+                conv_module=self,
+                input_len=x[idx].size(split_dim) + cache_len,
+                pad_len=pad_len,
+                dim=split_dim - 2,
+            )
+            if next_catch_size != 0:
+                assert next_catch_size <= x[idx].size(split_dim)
+                next_cache = (
+                    x[idx].transpose(0, split_dim)[-next_catch_size:].transpose(0, split_dim)
+                )
+
+            # Recursive.
+            x[idx] = self.memory_limit_conv(
+                x[idx],
+                split_dim=split_dim + 1,
+                padding=padding,
+                prev_cache=cache,
+            )
+
+            # Update cache.
+            cache = next_cache
+
+        output = torch.cat(x, split_dim)
+        return output
+
+    def forward(
+        self,
+        input: Union[Tensor, List[Tensor]],
+        memory_state: MemoryState = MemoryState.UNSET,
+    ) -> Tensor:
+        assert memory_state != MemoryState.UNSET
+        if memory_state != MemoryState.ACTIVE:
+            self.memory = None
+        if (
+            math.isinf(self.memory_limit)
+            and torch.is_tensor(input)
+            and get_sequence_parallel_group() is None
+        ):
+            return self.basic_forward(input, memory_state)
+        return self.slicing_forward(input, memory_state)
+
+    def basic_forward(self, input: Tensor, memory_state: MemoryState = MemoryState.UNSET):
+        mem_size = self.stride[0] - self.kernel_size[0]
+        if (self.memory is not None) and (memory_state == MemoryState.ACTIVE):
+            input = extend_head(input, memory=self.memory, times=-1)
+        else:
+            input = extend_head(input, times=self.temporal_padding * 2)
+        memory = (
+            input[:, :, mem_size:].detach()
+            if (mem_size != 0 and memory_state != MemoryState.DISABLED)
+            else None
+        )
+        if (
+            memory_state != MemoryState.DISABLED
+            and not self.training
+            and (self.memory_device is not None)
+        ):
+            self.memory = memory
+            if self.memory_device == "cpu" and self.memory is not None:
+                self.memory = self.memory.to("cpu")
+        return super().forward(input)
+
+    def slicing_forward(
+        self,
+        input: Union[Tensor, List[Tensor]],
+        memory_state: MemoryState = MemoryState.UNSET,
+    ) -> Tensor:
+        squeeze_out = False
+        if torch.is_tensor(input):
+            input = [input]
+            squeeze_out = True
+
+        cache_size = self.kernel_size[0] - self.stride[0]
+        cache = cache_send_recv(
+            input, cache_size=cache_size, memory=self.memory, times=self.temporal_padding * 2
+        )
+
+        # Single GPU inference - simplified memory management
+        if (
+            memory_state in [MemoryState.INITIALIZING, MemoryState.ACTIVE]  # use_slicing
+            and not self.training
+            and (self.memory_device is not None)
+            and cache_size != 0
+        ):
+            if cache_size > input[-1].size(2) and cache is not None and len(input) == 1:
+                input[0] = torch.cat([cache, input[0]], dim=2)
+                cache = None
+            if cache_size <= input[-1].size(2):
+                self.memory = input[-1][:, :, -cache_size:].detach().contiguous()
+                if self.memory_device == "cpu" and self.memory is not None:
+                    self.memory = self.memory.to("cpu")
+
+        padding = tuple(x for x in reversed(self.padding) for _ in range(2))
+        for i in range(len(input)):
+            # Prepare cache for next input slice.
+            next_cache = None
+            cache_size = 0
+            if i < len(input) - 1:
+                cache_len = cache.size(2) if cache is not None else 0
+                cache_size = get_cache_size(self, input[i].size(2) + cache_len, pad_len=0)
+            if cache_size != 0:
+                if cache_size > input[i].size(2) and cache is not None:
+                    input[i] = torch.cat([cache, input[i]], dim=2)
+                    cache = None
+                assert cache_size <= input[i].size(2), f"{cache_size} > {input[i].size(2)}"
+                next_cache = input[i][:, :, -cache_size:]
+
+            # Conv forward for this input slice.
+            input[i] = self.memory_limit_conv(
+                input[i],
+                padding=padding,
+                prev_cache=cache,
+            )
+
+            # Update cache.
+            cache = next_cache
+
+        return input[0] if squeeze_out else input
+
+    def tflops(self, args, kwargs, output) -> float:
+        if torch.is_tensor(output):
+            output_numel = output.numel()
+        elif isinstance(output, list):
+            output_numel = sum(o.numel() for o in output)
+        else:
+            raise NotImplementedError
+        return (2 * math.prod(self.kernel_size) * self.in_channels * (output_numel / 1e6)) / 1e6
+
+    def _load_from_state_dict(
+        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+    ):
+        if self.inflation_mode != "none":
+            state_dict = modify_state_dict(
+                self,
+                state_dict,
+                prefix,
+                inflate_weight_fn=inflate_weight,
+                inflate_bias_fn=inflate_bias,
+            )
+        super()._load_from_state_dict(
+            state_dict,
+            prefix,
+            local_metadata,
+            (strict and self.inflation_mode == "none"),
+            missing_keys,
+            unexpected_keys,
+            error_msgs,
+        )
+
+
+def init_causal_conv3d(
+    *args,
+    inflation_mode: _inflation_mode_t,
+    **kwargs,
+):
+    """
+    Initialize a Causal-3D convolution layer.
+    Parameters:
+        inflation_mode: Listed as below. It's compatible with all the 3D-VAE checkpoints we have.
+            - none: No inflation will be conducted.
+                    The loading logic of state dict will fall back to default.
+            - tail / replicate: Refer to the definition of `InflatedCausalConv3d`.
+    """
+    return InflatedCausalConv3d(*args, inflation_mode=inflation_mode, **kwargs)
+
+
+def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor:
+    input_dtype = x.dtype
+    if isinstance(norm_layer, (nn.LayerNorm, RMSNorm)):
+        if x.ndim == 4:
+            x = rearrange(x, "b c h w -> b h w c")
+            x = norm_layer(x)
+            x = rearrange(x, "b h w c -> b c h w")
+            return x.to(input_dtype)
+        if x.ndim == 5:
+            x = rearrange(x, "b c t h w -> b t h w c")
+            x = norm_layer(x)
+            x = rearrange(x, "b t h w c -> b c t h w")
+            return x.to(input_dtype)
+    if isinstance(norm_layer, (nn.GroupNorm, nn.BatchNorm2d, nn.SyncBatchNorm)):
+        if x.ndim <= 4:
+            return norm_layer(x).to(input_dtype)
+        if x.ndim == 5:
+            t = x.size(2)
+            x = rearrange(x, "b c t h w -> (b t) c h w")
+            memory_occupy = x.numel() * x.element_size() / 1024**3
+            if isinstance(norm_layer, nn.GroupNorm) and memory_occupy > get_norm_limit():
+                num_chunks = min(4 if x.element_size() == 2 else 2, norm_layer.num_groups)
+                assert norm_layer.num_groups % num_chunks == 0
+                num_groups_per_chunk = norm_layer.num_groups // num_chunks
+
+                x = list(x.chunk(num_chunks, dim=1))
+                weights = norm_layer.weight.chunk(num_chunks, dim=0)
+                biases = norm_layer.bias.chunk(num_chunks, dim=0)
+                for i, (w, b) in enumerate(zip(weights, biases)):
+                    x[i] = F.group_norm(x[i], num_groups_per_chunk, w, b, norm_layer.eps)
+                    x[i] = x[i].to(input_dtype)
+                # ADD BY NUMZ
+                # ADD BY NUMZ
+                x = torch.cat(x, dim=1)
+            else:
+                x = norm_layer(x)
+            x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+            return x.to(input_dtype)
+    raise NotImplementedError
+
+
+def remove_head(tensor: Tensor, times: int = 1) -> Tensor:
+    """
+    Remove duplicated first frame features in the up-sampling process.
+    """
+    # Single GPU inference - always process
+    if times == 0:
+        return tensor
+    return torch.cat(tensors=(tensor[:, :, :1], tensor[:, :, times + 1 :]), dim=2)
+
+
+def extend_head(tensor: Tensor, times: int = 2, memory: Optional[Tensor] = None) -> Tensor:
+    """
+    When memory is None:
+        - Duplicate first frame features in the down-sampling process.
+    When memory is not None:
+        - Concatenate memory features with the input features to keep temporal consistency.
+    """
+    if memory is not None:
+        return torch.cat((memory.to(tensor), tensor), dim=2)
+    assert times >= 0, "Invalid input for function 'extend_head'!"
+    if times == 0:
+        return tensor
+    else:
+        tile_repeat = [1] * tensor.ndim
+        tile_repeat[2] = times
+        return torch.cat(tensors=(torch.tile(tensor[:, :, :1], tile_repeat), tensor), dim=2)
+
+
+def inflate_weight(weight_2d: torch.Tensor, weight_3d: torch.Tensor, inflation_mode: str):
+    """
+    Inflate a 2D convolution weight matrix to a 3D one.
+    Parameters:
+        weight_2d:      The weight matrix of 2D conv to be inflated.
+        weight_3d:      The weight matrix of 3D conv to be initialized.
+        inflation_mode: the mode of inflation
+    """
+    assert inflation_mode in ["tail", "replicate"]
+    assert weight_3d.shape[:2] == weight_2d.shape[:2]
+    with torch.no_grad():
+        if inflation_mode == "replicate":
+            depth = weight_3d.size(2)
+            weight_3d.copy_(weight_2d.unsqueeze(2).repeat(1, 1, depth, 1, 1) / depth)
+        else:
+            weight_3d.fill_(0.0)
+            weight_3d[:, :, -1].copy_(weight_2d)
+    return weight_3d
+
+
+def inflate_bias(bias_2d: torch.Tensor, bias_3d: torch.Tensor, inflation_mode: str):
+    """
+    Inflate a 2D convolution bias tensor to a 3D one
+    Parameters:
+        bias_2d:        The bias tensor of 2D conv to be inflated.
+        bias_3d:        The bias tensor of 3D conv to be initialized.
+        inflation_mode: Placeholder to align `inflate_weight`.
+    """
+    assert bias_3d.shape == bias_2d.shape
+    with torch.no_grad():
+        bias_3d.copy_(bias_2d)
+    return bias_3d
+
+
+def modify_state_dict(layer, state_dict, prefix, inflate_weight_fn, inflate_bias_fn):
+    """
+    the main function to inflated 2D parameters to 3D.
+    """
+    weight_name = prefix + "weight"
+    bias_name = prefix + "bias"
+    if weight_name in state_dict:
+        weight_2d = state_dict[weight_name]
+        if weight_2d.dim() == 4:
+            # Assuming the 2D weights are 4D tensors (out_channels, in_channels, h, w)
+            weight_3d = inflate_weight_fn(
+                weight_2d=weight_2d,
+                weight_3d=layer.weight,
+                inflation_mode=layer.inflation_mode,
+            )
+            state_dict[weight_name] = weight_3d
+        else:
+            return state_dict
+            # It's a 3d state dict, should not do inflation on both bias and weight.
+    if bias_name in state_dict:
+        bias_2d = state_dict[bias_name]
+        if bias_2d.dim() == 1:
+            # Assuming the 2D biases are 1D tensors (out_channels,)
+            bias_3d = inflate_bias_fn(
+                bias_2d=bias_2d,
+                bias_3d=layer.bias,
+                inflation_mode=layer.inflation_mode,
+            )
+            state_dict[bias_name] = bias_3d
+    return state_dict

modules/seedvr/src/models/video_vae_v3/modules/context_parallel_lib.py

@@ -0,0 +1,66 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import List
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from .types import MemoryState
+
+# Single GPU inference - no distributed processing needed
+
+
+def causal_conv_slice_inputs(x, split_size, memory_state):
+    # Single GPU inference - no slicing needed, return full tensor
+    return x
+
+
+def causal_conv_gather_outputs(x):
+    # Single GPU inference - no gathering needed, return tensor as is
+    return x
+
+
+def get_output_len(conv_module, input_len, pad_len, dim=0):
+    dilated_kernerl_size = conv_module.dilation[dim] * (conv_module.kernel_size[dim] - 1) + 1
+    output_len = (input_len + pad_len - dilated_kernerl_size) // conv_module.stride[dim] + 1
+    return output_len
+
+
+def get_cache_size(conv_module, input_len, pad_len, dim=0):
+    dilated_kernerl_size = conv_module.dilation[dim] * (conv_module.kernel_size[dim] - 1) + 1
+    output_len = (input_len + pad_len - dilated_kernerl_size) // conv_module.stride[dim] + 1
+    remain_len = (
+        input_len + pad_len - ((output_len - 1) * conv_module.stride[dim] + dilated_kernerl_size)
+    )
+    overlap_len = dilated_kernerl_size - conv_module.stride[dim]
+    cache_len = overlap_len + remain_len  # >= 0
+
+    assert output_len > 0
+    return cache_len
+
+
+def cache_send_recv(tensor: List[Tensor], cache_size, times, memory=None):
+    # Single GPU inference - simplified cache handling
+    recv_buffer = None
+
+    # Handle memory buffer for single GPU case
+    if memory is not None:
+        recv_buffer = memory.to(tensor[0])
+    elif times > 0:
+        tile_repeat = [1] * tensor[0].ndim
+        tile_repeat[2] = times
+        recv_buffer = torch.tile(tensor[0][:, :, :1], tile_repeat)
+
+    return recv_buffer

modules/seedvr/src/models/video_vae_v3/modules/global_config.py

@@ -0,0 +1,28 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Optional
+
+_NORM_LIMIT = float("inf")
+
+
+def get_norm_limit():
+    return _NORM_LIMIT
+
+
+def set_norm_limit(value: Optional[float] = None):
+    global _NORM_LIMIT
+    if value is None:
+        value = float("inf")
+    _NORM_LIMIT = value

modules/seedvr/src/models/video_vae_v3/modules/inflated_layers.py

@@ -0,0 +1,106 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from functools import partial
+from typing import Literal, Optional
+from torch import Tensor
+from torch.nn import Conv3d
+
+from .inflated_lib import (
+    MemoryState,
+    extend_head,
+    inflate_bias,
+    inflate_weight,
+    modify_state_dict,
+)
+
+_inflation_mode_t = Literal["none", "tail", "replicate"]
+_memory_device_t = Optional[Literal["cpu", "same"]]
+
+
+class InflatedCausalConv3d(Conv3d):
+    def __init__(
+        self,
+        *args,
+        inflation_mode: _inflation_mode_t,
+        memory_device: _memory_device_t = "same",
+        **kwargs,
+    ):
+        self.inflation_mode = inflation_mode
+        self.memory = None
+        super().__init__(*args, **kwargs)
+        self.temporal_padding = self.padding[0]
+        self.memory_device = memory_device
+        self.padding = (0, *self.padding[1:])  # Remove temporal pad to keep causal.
+
+    def set_memory_device(self, memory_device: _memory_device_t):
+        self.memory_device = memory_device
+
+    def forward(self, input: Tensor, memory_state: MemoryState = MemoryState.DISABLED) -> Tensor:
+        mem_size = self.stride[0] - self.kernel_size[0]
+        if (self.memory is not None) and (memory_state == MemoryState.ACTIVE):
+            input = extend_head(input, memory=self.memory)
+        else:
+            input = extend_head(input, times=self.temporal_padding * 2)
+        memory = (
+            input[:, :, mem_size:].detach()
+            if (mem_size != 0 and memory_state != MemoryState.DISABLED)
+            else None
+        )
+        if (
+            memory_state != MemoryState.DISABLED
+            and not self.training
+            and (self.memory_device is not None)
+        ):
+            self.memory = memory
+            if self.memory_device == "cpu" and self.memory is not None:
+                self.memory = self.memory.to("cpu")
+        return super().forward(input)
+
+    def _load_from_state_dict(
+        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+    ):
+        if self.inflation_mode != "none":
+            state_dict = modify_state_dict(
+                self,
+                state_dict,
+                prefix,
+                inflate_weight_fn=partial(inflate_weight, position="tail"),
+                inflate_bias_fn=partial(inflate_bias, position="tail"),
+            )
+        super()._load_from_state_dict(
+            state_dict,
+            prefix,
+            local_metadata,
+            (strict and self.inflation_mode == "none"),
+            missing_keys,
+            unexpected_keys,
+            error_msgs,
+        )
+
+
+def init_causal_conv3d(
+    *args,
+    inflation_mode: _inflation_mode_t,
+    **kwargs,
+):
+    """
+    Initialize a Causal-3D convolution layer.
+    Parameters:
+        inflation_mode: Listed as below. It's compatible with all the 3D-VAE checkpoints we have.
+            - none: No inflation will be conducted.
+                    The loading logic of state dict will fall back to default.
+            - tail / replicate: Refer to the definition of `InflatedCausalConv3d`.
+    """
+    return InflatedCausalConv3d(*args, inflation_mode=inflation_mode, **kwargs)

modules/seedvr/src/models/video_vae_v3/modules/inflated_lib.py

@@ -0,0 +1,156 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from enum import Enum
+from typing import Optional
+import numpy as np
+import torch
+from diffusers.models.normalization import RMSNorm
+from einops import rearrange
+from torch import Tensor, nn
+
+from ....common.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class MemoryState(Enum):
+    """
+    State[Disabled]:        No memory bank will be enabled.
+    State[Initializing]:    The model is handling the first clip,
+                            need to reset / initialize the memory bank.
+    State[Active]:          There has been some data in the memory bank.
+    """
+
+    DISABLED = 0
+    INITIALIZING = 1
+    ACTIVE = 2
+
+
+def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor:
+    if isinstance(norm_layer, (nn.LayerNorm, RMSNorm)):
+        if x.ndim == 4:
+            x = rearrange(x, "b c h w -> b h w c")
+            x = norm_layer(x)
+            x = rearrange(x, "b h w c -> b c h w")
+            return x
+        if x.ndim == 5:
+            x = rearrange(x, "b c t h w -> b t h w c")
+            x = norm_layer(x)
+            x = rearrange(x, "b t h w c -> b c t h w")
+            return x
+    if isinstance(norm_layer, (nn.GroupNorm, nn.BatchNorm2d, nn.SyncBatchNorm)):
+        if x.ndim <= 4:
+            return norm_layer(x)
+        if x.ndim == 5:
+            t = x.size(2)
+            x = rearrange(x, "b c t h w -> (b t) c h w")
+            x = norm_layer(x)
+            x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+            return x
+    raise NotImplementedError
+
+
+def remove_head(tensor: Tensor, times: int = 1) -> Tensor:
+    """
+    Remove duplicated first frame features in the up-sampling process.
+    """
+    if times == 0:
+        return tensor
+    return torch.cat(tensors=(tensor[:, :, :1], tensor[:, :, times + 1 :]), dim=2)
+
+
+def extend_head(
+    tensor: Tensor, times: Optional[int] = 2, memory: Optional[Tensor] = None
+) -> Tensor:
+    """
+    When memory is None:
+        - Duplicate first frame features in the down-sampling process.
+    When memory is not None:
+        - Concatenate memory features with the input features to keep temporal consistency.
+    """
+    if times == 0:
+        return tensor
+    if memory is not None:
+        return torch.cat((memory.to(tensor), tensor), dim=2)
+    else:
+        tile_repeat = np.ones(tensor.ndim).astype(int)
+        tile_repeat[2] = times
+        return torch.cat(tensors=(torch.tile(tensor[:, :, :1], list(tile_repeat)), tensor), dim=2)
+
+
+def inflate_weight(weight_2d: torch.Tensor, weight_3d: torch.Tensor, inflation_mode: str):
+    """
+    Inflate a 2D convolution weight matrix to a 3D one.
+    Parameters:
+        weight_2d:      The weight matrix of 2D conv to be inflated.
+        weight_3d:      The weight matrix of 3D conv to be initialized.
+        inflation_mode: the mode of inflation
+    """
+    assert inflation_mode in ["constant", "replicate"]
+    assert weight_3d.shape[:2] == weight_2d.shape[:2]
+    with torch.no_grad():
+        if inflation_mode == "replicate":
+            depth = weight_3d.size(2)
+            weight_3d.copy_(weight_2d.unsqueeze(2).repeat(1, 1, depth, 1, 1) / depth)
+        else:
+            weight_3d.fill_(0.0)
+            weight_3d[:, :, -1].copy_(weight_2d)
+    return weight_3d
+
+
+def inflate_bias(bias_2d: torch.Tensor, bias_3d: torch.Tensor, inflation_mode: str):
+    """
+    Inflate a 2D convolution bias tensor to a 3D one
+    Parameters:
+        bias_2d:        The bias tensor of 2D conv to be inflated.
+        bias_3d:        The bias tensor of 3D conv to be initialized.
+        inflation_mode: Placeholder to align `inflate_weight`.
+    """
+    assert bias_3d.shape == bias_2d.shape
+    with torch.no_grad():
+        bias_3d.copy_(bias_2d)
+    return bias_3d
+
+
+def modify_state_dict(layer, state_dict, prefix, inflate_weight_fn, inflate_bias_fn):
+    """
+    the main function to inflated 2D parameters to 3D.
+    """
+    weight_name = prefix + "weight"
+    bias_name = prefix + "bias"
+    if weight_name in state_dict:
+        weight_2d = state_dict[weight_name]
+        if weight_2d.dim() == 4:
+            # Assuming the 2D weights are 4D tensors (out_channels, in_channels, h, w)
+            weight_3d = inflate_weight_fn(
+                weight_2d=weight_2d,
+                weight_3d=layer.weight,
+                inflation_mode=layer.inflation_mode,
+            )
+            state_dict[weight_name] = weight_3d
+        else:
+            return state_dict
+            # It's a 3d state dict, should not do inflation on both bias and weight.
+    if bias_name in state_dict:
+        bias_2d = state_dict[bias_name]
+        if bias_2d.dim() == 1:
+            # Assuming the 2D biases are 1D tensors (out_channels,)
+            bias_3d = inflate_bias_fn(
+                bias_2d=bias_2d,
+                bias_3d=layer.bias,
+                inflation_mode=layer.inflation_mode,
+            )
+            state_dict[bias_name] = bias_3d
+    return state_dict

modules/seedvr/src/models/video_vae_v3/modules/types.py

@@ -0,0 +1,76 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from enum import Enum
+from typing import Dict, Literal, NamedTuple, Optional
+import torch
+
+_receptive_field_t = Literal["half", "full"]
+_inflation_mode_t = Literal["none", "tail", "replicate"]
+_memory_device_t = Optional[Literal["cpu", "same"]]
+_gradient_checkpointing_t = Optional[Literal["half", "full"]]
+_selective_checkpointing_t = Optional[Literal["coarse", "fine"]]
+
+class DiagonalGaussianDistribution:
+    def __init__(self, mean: torch.Tensor, logvar: torch.Tensor):
+        self.mean = mean
+        self.logvar = torch.clamp(logvar, -30.0, 20.0)
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+
+    def mode(self) -> torch.Tensor:
+        return self.mean
+
+    def sample(self) -> torch.FloatTensor:
+        return self.mean + self.std * torch.randn_like(self.mean)
+
+    def kl(self) -> torch.Tensor:
+        return 0.5 * torch.sum(
+            self.mean**2 + self.var - 1.0 - self.logvar,
+            dim=list(range(1, self.mean.ndim)),
+        )
+
+class MemoryState(Enum):
+    """
+    State[Disabled]:        No memory bank will be enabled.
+    State[Initializing]:    The model is handling the first clip, need to reset the memory bank.
+    State[Active]:          There has been some data in the memory bank.
+    State[Unset]:           Error state, indicating users didn't pass correct memory state in.
+    """
+
+    DISABLED = 0
+    INITIALIZING = 1
+    ACTIVE = 2
+    UNSET = 3
+
+
+class QuantizerOutput(NamedTuple):
+    latent: torch.Tensor
+    extra_loss: torch.Tensor
+    statistics: Dict[str, torch.Tensor]
+
+
+class CausalAutoencoderOutput(NamedTuple):
+    sample: torch.Tensor
+    latent: torch.Tensor
+    posterior: Optional[DiagonalGaussianDistribution]
+
+
+class CausalEncoderOutput(NamedTuple):
+    latent: torch.Tensor
+    posterior: Optional[DiagonalGaussianDistribution]
+
+
+class CausalDecoderOutput(NamedTuple):
+    sample: torch.Tensor

modules/seedvr/src/models/video_vae_v3/modules/video_vae.py

@@ -0,0 +1,956 @@
+# Copyright (c) 2023 HuggingFace Team
+# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache License, Version 2.0 (the "License")
+#
+# This file has been modified by ByteDance Ltd. and/or its affiliates. on 1st June 2025
+#
+# Original file was released under Apache License, Version 2.0 (the "License"), with the full license text
+# available at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# This modified file is released under the same license.
+
+from contextlib import nullcontext
+from typing import Optional, Tuple, Literal, Callable, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+from einops import rearrange
+from ....common.half_precision_fixes import safe_pad_operation
+
+from ....common.distributed.advanced import get_sequence_parallel_world_size
+from ....common.logger import get_logger
+from .causal_inflation_lib import (
+    InflatedCausalConv3d,
+    causal_norm_wrapper,
+    init_causal_conv3d,
+    remove_head,
+)
+from .context_parallel_lib import (
+    causal_conv_gather_outputs,
+    causal_conv_slice_inputs,
+)
+from .global_config import set_norm_limit
+from .types import (
+    CausalAutoencoderOutput,
+    CausalDecoderOutput,
+    CausalEncoderOutput,
+    MemoryState,
+    _inflation_mode_t,
+    _memory_device_t,
+    _receptive_field_t,
+    _selective_checkpointing_t,
+)
+
+logger = get_logger(__name__)  # pylint: disable=invalid-name
+
+# Fake func, no checkpointing is required for inference
+def gradient_checkpointing(module: Union[Callable, nn.Module], *args, enabled: bool, **kwargs):
+    return module(*args, **kwargs)
+
+class ResnetBlock2D(nn.Module):
+    r"""
+    A Resnet block.
+
+    Parameters:
+        in_channels (`int`): The number of channels in the input.
+        out_channels (`int`, *optional*, default to be `None`):
+            The number of output channels for the first conv2d layer.
+            If None, same as `in_channels`.
+        dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.
+    """
+
+    def __init__(
+        self, *, in_channels: int, out_channels: Optional[int] = None, dropout: float = 0.0
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+
+        self.nonlinearity = nn.SiLU()
+
+        self.norm1 = torch.nn.GroupNorm(
+            num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+        )
+
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        self.norm2 = torch.nn.GroupNorm(
+            num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
+        )
+
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        self.use_in_shortcut = self.in_channels != out_channels
+
+        self.conv_shortcut = None
+        if self.use_in_shortcut:
+            self.conv_shortcut = nn.Conv2d(
+                in_channels, out_channels, kernel_size=1, stride=1, padding=0
+            )
+
+    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden = input_tensor
+
+        hidden = self.norm1(hidden)
+        hidden = self.nonlinearity(hidden)
+        hidden = self.conv1(hidden)
+
+        hidden = self.norm2(hidden)
+        hidden = self.nonlinearity(hidden)
+        hidden = self.dropout(hidden)
+        hidden = self.conv2(hidden)
+
+        if self.conv_shortcut is not None:
+            input_tensor = self.conv_shortcut(input_tensor)
+
+        output_tensor = input_tensor + hidden
+
+        return output_tensor
+
+class Upsample3D(nn.Module):
+    """A 3D upsampling layer."""
+
+    def __init__(
+        self,
+        channels: int,
+        inflation_mode: _inflation_mode_t = "tail",
+        temporal_up: bool = False,
+        spatial_up: bool = True,
+        slicing: bool = False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.conv = init_causal_conv3d(
+            self.channels, self.channels, kernel_size=3, padding=1, inflation_mode=inflation_mode
+        )
+
+        self.temporal_up = temporal_up
+        self.spatial_up = spatial_up
+        self.temporal_ratio = 2 if temporal_up else 1
+        self.spatial_ratio = 2 if spatial_up else 1
+        self.slicing = slicing
+
+        upscale_ratio = (self.spatial_ratio**2) * self.temporal_ratio
+        self.upscale_conv = nn.Conv3d(
+            self.channels, self.channels * upscale_ratio, kernel_size=1, padding=0
+        )
+        identity = (
+            torch.eye(self.channels).repeat(upscale_ratio, 1).reshape_as(self.upscale_conv.weight)
+        )
+
+        self.upscale_conv.weight.data.copy_(identity)
+        nn.init.zeros_(self.upscale_conv.bias)
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        memory_state: MemoryState,
+    ) -> torch.FloatTensor:
+        return gradient_checkpointing(
+            self.custom_forward,
+            hidden_states,
+            memory_state,
+            enabled=self.training and self.gradient_checkpointing,
+        )
+
+    def custom_forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        memory_state: MemoryState,
+    ) -> torch.FloatTensor:
+        assert hidden_states.shape[1] == self.channels
+
+        if self.slicing:
+            split_size = hidden_states.size(2) // 2
+            hidden_states = list(
+                hidden_states.split([split_size, hidden_states.size(2) - split_size], dim=2)
+            )
+        else:
+            hidden_states = [hidden_states]
+
+        for i in range(len(hidden_states)):
+            hidden_states[i] = self.upscale_conv(hidden_states[i])
+            hidden_states[i] = rearrange(
+                hidden_states[i],
+                "b (x y z c) f h w -> b c (f z) (h x) (w y)",
+                x=self.spatial_ratio,
+                y=self.spatial_ratio,
+                z=self.temporal_ratio,
+            )
+
+        # [Overridden] For causal temporal conv
+        if self.temporal_up and memory_state != MemoryState.ACTIVE:
+            hidden_states[0] = remove_head(hidden_states[0])
+
+        if self.slicing:
+            hidden_states = self.conv(hidden_states, memory_state=memory_state)
+            return torch.cat(hidden_states, dim=2)
+        else:
+            return self.conv(hidden_states[0], memory_state=memory_state)
+
+
+class Downsample3D(nn.Module):
+    """A 3D downsampling layer."""
+
+    def __init__(
+        self,
+        channels: int,
+        inflation_mode: _inflation_mode_t = "tail",
+        temporal_down: bool = False,
+        spatial_down: bool = True,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.temporal_down = temporal_down
+        self.spatial_down = spatial_down
+
+        self.temporal_ratio = 2 if temporal_down else 1
+        self.spatial_ratio = 2 if spatial_down else 1
+
+        self.temporal_kernel = 3 if temporal_down else 1
+        self.spatial_kernel = 3 if spatial_down else 1
+
+        self.conv = init_causal_conv3d(
+            self.channels,
+            self.channels,
+            kernel_size=(self.temporal_kernel, self.spatial_kernel, self.spatial_kernel),
+            stride=(self.temporal_ratio, self.spatial_ratio, self.spatial_ratio),
+            padding=((1 if self.temporal_down else 0), 0, 0),
+            inflation_mode=inflation_mode,
+        )
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        memory_state: MemoryState,
+    ) -> torch.FloatTensor:
+        return gradient_checkpointing(
+            self.custom_forward,
+            hidden_states,
+            memory_state,
+            enabled=self.training and self.gradient_checkpointing,
+        )
+
+    def custom_forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        memory_state: MemoryState,
+    ) -> torch.FloatTensor:
+
+        assert hidden_states.shape[1] == self.channels
+
+        if self.spatial_down:
+            hidden_states = safe_pad_operation(hidden_states, (0, 1, 0, 1), mode="constant", value=0)
+
+        hidden_states = self.conv(hidden_states, memory_state=memory_state)
+        return hidden_states
+
+
+class ResnetBlock3D(ResnetBlock2D):
+    def __init__(
+        self,
+        *args,
+        inflation_mode: _inflation_mode_t = "tail",
+        time_receptive_field: _receptive_field_t = "half",
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.conv1 = init_causal_conv3d(
+            self.in_channels,
+            self.out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            inflation_mode=inflation_mode,
+        )
+
+        self.conv2 = init_causal_conv3d(
+            self.out_channels,
+            self.out_channels,
+            kernel_size=(1, 3, 3) if time_receptive_field == "half" else (3, 3, 3),
+            stride=1,
+            padding=(0, 1, 1) if time_receptive_field == "half" else (1, 1, 1),
+            inflation_mode=inflation_mode,
+        )
+
+        if self.use_in_shortcut:
+            self.conv_shortcut = init_causal_conv3d(
+                self.in_channels,
+                self.out_channels,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                bias=(self.conv_shortcut.bias is not None),
+                inflation_mode=inflation_mode,
+            )
+        self.gradient_checkpointing = False
+
+    def forward(self, input_tensor: torch.Tensor, memory_state: MemoryState = MemoryState.UNSET):
+        return gradient_checkpointing(
+            self.custom_forward,
+            input_tensor,
+            memory_state,
+            enabled=self.training and self.gradient_checkpointing,
+        )
+
+    def custom_forward(
+        self, input_tensor: torch.Tensor, memory_state: MemoryState = MemoryState.UNSET
+    ):
+        assert memory_state != MemoryState.UNSET
+        hidden_states = input_tensor
+
+        hidden_states = causal_norm_wrapper(self.norm1, hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.conv1(hidden_states, memory_state=memory_state)
+
+        hidden_states = causal_norm_wrapper(self.norm2, hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.conv2(hidden_states, memory_state=memory_state)
+
+        if self.conv_shortcut is not None:
+            input_tensor = self.conv_shortcut(input_tensor, memory_state=memory_state)
+
+        output_tensor = input_tensor + hidden_states
+
+        return output_tensor
+
+
+class DownEncoderBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        add_downsample: bool = True,
+        inflation_mode: _inflation_mode_t = "tail",
+        time_receptive_field: _receptive_field_t = "half",
+        temporal_down: bool = True,
+        spatial_down: bool = True,
+    ):
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                ResnetBlock3D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    dropout=dropout,
+                    inflation_mode=inflation_mode,
+                    time_receptive_field=time_receptive_field,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        self.downsamplers = None
+        if add_downsample:
+            # Todo: Refactor this line before V5 Image VAE Training.
+            self.downsamplers = nn.ModuleList(
+                [
+                    Downsample3D(
+                        channels=out_channels,
+                        inflation_mode=inflation_mode,
+                        temporal_down=temporal_down,
+                        spatial_down=spatial_down,
+                    )
+                ]
+            )
+
+    def forward(
+        self, hidden_states: torch.FloatTensor, memory_state: MemoryState
+    ) -> torch.FloatTensor:
+        for resnet in self.resnets:
+            hidden_states = resnet(hidden_states, memory_state=memory_state)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states, memory_state=memory_state)
+
+        return hidden_states
+
+
+class UpDecoderBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        add_upsample: bool = True,
+        inflation_mode: _inflation_mode_t = "tail",
+        time_receptive_field: _receptive_field_t = "half",
+        temporal_up: bool = True,
+        spatial_up: bool = True,
+        slicing: bool = False,
+    ):
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            input_channels = in_channels if i == 0 else out_channels
+
+            resnets.append(
+                ResnetBlock3D(
+                    in_channels=input_channels,
+                    out_channels=out_channels,
+                    dropout=dropout,
+                    inflation_mode=inflation_mode,
+                    time_receptive_field=time_receptive_field,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        self.upsamplers = None
+        # Todo: Refactor this line before V5 Image VAE Training.
+        if add_upsample:
+            self.upsamplers = nn.ModuleList(
+                [
+                    Upsample3D(
+                        channels=out_channels,
+                        inflation_mode=inflation_mode,
+                        temporal_up=temporal_up,
+                        spatial_up=spatial_up,
+                        slicing=slicing,
+                    )
+                ]
+            )
+
+    def forward(
+        self, hidden_states: torch.FloatTensor, memory_state: MemoryState
+    ) -> torch.FloatTensor:
+        for resnet in self.resnets:
+            hidden_states = resnet(hidden_states, memory_state=memory_state)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states, memory_state=memory_state)
+
+        return hidden_states
+
+
+class UNetMidBlock3D(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        dropout: float = 0.0,
+        inflation_mode: _inflation_mode_t = "tail",
+        time_receptive_field: _receptive_field_t = "half",
+    ):
+        super().__init__()
+        self.resnets = nn.ModuleList(
+            [
+                ResnetBlock3D(
+                    in_channels=channels,
+                    out_channels=channels,
+                    dropout=dropout,
+                    inflation_mode=inflation_mode,
+                    time_receptive_field=time_receptive_field,
+                ),
+                ResnetBlock3D(
+                    in_channels=channels,
+                    out_channels=channels,
+                    dropout=dropout,
+                    inflation_mode=inflation_mode,
+                    time_receptive_field=time_receptive_field,
+                ),
+            ]
+        )
+
+    def forward(self, hidden_states: torch.Tensor, memory_state: MemoryState):
+        for resnet in self.resnets:
+            hidden_states = resnet(hidden_states, memory_state)
+        return hidden_states
+
+
+class Encoder3D(nn.Module):
+    r"""
+    The `Encoder` layer of a variational autoencoder that encodes
+    its input into a latent representation.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        block_out_channels: Tuple[int, ...] = (64,),
+        layers_per_block: int = 2,
+        double_z: bool = True,
+        temporal_down_num: int = 2,
+        inflation_mode: _inflation_mode_t = "tail",
+        time_receptive_field: _receptive_field_t = "half",
+        selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",),
+    ):
+        super().__init__()
+        self.layers_per_block = layers_per_block
+
+        self.temporal_down_num = temporal_down_num
+
+        self.conv_in = init_causal_conv3d(
+            in_channels,
+            block_out_channels[0],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            inflation_mode=inflation_mode,
+        )
+
+        self.down_blocks = nn.ModuleList([])
+
+        # down
+        output_channel = block_out_channels[0]
+        for i in range(len(block_out_channels)):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+            is_temporal_down_block = i >= len(block_out_channels) - self.temporal_down_num - 1
+            # Note: take the last one
+
+            down_block = DownEncoderBlock3D(
+                num_layers=self.layers_per_block,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                add_downsample=not is_final_block,
+                temporal_down=is_temporal_down_block,
+                spatial_down=True,
+                inflation_mode=inflation_mode,
+                time_receptive_field=time_receptive_field,
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        self.mid_block = UNetMidBlock3D(
+            channels=block_out_channels[-1],
+            inflation_mode=inflation_mode,
+            time_receptive_field=time_receptive_field,
+        )
+
+        # out
+        self.conv_norm_out = nn.GroupNorm(
+            num_channels=block_out_channels[-1], num_groups=32, eps=1e-6
+        )
+        self.conv_act = nn.SiLU()
+
+        conv_out_channels = 2 * out_channels if double_z else out_channels
+        self.conv_out = init_causal_conv3d(
+            block_out_channels[-1], conv_out_channels, 3, padding=1, inflation_mode=inflation_mode
+        )
+
+        assert len(selective_checkpointing) == len(self.down_blocks)
+        self.set_gradient_checkpointing(selective_checkpointing)
+
+    def set_gradient_checkpointing(self, checkpointing_types):
+        gradient_checkpointing = []
+        for down_block, sac_type in zip(self.down_blocks, checkpointing_types):
+            if sac_type == "coarse":
+                gradient_checkpointing.append(True)
+            elif sac_type == "fine":
+                for n, m in down_block.named_modules():
+                    if hasattr(m, "gradient_checkpointing"):
+                        m.gradient_checkpointing = True
+                        logger.debug(f"set gradient_checkpointing: {n}")
+                gradient_checkpointing.append(False)
+            else:
+                gradient_checkpointing.append(False)
+        self.gradient_checkpointing = gradient_checkpointing
+        logger.info(f"[Encoder3D] gradient_checkpointing: {checkpointing_types}")
+
+    def forward(self, sample: torch.FloatTensor, memory_state: MemoryState) -> torch.FloatTensor:
+        r"""The forward method of the `Encoder` class."""
+        sample = self.conv_in(sample, memory_state=memory_state)
+        # down
+        for down_block, sac in zip(self.down_blocks, self.gradient_checkpointing):
+            sample = gradient_checkpointing(
+                down_block,
+                sample,
+                memory_state=memory_state,
+                enabled=self.training and sac,
+            )
+
+        # middle
+        sample = self.mid_block(sample, memory_state=memory_state)
+
+        # post-process
+        sample = causal_norm_wrapper(self.conv_norm_out, sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample, memory_state=memory_state)
+
+        return sample
+
+
+class Decoder3D(nn.Module):
+    r"""
+    The `Decoder` layer of a variational autoencoder that
+    decodes its latent representation into an output sample.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        block_out_channels: Tuple[int, ...] = (64,),
+        layers_per_block: int = 2,
+        inflation_mode: _inflation_mode_t = "tail",
+        time_receptive_field: _receptive_field_t = "half",
+        temporal_up_num: int = 2,
+        slicing_up_num: int = 0,
+        selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",),
+    ):
+        super().__init__()
+        self.layers_per_block = layers_per_block
+        self.temporal_up_num = temporal_up_num
+
+        self.conv_in = init_causal_conv3d(
+            in_channels,
+            block_out_channels[-1],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            inflation_mode=inflation_mode,
+        )
+
+        self.up_blocks = nn.ModuleList([])
+
+        # mid
+        self.mid_block = UNetMidBlock3D(
+            channels=block_out_channels[-1],
+            inflation_mode=inflation_mode,
+            time_receptive_field=time_receptive_field,
+        )
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        output_channel = reversed_block_out_channels[0]
+        for i in range(len(reversed_block_out_channels)):
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+
+            is_final_block = i == len(block_out_channels) - 1
+            is_temporal_up_block = i < self.temporal_up_num
+            is_slicing_up_block = i >= len(block_out_channels) - slicing_up_num
+            # Note: Keep symmetric
+
+            up_block = UpDecoderBlock3D(
+                num_layers=self.layers_per_block + 1,
+                in_channels=prev_output_channel,
+                out_channels=output_channel,
+                add_upsample=not is_final_block,
+                temporal_up=is_temporal_up_block,
+                slicing=is_slicing_up_block,
+                inflation_mode=inflation_mode,
+                time_receptive_field=time_receptive_field,
+            )
+            self.up_blocks.append(up_block)
+
+        # out
+        self.conv_norm_out = nn.GroupNorm(
+            num_channels=block_out_channels[0], num_groups=32, eps=1e-6
+        )
+        self.conv_act = nn.SiLU()
+        self.conv_out = init_causal_conv3d(
+            block_out_channels[0], out_channels, 3, padding=1, inflation_mode=inflation_mode
+        )
+
+        assert len(selective_checkpointing) == len(self.up_blocks)
+        self.set_gradient_checkpointing(selective_checkpointing)
+
+    def set_gradient_checkpointing(self, checkpointing_types):
+        gradient_checkpointing = []
+        for up_block, sac_type in zip(self.up_blocks, checkpointing_types):
+            if sac_type == "coarse":
+                gradient_checkpointing.append(True)
+            elif sac_type == "fine":
+                for n, m in up_block.named_modules():
+                    if hasattr(m, "gradient_checkpointing"):
+                        m.gradient_checkpointing = True
+                        logger.debug(f"set gradient_checkpointing: {n}")
+                gradient_checkpointing.append(False)
+            else:
+                gradient_checkpointing.append(False)
+        self.gradient_checkpointing = gradient_checkpointing
+        logger.info(f"[Decoder3D] gradient_checkpointing: {checkpointing_types}")
+
+    def forward(self, sample: torch.FloatTensor, memory_state: MemoryState) -> torch.FloatTensor:
+        r"""The forward method of the `Decoder` class."""
+
+        sample = self.conv_in(sample, memory_state=memory_state)
+
+        # middle
+        sample = self.mid_block(sample, memory_state=memory_state)
+
+        # up
+        for up_block, sac in zip(self.up_blocks, self.gradient_checkpointing):
+            sample = gradient_checkpointing(
+                up_block,
+                sample,
+                memory_state=memory_state,
+                enabled=self.training and sac,
+            )
+
+        # post-process
+        sample = causal_norm_wrapper(self.conv_norm_out, sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample, memory_state=memory_state)
+
+        return sample
+
+
+class VideoAutoencoderKL(nn.Module):
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        block_out_channels: Tuple[int] = (64,),
+        layers_per_block: int = 1,
+        latent_channels: int = 4,
+        use_quant_conv: bool = True,
+        use_post_quant_conv: bool = True,
+        enc_selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",),
+        dec_selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",),
+        temporal_scale_num: int = 3,
+        slicing_up_num: int = 0,
+        inflation_mode: _inflation_mode_t = "tail",
+        time_receptive_field: _receptive_field_t = "half",
+        slicing_sample_min_size: int = None,
+        spatial_downsample_factor: int = 16,
+        temporal_downsample_factor: int = 8,
+        freeze_encoder: bool = False,
+    ):
+        super().__init__()
+        self.spatial_downsample_factor = spatial_downsample_factor
+        self.temporal_downsample_factor = temporal_downsample_factor
+        self.freeze_encoder = freeze_encoder
+        if slicing_sample_min_size is None:
+            slicing_sample_min_size = temporal_downsample_factor
+        self.slicing_sample_min_size = slicing_sample_min_size
+        self.slicing_latent_min_size = slicing_sample_min_size // (2**temporal_scale_num)
+
+        # pass init params to Encoder
+        self.encoder = Encoder3D(
+            in_channels=in_channels,
+            out_channels=latent_channels,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            double_z=True,
+            temporal_down_num=temporal_scale_num,
+            selective_checkpointing=enc_selective_checkpointing,
+            inflation_mode=inflation_mode,
+            time_receptive_field=time_receptive_field,
+        )
+
+        # pass init params to Decoder
+        self.decoder = Decoder3D(
+            in_channels=latent_channels,
+            out_channels=out_channels,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            # [Override] add temporal_up_num parameter
+            temporal_up_num=temporal_scale_num,
+            slicing_up_num=slicing_up_num,
+            selective_checkpointing=dec_selective_checkpointing,
+            inflation_mode=inflation_mode,
+            time_receptive_field=time_receptive_field,
+        )
+
+        self.quant_conv = (
+            init_causal_conv3d(
+                in_channels=2 * latent_channels,
+                out_channels=2 * latent_channels,
+                kernel_size=1,
+                inflation_mode=inflation_mode,
+            )
+            if use_quant_conv
+            else None
+        )
+        self.post_quant_conv = (
+            init_causal_conv3d(
+                in_channels=latent_channels,
+                out_channels=latent_channels,
+                kernel_size=1,
+                inflation_mode=inflation_mode,
+            )
+            if use_post_quant_conv
+            else None
+        )
+
+        self.use_slicing = False
+
+    def enable_slicing(self):
+        self.use_slicing = True
+
+    def disable_slicing(self):
+        self.use_slicing = False
+
+    def encode(self, x: torch.FloatTensor) -> CausalEncoderOutput:
+        if x.ndim == 4:
+            x = x.unsqueeze(2)
+        h = self.slicing_encode(x)
+        p = DiagonalGaussianDistribution(h)
+        z = p.sample()
+        return CausalEncoderOutput(z, p)
+
+    def decode(self, z: torch.FloatTensor) -> CausalDecoderOutput:
+        if z.ndim == 4:
+            z = z.unsqueeze(2)
+        x = self.slicing_decode(z)
+        return CausalDecoderOutput(x)
+
+    def _encode(self, x: torch.Tensor, memory_state: MemoryState) -> torch.Tensor:
+        x = causal_conv_slice_inputs(x, self.slicing_sample_min_size, memory_state=memory_state)
+        h = self.encoder(x, memory_state=memory_state)
+        h = self.quant_conv(h, memory_state=memory_state) if self.quant_conv is not None else h
+        h = causal_conv_gather_outputs(h)
+        return h
+
+    def _decode(self, z: torch.Tensor, memory_state: MemoryState) -> torch.Tensor:
+        z = causal_conv_slice_inputs(z, self.slicing_latent_min_size, memory_state=memory_state)
+        z = (
+            self.post_quant_conv(z, memory_state=memory_state)
+            if self.post_quant_conv is not None
+            else z
+        )
+        x = self.decoder(z, memory_state=memory_state)
+        x = causal_conv_gather_outputs(x)
+        return x
+
+    def slicing_encode(self, x: torch.Tensor) -> torch.Tensor:
+        sp_size = get_sequence_parallel_world_size()
+        if self.use_slicing and (x.shape[2] - 1) > self.slicing_sample_min_size * sp_size:
+            x_slices = x[:, :, 1:].split(split_size=self.slicing_sample_min_size * sp_size, dim=2)
+            encoded_slices = [
+                self._encode(
+                    torch.cat((x[:, :, :1], x_slices[0]), dim=2),
+                    memory_state=MemoryState.INITIALIZING,
+                )
+            ]
+            for x_idx in range(1, len(x_slices)):
+                encoded_slices.append(
+                    self._encode(x_slices[x_idx], memory_state=MemoryState.ACTIVE)
+                )
+            return torch.cat(encoded_slices, dim=2)
+        else:
+            return self._encode(x, memory_state=MemoryState.DISABLED)
+
+    def slicing_decode(self, z: torch.Tensor) -> torch.Tensor:
+        sp_size = get_sequence_parallel_world_size()
+        if self.use_slicing and (z.shape[2] - 1) > self.slicing_latent_min_size * sp_size:
+            z_slices = z[:, :, 1:].split(split_size=self.slicing_latent_min_size * sp_size, dim=2)
+            decoded_slices = [
+                self._decode(
+                    torch.cat((z[:, :, :1], z_slices[0]), dim=2),
+                    memory_state=MemoryState.INITIALIZING,
+                )
+            ]
+            for z_idx in range(1, len(z_slices)):
+                decoded_slices.append(
+                    self._decode(z_slices[z_idx], memory_state=MemoryState.ACTIVE)
+                )
+            return torch.cat(decoded_slices, dim=2)
+        else:
+            return self._decode(z, memory_state=MemoryState.DISABLED)
+
+    def forward(self, x: torch.FloatTensor) -> CausalAutoencoderOutput:
+        with torch.no_grad() if self.freeze_encoder else nullcontext():
+            z, p = self.encode(x)
+        x = self.decode(z).sample
+        return CausalAutoencoderOutput(x, z, p)
+
+    def preprocess(self, x: torch.Tensor):
+        # x should in [B, C, T, H, W], [B, C, H, W]
+        assert x.ndim == 4 or x.size(2) % self.temporal_downsample_factor == 1
+        return x
+
+    def postprocess(self, x: torch.Tensor):
+        # x should in [B, C, T, H, W], [B, C, H, W]
+        return x
+
+    def set_causal_slicing(
+        self,
+        *,
+        split_size: Optional[int],
+        memory_device: _memory_device_t,
+    ):
+        assert (
+            split_size is None or memory_device is not None
+        ), "if split_size is set, memory_device must not be None."
+        if split_size is not None:
+            self.enable_slicing()
+            self.slicing_sample_min_size = split_size
+            self.slicing_latent_min_size = split_size // self.temporal_downsample_factor
+        else:
+            self.disable_slicing()
+        for module in self.modules():
+            if isinstance(module, InflatedCausalConv3d):
+                module.set_memory_device(memory_device)
+
+    def set_memory_limit(self, conv_max_mem: Optional[float], norm_max_mem: Optional[float]):
+        set_norm_limit(norm_max_mem)
+        for m in self.modules():
+            if isinstance(m, InflatedCausalConv3d):
+                m.set_memory_limit(conv_max_mem if conv_max_mem is not None else float("inf"))
+
+
+class VideoAutoencoderKLWrapper(VideoAutoencoderKL):
+    def __init__(
+        self, *args, spatial_downsample_factor: int, temporal_downsample_factor: int, **kwargs
+    ):
+        self.spatial_downsample_factor = spatial_downsample_factor
+        self.temporal_downsample_factor = temporal_downsample_factor
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x) -> CausalAutoencoderOutput:
+        z, _, p = self.encode(x)
+        x, _ = self.decode(z)
+        return CausalAutoencoderOutput(x, z, None, p)
+
+    def encode(self, x) -> CausalEncoderOutput:
+        if x.ndim == 4:
+            x = x.unsqueeze(2)
+        p = super().encode(x).latent_dist
+        z = p.sample().squeeze(2)
+        return CausalEncoderOutput(z, None, p)
+
+    def decode(self, z) -> CausalDecoderOutput:
+        if z.ndim == 4:
+            z = z.unsqueeze(2)
+        x = super().decode(z).sample.squeeze(2)
+        return CausalDecoderOutput(x, None)
+
+    def preprocess(self, x):
+        # x should in [B, C, T, H, W], [B, C, H, W]
+        assert x.ndim == 4 or x.size(2) % 4 == 1
+        return x
+
+    def postprocess(self, x):
+        # x should in [B, C, T, H, W], [B, C, H, W]
+        return x
+
+    def set_causal_slicing(
+        self,
+        *,
+        split_size: Optional[int],
+        memory_device: Optional[Literal["cpu", "same"]],
+    ):
+        assert (
+            split_size is None or memory_device is not None
+        ), "if split_size is set, memory_device must not be None."
+        if split_size is not None:
+            self.enable_slicing()
+        else:
+            self.disable_slicing()
+        self.slicing_sample_min_size = split_size
+        if split_size is not None:
+            self.slicing_latent_min_size = split_size // self.temporal_downsample_factor
+        for module in self.modules():
+            if isinstance(module, InflatedCausalConv3d):
+                module.set_memory_device(memory_device)

modules/seedvr/src/models/video_vae_v3/s8_c16_t4_inflation_sd3.yaml

@@ -0,0 +1,28 @@
+act_fn: silu
+block_out_channels:
+  - 128
+  - 256
+  - 512
+  - 512
+down_block_types:
+  - DownEncoderBlock3D
+  - DownEncoderBlock3D
+  - DownEncoderBlock3D
+  - DownEncoderBlock3D
+in_channels: 3
+latent_channels: 16
+layers_per_block: 2
+norm_num_groups: 32
+out_channels: 3
+slicing_sample_min_size: 4
+temporal_scale_num: 2
+inflation_mode: pad
+up_block_types:
+  - UpDecoderBlock3D
+  - UpDecoderBlock3D
+  - UpDecoderBlock3D
+  - UpDecoderBlock3D
+spatial_downsample_factor: 8
+temporal_downsample_factor: 4
+use_quant_conv: False
+use_post_quant_conv: False

modules/seedvr/src/optimization/memory_manager.py

@@ -0,0 +1,131 @@
+"""
+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
+from typing import Tuple, Optional
+from src.common.cache import Cache
+from src.models.dit_v2.rope import RotaryEmbeddingBase
+
+
+def preinitialize_rope_cache(runner) -> None:
+    """
+    🚀 Pre-initialize RoPE cache to avoid OOM at first launch
+    
+    Args:
+        runner: The model runner containing DiT and VAE models
+    """
+
+    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(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():
+                        # 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())
+
+                    # Store in cache
+                    temp_cache(cache_key, compute_freqs)
+
+            except Exception as e:
+                print(f"    ❌ Error in module {name}: {e}")
+            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:
+        print(f"  ⚠️ Error during RoPE pre-init: {e}")
+        print("  🔄 Model will work but could have OOM at first launch")
+
+
+def clear_rope_cache(runner) -> None:
+    """
+    🧹 Clear RoPE cache to free VRAM
+    
+    Args:
+        runner: The model runner containing the cache
+    """
+    print("🧹 Cleaning RoPE cache...")
+
+    if hasattr(runner, 'cache') and hasattr(runner.cache, 'cache'):
+        # Count entries before cleanup
+        cache_size = len(runner.cache.cache)
+
+        # Free all tensors from cache
+        for key, value in runner.cache.cache.items():
+            if isinstance(value, (tuple, list)):
+                for item in value:
+                    if hasattr(item, 'cpu'):
+                        item.cpu()
+                        del item
+            elif hasattr(value, 'cpu'):
+                value.cpu()
+                del value
+
+        # Clear the cache
+        runner.cache.cache.clear()
+        print(f"  ✅ RoPE cache cleared ({cache_size} entries removed)")
+
+    if hasattr(runner, 'dit'):
+        cleared_lru_count = 0
+        for module in runner.dit.modules():
+            if isinstance(module, RotaryEmbeddingBase):
+                if hasattr(module.get_axial_freqs, 'cache_clear'):
+                    module.get_axial_freqs.cache_clear()
+                    cleared_lru_count += 1
+        if cleared_lru_count > 0:
+            print(f"  ✅ Cleared {cleared_lru_count} LRU caches from RoPE modules.")
+
+    print("🎯 RoPE cache cleanup completed!")

modules/seedvr/src/optimization/performance.py

@@ -0,0 +1,160 @@
+"""
+Performance optimization module for SeedVR2
+Contains optimized tensor operations and video processing functions
+
+Extracted from: seedvr2.py (lines 1633-1730)
+"""
+
+import torch
+from typing import List, Union
+
+
+def optimized_video_rearrange(video_tensors: List[torch.Tensor]) -> List[torch.Tensor]:
+    """
+    🚀 OPTIMIZED version of video rearrangement
+    Replaces slow loops with vectorized operations
+    
+    Transforms:
+    - 3D: c h w -> t c h w (with t=1)  
+    - 4D: c t h w -> t c h w
+    
+    Expected gains: 5-10x faster than naive loops
+    
+    Args:
+        video_tensors: List of video tensors to rearrange
+        
+    Returns:
+        List of rearranged tensors in t c h w format
+    """
+    if not video_tensors:
+        return []
+
+    # 🔍 Analyze dimensions to optimize processing
+    videos_3d = []
+    videos_4d = []
+    indices_3d = []
+    indices_4d = []
+
+    for i, video in enumerate(video_tensors):
+        if video.ndim == 3:
+            videos_3d.append(video)
+            indices_3d.append(i)
+        else:  # ndim == 4
+            videos_4d.append(video)
+            indices_4d.append(i)
+
+    # 🎯 Prepare final result
+    samples = [None] * len(video_tensors)
+
+    # 🚀 BATCH PROCESSING for 3D videos (c h w -> 1 c h w)
+    if videos_3d:
+        # Method 1: Stack + permute (faster than rearrange)
+        # c h w -> c 1 h w -> 1 c h w
+        batch_3d = torch.stack([v.unsqueeze(1) for v in videos_3d])  # [batch, c, 1, h, w]
+        batch_3d = batch_3d.permute(0, 2, 1, 3, 4)  # [batch, 1, c, h, w]
+
+        for i, idx in enumerate(indices_3d):
+            samples[idx] = batch_3d[i]  # [1, c, h, w]
+
+    # 🚀 BATCH PROCESSING for 4D videos (c t h w -> t c h w)
+    if videos_4d:
+        # Check if all 4D videos have the same shape for maximum optimization
+        shapes = [v.shape for v in videos_4d]
+        if len(set(shapes)) == 1:
+            # 🎯 MAXIMUM OPTIMIZATION: All shapes identical
+            # Stack + permute in single operation
+            batch_4d = torch.stack(videos_4d)  # [batch, c, t, h, w]
+            batch_4d = batch_4d.permute(0, 2, 1, 3, 4)  # [batch, t, c, h, w]
+
+            for i, idx in enumerate(indices_4d):
+                samples[idx] = batch_4d[i]  # [t, c, h, w]
+        else:
+            # 🔄 FALLBACK: Different shapes, optimized individual processing
+            for i, idx in enumerate(indices_4d):
+                # Use permute instead of rearrange (faster)
+                samples[idx] = videos_4d[i].permute(1, 0, 2, 3)  # c t h w -> t c h w
+
+    return samples
+
+
+def optimized_single_video_rearrange(video: torch.Tensor) -> torch.Tensor:
+    """
+    🚀 OPTIMIZED version for single video tensor
+    Replaces rearrange() with native PyTorch operations
+    
+    Transforms:
+    - 3D: c h w -> 1 c h w (add temporal dimension)
+    - 4D: c t h w -> t c h w (permute dimensions)
+    
+    Expected gains: 2-5x faster than rearrange()
+    
+    Args:
+        video: Input video tensor
+        
+    Returns:
+        Rearranged tensor with temporal dimension first
+    """
+    if video.ndim == 3:
+        # c h w -> 1 c h w (add temporal dimension t=1)
+        return video.unsqueeze(0)
+    else:  # ndim == 4
+        # c t h w -> t c h w (permute channels and temporal)
+        return video.permute(1, 0, 2, 3)
+
+
+def optimized_sample_to_image_format(sample: torch.Tensor) -> torch.Tensor:
+    """
+    🚀 OPTIMIZED version to convert sample to image format
+    Replaces rearrange() with native PyTorch operations
+    
+    Transforms:
+    - 3D: c h w -> 1 h w c (add temporal dimension + permute to image format)
+    - 4D: t c h w -> t h w c (permute to image format)
+    
+    Expected gains: 2-5x faster than rearrange()
+    
+    Args:
+        sample: Input sample tensor
+        
+    Returns:
+        Tensor in image format (channels last)
+    """
+    if sample.ndim == 3:
+        # c h w -> 1 h w c (add temporal dimension then permute)
+        return sample.unsqueeze(0).permute(0, 2, 3, 1)
+    else:  # ndim == 4
+        # t c h w -> t h w c (permute channels to last)
+        return sample.permute(0, 2, 3, 1)
+
+
+def temporal_latent_blending(latents1: torch.Tensor, latents2: torch.Tensor, blend_frames: int) -> torch.Tensor:
+    """
+    🎨 Temporal blending in latent space to avoid discontinuities
+    
+    Args:
+        latents1: Latents from previous batch (end frames)
+        latents2: Latents from current batch (start frames)
+        blend_frames: Number of frames to blend
+    
+    Returns:
+        Blended latents for smooth transition
+    """
+    if latents1.shape[0] != latents2.shape[0]:
+        # Adjust dimensions if necessary
+        min_frames = min(latents1.shape[0], latents2.shape[0])
+        latents1 = latents1[:min_frames]
+        latents2 = latents2[:min_frames]
+
+    # Create linear blending weights
+    # Frame 0: 100% latents1, 0% latents2
+    # Frame n: 0% latents1, 100% latents2
+    weights1 = torch.linspace(1.0, 0.0, blend_frames).view(-1, 1, 1, 1).to(latents1.device)
+    weights2 = torch.linspace(0.0, 1.0, blend_frames).view(-1, 1, 1, 1).to(latents2.device)
+
+    # Apply blending
+    blended_latents = weights1 * latents1 + weights2 * latents2
+
+    return blended_latents
+
+
+

modules/seedvr/src/utils/color_fix.py

@@ -0,0 +1,138 @@
+import torch
+from PIL import Image
+from torch import Tensor
+from torch.nn import functional as F
+from src.common.half_precision_fixes import safe_pad_operation, safe_interpolate_operation
+from torchvision.transforms import ToTensor, ToPILImage
+
+def adain_color_fix(target: Image, source: Image):
+    # Convert images to tensors
+    to_tensor = ToTensor()
+    target_tensor = to_tensor(target).unsqueeze(0)
+    source_tensor = to_tensor(source).unsqueeze(0)
+
+    # Apply adaptive instance normalization
+    result_tensor = adaptive_instance_normalization(target_tensor, source_tensor)
+
+    # Convert tensor back to image
+    to_image = ToPILImage()
+    result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0))
+
+    return result_image
+
+def wavelet_color_fix(target: Image, source: Image):
+    # Convert images to tensors
+    to_tensor = ToTensor()
+    target_tensor = to_tensor(target).unsqueeze(0)
+    source_tensor = to_tensor(source).unsqueeze(0)
+
+    # Apply wavelet reconstruction
+    result_tensor = wavelet_reconstruction(target_tensor, source_tensor)
+
+    # Convert tensor back to image
+    to_image = ToPILImage()
+    result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0))
+
+    return result_image
+
+def calc_mean_std(feat: Tensor, eps=1e-5):
+    """Calculate mean and std for adaptive_instance_normalization.
+    Args:
+        feat (Tensor): 4D tensor.
+        eps (float): A small value added to the variance to avoid
+            divide-by-zero. Default: 1e-5.
+    """
+    size = feat.size()
+    assert len(size) == 4, 'The input feature should be 4D tensor.'
+    b, c = size[:2]
+    feat_var = feat.view(b, c, -1).var(dim=2) + eps
+    feat_std = feat_var.sqrt().view(b, c, 1, 1)
+    feat_mean = feat.view(b, c, -1).mean(dim=2).view(b, c, 1, 1)
+    return feat_mean, feat_std
+
+def adaptive_instance_normalization(content_feat:Tensor, style_feat:Tensor):
+    """Adaptive instance normalization.
+    Adjust the reference features to have the similar color and illuminations
+    as those in the degradate features.
+    Args:
+        content_feat (Tensor): The reference feature.
+        style_feat (Tensor): The degradate features.
+    """
+    size = content_feat.size()
+    style_mean, style_std = calc_mean_std(style_feat)
+    content_mean, content_std = calc_mean_std(content_feat)
+    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
+    return normalized_feat * style_std.expand(size) + style_mean.expand(size)
+
+def wavelet_blur(image: Tensor, radius: int):
+    """
+    Apply wavelet blur to the input tensor.
+    """
+    # input shape: (1, 3, H, W)
+    # convolution kernel
+    kernel_vals = [
+        [0.0625, 0.125, 0.0625],
+        [0.125, 0.25, 0.125],
+        [0.0625, 0.125, 0.0625],
+    ]
+    kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device)
+    # add channel dimensions to the kernel to make it a 4D tensor
+    kernel = kernel[None, None]
+    # repeat the kernel across all input channels
+    kernel = kernel.repeat(3, 1, 1, 1)
+    image = safe_pad_operation(image, (radius, radius, radius, radius), mode='replicate')
+    # apply convolution
+    output = F.conv2d(image, kernel, groups=3, dilation=radius)
+    return output
+
+def wavelet_decomposition(image: Tensor, levels=5):
+    """
+    Apply wavelet decomposition to the input tensor.
+    This function only returns the low frequency & the high frequency.
+    """
+    high_freq = torch.zeros_like(image)
+    for i in range(levels):
+        radius = 2 ** i
+        low_freq = wavelet_blur(image, radius)
+        high_freq += (image - low_freq)
+        image = low_freq
+
+    return high_freq, low_freq
+
+
+
+def wavelet_reconstruction(content_feat:Tensor, style_feat:Tensor):
+    """
+    Apply wavelet decomposition, so that the content will have the same color as the style.
+    """
+    # Vérifier et ajuster les dimensions si nécessaire
+    if content_feat.shape != style_feat.shape:
+        # Redimensionner style_feat pour correspondre à content_feat
+        target_shape = content_feat.shape
+        if len(target_shape) >= 3:  # Au moins 3 dimensions
+            # Utiliser interpolation pour ajuster les dimensions spatiales
+            style_feat = safe_interpolate_operation(
+                style_feat,
+                size=target_shape[-2:],  # Dernières 2 dimensions (H, W)
+                mode='bilinear',
+                align_corners=False
+            )
+
+    # calculate the wavelet decomposition of the content feature
+    content_high_freq, content_low_freq = wavelet_decomposition(content_feat)
+    del content_low_freq
+    # calculate the wavelet decomposition of the style feature
+    style_high_freq, style_low_freq = wavelet_decomposition(style_feat)
+    del style_high_freq
+
+    # Vérification finale avant addition
+    if content_high_freq.shape != style_low_freq.shape:
+        style_low_freq = safe_interpolate_operation(
+            style_low_freq,
+            size=content_high_freq.shape[-2:],
+            mode='bilinear',
+            align_corners=False
+        )
+
+    # reconstruct the content feature with the style's high frequency
+    return content_high_freq + style_low_freq

modules/seedvr/test.py

@@ -0,0 +1,58 @@
+import os
+import argparse
+import numpy as np
+import torch
+from PIL import Image
+from huggingface_hub import snapshot_download
+from torchvision.transforms import ToPILImage
+from src.core.generation import generation_loop
+from src.core.model_manager import configure_runner
+
+
+
+device = 'cuda'
+dtype = torch.bfloat16
+model_dir = 'seedvr2_models'
+model = 'seedvr2_ema_3b_fp16.safetensors'
+resolution = 1024
+seed = 100
+cfg = 1.0
+input_image = '/home/vlado/generative/Samples/cutie-512.png'
+
+to_pil = ToPILImage()
+runner = None
+loaded_model = None
+
+
+def upscale_image(model_name:str, image_path:str):
+    global runner, loaded_model
+    if (runner is None) or (loaded_model != model_name):
+        runner = configure_runner(model_name, model_dir, device=device, dtype=dtype)
+        loaded_model = model_name
+
+    image = Image.open(image_path).convert("RGB")
+    image_tensor = np.array(image)
+    image_tensor = torch.from_numpy(image_tensor).to(device=device, dtype=dtype).unsqueeze(0) / 255.0
+
+    result_tensor = generation_loop(
+        runner=runner,
+        images=image_tensor,
+        cfg_scale=cfg,
+        seed=seed,
+        res_w=resolution,
+        batch_size=1,
+        temporal_overlap=0,
+        device=device,
+    )
+    image = to_pil(result_tensor.squeeze().permute((2, 0, 1)))
+
+    output_path = os.path.join('/tmp', os.path.basename(image_path))
+
+    image.save(output_path, quality=95)
+    return image
+
+
+if __name__ == "__main__":
+    output_image = upscale_image(model, input_image)
+    print('input:', input_image)
+    print('output:', output_image)

modules/upscaler.py

@@ -79,7 +79,7 @@ class Upscaler:
                 scalers.append(scaler)
                 loaded.append(model_path)
                 # shared.log.debug(f'Upscaler type={self.name} folder="{self.user_path}" model="{model[0]}" path="{model_path}"')
-        if not os.path.exists(self.user_path):
+        if self.user_path is None or not os.path.exists(self.user_path):
             return scalers
         self.find_folder(self.user_path, scalers, loaded)
         return scalers
@@ -148,7 +148,7 @@ class UpscalerData:
     scaler: Upscaler = None
     model: None
 
-    def __init__(self, name: str, path: str, upscaler: Upscaler = None, scale: int = 4, model=None):
+    def __init__(self, name: str, path: str = None, upscaler: Upscaler = None, scale: int = 4, model=None):
         self.name = name
         self.data_path = path
         self.local_data_path = path