Upload folder using huggingface_hub

models/__init__.py

@@ -0,0 +1,38 @@
+"""Inference-facing model registry for StreamDiffusionV2."""
+
+from .wan.wan_wrapper import (
+    CausalWanDiffusionWrapper,
+    WanDiffusionWrapper,
+    WanTextEncoder,
+    WanVAEWrapper,
+)
+
+
+DIFFUSION_NAME_TO_CLASS = {
+    "wan": WanDiffusionWrapper,
+    "causal_wan": CausalWanDiffusionWrapper,
+}
+
+
+TEXT_ENCODER_NAME_TO_CLASS = {
+    "wan": WanTextEncoder,
+    "causal_wan": WanTextEncoder,
+}
+
+
+VAE_NAME_TO_CLASS = {
+    "wan": WanVAEWrapper,
+    "causal_wan": WanVAEWrapper,
+}
+
+
+def get_diffusion_wrapper(model_name):
+    return DIFFUSION_NAME_TO_CLASS[model_name]
+
+
+def get_text_encoder_wrapper(model_name):
+    return TEXT_ENCODER_NAME_TO_CLASS[model_name]
+
+
+def get_vae_wrapper(model_name):
+    return VAE_NAME_TO_CLASS[model_name]

models/data.py

@@ -0,0 +1,17 @@
+"""Minimal dataset helpers used by inference entrypoints."""
+
+from torch.utils.data import Dataset
+
+
+class TextDataset(Dataset):
+    """Load one prompt per line from a UTF-8 text file."""
+
+    def __init__(self, data_path: str):
+        with open(data_path, "r", encoding="utf-8") as handle:
+            self.texts = [line.strip() for line in handle]
+
+    def __len__(self) -> int:
+        return len(self.texts)
+
+    def __getitem__(self, idx: int) -> str:
+        return self.texts[idx]

models/model_interface.py

@@ -0,0 +1,114 @@
+from models.scheduler import SchedulerInterface
+from abc import abstractmethod, ABC
+from typing import List, Optional
+import torch
+import types
+
+
+class DiffusionModelInterface(ABC, torch.nn.Module):
+    scheduler: SchedulerInterface
+
+    @abstractmethod
+    def forward(
+        self, noisy_image_or_video: torch.Tensor, conditional_dict: dict,
+        timestep: torch.Tensor, kv_cache: Optional[List[dict]] = None,
+        crossattn_cache: Optional[List[dict]] = None,
+        current_start: Optional[int] = None,
+        current_end: Optional[int] = None
+    ) -> torch.Tensor:
+        """
+        A method to run diffusion model.
+        Input:
+            - noisy_image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - timestep: a tensor with shape [B, F]  where the number of frame is 1 for images.
+            all data should be on the same device as the model.
+            - kv_cache: a list of dictionaries containing the key and value tensors for each attention layer.
+            - current_start: the start index of the current frame in the sequence.
+            - current_end: the end index of the current frame in the sequence.
+        Output: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+        We always expect a X0 prediction form for the output.
+        """
+        pass
+
+    def get_scheduler(self) -> SchedulerInterface:
+        """
+        Update the current scheduler with the interface's static method
+        """
+        scheduler = self.scheduler
+        scheduler.convert_x0_to_noise = types.MethodType(
+            SchedulerInterface.convert_x0_to_noise, scheduler)
+        scheduler.convert_noise_to_x0 = types.MethodType(
+            SchedulerInterface.convert_noise_to_x0, scheduler)
+        scheduler.convert_velocity_to_x0 = types.MethodType(
+            SchedulerInterface.convert_velocity_to_x0, scheduler)
+        self.scheduler = scheduler
+        return scheduler
+
+    def post_init(self):
+        """
+        A few custom initialization steps that should be called after the object is created.
+        Currently, the only one we have is to bind a few methods to scheduler.
+        We can gradually add more methods here if needed.
+        """
+        self.get_scheduler()
+
+    def set_module_grad(self, module_grad: dict) -> None:
+        """
+        Adjusts the state of each module in the object.
+
+        Parameters:
+        - module_grad (dict): A dictionary where each key is the name of a module (as an attribute of the object),
+          and each value is a bool indicating whether the module's parameters require gradients.
+
+        Functionality:
+        For each module name in the dictionary:
+        - Updates whether its parameters require gradients based on 'is_trainable'.
+        """
+        for k, is_trainable in module_grad.items():
+            getattr(self, k).requires_grad_(is_trainable)
+
+    @abstractmethod
+    def enable_gradient_checkpointing(self) -> None:
+        """
+        Activates gradient checkpointing for the current model (may be referred to as *activation checkpointing* or
+        *checkpoint activations* in other frameworks).
+        """
+        pass
+
+
+class VAEInterface(ABC, torch.nn.Module):
+    @abstractmethod
+    def decode_to_pixel(self, latent: torch.Tensor) -> torch.Tensor:
+        """
+        A method to decode a latent representation to an image or video.
+        Input: a tensor with shape [B, F // T, C, H // S, W // S] where T and S are temporal and spatial compression factors.
+        Output: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+        """
+        pass
+
+
+class TextEncoderInterface(ABC, torch.nn.Module):
+    @abstractmethod
+    def forward(self, text_prompts: List[str]) -> dict:
+        """
+        A method to tokenize text prompts with a tokenizer and encode them into a latent representation.
+        Input: a list of strings.
+        Output: a dictionary containing the encoded representation of the text prompts.
+        """
+        pass
+
+
+class InferencePipelineInterface(ABC):
+    @abstractmethod
+    def inference_with_trajectory(self, noise: torch.Tensor, conditional_dict: dict) -> torch.Tensor:
+        """
+        Run inference with the given diffusion / distilled generators.
+        Input:
+            - noise: a tensor sampled from N(0, 1) with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+        Output:
+            - output: a tensor with shape [B, T, F, C, H, W].
+            T is the total number of timesteps. output[0] is a pure noise and output[i] and i>0
+            represents the x0 prediction at each timestep.
+        """

models/scheduler.py

@@ -0,0 +1,103 @@
+from abc import abstractmethod, ABC
+import torch
+
+
+class SchedulerInterface(ABC):
+    """
+    Base class for diffusion noise schedule.
+    """
+    alphas_cumprod: torch.Tensor  # [T], alphas for defining the noise schedule
+
+    @abstractmethod
+    def add_noise(
+        self, clean_latent: torch.Tensor,
+        noise: torch.Tensor, timestep: torch.Tensor
+    ):
+        """
+        Diffusion forward corruption process.
+        Input:
+            - clean_latent: the clean latent with shape [B, C, H, W]
+            - noise: the noise with shape [B, C, H, W]
+            - timestep: the timestep with shape [B]
+        Output: the corrupted latent with shape [B, C, H, W]
+        """
+        pass
+
+    def convert_x0_to_noise(
+        self, x0: torch.Tensor, xt: torch.Tensor,
+        timestep: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Convert the diffusion network's x0 prediction to noise predidction.
+        x0: the predicted clean data with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        noise = (xt-sqrt(alpha_t)*x0) / sqrt(beta_t) (eq 11 in https://arxiv.org/abs/2311.18828)
+        """
+        # use higher precision for calculations
+        original_dtype = x0.dtype
+        x0, xt, alphas_cumprod = map(
+            lambda x: x.double().to(x0.device), [x0, xt,
+                                                 self.alphas_cumprod]
+        )
+
+        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        noise_pred = (xt - alpha_prod_t **
+                      (0.5) * x0) / beta_prod_t ** (0.5)
+        return noise_pred.to(original_dtype)
+
+    def convert_noise_to_x0(
+        self, noise: torch.Tensor, xt: torch.Tensor,
+        timestep: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Convert the diffusion network's noise prediction to x0 predidction.
+        noise: the predicted noise with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        x0 = (x_t - sqrt(beta_t) * noise) / sqrt(alpha_t) (eq 11 in https://arxiv.org/abs/2311.18828)
+        """
+        # use higher precision for calculations
+        original_dtype = noise.dtype
+        noise, xt, alphas_cumprod = map(
+            lambda x: x.double().to(noise.device), [noise, xt,
+                                                    self.alphas_cumprod]
+        )
+        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        x0_pred = (xt - beta_prod_t **
+                   (0.5) * noise) / alpha_prod_t ** (0.5)
+        return x0_pred.to(original_dtype)
+
+    def convert_velocity_to_x0(
+        self, velocity: torch.Tensor, xt: torch.Tensor,
+        timestep: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Convert the diffusion network's velocity prediction to x0 predidction.
+        velocity: the predicted noise with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        v = sqrt(alpha_t) * noise - sqrt(beta_t) x0
+        noise = (xt-sqrt(alpha_t)*x0) / sqrt(beta_t)
+        given v, x_t, we have
+        x0 = sqrt(alpha_t) * x_t - sqrt(beta_t) * v
+        see derivations https://chatgpt.com/share/679fb6c8-3a30-8008-9b0e-d1ae892dac56
+        """
+        # use higher precision for calculations
+        original_dtype = velocity.dtype
+        velocity, xt, alphas_cumprod = map(
+            lambda x: x.double().to(velocity.device), [velocity, xt,
+                                                       self.alphas_cumprod]
+        )
+        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        x0_pred = (alpha_prod_t ** 0.5) * xt - (beta_prod_t ** 0.5) * velocity
+        return x0_pred.to(original_dtype)

models/util.py

@@ -0,0 +1,17 @@
+"""Small runtime helpers used by the offline and online inference entrypoints."""
+
+import random
+
+import numpy as np
+import torch
+
+
+def set_seed(seed: int, deterministic: bool = False) -> None:
+    """Seed Python, NumPy, and PyTorch for reproducible inference."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    if deterministic:
+        torch.use_deterministic_algorithms(True)

models/wan/causal_model.py

@@ -0,0 +1,1013 @@
+from models.wan.wan_base.modules.attention import attention
+from models.wan.wan_base.modules.model import (
+    WanRMSNorm,
+    rope_apply,
+    WanLayerNorm,
+    WAN_CROSSATTENTION_CLASSES,
+    Head,
+    rope_params,
+    MLPProj,
+    sinusoidal_embedding_1d
+)
+from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from torch.nn.attention.flex_attention import BlockMask
+from diffusers.models.modeling_utils import ModelMixin
+import torch.nn.functional as F
+import torch.nn as nn
+import torch
+import math
+from collections import OrderedDict
+import torch.distributed as dist
+import warnings
+
+try:
+    import flash_attn_interface
+    FLASH_ATTN_AVAILABLE = True
+except (ImportError, ModuleNotFoundError):
+    try:
+        from flash_attn import flash_attn_interface
+        FLASH_ATTN_AVAILABLE = True
+    except (ImportError, ModuleNotFoundError):
+        flash_attn_interface = None
+        FLASH_ATTN_AVAILABLE = False
+
+# wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention
+# see https://github.com/pytorch/pytorch/issues/133254
+# change to default for other models
+flex_attention = torch.compile(
+    flex_attention, dynamic=False, mode="max-autotune")
+
+_CAUSAL_ROPE_FREQ_CACHE = OrderedDict()
+_CAUSAL_ROPE_FREQ_CACHE_SIZE = 16
+
+
+def _causal_rope_cache_key(freqs, f, h, w, start_frame, device):
+    return (
+        freqs,
+        device.type,
+        device.index,
+        f,
+        h,
+        w,
+        start_frame,
+    )
+
+
+def _get_causal_rope_freqs(freqs_source, freqs_parts, f, h, w, start_frame, device):
+    key = _causal_rope_cache_key(freqs_source, f, h, w, start_frame, device)
+    cached = _CAUSAL_ROPE_FREQ_CACHE.get(key)
+    if cached is not None:
+        _CAUSAL_ROPE_FREQ_CACHE.move_to_end(key)
+        return cached
+
+    temporal, height, width = freqs_parts
+    temporal_freqs = temporal[start_frame:start_frame + f].repeat_interleave(h * w, dim=0)
+    height_freqs = height[:h].repeat_interleave(w, dim=0).repeat(f, 1)
+    width_freqs = width[:w].repeat(h, 1).repeat(f, 1)
+    rope_freqs = torch.cat([temporal_freqs, height_freqs, width_freqs], dim=-1).unsqueeze(1)
+
+    _CAUSAL_ROPE_FREQ_CACHE[key] = rope_freqs
+    if len(_CAUSAL_ROPE_FREQ_CACHE) > _CAUSAL_ROPE_FREQ_CACHE_SIZE:
+        _CAUSAL_ROPE_FREQ_CACHE.popitem(last=False)
+    return rope_freqs
+
+
+def _prepare_causal_rope_cache(grid_sizes, freqs, start_frame=0):
+    c = freqs.shape[1]
+    freqs_parts = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+    if isinstance(start_frame, torch.Tensor):
+        start_frames = start_frame.tolist()
+    else:
+        start_frames = [int(start_frame)] * grid_sizes.shape[0]
+
+    rope_cache = []
+    for grid_size, sf in zip(grid_sizes.tolist(), start_frames):
+        f, h, w = grid_size
+        seq_len = f * h * w
+        rope_freqs = _get_causal_rope_freqs(freqs, freqs_parts, f, h, w, sf, freqs.device)
+        rope_cache.append((seq_len, rope_freqs))
+    return rope_cache
+
+
+def causal_rope_apply(x, grid_sizes, freqs, start_frame=0, rope_cache=None):
+    n = x.size(2)
+
+    if rope_cache is None:
+        rope_cache = _prepare_causal_rope_cache(grid_sizes, freqs, start_frame=start_frame)
+
+    output = x.clone()
+
+    for i, (seq_len, freqs_i) in enumerate(rope_cache):
+        x_i = torch.view_as_complex(
+            x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)
+        )
+        output[i, :seq_len] = torch.view_as_real(x_i * freqs_i).flatten(2).type_as(x)
+
+    return output
+
+
+def attention_with_kvcache_fallback(q, k_cache, v_cache, cache_seqlens):
+    out_dtype = q.dtype
+    max_seq_len = k_cache.shape[1]
+
+    def prepare_inputs(q_tensor, k_tensor, v_tensor):
+        if q_tensor.device.type == "cpu" and q_tensor.dtype in (torch.float16, torch.bfloat16):
+            q_tensor = q_tensor.float()
+            k_tensor = k_tensor.float()
+            v_tensor = v_tensor.float()
+        return q_tensor, k_tensor, v_tensor
+
+    # Fast path: every sample uses the same fully valid cache span.
+    if torch.all(cache_seqlens == max_seq_len):
+        q_all = q.transpose(1, 2)
+        k_all = k_cache.transpose(1, 2)
+        v_all = v_cache.transpose(1, 2)
+        q_all, k_all, v_all = prepare_inputs(q_all, k_all, v_all)
+        x = F.scaled_dot_product_attention(
+            q_all,
+            k_all,
+            v_all,
+            attn_mask=None,
+            dropout_p=0.0,
+            # Keep parity with flash_attn_with_kvcache(..., causal=False).
+            is_causal=False,
+        )
+        return x.transpose(1, 2).to(out_dtype).contiguous()
+
+    outputs = []
+    for batch_idx, seq_len in enumerate(cache_seqlens.tolist()):
+        q_i = q[batch_idx:batch_idx + 1].transpose(1, 2)
+        k_i = k_cache[batch_idx:batch_idx + 1, :seq_len].transpose(1, 2)
+        v_i = v_cache[batch_idx:batch_idx + 1, :seq_len].transpose(1, 2)
+        q_i, k_i, v_i = prepare_inputs(q_i, k_i, v_i)
+
+        x_i = F.scaled_dot_product_attention(
+            q_i,
+            k_i,
+            v_i,
+            attn_mask=None,
+            dropout_p=0.0,
+            # Keep parity with flash_attn_with_kvcache(..., causal=False).
+            is_causal=False,
+        )
+        outputs.append(x_i.transpose(1, 2).to(out_dtype))
+
+    return torch.cat(outputs, dim=0).contiguous()
+
+
+class CausalWanSelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 eps=1e-6):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+        
+        self.sink_size = 3
+        self.adapt_sink_thr = -1
+        self.evict_idx = None
+
+        # layers
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(
+        self,
+        x,
+        seq_lens,
+        grid_sizes,
+        freqs,
+        block_mask,
+        kv_cache=None,
+        current_start=0,
+        current_end=0,
+        causal_rope_cache=None,
+    ):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, num_heads, C / num_heads]
+            seq_lens(Tensor): Shape [B]
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+            block_mask (BlockMask)
+        """
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+
+        # query, key, value function
+        def qkv_fn(x):
+            q = self.norm_q(self.q(x)).view(b, s, n, d)
+            k = self.norm_k(self.k(x)).view(b, s, n, d)
+            v = self.v(x).view(b, s, n, d)
+            return q, k, v
+
+        q, k, v = qkv_fn(x)
+
+        if kv_cache is None:
+            roped_query = rope_apply(q, grid_sizes, freqs).type_as(v)
+            roped_key = rope_apply(k, grid_sizes, freqs).type_as(v)
+
+            padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]
+            padded_roped_query = torch.cat(
+                [roped_query,
+                 torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]],
+                             device=q.device, dtype=v.dtype)],
+                dim=1
+            )
+
+            padded_roped_key = torch.cat(
+                [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]],
+                                        device=k.device, dtype=v.dtype)],
+                dim=1
+            )
+
+            padded_v = torch.cat(
+                [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]],
+                                device=v.device, dtype=v.dtype)],
+                dim=1
+            )
+
+            x = flex_attention(
+                query=padded_roped_query.transpose(2, 1),
+                key=padded_roped_key.transpose(2, 1),
+                value=padded_v.transpose(2, 1),
+                block_mask=block_mask
+            )[:, :, :-padded_length].transpose(2, 1)
+        else:
+            frame_seqlen = math.prod(grid_sizes[0][1:]).item()
+            current_start_frame = current_start // frame_seqlen
+            roped_query = causal_rope_apply(
+                q,
+                grid_sizes,
+                freqs,
+                start_frame=current_start_frame,
+                rope_cache=causal_rope_cache,
+            ).type_as(v)
+            roped_key = causal_rope_apply(
+                k,
+                grid_sizes,
+                freqs,
+                start_frame=current_start_frame,
+                rope_cache=causal_rope_cache,
+            ).type_as(v)
+
+            seq_lens = []
+            kv_cache_size = kv_cache["k"].shape[1]
+            cache_bs = kv_cache['k'].shape[0]
+            
+            # Ring-buffer queue init
+            if self.evict_idx is None:
+                self.evict_idx = [[]]
+
+            if len(self.evict_idx) < cache_bs:
+                self.evict_idx = [self.evict_idx[0].copy() for _ in range(cache_bs)]
+                
+            for i, c_start in enumerate(current_start):
+                num_new_tokens = roped_query.shape[1]
+                current_end = c_start + roped_query.shape[1]
+                sink_tokens = self.sink_size * frame_seqlen
+                
+                if sink_tokens > 0 and self.adapt_sink_thr > -1 and v.shape[1] <= frame_seqlen:
+                    # Caculate similarity between new keys/values and the oldest ones in the cache
+                    k_sink_mean = kv_cache["k"][i:i+1, :sink_tokens].reshape(self.sink_size, frame_seqlen, -1).mean(1)
+                    k_new_mean = roped_key[i:i+1].reshape(1, frame_seqlen, -1).mean(1)
+                    k_cos_sim = torch.cosine_similarity(k_sink_mean, k_new_mean, dim=-1)
+
+                    v_sink_mean = kv_cache["v"][i:i+1, :sink_tokens].reshape(self.sink_size, frame_seqlen, -1).mean(1)
+                    v_new_mean = v[i:i+1].reshape(1, frame_seqlen, -1).mean(1)
+                    v_cos_sim = torch.cosine_similarity(v_sink_mean, v_new_mean, dim=-1)
+
+                    avg_cos_sim = (k_cos_sim + v_cos_sim)/2
+                    # When the similarity is low, refresh the sink
+                    if avg_cos_sim.min() < self.adapt_sink_thr:
+                        idx = torch.argmin(avg_cos_sim).item()
+                        temp_evict_idx = (idx+1) * frame_seqlen
+                        self.evict_idx[i].insert(0, temp_evict_idx)
+
+                # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache
+                if current_end > kv_cache_size or kv_cache["local_end_index"][i]>=kv_cache_size:
+                    kv_cache["global_end_index"][i].fill_(c_start)
+                    kv_cache["local_end_index"][i].fill_(kv_cache_size)
+
+                    target_end = self.evict_idx[i][0]
+
+                    # current_step = kv_cache['current_step']
+                    # Update the buffer
+                    if cache_bs==1 and kv_cache['current_step'] > 1:
+                        kv_cache['current_step']-=1
+                    else:
+                        evict_idx = self.evict_idx[i].pop(0)
+                        if evict_idx > sink_tokens:
+                            self.evict_idx[i].append(evict_idx)
+                        kv_cache['current_step']=kv_cache['total_steps']
+
+                    # print(f"self.evict_idx: {self.evict_idx[i]}, total steps: {kv_cache['total_steps']}, current step: {current_step}, target: {target_end-num_new_tokens}:{target_end}, kv size:{kv_cache_size}")
+
+                    # Newly added cache covers the oldest one
+                    kv_cache["k"][i:i+1, target_end-num_new_tokens:target_end] = roped_key[i:i+1]
+                    kv_cache["v"][i:i+1, target_end-num_new_tokens:target_end] = v[i:i+1]
+
+                    local_end_index = kv_cache["local_end_index"][i].item()
+
+                else:
+                    local_end_index = kv_cache["local_end_index"][i].item() + current_end - kv_cache["global_end_index"][i].item()
+
+                    rolling_end = (current_end + num_new_tokens).item()
+                    if rolling_end > self.sink_size * frame_seqlen and rolling_end <= kv_cache_size \
+                        and (not self.evict_idx[i] or self.evict_idx[i][-1] != rolling_end):
+                        self.evict_idx[i].append(rolling_end)
+
+                    local_start_index = local_end_index - num_new_tokens
+                    # print(f"target: {local_start_index}:{local_end_index}")
+                    kv_cache["k"][i:i+1, local_start_index:local_end_index] = roped_key[i:i+1]
+                    kv_cache["v"][i:i+1, local_start_index:local_end_index] = v[i:i+1]
+
+                seq_lens.append(local_end_index)
+
+                kv_cache["global_end_index"][i].fill_(current_end)
+                kv_cache["local_end_index"][i].fill_(local_end_index)
+            
+            seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=roped_query.device)
+
+            max_seq_len = int(seq_lens.max().item())
+            k_cache = kv_cache["k"][:, :max_seq_len]
+            v_cache = kv_cache["v"][:, :max_seq_len]
+
+            if FLASH_ATTN_AVAILABLE:
+                try:
+                    with torch.cuda.device(roped_query.device):
+                        x = flash_attn_interface.flash_attn_with_kvcache(
+                            q=roped_query,
+                            k_cache=k_cache,
+                            v_cache=v_cache,
+                            cache_seqlens=seq_lens,
+                        )
+                except RuntimeError as exc:
+                    if "DeviceType::CUDA" not in str(exc):
+                        raise
+                    warnings.warn(
+                        "flash_attn_with_kvcache failed on the current GPU; "
+                        "falling back to scaled_dot_product_attention.",
+                        stacklevel=2,
+                    )
+                    x = attention_with_kvcache_fallback(
+                        q=roped_query,
+                        k_cache=k_cache,
+                        v_cache=v_cache,
+                        cache_seqlens=seq_lens,
+                    )
+            else:
+                warnings.warn(
+                    "flash_attn is not installed; falling back to "
+                    "scaled_dot_product_attention for KV-cache attention.",
+                    stacklevel=2,
+                )
+                x = attention_with_kvcache_fallback(
+                    q=roped_query,
+                    k_cache=k_cache,
+                    v_cache=v_cache,
+                    cache_seqlens=seq_lens,
+                )
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class CausalWanAttentionBlock(nn.Module):
+
+    def __init__(self,
+                 cross_attn_type,
+                 dim,
+                 ffn_dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 cross_attn_norm=False,
+                 eps=1e-6):
+        super().__init__()
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        # layers
+        self.norm1 = WanLayerNorm(dim, eps)
+        self.self_attn = CausalWanSelfAttention(dim, num_heads, window_size, qk_norm,
+                                                eps)
+        self.norm3 = WanLayerNorm(
+            dim, eps,
+            elementwise_affine=True) if cross_attn_norm else nn.Identity()
+        self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim,
+                                                                      num_heads,
+                                                                      (-1, -1),
+                                                                      qk_norm,
+                                                                      eps)
+        self.norm2 = WanLayerNorm(dim, eps)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
+            nn.Linear(ffn_dim, dim))
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        x,
+        e,
+        seq_lens,
+        grid_sizes,
+        freqs,
+        context,
+        context_lens,
+        block_mask,
+        kv_cache=None,
+        crossattn_cache=None,
+        current_start=0,
+        current_end=0,
+        causal_rope_cache=None,
+    ):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+            e(Tensor): Shape [B, F, 6, C]
+            seq_lens(Tensor): Shape [B], length of each sequence in batch
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
+        # assert e.dtype == torch.float32
+        # with amp.autocast(dtype=torch.float32):
+        e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2)
+        # assert e[0].dtype == torch.float32
+
+        # self-attention
+        y = self.self_attn(
+            (self.norm1(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen))
+             * (1 + e[1]) + e[0]).flatten(1, 2),
+            seq_lens, grid_sizes,
+            freqs, block_mask, kv_cache, current_start, current_end, causal_rope_cache)
+
+        # with amp.autocast(dtype=torch.float32):
+        x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen))
+                 * e[2]).flatten(1, 2)
+
+        # cross-attention & ffn function
+        def cross_attn_ffn(x, context, context_lens, e, crossattn_cache=None):
+            x = x + self.cross_attn(self.norm3(x), context,
+                                    context_lens, crossattn_cache=crossattn_cache)
+            y = self.ffn(
+                (self.norm2(x).unflatten(dim=1, sizes=(num_frames,
+                 frame_seqlen)) * (1 + e[4]) + e[3]).flatten(1, 2)
+            )
+            # with amp.autocast(dtype=torch.float32):
+            x = x + (y.unflatten(dim=1, sizes=(num_frames,
+                     frame_seqlen)) * e[5]).flatten(1, 2)
+            return x
+
+        x = cross_attn_ffn(x, context, context_lens, e, crossattn_cache)
+        return x
+
+
+class CausalHead(nn.Module):
+
+    def __init__(self, dim, out_dim, patch_size, eps=1e-6):
+        super().__init__()
+        self.dim = dim
+        self.out_dim = out_dim
+        self.patch_size = patch_size
+        self.eps = eps
+
+        # layers
+        out_dim = math.prod(patch_size) * out_dim
+        self.norm = WanLayerNorm(dim, eps)
+        self.head = nn.Linear(dim, out_dim)
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x, e):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            e(Tensor): Shape [B, F, 1, C]
+        """
+        # assert e.dtype == torch.float32
+        # with amp.autocast(dtype=torch.float32):
+        num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
+        e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2)
+        x = (self.head(
+            self.norm(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
+            (1 + e[1]) + e[0]))
+        return x
+
+
+class CausalWanModel(ModelMixin, ConfigMixin):
+    r"""
+    Wan diffusion backbone supporting both text-to-video and image-to-video.
+    """
+
+    ignore_for_config = [
+        'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size'
+    ]
+    _no_split_modules = ['WanAttentionBlock']
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(self,
+                 model_type='t2v',
+                 patch_size=(1, 2, 2),
+                 text_len=512,
+                 in_dim=16,
+                 dim=2048,
+                 ffn_dim=8192,
+                 freq_dim=256,
+                 text_dim=4096,
+                 out_dim=16,
+                 num_heads=16,
+                 num_layers=32,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 cross_attn_norm=True,
+                 eps=1e-6):
+        r"""
+        Initialize the diffusion model backbone.
+
+        Args:
+            model_type (`str`, *optional*, defaults to 't2v'):
+                Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+            patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+                3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+            text_len (`int`, *optional*, defaults to 512):
+                Fixed length for text embeddings
+            in_dim (`int`, *optional*, defaults to 16):
+                Input video channels (C_in)
+            dim (`int`, *optional*, defaults to 2048):
+                Hidden dimension of the transformer
+            ffn_dim (`int`, *optional*, defaults to 8192):
+                Intermediate dimension in feed-forward network
+            freq_dim (`int`, *optional*, defaults to 256):
+                Dimension for sinusoidal time embeddings
+            text_dim (`int`, *optional*, defaults to 4096):
+                Input dimension for text embeddings
+            out_dim (`int`, *optional*, defaults to 16):
+                Output video channels (C_out)
+            num_heads (`int`, *optional*, defaults to 16):
+                Number of attention heads
+            num_layers (`int`, *optional*, defaults to 32):
+                Number of transformer blocks
+            window_size (`tuple`, *optional*, defaults to (-1, -1)):
+                Window size for local attention (-1 indicates global attention)
+            qk_norm (`bool`, *optional*, defaults to True):
+                Enable query/key normalization
+            cross_attn_norm (`bool`, *optional*, defaults to False):
+                Enable cross-attention normalization
+            eps (`float`, *optional*, defaults to 1e-6):
+                Epsilon value for normalization layers
+        """
+
+        super().__init__()
+
+        assert model_type in ['t2v', 'i2v']
+        self.model_type = model_type
+
+        self.patch_size = patch_size
+        self.text_len = text_len
+        self.in_dim = in_dim
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.freq_dim = freq_dim
+        self.text_dim = text_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        # embeddings
+        self.patch_embedding = nn.Conv3d(
+            in_dim, dim, kernel_size=patch_size, stride=patch_size)
+        self.text_embedding = nn.Sequential(
+            nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
+            nn.Linear(dim, dim))
+
+        self.time_embedding = nn.Sequential(
+            nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.time_projection = nn.Sequential(
+            nn.SiLU(), nn.Linear(dim, dim * 6))
+
+        # blocks
+        cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'
+        self.blocks = nn.ModuleList([
+            CausalWanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,
+                                    window_size, qk_norm, cross_attn_norm, eps)
+            for _ in range(num_layers)
+        ])
+
+        # head
+        self.head = CausalHead(dim, out_dim, patch_size, eps)
+
+        # buffers (don't use register_buffer otherwise dtype will be changed in to())
+        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+        d = dim // num_heads
+        self.freqs = torch.cat([
+            rope_params(1024, d - 4 * (d // 6)),
+            rope_params(1024, 2 * (d // 6)),
+            rope_params(1024, 2 * (d // 6))
+        ],
+            dim=1)
+
+        if model_type == 'i2v':
+            self.img_emb = MLPProj(1280, dim)
+
+        # initialize weights
+        self.init_weights()
+
+        self.gradient_checkpointing = False
+
+        self.block_mask = None
+
+        self.num_frame_per_block = 1
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        self.gradient_checkpointing = value
+
+    @staticmethod
+    def _prepare_blockwise_causal_attn_mask(
+        device: torch.device | str, num_frames: int = 21,
+        frame_seqlen: int = 1560, num_frame_per_block=1
+    ) -> BlockMask:
+        """
+        we will divide the token sequence into the following format
+        [1 latent frame] [1 latent frame] ... [1 latent frame]
+        We use flexattention to construct the attention mask
+        """
+        total_length = num_frames * frame_seqlen
+
+        # we do right padding to get to a multiple of 128
+        padded_length = math.ceil(total_length / 128) * 128 - total_length
+
+        ends = torch.zeros(total_length + padded_length,
+                           device=device, dtype=torch.long)
+
+        # Block-wise causal mask will attend to all elements that are before the end of the current chunk
+        frame_indices = torch.arange(
+            start=0,
+            end=total_length,
+            step=frame_seqlen * num_frame_per_block,
+            device=device
+        )
+
+        for tmp in frame_indices:
+            ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \
+                frame_seqlen * num_frame_per_block
+
+        def attention_mask(b, h, q_idx, kv_idx):
+            return (kv_idx < ends[q_idx]) | (q_idx == kv_idx)
+            # return ((kv_idx < total_length) & (q_idx < total_length))  | (q_idx == kv_idx) # bidirectional mask
+
+        block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length,
+                                       KV_LEN=total_length + padded_length, _compile=False, device=device)
+
+        import torch.distributed as dist
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(
+                f" cache a block wise causal mask with block size of {num_frame_per_block} frames")
+            print(block_mask)
+
+        return block_mask
+
+    def _forward_inference(
+        self,
+        x,
+        t,
+        context,
+        seq_len,
+        clip_fea=None,
+        y=None,
+        kv_cache: dict = None,
+        crossattn_cache: dict = None,
+        current_start: int = 0,
+        current_end: int = 0,
+        block_mode: str = 'input',
+        block_num: int = [-1],
+        patched_x_shape: torch.Tensor = None,
+    ):
+        r"""
+        Run the diffusion model with kv caching.
+        See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details.
+        This function will be run for num_frame times.
+        Process the latent frames one by one (1560 tokens each)
+
+        Args:
+            x (List[Tensor]):
+                List of input video tensors, each with shape [C_in, F, H, W]
+            t (Tensor):
+                Diffusion timesteps tensor of shape [B]
+            context (List[Tensor]):
+                List of text embeddings each with shape [L, C]
+            seq_len (`int`):
+                Maximum sequence length for positional encoding
+            clip_fea (Tensor, *optional*):
+                CLIP image features for image-to-video mode
+            y (List[Tensor], *optional*):
+                Conditional video inputs for image-to-video mode, same shape as x
+
+        Returns:
+            List[Tensor]:
+                List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+        """
+        if self.model_type == 'i2v':
+            assert clip_fea is not None and y is not None
+        # params
+        device = self.patch_embedding.weight.device
+        if self.freqs.device != device:
+            self.freqs = self.freqs.to(device)
+
+        if block_mode == 'input':
+            if y is not None:
+                x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+            # embeddings
+            x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+            bsz, cch, tlen, hh, ww = x[0].shape
+            patched_x_shape = torch.tensor([bsz, cch, tlen, hh, ww], dtype=torch.int64, device=device)
+        else:
+            bsz, cch, tlen, hh, ww = [int(i) for i in patched_x_shape.tolist()]
+            x = [u.permute(1,0).reshape(bsz, cch, tlen, hh, ww) for u in x]
+            
+        grid_sizes = torch.stack(
+            [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+        x = [u.flatten(2).transpose(1, 2) for u in x]
+        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+        assert seq_lens.max() <= seq_len
+        x = torch.cat(x)
+        """
+        torch.cat([
+            torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
+                      dim=1) for u in x
+        ])
+        """
+
+        # time embeddings
+        # with amp.autocast(dtype=torch.float32):
+        e = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x))
+        e0 = self.time_projection(e).unflatten(
+            1, (6, self.dim)).unflatten(dim=0, sizes=t.shape)
+        # assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+        # context
+        context_lens = None
+        context = self.text_embedding(
+            torch.stack([
+                torch.cat(
+                    [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+                for u in context
+            ]))
+
+        if clip_fea is not None:
+            context_clip = self.img_emb(clip_fea)  # bs x 257 x dim
+            context = torch.concat([context_clip, context], dim=1)
+
+        # arguments
+        kwargs = dict(
+            e=e0,
+            seq_lens=seq_lens,
+            grid_sizes=grid_sizes,
+            freqs=self.freqs,
+            context=context,
+            context_lens=context_lens,
+            block_mask=self.block_mask
+        )
+        if kv_cache is not None:
+            kwargs["causal_rope_cache"] = _prepare_causal_rope_cache(
+                grid_sizes,
+                self.freqs,
+                start_frame=current_start // math.prod(grid_sizes[0][1:]).item(),
+            )
+
+        def create_custom_forward(module):
+            def custom_forward(*inputs, **kwargs):
+                return module(*inputs, **kwargs)
+            return custom_forward
+        
+        for block_index, block in enumerate(self.blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                assert False
+            else:
+                if (block_mode == 'output' or block_mode == 'middle') and block_index < block_num[0]:
+                    continue
+                if (block_mode == 'input' or block_mode == 'middle') and block_index == block_num[-1]:
+                    return x, patched_x_shape
+                kwargs.update(
+                    {
+                        "kv_cache": kv_cache[block_index],
+                        "crossattn_cache": crossattn_cache[block_index],
+                        "current_start": current_start,
+                        "current_end": current_end
+                    }
+                )
+                x = block(x, **kwargs)
+        if block_mode == 'input' and block_num[-1] == len(self.blocks):
+            return x, patched_x_shape
+
+        # head
+        x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2))
+
+        # unpatchify
+        x = self.unpatchify(x, grid_sizes)
+        return torch.stack(x)
+
+    def _forward_train(
+        self,
+        x,
+        t,
+        context,
+        seq_len,
+        clip_fea=None,
+        y=None,
+    ):
+        r"""
+        Forward pass through the diffusion model
+
+        Args:
+            x (List[Tensor]):
+                List of input video tensors, each with shape [C_in, F, H, W]
+            t (Tensor):
+                Diffusion timesteps tensor of shape [B]
+            context (List[Tensor]):
+                List of text embeddings each with shape [L, C]
+            seq_len (`int`):
+                Maximum sequence length for positional encoding
+            clip_fea (Tensor, *optional*):
+                CLIP image features for image-to-video mode
+            y (List[Tensor], *optional*):
+                Conditional video inputs for image-to-video mode, same shape as x
+
+        Returns:
+            List[Tensor]:
+                List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+        """
+        if self.model_type == 'i2v':
+            assert clip_fea is not None and y is not None
+        # params
+        device = self.patch_embedding.weight.device
+        if self.freqs.device != device:
+            self.freqs = self.freqs.to(device)
+
+        # Construct blockwise causal attn mask
+        if self.block_mask is None:
+            self.block_mask = self._prepare_blockwise_causal_attn_mask(
+                device, num_frames=x.shape[2],
+                frame_seqlen=x.shape[-2] *
+                x.shape[-1] // (self.patch_size[1] * self.patch_size[2]),
+                num_frame_per_block=self.num_frame_per_block
+            )
+
+        if y is not None:
+            x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+        # embeddings
+        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+        grid_sizes = torch.stack(
+            [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+        x = [u.flatten(2).transpose(1, 2) for u in x]
+        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+        assert seq_lens.max() <= seq_len
+        x = torch.cat([
+            torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
+                      dim=1) for u in x
+        ])
+
+        # time embeddings
+        # with amp.autocast(dtype=torch.float32):
+        e = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x))
+        e0 = self.time_projection(e).unflatten(
+            1, (6, self.dim)).unflatten(dim=0, sizes=t.shape)
+        # assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+        # context
+        context_lens = None
+        context = self.text_embedding(
+            torch.stack([
+                torch.cat(
+                    [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+                for u in context
+            ]))
+
+        if clip_fea is not None:
+            context_clip = self.img_emb(clip_fea)  # bs x 257 x dim
+            context = torch.concat([context_clip, context], dim=1)
+
+        # arguments
+        kwargs = dict(
+            e=e0,
+            seq_lens=seq_lens,
+            grid_sizes=grid_sizes,
+            freqs=self.freqs,
+            context=context,
+            context_lens=context_lens,
+            block_mask=self.block_mask)
+
+        def create_custom_forward(module):
+            def custom_forward(*inputs, **kwargs):
+                return module(*inputs, **kwargs)
+            return custom_forward
+
+        for block_index, block in enumerate(self.blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                x = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    x, **kwargs,
+                    use_reentrant=False,
+                )
+            else:
+                x = block(x, **kwargs)
+
+        # head
+        x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2))
+
+        # unpatchify
+        x = self.unpatchify(x, grid_sizes)
+        return torch.stack(x)
+
+    def forward(
+        self,
+        *args,
+        **kwargs
+    ):
+        if kwargs.get('kv_cache', None) is not None:
+            return self._forward_inference(*args, **kwargs)
+        else:
+            return self._forward_train(*args, **kwargs)
+
+    def unpatchify(self, x, grid_sizes):
+        r"""
+        Reconstruct video tensors from patch embeddings.
+
+        Args:
+            x (List[Tensor]):
+                List of patchified features, each with shape [L, C_out * prod(patch_size)]
+            grid_sizes (Tensor):
+                Original spatial-temporal grid dimensions before patching,
+                    shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
+
+        Returns:
+            List[Tensor]:
+                Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
+        """
+
+        c = self.out_dim
+        out = []
+        for u, v in zip(x, grid_sizes.tolist()):
+            u = u[:math.prod(v)].view(*v, *self.patch_size, c)
+            u = torch.einsum('fhwpqrc->cfphqwr', u)
+            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
+            out.append(u)
+        return out
+
+    def init_weights(self):
+        r"""
+        Initialize model parameters using Xavier initialization.
+        """
+
+        # basic init
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+        # init embeddings
+        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
+        for m in self.text_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=.02)
+        for m in self.time_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=.02)
+
+        # init output layer
+        nn.init.zeros_(self.head.head.weight)

models/wan/causal_stream_inference.py

@@ -0,0 +1,420 @@
+from models import (
+    get_diffusion_wrapper,
+    get_text_encoder_wrapper,
+    get_vae_wrapper
+)
+from models.wan.taehv_wrapper import TAEHVWanVAEWrapper
+from typing import List
+import torch
+import torch.distributed as dist
+import logging
+
+LOGGER = logging.getLogger(__name__)
+
+class CausalStreamInferencePipeline(torch.nn.Module):
+    def __init__(self, args, device):
+        super().__init__()
+        model_type = args.model_type
+        self.device = device
+        # Step 1: Initialize all models
+        self.generator_model_name = getattr(
+            args, "generator_name", args.model_name)
+        self.generator = get_diffusion_wrapper(
+            model_name=self.generator_model_name)(model_type=model_type)
+        self.text_encoder = get_text_encoder_wrapper(
+            model_name=args.model_name)(model_type=model_type)
+        if getattr(args, "use_taehv", False):
+            LOGGER.info("Using TAEHV VAE wrapper for Wan inference")
+            self.vae = TAEHVWanVAEWrapper(
+                model_type=model_type,
+                checkpoint_path=getattr(args, "taehv_checkpoint_path", None),
+                use_tensorrt=getattr(args, "use_tensorrt", False),
+            )
+        else:
+            self.vae = get_vae_wrapper(model_name=args.model_name)(model_type=model_type)
+
+        # Step 2: Initialize all causal hyperparmeters
+        self._init_denoising_step_list(args, device)
+
+        if model_type == "T2V-1.3B":
+            self.num_transformer_blocks = 30
+            self.num_heads = 12
+        elif model_type == "T2V-14B":
+            self.num_transformer_blocks = 40
+            self.num_heads = 40
+        else:
+            raise ValueError(f"Model type {model_type} not supported")
+        scale_size = 16
+        self.height = args.height//scale_size*2
+        self.width = args.width//scale_size*2
+        self.frame_seq_length = (args.height//scale_size) * (args.width//scale_size)
+        self.num_kv_cache = args.num_kv_cache
+        self.kv_cache_length = self.frame_seq_length*self.num_kv_cache
+        self.num_sink_tokens = args.num_sink_tokens
+        self.adapt_sink_threshold = args.adapt_sink_threshold
+
+        self.conditional_dict = None
+        self.kv_cache1 = None
+        self.kv_cache2 = None
+        self.hidden_states = None
+        self.block_x = None
+        self.args = args
+        self.num_frame_per_block = getattr(
+            args, "num_frame_per_block", 1)
+
+        LOGGER.info("KV inference with %s frames per block", self.num_frame_per_block)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        self.generator.model.to(self.device)
+
+    def _init_denoising_step_list(self, args, device):
+        self.denoising_step_list = torch.tensor(
+            args.denoising_step_list, dtype=torch.long, device=device)
+        assert self.denoising_step_list[-1] == 0
+        if not args.t2v:
+            # remove the last timestep (which equals zero)
+            self.denoising_step_list = self.denoising_step_list[:-1]
+
+        self.scheduler = self.generator.get_scheduler()
+        if args.warp_denoising_step:  # Warp the denoising step according to the scheduler time shift
+            timesteps = torch.cat(
+                (self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))
+            ).to(device)
+            self.denoising_step_list = timesteps[1000 - self.denoising_step_list]
+
+    def _initialize_kv_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+        """
+        kv_cache1 = []
+        
+        for i in range(self.num_transformer_blocks):
+            cache_length = self.kv_cache_length
+            self.generator.model.blocks[i].self_attn.sink_size = self.num_sink_tokens
+            self.generator.model.blocks[i].self_attn.adapt_sink_thr = self.adapt_sink_threshold
+
+            kv_cache1.append({
+                "k": torch.zeros([batch_size, cache_length, self.num_heads, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, cache_length, self.num_heads, 128], dtype=dtype, device=device),
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "total_steps": len(self.denoising_step_list),
+                "current_step": len(self.denoising_step_list),
+            })
+
+        self.kv_cache1 = kv_cache1  # always store the clean cache
+
+    def _initialize_crossattn_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU cross-attention cache for the Wan model.
+        """
+        crossattn_cache = []
+
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache.append({
+                "k": torch.zeros([batch_size, 512, self.num_heads, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, 512, self.num_heads, 128], dtype=dtype, device=device),
+                "is_init": False,
+            })
+
+        self.crossattn_cache = crossattn_cache  # always store the clean cache
+    
+    def prepare(
+        self,
+        text_prompts: List[str],
+        device: torch.device,
+        dtype: torch.dtype,
+        block_mode: str='input',
+        noise: torch.Tensor = None,
+        current_start: int = 0,
+        current_end: int = None,
+        block_num: torch.Tensor = None,
+        batch_denoise: bool=True,
+    ):
+        self.device = device
+        batch_size = noise.shape[0]
+
+        self.conditional_dict = self.text_encoder(
+            text_prompts=text_prompts
+        )
+
+        # Step 1: Initialize KV cache
+        if self.kv_cache1 is None:
+            self._initialize_kv_cache(
+                batch_size=batch_size,
+                dtype=dtype,
+                device=device
+            )
+
+            self._initialize_crossattn_cache(
+                batch_size=batch_size,
+                dtype=dtype,
+                device=device
+            )
+        else:
+            # reset cross attn cache
+            for block_index in range(self.num_transformer_blocks):
+                self.crossattn_cache[block_index]["is_init"] = False
+        
+        current_start = torch.tensor([current_start], dtype=torch.long, device=device)
+        current_end = torch.tensor([current_end], dtype=torch.long, device=device)
+
+        for index, current_timestep in enumerate(self.denoising_step_list):
+            # set current timestep
+            timestep = torch.ones(
+                [batch_size, noise.shape[1]], device=noise.device, dtype=torch.int64) * current_timestep
+
+            if index < len(self.denoising_step_list) - 1:
+                denoised_pred = self.generator(
+                    noisy_image_or_video=noise,
+                    conditional_dict=self.conditional_dict,
+                    timestep=timestep,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start,
+                    current_end=current_end
+                )
+                next_timestep = self.denoising_step_list[index + 1]
+                noise = self.scheduler.add_noise(
+                    denoised_pred.flatten(0, 1),
+                    torch.randn_like(denoised_pred.flatten(0, 1)),
+                    next_timestep *
+                    torch.ones([batch_size], device=noise.device,
+                                dtype=torch.long)
+                ).unflatten(0, denoised_pred.shape[:2])
+            else:
+                # for getting real output
+                denoised_pred = self.generator(
+                    noisy_image_or_video=noise,
+                    conditional_dict=self.conditional_dict,
+                    timestep=timestep,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start,
+                    current_end=current_end
+                )
+
+        if not batch_denoise:
+            return denoised_pred
+
+        # Pre-allocate hidden_states tensor to avoid memory allocation during inference
+        self.batch_size = len(self.denoising_step_list)
+
+        # Determine which blocks to keep based on block_num range
+        blocks_to_keep = []
+        if block_num is not None:
+            start_block, end_block = block_num[0].item(), block_num[1].item()
+            blocks_to_keep = list(range(start_block, end_block))
+        else:
+            blocks_to_keep = list(range(self.num_transformer_blocks))
+
+        # Process only the blocks in the specified range
+        for i in range(self.num_transformer_blocks):
+            if dist.is_initialized():
+                dist.broadcast(self.crossattn_cache[i]['k'], src=0)
+                dist.broadcast(self.crossattn_cache[i]['v'], src=0)
+                dist.broadcast(self.kv_cache1[i]['k'], src=0)
+                dist.broadcast(self.kv_cache1[i]['v'], src=0)
+
+            self.kv_cache1[i]['k'] = self.kv_cache1[i]['k'].repeat(self.batch_size, 1, 1, 1)
+            self.kv_cache1[i]['v'] = self.kv_cache1[i]['v'].repeat(self.batch_size, 1, 1, 1)
+
+            self.kv_cache1[i]['global_end_index'] = self.kv_cache1[i]['global_end_index'].repeat(self.batch_size)
+            self.kv_cache1[i]['local_end_index'] = self.kv_cache1[i]['local_end_index'].repeat(self.batch_size)
+
+            self.crossattn_cache[i]['k'] = self.crossattn_cache[i]['k'].expand(self.batch_size, -1, -1, -1)
+            self.crossattn_cache[i]['v'] = self.crossattn_cache[i]['v'].expand(self.batch_size, -1, -1, -1)
+        
+        # Remove blocks outside the range
+        if block_num is not None:
+            for i in range(self.num_transformer_blocks):
+                if i not in blocks_to_keep:
+                    self.kv_cache1[i]['k'] = self.kv_cache1[i]['k'].cpu()
+                    self.kv_cache1[i]['v'] = self.kv_cache1[i]['v'].cpu()
+
+        self.hidden_states = torch.zeros(
+            (self.batch_size, self.num_frame_per_block, *noise.shape[2:]), dtype=noise.dtype, device=device
+        )
+
+        if block_mode in ['output', 'middle']:
+            self.block_x = torch.zeros(
+                (self.batch_size, self.frame_seq_length, self.num_heads*128), dtype=noise.dtype, device=device
+            )
+        else:
+            self.block_x = None
+
+        self.kv_cache_starts = torch.ones(self.batch_size, dtype=torch.long, device=device) * current_end
+        self.kv_cache_ends = torch.ones(self.batch_size, dtype=torch.long, device=device) * current_end + self.frame_seq_length
+
+        self.timestep = self.denoising_step_list
+
+        self.conditional_dict['prompt_embeds'] = self.conditional_dict['prompt_embeds'].repeat(self.batch_size, 1, 1)
+    
+        return denoised_pred
+    
+    def inference_stream(
+        self, 
+        noise: torch.Tensor, 
+        current_start: int, 
+        current_end: int, 
+        current_step: int,
+        ) -> torch.Tensor:
+
+        self.hidden_states[1:] = self.hidden_states[:-1].clone()
+        self.hidden_states[0] = noise[0]
+
+        self.kv_cache_starts[1:] = self.kv_cache_starts[:-1].clone()
+        self.kv_cache_starts[0] = current_start
+        
+        self.kv_cache_ends[1:] = self.kv_cache_ends[:-1].clone()
+        self.kv_cache_ends[0] = current_end
+
+        if current_step is not None:
+            self.timestep[0] = current_step
+        
+        self.hidden_states = self.generator(
+            noisy_image_or_video=self.hidden_states,
+            conditional_dict=self.conditional_dict,
+            timestep=self.timestep.unsqueeze(1).expand(-1, self.hidden_states.shape[1]),
+            kv_cache=self.kv_cache1,
+            crossattn_cache=self.crossattn_cache,
+            current_start=self.kv_cache_starts,
+            current_end=self.kv_cache_ends,
+        )
+
+        for i in range(len(self.denoising_step_list) - 1):
+            self.hidden_states[[i]] = self.scheduler.add_noise(
+                self.hidden_states[[i]],
+                torch.randn_like(self.hidden_states[[i]]),
+                self.denoising_step_list[i + 1] *
+                torch.ones([1], device=self.hidden_states.device,
+                            dtype=torch.long)
+            )
+
+        return self.hidden_states
+    
+    def inference_wo_batch(
+        self, 
+        noise: torch.Tensor, 
+        current_start: int, 
+        current_end: int, 
+        current_step: int,
+        ) -> torch.Tensor:
+
+        batch_size = noise.shape[0]
+
+        current_start = torch.ones(batch_size, dtype=torch.long, device=self.device) * current_start
+        current_end = torch.ones(batch_size, dtype=torch.long, device=self.device) * current_end
+
+        # Step 2.1: Spatial denoising loop
+        self.denoising_step_list[0] = current_step
+        for index, current_timestep in enumerate(self.denoising_step_list):
+            # set current timestep
+            timestep = torch.ones(
+                [batch_size, noise.shape[1]], device=noise.device, dtype=torch.int64) * current_timestep
+
+            if index < len(self.denoising_step_list) - 1:
+                denoised_pred = self.generator(
+                    noisy_image_or_video=noise,
+                    conditional_dict=self.conditional_dict,
+                    timestep=timestep,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start,
+                    current_end=current_end
+                )
+                next_timestep = self.denoising_step_list[index + 1]
+                noise = self.scheduler.add_noise(
+                    denoised_pred.flatten(0, 1),
+                    torch.randn_like(denoised_pred.flatten(0, 1)),
+                    next_timestep *
+                    torch.ones([batch_size], device=noise.device,
+                                dtype=torch.long)
+                ).unflatten(0, denoised_pred.shape[:2])
+            else:
+                # for getting real output
+                denoised_pred = self.generator(
+                    noisy_image_or_video=noise,
+                    conditional_dict=self.conditional_dict,
+                    timestep=timestep,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start,
+                    current_end=current_end
+                )
+
+        return denoised_pred
+
+    def inference(
+        self, 
+        noise: torch.Tensor, 
+        current_start: int, 
+        current_end: int, 
+        current_step: int, 
+        block_mode: str='input', 
+        block_num=None,
+        patched_x_shape: torch.Tensor=None, 
+        block_x: torch.Tensor=None,
+        ) -> torch.Tensor:
+
+        if block_mode == 'input':
+            self.hidden_states[1:] = self.hidden_states[:-1].clone()
+            self.hidden_states[0] = noise[0]
+
+            self.kv_cache_starts[1:] = self.kv_cache_starts[:-1].clone()
+            self.kv_cache_starts[0] = current_start
+            
+            self.kv_cache_ends[1:] = self.kv_cache_ends[:-1].clone()
+            self.kv_cache_ends[0] = current_end
+        else:
+            self.block_x.copy_(block_x)
+            self.hidden_states.copy_(noise)
+            self.kv_cache_starts.copy_(current_start)
+            self.kv_cache_ends.copy_(current_end)
+
+        if current_step is not None:
+            self.timestep[0] = current_step
+        
+        if block_mode == 'output':
+            denoised_pred = self.generator.forward_output(
+                noisy_image_or_video=self.hidden_states,
+                conditional_dict=self.conditional_dict,
+                timestep=self.timestep.unsqueeze(1).expand(-1, self.hidden_states.shape[1]),
+                kv_cache=self.kv_cache1,
+                crossattn_cache=self.crossattn_cache,
+                current_start=self.kv_cache_starts,
+                current_end=self.kv_cache_ends,
+                block_mode=block_mode,
+                block_num=block_num,
+                patched_x_shape=patched_x_shape,
+                block_x=self.block_x
+            )
+
+            for i in range(len(self.denoising_step_list) - 1):
+                denoised_pred[[i]] = self.scheduler.add_noise(
+                    denoised_pred[[i]],
+                    torch.randn_like(denoised_pred[[i]]),
+                    self.denoising_step_list[i + 1] *
+                    torch.ones([1], device=denoised_pred.device,
+                                dtype=torch.long)
+                )
+            patched_x_shape = None
+
+        else:
+            denoised_pred, patched_x_shape = self.generator.forward_input(
+                noisy_image_or_video=self.hidden_states,
+                conditional_dict=self.conditional_dict,
+                timestep=self.timestep.unsqueeze(1).expand(-1, self.hidden_states.shape[1]),
+                kv_cache=self.kv_cache1,
+                crossattn_cache=self.crossattn_cache,
+                current_start=self.kv_cache_starts,
+                current_end=self.kv_cache_ends,
+                block_mode=block_mode,
+                block_num=block_num,
+                patched_x_shape=patched_x_shape,
+                block_x=self.block_x,
+            ) 
+
+        return denoised_pred, patched_x_shape

models/wan/flow_match.py

@@ -0,0 +1,83 @@
+"""
+The following code is copied from https://github.com/modelscope/DiffSynth-Studio/blob/main/diffsynth/schedulers/flow_match.py
+"""
+import torch
+
+
+class FlowMatchScheduler():
+
+    def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.set_timesteps(num_inference_steps)
+
+    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False):
+        sigma_start = self.sigma_min + \
+            (self.sigma_max - self.sigma_min) * denoising_strength
+        if self.extra_one_step:
+            self.sigmas = torch.linspace(
+                sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]
+        else:
+            self.sigmas = torch.linspace(
+                sigma_start, self.sigma_min, num_inference_steps)
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        self.sigmas = self.shift * self.sigmas / \
+            (1 + (self.shift - 1) * self.sigmas)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        if training:
+            x = self.timesteps
+            y = torch.exp(-2 * ((x - num_inference_steps / 2) /
+                          num_inference_steps) ** 2)
+            y_shifted = y - y.min()
+            bsmntw_weighing = y_shifted * \
+                (num_inference_steps / y_shifted.sum())
+            self.linear_timesteps_weights = bsmntw_weighing
+
+    def step(self, model_output, timestep, sample, to_final=False):
+        self.sigmas = self.sigmas.to(model_output.device)
+        self.timesteps = self.timesteps.to(model_output.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        if to_final or (timestep_id + 1 >= len(self.timesteps)).any():
+            sigma_ = 1 if (
+                self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        return prev_sample
+
+    def add_noise(self, original_samples, noise, timestep):
+        """
+        Diffusion forward corruption process.
+        Input:
+            - clean_latent: the clean latent with shape [B, C, H, W]
+            - noise: the noise with shape [B, C, H, W]
+            - timestep: the timestep with shape [B]
+        Output: the corrupted latent with shape [B, C, H, W]
+        """
+        self.sigmas = self.sigmas.to(noise.device)
+        self.timesteps = self.timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample.type_as(noise)
+
+    def training_target(self, sample, noise, timestep):
+        target = noise - sample
+        return target
+
+    def training_weight(self, timestep):
+        timestep_id = torch.argmin(
+            (self.timesteps - timestep.to(self.timesteps.device)).abs())
+        weights = self.linear_timesteps_weights[timestep_id]
+        return weights

models/wan/taehv_wrapper.py

@@ -0,0 +1,505 @@
+"""TAEHV-based Wan VAE wrapper used by offline inference with --use_taehv."""
+
+from __future__ import annotations
+
+import logging
+import os
+import urllib.request
+from collections import namedtuple
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from models.model_interface import VAEInterface
+from models.wan.wan_wrapper import WanVAEWrapper
+
+DecoderResult = namedtuple("DecoderResult", ("frame", "memory"))
+TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index"))
+LOGGER = logging.getLogger(__name__)
+
+try:
+    import tensorrt as trt
+except ImportError:
+    trt = None
+
+def _resolve_project_root() -> Path:
+    env_root = os.environ.get("STREAMDIFFUSIONV2_ROOT")
+    if env_root:
+        return Path(env_root).expanduser().resolve()
+
+    repo_root = Path(__file__).resolve().parents[2]
+    if (repo_root / "ckpts").exists():
+        return repo_root
+
+    cwd = Path.cwd().resolve()
+    if (cwd / "ckpts").exists():
+        return cwd
+
+    return repo_root
+
+
+PROJECT_ROOT = _resolve_project_root()
+DEFAULT_TAEHV_CHECKPOINT = str(PROJECT_ROOT / "ckpts" / "taew2_1.pth")
+DEFAULT_TAEHV_URL = "https://github.com/madebyollin/taehv/raw/main/taew2_1.pth"
+
+
+def conv(n_in, n_out, **kwargs):
+    return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs)
+
+
+class Clamp(nn.Module):
+    def forward(self, x):
+        return torch.tanh(x / 3) * 3
+
+
+class MemBlock(nn.Module):
+    def __init__(self, n_in, n_out):
+        super().__init__()
+        self.conv = nn.Sequential(
+            conv(n_in * 2, n_out),
+            nn.ReLU(inplace=True),
+            conv(n_out, n_out),
+            nn.ReLU(inplace=True),
+            conv(n_out, n_out),
+        )
+        self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
+        self.act = nn.ReLU(inplace=True)
+
+    def forward(self, x, past):
+        return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x))
+
+
+class TPool(nn.Module):
+    def __init__(self, n_f, stride):
+        super().__init__()
+        self.stride = stride
+        self.conv = nn.Conv2d(n_f * stride, n_f, 1, bias=False)
+
+    def forward(self, x):
+        _nt, c, h, w = x.shape
+        return self.conv(x.reshape(-1, self.stride * c, h, w))
+
+
+class TGrow(nn.Module):
+    def __init__(self, n_f, stride):
+        super().__init__()
+        self.stride = stride
+        self.conv = nn.Conv2d(n_f, n_f * stride, 1, bias=False)
+
+    def forward(self, x):
+        _nt, c, h, w = x.shape
+        x = self.conv(x)
+        return x.reshape(-1, c, h, w)
+
+
+def apply_model_with_memblocks(model, x, parallel, show_progress_bar):
+    assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor"
+    n, t, c, h, w = x.shape
+    if parallel:
+        x = x.reshape(n * t, c, h, w)
+        for block in model:
+            if isinstance(block, MemBlock):
+                nt, c, h, w = x.shape
+                t = nt // n
+                _x = x.reshape(n, t, c, h, w)
+                mem = F.pad(_x, (0, 0, 0, 0, 0, 0, 1, 0), value=0)[:, :t].reshape(x.shape)
+                x = block(x, mem)
+            else:
+                x = block(x)
+        nt, c, h, w = x.shape
+        t = nt // n
+        x = x.view(n, t, c, h, w)
+    else:
+        out = []
+        work_queue = [TWorkItem(xt, 0) for xt in x.reshape(n, t * c, h, w).chunk(t, dim=1)]
+        mem = [None] * len(model)
+        while work_queue:
+            xt, block_index = work_queue.pop(0)
+            if block_index == len(model):
+                out.append(xt)
+                continue
+
+            block = model[block_index]
+            if isinstance(block, MemBlock):
+                if mem[block_index] is None:
+                    xt_new = block(xt, xt * 0)
+                    mem[block_index] = xt
+                else:
+                    xt_new = block(xt, mem[block_index])
+                    mem[block_index].copy_(xt)
+                work_queue.insert(0, TWorkItem(xt_new, block_index + 1))
+            elif isinstance(block, TPool):
+                if mem[block_index] is None:
+                    mem[block_index] = []
+                mem[block_index].append(xt)
+                if len(mem[block_index]) == block.stride:
+                    n, c, h, w = xt.shape
+                    xt = block(torch.cat(mem[block_index], 1).view(n * block.stride, c, h, w))
+                    mem[block_index] = []
+                    work_queue.insert(0, TWorkItem(xt, block_index + 1))
+            elif isinstance(block, TGrow):
+                xt = block(xt)
+                n_out, c_out, h_out, w_out = xt.shape
+                batch_size = n_out // block.stride
+                grown = xt.view(batch_size, block.stride * c_out, h_out, w_out)
+                for xt_next in reversed(grown.chunk(block.stride, dim=1)):
+                    work_queue.insert(0, TWorkItem(xt_next, block_index + 1))
+            else:
+                xt = block(xt)
+                work_queue.insert(0, TWorkItem(xt, block_index + 1))
+        x = torch.stack(out, 1)
+    return x
+
+
+class TAEHV(nn.Module):
+    latent_channels = 16
+    image_channels = 3
+
+    def __init__(self, checkpoint_path=DEFAULT_TAEHV_CHECKPOINT, decoder_time_upscale=(True, True), decoder_space_upscale=(True, True, True)):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            conv(TAEHV.image_channels, 64),
+            nn.ReLU(inplace=True),
+            TPool(64, 2),
+            conv(64, 64, stride=2, bias=False),
+            MemBlock(64, 64),
+            MemBlock(64, 64),
+            MemBlock(64, 64),
+            TPool(64, 2),
+            conv(64, 64, stride=2, bias=False),
+            MemBlock(64, 64),
+            MemBlock(64, 64),
+            MemBlock(64, 64),
+            TPool(64, 1),
+            conv(64, 64, stride=2, bias=False),
+            MemBlock(64, 64),
+            MemBlock(64, 64),
+            MemBlock(64, 64),
+            conv(64, TAEHV.latent_channels),
+        )
+        n_f = [256, 128, 64, 64]
+        self.frames_to_trim = 2 ** sum(decoder_time_upscale) - 1
+        self.decoder = nn.Sequential(
+            Clamp(),
+            conv(TAEHV.latent_channels, n_f[0]),
+            nn.ReLU(inplace=True),
+            MemBlock(n_f[0], n_f[0]),
+            MemBlock(n_f[0], n_f[0]),
+            MemBlock(n_f[0], n_f[0]),
+            nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1),
+            TGrow(n_f[0], 1),
+            conv(n_f[0], n_f[1], bias=False),
+            MemBlock(n_f[1], n_f[1]),
+            MemBlock(n_f[1], n_f[1]),
+            MemBlock(n_f[1], n_f[1]),
+            nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1),
+            TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1),
+            conv(n_f[1], n_f[2], bias=False),
+            MemBlock(n_f[2], n_f[2]),
+            MemBlock(n_f[2], n_f[2]),
+            MemBlock(n_f[2], n_f[2]),
+            nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1),
+            TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1),
+            conv(n_f[2], n_f[3], bias=False),
+            nn.ReLU(inplace=True),
+            conv(n_f[3], TAEHV.image_channels),
+        )
+        self.load_state_dict(self.patch_tgrow_layers(torch.load(checkpoint_path, map_location="cpu", weights_only=True)))
+
+    def patch_tgrow_layers(self, state_dict):
+        new_state_dict = self.state_dict()
+        for index, layer in enumerate(self.decoder):
+            if isinstance(layer, TGrow):
+                key = f"decoder.{index}.conv.weight"
+                if state_dict[key].shape[0] > new_state_dict[key].shape[0]:
+                    state_dict[key] = state_dict[key][-new_state_dict[key].shape[0]:]
+        return state_dict
+
+    def encode_video(self, x, parallel=True, show_progress_bar=False):
+        return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar)
+
+    def decode_video(self, x, parallel=True, show_progress_bar=False):
+        return apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar)
+
+
+class TAEHVParallelDecoderModule(nn.Module):
+    """Shape-specialized decoder graph that is friendly to ONNX/TensorRT export."""
+
+    def __init__(self, decoder: nn.Module):
+        super().__init__()
+        self.decoder = decoder
+
+    def forward(self, latent: torch.Tensor) -> torch.Tensor:
+        return apply_model_with_memblocks(
+            self.decoder,
+            latent,
+            parallel=True,
+            show_progress_bar=False,
+        )
+
+
+class TAEHVTensorRTDecoder:
+    """Cache TensorRT engines for the common fixed TAEHV decoder shapes."""
+
+    def __init__(
+        self,
+        decoder_module: nn.Module,
+        cache_dir: str | Path,
+        workspace_bytes: int,
+    ) -> None:
+        if trt is None:
+            raise ImportError("TensorRT is not installed, but TensorRT decode was requested.")
+
+        self.decoder_module = decoder_module.eval()
+        self.cache_dir = Path(cache_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.workspace_bytes = int(workspace_bytes)
+        self.logger = trt.Logger(trt.Logger.WARNING)
+        self.runtime = trt.Runtime(self.logger)
+        self._engine_cache: dict[tuple[int, ...], tuple[object, object, str, str]] = {}
+        self._stream_cache: dict[int, torch.cuda.Stream] = {}
+
+    def _shape_tag(self, shape: tuple[int, ...]) -> str:
+        return "x".join(str(dim) for dim in shape)
+
+    def _engine_path(self, shape: tuple[int, ...]) -> Path:
+        return self.cache_dir / f"taehv_decoder_trt_{self._shape_tag(shape)}.plan"
+
+    def _onnx_path(self, shape: tuple[int, ...]) -> Path:
+        return self.cache_dir / f"taehv_decoder_trt_{self._shape_tag(shape)}.onnx"
+
+    def _build_engine(self, shape: tuple[int, ...], device: torch.device) -> bytes:
+        onnx_path = self._onnx_path(shape)
+        engine_path = self._engine_path(shape)
+
+        sample = torch.randn(shape, device=device, dtype=torch.float16)
+        with torch.no_grad():
+            torch.onnx.export(
+                self.decoder_module,
+                (sample,),
+                onnx_path.as_posix(),
+                input_names=["latent"],
+                output_names=["video"],
+                opset_version=18,
+                do_constant_folding=True,
+            )
+
+        builder = trt.Builder(self.logger)
+        network = builder.create_network(
+            1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+        )
+        parser = trt.OnnxParser(network, self.logger)
+        with open(onnx_path, "rb") as handle:
+            if not parser.parse(handle.read()):
+                errors = "\n".join(str(parser.get_error(i)) for i in range(parser.num_errors))
+                raise RuntimeError(f"Failed to parse ONNX for TAEHV TensorRT engine:\n{errors}")
+
+        config = builder.create_builder_config()
+        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, self.workspace_bytes)
+        config.set_flag(trt.BuilderFlag.FP16)
+        serialized = builder.build_serialized_network(network, config)
+        if serialized is None:
+            raise RuntimeError("TensorRT failed to build the TAEHV decoder engine")
+        engine_bytes = bytes(serialized)
+        engine_path.write_bytes(engine_bytes)
+        return engine_bytes
+
+    def _load_engine(self, shape: tuple[int, ...], device: torch.device):
+        cached = self._engine_cache.get(shape)
+        if cached is not None:
+            return cached
+
+        engine_path = self._engine_path(shape)
+        if engine_path.exists():
+            engine_bytes = engine_path.read_bytes()
+        else:
+            LOGGER.info("Building TensorRT engine for TAEHV decoder shape %s", shape)
+            engine_bytes = self._build_engine(shape, device)
+
+        engine = self.runtime.deserialize_cuda_engine(engine_bytes)
+        if engine is None:
+            raise RuntimeError(f"Failed to deserialize TensorRT engine for TAEHV shape {shape}")
+        context = engine.create_execution_context()
+        input_name = ""
+        output_name = ""
+        for index in range(engine.num_io_tensors):
+            name = engine.get_tensor_name(index)
+            mode = engine.get_tensor_mode(name)
+            if mode == trt.TensorIOMode.INPUT:
+                input_name = name
+            elif mode == trt.TensorIOMode.OUTPUT:
+                output_name = name
+        if not input_name or not output_name:
+            raise RuntimeError("Failed to discover TensorRT decoder tensor names")
+
+        cached = (engine, context, input_name, output_name)
+        self._engine_cache[shape] = cached
+        return cached
+
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        latent = latent.contiguous().to(dtype=torch.float16)
+        shape = tuple(int(dim) for dim in latent.shape)
+        _engine, context, input_name, output_name = self._load_engine(shape, latent.device)
+
+        if not context.set_input_shape(input_name, shape):
+            raise RuntimeError(f"TensorRT rejected decoder input shape {shape}")
+        output_shape = tuple(int(dim) for dim in context.get_tensor_shape(output_name))
+        output = torch.empty(output_shape, device=latent.device, dtype=latent.dtype)
+
+        context.set_tensor_address(input_name, int(latent.data_ptr()))
+        context.set_tensor_address(output_name, int(output.data_ptr()))
+        device_index = latent.device.index
+        if device_index is None:
+            raise RuntimeError("TensorRT decode requires an explicit CUDA device index")
+        stream = self._stream_cache.get(device_index)
+        if stream is None:
+            stream = torch.cuda.Stream(device=latent.device)
+            self._stream_cache[device_index] = stream
+        current_stream = torch.cuda.current_stream(device=latent.device)
+        stream.wait_stream(current_stream)
+        ok = context.execute_async_v3(stream.cuda_stream)
+        if not ok:
+            raise RuntimeError("TensorRT execution failed for the TAEHV decoder")
+        current_stream.wait_stream(stream)
+        return output
+
+
+class TAEHVWanVAEWrapper(VAEInterface):
+    """Wan stream encoder with a TAEHV decoder for faster pixel reconstruction."""
+
+    def __init__(
+        self,
+        model_type: str = "T2V-1.3B",
+        checkpoint_path: str | None = None,
+        auto_download: bool = True,
+        parallel_decode: bool = False,
+        use_tensorrt: bool = False,
+        tensorrt_cache_dir: str | None = None,
+        tensorrt_workspace_bytes: int = 4 << 30,
+    ):
+        super().__init__()
+        self.checkpoint_path = checkpoint_path or DEFAULT_TAEHV_CHECKPOINT
+        self.use_tensorrt = use_tensorrt
+        self.parallel_decode = parallel_decode or use_tensorrt
+        self.model = SimpleNamespace(first_encode=True, first_decode=True)
+        self.decode_context_latents = 3
+        self._decode_latent_cache: torch.Tensor | None = None
+        self.encoder_vae = WanVAEWrapper(model_type=model_type)
+        self.taehv = TAEHV(checkpoint_path=self._resolve_checkpoint(auto_download))
+        self._tensorrt_cache_dir = tensorrt_cache_dir or (PROJECT_ROOT / "ckpts" / "taehv_trt")
+        self._tensorrt_workspace_bytes = int(tensorrt_workspace_bytes)
+        self._tensorrt_decoder: TAEHVTensorRTDecoder | None = None
+        self._tensorrt_failed = False
+
+    def _resolve_checkpoint(self, auto_download: bool) -> str:
+        if os.path.exists(self.checkpoint_path):
+            return self.checkpoint_path
+
+        if not auto_download:
+            raise FileNotFoundError(f"TAEHV checkpoint not found: {self.checkpoint_path}")
+
+        os.makedirs(os.path.dirname(self.checkpoint_path), exist_ok=True)
+        urllib.request.urlretrieve(DEFAULT_TAEHV_URL, self.checkpoint_path)
+        return self.checkpoint_path
+
+    def to(self, *args, **kwargs):
+        device = kwargs.get("device")
+        dtype = kwargs.get("dtype")
+
+        if args:
+            if len(args) >= 1 and not isinstance(args[0], torch.dtype):
+                device = args[0]
+            if len(args) >= 2 and isinstance(args[1], torch.dtype):
+                dtype = args[1]
+            elif len(args) == 1 and isinstance(args[0], torch.dtype):
+                dtype = args[0]
+
+        encoder_kwargs = {}
+        if device is not None:
+            encoder_kwargs["device"] = device
+        if dtype is not None:
+            encoder_kwargs["dtype"] = dtype
+        self.encoder_vae.to(**encoder_kwargs)
+
+        taehv_kwargs = {"dtype": torch.float16}
+        if device is not None:
+            taehv_kwargs["device"] = device
+        self.taehv.to(**taehv_kwargs)
+        if self.use_tensorrt and not self._tensorrt_failed:
+            self._tensorrt_decoder = TAEHVTensorRTDecoder(
+                decoder_module=TAEHVParallelDecoderModule(self.taehv.decoder).to(**taehv_kwargs).eval(),
+                cache_dir=self._tensorrt_cache_dir,
+                workspace_bytes=self._tensorrt_workspace_bytes,
+            )
+        return self
+
+    def _pixels_to_unit_range(self, video: torch.Tensor) -> torch.Tensor:
+        return (video * 0.5 + 0.5).clamp(0, 1)
+
+    def _unit_range_to_pixels(self, video: torch.Tensor) -> torch.Tensor:
+        return video.mul(2).sub(1).clamp(-1, 1)
+
+    def _to_ntchw(self, video: torch.Tensor) -> torch.Tensor:
+        return video.permute(0, 2, 1, 3, 4).contiguous()
+
+    def _to_ncthw(self, video: torch.Tensor) -> torch.Tensor:
+        return video.permute(0, 2, 1, 3, 4).contiguous()
+
+    def _decode_video(self, latent: torch.Tensor) -> torch.Tensor:
+        if self.use_tensorrt and self._tensorrt_decoder is not None and not self._tensorrt_failed:
+            try:
+                return self._tensorrt_decoder.decode(latent)
+            except Exception as exc:
+                self._tensorrt_failed = True
+                self._tensorrt_decoder = None
+                LOGGER.warning(
+                    "TAEHV TensorRT decode failed for shape %s, falling back to PyTorch parallel decode: %s",
+                    tuple(latent.shape),
+                    exc,
+                )
+        return self.taehv.decode_video(
+            latent,
+            parallel=self.parallel_decode,
+            show_progress_bar=False,
+        )
+
+    def decode_to_pixel(self, latent: torch.Tensor) -> torch.Tensor:
+        video = self._decode_video(latent)
+        return self._unit_range_to_pixels(video)
+
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        return self.decode_to_pixel(latent)
+
+    def stream_encode(self, video: torch.Tensor, is_scale=False) -> torch.Tensor:
+        self.encoder_vae.model.first_encode = self.model.first_encode
+        latent = self.encoder_vae.stream_encode(video, is_scale=is_scale)
+        self.model.first_encode = self.encoder_vae.model.first_encode
+        return latent
+
+    def stream_decode_to_pixel(self, latent: torch.Tensor) -> torch.Tensor:
+        model_dtype = next(self.taehv.parameters()).dtype
+        latent = latent.to(dtype=model_dtype)
+
+        # Match Self-Forcing's TAEHV usage: keep a short latent prefix so each
+        # incremental decode sees recent temporal context, then trim the
+        # already-emitted frames from the pixel output.
+        if self.model.first_decode:
+            self.model.first_decode = False
+            self._decode_latent_cache = None
+            decode_latent = latent
+            emitted_latents = max(latent.shape[1] - 1, 0)
+        else:
+            context = self._decode_latent_cache
+            decode_latent = latent if context is None else torch.cat([context.to(device=latent.device, dtype=model_dtype), latent], dim=1)
+            emitted_latents = latent.shape[1]
+
+        self._decode_latent_cache = decode_latent[:, -min(self.decode_context_latents, decode_latent.shape[1]):].detach().clone()
+        video = self._decode_video(decode_latent)
+        if emitted_latents:
+            video = video[:, -emitted_latents * 4:, :, :, :]
+        else:
+            video = video[:, 0:0, :, :, :]
+        return self._unit_range_to_pixels(video).to(dtype=latent.dtype)

models/wan/wan_base/README.md

@@ -0,0 +1,2 @@
+Code in this folder is modified from https://github.com/Wan-Video/Wan2.1
+Apache-2.0 License 

models/wan/wan_base/__init__.py

@@ -0,0 +1 @@
+"""Minimal Wan base package used by StreamDiffusionV2 inference."""

models/wan/wan_base/modules/__init__.py

@@ -0,0 +1,16 @@
+from .attention import flash_attention
+from .model import WanModel
+from .t5 import T5Decoder, T5Encoder, T5EncoderModel, T5Model
+from .tokenizers import HuggingfaceTokenizer
+from .vae import WanVAE
+
+__all__ = [
+    'WanVAE',
+    'WanModel',
+    'T5Model',
+    'T5Encoder',
+    'T5Decoder',
+    'T5EncoderModel',
+    'HuggingfaceTokenizer',
+    'flash_attention',
+]

models/wan/wan_base/modules/attention.py

@@ -0,0 +1,310 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+
+try:
+    import flash_attn_interface
+    FLASH_ATTN_3_AVAILABLE = True
+except (ImportError, ModuleNotFoundError):
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+    FLASH_ATTN_2_AVAILABLE = hasattr(flash_attn, "flash_attn_varlen_func")
+except (ImportError, ModuleNotFoundError):
+    FLASH_ATTN_2_AVAILABLE = False
+
+import warnings
+
+__all__ = [
+    'flash_attention',
+    'attention',
+]
+
+
+def _prepare_sdpa_inputs(q, k, v, dtype):
+    q = q.transpose(1, 2)
+    k = k.transpose(1, 2)
+    v = v.transpose(1, 2)
+
+    if q.device.type == 'cpu' and dtype in (torch.float16, torch.bfloat16):
+        q = q.float()
+        k = k.float()
+        v = v.float()
+    else:
+        q = q.to(dtype)
+        k = k.to(dtype)
+        v = v.to(dtype)
+
+    return q, k, v
+
+
+def _build_length_mask(batch_size, q_len, k_len, device, q_lens, k_lens, causal, window_size):
+    mask = torch.ones((batch_size, q_len, k_len), dtype=torch.bool, device=device)
+
+    q_idx = torch.arange(q_len, device=device).view(1, q_len, 1)
+    k_idx = torch.arange(k_len, device=device).view(1, 1, k_len)
+
+    if q_lens is not None:
+        q_lens = q_lens.to(device=device, dtype=torch.long)
+        mask = mask & (q_idx < q_lens.view(batch_size, 1, 1))
+
+    if k_lens is not None:
+        k_lens = k_lens.to(device=device, dtype=torch.long)
+        mask = mask & (k_idx < k_lens.view(batch_size, 1, 1))
+
+    if causal:
+        mask = mask & (k_idx <= q_idx)
+
+    if window_size != (-1, -1):
+        left, right = window_size
+        if left >= 0:
+            mask = mask & (k_idx >= q_idx - left)
+        if right >= 0:
+            mask = mask & (k_idx <= q_idx + right)
+
+    return mask.unsqueeze(1)
+
+
+def _merge_sdpa_masks(length_mask, attn_mask, dtype):
+    if attn_mask is None:
+        return length_mask
+
+    if attn_mask.dtype == torch.bool:
+        return length_mask & attn_mask
+
+    additive_mask = torch.zeros_like(length_mask, dtype=dtype)
+    additive_mask = additive_mask.masked_fill(~length_mask, float('-inf'))
+    return additive_mask + attn_mask.to(dtype)
+
+
+def _sdpa_attention_fallback(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    dtype=torch.bfloat16,
+    attn_mask=None,
+):
+    out_dtype = q.dtype
+    batch_size, q_len, k_len = q.size(0), q.size(1), k.size(1)
+
+    q, k, v = _prepare_sdpa_inputs(q, k, v, dtype)
+
+    total_scale = 1.0
+    if q_scale is not None:
+        total_scale *= q_scale
+    if softmax_scale is not None:
+        total_scale *= softmax_scale
+    if total_scale != 1.0:
+        q = q * total_scale
+
+    mask = _build_length_mask(
+        batch_size=batch_size,
+        q_len=q_len,
+        k_len=k_len,
+        device=q.device,
+        q_lens=q_lens,
+        k_lens=k_lens,
+        causal=causal,
+        window_size=window_size,
+    )
+    mask = _merge_sdpa_masks(mask, attn_mask, q.dtype)
+
+    out = torch.nn.functional.scaled_dot_product_attention(
+        q,
+        k,
+        v,
+        attn_mask=mask,
+        is_causal=False,
+        dropout_p=dropout_p,
+    )
+
+    if q_lens is not None:
+        q_valid = (
+            torch.arange(q_len, device=out.device).view(1, q_len, 1)
+            < q_lens.to(device=out.device, dtype=torch.long).view(batch_size, 1, 1)
+        ).unsqueeze(1)
+        out = out.masked_fill(~q_valid, 0)
+
+    return out.transpose(1, 2).contiguous().to(out_dtype)
+
+
+def flash_attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    version=None,
+):
+    """
+    q:              [B, Lq, Nq, C1].
+    k:              [B, Lk, Nk, C1].
+    v:              [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+    q_lens:         [B].
+    k_lens:         [B].
+    dropout_p:      float. Dropout probability.
+    softmax_scale:  float. The scaling of QK^T before applying softmax.
+    causal:         bool. Whether to apply causal attention mask.
+    window_size:    (left right). If not (-1, -1), apply sliding window local attention.
+    deterministic:  bool. If True, slightly slower and uses more memory.
+    dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+    """
+    half_dtypes = (torch.float16, torch.bfloat16)
+    assert dtype in half_dtypes
+    if not (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE):
+        warnings.warn(
+            'flash_attn is not installed; falling back to scaled_dot_product_attention.',
+            stacklevel=2,
+        )
+        return _sdpa_attention_fallback(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            dtype=dtype,
+        )
+
+    assert q.device.type == 'cuda' and q.size(-1) <= 256
+
+    # params
+    b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+    def half(x):
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    # preprocess query
+    if q_lens is None:
+        q = half(q.flatten(0, 1))
+        q_lens = torch.tensor(
+            [lq] * b, dtype=torch.int32).to(
+                device=q.device, non_blocking=True)
+    else:
+        q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
+
+    # preprocess key, value
+    if k_lens is None:
+        k = half(k.flatten(0, 1))
+        v = half(v.flatten(0, 1))
+        k_lens = torch.tensor(
+            [lk] * b, dtype=torch.int32).to(
+                device=k.device, non_blocking=True)
+    else:
+        k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
+        v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
+
+    q = q.to(v.dtype)
+    k = k.to(v.dtype)
+
+    if q_scale is not None:
+        q = q * q_scale
+
+    if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+        warnings.warn(
+            'Flash attention 3 is not available, use flash attention 2 instead.'
+        )
+
+    # apply attention
+    if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+        # Note: dropout_p, window_size are not supported in FA3 now.
+        x = flash_attn_interface.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            deterministic=deterministic)[0].unflatten(0, (b, lq))
+    else:
+        assert FLASH_ATTN_2_AVAILABLE
+        x = flash_attn.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic).unflatten(0, (b, lq))
+
+    # output
+    return x.type(out_dtype)
+
+
+def attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    fa_version=None,
+    attn_mask=None,
+):
+    if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+        return flash_attention(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+            dtype=dtype,
+            version=fa_version,
+        )
+    else:
+        return _sdpa_attention_fallback(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            dtype=dtype,
+            attn_mask=attn_mask,
+        )

models/wan/wan_base/modules/model.py

@@ -0,0 +1,653 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import math
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+from .attention import flash_attention
+
+__all__ = ['WanModel']
+
+
+def sinusoidal_embedding_1d(dim, position):
+    # preprocess
+    assert dim % 2 == 0
+    half = dim // 2
+    position = position.type(torch.float64)
+
+    # calculation
+    sinusoid = torch.outer(
+        position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x
+
+
+# @amp.autocast(enabled=False)
+def rope_params(max_seq_len, dim, theta=10000):
+    assert dim % 2 == 0
+    freqs = torch.outer(
+        torch.arange(max_seq_len),
+        1.0 / torch.pow(theta,
+                        torch.arange(0, dim, 2).to(torch.float64).div(dim)))
+    freqs = torch.polar(torch.ones_like(freqs), freqs)
+    return freqs
+
+
+# @amp.autocast(enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+    n, c = x.size(2), x.size(3) // 2
+
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+    # loop over samples
+    output = []
+    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+        seq_len = f * h * w
+
+        # precompute multipliers
+        x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(
+            seq_len, n, -1, 2))
+        freqs_i = torch.cat([
+            freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+        ],
+            dim=-1).reshape(seq_len, 1, -1)
+
+        # apply rotary embedding
+        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+        x_i = torch.cat([x_i, x[i, seq_len:]])
+
+        # append to collection
+        output.append(x_i)
+    return torch.stack(output).type_as(x)
+
+
+class WanRMSNorm(nn.Module):
+
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        return self._norm(x.float()).type_as(x) * self.weight
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+
+class WanLayerNorm(nn.LayerNorm):
+
+    def __init__(self, dim, eps=1e-6, elementwise_affine=False):
+        super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        return super().forward(x).type_as(x)
+
+
+class WanSelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 eps=1e-6):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # layers
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(self, x, seq_lens, grid_sizes, freqs):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, num_heads, C / num_heads]
+            seq_lens(Tensor): Shape [B]
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+
+        # query, key, value function
+        def qkv_fn(x):
+            q = self.norm_q(self.q(x)).view(b, s, n, d)
+            k = self.norm_k(self.k(x)).view(b, s, n, d)
+            v = self.v(x).view(b, s, n, d)
+            return q, k, v
+
+        q, k, v = qkv_fn(x)
+
+        x = flash_attention(
+            q=rope_apply(q, grid_sizes, freqs),
+            k=rope_apply(k, grid_sizes, freqs),
+            v=v,
+            k_lens=seq_lens,
+            window_size=self.window_size)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanT2VCrossAttention(WanSelfAttention):
+
+    def forward(self, x, context, context_lens, crossattn_cache=None):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+            crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding.
+        """
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q(self.q(x)).view(b, -1, n, d)
+
+        if crossattn_cache is not None:
+            if not crossattn_cache["is_init"]:
+                crossattn_cache["is_init"] = True
+                k = self.norm_k(self.k(context)).view(b, -1, n, d)
+                v = self.v(context).view(b, -1, n, d)
+                crossattn_cache["k"] = k
+                crossattn_cache["v"] = v
+            else:
+                k = crossattn_cache["k"]
+                v = crossattn_cache["v"]
+        else:
+            k = self.norm_k(self.k(context)).view(b, -1, n, d)
+            v = self.v(context).view(b, -1, n, d)
+
+        # compute attention
+        x = flash_attention(q, k, v, k_lens=context_lens)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanI2VCrossAttention(WanSelfAttention):
+
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 eps=1e-6):
+        super().__init__(dim, num_heads, window_size, qk_norm, eps)
+
+        self.k_img = nn.Linear(dim, dim)
+        self.v_img = nn.Linear(dim, dim)
+        # self.alpha = nn.Parameter(torch.zeros((1, )))
+        self.norm_k_img = WanRMSNorm(
+            dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(self, x, context, context_lens):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        context_img = context[:, :257]
+        context = context[:, 257:]
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q(self.q(x)).view(b, -1, n, d)
+        k = self.norm_k(self.k(context)).view(b, -1, n, d)
+        v = self.v(context).view(b, -1, n, d)
+        k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d)
+        v_img = self.v_img(context_img).view(b, -1, n, d)
+        img_x = flash_attention(q, k_img, v_img, k_lens=None)
+        # compute attention
+        x = flash_attention(q, k, v, k_lens=context_lens)
+
+        # output
+        x = x.flatten(2)
+        img_x = img_x.flatten(2)
+        x = x + img_x
+        x = self.o(x)
+        return x
+
+
+WAN_CROSSATTENTION_CLASSES = {
+    't2v_cross_attn': WanT2VCrossAttention,
+    'i2v_cross_attn': WanI2VCrossAttention,
+}
+
+
+class WanAttentionBlock(nn.Module):
+
+    def __init__(self,
+                 cross_attn_type,
+                 dim,
+                 ffn_dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 cross_attn_norm=False,
+                 eps=1e-6):
+        super().__init__()
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        # layers
+        self.norm1 = WanLayerNorm(dim, eps)
+        self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm,
+                                          eps)
+        self.norm3 = WanLayerNorm(
+            dim, eps,
+            elementwise_affine=True) if cross_attn_norm else nn.Identity()
+        self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim,
+                                                                      num_heads,
+                                                                      (-1, -1),
+                                                                      qk_norm,
+                                                                      eps)
+        self.norm2 = WanLayerNorm(dim, eps)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
+            nn.Linear(ffn_dim, dim))
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        x,
+        e,
+        seq_lens,
+        grid_sizes,
+        freqs,
+        context,
+        context_lens,
+    ):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+            e(Tensor): Shape [B, 6, C]
+            seq_lens(Tensor): Shape [B], length of each sequence in batch
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        # assert e.dtype == torch.float32
+        # with amp.autocast(dtype=torch.float32):
+        e = (self.modulation + e).chunk(6, dim=1)
+        # assert e[0].dtype == torch.float32
+
+        # self-attention
+        y = self.self_attn(
+            self.norm1(x) * (1 + e[1]) + e[0], seq_lens, grid_sizes,
+            freqs)
+        # with amp.autocast(dtype=torch.float32):
+        x = x + y * e[2]
+
+        # cross-attention & ffn function
+        def cross_attn_ffn(x, context, context_lens, e):
+            x = x + self.cross_attn(self.norm3(x), context, context_lens)
+            y = self.ffn(self.norm2(x) * (1 + e[4]) + e[3])
+            # with amp.autocast(dtype=torch.float32):
+            x = x + y * e[5]
+            return x
+
+        x = cross_attn_ffn(x, context, context_lens, e)
+        return x
+
+
+class Head(nn.Module):
+
+    def __init__(self, dim, out_dim, patch_size, eps=1e-6):
+        super().__init__()
+        self.dim = dim
+        self.out_dim = out_dim
+        self.patch_size = patch_size
+        self.eps = eps
+
+        # layers
+        out_dim = math.prod(patch_size) * out_dim
+        self.norm = WanLayerNorm(dim, eps)
+        self.head = nn.Linear(dim, out_dim)
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x, e):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            e(Tensor): Shape [B, C]
+        """
+        # assert e.dtype == torch.float32
+        # with amp.autocast(dtype=torch.float32):
+        e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
+        x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))
+        return x
+
+
+class MLPProj(torch.nn.Module):
+
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+
+        self.proj = torch.nn.Sequential(
+            torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim),
+            torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim),
+            torch.nn.LayerNorm(out_dim))
+
+    def forward(self, image_embeds):
+        clip_extra_context_tokens = self.proj(image_embeds)
+        return clip_extra_context_tokens
+
+
+class WanModel(ModelMixin, ConfigMixin):
+    r"""
+    Wan diffusion backbone supporting both text-to-video and image-to-video.
+    """
+
+    ignore_for_config = [
+        'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size'
+    ]
+    _no_split_modules = ['WanAttentionBlock']
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(self,
+                 model_type='t2v',
+                 patch_size=(1, 2, 2),
+                 text_len=512,
+                 in_dim=16,
+                 dim=2048,
+                 ffn_dim=8192,
+                 freq_dim=256,
+                 text_dim=4096,
+                 out_dim=16,
+                 num_heads=16,
+                 num_layers=32,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 cross_attn_norm=True,
+                 eps=1e-6):
+        r"""
+        Initialize the diffusion model backbone.
+
+        Args:
+            model_type (`str`, *optional*, defaults to 't2v'):
+                Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+            patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+                3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+            text_len (`int`, *optional*, defaults to 512):
+                Fixed length for text embeddings
+            in_dim (`int`, *optional*, defaults to 16):
+                Input video channels (C_in)
+            dim (`int`, *optional*, defaults to 2048):
+                Hidden dimension of the transformer
+            ffn_dim (`int`, *optional*, defaults to 8192):
+                Intermediate dimension in feed-forward network
+            freq_dim (`int`, *optional*, defaults to 256):
+                Dimension for sinusoidal time embeddings
+            text_dim (`int`, *optional*, defaults to 4096):
+                Input dimension for text embeddings
+            out_dim (`int`, *optional*, defaults to 16):
+                Output video channels (C_out)
+            num_heads (`int`, *optional*, defaults to 16):
+                Number of attention heads
+            num_layers (`int`, *optional*, defaults to 32):
+                Number of transformer blocks
+            window_size (`tuple`, *optional*, defaults to (-1, -1)):
+                Window size for local attention (-1 indicates global attention)
+            qk_norm (`bool`, *optional*, defaults to True):
+                Enable query/key normalization
+            cross_attn_norm (`bool`, *optional*, defaults to False):
+                Enable cross-attention normalization
+            eps (`float`, *optional*, defaults to 1e-6):
+                Epsilon value for normalization layers
+        """
+
+        super().__init__()
+
+        assert model_type in ['t2v', 'i2v']
+        self.model_type = model_type
+
+        self.patch_size = patch_size
+        self.text_len = text_len
+        self.in_dim = in_dim
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.freq_dim = freq_dim
+        self.text_dim = text_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        # embeddings
+        self.patch_embedding = nn.Conv3d(
+            in_dim, dim, kernel_size=patch_size, stride=patch_size)
+        self.text_embedding = nn.Sequential(
+            nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
+            nn.Linear(dim, dim))
+
+        self.time_embedding = nn.Sequential(
+            nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.time_projection = nn.Sequential(
+            nn.SiLU(), nn.Linear(dim, dim * 6))
+
+        # blocks
+        cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'
+        self.blocks = nn.ModuleList([
+            WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,
+                              window_size, qk_norm, cross_attn_norm, eps)
+            for _ in range(num_layers)
+        ])
+
+        # head
+        self.head = Head(dim, out_dim, patch_size, eps)
+
+        # buffers (don't use register_buffer otherwise dtype will be changed in to())
+        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+        d = dim // num_heads
+        self.freqs = torch.cat([
+            rope_params(1024, d - 4 * (d // 6)),
+            rope_params(1024, 2 * (d // 6)),
+            rope_params(1024, 2 * (d // 6))
+        ],
+            dim=1)
+
+        if model_type == 'i2v':
+            self.img_emb = MLPProj(1280, dim)
+
+        # initialize weights
+        self.init_weights()
+
+        self.gradient_checkpointing = False
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        self.gradient_checkpointing = value
+
+    def forward(
+        self,
+        x,
+        t,
+        context,
+        seq_len,
+        clip_fea=None,
+        y=None,
+    ):
+        r"""
+        Forward pass through the diffusion model
+
+        Args:
+            x (List[Tensor]):
+                List of input video tensors, each with shape [C_in, F, H, W]
+            t (Tensor):
+                Diffusion timesteps tensor of shape [B]
+            context (List[Tensor]):
+                List of text embeddings each with shape [L, C]
+            seq_len (`int`):
+                Maximum sequence length for positional encoding
+            clip_fea (Tensor, *optional*):
+                CLIP image features for image-to-video mode
+            y (List[Tensor], *optional*):
+                Conditional video inputs for image-to-video mode, same shape as x
+
+        Returns:
+            List[Tensor]:
+                List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+        """
+        if self.model_type == 'i2v':
+            assert clip_fea is not None and y is not None
+        # params
+        device = self.patch_embedding.weight.device
+        if self.freqs.device != device:
+            self.freqs = self.freqs.to(device)
+
+        if y is not None:
+            x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+        # embeddings
+        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+        grid_sizes = torch.stack(
+            [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+        x = [u.flatten(2).transpose(1, 2) for u in x]
+        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+        assert seq_lens.max() <= seq_len
+        x = torch.cat([
+            torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
+                      dim=1) for u in x
+        ])
+
+        # time embeddings
+        # with amp.autocast(dtype=torch.float32):
+        e = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim, t).type_as(x))
+        e0 = self.time_projection(e).unflatten(1, (6, self.dim))
+        # assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+        # context
+        context_lens = None
+        context = self.text_embedding(
+            torch.stack([
+                torch.cat(
+                    [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+                for u in context
+            ]))
+
+        if clip_fea is not None:
+            context_clip = self.img_emb(clip_fea)  # bs x 257 x dim
+            context = torch.concat([context_clip, context], dim=1)
+
+        # arguments
+        kwargs = dict(
+            e=e0,
+            seq_lens=seq_lens,
+            grid_sizes=grid_sizes,
+            freqs=self.freqs,
+            context=context,
+            context_lens=context_lens)
+
+        def create_custom_forward(module):
+            def custom_forward(*inputs, **kwargs):
+                return module(*inputs, **kwargs)
+            return custom_forward
+
+        for block in self.blocks:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                x = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    x, **kwargs,
+                    use_reentrant=False,
+                )
+            else:
+                x = block(x, **kwargs)
+
+        # head
+        x = self.head(x, e)
+
+        # unpatchify
+        x = self.unpatchify(x, grid_sizes)
+        return torch.stack(x)
+
+    def unpatchify(self, x, grid_sizes):
+        r"""
+        Reconstruct video tensors from patch embeddings.
+
+        Args:
+            x (List[Tensor]):
+                List of patchified features, each with shape [L, C_out * prod(patch_size)]
+            grid_sizes (Tensor):
+                Original spatial-temporal grid dimensions before patching,
+                    shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
+
+        Returns:
+            List[Tensor]:
+                Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
+        """
+
+        c = self.out_dim
+        out = []
+        for u, v in zip(x, grid_sizes.tolist()):
+            u = u[:math.prod(v)].view(*v, *self.patch_size, c)
+            u = torch.einsum('fhwpqrc->cfphqwr', u)
+            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
+            out.append(u)
+        return out
+
+    def init_weights(self):
+        r"""
+        Initialize model parameters using Xavier initialization.
+        """
+
+        # basic init
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+        # init embeddings
+        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
+        for m in self.text_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=.02)
+        for m in self.time_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=.02)
+
+        # init output layer
+        nn.init.zeros_(self.head.head.weight)

models/wan/wan_base/modules/t5.py

@@ -0,0 +1,515 @@
+# Modified from transformers.models.t5.modeling_t5
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .tokenizers import HuggingfaceTokenizer
+
+__all__ = [
+    'T5Model',
+    'T5Encoder',
+    'T5Decoder',
+    'T5EncoderModel',
+]
+
+
+def fp16_clamp(x):
+    if x.dtype == torch.float16 and torch.isinf(x).any():
+        clamp = torch.finfo(x.dtype).max - 1000
+        x = torch.clamp(x, min=-clamp, max=clamp)
+    return x
+
+
+def init_weights(m):
+    if isinstance(m, T5LayerNorm):
+        nn.init.ones_(m.weight)
+    elif isinstance(m, T5Model):
+        nn.init.normal_(m.token_embedding.weight, std=1.0)
+    elif isinstance(m, T5FeedForward):
+        nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5)
+        nn.init.normal_(m.fc1.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5)
+    elif isinstance(m, T5Attention):
+        nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5)
+        nn.init.normal_(m.k.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.v.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5)
+    elif isinstance(m, T5RelativeEmbedding):
+        nn.init.normal_(
+            m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5)
+
+
+class GELU(nn.Module):
+
+    def forward(self, x):
+        return 0.5 * x * (1.0 + torch.tanh(
+            math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
+
+
+class T5LayerNorm(nn.Module):
+
+    def __init__(self, dim, eps=1e-6):
+        super(T5LayerNorm, self).__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) +
+                            self.eps)
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            x = x.type_as(self.weight)
+        return self.weight * x
+
+
+class T5Attention(nn.Module):
+
+    def __init__(self, dim, dim_attn, num_heads, dropout=0.1):
+        assert dim_attn % num_heads == 0
+        super(T5Attention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.num_heads = num_heads
+        self.head_dim = dim_attn // num_heads
+
+        # layers
+        self.q = nn.Linear(dim, dim_attn, bias=False)
+        self.k = nn.Linear(dim, dim_attn, bias=False)
+        self.v = nn.Linear(dim, dim_attn, bias=False)
+        self.o = nn.Linear(dim_attn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, context=None, mask=None, pos_bias=None):
+        """
+        x:          [B, L1, C].
+        context:    [B, L2, C] or None.
+        mask:       [B, L2] or [B, L1, L2] or None.
+        """
+        # check inputs
+        context = x if context is None else context
+        b, n, c = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.q(x).view(b, -1, n, c)
+        k = self.k(context).view(b, -1, n, c)
+        v = self.v(context).view(b, -1, n, c)
+
+        # attention bias
+        attn_bias = x.new_zeros(b, n, q.size(1), k.size(1))
+        if pos_bias is not None:
+            attn_bias += pos_bias
+        if mask is not None:
+            assert mask.ndim in [2, 3]
+            mask = mask.view(b, 1, 1,
+                             -1) if mask.ndim == 2 else mask.unsqueeze(1)
+            attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min)
+
+        # compute attention (T5 does not use scaling)
+        attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias
+        attn = F.softmax(attn.float(), dim=-1).type_as(attn)
+        x = torch.einsum('bnij,bjnc->binc', attn, v)
+
+        # output
+        x = x.reshape(b, -1, n * c)
+        x = self.o(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5FeedForward(nn.Module):
+
+    def __init__(self, dim, dim_ffn, dropout=0.1):
+        super(T5FeedForward, self).__init__()
+        self.dim = dim
+        self.dim_ffn = dim_ffn
+
+        # layers
+        self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU())
+        self.fc1 = nn.Linear(dim, dim_ffn, bias=False)
+        self.fc2 = nn.Linear(dim_ffn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = self.fc1(x) * self.gate(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5SelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5SelfAttention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.norm1 = T5LayerNorm(dim)
+        self.attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm2 = T5LayerNorm(dim)
+        self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+        self.pos_embedding = None if shared_pos else T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=True)
+
+    def forward(self, x, mask=None, pos_bias=None):
+        e = pos_bias if self.shared_pos else self.pos_embedding(
+            x.size(1), x.size(1))
+        x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e))
+        x = fp16_clamp(x + self.ffn(self.norm2(x)))
+        return x
+
+
+class T5CrossAttention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5CrossAttention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.norm1 = T5LayerNorm(dim)
+        self.self_attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm2 = T5LayerNorm(dim)
+        self.cross_attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm3 = T5LayerNorm(dim)
+        self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+        self.pos_embedding = None if shared_pos else T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=False)
+
+    def forward(self,
+                x,
+                mask=None,
+                encoder_states=None,
+                encoder_mask=None,
+                pos_bias=None):
+        e = pos_bias if self.shared_pos else self.pos_embedding(
+            x.size(1), x.size(1))
+        x = fp16_clamp(x + self.self_attn(self.norm1(x), mask=mask, pos_bias=e))
+        x = fp16_clamp(x + self.cross_attn(
+            self.norm2(x), context=encoder_states, mask=encoder_mask))
+        x = fp16_clamp(x + self.ffn(self.norm3(x)))
+        return x
+
+
+class T5RelativeEmbedding(nn.Module):
+
+    def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128):
+        super(T5RelativeEmbedding, self).__init__()
+        self.num_buckets = num_buckets
+        self.num_heads = num_heads
+        self.bidirectional = bidirectional
+        self.max_dist = max_dist
+
+        # layers
+        self.embedding = nn.Embedding(num_buckets, num_heads)
+
+    def forward(self, lq, lk):
+        device = self.embedding.weight.device
+        # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \
+        #     torch.arange(lq).unsqueeze(1).to(device)
+        rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \
+            torch.arange(lq, device=device).unsqueeze(1)
+        rel_pos = self._relative_position_bucket(rel_pos)
+        rel_pos_embeds = self.embedding(rel_pos)
+        rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze(
+            0)  # [1, N, Lq, Lk]
+        return rel_pos_embeds.contiguous()
+
+    def _relative_position_bucket(self, rel_pos):
+        # preprocess
+        if self.bidirectional:
+            num_buckets = self.num_buckets // 2
+            rel_buckets = (rel_pos > 0).long() * num_buckets
+            rel_pos = torch.abs(rel_pos)
+        else:
+            num_buckets = self.num_buckets
+            rel_buckets = 0
+            rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos))
+
+        # embeddings for small and large positions
+        max_exact = num_buckets // 2
+        rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) /
+                                     math.log(self.max_dist / max_exact) *
+                                     (num_buckets - max_exact)).long()
+        rel_pos_large = torch.min(
+            rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1))
+        rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large)
+        return rel_buckets
+
+
+class T5Encoder(nn.Module):
+
+    def __init__(self,
+                 vocab,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_layers,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5Encoder, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \
+            else nn.Embedding(vocab, dim)
+        self.pos_embedding = T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=True) if shared_pos else None
+        self.dropout = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList([
+            T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,
+                            shared_pos, dropout) for _ in range(num_layers)
+        ])
+        self.norm = T5LayerNorm(dim)
+
+        # initialize weights
+        self.apply(init_weights)
+
+    def forward(self, ids, mask=None):
+        x = self.token_embedding(ids)
+        x = self.dropout(x)
+        e = self.pos_embedding(x.size(1),
+                               x.size(1)) if self.shared_pos else None
+        for block in self.blocks:
+            x = block(x, mask, pos_bias=e)
+        x = self.norm(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5Decoder(nn.Module):
+
+    def __init__(self,
+                 vocab,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_layers,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5Decoder, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \
+            else nn.Embedding(vocab, dim)
+        self.pos_embedding = T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=False) if shared_pos else None
+        self.dropout = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList([
+            T5CrossAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,
+                             shared_pos, dropout) for _ in range(num_layers)
+        ])
+        self.norm = T5LayerNorm(dim)
+
+        # initialize weights
+        self.apply(init_weights)
+
+    def forward(self, ids, mask=None, encoder_states=None, encoder_mask=None):
+        b, s = ids.size()
+
+        # causal mask
+        if mask is None:
+            mask = torch.tril(torch.ones(1, s, s).to(ids.device))
+        elif mask.ndim == 2:
+            mask = torch.tril(mask.unsqueeze(1).expand(-1, s, -1))
+
+        # layers
+        x = self.token_embedding(ids)
+        x = self.dropout(x)
+        e = self.pos_embedding(x.size(1),
+                               x.size(1)) if self.shared_pos else None
+        for block in self.blocks:
+            x = block(x, mask, encoder_states, encoder_mask, pos_bias=e)
+        x = self.norm(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5Model(nn.Module):
+
+    def __init__(self,
+                 vocab_size,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 encoder_layers,
+                 decoder_layers,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5Model, self).__init__()
+        self.vocab_size = vocab_size
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.encoder_layers = encoder_layers
+        self.decoder_layers = decoder_layers
+        self.num_buckets = num_buckets
+
+        # layers
+        self.token_embedding = nn.Embedding(vocab_size, dim)
+        self.encoder = T5Encoder(self.token_embedding, dim, dim_attn, dim_ffn,
+                                 num_heads, encoder_layers, num_buckets,
+                                 shared_pos, dropout)
+        self.decoder = T5Decoder(self.token_embedding, dim, dim_attn, dim_ffn,
+                                 num_heads, decoder_layers, num_buckets,
+                                 shared_pos, dropout)
+        self.head = nn.Linear(dim, vocab_size, bias=False)
+
+        # initialize weights
+        self.apply(init_weights)
+
+    def forward(self, encoder_ids, encoder_mask, decoder_ids, decoder_mask):
+        x = self.encoder(encoder_ids, encoder_mask)
+        x = self.decoder(decoder_ids, decoder_mask, x, encoder_mask)
+        x = self.head(x)
+        return x
+
+
+def _t5(name,
+        encoder_only=False,
+        decoder_only=False,
+        return_tokenizer=False,
+        tokenizer_kwargs={},
+        dtype=torch.float32,
+        device='cpu',
+        **kwargs):
+    # sanity check
+    assert not (encoder_only and decoder_only)
+
+    # params
+    if encoder_only:
+        model_cls = T5Encoder
+        kwargs['vocab'] = kwargs.pop('vocab_size')
+        kwargs['num_layers'] = kwargs.pop('encoder_layers')
+        _ = kwargs.pop('decoder_layers')
+    elif decoder_only:
+        model_cls = T5Decoder
+        kwargs['vocab'] = kwargs.pop('vocab_size')
+        kwargs['num_layers'] = kwargs.pop('decoder_layers')
+        _ = kwargs.pop('encoder_layers')
+    else:
+        model_cls = T5Model
+
+    # init model
+    with torch.device(device):
+        model = model_cls(**kwargs)
+
+    # set device
+    model = model.to(dtype=dtype, device=device)
+
+    # init tokenizer
+    if return_tokenizer:
+        from .tokenizers import HuggingfaceTokenizer
+        tokenizer = HuggingfaceTokenizer(f'google/{name}', **tokenizer_kwargs)
+        return model, tokenizer
+    else:
+        return model
+
+
+def umt5_xxl(**kwargs):
+    cfg = dict(
+        vocab_size=256384,
+        dim=4096,
+        dim_attn=4096,
+        dim_ffn=10240,
+        num_heads=64,
+        encoder_layers=24,
+        decoder_layers=24,
+        num_buckets=32,
+        shared_pos=False,
+        dropout=0.1)
+    cfg.update(**kwargs)
+    return _t5('umt5-xxl', **cfg)
+
+
+class T5EncoderModel:
+
+    def __init__(
+        self,
+        text_len,
+        dtype=torch.bfloat16,
+        device=None,
+        checkpoint_path=None,
+        tokenizer_path=None,
+        shard_fn=None,
+    ):
+        self.text_len = text_len
+        self.dtype = dtype
+        if device is None:
+            device = torch.cuda.current_device()
+        self.device = device
+        self.checkpoint_path = checkpoint_path
+        self.tokenizer_path = tokenizer_path
+
+        # init model
+        model = umt5_xxl(
+            encoder_only=True,
+            return_tokenizer=False,
+            dtype=dtype,
+            device=device).eval().requires_grad_(False)
+        logging.info(f'loading {checkpoint_path}')
+        model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))
+        self.model = model
+        if shard_fn is not None:
+            self.model = shard_fn(self.model, sync_module_states=False)
+        else:
+            self.model.to(self.device)
+        # init tokenizer
+        self.tokenizer = HuggingfaceTokenizer(
+            name=tokenizer_path, seq_len=text_len, clean='whitespace')
+
+    def __call__(self, texts, device):
+        ids, mask = self.tokenizer(
+            texts, return_mask=True, add_special_tokens=True)
+        ids = ids.to(device)
+        mask = mask.to(device)
+        seq_lens = mask.gt(0).sum(dim=1).long()
+        context = self.model(ids, mask)
+        return [u[:v] for u, v in zip(context, seq_lens)]

models/wan/wan_base/modules/tokenizers.py

@@ -0,0 +1,82 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import html
+import string
+
+import ftfy
+import regex as re
+from transformers import AutoTokenizer
+
+__all__ = ['HuggingfaceTokenizer']
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+def canonicalize(text, keep_punctuation_exact_string=None):
+    text = text.replace('_', ' ')
+    if keep_punctuation_exact_string:
+        text = keep_punctuation_exact_string.join(
+            part.translate(str.maketrans('', '', string.punctuation))
+            for part in text.split(keep_punctuation_exact_string))
+    else:
+        text = text.translate(str.maketrans('', '', string.punctuation))
+    text = text.lower()
+    text = re.sub(r'\s+', ' ', text)
+    return text.strip()
+
+
+class HuggingfaceTokenizer:
+
+    def __init__(self, name, seq_len=None, clean=None, **kwargs):
+        assert clean in (None, 'whitespace', 'lower', 'canonicalize')
+        self.name = name
+        self.seq_len = seq_len
+        self.clean = clean
+
+        # init tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(name, **kwargs)
+        self.vocab_size = self.tokenizer.vocab_size
+
+    def __call__(self, sequence, **kwargs):
+        return_mask = kwargs.pop('return_mask', False)
+
+        # arguments
+        _kwargs = {'return_tensors': 'pt'}
+        if self.seq_len is not None:
+            _kwargs.update({
+                'padding': 'max_length',
+                'truncation': True,
+                'max_length': self.seq_len
+            })
+        _kwargs.update(**kwargs)
+
+        # tokenization
+        if isinstance(sequence, str):
+            sequence = [sequence]
+        if self.clean:
+            sequence = [self._clean(u) for u in sequence]
+        ids = self.tokenizer(sequence, **_kwargs)
+
+        # output
+        if return_mask:
+            return ids.input_ids, ids.attention_mask
+        else:
+            return ids.input_ids
+
+    def _clean(self, text):
+        if self.clean == 'whitespace':
+            text = whitespace_clean(basic_clean(text))
+        elif self.clean == 'lower':
+            text = whitespace_clean(basic_clean(text)).lower()
+        elif self.clean == 'canonicalize':
+            text = canonicalize(basic_clean(text))
+        return text

models/wan/wan_base/modules/vae.py

@@ -0,0 +1,754 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+__all__ = [
+    'WanVAE',
+]
+
+CACHE_T = 2
+
+
+class CausalConv3d(nn.Conv3d):
+    """
+    Causal 3d convolusion.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._padding = (self.padding[2], self.padding[2], self.padding[1],
+                         self.padding[1], 2 * self.padding[0], 0)
+        self.padding = (0, 0, 0)
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+
+        return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+
+    def __init__(self, dim, channel_first=True, images=True, bias=False):
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.
+
+    def forward(self, x):
+        return F.normalize(
+            x, dim=(1 if self.channel_first else
+                    -1)) * self.scale * self.gamma + self.bias
+
+
+class Upsample(nn.Upsample):
+
+    def forward(self, x):
+        """
+        Fix bfloat16 support for nearest neighbor interpolation.
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+
+    def __init__(self, dim, mode):
+        assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d',
+                        'downsample3d')
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        # layers
+        if mode == 'upsample2d':
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
+        elif mode == 'upsample3d':
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
+            self.time_conv = CausalConv3d(
+                dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+
+        elif mode == 'downsample2d':
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == 'downsample3d':
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = CausalConv3d(
+                dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        b, c, t, h, w = x.size()
+        if self.mode == 'upsample3d':
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = 'Rep'
+                    feat_idx[0] += 1
+                else:
+
+                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    if cache_x.shape[2] < CACHE_T and feat_cache[
+                            idx] is not None and feat_cache[idx] != 'Rep':
+                        # cache last frame of last two chunk
+                        cache_x = torch.cat([
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                cache_x.device), cache_x
+                        ],
+                            dim=2)
+                    if cache_x.shape[2] < CACHE_T and feat_cache[
+                            idx] is not None and feat_cache[idx] == 'Rep':
+                        cache_x = torch.cat([
+                            torch.zeros_like(cache_x).to(cache_x.device),
+                            cache_x
+                        ],
+                            dim=2)
+                    if feat_cache[idx] == 'Rep':
+                        x = self.time_conv(x)
+                    else:
+                        x = self.time_conv(x, feat_cache[idx])
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+
+                    x = x.reshape(b, 2, c, t, h, w)
+                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),
+                                    3)
+                    x = x.reshape(b, c, t * 2, h, w)
+        t = x.shape[2]
+        x = rearrange(x, 'b c t h w -> (b t) c h w')
+        x = self.resample(x)
+        x = rearrange(x, '(b t) c h w -> b c t h w', t=t)
+
+        if self.mode == 'downsample3d':
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = x.clone()
+                    feat_idx[0] += 1
+                else:
+
+                    cache_x = x[:, :, -1:, :, :].clone()
+                    # if cache_x.shape[2] < CACHE_T and feat_cache[idx] is not None and feat_cache[idx]!='Rep':
+                    #     # cache last frame of last two chunk
+                    #     cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+
+                    x = self.time_conv(
+                        torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+        return x
+
+    def init_weight(self, conv):
+        conv_weight = conv.weight
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        one_matrix = torch.eye(c1, c2)
+        init_matrix = one_matrix
+        nn.init.zeros_(conv_weight)
+        # conv_weight.data[:,:,-1,1,1] = init_matrix * 0.5
+        conv_weight.data[:, :, 1, 0, 0] = init_matrix  # * 0.5
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def init_weight2(self, conv):
+        conv_weight = conv.weight.data
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1 // 2, c2)
+        # init_matrix = repeat(init_matrix, 'o ... -> (o 2) ...').permute(1,0,2).contiguous().reshape(c1,c2)
+        conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
+        conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+
+    def __init__(self, in_dim, out_dim, dropout=0.0):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        # layers
+        self.residual = nn.Sequential(
+            RMS_norm(in_dim, images=False), nn.SiLU(),
+            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout),
+            CausalConv3d(out_dim, out_dim, 3, padding=1))
+        self.shortcut = CausalConv3d(in_dim, out_dim, 1) \
+            if in_dim != out_dim else nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < CACHE_T and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat([
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device), cache_x
+                    ],
+                        dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    Causal self-attention with a single head.
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+        # layers
+        self.norm = RMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+
+        # zero out the last layer params
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x):
+        identity = x
+        b, c, t, h, w = x.size()
+        x = rearrange(x, 'b c t h w -> (b t) c h w')
+        x = self.norm(x)
+        # compute query, key, value
+        q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3,
+                                         -1).permute(0, 1, 3,
+                                                     2).contiguous().chunk(
+                                                         3, dim=-1)
+
+        # apply attention
+        x = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+        )
+        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+
+        # output
+        x = self.proj(x)
+        x = rearrange(x, '(b t) c h w-> b c t h w', t=t)
+        return x + identity
+
+
+class Encoder3d(nn.Module):
+
+    def __init__(self,
+                 dim=128,
+                 z_dim=4,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_downsample=[True, True, False],
+                 dropout=0.0):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+
+        # dimensions
+        dims = [dim * u for u in [1] + dim_mult]
+        scale = 1.0
+
+        # init block
+        self.conv1 = CausalConv3d(3, dims[0], 3, padding=1)
+
+        # downsample blocks
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            for _ in range(num_res_blocks):
+                downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    downsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # downsample block
+            if i != len(dim_mult) - 1:
+                mode = 'downsample3d' if temperal_downsample[
+                    i] else 'downsample2d'
+                downsamples.append(Resample(out_dim, mode=mode))
+                scale /= 2.0
+        self.downsamples = nn.Sequential(*downsamples)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(out_dim, out_dim, dropout), AttentionBlock(out_dim),
+            ResidualBlock(out_dim, out_dim, dropout))
+
+        # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False), nn.SiLU(),
+            CausalConv3d(out_dim, z_dim, 3, padding=1))
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < CACHE_T and feat_cache[idx] is not None:
+                # cache last frame of last two chunk
+                cache_x = torch.cat([
+                    feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                        cache_x.device), cache_x
+                ],
+                    dim=2)
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        # downsamples
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        # middle
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        # head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < CACHE_T and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat([
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device), cache_x
+                    ],
+                        dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+class Decoder3d(nn.Module):
+
+    def __init__(self,
+                 dim=128,
+                 z_dim=4,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_upsample=[False, True, True],
+                 dropout=0.0):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_upsample = temperal_upsample
+
+        # dimensions
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2**(len(dim_mult) - 2)
+
+        # init block
+        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(dims[0], dims[0], dropout), AttentionBlock(dims[0]),
+            ResidualBlock(dims[0], dims[0], dropout))
+
+        # upsample blocks
+        upsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            if i == 1 or i == 2 or i == 3:
+                in_dim = in_dim // 2
+            for _ in range(num_res_blocks + 1):
+                upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    upsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # upsample block
+            if i != len(dim_mult) - 1:
+                mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d'
+                upsamples.append(Resample(out_dim, mode=mode))
+                scale *= 2.0
+        self.upsamples = nn.Sequential(*upsamples)
+
+        # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False), nn.SiLU(),
+            CausalConv3d(out_dim, 3, 3, padding=1))
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        # conv1
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < CACHE_T and feat_cache[idx] is not None:
+                # cache last frame of last two chunk
+                cache_x = torch.cat([
+                    feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                        cache_x.device), cache_x
+                ],
+                    dim=2)
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        # middle
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        # upsamples
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        # head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < CACHE_T and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat([
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device), cache_x
+                    ],
+                        dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+def count_conv3d(model):
+    count = 0
+    for m in model.modules():
+        if isinstance(m, CausalConv3d):
+            count += 1
+    return count
+
+
+class WanVAE_(nn.Module):
+
+    def __init__(self,
+                 dim=128,
+                 z_dim=4,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_downsample=[True, True, False],
+                 dropout=0.0):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        # modules
+        self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks,
+                                 attn_scales, self.temperal_downsample, dropout)
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks,
+                                 attn_scales, self.temperal_upsample, dropout)
+        self.first_encode = True
+        self.first_decode = True
+
+    def forward(self, x):
+        mu, log_var = self.encode(x)
+        z = self.reparameterize(mu, log_var)
+        x_recon = self.decode(z)
+        return x_recon, mu, log_var
+
+    def encode(self, x, scale):
+        self.clear_cache()
+        # cache
+        t = x.shape[2]
+        iter_ = 1 + (t - 1) // 4
+        # 对encode输入的x，按时间拆分为1、4、4、4....
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(
+                    x[:, :, :1, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx)
+            else:
+                out_ = self.encoder(
+                    x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx)
+                out = torch.cat([out, out_], 2)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            mu = (mu - scale[0]) * scale[1]
+        self.clear_cache()
+        return mu
+
+    def stream_encode(self, x, scale):
+        # cache
+        t = x.shape[2]
+        if self.first_encode:
+            self.first_encode = False
+            self.clear_cache_encode()
+            self._enc_conv_idx = [0]
+            out = self.encoder(
+                x[:, :, :1, :, :],
+                feat_cache=self._enc_feat_map,
+                feat_idx=self._enc_conv_idx,
+                )
+            self._enc_conv_idx = [0]
+            out_ = self.encoder(
+                x[:, :, 1:, :, :],
+                feat_cache=self._enc_feat_map,
+                feat_idx=self._enc_conv_idx,
+                )
+            out = torch.cat([out, out_], 2)
+        else:
+            out=[]
+            for i in range(t//4):
+                self._enc_conv_idx = [0]
+                out.append(self.encoder(
+                    x[:, :, i*4:(i+1)*4, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                    ))
+            out = torch.cat(out, 2)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if scale is not None:
+            if isinstance(scale[0], torch.Tensor):
+                mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+                    1, self.z_dim, 1, 1, 1)
+            else:
+                mu = (mu - scale[0]) * scale[1]
+        # self.clear_cache()
+        return mu
+
+    def decode(self, z, scale):
+        self.clear_cache()
+        # z: [b,c,t,h,w]
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            z = z / scale[1] + scale[0]
+        iter_ = z.shape[2]
+        x = self.conv2(z)
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx)
+            else:
+                out_ = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx)
+                out = torch.cat([out, out_], 2)
+        self.clear_cache()
+        return out
+
+    def stream_decode(self, z, scale):
+        # z: [b,c,t,h,w]
+        t=z.shape[2]
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            z = z / scale[1] + scale[0]
+        x = self.conv2(z)
+        if self.first_decode:
+            self.first_decode = False
+            self.clear_cache_decode()
+            self.first_batch = False
+            self._conv_idx = [0]
+            out = self.decoder(
+                x[:, :, :1, :, :],
+                feat_cache=self._feat_map,
+                feat_idx=self._conv_idx,
+                )
+            self._conv_idx = [0]
+            out_ = self.decoder(
+                x[:, :, 1:, :, :],
+                feat_cache=self._feat_map,
+                feat_idx=self._conv_idx,
+                )
+            out = torch.cat([out, out_], 2)
+        else:
+            out = []
+            for i in range(t):
+                self._conv_idx = [0]
+                out.append(self.decoder(
+                    x[:, :, i:(i+1), :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                    ))
+            out = torch.cat(out, 2)
+        # self.clear_cache()
+        return out
+
+    def reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return eps * std + mu
+
+    def sample(self, imgs, deterministic=False):
+        mu, log_var = self.encode(imgs)
+        if deterministic:
+            return mu
+        std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+        return mu + std * torch.randn_like(std)
+
+    def clear_cache(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        # cache encode
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+    def clear_cache_decode(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+    
+    def clear_cache_encode(self):
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+    
+
+
+def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **kwargs):
+    """
+    Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.
+    """
+    # params
+    cfg = dict(
+        dim=96,
+        z_dim=z_dim,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[False, True, True],
+        dropout=0.0)
+    cfg.update(**kwargs)
+
+    # init model
+    with torch.device('meta'):
+        model = WanVAE_(**cfg)
+
+    # load checkpoint
+    logging.info(f'loading {pretrained_path}')
+    model.load_state_dict(
+        torch.load(pretrained_path, map_location=device), assign=True)
+
+    return model
+
+
+class WanVAE:
+
+    def __init__(self,
+                 z_dim=16,
+                 vae_pth='cache/vae_step_411000.pth',
+                 dtype=torch.float,
+                 device="cuda"):
+        self.dtype = dtype
+        self.device = device
+
+        mean = [
+            -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
+            0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921
+        ]
+        std = [
+            2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,
+            3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160
+        ]
+        self.mean = torch.tensor(mean, dtype=dtype, device=device)
+        self.std = torch.tensor(std, dtype=dtype, device=device)
+        self.scale = [self.mean, 1.0 / self.std]
+
+        # init model
+        self.model = _video_vae(
+            pretrained_path=vae_pth,
+            z_dim=z_dim,
+        ).eval().requires_grad_(False).to(device)
+
+    def encode(self, videos):
+        """
+        videos: A list of videos each with shape [C, T, H, W].
+        """
+        with amp.autocast(dtype=self.dtype):
+            return [
+                self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0)
+                for u in videos
+            ]
+
+    def decode(self, zs):
+        with amp.autocast(dtype=self.dtype):
+            return [
+                self.model.decode(u.unsqueeze(0),
+                                  self.scale).float().clamp_(-1, 1).squeeze(0)
+                for u in zs
+            ]

models/wan/wan_wrapper.py

@@ -0,0 +1,351 @@
+from models.model_interface import (
+    DiffusionModelInterface,
+    TextEncoderInterface,
+    VAEInterface
+)
+from models.wan.wan_base.modules.tokenizers import HuggingfaceTokenizer
+from models.wan.wan_base.modules.model import WanModel
+from models.wan.wan_base.modules.vae import _video_vae
+from models.wan.wan_base.modules.t5 import umt5_xxl
+from models.wan.flow_match import FlowMatchScheduler
+from models.wan.causal_model import CausalWanModel
+from typing import List, Tuple, Dict, Optional
+import torch
+import os
+import torch.distributed as dist
+import time
+from pathlib import Path
+
+
+def _resolve_project_root() -> Path:
+    env_root = os.environ.get("STREAMDIFFUSIONV2_ROOT")
+    if env_root:
+        return Path(env_root).expanduser().resolve()
+
+    repo_root = Path(__file__).resolve().parents[2]
+    if (repo_root / "wan_models").exists():
+        return repo_root
+
+    cwd = Path.cwd().resolve()
+    if (cwd / "wan_models").exists():
+        return cwd
+
+    return repo_root
+
+
+PROJECT_ROOT = _resolve_project_root()
+
+
+class WanTextEncoder(TextEncoderInterface):
+    def __init__(self, model_type="T2V-1.3B") -> None:
+        super().__init__()
+
+        self.text_encoder = umt5_xxl(
+            encoder_only=True,
+            return_tokenizer=False,
+            dtype=torch.float32,
+            device=torch.device('cpu')
+        ).eval().requires_grad_(False)
+        self.text_encoder.load_state_dict(
+            torch.load(
+                PROJECT_ROOT / f"wan_models/Wan2.1-{model_type}/models_t5_umt5-xxl-enc-bf16.pth",
+                map_location='cpu', weights_only=False
+            )
+        )
+
+        self.tokenizer = HuggingfaceTokenizer(
+            name=str(PROJECT_ROOT / f"wan_models/Wan2.1-{model_type}/google/umt5-xxl/"), seq_len=512, clean='whitespace')
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def forward(self, text_prompts: List[str]) -> dict:
+        ids, mask = self.tokenizer(
+            text_prompts, return_mask=True, add_special_tokens=True)
+        ids = ids.to(self.device)
+        mask = mask.to(self.device)
+        seq_lens = mask.gt(0).sum(dim=1).long()
+        context = self.text_encoder(ids, mask)
+
+        for u, v in zip(context, seq_lens):
+            u[v:] = 0.0  # set padding to 0.0
+
+        return {
+            "prompt_embeds": context
+        }
+
+
+class WanVAEWrapper(VAEInterface):
+    def __init__(self, model_type="T2V-1.3B"):
+        super().__init__()
+        mean = [
+            -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
+            0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921
+        ]
+        std = [
+            2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,
+            3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160
+        ]
+        self.mean = torch.tensor(mean, dtype=torch.float32)
+        self.std = torch.tensor(std, dtype=torch.float32)
+
+        # init model
+        self.model = _video_vae(
+            pretrained_path=str(PROJECT_ROOT / f"wan_models/Wan2.1-{model_type}/Wan2.1_VAE.pth"),
+            z_dim=16,
+        ).eval().requires_grad_(False)
+
+    def decode_to_pixel(self, latent: torch.Tensor) -> torch.Tensor:
+        # from [batch_size, num_frames, num_channels, height, width]
+        # to [batch_size, num_channels, num_frames, height, width]
+        zs = latent.permute(0, 2, 1, 3, 4)
+
+        device, dtype = latent.device, latent.dtype
+        scale = [self.mean.to(device=device, dtype=dtype),
+                 1.0 / self.std.to(device=device, dtype=dtype)]
+
+        output = [
+            self.model.decode(u.unsqueeze(0),
+                              scale).float().clamp_(-1, 1).squeeze(0)
+            for u in zs
+        ]
+        output = torch.stack(output, dim=0)
+        # from [batch_size, num_channels, num_frames, height, width]
+        # to [batch_size, num_frames, num_channels, height, width]
+        output = output.permute(0, 2, 1, 3, 4)
+        return output
+
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        # from [batch_size, num_frames, num_channels, height, width]
+        # to [batch_size, num_channels, num_frames, height, width]
+        zs = latent.permute(0, 2, 1, 3, 4)
+
+        device, dtype = latent.device, latent.dtype
+        scale = [self.mean.to(device=device, dtype=dtype),
+                 1.0 / self.std.to(device=device, dtype=dtype)]
+
+        output = self.model.decode(zs, scale).clamp_(-1, 1)
+        # from [batch_size, num_channels, num_frames, height, width]
+        # to [batch_size, num_frames, num_channels, height, width]
+        # output = output.permute(0, 2, 1, 3, 4)
+        return output
+    
+    def stream_encode(self, video: torch.Tensor, is_scale=False) -> torch.Tensor:
+        if is_scale:
+            device, dtype = video.device, video.dtype
+            scale = [self.mean.to(device=device, dtype=dtype),
+                    1.0 / self.std.to(device=device, dtype=dtype)]
+        else:
+            scale = None
+        return self.model.stream_encode(video, scale)
+    
+    def stream_decode_to_pixel(self, latent: torch.Tensor) -> torch.Tensor:
+        zs = latent.permute(0, 2, 1, 3, 4)
+        zs = zs.to(device=latent.device, dtype=torch.bfloat16)
+        device, dtype = latent.device, latent.dtype
+        scale = [self.mean.to(device=device, dtype=dtype),
+                 1.0 / self.std.to(device=device, dtype=dtype)]
+        output = self.model.stream_decode(zs, scale).float().clamp_(-1, 1)
+        output = output.permute(0, 2, 1, 3, 4)
+        return output
+
+
+class WanDiffusionWrapper(DiffusionModelInterface):
+    def __init__(self, model_type="T2V-1.3B"):
+        super().__init__()
+
+        self.model = WanModel.from_pretrained(str(PROJECT_ROOT / f"wan_models/Wan2.1-{model_type}/"))
+        self.model.eval()
+
+        self.uniform_timestep = True
+
+        self.scheduler = FlowMatchScheduler(
+            shift=8.0, sigma_min=0.0, extra_one_step=True
+        )
+        self.scheduler.set_timesteps(1000, training=True)
+
+        self.seq_len = 32760  # [1, 21, 16, 60, 104]
+        super().post_init()
+
+    def enable_gradient_checkpointing(self) -> None:
+        self.model.enable_gradient_checkpointing()
+
+    def _convert_flow_pred_to_x0(self, flow_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
+        """
+        Convert flow matching's prediction to x0 prediction.
+        flow_pred: the prediction with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        pred = noise - x0
+        x_t = (1-sigma_t) * x0 + sigma_t * noise
+        we have x0 = x_t - sigma_t * pred
+        see derivations https://chatgpt.com/share/67bf8589-3d04-8008-bc6e-4cf1a24e2d0e
+        """
+        # use higher precision for calculations
+        original_dtype = flow_pred.dtype
+        flow_pred, xt, sigmas, timesteps = map(
+            lambda x: x.double().to(flow_pred.device), [flow_pred, xt,
+                                                        self.scheduler.sigmas,
+                                                        self.scheduler.timesteps]
+        )
+
+        timestep_id = torch.argmin(
+            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        x0_pred = xt - sigma_t * flow_pred
+        return x0_pred.to(original_dtype)
+
+    @staticmethod
+    def _convert_x0_to_flow_pred(scheduler, x0_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
+        """
+        Convert x0 prediction to flow matching's prediction.
+        x0_pred: the x0 prediction with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        pred = (x_t - x_0) / sigma_t
+        """
+        # use higher precision for calculations
+        original_dtype = x0_pred.dtype
+        x0_pred, xt, sigmas, timesteps = map(
+            lambda x: x.double().to(x0_pred.device), [x0_pred, xt,
+                                                      scheduler.sigmas,
+                                                      scheduler.timesteps]
+        )
+        timestep_id = torch.argmin(
+            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        flow_pred = (xt - x0_pred) / sigma_t
+        return flow_pred.to(original_dtype)
+
+    def forward(
+        self, noisy_image_or_video: torch.Tensor, conditional_dict: dict,
+        timestep: torch.Tensor, kv_cache: Optional[List[dict]] = None,
+        crossattn_cache: Optional[List[dict]] = None,
+        current_start: Optional[int] = None,
+        current_end: Optional[int] = None
+    ) -> torch.Tensor:
+        prompt_embeds = conditional_dict["prompt_embeds"]
+
+        # [B, F] -> [B]
+        if self.uniform_timestep:
+            input_timestep = timestep[:, 0]
+        else:
+            input_timestep = timestep
+
+        if kv_cache is not None:
+            flow_pred = self.model(
+                noisy_image_or_video.permute(0, 2, 1, 3, 4),
+                t=input_timestep, context=prompt_embeds,
+                seq_len=self.seq_len,
+                kv_cache=kv_cache,
+                crossattn_cache=crossattn_cache,
+                current_start=current_start,
+                current_end=current_end
+            ).permute(0, 2, 1, 3, 4)
+        else:
+            flow_pred = self.model(
+                noisy_image_or_video.permute(0, 2, 1, 3, 4),
+                t=input_timestep, context=prompt_embeds,
+                seq_len=self.seq_len
+            ).permute(0, 2, 1, 3, 4)
+
+        pred_x0 = self._convert_flow_pred_to_x0(
+            flow_pred=flow_pred.flatten(0, 1),
+            xt=noisy_image_or_video.flatten(0, 1),
+            timestep=timestep.flatten(0, 1)
+        ).unflatten(0, flow_pred.shape[:2])
+
+        return pred_x0
+
+    def forward_input(
+        self, noisy_image_or_video: torch.Tensor, conditional_dict: dict,
+        timestep: torch.Tensor,block_mode: str='input', block_num = None, kv_cache: Optional[List[dict]] = None,
+        crossattn_cache: Optional[List[dict]] = None,
+        current_start: Optional[int] = None,
+        current_end: Optional[int] = None,
+        patched_x_shape: torch.Tensor = None,
+        block_x: torch.Tensor = None,
+    ) -> torch.Tensor:
+        assert kv_cache is not None, "kv_cache must be provided"
+
+        prompt_embeds = conditional_dict["prompt_embeds"]
+
+        # [B, F] -> [B]
+        if self.uniform_timestep:
+            input_timestep = timestep[:, 0]
+        else:
+            input_timestep = timestep
+        
+        if block_x is not None and block_mode == 'middle':
+            noisy_image_or_video = block_x
+        else:
+            noisy_image_or_video = noisy_image_or_video.permute(0, 2, 1, 3, 4)
+
+        output, patched_x_shape = self.model(
+            noisy_image_or_video,
+            t=input_timestep, context=prompt_embeds,
+            seq_len=self.seq_len,
+            kv_cache=kv_cache,
+            crossattn_cache=crossattn_cache,
+            current_start=current_start,
+            current_end=current_end,
+            block_mode=block_mode,
+            block_num=block_num,
+            patched_x_shape=patched_x_shape,
+        )
+
+        return output, patched_x_shape
+
+    def forward_output(
+        self, noisy_image_or_video: torch.Tensor, conditional_dict: dict,
+        timestep: torch.Tensor, block_mode: str='output', block_num = None, kv_cache: Optional[List[dict]] = None,
+        crossattn_cache: Optional[List[dict]] = None,
+        current_start: Optional[int] = None,
+        current_end: Optional[int] = None,
+        patched_x_shape: torch.Tensor = None,
+        block_x: torch.Tensor = None,
+    ) -> torch.Tensor:
+        assert kv_cache is not None, "kv_cache must be provided"
+
+        prompt_embeds = conditional_dict["prompt_embeds"]
+
+        # [B, F] -> [B]
+        if self.uniform_timestep:
+            input_timestep = timestep[:, 0]
+        else:
+            input_timestep = timestep
+
+        flow_pred = self.model(
+            block_x,
+            t=input_timestep, context=prompt_embeds,
+            seq_len=self.seq_len,
+            kv_cache=kv_cache,
+            crossattn_cache=crossattn_cache,
+            current_start=current_start,
+            current_end=current_end,
+            block_mode=block_mode,
+            block_num=block_num,
+            patched_x_shape=patched_x_shape,
+        ).permute(0, 2, 1, 3, 4)
+
+        pred_x0 = self._convert_flow_pred_to_x0(
+            flow_pred=flow_pred.flatten(0, 1),
+            xt=noisy_image_or_video.flatten(0, 1),
+            timestep=timestep.flatten(0, 1)
+        ).unflatten(0, flow_pred.shape[:2])
+
+        return pred_x0
+
+
+class CausalWanDiffusionWrapper(WanDiffusionWrapper):
+    def __init__(self, model_type="T2V-1.3B"):
+        super().__init__()
+
+        self.model = CausalWanModel.from_pretrained(
+            str(PROJECT_ROOT / f"wan_models/Wan2.1-{model_type}/"))
+        self.model.eval()
+
+        self.uniform_timestep = False

requirements.txt

@@ -1 +1,12 @@
-streamdiffusionv2
+diffusers==0.35.1
+accelerate
+av
+einops
+ftfy
+imageio
+imageio-ffmpeg
+omegaconf
+markdown2
+numpy<2
+sentencepiece
+regex

streamdiffusionv2/__init__.py

@@ -0,0 +1,22 @@
+"""Public StreamDiffusionV2 Python API."""
+
+from streamdiffusionv2.pipeline import (
+    DenoisedChunk,
+    EncodedChunk,
+    StreamDiffusionV2Pipeline,
+    VideoChunk,
+    export_video,
+    load_video,
+)
+from streamv2v.api import StreamVideoToVideo, run_video_to_video
+
+__all__ = [
+    "DenoisedChunk",
+    "EncodedChunk",
+    "StreamDiffusionV2Pipeline",
+    "StreamVideoToVideo",
+    "VideoChunk",
+    "export_video",
+    "load_video",
+    "run_video_to_video",
+]

streamdiffusionv2/pipeline.py

@@ -0,0 +1,437 @@
+"""Readable staged video-to-video API for StreamDiffusionV2."""
+
+from __future__ import annotations
+
+from contextlib import ExitStack
+from dataclasses import dataclass
+from importlib.resources import as_file, files
+from pathlib import Path
+from typing import Literal
+
+from diffusers.utils import export_to_video as diffusers_export_to_video
+import numpy as np
+import torch
+
+from models.util import set_seed
+from streamv2v.inference import (
+    SingleGPUInferencePipeline as StreamBatchInferencePipeline,
+    compute_noise_scale_and_step,
+)
+from streamv2v.inference_common import load_mp4_as_tensor, merge_cli_config, normalize_acceleration_flags
+from streamv2v.inference_wo_batch import SingleGPUInferencePipeline as StreamNoBatchInferencePipeline
+
+
+SingleMode = Literal["single", "single-wo"]
+
+
+@dataclass
+class VideoChunk:
+    """One video chunk prepared for the encode -> denoise -> decode loop."""
+
+    frames: torch.Tensor
+    start_idx: int
+    end_idx: int
+    current_start: int
+    current_end: int
+
+
+@dataclass
+class EncodedChunk:
+    """Encoded latent chunk plus the schedule metadata needed for denoising."""
+
+    noisy_latents: torch.Tensor
+    current_start: int
+    current_end: int
+    noise_scale: float
+    current_step: int | None = None
+
+
+@dataclass
+class DenoisedChunk:
+    """Denoised latent chunk ready for VAE decoding."""
+
+    denoised_pred: torch.Tensor
+    last_frame_only: bool
+
+
+def _resolve_default_config_path(resource_stack: ExitStack) -> str:
+    resource = files("streamv2v.configs").joinpath("wan_causal_dmd_v2v.yaml")
+    return str(resource_stack.enter_context(as_file(resource)))
+
+
+def _resolve_device(device: str | torch.device | None) -> torch.device:
+    cuda_available = torch.cuda.is_available()
+    if device is None:
+        return torch.device("cuda" if cuda_available else "cpu")
+    resolved = torch.device(device)
+    if resolved.type == "cuda" and not cuda_available:
+        raise RuntimeError("CUDA is not available in the current Python environment")
+    if resolved.type == "cuda" and resolved.index is not None:
+        torch.cuda.set_device(resolved.index)
+    return resolved
+
+
+def _normalize_video_tensor(
+    video: str | Path | torch.Tensor,
+    *,
+    height: int,
+    width: int,
+    device: torch.device,
+) -> torch.Tensor:
+    if isinstance(video, (str, Path)):
+        tensor = load_mp4_as_tensor(str(video), resize_hw=(height, width)).unsqueeze(0)
+    else:
+        tensor = video
+        if tensor.ndim == 4:
+            tensor = tensor.unsqueeze(0)
+        if tensor.ndim != 5:
+            raise ValueError("video tensor must have shape [B, C, T, H, W] or [C, T, H, W]")
+    if tensor.dtype != torch.bfloat16:
+        tensor = tensor.to(dtype=torch.bfloat16)
+    return tensor.to(device)
+
+
+def load_video(video_path: str, *, height: int = 480, width: int = 832) -> torch.Tensor:
+    """Load a video file as a normalized tensor with shape [C, T, H, W]."""
+    return load_mp4_as_tensor(video_path, resize_hw=(height, width))
+
+
+def export_video(video: np.ndarray, output_path: str, *, fps: int = 16) -> str:
+    """Write a `[T, H, W, C]` float video array to an mp4 file."""
+    output_file = Path(output_path)
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+    diffusers_export_to_video(video, str(output_file), fps=fps)
+    return str(output_file)
+
+
+class StreamDiffusionV2Pipeline:
+    """Readable staged single-GPU API that mirrors the offline inference flow."""
+
+    def __init__(
+        self,
+        checkpoint_folder: str,
+        *,
+        mode: SingleMode = "single",
+        config_path: str | None = None,
+        device: str | torch.device | None = None,
+        noise_scale: float = 0.8,
+        height: int = 480,
+        width: int = 832,
+        fps: int = 16,
+        step: int = 2,
+        seed: int = 0,
+        model_type: str = "T2V-1.3B",
+        use_taehv: bool = False,
+        use_tensorrt: bool = False,
+        fast: bool = False,
+        profile: bool = False,
+    ) -> None:
+        if mode not in {"single", "single-wo"}:
+            raise ValueError("StreamDiffusionV2Pipeline only supports 'single' and 'single-wo'")
+
+        self._resource_stack = ExitStack()
+        self.mode = mode
+        self.device = _resolve_device(device)
+        self.checkpoint_folder = checkpoint_folder
+        self.noise_scale = float(noise_scale)
+        self.height = int(height)
+        self.width = int(width)
+        self.fps = int(fps)
+        self.seed = int(seed)
+        self.step = int(step)
+        self.profile = bool(profile)
+        self.model_type = model_type
+        self.prompt: str | None = None
+
+        resolved_config_path = config_path or _resolve_default_config_path(self._resource_stack)
+        self.config_path = resolved_config_path
+        flags = normalize_acceleration_flags(
+            {
+                "use_taehv": use_taehv,
+                "use_tensorrt": use_tensorrt,
+                "fast": fast,
+            }
+        )
+        self.use_taehv = bool(flags["use_taehv"])
+        self.use_tensorrt = bool(flags["use_tensorrt"])
+        self.fast = bool(flags["fast"])
+        config_args = {
+            "config_path": resolved_config_path,
+            "checkpoint_folder": checkpoint_folder,
+            "noise_scale": noise_scale,
+            "height": height,
+            "width": width,
+            "fps": fps,
+            "step": step,
+            "seed": seed,
+            "model_type": model_type,
+            "profile": profile,
+            "use_taehv": self.use_taehv,
+            "use_tensorrt": self.use_tensorrt,
+            "fast": self.fast,
+            "t2v": False,
+            "target_fps": None,
+            "fixed_noise_scale": False,
+            "num_frames": 81,
+        }
+        self.config = merge_cli_config(resolved_config_path, config_args)
+
+        manager_cls = (
+            StreamBatchInferencePipeline if mode == "single" else StreamNoBatchInferencePipeline
+        )
+        torch.set_grad_enabled(False)
+        set_seed(self.seed)
+        self.pipeline_manager = manager_cls(self.config, self.device)
+        self.pipeline_manager.load_model(checkpoint_folder)
+        self.chunk_size = 4 * self.config.num_frame_per_block
+        self.num_steps = len(self.pipeline_manager.pipeline.denoising_step_list)
+        self._next_chunk_index = 0
+
+    def close(self) -> None:
+        self._resource_stack.close()
+
+    def __enter__(self) -> "StreamDiffusionV2Pipeline":
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.close()
+
+    def enable_acceleration(
+        self,
+        *,
+        use_taehv: bool = False,
+        use_tensorrt: bool = False,
+        fast: bool = False,
+    ) -> "StreamDiffusionV2Pipeline":
+        """Rebuild the pipeline with the requested acceleration flags."""
+        replacement = StreamDiffusionV2Pipeline(
+            checkpoint_folder=self.checkpoint_folder,
+            mode=self.mode,
+            config_path=self.config_path,
+            device=self.device,
+            noise_scale=self.noise_scale,
+            height=self.height,
+            width=self.width,
+            fps=self.fps,
+            step=self.step,
+            seed=self.seed,
+            model_type=self.model_type,
+            use_taehv=use_taehv,
+            use_tensorrt=use_tensorrt,
+            fast=fast,
+            profile=self.profile,
+        )
+        self.close()
+        self.__dict__.update(replacement.__dict__)
+        return self
+
+    def prepare(self, prompt: str) -> None:
+        """Reset the stream state and store the prompt for the next denoising pass."""
+        self.prompt = prompt
+        self.pipeline_manager.reset_stream_state(reset_vae_flags=True)
+        self.pipeline_manager.processed = 0
+        self._next_chunk_index = 0
+
+    def chunk_video(self, video: str | Path | torch.Tensor) -> list[VideoChunk]:
+        """Split a full input video into the same chunks used by the offline inference loop."""
+        input_video = _normalize_video_tensor(
+            video,
+            height=self.height,
+            width=self.width,
+            device=self.device,
+        )
+        _, _, total_frames, _, _ = input_video.shape
+        if total_frames < 1 + self.chunk_size:
+            raise ValueError(f"video must contain at least {1 + self.chunk_size} frames")
+
+        chunks: list[VideoChunk] = []
+        start_idx = 0
+        end_idx = 1 + self.chunk_size
+        current_start = 0
+        current_end = self.pipeline_manager.pipeline.frame_seq_length * (1 + self.chunk_size // 4)
+
+        chunks.append(
+            VideoChunk(
+                frames=input_video[:, :, start_idx:end_idx],
+                start_idx=start_idx,
+                end_idx=end_idx,
+                current_start=current_start,
+                current_end=current_end,
+            )
+        )
+
+        while True:
+            start_idx = end_idx
+            end_idx = end_idx + self.chunk_size
+            if end_idx > total_frames:
+                break
+            current_start = current_end
+            current_end = current_end + (self.chunk_size // 4) * self.pipeline_manager.pipeline.frame_seq_length
+            chunks.append(
+                VideoChunk(
+                    frames=input_video[:, :, start_idx:end_idx],
+                    start_idx=start_idx,
+                    end_idx=end_idx,
+                    current_start=current_start,
+                    current_end=current_end,
+                )
+            )
+        return chunks
+
+    @torch.inference_mode()
+    def encode_chunk(
+        self,
+        input_video: str | Path | torch.Tensor,
+        chunk: VideoChunk,
+        *,
+        previous_noise_scale: float | None = None,
+        initial_noise_scale: float | None = None,
+    ) -> EncodedChunk:
+        """Encode one chunk in the same style as the offline inference loop."""
+        full_video = _normalize_video_tensor(
+            input_video,
+            height=self.height,
+            width=self.width,
+            device=self.device,
+        )
+        noise_scale = self.noise_scale if previous_noise_scale is None else float(previous_noise_scale)
+        init_noise_scale = self.noise_scale if initial_noise_scale is None else float(initial_noise_scale)
+        current_step = None
+
+        if chunk.start_idx != 0:
+            noise_scale, current_step = compute_noise_scale_and_step(
+                full_video,
+                chunk.end_idx,
+                self.chunk_size,
+                noise_scale,
+                init_noise_scale,
+            )
+
+        latents = self.pipeline_manager._timed_stream_encode(chunk.frames)
+        latents = latents.transpose(2, 1).contiguous().to(dtype=torch.bfloat16)
+        noise = torch.randn_like(latents)
+        return EncodedChunk(
+            noisy_latents=noise * noise_scale + latents * (1 - noise_scale),
+            current_start=chunk.current_start,
+            current_end=chunk.current_end,
+            noise_scale=float(noise_scale),
+            current_step=current_step,
+        )
+
+    @torch.inference_mode()
+    def encode_video(self, video: str | Path | torch.Tensor) -> list[EncodedChunk]:
+        """Encode a full input video into noisy latent chunks."""
+        chunks: list[EncodedChunk] = []
+        noise_scale = float(self.noise_scale)
+        init_noise_scale = noise_scale
+        video_chunks = self.chunk_video(video)
+        full_video = _normalize_video_tensor(
+            video,
+            height=self.height,
+            width=self.width,
+            device=self.device,
+        )
+        for chunk in video_chunks:
+            encoded_chunk = self.encode_chunk(
+                full_video,
+                chunk,
+                previous_noise_scale=noise_scale,
+                initial_noise_scale=init_noise_scale,
+            )
+            noise_scale = encoded_chunk.noise_scale
+            chunks.append(encoded_chunk)
+        return chunks
+
+    @torch.inference_mode()
+    def denoise_chunks(self, chunks: list[EncodedChunk]) -> list[DenoisedChunk]:
+        """Run DiT denoising over the encoded chunks."""
+        if not chunks:
+            raise ValueError("chunks must not be empty")
+        if self.prompt is None:
+            raise RuntimeError("Call prepare(prompt) before denoise_chunks(...)")
+
+        self.prepare(self.prompt)
+        outputs: list[DenoisedChunk] = []
+        for chunk in chunks:
+            denoised_chunk = self.denoise_chunk(chunk)
+            if denoised_chunk is not None:
+                outputs.append(denoised_chunk)
+        return outputs
+
+    @torch.inference_mode()
+    def denoise_chunk(self, chunk: EncodedChunk) -> DenoisedChunk | None:
+        """Run DiT on one encoded chunk and return a decodable latent when available."""
+        if self.prompt is None:
+            raise RuntimeError("Call prepare(prompt) before denoise_chunk(...)")
+
+        if self._next_chunk_index == 0:
+            if self.mode == "single":
+                denoised_pred = self.pipeline_manager.prepare_pipeline(
+                    text_prompts=[self.prompt],
+                    noise=chunk.noisy_latents,
+                    current_start=chunk.current_start,
+                    current_end=chunk.current_end,
+                )
+            else:
+                denoised_pred = self.pipeline_manager.prepare_pipeline(
+                    text_prompts=[self.prompt],
+                    noise=chunk.noisy_latents,
+                    current_start=chunk.current_start,
+                    current_end=chunk.current_end,
+                    batch_denoise=False,
+                )
+            self._next_chunk_index += 1
+            return DenoisedChunk(denoised_pred=denoised_pred, last_frame_only=False)
+
+        current_start = chunk.current_start
+        current_end = chunk.current_end
+
+        if current_start // self.pipeline_manager.pipeline.frame_seq_length >= self.pipeline_manager.t_refresh:
+            current_start = self.pipeline_manager.pipeline.kv_cache_length - self.pipeline_manager.pipeline.frame_seq_length
+            current_end = current_start + (self.chunk_size // 4) * self.pipeline_manager.pipeline.frame_seq_length
+
+        if self.mode == "single":
+            denoised_pred = self.pipeline_manager.pipeline.inference_stream(
+                noise=chunk.noisy_latents,
+                current_start=current_start,
+                current_end=current_end,
+                current_step=chunk.current_step,
+            )
+            self.pipeline_manager.processed += 1
+            self._next_chunk_index += 1
+            if self.pipeline_manager.processed < self.num_steps:
+                return None
+            return DenoisedChunk(denoised_pred=denoised_pred, last_frame_only=True)
+
+        denoised_pred = self.pipeline_manager.pipeline.inference_wo_batch(
+            noise=chunk.noisy_latents,
+            current_start=current_start,
+            current_end=current_end,
+            current_step=chunk.current_step,
+        )
+        self.pipeline_manager.processed += 1
+        self._next_chunk_index += 1
+        return DenoisedChunk(denoised_pred=denoised_pred, last_frame_only=True)
+
+    @torch.inference_mode()
+    def decode_chunks(self, chunks: list[DenoisedChunk]) -> np.ndarray:
+        """Decode denoised latent chunks into a `[T, H, W, C]` video array."""
+        if not chunks:
+            raise ValueError("chunks must not be empty")
+        decoded = [self.decode_chunk(chunk) for chunk in chunks]
+        return np.concatenate(decoded, axis=0)
+
+    @torch.inference_mode()
+    def decode_chunk(self, chunk: DenoisedChunk) -> np.ndarray:
+        """Decode one denoised latent chunk into `[T, H, W, C]` frames."""
+        return self.pipeline_manager._decode_video_array(
+            chunk.denoised_pred,
+            last_frame_only=chunk.last_frame_only,
+        )
+
+    @torch.inference_mode()
+    def __call__(self, video: str | Path | torch.Tensor) -> np.ndarray:
+        """Run the full staged pipeline after `prepare(prompt)` has been called."""
+        encoded = self.encode_video(video)
+        denoised = self.denoise_chunks(encoded)
+        return self.decode_chunks(denoised)

streamv2v/__init__.py

@@ -0,0 +1,14 @@
+"""StreamDiffusionV2 inference package."""
+
+from streamv2v.api import StreamVideoToVideo, run_video_to_video
+from streamv2v.inference_common import load_mp4_as_tensor
+
+__all__ = [
+    "StreamVideoToVideo",
+    "load_mp4_as_tensor",
+    "run_video_to_video",
+    "inference",
+    "inference_common",
+    "inference_pipe",
+    "inference_wo_batch",
+]

streamv2v/api.py

@@ -0,0 +1,309 @@
+"""Public Python API for simple offline video-to-video inference."""
+
+from __future__ import annotations
+
+from contextlib import ExitStack
+from importlib.resources import as_file, files
+import os
+from pathlib import Path
+import shutil
+import socket
+import subprocess
+import sys
+import tempfile
+from typing import Literal, Sequence
+
+import torch
+
+from streamv2v.inference_common import load_mp4_as_tensor, normalize_acceleration_flags
+
+
+InferenceMode = Literal["single", "single-wo", "pipe"]
+
+_SINGLE_MODE_TO_MODULE = {
+    "single": "streamv2v.inference",
+    "single-wo": "streamv2v.inference_wo_batch",
+}
+
+
+def _resolve_default_config_path(resource_stack: ExitStack) -> str:
+    resource = files("streamv2v.configs").joinpath("wan_causal_dmd_v2v.yaml")
+    return str(resource_stack.enter_context(as_file(resource)))
+
+
+def _normalize_gpu_ids(gpu_ids: int | Sequence[int] | None) -> list[int] | None:
+    if gpu_ids is None:
+        return None
+    if isinstance(gpu_ids, int):
+        return [gpu_ids]
+    return [int(gpu_id) for gpu_id in gpu_ids]
+
+
+def _normalize_device_gpu_id(device: str | torch.device | None) -> list[int] | None:
+    if device is None:
+        return None
+    device_str = str(device)
+    if not device_str.startswith("cuda:"):
+        return None
+    return [int(device_str.split(":", 1)[1])]
+
+
+def _resolve_single_gpu_id(gpu_ids: list[int] | None) -> int | None:
+    if gpu_ids is None:
+        return None
+    if len(gpu_ids) != 1:
+        raise ValueError("single and single-wo modes accept exactly one GPU id")
+    return int(gpu_ids[0])
+
+
+def _pick_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return int(sock.getsockname()[1])
+
+
+def _build_common_args(
+    *,
+    config_path: str,
+    checkpoint_folder: str,
+    video_path: str,
+    prompt_file_path: str,
+    output_folder: str,
+    noise_scale: float,
+    height: int,
+    width: int,
+    fps: int,
+    step: int,
+    seed: int,
+    model_type: str,
+    profile: bool,
+    use_taehv: bool,
+    use_tensorrt: bool,
+    fast: bool,
+) -> list[str]:
+    args = [
+        "--config_path",
+        config_path,
+        "--checkpoint_folder",
+        checkpoint_folder,
+        "--output_folder",
+        output_folder,
+        "--prompt_file_path",
+        prompt_file_path,
+        "--video_path",
+        video_path,
+        "--noise_scale",
+        str(noise_scale),
+        "--height",
+        str(height),
+        "--width",
+        str(width),
+        "--fps",
+        str(fps),
+        "--step",
+        str(step),
+        "--seed",
+        str(seed),
+        "--model_type",
+        model_type,
+    ]
+    if profile:
+        args.append("--profile")
+    if use_taehv:
+        args.append("--use_taehv")
+    if use_tensorrt:
+        args.append("--use_tensorrt")
+    if fast:
+        args.append("--fast")
+    return args
+
+
+class StreamVideoToVideo:
+    """Convenience wrapper around the offline Python entrypoints."""
+
+    def __init__(
+        self,
+        checkpoint_folder: str,
+        mode: InferenceMode = "single",
+        *,
+        config_path: str | None = None,
+        device: str | torch.device | None = None,
+        gpu_ids: int | Sequence[int] | None = None,
+        num_gpus: int | None = None,
+        noise_scale: float = 0.8,
+        height: int = 480,
+        width: int = 832,
+        fps: int = 16,
+        step: int = 2,
+        seed: int = 0,
+        model_type: str = "T2V-1.3B",
+        use_taehv: bool = False,
+        use_tensorrt: bool = False,
+        fast: bool = False,
+        profile: bool = False,
+        schedule_block: bool = False,
+    ) -> None:
+        self.checkpoint_folder = checkpoint_folder
+        self.mode = mode
+        self.config_path = config_path
+        self.device = device
+        self.gpu_ids = gpu_ids
+        self.num_gpus = num_gpus
+        self.noise_scale = noise_scale
+        self.height = height
+        self.width = width
+        self.fps = fps
+        self.step = step
+        self.seed = seed
+        self.model_type = model_type
+        self.use_taehv = use_taehv
+        self.use_tensorrt = use_tensorrt
+        self.fast = fast
+        self.profile = profile
+        self.schedule_block = schedule_block
+
+    def generate(self, video_path: str, prompt: str) -> torch.Tensor:
+        with tempfile.TemporaryDirectory(prefix="streamv2v_generate_") as temp_dir:
+            output_path = os.path.join(temp_dir, "output.mp4")
+            self.run_video(video_path=video_path, prompt=prompt, output_path=output_path)
+            return load_mp4_as_tensor(output_path, normalize=False)
+
+    def run_video(self, video_path: str, prompt: str, output_path: str) -> str:
+        return run_video_to_video(
+            checkpoint_folder=self.checkpoint_folder,
+            video_path=video_path,
+            prompt=prompt,
+            output_path=output_path,
+            mode=self.mode,
+            config_path=self.config_path,
+            device=self.device,
+            gpu_ids=self.gpu_ids,
+            num_gpus=self.num_gpus,
+            noise_scale=self.noise_scale,
+            height=self.height,
+            width=self.width,
+            fps=self.fps,
+            step=self.step,
+            seed=self.seed,
+            model_type=self.model_type,
+            use_taehv=self.use_taehv,
+            use_tensorrt=self.use_tensorrt,
+            fast=self.fast,
+            profile=self.profile,
+            schedule_block=self.schedule_block,
+        )
+
+
+def run_video_to_video(
+    *,
+    checkpoint_folder: str,
+    video_path: str,
+    prompt: str,
+    output_path: str,
+    mode: InferenceMode = "single",
+    config_path: str | None = None,
+    device: str | torch.device | None = None,
+    gpu_ids: int | Sequence[int] | None = None,
+    num_gpus: int | None = None,
+    noise_scale: float = 0.8,
+    height: int = 480,
+    width: int = 832,
+    fps: int = 16,
+    step: int = 2,
+    seed: int = 0,
+    model_type: str = "T2V-1.3B",
+    use_taehv: bool = False,
+    use_tensorrt: bool = False,
+    fast: bool = False,
+    profile: bool = False,
+    schedule_block: bool = False,
+) -> str:
+    """Run offline video-to-video inference from Python."""
+    flags = normalize_acceleration_flags(
+        {
+            "use_taehv": use_taehv,
+            "use_tensorrt": use_tensorrt,
+            "fast": fast,
+        }
+    )
+    use_taehv = bool(flags["use_taehv"])
+    use_tensorrt = bool(flags["use_tensorrt"])
+    fast = bool(flags["fast"])
+
+    requested_gpu_ids = _normalize_gpu_ids(gpu_ids)
+    device_gpu_ids = _normalize_device_gpu_id(device)
+    if requested_gpu_ids is None and device_gpu_ids is not None:
+        requested_gpu_ids = device_gpu_ids
+
+    if mode == "pipe":
+        if num_gpus is None:
+            num_gpus = len(requested_gpu_ids) if requested_gpu_ids is not None else 2
+        if requested_gpu_ids is not None and len(requested_gpu_ids) != num_gpus:
+            raise ValueError("num_gpus must match len(gpu_ids) for pipe mode")
+    elif num_gpus is not None and num_gpus != 1:
+        raise ValueError("num_gpus is only used for pipe mode")
+
+    resource_stack = ExitStack()
+    try:
+        resolved_config_path = config_path or _resolve_default_config_path(resource_stack)
+        output_file = Path(output_path)
+        output_file.parent.mkdir(parents=True, exist_ok=True)
+
+        with tempfile.TemporaryDirectory(prefix="streamv2v_api_") as temp_dir:
+            temp_dir_path = Path(temp_dir)
+            prompt_path = temp_dir_path / "prompt.txt"
+            prompt_path.write_text(prompt + "\n", encoding="utf-8")
+            temp_output_dir = temp_dir_path / "outputs"
+            temp_output_dir.mkdir(parents=True, exist_ok=True)
+
+            common_args = _build_common_args(
+                config_path=resolved_config_path,
+                checkpoint_folder=checkpoint_folder,
+                video_path=video_path,
+                prompt_file_path=str(prompt_path),
+                output_folder=str(temp_output_dir),
+                noise_scale=noise_scale,
+                height=height,
+                width=width,
+                fps=fps,
+                step=step,
+                seed=seed,
+                model_type=model_type,
+                profile=profile,
+                use_taehv=use_taehv,
+                use_tensorrt=use_tensorrt,
+                fast=fast,
+            )
+
+            env = os.environ.copy()
+            if requested_gpu_ids is not None and mode == "pipe":
+                env["CUDA_VISIBLE_DEVICES"] = ",".join(str(gpu_id) for gpu_id in requested_gpu_ids)
+
+            if mode == "pipe":
+                cmd = [
+                    sys.executable,
+                    "-m",
+                    "torch.distributed.run",
+                    f"--nproc_per_node={num_gpus}",
+                    f"--master_port={_pick_free_port()}",
+                    "-m",
+                    "streamv2v.inference_pipe",
+                    *common_args,
+                ]
+                if schedule_block:
+                    cmd.append("--schedule_block")
+            else:
+                module_name = _SINGLE_MODE_TO_MODULE.get(mode)
+                if module_name is None:
+                    raise ValueError(f"Unsupported mode: {mode}")
+                cmd = [sys.executable, "-m", module_name, *common_args]
+                single_gpu_id = _resolve_single_gpu_id(requested_gpu_ids)
+                if single_gpu_id is not None:
+                    cmd.extend(["--gpu_id", str(single_gpu_id)])
+
+            subprocess.run(cmd, env=env, check=True)
+            generated_path = temp_output_dir / "output_000.mp4"
+            shutil.copy2(generated_path, output_file)
+        return str(output_file)
+    finally:
+        resource_stack.close()

streamv2v/communication/__init__.py

@@ -0,0 +1,27 @@
+"""
+Communication module for distributed inference pipeline.
+
+This module provides abstractions for distributed communication operations,
+model data transfer, and buffer management in the StreamDiffusionV2 pipeline.
+"""
+
+from .distributed_communicator import DistributedCommunicator
+from .model_data_transfer import ModelDataTransfer
+from .buffer_manager import BufferManager
+from .data_containers import LatentData, KVCacheData, CommunicationConfig
+from .kv_cache_manager import KVCacheManager
+from .utils import CommunicationTags, init_distributed, setup_logging, compute_balanced_split
+
+__all__ = [
+    'DistributedCommunicator',
+    'ModelDataTransfer', 
+    'BufferManager',
+    'LatentData',
+    'KVCacheData',
+    'CommunicationConfig',
+    'KVCacheManager',
+    'CommunicationTags',
+    'init_distributed',
+    'setup_logging',
+    'compute_balanced_split'
+]

streamv2v/communication/buffer_manager.py

@@ -0,0 +1,252 @@
+"""
+Buffer manager for efficient GPU memory management.
+
+This module provides a buffer pool manager to avoid repeated GPU memory allocations
+during distributed communication operations.
+"""
+
+import torch
+from typing import Dict, List, Tuple, Optional
+import threading
+import logging
+from .data_containers import CommunicationConfig
+
+
+BufferKey = Tuple[Tuple[int, ...], torch.dtype]
+
+
+class BufferManager:
+    """
+    Manages GPU buffer pools to avoid repeated allocations.
+    
+    This class maintains pools of pre-allocated GPU tensors that can be reused
+    across communication operations, reducing memory allocation overhead.
+    """
+    
+    def __init__(self, device: torch.device, config: Optional[CommunicationConfig] = None):
+        """
+        Initialize the buffer manager.
+        
+        Args:
+            device: GPU device for buffer allocation
+            config: Communication configuration
+        """
+        self.device = device
+        self.config = config or CommunicationConfig()
+        
+        # Buffer pools: {(shape, dtype): [tensor1, tensor2, ...]}
+        self.free_buffers: Dict[BufferKey, List[torch.Tensor]] = {}
+        self.free_buffers_origin: Dict[BufferKey, List[torch.Tensor]] = {}
+        self.free_buffers_kv: Dict[BufferKey, List[torch.Tensor]] = {}
+        self.free_buffers_misc: Dict[BufferKey, List[torch.Tensor]] = {}
+        
+        # Thread safety
+        self._lock = threading.Lock()
+        
+        # Statistics
+        self.allocation_count = 0
+        self.reuse_count = 0
+        self.total_allocated_memory = 0
+        
+        # Setup logging
+        self.logger = logging.getLogger(f"BufferManager_{device}")
+        self.logger.propagate = False
+        if not self.logger.handlers:
+            handler = logging.StreamHandler()
+            # handler.setLevel(logging.DEBUG)
+            formatter = logging.Formatter(
+                f'[BufferManager {device}] %(asctime)s - %(name)s - %(levelname)s - %(message)s'
+            )
+            handler.setFormatter(formatter)
+            self.logger.addHandler(handler)
+        # self.logger.setLevel(logging.DEBUG)
+    
+    def get_buffer(self, shape: Tuple[int, ...], dtype: torch.dtype, 
+                   buffer_type: str = "latent") -> torch.Tensor:
+        """
+        Get or allocate a buffer with the specified shape and dtype.
+        
+        Args:
+            shape: Tensor shape
+            dtype: Tensor data type
+            buffer_type: Type of buffer ("latent", "origin", "kv")
+            
+        Returns:
+            Tensor buffer
+        """
+        with self._lock:
+            # Select the appropriate buffer pool
+            if buffer_type == "latent":
+                buffer_pool = self.free_buffers
+            elif buffer_type == "origin":
+                buffer_pool = self.free_buffers_origin
+            elif buffer_type == "kv":
+                buffer_pool = self.free_buffers_kv
+            elif buffer_type == "misc":
+                buffer_pool = self.free_buffers_misc
+            else:
+                raise ValueError(f"Unknown buffer type: {buffer_type}")
+            
+            # Try to reuse existing buffer
+            key = (tuple(shape), dtype)
+
+            if self.config.enable_buffer_reuse and key in buffer_pool and len(buffer_pool[key]) > 0:
+                buffer = buffer_pool[key].pop()
+                self.reuse_count += 1
+                self.logger.debug(f"Reused buffer of shape {shape}, dtype {dtype}, type {buffer_type}")
+                return buffer
+            
+            # Allocate new buffer
+            buffer = torch.empty(shape, dtype=dtype, device=self.device)
+            self.allocation_count += 1
+            self.total_allocated_memory += buffer.numel() * buffer.element_size()
+            
+            self.logger.debug(f"Allocated new buffer of shape {shape}, dtype {dtype}, type {buffer_type}")
+            return buffer
+    
+    def return_buffer(self, tensor: torch.Tensor, buffer_type: str = "latent") -> None:
+        """
+        Return a buffer to the pool for reuse.
+        
+        Args:
+            tensor: Tensor to return
+            buffer_type: Type of buffer ("latent", "origin", "kv")
+        """
+        if not self.config.enable_buffer_reuse:
+            return
+        
+        with self._lock:
+            # Select the appropriate buffer pool
+            if buffer_type == "latent":
+                buffer_pool = self.free_buffers
+            elif buffer_type == "origin":
+                buffer_pool = self.free_buffers_origin
+            elif buffer_type == "kv":
+                buffer_pool = self.free_buffers_kv
+            elif buffer_type == "misc":
+                buffer_pool = self.free_buffers_misc
+            else:
+                raise ValueError(f"Unknown buffer type: {buffer_type}")
+            
+            key = (tuple(tensor.shape), tensor.dtype)
+            
+            # Initialize pool for this shape if it doesn't exist
+            if key not in buffer_pool:
+                buffer_pool[key] = []
+            
+            # Add buffer to pool if not at capacity
+            if len(buffer_pool[key]) < self.config.buffer_pool_size:
+                buffer_pool[key].append(tensor)
+                self.logger.debug(
+                    f"Returned buffer of shape {tuple(tensor.shape)}, dtype {tensor.dtype}, type {buffer_type}"
+                )
+            else:
+                self.logger.debug(
+                    f"Buffer pool full for shape {tuple(tensor.shape)}, dtype {tensor.dtype}, type {buffer_type}, discarding"
+                )
+    
+    def clear_buffers(self, buffer_type: Optional[str] = None) -> None:
+        """
+        Clear buffer pools to free memory.
+        
+        Args:
+            buffer_type: Specific buffer type to clear, or None to clear all
+        """
+        with self._lock:
+            if buffer_type is None:
+                # Clear all buffer pools
+                self.free_buffers.clear()
+                self.free_buffers_origin.clear()
+                self.free_buffers_kv.clear()
+                self.free_buffers_misc.clear()
+                self.logger.info("Cleared all buffer pools")
+            else:
+                # Clear specific buffer pool
+                if buffer_type == "latent":
+                    self.free_buffers.clear()
+                elif buffer_type == "origin":
+                    self.free_buffers_origin.clear()
+                elif buffer_type == "kv":
+                    self.free_buffers_kv.clear()
+                elif buffer_type == "misc":
+                    self.free_buffers_misc.clear()
+                else:
+                    raise ValueError(f"Unknown buffer type: {buffer_type}")
+                self.logger.info(f"Cleared {buffer_type} buffer pool")
+    
+    def get_statistics(self) -> Dict[str, any]:
+        """
+        Get buffer manager statistics.
+        
+        Returns:
+            Dictionary containing statistics
+        """
+        with self._lock:
+            total_free_buffers = sum(len(pool) for pool in self.free_buffers.values())
+            total_free_buffers_origin = sum(len(pool) for pool in self.free_buffers_origin.values())
+            total_free_buffers_kv = sum(len(pool) for pool in self.free_buffers_kv.values())
+            total_free_buffers_misc = sum(len(pool) for pool in self.free_buffers_misc.values())
+
+            return {
+                "allocation_count": self.allocation_count,
+                "reuse_count": self.reuse_count,
+                "total_allocated_memory_bytes": self.total_allocated_memory,
+                "total_free_buffers": total_free_buffers,
+                "total_free_buffers_origin": total_free_buffers_origin,
+                "total_free_buffers_kv": total_free_buffers_kv,
+                "total_free_buffers_misc": total_free_buffers_misc,
+                "reuse_rate": self.reuse_count / max(1, self.allocation_count),
+                "buffer_pool_size": self.config.buffer_pool_size,
+                "enable_buffer_reuse": self.config.enable_buffer_reuse
+            }
+    
+    def print_statistics(self) -> None:
+        """Print buffer manager statistics."""
+        stats = self.get_statistics()
+        self.logger.info("Buffer Manager Statistics:")
+        for key, value in stats.items():
+            self.logger.info(f"  {key}: {value}")
+    
+    def preallocate_buffers(self, common_shapes: List[Tuple[Tuple[int, ...], torch.dtype, str]], 
+                          count_per_shape: int = 5) -> None:
+        """
+        Preallocate buffers for common shapes to reduce allocation overhead.
+        
+        Args:
+            common_shapes: List of (shape, dtype, buffer_type) tuples
+            count_per_shape: Number of buffers to preallocate per shape
+        """
+        with self._lock:
+            for shape, dtype, buffer_type in common_shapes:
+                for _ in range(count_per_shape):
+                    buffer = torch.empty(shape, dtype=dtype, device=self.device)
+                    
+                    # Select the appropriate buffer pool
+                    if buffer_type == "latent":
+                        buffer_pool = self.free_buffers
+                    elif buffer_type == "origin":
+                        buffer_pool = self.free_buffers_origin
+                    elif buffer_type == "kv":
+                        buffer_pool = self.free_buffers_kv
+                    elif buffer_type == "misc":
+                        buffer_pool = self.free_buffers_misc
+                    else:
+                        raise ValueError(f"Unknown buffer type: {buffer_type}")
+                    
+                    # Initialize pool for this shape if it doesn't exist
+                    key = (tuple(shape), dtype)
+                    if key not in buffer_pool:
+                        buffer_pool[key] = []
+                    
+                    buffer_pool[key].append(buffer)
+                    self.allocation_count += 1
+                    self.total_allocated_memory += buffer.numel() * buffer.element_size()
+            
+            self.logger.info(f"Preallocated {len(common_shapes) * count_per_shape} buffers")
+    
+    def __del__(self):
+        """Cleanup when the buffer manager is destroyed."""
+        try:
+            self.clear_buffers()
+        except Exception:
+            pass

streamv2v/communication/data_containers.py

@@ -0,0 +1,148 @@
+"""
+Data containers for communication operations.
+
+This module defines data structures used for communication between distributed ranks.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, List, Dict, Any
+import torch
+
+
+@dataclass
+class LatentData:
+    """
+    Container for latent data and related information.
+    
+    This class encapsulates all the data that needs to be transferred between ranks
+    during the inference pipeline.
+    """
+    chunk_idx: int
+    latents: torch.Tensor
+    original_latents: torch.Tensor
+    current_start: torch.Tensor
+    current_end: torch.Tensor
+    current_step: int
+    patched_x_shape: torch.Tensor
+    
+    def __post_init__(self):
+        """Validate tensor shapes and types after initialization."""
+        if not isinstance(self.latents, torch.Tensor):
+            raise TypeError("latents must be a torch.Tensor")
+        if not isinstance(self.original_latents, torch.Tensor):
+            raise TypeError("original_latents must be a torch.Tensor")
+        if not isinstance(self.current_start, torch.Tensor):
+            raise TypeError("current_start must be a torch.Tensor")
+        if not isinstance(self.current_end, torch.Tensor):
+            raise TypeError("current_end must be a torch.Tensor")
+        if not isinstance(self.patched_x_shape, torch.Tensor):
+            raise TypeError("patched_x_shape must be a torch.Tensor")
+
+
+@dataclass
+class KVCacheData:
+    """
+    Container for KV cache data.
+    
+    This class encapsulates key-value cache information for transformer blocks.
+    """
+    block_index: int
+    k_cache: torch.Tensor
+    v_cache: torch.Tensor
+    global_end_index: torch.Tensor
+    local_end_index: torch.Tensor
+    
+    def __post_init__(self):
+        """Validate tensor shapes and types after initialization."""
+        if not isinstance(self.k_cache, torch.Tensor):
+            raise TypeError("k_cache must be a torch.Tensor")
+        if not isinstance(self.v_cache, torch.Tensor):
+            raise TypeError("v_cache must be a torch.Tensor")
+        if not isinstance(self.global_end_index, torch.Tensor):
+            raise TypeError("global_end_index must be a torch.Tensor")
+        if not isinstance(self.local_end_index, torch.Tensor):
+            raise TypeError("local_end_index must be a torch.Tensor")
+
+
+@dataclass
+class CommunicationConfig:
+    """
+    Configuration for communication operations.
+    
+    This class holds configuration parameters for distributed communication.
+    """
+    max_outstanding: int = 1
+    buffer_pool_size: int = 10
+    enable_buffer_reuse: bool = True
+    communication_timeout: float = 30.0
+    enable_async_communication: bool = True
+    
+    def __post_init__(self):
+        """Validate configuration parameters."""
+        if self.max_outstanding < 1:
+            raise ValueError("max_outstanding must be at least 1")
+        if self.buffer_pool_size < 1:
+            raise ValueError("buffer_pool_size must be at least 1")
+        if self.communication_timeout <= 0:
+            raise ValueError("communication_timeout must be positive")
+
+
+@dataclass
+class BlockInterval:
+    """
+    Container for block interval information.
+    
+    This class represents a block interval [start, end) for a specific rank.
+    """
+    start: int
+    end: int
+    rank: int
+    
+    def __post_init__(self):
+        """Validate block interval parameters."""
+        if self.start < 0:
+            raise ValueError("start must be non-negative")
+        if self.end <= self.start:
+            raise ValueError("end must be greater than start")
+        if self.rank < 0:
+            raise ValueError("rank must be non-negative")
+    
+    @property
+    def size(self) -> int:
+        """Get the size of the block interval."""
+        return self.end - self.start
+    
+    def contains(self, block_index: int) -> bool:
+        """Check if the block interval contains the given block index."""
+        return self.start <= block_index < self.end
+
+
+@dataclass
+class PerformanceMetrics:
+    """
+    Container for performance metrics.
+    
+    This class holds timing and performance information for communication operations.
+    """
+    dit_time: float
+    total_time: float
+    communication_time: float
+    buffer_allocation_time: float
+    
+    def __post_init__(self):
+        """Validate performance metrics."""
+        if self.dit_time < 0:
+            raise ValueError("dit_time must be non-negative")
+        if self.total_time < 0:
+            raise ValueError("total_time must be non-negative")
+        if self.communication_time < 0:
+            raise ValueError("communication_time must be non-negative")
+        if self.buffer_allocation_time < 0:
+            raise ValueError("buffer_allocation_time must be non-negative")
+    
+    @property
+    def efficiency(self) -> float:
+        """Calculate communication efficiency (computation time / total time)."""
+        if self.total_time == 0:
+            return 0.0
+        return (self.total_time - self.communication_time) / self.total_time

streamv2v/communication/distributed_communicator.py

@@ -0,0 +1,392 @@
+"""
+Distributed communication abstraction layer.
+
+This module provides a high-level interface for distributed communication operations,
+encapsulating the low-level PyTorch distributed primitives.
+"""
+
+import torch
+import torch.distributed as dist
+from typing import List, Tuple, Optional, Any
+import logging
+import time
+from .utils import CommunicationTags, get_next_rank, get_prev_rank, CommunicationTimer
+from .data_containers import CommunicationConfig
+
+
+class DistributedCommunicator:
+    """
+    High-level interface for distributed communication operations.
+    
+    This class encapsulates all distributed communication operations, providing
+    a clean interface for sending and receiving tensors between ranks.
+    """
+    
+    def __init__(self, rank: int, world_size: int, device: torch.device, 
+                 config: Optional[CommunicationConfig] = None):
+        """
+        Initialize the distributed communicator.
+        
+        Args:
+            rank: Current rank
+            world_size: Total number of ranks
+            device: GPU device for communication
+            config: Communication configuration
+        """
+        self.rank = rank
+        self.world_size = world_size
+        self.device = device
+        self.config = config or CommunicationConfig()
+        
+        # Track outstanding operations
+        self.outstanding_operations: List[Any] = []
+        
+        # Setup logging
+        self.logger = logging.getLogger(f"DistributedCommunicator_rank_{rank}")
+        self.logger.propagate = False
+        if not self.logger.handlers:
+            handler = logging.StreamHandler()
+            formatter = logging.Formatter(
+                f'[Rank {rank}] %(asctime)s - %(name)s - %(levelname)s - %(message)s'
+            )
+            handler.setFormatter(formatter)
+            self.logger.addHandler(handler)
+        
+        # Validate distributed is initialized
+        if not dist.is_initialized():
+            raise RuntimeError("Distributed not initialized. Call init_distributed() first.")
+    
+    def send_tensor_async(self, tensor: torch.Tensor, dst: int, tag: int) -> Any:
+        """
+        Asynchronously send a tensor to the specified destination.
+        
+        Args:
+            tensor: Tensor to send
+            dst: Destination rank
+            tag: Communication tag
+            
+        Returns:
+            Work object for the send operation
+        """
+        if tensor.device != self.device:
+            raise ValueError(f"Tensor device {tensor.device} doesn't match communicator device {self.device}")
+        
+        work = dist.isend(tensor, dst=dst, tag=tag)
+        self.outstanding_operations.append(work)
+        
+        self.logger.debug(f"Started async send to rank {dst} with tag {tag}, tensor shape: {tensor.shape}")
+        return work
+    
+    def recv_tensor(self, src: int, tag: int, shape: Tuple[int, ...], 
+                   dtype: torch.dtype) -> torch.Tensor:
+        """
+        Receive a tensor from the specified source.
+        
+        Args:
+            src: Source rank
+            tag: Communication tag
+            shape: Expected tensor shape
+            dtype: Expected tensor dtype
+            
+        Returns:
+            Received tensor
+        """
+        tensor = torch.empty(shape, dtype=dtype, device=self.device)
+        
+        with CommunicationTimer(f"recv_tensor from rank {src}", self.logger):
+            dist.recv(tensor, src=src, tag=tag)
+        
+        self.logger.debug(f"Received tensor from rank {src} with tag {tag}, shape: {tensor.shape}")
+        return tensor
+    
+    def send_header_and_tensor_async(self, header: torch.Tensor, tensor: torch.Tensor,
+                                   dst: int, tag_header: int, tag_tensor: int) -> Tuple[Any, Any]:
+        """
+        Asynchronously send a header and tensor pair.
+        
+        Args:
+            header: Header tensor containing metadata
+            tensor: Data tensor
+            dst: Destination rank
+            tag_header: Tag for header
+            tag_tensor: Tag for tensor
+            
+        Returns:
+            Tuple of (header_work, tensor_work)
+        """
+        if header.device != self.device or tensor.device != self.device:
+            raise ValueError("Header and tensor must be on the same device as communicator")
+        
+        header_work = dist.isend(header, dst=dst, tag=tag_header)
+        tensor_work = dist.isend(tensor, dst=dst, tag=tag_tensor)
+        
+        self.outstanding_operations.extend([header_work, tensor_work])
+        
+        self.logger.debug(f"Started async send of header+tensor to rank {dst}, "
+                         f"header shape: {header.shape}, tensor shape: {tensor.shape}")
+        return header_work, tensor_work
+    
+    def recv_header_and_tensor(self, src: int, tag_header: int, tag_tensor: int, header_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Receive a header and tensor pair.
+        
+        Args:
+            src: Source rank
+            tag_header: Tag for header
+            tag_tensor: Tag for tensor
+            header_len: Length of header tensor to receive
+            
+        Returns:
+            Tuple of (header, tensor)
+        """
+        with CommunicationTimer(f"recv_header_and_tensor from rank {src}", self.logger):
+            # First receive the header to get tensor shape (length can vary)
+            header = torch.empty(header_len, dtype=torch.int64, device=self.device)
+            dist.recv(header, src=src, tag=tag_header)
+            
+            # Parse header to get tensor shape
+            chunk_idx, shape = self._parse_header(header)
+            
+            # Receive the tensor
+            tensor = torch.empty(shape, dtype=torch.bfloat16, device=self.device)
+            dist.recv(tensor, src=src, tag=tag_tensor)
+        
+        self.logger.debug(f"Received header+tensor from rank {src}, "
+                         f"header: {header.tolist()}, tensor shape: {tensor.shape}")
+        return header, tensor
+    
+    def send_latent_data_async(self, chunk_idx: int, latents: torch.Tensor,
+                             original_latents: torch.Tensor, patched_x_shape: torch.Tensor,
+                             current_start: torch.Tensor, current_end: torch.Tensor,
+                             current_step: int) -> List[Any]:
+        """
+        Asynchronously send all latent data components.
+        
+        Args:
+            chunk_idx: Chunk index
+            latents: Latent tensor
+            original_latents: Original latent tensor
+            patched_x_shape: Patched x shape tensor
+            current_start: Current start indices
+            current_end: Current end indices
+            current_step: Current step
+            
+        Returns:
+            List of work objects for all send operations
+        """
+        dst = get_next_rank(self.rank, self.world_size)
+        work_objects = []
+        
+        # Create headers
+        latent_header = self._create_header(chunk_idx, latents.shape)
+        origin_header = self._create_header(chunk_idx, original_latents.shape)
+        
+        # Create start/end/step tensor
+        start_end_step = torch.cat([
+            current_start, 
+            current_end, 
+            torch.tensor([current_step], dtype=torch.int64, device=self.device)
+        ], dim=0)
+        
+        # Send all components asynchronously
+        work_objects.extend(self.send_header_and_tensor_async(
+            latent_header, latents, dst, CommunicationTags.LATENT_HDR, CommunicationTags.LATENT_PAY
+        ))
+        
+        work_objects.extend(self.send_header_and_tensor_async(
+            origin_header, original_latents, dst, 
+            CommunicationTags.LATENT_ORIGIN_HDR, CommunicationTags.LATENT_ORIGIN_PAY
+        ))
+        
+        work_objects.append(self.send_tensor_async(
+            patched_x_shape, dst, CommunicationTags.PATCHED_X_SHAPE
+        ))
+        
+        work_objects.append(self.send_tensor_async(
+            start_end_step, dst, CommunicationTags.START_END_STEP
+        ))
+        
+        self.logger.debug(f"Started async send of latent data to rank {dst}, chunk_idx: {chunk_idx}")
+        return work_objects
+    
+    def recv_latent_data_async(self, num_steps: int, buffer_manager) -> Tuple[int, torch.Tensor, torch.Tensor, 
+                                                                           torch.Tensor, torch.Tensor, int, torch.Tensor]:
+        """
+        Asynchronously receive all latent data components.
+        
+        Args:
+            num_steps: Number of denoising steps
+            buffer_manager: Buffer manager for tensor allocation
+            
+        Returns:
+            Tuple of (chunk_idx, latents, original_latents, current_start, current_end, current_step, patched_x_shape)
+        """
+        src = get_prev_rank(self.rank, self.world_size)
+        
+        with CommunicationTimer(f"recv_latent_data_async from rank {src}", self.logger):
+            # Receive latent header (length 4): [i, bsz, slen, cch]
+            latent_header = buffer_manager.get_buffer((4,), torch.int64, "misc")
+            dist.recv(latent_header, src=src, tag=CommunicationTags.LATENT_HDR)
+            chunk_idx, latent_shape = self._parse_header(latent_header)
+            # header no longer needed
+            buffer_manager.return_buffer(latent_header, "misc")
+            # Allocate or reuse buffer for latents: shape (bsz, slen, cch)
+            latents = buffer_manager.get_buffer(tuple(latent_shape), torch.bfloat16, "latent")
+            dist.recv(latents, src=src, tag=CommunicationTags.LATENT_PAY)
+
+            # Receive original latent header (length 6): [i, bsz, cch, tlen, hh, ww]
+            origin_header = buffer_manager.get_buffer((6,), torch.int64, "misc")
+            dist.recv(origin_header, src=src, tag=CommunicationTags.LATENT_ORIGIN_HDR)
+            _, origin_shape = self._parse_header(origin_header)
+            # header no longer needed
+            buffer_manager.return_buffer(origin_header, "misc")
+            # Allocate or reuse buffer for original latents: shape (bsz, cch, tlen, hh, ww)
+            original_latents = buffer_manager.get_buffer(tuple(origin_shape), torch.bfloat16, "origin")
+            dist.recv(original_latents, src=src, tag=CommunicationTags.LATENT_ORIGIN_PAY)
+
+            # Receive patched_x_shape (length 5, int64)
+            patched_x_shape = buffer_manager.get_buffer((5,), torch.int64, "misc")
+            dist.recv(patched_x_shape, src=src, tag=CommunicationTags.PATCHED_X_SHAPE)
+
+            # Receive start_end_step (length 2*num_steps+1, int64)
+            start_end_step = buffer_manager.get_buffer((2 * num_steps + 1,), torch.int64, "misc")
+            dist.recv(start_end_step, src=src, tag=CommunicationTags.START_END_STEP)
+
+            # Parse start/end/step into dedicated misc buffers, then release the combined vector
+            current_start = buffer_manager.get_buffer((num_steps,), torch.int64, "misc")
+            current_end = buffer_manager.get_buffer((num_steps,), torch.int64, "misc")
+            current_start.copy_(start_end_step[:num_steps])
+            current_end.copy_(start_end_step[num_steps:-1])
+            current_step = int(start_end_step[-1].item())
+            # Release the temporary combined buffer
+            buffer_manager.return_buffer(start_end_step, "misc")
+        
+        self.logger.debug(f"Received latent data from rank {src}, chunk_idx: {chunk_idx}")
+        return chunk_idx, latents, original_latents, current_start, current_end, current_step, patched_x_shape
+
+    def send_prompt_async(self, prompt: str, device: torch.device) -> List[Any]:
+        work_objects = []
+        dst = get_next_rank(self.rank, self.world_size)
+
+        # Encode to bytes
+        encoded = prompt.encode("utf-8")
+        data = torch.ByteTensor(list(encoded)).to(device)
+
+        # Send length first
+        length = torch.tensor([len(data)], dtype=torch.int64, device=data.device)
+        work_objects.append(dist.isend(length, dst=dst, tag=CommunicationTags.UPDATED_PROMPT_LENGTH))
+
+        # Then send the content
+        work_objects.append(dist.isend(data, dst=dst, tag=CommunicationTags.UPDATED_PROMPT))
+
+        return work_objects
+
+    def recv_prompt_async(self) -> str:
+        src = get_prev_rank(self.rank, self.world_size)
+
+        # Receive length first
+        length = torch.empty(1, dtype=torch.int64, device=self.device)
+        dist.recv(length, src=src, tag=CommunicationTags.UPDATED_PROMPT_LENGTH)
+
+        # Then receive the content
+        prompt = torch.empty(length.item(), dtype=torch.uint8, device=self.device)
+        dist.recv(prompt, src=src, tag=CommunicationTags.UPDATED_PROMPT)
+
+        return bytes(prompt.cpu().tolist()).decode("utf-8")    
+    
+    def broadcast_tensor(self, tensor: torch.Tensor, src: int) -> None:
+        """
+        Broadcast a tensor from source to all ranks.
+        
+        Args:
+            tensor: Tensor to broadcast
+            src: Source rank
+        """
+        with CommunicationTimer(f"broadcast_tensor from rank {src}", self.logger):
+            dist.broadcast(tensor, src=src)
+        
+        self.logger.debug(f"Broadcasted tensor from rank {src}, shape: {tensor.shape}")
+    
+    def all_gather_tensors(self, tensor: torch.Tensor) -> List[torch.Tensor]:
+        """
+        Gather tensors from all ranks.
+        
+        Args:
+            tensor: Local tensor to gather
+            
+        Returns:
+            List of tensors from all ranks
+        """
+        with CommunicationTimer("all_gather_tensors", self.logger):
+            gather_list = [torch.zeros_like(tensor) for _ in range(self.world_size)]
+            dist.all_gather(gather_list, tensor)
+        
+        self.logger.debug(f"Gathered tensors from all ranks, local shape: {tensor.shape}")
+        return gather_list
+    
+    def wait_for_outstanding(self, max_outstanding: Optional[int] = None) -> None:
+        """
+        Wait for outstanding operations to complete.
+        
+        Args:
+            max_outstanding: Maximum number of outstanding operations to keep
+        """
+        max_outstanding = max_outstanding or self.config.max_outstanding
+        
+        while len(self.outstanding_operations) >= max_outstanding:
+            if not self.outstanding_operations:
+                break
+            
+            # Wait for the oldest operation
+            oldest_operations = self.outstanding_operations.pop(0)
+            
+            # Handle both single work objects and lists of work objects
+            if isinstance(oldest_operations, (list, tuple)):
+                for work in oldest_operations:
+                    try:
+                        work.wait()
+                    except Exception as e:
+                        self.logger.error(f"Error waiting for outstanding operation: {e}")
+                        raise
+            else:
+                try:
+                    oldest_operations.wait()
+                except Exception as e:
+                    self.logger.error(f"Error waiting for outstanding operation: {e}")
+                    raise
+        
+        self.logger.debug(f"Outstanding operations: {len(self.outstanding_operations)}")
+    
+    def barrier(self) -> None:
+        """Synchronize all ranks."""
+        with CommunicationTimer("barrier", self.logger):
+            dist.barrier()
+    
+    def _create_header(self, chunk_idx: int, shape: Tuple[int, ...]) -> torch.Tensor:
+        """Create a header tensor for communication."""
+        header_data = [chunk_idx] + list(shape)
+        return torch.tensor(header_data, dtype=torch.int64, device=self.device)
+    
+    def _parse_header(self, header: torch.Tensor) -> Tuple[int, Tuple[int, ...]]:
+        """Parse a header tensor to extract metadata."""
+        header_list = header.tolist()
+        chunk_idx = int(header_list[0])
+        shape = tuple(int(x) for x in header_list[1:])
+        return chunk_idx, shape
+    
+    def get_statistics(self) -> dict:
+        """Get communication statistics."""
+        return {
+            "rank": self.rank,
+            "world_size": self.world_size,
+            "outstanding_operations": len(self.outstanding_operations),
+            "max_outstanding": self.config.max_outstanding,
+            "device": str(self.device)
+        }
+    
+    def print_statistics(self) -> None:
+        """Print communication statistics."""
+        stats = self.get_statistics()
+        self.logger.info("Distributed Communicator Statistics:")
+        for key, value in stats.items():
+            self.logger.info(f"  {key}: {value}")

streamv2v/communication/kv_cache_manager.py

@@ -0,0 +1,263 @@
+"""
+KV Cache management for distributed inference.
+
+This module provides functionality for managing and rebalancing KV caches
+across distributed ranks during inference.
+"""
+
+import torch
+import torch.distributed as dist
+from typing import List, Dict, Tuple, Optional
+import logging
+from .utils import CommunicationTags, CommunicationTimer
+from .data_containers import KVCacheData, BlockInterval
+
+
+class KVCacheManager:
+    """
+    Manages KV cache operations for distributed inference.
+    
+    This class handles KV cache broadcasting, rebalancing, and ownership
+    management across distributed ranks.
+    """
+    
+    def __init__(self, pipeline, device: torch.device):
+        """
+        Initialize the KV cache manager.
+        
+        Args:
+            pipeline: The inference pipeline containing KV caches
+            device: GPU device for operations
+        """
+        self.pipeline = pipeline
+        self.device = device
+        self.frame_seq_length = pipeline.frame_seq_length
+        self.time_step_length = len(pipeline.denoising_step_list)
+        
+        # Setup logging
+        self.logger = logging.getLogger(f"KVCacheManager_{device}")
+        self.logger.propagate = False
+        if not self.logger.handlers:
+            handler = logging.StreamHandler()
+            formatter = logging.Formatter(
+                f'[KVCacheManager {device}] %(asctime)s - %(name)s - %(levelname)s - %(message)s'
+            )
+            handler.setFormatter(formatter)
+            self.logger.addHandler(handler)
+    
+    def broadcast_kv_blocks(self, block_indices: List[int], donor_rank: int) -> None:
+        """
+        Broadcast kv_cache1 entries for the specified block indices from donor_rank to all ranks.
+        
+        This ensures the receiver rank has the up-to-date KV cache when ownership moves.
+        
+        Args:
+            block_indices: List of block indices to broadcast
+            donor_rank: Rank that owns the KV cache data
+        """
+        if len(block_indices) == 0:
+            return
+        
+        rank = dist.get_rank()
+        
+        with CommunicationTimer(f"broadcast_kv_blocks from rank {donor_rank}", self.logger):
+            for bi in block_indices:
+                # Broadcast key cache
+                if self.pipeline.kv_cache1[bi]['k'].device != self.device:
+                    self.pipeline.kv_cache1[bi]['k'] = self.pipeline.kv_cache1[bi]['k'].to(self.device)
+                    self.pipeline.kv_cache1[bi]['v'] = self.pipeline.kv_cache1[bi]['v'].to(self.device)
+                
+                dist.barrier()
+
+                dist.broadcast(self.pipeline.kv_cache1[bi]['k'], src=donor_rank)
+                # Broadcast value cache
+                dist.broadcast(self.pipeline.kv_cache1[bi]['v'], src=donor_rank)
+                # Broadcast global end index
+                dist.broadcast(self.pipeline.kv_cache1[bi]['global_end_index'], src=donor_rank)
+                # Broadcast local end index
+                dist.broadcast(self.pipeline.kv_cache1[bi]['local_end_index'], src=donor_rank)
+                
+                # Adjust global_end_index for the receiving rank
+                if donor_rank > rank:
+                    self.pipeline.kv_cache1[bi]['global_end_index'] += self.frame_seq_length * (donor_rank - rank) * self.time_step_length
+        
+        self.logger.debug(f"Broadcasted KV cache for blocks {block_indices} from rank {donor_rank}")
+    
+    def compute_block_owners(self, block_intervals: torch.Tensor, total_blocks: int) -> torch.Tensor:
+        """
+        Given block intervals in [start, end) format for all ranks, return a tensor
+        where each entry is the owner rank of that block index.
+        
+        Args:
+            block_intervals: Block intervals for all ranks [world_size, 2]
+            total_blocks: Total number of blocks
+            
+        Returns:
+            Tensor of length total_blocks with owner ranks
+        """
+        world_size = block_intervals.shape[0]
+        owners = torch.full((total_blocks,), -1, dtype=torch.int64, device=block_intervals.device)
+        
+        for r in range(world_size):
+            s = int(block_intervals[r, 0].item())
+            e = int(block_intervals[r, 1].item())
+            if e > s:
+                owners[s:e] = r
+        
+        self.logger.debug(f"Computed block owners: {owners.tolist()}")
+        return owners
+    
+    def rebalance_kv_cache_by_diff(self, old_block_intervals: torch.Tensor, 
+                                  new_block_intervals: torch.Tensor, total_blocks: int) -> None:
+        """
+        Compare ownership from old to new intervals and broadcast KV cache for blocks whose owner changes.
+        
+        For each moved block i, use the previous owner's rank as src to broadcast
+        pipeline.kv_cache1[i]['k'/'v'/...] to all ranks so the new owner has the correct state.
+        
+        Args:
+            old_block_intervals: Previous block intervals [world_size, 2]
+            new_block_intervals: New block intervals [world_size, 2]
+            total_blocks: Total number of blocks
+        """
+        with CommunicationTimer("rebalance_kv_cache_by_diff", self.logger):
+            old_owners = self.compute_block_owners(old_block_intervals, total_blocks)
+            new_owners = self.compute_block_owners(new_block_intervals, total_blocks)
+            
+            # Find blocks that changed ownership
+            moved_by_src = {}
+            for i in range(total_blocks):
+                o = int(old_owners[i].item())
+                n = int(new_owners[i].item())
+                if o != n and o >= 0:
+                    if o not in moved_by_src:
+                        moved_by_src[o] = []
+                    moved_by_src[o].append(i)
+            
+            # Synchronize before broadcasting
+            dist.barrier()
+            
+            # Broadcast per donor rank (can batch multiple blocks per src)
+            for src, blocks in moved_by_src.items():
+                self.broadcast_kv_blocks(blocks, donor_rank=src)
+        
+        self.logger.info(f"Rebalanced KV cache: {len(moved_by_src)} ranks had ownership changes")
+    
+    def get_kv_cache_statistics(self, block_intervals: torch.Tensor, total_blocks: int) -> Dict[str, any]:
+        """
+        Get statistics about KV cache distribution.
+        
+        Args:
+            block_intervals: Current block intervals [world_size, 2]
+            total_blocks: Total number of blocks
+            
+        Returns:
+            Dictionary containing KV cache statistics
+        """
+        owners = self.compute_block_owners(block_intervals, total_blocks)
+        
+        # Count blocks per rank
+        block_counts = {}
+        for rank in range(block_intervals.shape[0]):
+            block_counts[rank] = int((owners == rank).sum().item())
+        
+        # Calculate memory usage per rank (approximate)
+        memory_per_block = 0
+        if hasattr(self.pipeline, 'kv_cache1') and len(self.pipeline.kv_cache1) > 0:
+            # Estimate memory per block based on first block
+            first_block = self.pipeline.kv_cache1[0]
+            if 'k' in first_block and 'v' in first_block:
+                k_memory = first_block['k'].numel() * first_block['k'].element_size()
+                v_memory = first_block['v'].numel() * first_block['v'].element_size()
+                memory_per_block = k_memory + v_memory
+        
+        memory_usage = {rank: block_counts[rank] * memory_per_block for rank in block_counts}
+        
+        return {
+            "block_counts": block_counts,
+            "memory_usage_bytes": memory_usage,
+            "total_blocks": total_blocks,
+            "memory_per_block_bytes": memory_per_block,
+            "frame_seq_length": self.frame_seq_length
+        }
+    
+    def print_kv_cache_statistics(self, block_intervals: torch.Tensor, total_blocks: int) -> None:
+        """
+        Print KV cache statistics.
+        
+        Args:
+            block_intervals: Current block intervals [world_size, 2]
+            total_blocks: Total number of blocks
+        """
+        stats = self.get_kv_cache_statistics(block_intervals, total_blocks)
+        
+        self.logger.info("KV Cache Statistics:")
+        self.logger.info(f"  Total blocks: {stats['total_blocks']}")
+        self.logger.info(f"  Memory per block: {stats['memory_per_block_bytes']} bytes")
+        self.logger.info(f"  Frame sequence length: {stats['frame_seq_length']}")
+        
+        self.logger.info("  Block distribution:")
+        for rank, count in stats['block_counts'].items():
+            memory_mb = stats['memory_usage_bytes'][rank] / (1024 * 1024)
+            self.logger.info(f"    Rank {rank}: {count} blocks, {memory_mb:.2f} MB")
+    
+    def validate_kv_cache_consistency(self, block_intervals: torch.Tensor, total_blocks: int) -> bool:
+        """
+        Validate that KV cache ownership is consistent with block intervals.
+        
+        Args:
+            block_intervals: Current block intervals [world_size, 2]
+            total_blocks: Total number of blocks
+            
+        Returns:
+            True if consistent, False otherwise
+        """
+        owners = self.compute_block_owners(block_intervals, total_blocks)
+        
+        # Check that all blocks have owners
+        unowned_blocks = (owners == -1).sum().item()
+        if unowned_blocks > 0:
+            self.logger.error(f"Found {unowned_blocks} unowned blocks")
+            return False
+        
+        # Check that block intervals are contiguous and non-overlapping
+        for rank in range(block_intervals.shape[0]):
+            start = int(block_intervals[rank, 0].item())
+            end = int(block_intervals[rank, 1].item())
+            
+            if start < 0 or end > total_blocks or start >= end:
+                self.logger.error(f"Invalid block interval for rank {rank}: [{start}, {end})")
+                return False
+            
+            # Check that all blocks in this interval are owned by this rank
+            for block_idx in range(start, end):
+                if int(owners[block_idx].item()) != rank:
+                    self.logger.error(f"Block {block_idx} not owned by rank {rank}")
+                    return False
+        
+        self.logger.debug("KV cache consistency validation passed")
+        return True
+    
+    def cleanup_kv_cache(self, block_intervals: torch.Tensor, total_blocks: int) -> None:
+        """
+        Clean up KV cache for blocks not owned by current rank.
+        
+        Args:
+            block_intervals: Current block intervals [world_size, 2]
+            total_blocks: Total number of blocks
+        """
+        rank = dist.get_rank()
+        owners = self.compute_block_owners(block_intervals, total_blocks)
+        
+        cleaned_blocks = 0
+        for block_idx in range(total_blocks):
+            if int(owners[block_idx].item()) != rank:
+                # Clear KV cache for blocks not owned by this rank
+                if hasattr(self.pipeline, 'kv_cache1') and block_idx < len(self.pipeline.kv_cache1):
+                    if 'k' in self.pipeline.kv_cache1[block_idx]:
+                        self.pipeline.kv_cache1[block_idx]['k'].zero_()
+                    if 'v' in self.pipeline.kv_cache1[block_idx]:
+                        self.pipeline.kv_cache1[block_idx]['v'].zero_()
+                    cleaned_blocks += 1
+        
+        self.logger.info(f"Cleaned up KV cache for {cleaned_blocks} blocks not owned by rank {rank}")

streamv2v/communication/model_data_transfer.py

@@ -0,0 +1,273 @@
+"""
+Model data transfer abstraction layer.
+
+This module provides high-level interfaces for transferring model data
+between distributed ranks during inference.
+"""
+
+import torch
+from typing import List, Tuple, Optional, Any
+import logging
+from .distributed_communicator import DistributedCommunicator
+from .buffer_manager import BufferManager
+from .kv_cache_manager import KVCacheManager
+from .data_containers import LatentData, CommunicationConfig, PerformanceMetrics
+from .utils import CommunicationTimer
+
+
+class ModelDataTransfer:
+    """
+    High-level interface for model data transfer operations.
+    
+    This class encapsulates all model-related data transfer operations,
+    providing a clean interface for sending and receiving latent data,
+    KV caches, and other model state between ranks.
+    """
+    
+    def __init__(self, communicator: DistributedCommunicator, 
+                 buffer_manager: BufferManager,
+                 kv_cache_manager: Optional[KVCacheManager] = None,
+                 config: Optional[CommunicationConfig] = None):
+        """
+        Initialize the model data transfer manager.
+        
+        Args:
+            communicator: Distributed communicator instance
+            buffer_manager: Buffer manager for tensor allocation
+            kv_cache_manager: KV cache manager (optional)
+            config: Communication configuration
+        """
+        self.comm = communicator
+        self.buffer_mgr = buffer_manager
+        self.kv_cache_mgr = kv_cache_manager
+        self.config = config or CommunicationConfig()
+        
+        # Setup logging
+        self.logger = logging.getLogger(f"ModelDataTransfer_rank_{communicator.rank}")
+        self.logger.propagate = False
+        if not self.logger.handlers:
+            handler = logging.StreamHandler()
+            formatter = logging.Formatter(
+                f'[Rank {communicator.rank}] %(asctime)s - %(name)s - %(levelname)s - %(message)s'
+            )
+            handler.setFormatter(formatter)
+            self.logger.addHandler(handler)
+        
+        # Performance tracking
+        self.transfer_count = 0
+        self.total_transfer_time = 0.0
+    
+    def send_latent_data_async(self, chunk_idx: int, latents: torch.Tensor,
+                             original_latents: torch.Tensor, patched_x_shape: torch.Tensor,
+                             current_start: torch.Tensor, current_end: torch.Tensor,
+                             current_step: int) -> List[Any]:
+        """
+        Asynchronously send latent data to the next rank.
+        
+        Args:
+            chunk_idx: Chunk index
+            latents: Latent tensor
+            original_latents: Original latent tensor
+            patched_x_shape: Patched x shape tensor
+            current_start: Current start indices
+            current_end: Current end indices
+            current_step: Current step
+            
+        Returns:
+            List of work objects for all send operations
+        """
+        with CommunicationTimer(f"send_latent_data_async chunk_{chunk_idx}", self.logger):
+            work_objects = self.comm.send_latent_data_async(
+                chunk_idx=chunk_idx,
+                latents=latents,
+                original_latents=original_latents,
+                patched_x_shape=patched_x_shape,
+                current_start=current_start,
+                current_end=current_end,
+                current_step=current_step
+            )
+        
+        self.transfer_count += 1
+        self.logger.debug(f"Sent latent data for chunk {chunk_idx}")
+        return work_objects
+    
+    def receive_latent_data_async(self, num_steps: int) -> LatentData:
+        """
+        Asynchronously receive latent data from the previous rank.
+        
+        Args:
+            num_steps: Number of denoising steps
+            
+        Returns:
+            LatentData object containing all received data
+        """
+        with CommunicationTimer("receive_latent_data_async", self.logger):
+            chunk_idx, latents, original_latents, current_start, current_end, current_step, patched_x_shape = \
+                self.comm.recv_latent_data_async(num_steps, self.buffer_mgr)
+        
+        self.transfer_count += 1
+        self.logger.debug(f"Received latent data for chunk {chunk_idx}")
+        
+        return LatentData(
+            chunk_idx=chunk_idx,
+            latents=latents,
+            original_latents=original_latents,
+            current_start=current_start,
+            current_end=current_end,
+            current_step=current_step,
+            patched_x_shape=patched_x_shape
+        )
+
+    def release_latent_data(self, latent_data: Optional[LatentData]) -> None:
+        """Return received latent-data buffers to the buffer pool."""
+        if latent_data is None or self.buffer_mgr is None:
+            return
+
+        self.buffer_mgr.return_buffer(latent_data.latents, "latent")
+        self.buffer_mgr.return_buffer(latent_data.original_latents, "origin")
+        self.buffer_mgr.return_buffer(latent_data.patched_x_shape, "misc")
+        self.buffer_mgr.return_buffer(latent_data.current_start, "misc")
+        self.buffer_mgr.return_buffer(latent_data.current_end, "misc")
+
+    def send_prompt_async(self, prompt: str, device: torch.device) -> List[Any]:
+        return self.comm.send_prompt_async(prompt, device)
+
+    def recv_prompt_async(self) -> str:
+        return self.comm.recv_prompt_async()
+    
+    def send_kv_cache_blocks(self, block_indices: List[int], donor_rank: int) -> None:
+        """
+        Send KV cache blocks to all ranks.
+        
+        Args:
+            block_indices: List of block indices to send
+            donor_rank: Rank that owns the KV cache data
+        """
+        if self.kv_cache_mgr is None:
+            raise RuntimeError("KV cache manager not initialized")
+        
+        with CommunicationTimer(f"send_kv_cache_blocks {len(block_indices)} blocks", self.logger):
+            self.kv_cache_mgr.broadcast_kv_blocks(block_indices, donor_rank)
+        
+        self.logger.debug(f"Sent KV cache blocks {block_indices} from rank {donor_rank}")
+    
+    def rebalance_kv_cache(self, old_intervals: torch.Tensor, 
+                          new_intervals: torch.Tensor, total_blocks: int) -> None:
+        """
+        Rebalance KV cache ownership based on new block intervals.
+        
+        Args:
+            old_intervals: Previous block intervals [world_size, 2]
+            new_intervals: New block intervals [world_size, 2]
+            total_blocks: Total number of blocks
+        """
+        if self.kv_cache_mgr is None:
+            raise RuntimeError("KV cache manager not initialized")
+        
+        with CommunicationTimer("rebalance_kv_cache", self.logger):
+            self.kv_cache_mgr.rebalance_kv_cache_by_diff(old_intervals, new_intervals, total_blocks)
+        
+        self.logger.info("Rebalanced KV cache ownership")
+    
+    def broadcast_tensor(self, tensor: torch.Tensor, src: int) -> None:
+        """
+        Broadcast a tensor from source to all ranks.
+        
+        Args:
+            tensor: Tensor to broadcast
+            src: Source rank
+        """
+        with CommunicationTimer(f"broadcast_tensor from rank {src}", self.logger):
+            self.comm.broadcast_tensor(tensor, src)
+        
+        self.logger.debug(f"Broadcasted tensor from rank {src}, shape: {tensor.shape}")
+    
+    def all_gather_tensors(self, tensor: torch.Tensor) -> List[torch.Tensor]:
+        """
+        Gather tensors from all ranks.
+        
+        Args:
+            tensor: Local tensor to gather
+            
+        Returns:
+            List of tensors from all ranks
+        """
+        with CommunicationTimer("all_gather_tensors", self.logger):
+            gather_list = self.comm.all_gather_tensors(tensor)
+        
+        self.logger.debug(f"Gathered tensors from all ranks, local shape: {tensor.shape}")
+        return gather_list
+    
+    def wait_for_outstanding(self, max_outstanding: Optional[int] = None) -> None:
+        """
+        Wait for outstanding operations to complete.
+        
+        Args:
+            max_outstanding: Maximum number of outstanding operations to keep
+        """
+        with CommunicationTimer("wait_for_outstanding", self.logger):
+            self.comm.wait_for_outstanding(max_outstanding)
+    
+    def barrier(self) -> None:
+        """Synchronize all ranks."""
+        with CommunicationTimer("barrier", self.logger):
+            self.comm.barrier()
+    
+    def get_performance_metrics(self) -> PerformanceMetrics:
+        """
+        Get performance metrics for data transfer operations.
+        
+        Returns:
+            PerformanceMetrics object containing timing information
+        """
+        # This is a simplified version - in practice, you'd want to track
+        # more detailed timing information
+        avg_transfer_time = self.total_transfer_time / max(1, self.transfer_count)
+        
+        return PerformanceMetrics(
+            dit_time=0.0,  # Would be filled by caller
+            total_time=0.0,  # Would be filled by caller
+            communication_time=avg_transfer_time,
+            buffer_allocation_time=0.0  # Would be tracked by buffer manager
+        )
+    
+    def get_statistics(self) -> dict:
+        """
+        Get transfer statistics.
+        
+        Returns:
+            Dictionary containing transfer statistics
+        """
+        return {
+            "transfer_count": self.transfer_count,
+            "total_transfer_time": self.total_transfer_time,
+            "avg_transfer_time": self.total_transfer_time / max(1, self.transfer_count),
+            "communicator_stats": self.comm.get_statistics(),
+            "buffer_manager_stats": self.buffer_mgr.get_statistics() if self.buffer_mgr else None
+        }
+    
+    def print_statistics(self) -> None:
+        """Print transfer statistics."""
+        stats = self.get_statistics()
+        self.logger.info("Model Data Transfer Statistics:")
+        for key, value in stats.items():
+            if key == "communicator_stats" or key == "buffer_manager_stats":
+                if value:
+                    self.logger.info(f"  {key}:")
+                    for sub_key, sub_value in value.items():
+                        self.logger.info(f"    {sub_key}: {sub_value}")
+            else:
+                self.logger.info(f"  {key}: {value}")
+    
+    def cleanup(self) -> None:
+        """Clean up resources."""
+        if self.buffer_mgr:
+            self.buffer_mgr.clear_buffers()
+        self.logger.info("Model data transfer cleanup completed")
+    
+    def __del__(self):
+        """Cleanup when the transfer manager is destroyed."""
+        try:
+            self.cleanup()
+        except Exception:
+            pass

streamv2v/communication/test_communication.py

@@ -0,0 +1,393 @@
+"""
+Test cases for the communication module.
+
+This module provides comprehensive tests for all communication abstractions.
+"""
+
+import unittest
+import logging
+import torch
+import torch.distributed as dist
+import tempfile
+import os
+import sys
+from unittest.mock import Mock, patch, MagicMock
+
+# Add the parent directory to the path to import our modules
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from communication.data_containers import LatentData, KVCacheData, CommunicationConfig, BlockInterval, PerformanceMetrics
+from communication.buffer_manager import BufferManager
+from communication.utils import CommunicationTags, setup_logging, compute_balanced_split
+from communication.distributed_communicator import DistributedCommunicator
+from communication.kv_cache_manager import KVCacheManager
+from communication.model_data_transfer import ModelDataTransfer
+
+
+class TestDataContainers(unittest.TestCase):
+    """Test cases for data container classes."""
+    
+    def setUp(self):
+        """Set up test fixtures."""
+        self.device = torch.device('cpu')
+        self.sample_latents = torch.randn(1, 4, 16, 16, device=self.device)
+        self.sample_original_latents = torch.randn(1, 4, 16, 16, 16, device=self.device)
+        self.sample_current_start = torch.tensor([0, 1, 2], device=self.device)
+        self.sample_current_end = torch.tensor([1, 2, 3], device=self.device)
+        self.sample_patched_x_shape = torch.tensor([1, 4, 16, 16, 16], device=self.device)
+    
+    def test_latent_data_creation(self):
+        """Test LatentData creation and validation."""
+        latent_data = LatentData(
+            chunk_idx=0,
+            latents=self.sample_latents,
+            original_latents=self.sample_original_latents,
+            current_start=self.sample_current_start,
+            current_end=self.sample_current_end,
+            current_step=100,
+            patched_x_shape=self.sample_patched_x_shape
+        )
+        
+        self.assertEqual(latent_data.chunk_idx, 0)
+        self.assertEqual(latent_data.current_step, 100)
+        self.assertTrue(torch.equal(latent_data.latents, self.sample_latents))
+    
+    def test_latent_data_validation(self):
+        """Test LatentData validation with invalid inputs."""
+        with self.assertRaises(TypeError):
+            LatentData(
+                chunk_idx=0,
+                latents="invalid",  # Should be torch.Tensor
+                original_latents=self.sample_original_latents,
+                current_start=self.sample_current_start,
+                current_end=self.sample_current_end,
+                current_step=100,
+                patched_x_shape=self.sample_patched_x_shape
+            )
+    
+    def test_communication_config(self):
+        """Test CommunicationConfig creation and validation."""
+        config = CommunicationConfig(
+            max_outstanding=5,
+            buffer_pool_size=20,
+            enable_buffer_reuse=True,
+            communication_timeout=60.0
+        )
+        
+        self.assertEqual(config.max_outstanding, 5)
+        self.assertEqual(config.buffer_pool_size, 20)
+        self.assertTrue(config.enable_buffer_reuse)
+        self.assertEqual(config.communication_timeout, 60.0)
+    
+    def test_communication_config_validation(self):
+        """Test CommunicationConfig validation with invalid inputs."""
+        with self.assertRaises(ValueError):
+            CommunicationConfig(max_outstanding=0)  # Should be at least 1
+        
+        with self.assertRaises(ValueError):
+            CommunicationConfig(buffer_pool_size=0)  # Should be at least 1
+        
+        with self.assertRaises(ValueError):
+            CommunicationConfig(communication_timeout=0)  # Should be positive
+    
+    def test_block_interval(self):
+        """Test BlockInterval creation and methods."""
+        interval = BlockInterval(start=0, end=10, rank=0)
+        
+        self.assertEqual(interval.start, 0)
+        self.assertEqual(interval.end, 10)
+        self.assertEqual(interval.rank, 0)
+        self.assertEqual(interval.size, 10)
+        self.assertTrue(interval.contains(5))
+        self.assertFalse(interval.contains(10))
+        self.assertFalse(interval.contains(-1))
+    
+    def test_block_interval_validation(self):
+        """Test BlockInterval validation with invalid inputs."""
+        with self.assertRaises(ValueError):
+            BlockInterval(start=-1, end=10, rank=0)  # Start should be non-negative
+        
+        with self.assertRaises(ValueError):
+            BlockInterval(start=10, end=5, rank=0)  # End should be greater than start
+        
+        with self.assertRaises(ValueError):
+            BlockInterval(start=0, end=10, rank=-1)  # Rank should be non-negative
+    
+    def test_performance_metrics(self):
+        """Test PerformanceMetrics creation and methods."""
+        metrics = PerformanceMetrics(
+            dit_time=1.0,
+            total_time=2.0,
+            communication_time=0.5,
+            buffer_allocation_time=0.1
+        )
+        
+        self.assertEqual(metrics.dit_time, 1.0)
+        self.assertEqual(metrics.total_time, 2.0)
+        self.assertEqual(metrics.communication_time, 0.5)
+        self.assertEqual(metrics.buffer_allocation_time, 0.1)
+        self.assertEqual(metrics.efficiency, 0.75)  # (2.0 - 0.5) / 2.0
+
+
+class TestBufferManager(unittest.TestCase):
+    """Test cases for BufferManager."""
+    
+    def setUp(self):
+        """Set up test fixtures."""
+        self.device = torch.device('cpu')
+        self.config = CommunicationConfig(buffer_pool_size=5)
+        self.buffer_manager = BufferManager(self.device, self.config)
+    
+    def test_buffer_allocation(self):
+        """Test buffer allocation and reuse."""
+        shape = (1, 4, 16, 16)
+        dtype = torch.float32
+        
+        # Allocate a buffer
+        buffer1 = self.buffer_manager.get_buffer(shape, dtype, "latent")
+        self.assertEqual(buffer1.shape, shape)
+        self.assertEqual(buffer1.dtype, dtype)
+        self.assertEqual(buffer1.device, self.device)
+        
+        # Return the buffer
+        self.buffer_manager.return_buffer(buffer1, "latent")
+        
+        # Get another buffer of the same shape - should reuse
+        buffer2 = self.buffer_manager.get_buffer(shape, dtype, "latent")
+        self.assertEqual(buffer2.shape, shape)
+        self.assertEqual(buffer2.dtype, dtype)
+    
+    def test_buffer_statistics(self):
+        """Test buffer manager statistics."""
+        shape = (1, 4, 16, 16)
+        dtype = torch.float32
+        
+        # Allocate and return some buffers
+        buffer1 = self.buffer_manager.get_buffer(shape, dtype, "latent")
+        self.buffer_manager.return_buffer(buffer1, "latent")
+        
+        buffer2 = self.buffer_manager.get_buffer(shape, dtype, "latent")
+        self.buffer_manager.return_buffer(buffer2, "latent")
+        
+        stats = self.buffer_manager.get_statistics()
+        self.assertEqual(stats['allocation_count'], 1)
+        self.assertEqual(stats['reuse_count'], 1)
+        self.assertGreater(stats['total_allocated_memory_bytes'], 0)
+    
+    def test_buffer_cleanup(self):
+        """Test buffer cleanup."""
+        shape = (1, 4, 16, 16)
+        dtype = torch.float32
+        
+        # Allocate and return some buffers
+        buffer1 = self.buffer_manager.get_buffer(shape, dtype, "latent")
+        self.buffer_manager.return_buffer(buffer1, "latent")
+        
+        # Clear buffers
+        self.buffer_manager.clear_buffers("latent")
+        
+        stats = self.buffer_manager.get_statistics()
+        self.assertEqual(stats['total_free_buffers'], 0)
+
+
+class TestUtils(unittest.TestCase):
+    """Test cases for utility functions."""
+    
+    def test_compute_balanced_split(self):
+        """Test the compute_balanced_split function."""
+        total_blocks = 30
+        rank_times = [1.0, 2.0, 1.5]  # Rank 1 is slower
+        dit_times = [0.8, 1.6, 1.2]
+        current_block_nums = [[0, 10], [10, 20], [20, 30]]
+        
+        new_block_nums = compute_balanced_split(total_blocks, rank_times, dit_times, current_block_nums)
+        
+        # Should have same number of ranks
+        self.assertEqual(len(new_block_nums), len(current_block_nums))
+        
+        # Should sum to total_blocks
+        total_allocated = sum(end - start for start, end in new_block_nums)
+        self.assertEqual(total_allocated, total_blocks)
+        
+        # Should be contiguous
+        for i in range(len(new_block_nums) - 1):
+            self.assertEqual(new_block_nums[i][1], new_block_nums[i + 1][0])
+    
+    def test_compute_balanced_split_edge_cases(self):
+        """Test compute_balanced_split with edge cases."""
+        # Empty input
+        result = compute_balanced_split(0, [], [], [])
+        self.assertEqual(result, [])
+        
+        # Single rank
+        result = compute_balanced_split(10, [1.0], [0.8], [[0, 10]])
+        self.assertEqual(result, [[0, 10]])
+        
+        # Invalid input lengths
+        result = compute_balanced_split(10, [1.0], [0.8], [[0, 10], [10, 20]])
+        self.assertEqual(result, [[0, 10], [10, 20]])  # Should return original
+
+
+class TestDistributedCommunicator(unittest.TestCase):
+    """Test cases for DistributedCommunicator."""
+    
+    def setUp(self):
+        """Set up test fixtures."""
+        self.device = torch.device('cpu')
+        self.config = CommunicationConfig()
+        
+        # Mock distributed environment
+        with patch('torch.distributed.is_initialized', return_value=True):
+            self.communicator = DistributedCommunicator(0, 2, self.device, self.config)
+    
+    def test_communicator_initialization(self):
+        """Test communicator initialization."""
+        self.assertEqual(self.communicator.rank, 0)
+        self.assertEqual(self.communicator.world_size, 2)
+        self.assertEqual(self.communicator.device, self.device)
+    
+    def test_communicator_initialization_without_distributed(self):
+        """Test communicator initialization without distributed."""
+        with patch('torch.distributed.is_initialized', return_value=False):
+            with self.assertRaises(RuntimeError):
+                DistributedCommunicator(0, 2, self.device, self.config)
+    
+    def test_create_header(self):
+        """Test header creation and parsing."""
+        chunk_idx = 5
+        shape = (1, 4, 16, 16)
+        
+        header = self.communicator._create_header(chunk_idx, shape)
+        self.assertEqual(header.shape, (5,))  # chunk_idx + 4 shape dimensions
+        self.assertEqual(header.dtype, torch.int64)
+        
+        parsed_chunk_idx, parsed_shape = self.communicator._parse_header(header)
+        self.assertEqual(parsed_chunk_idx, chunk_idx)
+        self.assertEqual(parsed_shape, shape)
+    
+    def test_communicator_statistics(self):
+        """Test communicator statistics."""
+        stats = self.communicator.get_statistics()
+        
+        self.assertEqual(stats['rank'], 0)
+        self.assertEqual(stats['world_size'], 2)
+        self.assertEqual(stats['outstanding_operations'], 0)
+        self.assertEqual(stats['max_outstanding'], 1)
+
+
+class TestKVCacheManager(unittest.TestCase):
+    """Test cases for KVCacheManager."""
+    
+    def setUp(self):
+        """Set up test fixtures."""
+        self.device = torch.device('cpu')
+        
+        # Mock pipeline with KV cache
+        self.mock_pipeline = Mock()
+        self.mock_pipeline.frame_seq_length = 16
+        self.mock_pipeline.denoising_step_list = [700, 500, 0]
+        self.mock_pipeline.kv_cache1 = [
+            {
+                'k': torch.randn(1, 8, 16, 64, device=self.device),
+                'v': torch.randn(1, 8, 16, 64, device=self.device),
+                'global_end_index': torch.tensor([16], device=self.device),
+                'local_end_index': torch.tensor([16], device=self.device)
+            }
+            for _ in range(30)
+        ]
+        
+        self.kv_cache_manager = KVCacheManager(self.mock_pipeline, self.device)
+    
+    def test_compute_block_owners(self):
+        """Test block owner computation."""
+        block_intervals = torch.tensor([[0, 10], [10, 20], [20, 30]], device=self.device)
+        total_blocks = 30
+        
+        owners = self.kv_cache_manager.compute_block_owners(block_intervals, total_blocks)
+        
+        self.assertEqual(owners.shape, (30,))
+        self.assertTrue(torch.all(owners[:10] == 0))
+        self.assertTrue(torch.all(owners[10:20] == 1))
+        self.assertTrue(torch.all(owners[20:30] == 2))
+    
+    def test_kv_cache_statistics(self):
+        """Test KV cache statistics."""
+        block_intervals = torch.tensor([[0, 10], [10, 20], [20, 30]], device=self.device)
+        total_blocks = 30
+        
+        stats = self.kv_cache_manager.get_kv_cache_statistics(block_intervals, total_blocks)
+        
+        self.assertEqual(stats['total_blocks'], 30)
+        self.assertEqual(stats['block_counts'][0], 10)
+        self.assertEqual(stats['block_counts'][1], 10)
+        self.assertEqual(stats['block_counts'][2], 10)
+        self.assertGreater(stats['memory_per_block_bytes'], 0)
+    
+    def test_validate_kv_cache_consistency(self):
+        """Test KV cache consistency validation."""
+        block_intervals = torch.tensor([[0, 10], [10, 20], [20, 30]], device=self.device)
+        total_blocks = 30
+        
+        is_consistent = self.kv_cache_manager.validate_kv_cache_consistency(block_intervals, total_blocks)
+        self.assertTrue(is_consistent)
+        
+        # Test with invalid intervals
+        invalid_intervals = torch.tensor([[0, 10], [10, 20], [20, 25]], device=self.device)  # Missing blocks
+        is_consistent = self.kv_cache_manager.validate_kv_cache_consistency(invalid_intervals, total_blocks)
+        self.assertFalse(is_consistent)
+
+
+class TestModelDataTransfer(unittest.TestCase):
+    """Test cases for ModelDataTransfer."""
+    
+    def setUp(self):
+        """Set up test fixtures."""
+        self.device = torch.device('cpu')
+        self.config = CommunicationConfig()
+        
+        # Mock components
+        with patch('torch.distributed.is_initialized', return_value=True):
+            self.communicator = DistributedCommunicator(0, 2, self.device, self.config)
+        
+        self.buffer_manager = BufferManager(self.device, self.config)
+        self.mock_pipeline = Mock()
+        self.mock_pipeline.frame_seq_length = 16
+        self.mock_pipeline.denoising_step_list = [700, 500, 0]
+        self.mock_pipeline.kv_cache1 = []
+        self.kv_cache_manager = KVCacheManager(self.mock_pipeline, self.device)
+        
+        self.data_transfer = ModelDataTransfer(
+            self.communicator,
+            self.buffer_manager,
+            self.kv_cache_manager,
+            self.config
+        )
+    
+    def test_data_transfer_initialization(self):
+        """Test data transfer initialization."""
+        self.assertEqual(self.data_transfer.comm, self.communicator)
+        self.assertEqual(self.data_transfer.buffer_mgr, self.buffer_manager)
+        self.assertEqual(self.data_transfer.kv_cache_mgr, self.kv_cache_manager)
+        self.assertEqual(self.data_transfer.transfer_count, 0)
+    
+    def test_data_transfer_statistics(self):
+        """Test data transfer statistics."""
+        stats = self.data_transfer.get_statistics()
+        
+        self.assertEqual(stats['transfer_count'], 0)
+        self.assertEqual(stats['total_transfer_time'], 0.0)
+        self.assertIsNotNone(stats['communicator_stats'])
+        self.assertIsNotNone(stats['buffer_manager_stats'])
+    
+    def test_cleanup(self):
+        """Test data transfer cleanup."""
+        # Should not raise any exceptions
+        self.data_transfer.cleanup()
+
+
+if __name__ == '__main__':
+    # Set up logging for tests
+    logging.basicConfig(level=logging.INFO)
+    
+    # Run tests
+    unittest.main(verbosity=2)

streamv2v/communication/utils.py

@@ -0,0 +1,288 @@
+"""
+Utility functions and constants for communication operations.
+
+This module provides utility functions and constants used across the communication module.
+"""
+
+import torch
+import torch.distributed as dist
+from typing import List, Tuple, Optional
+import time
+import logging
+
+# Communication tags for different types of data
+class CommunicationTags:
+    """Constants for communication tags."""
+    LATENT_HDR = 11001
+    LATENT_PAY = 11002
+    START_END_STEP = 11003
+    PATCHED_X_SHAPE = 11004
+    LATENT_ORIGIN_HDR = 11005
+    LATENT_ORIGIN_PAY = 11006
+    KV_CACHE_K = 11007
+    KV_CACHE_V = 11008
+    KV_CACHE_GLOBAL_END = 11009
+    KV_CACHE_LOCAL_END = 11010
+    BLOCK_INTERVALS = 11011
+    PERFORMANCE_METRICS = 11012
+    UPDATED_PROMPT_LENGTH = 11013
+    UPDATED_PROMPT = 11014
+
+
+def init_distributed():
+    """
+    Initialize distributed communication.
+    
+    This function initializes the distributed process group if not already initialized.
+    """
+    if not dist.is_initialized():
+        backend = "nccl"
+        dist.init_process_group(backend=backend)
+
+
+def get_rank_info() -> Tuple[int, int]:
+    """
+    Get current rank and world size.
+    
+    Returns:
+        Tuple of (rank, world_size)
+    """
+    if not dist.is_initialized():
+        raise RuntimeError("Distributed not initialized")
+    return dist.get_rank(), dist.get_world_size()
+
+
+def get_next_rank(rank: int, world_size: int) -> int:
+    """
+    Get the next rank in the ring topology.
+    
+    Args:
+        rank: Current rank
+        world_size: Total number of ranks
+        
+    Returns:
+        Next rank in the ring
+    """
+    return (rank + 1) % world_size
+
+
+def get_prev_rank(rank: int, world_size: int) -> int:
+    """
+    Get the previous rank in the ring topology.
+    
+    Args:
+        rank: Current rank
+        world_size: Total number of ranks
+        
+    Returns:
+        Previous rank in the ring
+    """
+    return (rank - 1) % world_size
+
+
+def create_tensor_header(shape: Tuple[int, ...], dtype: torch.dtype, 
+                        chunk_idx: int, device: torch.device) -> torch.Tensor:
+    """
+    Create a header tensor for communication.
+    
+    Args:
+        shape: Shape of the tensor to be sent
+        dtype: Data type of the tensor
+        chunk_idx: Chunk index
+        device: Device where the header will be created
+        
+    Returns:
+        Header tensor containing metadata
+    """
+    header_data = [chunk_idx] + list(shape)
+    return torch.tensor(header_data, dtype=torch.int64, device=device)
+
+
+def parse_tensor_header(header: torch.Tensor) -> Tuple[int, Tuple[int, ...]]:
+    """
+    Parse a header tensor to extract metadata.
+    
+    Args:
+        header: Header tensor
+        
+    Returns:
+        Tuple of (chunk_idx, shape)
+    """
+    header_list = header.tolist()
+    chunk_idx = int(header_list[0])
+    shape = tuple(int(x) for x in header_list[1:])
+    return chunk_idx, shape
+
+
+def validate_tensor_for_communication(tensor: torch.Tensor, 
+                                    expected_device: torch.device,
+                                    expected_dtype: torch.dtype) -> None:
+    """
+    Validate tensor properties for communication.
+    
+    Args:
+        tensor: Tensor to validate
+        expected_device: Expected device
+        expected_dtype: Expected data type
+        
+    Raises:
+        ValueError: If tensor properties don't match expectations
+    """
+    if not isinstance(tensor, torch.Tensor):
+        raise ValueError("Input must be a torch.Tensor")
+    
+    if tensor.device != expected_device:
+        raise ValueError(f"Tensor device {tensor.device} doesn't match expected {expected_device}")
+    
+    if tensor.dtype != expected_dtype:
+        raise ValueError(f"Tensor dtype {tensor.dtype} doesn't match expected {expected_dtype}")
+
+
+def compute_balanced_split(total_blocks: int, rank_times: List[float], 
+                          dit_times: List[float], 
+                          current_block_nums: List[List[int]]) -> List[List[int]]:
+    """
+    Compute new block splits for all ranks to balance total rank times.
+    
+    This function is moved from the original file to provide better organization.
+    
+    Args:
+        total_blocks: Total number of DiT blocks
+        rank_times: List of total iteration times for each rank [t_rank0, t_rank1, ..., t_rankN] (DiT + VAE time)
+        dit_times: List of pure DiT inference times for each rank [dit_rank0, dit_rank1, ..., dit_rankN] (DiT time only)
+        current_block_nums: List of current block_num format for each rank [[rank0_blocks], [rank1_blocks], ...]
+        
+    Returns:
+        List of new block_num format for each rank, matching the original format:
+        - For world_size == 2: [[end_idx_rank0], [start_idx_rank1]]
+        - For world_size > 2: [[end_idx_rank0], [start1, end1], [start2, end2], ..., [start_idx_last]]
+        Note: Numbers are shared across ranks (rank0_end = rank1_start, rank1_end = rank2_start, etc.)
+    """
+    num_ranks = len(rank_times)
+    if num_ranks == 0 or num_ranks != len(current_block_nums) or num_ranks != len(dit_times):
+        return current_block_nums
+    
+    # Edge case: if we have more ranks than blocks, we can't guarantee 1 block per rank
+    if num_ranks > total_blocks:
+        # Fall back to original behavior for this edge case
+        return current_block_nums
+
+    # Step 1: Calculate total DiT time and per-block DiT time
+    total_dit_time = sum(dit_times)
+    dit_time_per_block = total_dit_time / total_blocks
+    
+    # Step 2: Calculate average rank time
+    avg_rank_time = sum(rank_times) / num_ranks
+    
+    # Step 3: Extract current block counts from current_block_nums (all ranks use [start, end) now)
+    current_block_counts = []
+    for block_num in current_block_nums:
+        # block_num: [start, end) exclusive end
+        start_idx, end_idx = int(block_num[0]), int(block_num[1])
+        current_block_counts.append(max(0, end_idx - start_idx))
+    
+    # Step 4: Calculate target block counts based on time differences
+    target_blocks = []
+    for i in range(num_ranks):
+        time_diff = avg_rank_time - rank_times[i]  # positive = needs more time, negative = needs less time
+        block_adjustment = time_diff / dit_time_per_block  # convert time difference to block count
+        target_count = current_block_counts[i] + block_adjustment
+        # Ensure each rank gets at least 1 block (minimum allocation)
+        target_count = max(1, int(round(target_count)))
+        target_blocks.append(target_count)
+    
+    # Step 5: Adjust to ensure total blocks sum to total_blocks while maintaining minimum 1 block per rank
+    current_total = sum(target_blocks)
+    if current_total != total_blocks:
+        diff = total_blocks - current_total
+        # When adding, give to ranks with smallest counts first; when removing, take from largest counts first
+        if diff > 0:
+            order = sorted(range(num_ranks), key=lambda i: (target_blocks[i], i))
+        else:
+            order = sorted(range(num_ranks), key=lambda i: (target_blocks[i], i), reverse=True)
+        i = 0
+        while diff != 0 and num_ranks > 0:
+            idx = order[i % num_ranks]
+            if diff > 0:
+                target_blocks[idx] += 1
+                diff -= 1
+            else:
+                # Only remove blocks if rank has more than 1 block (maintain minimum allocation)
+                if target_blocks[idx] > 1:
+                    target_blocks[idx] -= 1
+                    diff += 1
+            i += 1
+    
+    # Step 6: Convert target block counts to contiguous [start, end) intervals from 0 to total_blocks
+    new_block_nums = []
+    running_start = 0
+    for i in range(num_ranks):
+        block_count = int(target_blocks[i])
+        start_idx = running_start
+        end_idx = start_idx + block_count
+        # Guard (should not trigger if sums are correct)
+        if end_idx > total_blocks:
+            end_idx = total_blocks
+        new_block_nums.append([start_idx, end_idx])
+        running_start = end_idx
+    
+    return new_block_nums
+
+
+def setup_logging(rank: int, log_level: int = logging.INFO) -> logging.Logger:
+    """
+    Setup logging for the current rank.
+    
+    Args:
+        rank: Current rank
+        log_level: Logging level
+        
+    Returns:
+        Configured logger
+    """
+    logger = logging.getLogger(f"rank_{rank}")
+    logger.setLevel(log_level)
+    # Prevent messages from propagating to the root logger (avoid double prints)
+    logger.propagate = False
+    
+    if not logger.handlers:
+        handler = logging.StreamHandler()
+        formatter = logging.Formatter(
+            f'[Rank {rank}] %(asctime)s - %(name)s - %(levelname)s - %(message)s'
+        )
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)
+    
+    return logger
+
+
+class CommunicationTimer:
+    """
+    Timer for measuring communication performance.
+    
+    This class provides context manager functionality for timing communication operations.
+    """
+    
+    def __init__(self, operation_name: str, logger: Optional[logging.Logger] = None):
+        self.operation_name = operation_name
+        self.logger = logger
+        self.start_time = None
+        self.end_time = None
+    
+    def __enter__(self):
+        self.start_time = time.time()
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.end_time = time.time()
+        duration = self.end_time - self.start_time
+        
+        if self.logger:
+            self.logger.info(f"{self.operation_name} took {duration:.4f} seconds")
+    
+    @property
+    def duration(self) -> float:
+        """Get the duration of the timed operation."""
+        if self.start_time is None or self.end_time is None:
+            return 0.0
+        return self.end_time - self.start_time

streamv2v/configs/__init__.py

@@ -0,0 +1,2 @@
+"""Packaged default configs for the public Python API."""
+

streamv2v/configs/wan_causal_dmd_v2v.yaml

@@ -0,0 +1,57 @@
+model_name: wan
+generator_name: causal_wan
+generator_ckpt: "ckpts/autoregressive_checkpoint/model.pt"
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+text_encoder_fsdp_wrap_strategy: size
+generator_grad:
+  model: true
+real_score_grad:
+  model: false
+fake_score_grad:
+  model: true
+denoising_step_list:
+- 700
+- 500
+- 400
+- 200
+- 0
+num_train_timestep: 1000
+timestep_shift: 8.0
+real_guidance_scale: 3.5
+generator_task: causal_video
+real_task_type: bidirectional_video
+fake_task_type: bidirectional_video
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+wandb_host: WANDB_HOST
+wandb_key: WANDB_KEY
+wandb_entity: tyin
+wandb_project: causvid
+wandb_name: wan_causal_dmd
+sharding_strategy: hybrid_full
+lr: 2.0e-06
+beta1: 0.9
+beta2: 0.999
+data_path: mixkit_latents_lmdb
+batch_size: 1
+log_iters: 200
+negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+dfake_gen_update_ratio: 5
+image_or_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+output_path: /mnt/localssd/wan_causal_dmd
+distillation_loss: dmd
+gradient_checkpointing: true
+backward_simulation: false
+num_frame_per_block: 1
+num_kv_cache: 6
+num_sink_tokens: 3
+adapt_sink_threshold: 0.2
+warp_denoising_step: false

streamv2v/configs/wan_causal_dmd_v2v_fast.yaml

@@ -0,0 +1,57 @@
+model_name: wan
+generator_name: causal_wan
+generator_ckpt: "ckpts/autoregressive_checkpoint/model.pt"
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+text_encoder_fsdp_wrap_strategy: size
+generator_grad:
+  model: true
+real_score_grad:
+  model: false
+fake_score_grad:
+  model: true
+denoising_step_list:
+- 700
+- 500
+- 400
+- 200
+- 0
+num_train_timestep: 1000
+timestep_shift: 8.0
+real_guidance_scale: 3.5
+generator_task: causal_video
+real_task_type: bidirectional_video
+fake_task_type: bidirectional_video
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+wandb_host: WANDB_HOST
+wandb_key: WANDB_KEY
+wandb_entity: tyin
+wandb_project: causvid
+wandb_name: wan_causal_dmd
+sharding_strategy: hybrid_full
+lr: 2.0e-06
+beta1: 0.9
+beta2: 0.999
+data_path: mixkit_latents_lmdb
+batch_size: 1
+log_iters: 200
+negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+dfake_gen_update_ratio: 5
+image_or_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+output_path: /mnt/localssd/wan_causal_dmd
+distillation_loss: dmd
+gradient_checkpointing: true
+backward_simulation: false
+num_frame_per_block: 1
+num_kv_cache: 5
+num_sink_tokens: 2
+adapt_sink_threshold: -1
+warp_denoising_step: false

streamv2v/inference.py

@@ -0,0 +1,525 @@
+"""
+Single GPU Inference Pipeline - Refactored from inference_pipe.py
+
+This file extracts core logic from multi-GPU inference code to implement a complete 
+inference pipeline on a single GPU:
+1. VAE encode input video
+2. DiT inference (using input mode, processing all 30 blocks)
+3. VAE decode output video
+"""
+
+from models.wan.causal_stream_inference import CausalStreamInferencePipeline
+from models.util import set_seed
+from diffusers.utils import export_to_video
+from models.data import TextDataset
+import argparse
+from dataclasses import dataclass
+import torch
+import os
+import time
+import numpy as np
+import logging
+from typing import List
+
+try:
+    from streamv2v.inference_common import (
+        load_generator_state_dict,
+        load_mp4_as_tensor,
+        merge_cli_config,
+    )
+except ModuleNotFoundError:
+    from inference_common import (
+        load_generator_state_dict,
+        load_mp4_as_tensor,
+        merge_cli_config,
+    )
+
+LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class SingleGPUStreamSession:
+    prompt: str
+    noise_scale: float
+    init_noise_scale: float
+    chunk_size: int
+    current_start: int
+    current_end: int
+    last_image: torch.Tensor
+    processed: int = 0
+
+def compute_noise_scale_and_step(input_video_original: torch.Tensor, end_idx: int, chunk_size: int, noise_scale: float, init_noise_scale: float):
+    """Compute adaptive noise scale and current step based on video content."""
+    l2_dist=(input_video_original[:,:,end_idx-chunk_size:end_idx]-input_video_original[:,:,end_idx-chunk_size-1:end_idx-1])**2
+    l2_dist = (torch.sqrt(l2_dist.mean(dim=(0,1,3,4))).max()/0.2).clamp(0,1)
+    new_noise_scale = (init_noise_scale-0.1*l2_dist.item())*0.9+noise_scale*0.1
+    current_step = int(1000*new_noise_scale)-100
+    return new_noise_scale, current_step
+
+class SingleGPUInferencePipeline:
+    """
+    Single GPU Inference Pipeline Manager
+    
+    This class encapsulates the complete inference logic on a single GPU, 
+    including encoding, inference, and decoding.
+    """
+    
+    def __init__(self, config, device: torch.device):
+        """
+        Initialize the single GPU inference pipeline manager.
+        
+        Args:
+            config: Configuration object
+            device: GPU device
+        """
+        self.config = config
+        self.device = device
+        
+        # Setup logging
+        self.logger = logging.getLogger("SingleGPUInference")
+        self.logger.setLevel(logging.INFO)
+        # Prevent messages from propagating to the root logger (avoid double prints)
+        self.logger.propagate = False
+        
+        if not self.logger.handlers:
+            handler = logging.StreamHandler()
+            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+            handler.setFormatter(formatter)
+            self.logger.addHandler(handler)
+        
+        # Initialize pipeline
+        self.pipeline = CausalStreamInferencePipeline(config, device=str(device))
+        self.pipeline.to(device=str(device), dtype=torch.bfloat16)
+        
+        # Performance tracking
+        self.t_dit = 100.0
+        self.t_total = 100.0
+        self.processed = 0
+        self.processed_offset = 3
+        self.base_chunk_size = 4
+        self.t_refresh = 50
+
+        self.t2v = config.t2v
+        self.profile = bool(config.get("profile", False))
+        self.encode_fps_list: list[float] = []
+        self.decode_fps_list: list[float] = []
+
+        
+        self.logger.info("Single GPU inference pipeline manager initialized")
+    
+    def load_model(self, checkpoint_folder: str):
+        """Load the model from checkpoint."""
+        ckpt_path, state_dict = load_generator_state_dict(checkpoint_folder)
+        self.logger.info(f"Loading checkpoint from {ckpt_path}")
+
+        # Load into the pipeline generator
+        try:
+            self.pipeline.generator.load_state_dict(state_dict, strict=True)
+        except RuntimeError as e:
+            # Try non-strict load as a fallback and report
+            self.logger.warning(f"Strict load_state_dict failed: {e}; retrying with strict=False")
+            self.pipeline.generator.load_state_dict(state_dict, strict=False)
+    
+    def prepare_pipeline(self, text_prompts: list, noise: torch.Tensor, current_start: int, current_end: int):
+        """Prepare the pipeline for inference."""
+        # Use the original prepare method which now handles distributed environment gracefully
+        denoised_pred = self.pipeline.prepare(
+            text_prompts=text_prompts,
+            device=self.device,
+            dtype=torch.bfloat16,
+            block_mode='input',
+            noise=noise,
+            current_start=current_start,
+            current_end=current_end
+        )
+        return denoised_pred
+
+    def _sync_for_timing(self):
+        if self.profile:
+            torch.cuda.synchronize()
+
+    def _record_stage_fps(self, values: list[float], num_frames: int, elapsed: float) -> None:
+        if self.profile and elapsed > 0 and num_frames > 0:
+            values.append(num_frames / elapsed)
+
+    def _timed_stream_encode(self, images: torch.Tensor) -> torch.Tensor:
+        self._sync_for_timing()
+        start_time = time.time()
+        latents = self.pipeline.vae.stream_encode(images)
+        self._sync_for_timing()
+        self._record_stage_fps(self.encode_fps_list, int(images.shape[2]), time.time() - start_time)
+        return latents
+
+    def _timed_stream_decode(self, denoised_pred: torch.Tensor) -> torch.Tensor:
+        self._sync_for_timing()
+        start_time = time.time()
+        video = self.pipeline.vae.stream_decode_to_pixel(denoised_pred)
+        self._sync_for_timing()
+        self._record_stage_fps(self.decode_fps_list, int(video.shape[1]), time.time() - start_time)
+        return video
+
+    def reset_stream_state(self, reset_vae_flags: bool = True) -> None:
+        """Reset cached model state before starting a new streaming session."""
+        if reset_vae_flags:
+            self.pipeline.vae.model.first_encode = True
+            self.pipeline.vae.model.first_decode = True
+
+        self.pipeline.kv_cache1 = None
+        self.pipeline.crossattn_cache = None
+        self.pipeline.block_x = None
+        self.pipeline.hidden_states = None
+        self.processed = 0
+
+    def _encode_noisy_latents(self, images: torch.Tensor, noise_scale: float) -> torch.Tensor:
+        latents = self._timed_stream_encode(images)
+        latents = latents.transpose(2, 1).contiguous().to(dtype=torch.bfloat16)
+        noise = torch.randn_like(latents)
+        return noise * noise_scale + latents * (1 - noise_scale)
+
+    def _decode_video_array(self, denoised_pred: torch.Tensor, last_frame_only: bool = False) -> np.ndarray:
+        if last_frame_only:
+            denoised_pred = denoised_pred[[-1]]
+
+        video = self._timed_stream_decode(denoised_pred)
+        video = (video * 0.5 + 0.5).clamp(0, 1)
+        video = video[0].permute(0, 2, 3, 1).contiguous()
+        return video.detach().cpu().float().numpy()
+
+    def start_stream_session(self, prompt: str, images: torch.Tensor, noise_scale: float) -> tuple[SingleGPUStreamSession, np.ndarray]:
+        """Initialize a streaming session and return the first decoded frames."""
+        self.reset_stream_state(reset_vae_flags=True)
+
+        chunk_size = self.base_chunk_size * self.pipeline.num_frame_per_block
+        current_start = 0
+        current_end = self.pipeline.frame_seq_length * (1 + chunk_size // self.base_chunk_size)
+
+        noisy_latents = self._encode_noisy_latents(images, noise_scale)
+        denoised_pred = self.prepare_pipeline(
+            text_prompts=[prompt],
+            noise=noisy_latents,
+            current_start=current_start,
+            current_end=current_end,
+        )
+        initial_video = self._decode_video_array(denoised_pred, last_frame_only=False)
+
+        session = SingleGPUStreamSession(
+            prompt=prompt,
+            noise_scale=noise_scale,
+            init_noise_scale=noise_scale,
+            chunk_size=chunk_size,
+            current_start=current_end,
+            current_end=current_end + (chunk_size // self.base_chunk_size) * self.pipeline.frame_seq_length,
+            last_image=images[:, :, [-1]],
+            processed=0,
+        )
+        return session, initial_video
+
+    def run_stream_batch(self, session: SingleGPUStreamSession, images: torch.Tensor, queue_wait_time: float | None = None) -> List[np.ndarray]:
+        """Process one or more chunk-aligned frame groups for an active streaming session."""
+        num_frames = images.shape[2]
+        input_batch = num_frames // session.chunk_size
+        noise_scale, current_step = compute_noise_scale_and_step(
+            input_video_original=torch.cat([session.last_image, images], dim=2),
+            end_idx=num_frames + 1,
+            chunk_size=num_frames,
+            noise_scale=float(session.noise_scale),
+            init_noise_scale=float(session.init_noise_scale),
+        )
+        noisy_latents = self._encode_noisy_latents(images, noise_scale)
+
+        outputs: List[np.ndarray] = []
+        num_steps = len(self.pipeline.denoising_step_list)
+
+        for batch_idx in range(input_batch):
+            if session.current_start // self.pipeline.frame_seq_length >= self.t_refresh:
+                session.current_start = self.pipeline.kv_cache_length - self.pipeline.frame_seq_length
+                session.current_end = session.current_start + (session.chunk_size // self.base_chunk_size) * self.pipeline.frame_seq_length
+
+            denoised_pred = self.pipeline.inference_stream(
+                noise=noisy_latents[:, batch_idx].unsqueeze(1),
+                current_start=session.current_start,
+                current_end=session.current_end,
+                current_step=current_step,
+            )
+
+            session.processed += 1
+            self.processed = session.processed
+
+            if session.processed >= num_steps:
+                outputs.append(self._decode_video_array(denoised_pred, last_frame_only=True))
+
+            session.current_start = session.current_end
+            session.current_end += (session.chunk_size // self.base_chunk_size) * self.pipeline.frame_seq_length
+
+        session.last_image = images[:, :, [-1]]
+        session.noise_scale = noise_scale
+        return outputs
+    
+    def run_inference(
+        self, 
+        input_video_original: torch.Tensor, 
+        prompts: list, 
+        num_chunks: int, 
+        chunk_size: int, 
+        noise_scale: float, 
+        output_folder: str, 
+        fps: int, 
+        target_fps:int,  
+        num_steps: int,
+        ):
+        """
+        Run the complete single GPU inference pipeline.
+        
+        This method integrates the complete encoding, inference, and decoding pipeline.
+        """
+        self.logger.info("Starting single GPU inference pipeline")
+        
+        os.makedirs(output_folder, exist_ok=True)
+        results = {}
+        save_results = 0
+
+        fps_list = []
+        dit_fps_list = []
+        self.encode_fps_list = []
+        self.decode_fps_list = []
+        
+        # Initialize variables
+        start_idx = 0
+        if self.t2v:
+            end_idx = 1 + chunk_size - 4
+        else:
+            end_idx = 1 + chunk_size
+        current_start = 0
+        current_end = self.pipeline.frame_seq_length * (1+(end_idx-1)//4)
+        
+        self._sync_for_timing()
+        start_time = time.time()
+        
+        # Process first chunk (initialization)
+        if not self.t2v:
+            inp = input_video_original[:, :, start_idx:end_idx]
+            
+            # VAE encoding
+            latents = self._timed_stream_encode(inp)
+            latents = latents.transpose(2, 1).contiguous().to(dtype=torch.bfloat16)
+            
+            noise = torch.randn_like(latents)
+            noisy_latents = noise * noise_scale + latents * (1 - noise_scale)
+        else:
+            noisy_latents = torch.randn(1,self.pipeline.num_frame_per_block,16,self.pipeline.height,self.pipeline.width, device=self.device, dtype=torch.bfloat16)
+        
+            
+        # Prepare pipeline
+        denoised_pred = self.prepare_pipeline(
+            text_prompts=prompts,
+            noise=noisy_latents,
+            current_start=current_start,
+            current_end=current_end
+        )
+        
+        # Save first result - only start decoding after num_steps
+        video = self._timed_stream_decode(denoised_pred)
+        video = (video * 0.5 + 0.5).clamp(0, 1)
+        video = video[0].permute(0, 2, 3, 1).contiguous()
+        results[save_results] = video.cpu().float().numpy()
+        self.logger.info(
+            "Prepared initial chunk: start=%s, end=%s, start_idx=%s, save_results=%s, frames=%s",
+            current_start,
+            current_end,
+            start_idx,
+            save_results,
+            video.shape[0],
+        )
+        save_results += 1
+        
+        init_noise_scale = noise_scale
+
+        # Process remaining chunks
+        while self.processed < num_chunks + num_steps - 1:
+            # Update indices
+            start_idx = end_idx
+            end_idx = end_idx + chunk_size
+            current_start = current_end
+            current_end = current_end + (chunk_size // 4) * self.pipeline.frame_seq_length
+
+            if not self.t2v and end_idx <= input_video_original.shape[2]:
+                inp = input_video_original[:, :, start_idx:end_idx]
+                
+                noise_scale, current_step = compute_noise_scale_and_step(
+                    input_video_original, end_idx, chunk_size, noise_scale, init_noise_scale
+                )
+                
+                # VAE encoding
+                latents = self._timed_stream_encode(inp)
+                latents = latents.transpose(2, 1).contiguous().to(dtype=torch.bfloat16)
+                
+                noise = torch.randn_like(latents)
+                noisy_latents = noise * noise_scale + latents * (1 - noise_scale)
+            else:
+                noisy_latents = torch.randn(1,self.pipeline.num_frame_per_block,16,self.pipeline.height,self.pipeline.width, device=self.device, dtype=torch.bfloat16)
+                current_step = None # Use default steps
+
+            self._sync_for_timing()
+            dit_start_time = time.time()
+                
+            # DiT inference - using input mode to process all 30 blocks
+            denoised_pred = self.pipeline.inference_stream(
+                noise=noisy_latents,
+                current_start=current_start,
+                current_end=current_end,
+                current_step=current_step,
+            )
+
+            if self.processed > self.processed_offset:
+                self._sync_for_timing()
+                if self.profile:
+                    dit_fps_list.append(chunk_size / (time.time() - dit_start_time))
+            
+            self.processed += 1
+            
+            # VAE decoding - only start decoding after num_steps
+            if self.processed >= num_steps:
+                if self.t2v and self.processed == num_steps:
+                    continue
+                video = self._timed_stream_decode(denoised_pred[[-1]])
+                video = (video * 0.5 + 0.5).clamp(0, 1)
+                video = video[0].permute(0, 2, 3, 1).contiguous()
+                results[save_results] = video.cpu().float().numpy()
+                save_results += 1
+            
+                # Update timing
+                if self.profile:
+                    self._sync_for_timing()
+                    end_time = time.time()
+                    t = end_time - start_time
+                    fps_test = chunk_size / t
+                    fps_list.append(fps_test)
+                    self.logger.info(f"Processed {self.processed}, time: {t:.4f} s, FPS: {fps_test:.4f}")
+                else:
+                    fps_test = None
+
+                if self.processed == num_steps + self.processed_offset and target_fps is not None and fps_test is not None and fps_test < target_fps:
+                    max_chunk_size = (self.pipeline.num_kv_cache - self.pipeline.num_sink_tokens - 1) * self.base_chunk_size
+                    num_chunks=(num_chunks-self.processed-num_steps+1)//(max_chunk_size//chunk_size)+self.processed-num_steps+1
+                    self.pipeline.hidden_states=self.pipeline.hidden_states.repeat(1,max_chunk_size//chunk_size,1,1,1)
+                    chunk_size = max_chunk_size
+                    self.logger.info(f"Adjust chunk size to {chunk_size}")
+
+                if self.profile:
+                    start_time = end_time
+        
+
+        # Save final video
+        video_list = [results[i] for i in range(num_chunks)]
+        video = np.concatenate(video_list, axis=0)
+        if self.profile and fps_list:
+            fps_avg = np.mean(np.array(fps_list))
+            dit_avg = np.mean(np.array(dit_fps_list)) if dit_fps_list else 0.0
+            encode_avg = np.mean(np.array(self.encode_fps_list)) if self.encode_fps_list else 0.0
+            decode_avg = np.mean(np.array(self.decode_fps_list)) if self.decode_fps_list else 0.0
+            self.logger.info(f"VAE Encode Average FPS: {encode_avg:.4f}")
+            self.logger.info(f"DiT Average FPS: {dit_avg:.4f}")
+            self.logger.info(f"VAE Decode Average FPS: {decode_avg:.4f}")
+            self.logger.info(f"Video shape: {video.shape}, Average FPS: {fps_avg:.4f}")
+        else:
+            self.logger.info(f"Video shape: {video.shape}")
+        
+        output_path = os.path.join(output_folder, f"output_{0:03d}.mp4")
+        export_to_video(video, output_path, fps=fps)
+        self.logger.info(f"Video saved to: {output_path}")
+        
+        self.logger.info("Single GPU inference pipeline completed")
+
+
+def main():
+    """Main function for the single GPU inference pipeline."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_path", type=str, required=True, help="Configuration file path")
+    parser.add_argument("--checkpoint_folder", type=str, required=True, help="Checkpoint folder path")
+    parser.add_argument("--output_folder", type=str, required=True, help="Output folder path")
+    parser.add_argument("--prompt_file_path", type=str, required=True, help="Prompt file path")
+    parser.add_argument("--video_path", type=str, required=False, default=None, help="Input video path")
+    parser.add_argument("--noise_scale", type=float, default=0.8, help="Noise scale")
+    parser.add_argument("--height", type=int, default=480, help="Video height")
+    parser.add_argument("--width", type=int, default=832, help="Video width")
+    parser.add_argument("--fps", type=int, default=16, help="Output video fps")
+    parser.add_argument("--step", type=int, default=2, help="Step")
+    parser.add_argument("--seed", type=int, default=0, help="Random seed")
+    parser.add_argument("--gpu_id", type=int, default=None, help="CUDA device index for single-GPU inference")
+    parser.add_argument("--model_type", type=str, default="T2V-1.3B", help="Model type (e.g., T2V-1.3B)")
+    parser.add_argument("--num_frames", type=int, default=81, help="Video length (number of frames)")
+    parser.add_argument("--fixed_noise_scale", action="store_true", default=False)
+    parser.add_argument("--t2v", action="store_true", default=False)
+    parser.add_argument("--target_fps", type=int, required=False, default=None, help="Video length (number of frames)")
+    parser.add_argument("--profile", action="store_true", default=False, help="Enable synchronized throughput logging")
+    parser.add_argument("--use_taehv", action="store_true", default=False, help="Use the lightweight TAEHV VAE for encode/decode")
+    parser.add_argument("--use_tensorrt", "--use_taehv_tensorrt", dest="use_tensorrt", action="store_true", default=False, help="Enable available TensorRT acceleration paths")
+    parser.add_argument("--fast", action="store_true", default=False, help="Enable the fast path: --use_taehv --use_tensorrt")
+    args = parser.parse_args()
+    
+    torch.set_grad_enabled(False)
+    
+    # Auto-detect device
+    if torch.cuda.is_available():
+        if args.gpu_id is not None:
+            torch.cuda.set_device(args.gpu_id)
+            device = torch.device(f"cuda:{args.gpu_id}")
+        else:
+            device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+    
+    # Load configuration
+    config = merge_cli_config(args.config_path, args)
+
+    set_seed(args.seed)
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+    LOGGER.info("Denoising Step List: %s", list(config.denoising_step_list))
+    
+    # Load input video
+    if not args.t2v:
+        input_video_original = load_mp4_as_tensor(args.video_path, resize_hw=(args.height, args.width)).unsqueeze(0)
+        LOGGER.info("Input video tensor shape: %s", tuple(input_video_original.shape))
+        b, c, t, h, w = input_video_original.shape
+        if input_video_original.dtype != torch.bfloat16:
+            input_video_original = input_video_original.to(dtype=torch.bfloat16).to(device)
+    else:
+        input_video_original = None
+        t = args.num_frames
+    
+    # Calculate number of chunks
+    chunk_size = 4 * config.num_frame_per_block
+    num_chunks = (t - 1) // chunk_size
+
+    if args.t2v:
+        num_chunks+=1
+    # Initialize pipeline manager
+    pipeline_manager = SingleGPUInferencePipeline(config, device)
+    pipeline_manager.load_model(args.checkpoint_folder)
+    
+    # Load prompts
+    dataset = TextDataset(args.prompt_file_path)
+    prompts = [dataset[0]]
+    num_steps = len(pipeline_manager.pipeline.denoising_step_list)
+    
+    # Run inference
+    try:
+        pipeline_manager.run_inference(
+            input_video_original, 
+            prompts, 
+            num_chunks, 
+            chunk_size, 
+            args.noise_scale, 
+            args.output_folder, 
+            args.fps, 
+            args.target_fps, 
+            num_steps,
+        )
+    except Exception as e:
+        LOGGER.exception("Error occurred during inference: %s", e)
+        raise
+
+
+if __name__ == "__main__":
+    main()

streamv2v/inference_common.py

@@ -0,0 +1,142 @@
+"""Shared helpers for the StreamDiffusionV2 inference entrypoints."""
+
+import os
+from typing import Any
+
+import av
+import numpy as np
+import torch
+import torchvision
+import torchvision.transforms.functional as TF
+from einops import rearrange
+from omegaconf import OmegaConf
+
+
+def _read_video_with_av(video_path: str) -> torch.Tensor:
+    """Read a video with PyAV when torchvision's legacy video API is absent."""
+    frames = []
+    with av.open(video_path) as container:
+        stream = container.streams.video[0]
+        for frame in container.decode(stream):
+            frames.append(frame.to_rgb().to_ndarray())
+
+    if not frames:
+        raise ValueError(f"No video frames decoded from {video_path}")
+
+    video = np.stack(frames, axis=0)
+    return torch.from_numpy(video).permute(0, 3, 1, 2).contiguous()
+
+
+def load_mp4_as_tensor(
+    video_path: str,
+    max_frames: int = None,
+    resize_hw: tuple[int, int] = None,
+    normalize: bool = True,
+) -> torch.Tensor:
+    """Load an mp4 video as a tensor with shape [C, T, H, W]."""
+    assert os.path.exists(video_path), f"Video file not found: {video_path}"
+
+    if hasattr(torchvision.io, "read_video"):
+        video, _, _ = torchvision.io.read_video(video_path, output_format="TCHW")
+    else:
+        video = _read_video_with_av(video_path)
+    if max_frames is not None:
+        video = video[:max_frames]
+
+    video = rearrange(video, "t c h w -> c t h w")
+    if resize_hw is not None:
+        _, t, _, _ = video.shape
+        video = torch.stack(
+            [TF.resize(video[:, i], resize_hw, antialias=True) for i in range(t)],
+            dim=1,
+        )
+    if video.dtype != torch.float32:
+        video = video.float()
+    if normalize:
+        video = video / 127.5 - 1.0
+
+    return video
+
+
+def resolve_config_path(config_path: str, args) -> str:
+    """Select an alternate config file when runtime flags imply one."""
+    fast = bool(args.get("fast", False)) if isinstance(args, dict) else bool(getattr(args, "fast", False))
+    if not fast:
+        return config_path
+
+    base_name = os.path.basename(config_path)
+    if base_name != "wan_causal_dmd_v2v.yaml":
+        return config_path
+
+    fast_config_path = os.path.join(os.path.dirname(config_path), "wan_causal_dmd_v2v_fast.yaml")
+    return fast_config_path if os.path.exists(fast_config_path) else config_path
+
+
+def merge_cli_config(config_path: str, args) -> OmegaConf:
+    """Load a YAML config and overlay CLI arguments onto it."""
+    config_path = resolve_config_path(config_path, args)
+    config = OmegaConf.load(config_path)
+    cli_config = OmegaConf.create(vars(args) if not isinstance(args, dict) else args)
+    config = OmegaConf.merge(config, cli_config)
+    config = normalize_acceleration_flags(config)
+
+    # CLI --step should always select the first N non-zero denoising steps from
+    # the canonical YAML schedule, then append the terminal zero step back.
+    full_denoising_list = list(config.denoising_step_list)
+    non_terminal_steps = [step for step in full_denoising_list if int(step) != 0]
+    step_value = int(config.step)
+    config.denoising_step_list = non_terminal_steps[:step_value]
+    config.denoising_step_list.append(0)
+    return config
+
+
+def load_generator_state_dict(checkpoint_folder: str):
+    """Load the generator weights from a checkpoint folder."""
+    ckpt_path = os.path.join(checkpoint_folder, "model.pt")
+    checkpoint = torch.load(ckpt_path, map_location="cpu")
+
+    def add_model_prefix(state_dict):
+        return {
+            key if key.startswith("model.") else f"model.{key}": value
+            for key, value in state_dict.items()
+        }
+
+    if isinstance(checkpoint, dict):
+        for key in ("generator", "generator_ema", "state_dict"):
+            if key in checkpoint:
+                return ckpt_path, add_model_prefix(checkpoint[key])
+
+    return ckpt_path, add_model_prefix(checkpoint)
+
+
+def _get_flag(config: Any, key: str, default=False):
+    if isinstance(config, dict):
+        return config.get(key, default)
+    return getattr(config, key, default)
+
+
+def _set_flag(config: Any, key: str, value) -> None:
+    if isinstance(config, dict):
+        config[key] = value
+    else:
+        setattr(config, key, value)
+
+
+def normalize_acceleration_flags(config):
+    """Apply shared CLI/runtime flag semantics for fast and TensorRT modes."""
+    use_taehv = bool(_get_flag(config, "use_taehv", False))
+    use_tensorrt = bool(_get_flag(config, "use_tensorrt", False))
+    fast = bool(_get_flag(config, "fast", False))
+
+    if fast:
+        use_taehv = True
+        use_tensorrt = True
+
+    # The current TensorRT path is implemented on top of the TAEHV decoder.
+    if use_tensorrt:
+        use_taehv = True
+
+    _set_flag(config, "use_taehv", use_taehv)
+    _set_flag(config, "use_tensorrt", use_tensorrt)
+    _set_flag(config, "fast", fast)
+    return config

streamv2v/inference_pipe.py

@@ -0,0 +1,1022 @@
+"""
+Refactored multi-rank inference pipeline with communication abstractions.
+
+This is a refactored version of inference_pipe_multi.py that uses the new
+communication abstraction layers for better code organization and maintainability.
+"""
+
+from models.wan.causal_stream_inference import CausalStreamInferencePipeline
+from models.util import set_seed
+from diffusers.utils import export_to_video
+from models.data import TextDataset
+import argparse
+from dataclasses import dataclass
+import torch
+import torch.distributed as dist
+import os
+import time
+import numpy as np
+import logging
+
+try:
+    from streamv2v.inference import compute_noise_scale_and_step
+    from streamv2v.communication import (
+        DistributedCommunicator,
+        ModelDataTransfer,
+        BufferManager,
+        KVCacheManager,
+        CommunicationConfig,
+        init_distributed,
+        setup_logging,
+        compute_balanced_split
+    )
+    from streamv2v.inference_common import (
+        load_generator_state_dict,
+        load_mp4_as_tensor,
+        merge_cli_config,
+    )
+except ModuleNotFoundError:
+    from inference import compute_noise_scale_and_step
+    from communication import (
+        DistributedCommunicator,
+        ModelDataTransfer,
+        BufferManager,
+        KVCacheManager,
+        CommunicationConfig,
+        init_distributed,
+        setup_logging,
+        compute_balanced_split
+    )
+    from inference_common import (
+        load_generator_state_dict,
+        load_mp4_as_tensor,
+        merge_cli_config,
+    )
+
+LOGGER = logging.getLogger(__name__)
+
+
+def compute_default_block_distribution(total_blocks: int, world_size: int) -> list[list[int]]:
+    """Split transformer blocks into contiguous ranges for each rank."""
+    if world_size == 2:
+        midpoint = total_blocks // 2
+        return [[0, midpoint], [midpoint, total_blocks]]
+
+    base = total_blocks // world_size
+    rem = total_blocks % world_size
+    start = 0
+    block_ranges = []
+    for rank in range(world_size):
+        size = base + (1 if rank < rem else 0)
+        end = start + size if rank < world_size - 1 else total_blocks
+        block_ranges.append([start, end])
+        start = end
+    return block_ranges
+
+
+@dataclass
+class MultiGPUDemoInputSession:
+    prompt: str
+    noise_scale: float
+    init_noise_scale: float
+    chunk_size: int
+    current_start: int
+    current_end: int
+    last_image: torch.Tensor
+    chunk_idx: int = 0
+    input_batch: int = 0
+    current_step: int = 0
+    noisy_latents: torch.Tensor | None = None
+
+
+class InferencePipelineManager:
+    """
+    Manages the inference pipeline with communication abstractions.
+    
+    This class encapsulates the main inference logic and uses the communication
+    abstractions for distributed operations.
+    """
+    
+    def __init__(self, config, device: torch.device, rank: int, world_size: int):
+        """
+        Initialize the inference pipeline manager.
+        
+        Args:
+            config: Configuration object
+            device: GPU device
+            rank: Current rank
+            world_size: Total number of ranks
+        """
+        self.config = config
+        self.device = device
+        self.rank = rank
+        self.world_size = world_size
+
+        self.com_stream = torch.cuda.Stream()
+        self.control_stream = torch.cuda.Stream()
+        
+        # Setup logging
+        self.logger = setup_logging(rank)
+        
+        # Initialize communication components
+        comm_config = CommunicationConfig(
+            max_outstanding=config.get('max_outstanding', 1),
+            buffer_pool_size=config.get('buffer_pool_size', 10),
+            enable_buffer_reuse=config.get('enable_buffer_reuse', True)
+        )
+        
+        self.communicator = DistributedCommunicator(rank, world_size, device, comm_config)
+        self.buffer_manager = BufferManager(device, comm_config)
+        
+        # Initialize pipeline
+        self.pipeline = CausalStreamInferencePipeline(config, device=str(device))
+        self.pipeline.to(device=str(device), dtype=torch.bfloat16)
+        
+        # Initialize KV cache manager
+        self.kv_cache_manager = KVCacheManager(self.pipeline, device)
+        
+        # Initialize model data transfer
+        self.data_transfer = ModelDataTransfer(
+            self.communicator, 
+            self.buffer_manager, 
+            self.kv_cache_manager, 
+            comm_config
+        )
+        
+        # Performance tracking
+        self.t_dit = 100.0
+        self.t_total = 100.0
+        self.processed = 0
+        self.schedule_step = (self.world_size + len(config.denoising_step_list)) * 2
+        self.processed_offset = 3
+        self.base_chunk_size = 4
+        self.t_refresh = 50
+        self.profile = bool(config.get('profile', False))
+        self.encode_fps_list: list[float] = []
+        self.decode_fps_list: list[float] = []
+        
+        self.logger.info(f"Initialized InferencePipelineManager for rank {rank}")
+    
+    def load_model(self, checkpoint_folder: str):
+        """Load the model from checkpoint."""
+        ckpt_path, state_dict = load_generator_state_dict(checkpoint_folder)
+        try:
+            self.pipeline.generator.load_state_dict(state_dict, strict=True)
+        except RuntimeError as exc:
+            self.logger.warning(f"Strict load_state_dict failed: {exc}; retrying with strict=False")
+            self.pipeline.generator.load_state_dict(state_dict, strict=False)
+        self.logger.info(f"Model loaded successfully from {ckpt_path}")
+    
+    def prepare_pipeline(self, text_prompts: list, noise: torch.Tensor, 
+                        block_mode: str, current_start: int, current_end: int, block_num: torch.Tensor):
+        """Prepare the pipeline for inference."""
+        denoised_pred = self.pipeline.prepare(
+            text_prompts=text_prompts, 
+            device=self.device, 
+            dtype=torch.bfloat16, 
+            noise=noise, 
+            block_mode=block_mode, 
+            current_start=current_start, 
+            current_end=current_end,
+            block_num=block_num
+        )
+        
+        # Broadcast the prepared result from rank 0
+        self.data_transfer.broadcast_tensor(denoised_pred, src=0)
+        return denoised_pred
+
+    def _wait_for_outstanding(self, outstanding: list) -> None:
+        """Keep the number of queued async sends bounded."""
+        while len(outstanding) >= self.config.get('max_outstanding', 1):
+            oldest = outstanding.pop(0)
+            for work in oldest:
+                work.wait()
+
+    def _drain_outstanding(self, outstanding: list) -> None:
+        """Wait for all queued async sends to complete."""
+        while outstanding:
+            oldest = outstanding.pop(0)
+            for work in oldest:
+                work.wait()
+
+    def _maybe_schedule_blocks(self, schedule_block: bool, threshold: int, block_num: torch.Tensor, total_blocks: int) -> bool:
+        """Run one-time block rebalancing when the warmup threshold is reached."""
+        if schedule_block and self.processed >= threshold:
+            self._handle_block_scheduling(block_num, total_blocks)
+            return False
+        return schedule_block
+
+    def _receive_latent_data(self, previous_latent_data, num_steps: int):
+        """Release the previous payload and receive the next one from the upstream rank."""
+        with torch.cuda.stream(self.com_stream):
+            if previous_latent_data is not None:
+                self.data_transfer.release_latent_data(previous_latent_data)
+            latent_data = self.data_transfer.receive_latent_data_async(num_steps)
+        torch.cuda.current_stream().wait_stream(self.com_stream)
+        return latent_data
+
+    def _run_worker_stage(self, role: str, latent_data, block_num: torch.Tensor):
+        """Execute the local DiT blocks for a middle or output rank."""
+        return self.pipeline.inference(
+            noise=latent_data.original_latents,
+            current_start=latent_data.current_start,
+            current_end=latent_data.current_end,
+            current_step=latent_data.current_step,
+            block_mode=role,
+            block_num=block_num,
+            patched_x_shape=latent_data.patched_x_shape,
+            block_x=latent_data.latents,
+        )
+
+    def _send_worker_result(self, role: str, outstanding: list, latent_data, denoised_pred: torch.Tensor) -> None:
+        """Forward the payload that should continue around the pipeline ring."""
+        if role == 'output':
+            latents = latent_data.latents
+            original_latents = denoised_pred
+        else:
+            latents = denoised_pred
+            original_latents = latent_data.original_latents
+
+        with torch.cuda.stream(self.com_stream):
+            work_objects = self.data_transfer.send_latent_data_async(
+                chunk_idx=latent_data.chunk_idx,
+                latents=latents,
+                original_latents=original_latents,
+                patched_x_shape=latent_data.patched_x_shape,
+                current_start=latent_data.current_start,
+                current_end=latent_data.current_end,
+                current_step=latent_data.current_step
+            )
+            outstanding.append(work_objects)
+
+    def _decode_prediction(self, denoised_pred: torch.Tensor) -> np.ndarray:
+        """Decode the newest latent prediction into pixel-space frames."""
+        video = self._timed_stream_decode(denoised_pred[[-1]])
+        video = (video * 0.5 + 0.5).clamp(0, 1)
+        video = video[0].permute(0, 2, 3, 1).contiguous()
+        return video.cpu().float().numpy()
+
+    def _rank_loop_complete(self, num_chunks: int, num_steps: int) -> bool:
+        """Return whether a non-output rank has processed all required chunks."""
+        return (
+            self.processed + self.processed_offset
+            >= num_chunks + num_steps * self.world_size + self.world_size - self.rank - 1
+        )
+
+    def _safe_mean(self, values: list) -> float:
+        if not values:
+            return 0.0
+        return float(np.mean(np.array(values)))
+
+    def _record_stage_fps(self, values: list[float], num_frames: int, elapsed: float) -> None:
+        if self.profile and elapsed > 0 and num_frames > 0:
+            values.append(num_frames / elapsed)
+
+    def _timing_enabled(self, schedule_block: bool = False) -> bool:
+        """Only force GPU synchronization when profiling or schedule calibration needs it."""
+        return self.profile or schedule_block
+
+    def _sync_for_timing(self, schedule_block: bool = False) -> None:
+        if self._timing_enabled(schedule_block):
+            torch.cuda.synchronize()
+
+    def _timed_stream_encode(self, images: torch.Tensor) -> torch.Tensor:
+        self._sync_for_timing()
+        start_time = time.time()
+        latents = self.pipeline.vae.stream_encode(images)
+        self._sync_for_timing()
+        self._record_stage_fps(self.encode_fps_list, int(images.shape[2]), time.time() - start_time)
+        return latents
+
+    def _timed_stream_decode(self, denoised_pred: torch.Tensor) -> torch.Tensor:
+        self._sync_for_timing()
+        start_time = time.time()
+        video = self.pipeline.vae.stream_decode_to_pixel(denoised_pred)
+        self._sync_for_timing()
+        self._record_stage_fps(self.decode_fps_list, int(video.shape[1]), time.time() - start_time)
+        return video
+
+    def reset_stream_state(self, reset_encode: bool = False, reset_decode: bool = False) -> None:
+        """Reset cached inference state before starting a new prompt/session."""
+        self.pipeline.kv_cache1 = None
+        self.pipeline.crossattn_cache = None
+        self.pipeline.block_x = None
+        self.pipeline.hidden_states = None
+        self.processed = 0
+
+        if reset_encode:
+            self.pipeline.vae.model.first_encode = True
+        if reset_decode:
+            self.pipeline.vae.model.first_decode = True
+
+    def _broadcast_initial_noise(self, noisy_latents: torch.Tensor) -> None:
+        latents_shape = torch.tensor(noisy_latents.shape, dtype=torch.int64, device=self.device)
+        self.communicator.broadcast_tensor(latents_shape, src=0)
+        self.communicator.broadcast_tensor(noisy_latents, src=0)
+
+    def _receive_initial_noise(self) -> torch.Tensor:
+        latents_shape = torch.zeros(5, dtype=torch.int64, device=self.device)
+        self.communicator.broadcast_tensor(latents_shape, src=0)
+        noisy_latents = torch.zeros(tuple(latents_shape.tolist()), dtype=torch.bfloat16, device=self.device)
+        self.communicator.broadcast_tensor(noisy_latents, src=0)
+        return noisy_latents
+
+    def get_demo_chunk_size(self) -> int:
+        """Return the demo stream chunk size in frames."""
+        return self.base_chunk_size * self.pipeline.num_frame_per_block
+
+    def get_demo_first_batch_num_frames(self) -> int:
+        """Return the number of frames required to initialize a demo stream."""
+        return 1 + self.get_demo_chunk_size()
+
+    def prepare_demo_input_session(self, images: torch.Tensor, prompt: str, block_num: torch.Tensor, noise_scale: float) -> None:
+        """Initialize rank 0 for demo streaming and broadcast the first noisy latents."""
+        self.reset_stream_state(reset_encode=True)
+        torch.cuda.empty_cache()
+
+        latents = self._timed_stream_encode(images)
+        latents = latents.transpose(2, 1).contiguous().to(dtype=torch.bfloat16)
+        noise = torch.randn_like(latents)
+        noisy_latents = noise * noise_scale + latents * (1 - noise_scale)
+
+        self._broadcast_initial_noise(noisy_latents)
+        self.prepare_pipeline(
+            text_prompts=[prompt],
+            noise=noisy_latents,
+            block_mode='input',
+            current_start=0,
+            current_end=self.pipeline.frame_seq_length * 2,
+            block_num=block_num,
+        )
+        torch.cuda.empty_cache()
+        dist.barrier()
+
+    def start_demo_input_stream_session(
+        self,
+        prompt: str,
+        images: torch.Tensor,
+        block_num: torch.Tensor,
+        noise_scale: float,
+    ) -> MultiGPUDemoInputSession:
+        """Initialize rank 0 and return the demo stream session state."""
+        chunk_size = self.get_demo_chunk_size()
+        self.prepare_demo_input_session(images, prompt, block_num, noise_scale)
+        current_start = self.pipeline.frame_seq_length * (1 + chunk_size // self.base_chunk_size)
+        current_end = current_start + (chunk_size // self.base_chunk_size) * self.pipeline.frame_seq_length
+        return MultiGPUDemoInputSession(
+            prompt=prompt,
+            noise_scale=noise_scale,
+            init_noise_scale=noise_scale,
+            chunk_size=chunk_size,
+            current_start=current_start,
+            current_end=current_end,
+            last_image=images[:, :, [-1]],
+        )
+
+    def prepare_demo_worker_session(self, prompt: str, block_mode: str, block_num: torch.Tensor, decode_initial: bool = False):
+        """Initialize a non-input rank for demo streaming from the broadcast first chunk."""
+        self.reset_stream_state(reset_decode=(block_mode == 'output'))
+        torch.cuda.empty_cache()
+
+        noisy_latents = self._receive_initial_noise()
+        denoised_pred = self.prepare_pipeline(
+            text_prompts=[prompt],
+            noise=noisy_latents,
+            block_mode=block_mode,
+            current_start=0,
+            current_end=self.pipeline.frame_seq_length * 2,
+            block_num=block_num,
+        )
+        torch.cuda.empty_cache()
+        dist.barrier()
+
+        if decode_initial:
+            return self._decode_prediction(denoised_pred)
+        return None
+
+    def maybe_refresh_demo_input_window(self, session: MultiGPUDemoInputSession) -> None:
+        """Wrap the KV-cache window once the streaming refresh threshold is reached."""
+        if session.current_start // self.pipeline.frame_seq_length >= self.t_refresh:
+            session.current_start = self.pipeline.kv_cache_length - self.pipeline.frame_seq_length
+            session.current_end = session.current_start + (session.chunk_size // self.base_chunk_size) * self.pipeline.frame_seq_length
+
+    def prepare_demo_input_batch(self, session: MultiGPUDemoInputSession, images: torch.Tensor) -> None:
+        """Encode one demo chunk and update the session with the current denoising step."""
+        num_frames = images.shape[2]
+        session.input_batch = num_frames // session.chunk_size
+        session.noise_scale, session.current_step = compute_noise_scale_and_step(
+            input_video_original=torch.cat([session.last_image, images], dim=2),
+            end_idx=num_frames + 1,
+            chunk_size=num_frames,
+            noise_scale=float(session.noise_scale),
+            init_noise_scale=float(session.init_noise_scale),
+        )
+
+        latents = self._timed_stream_encode(images)
+        latents = latents.transpose(2, 1).contiguous().to(dtype=torch.bfloat16)
+        noise = torch.randn_like(latents)
+        session.noisy_latents = noise * session.noise_scale + latents * (1 - session.noise_scale)
+
+    def run_demo_input_step(
+        self,
+        session: MultiGPUDemoInputSession,
+        block_num: torch.Tensor,
+        previous_latent_data=None,
+    ):
+        """Run one rank-0 demo step from the current session batch."""
+        if session.noisy_latents is None or session.input_batch <= 0:
+            raise RuntimeError("demo input batch was not prepared before run_demo_input_step")
+
+        denoised_pred, patched_x_shape = self.run_input_stage(
+            noisy_latents=session.noisy_latents[:, -session.input_batch].unsqueeze(1),
+            current_start=session.current_start,
+            current_end=session.current_end,
+            current_step=session.current_step,
+            block_num=block_num,
+            previous_latent_data=previous_latent_data,
+        )
+        session.input_batch -= 1
+        return denoised_pred, patched_x_shape
+
+    def advance_demo_input_stream_session(self, session: MultiGPUDemoInputSession, images: torch.Tensor) -> None:
+        """Advance the demo stream session after a chunk has been queued downstream."""
+        session.last_image = images[:, :, [-1]]
+        session.chunk_idx += 1
+        session.current_start = session.current_end
+        session.current_end += (session.chunk_size // self.base_chunk_size) * self.pipeline.frame_seq_length
+
+    def send_demo_input_prompt_update(
+        self,
+        prompt: str,
+        device: torch.device,
+        num_steps: int,
+        chunk_idx: int,
+        denoised_pred: torch.Tensor,
+        patched_x_shape: torch.Tensor,
+        current_step: int,
+    ) -> None:
+        """Signal a prompt restart from rank 0 and drain in-flight returns from downstream ranks."""
+        with torch.cuda.stream(self.com_stream):
+            self.data_transfer.send_latent_data_async(
+                chunk_idx=-1,
+                latents=denoised_pred.new_zeros([1] * denoised_pred.ndim),
+                original_latents=self.pipeline.hidden_states.new_zeros([1] * self.pipeline.hidden_states.ndim),
+                patched_x_shape=patched_x_shape,
+                current_start=self.pipeline.kv_cache_starts,
+                current_end=self.pipeline.kv_cache_ends,
+                current_step=int(current_step),
+            )
+            self.data_transfer.send_prompt_async(prompt, device)
+            for _ in range(min(chunk_idx, self.world_size - 1)):
+                pending_data = self.data_transfer.receive_latent_data_async(num_steps)
+                self.data_transfer.release_latent_data(pending_data)
+
+    def send_demo_middle_prompt_update(
+        self,
+        prompt: str,
+        device: torch.device,
+        denoised_pred: torch.Tensor | None,
+        latent_data,
+    ) -> None:
+        """Forward a prompt restart from a middle rank to the next rank."""
+        sentinel_source = denoised_pred if denoised_pred is not None else latent_data.latents
+        with torch.cuda.stream(self.com_stream):
+            self.data_transfer.send_latent_data_async(
+                chunk_idx=-1,
+                latents=sentinel_source.new_zeros([1] * sentinel_source.ndim),
+                original_latents=latent_data.original_latents,
+                patched_x_shape=latent_data.patched_x_shape,
+                current_start=latent_data.current_start,
+                current_end=latent_data.current_end,
+                current_step=int(latent_data.current_step),
+            )
+            self.data_transfer.send_prompt_async(prompt, device)
+
+    def run_input_stage(self, noisy_latents: torch.Tensor, current_start: int, current_end: int, current_step: int, block_num: torch.Tensor, previous_latent_data=None):
+        """Run the rank-0 stage for one streaming chunk."""
+        if previous_latent_data is not None and self.processed >= self.world_size:
+            self.pipeline.hidden_states.copy_(previous_latent_data.original_latents)
+            self.pipeline.kv_cache_starts.copy_(previous_latent_data.current_start)
+            self.pipeline.kv_cache_ends.copy_(previous_latent_data.current_end)
+
+        return self.pipeline.inference(
+            noise=noisy_latents,
+            current_start=current_start,
+            current_end=current_end,
+            current_step=current_step,
+            block_mode='input',
+            block_num=block_num,
+        )
+    
+    def run_rank_0_loop(self, input_video_original: torch.Tensor, prompts: list, 
+                       num_chunks: int, num_steps: int, chunk_size: int,
+                       block_num: torch.Tensor, noise_scale: float, 
+                       schedule_block: bool, total_blocks: int):
+        """
+        Run the main loop for rank 0 (encoder + async send).
+        
+        This method encapsulates the rank 0 logic using the communication abstractions.
+        """
+        self.logger.info("Starting rank 0 inference loop")
+        
+        # Initialize variables
+        start_idx = 0
+        end_idx = 1 + chunk_size
+        current_start = 0
+        current_end = self.pipeline.frame_seq_length * (1+chunk_size//self.base_chunk_size)
+        init_noise_scale = noise_scale
+        
+        outstanding = []
+        latent_data = None
+        
+        self._sync_for_timing(schedule_block)
+        start_time = time.time()
+        
+        while True:
+            # Process new chunk if available
+            start_idx = end_idx
+            end_idx = end_idx + chunk_size
+            current_start = current_end
+            current_end = current_end + (chunk_size // self.base_chunk_size) * self.pipeline.frame_seq_length
+
+            if schedule_block:
+                self._sync_for_timing(schedule_block)
+                start_vae = time.time()
+
+            if end_idx <= input_video_original.shape[2]:
+                inp = input_video_original[:, :, start_idx:end_idx]
+
+                noise_scale, current_step = compute_noise_scale_and_step(
+                    input_video_original, end_idx, chunk_size, noise_scale, init_noise_scale
+                )
+                
+                latents = self._timed_stream_encode(inp)
+                latents = latents.transpose(2, 1).contiguous().to(dtype=torch.bfloat16)
+                
+                noise = torch.randn_like(latents)
+                noisy_latents = noise * noise_scale + latents * (1 - noise_scale)
+
+            # if current_start//self.pipeline.frame_seq_length >= self.t_refresh:
+            #     current_start = self.pipeline.kv_cache_length - self.pipeline.frame_seq_length
+            #     current_end = current_start + (chunk_size // self.base_chunk_size) * self.pipeline.frame_seq_length
+            
+            # Measure DiT time if scheduling is enabled
+            if schedule_block:
+                self._sync_for_timing(schedule_block)
+                start_dit = time.time()
+                t_vae = start_dit - start_vae
+            
+            # Run inference
+            denoised_pred, patched_x_shape = self.pipeline.inference(
+                noise=noisy_latents,
+                current_start=current_start,
+                current_end=current_end,
+                current_step=current_step,
+                block_mode='input',
+                block_num=block_num[self.rank],
+            )
+            
+            # Update DiT timing
+            if schedule_block:
+                self._sync_for_timing(schedule_block)
+                temp = time.time() - start_dit
+                if temp < self.t_dit:
+                    self.t_dit = temp
+            
+            self.processed += 1
+            
+            with torch.cuda.stream(self.com_stream):
+                if self.processed >= self.world_size:
+                    if latent_data is not None:
+                        self.data_transfer.release_latent_data(latent_data)
+
+                    # Receive data from previous rank
+                    latent_data = self.data_transfer.receive_latent_data_async(num_steps)
+            
+            torch.cuda.current_stream().wait_stream(self.com_stream)
+            
+            # Wait for outstanding operations
+            self._wait_for_outstanding(outstanding)
+            
+            # Send data to next rank
+            with torch.cuda.stream(self.com_stream):
+                work_objects = self.data_transfer.send_latent_data_async(
+                    chunk_idx=start_idx,
+                    latents=denoised_pred,
+                    original_latents=self.pipeline.hidden_states,
+                    patched_x_shape=patched_x_shape,
+                    current_start=self.pipeline.kv_cache_starts,
+                    current_end=self.pipeline.kv_cache_ends,
+                    current_step=current_step
+                )
+                outstanding.append(work_objects)
+                # Handle block scheduling
+                if schedule_block and self.processed >= self.schedule_step:
+                    self._handle_block_scheduling(block_num, total_blocks)
+                    schedule_block = False
+
+            # Update timing and check completion
+            if self._timing_enabled(schedule_block):
+                self._sync_for_timing(schedule_block)
+                end_time = time.time()
+                t = end_time - start_time
+                self.logger.info(f"Encode {self.processed}, time: {t:.4f} s, fps: {inp.shape[2]/t:.4f}")
+
+                if schedule_block:
+                    t_total = self.t_dit + t_vae
+                    if t_total < self.t_total:
+                        self.t_total = t_total
+                start_time = end_time
+
+            if self.processed >= self.world_size:
+                self.pipeline.hidden_states.copy_(latent_data.original_latents)
+                self.pipeline.kv_cache_starts.copy_(latent_data.current_start)
+                self.pipeline.kv_cache_ends.copy_(latent_data.current_end)
+
+            if self.processed + self.processed_offset >= num_chunks + num_steps * self.world_size + self.world_size - self.rank - 1:
+                break
+
+        if latent_data is not None:
+            self.data_transfer.release_latent_data(latent_data)
+        self._drain_outstanding(outstanding)
+        self.logger.info(f"VAE Encode Average FPS: {self._safe_mean(self.encode_fps_list):.4f}")
+        self.logger.info("Rank 0 inference loop completed")
+    
+    def run_final_rank_loop(self, num_chunks: int, num_steps: int, chunk_size: int,
+                           block_num: torch.Tensor, output_folder: str, fps: int,
+                           schedule_block: bool, total_blocks: int, results: dict):
+        """Run the worker loop for the output rank."""
+        self.run_worker_rank_loop(
+            role='output',
+            num_chunks=num_chunks,
+            num_steps=num_steps,
+            chunk_size=chunk_size,
+            block_num=block_num,
+            schedule_block=schedule_block,
+            total_blocks=total_blocks,
+            output_folder=output_folder,
+            fps=fps,
+            results=results,
+        )
+    
+    def run_middle_rank_loop(self, num_chunks: int, num_steps: int, chunk_size: int,
+                            block_num: torch.Tensor, schedule_block: bool, total_blocks: int):
+        """Run the worker loop for a middle rank."""
+        self.run_worker_rank_loop(
+            role='middle',
+            num_chunks=num_chunks,
+            num_steps=num_steps,
+            chunk_size=chunk_size,
+            block_num=block_num,
+            schedule_block=schedule_block,
+            total_blocks=total_blocks,
+        )
+
+    def run_worker_rank_loop(
+        self,
+        role: str,
+        num_chunks: int,
+        num_steps: int,
+        chunk_size: int,
+        block_num: torch.Tensor,
+        schedule_block: bool,
+        total_blocks: int,
+        output_folder: str = None,
+        fps: int = None,
+        results: dict = None,
+    ):
+        """Run the shared receive -> infer -> forward loop for middle and output ranks."""
+        if role not in {'middle', 'output'}:
+            raise ValueError(f"Unsupported worker role: {role}")
+
+        self.logger.info(f"Starting {role} rank inference loop")
+
+        if role == 'output':
+            if output_folder is None or fps is None or results is None:
+                raise ValueError("output rank requires output_folder, fps, and results")
+            os.makedirs(output_folder, exist_ok=True)
+            save_results = 1
+
+        outstanding = []
+        fps_list = []
+        latent_data = None
+
+        self._sync_for_timing(schedule_block)
+        start_time = time.time()
+
+        while True:
+            latent_data = self._receive_latent_data(latent_data, num_steps)
+            schedule_block = self._maybe_schedule_blocks(
+                schedule_block,
+                self.schedule_step - self.rank,
+                block_num,
+                total_blocks,
+            )
+
+            if schedule_block:
+                self._sync_for_timing(schedule_block)
+                start_dit = time.time()
+
+            denoised_pred, _ = self._run_worker_stage(role, latent_data, block_num[self.rank])
+
+            if schedule_block:
+                self._sync_for_timing(schedule_block)
+                temp = time.time() - start_dit
+                if temp < self.t_dit:
+                    self.t_dit = temp
+
+            self.processed += 1
+            self._wait_for_outstanding(outstanding)
+            self._send_worker_result(role, outstanding, latent_data, denoised_pred)
+
+            if role == 'output':
+                if self.processed >= num_steps * self.world_size - 1:
+                    if schedule_block:
+                        self._sync_for_timing(schedule_block)
+                        start_vae = time.time()
+
+                    video = self._timed_stream_decode(denoised_pred[[-1]])
+                    video = (video * 0.5 + 0.5).clamp(0, 1)
+                    video = video[0].permute(0, 2, 3, 1).contiguous()
+                    results[save_results] = video.cpu().float().numpy()
+
+                    if self._timing_enabled(schedule_block):
+                        self._sync_for_timing(schedule_block)
+                        end_time = time.time()
+                        elapsed = end_time - start_time
+                        fps_test = video.shape[0] / elapsed
+                        if self.processed > self.schedule_step:
+                            fps_list.append(fps_test)
+                        self.logger.info(f"Decode {self.processed}, time: {elapsed:.4f} s, FPS: {fps_test:.4f}")
+
+                        if schedule_block:
+                            t_vae = end_time - start_vae
+                            t_total = t_vae + self.t_dit
+                            if t_total < self.t_total:
+                                self.t_total = t_total
+                        start_time = end_time
+                    save_results += 1
+
+                if save_results >= num_chunks:
+                    break
+            else:
+                if self._timing_enabled(schedule_block):
+                    self._sync_for_timing(schedule_block)
+                    end_time = time.time()
+                    elapsed = end_time - start_time
+                    fps_test = chunk_size / elapsed
+
+                    if self.processed > self.schedule_step:
+                        fps_list.append(fps_test)
+
+                    if schedule_block:
+                        t_total = self.t_dit
+                        if t_total < self.t_total:
+                            self.t_total = t_total
+
+                    self.logger.info(f"Middle {self.processed}, time: {elapsed:.4f} s, fps: {fps_test:.4f}")
+                    start_time = end_time
+
+                if self._rank_loop_complete(num_chunks, num_steps):
+                    break
+
+        if latent_data is not None:
+            self.data_transfer.release_latent_data(latent_data)
+        self._drain_outstanding(outstanding)
+
+        if role == 'output':
+            video_list = [results[i] for i in range(num_chunks)]
+            video = np.concatenate(video_list, axis=0)
+            fps_avg = self._safe_mean(fps_list)
+            self.logger.info(f"Video shape: {video.shape}, Average FPS: {fps_avg:.4f}")
+            self.logger.info(f"VAE Decode Average FPS: {self._safe_mean(self.decode_fps_list):.4f}")
+
+            output_path = os.path.join(output_folder, f"output_{0:03d}.mp4")
+            export_to_video(video, output_path, fps=fps)
+            self.logger.info(f"Video saved to: {output_path} (Press Ctrl+C to force exit)")
+            return
+
+        self.logger.info(f"DiT Average FPS: {self._safe_mean(fps_list):.4f}")
+        self.logger.info(f"Rank {self.rank} inference loop completed")
+    
+    def _handle_block_scheduling(self, block_num: torch.Tensor, total_blocks: int):
+        """Handle block scheduling and rebalancing."""
+        self.logger.info(f"Scheduling block in {self.processed}")
+        
+        # Gather timing information from all ranks
+        t_total_tensor = torch.tensor(self.t_total, dtype=torch.float32, device=self.device)
+        t_dit_tensor = torch.tensor(self.t_dit, dtype=torch.float32, device=self.device)
+        
+        gather_blocks = [torch.zeros_like(t_dit_tensor, dtype=torch.float32, device=self.device) 
+                        for _ in range(self.world_size)]
+        
+        dist.all_gather(gather_blocks, t_dit_tensor)
+        t_dit_list = [t_dit_i.item() for t_dit_i in gather_blocks]
+        
+        dist.all_gather(gather_blocks, t_total_tensor)
+        t_list = [t_i.item() for t_i in gather_blocks]
+        
+        # Compute new block distribution
+        new_block_num = torch.tensor(
+            compute_balanced_split(total_blocks, t_list, t_dit_list, block_num.tolist()),
+            dtype=torch.int64, device=self.device
+        )
+
+        self.logger.info(f"New block distribution: {new_block_num[self.rank].tolist()}")
+        
+        # Broadcast new block distribution
+        dist.broadcast(new_block_num, src=self.world_size - 1)
+        
+        # Rebalance KV cache
+        self.data_transfer.rebalance_kv_cache(block_num, new_block_num, total_blocks)
+        
+        # Update block_num
+        block_num.copy_(new_block_num)
+
+        start_block, end_block = block_num[self.rank][0].item(), block_num[self.rank][1].item()
+        blocks_to_keep = list(range(start_block, end_block))
+        for i in range(self.pipeline.num_transformer_blocks):
+            if i not in blocks_to_keep:
+                self.pipeline.kv_cache1[i]['k'] = self.pipeline.kv_cache1[i]['k'].cpu()
+                self.pipeline.kv_cache1[i]['v'] = self.pipeline.kv_cache1[i]['v'].cpu()
+
+        self.logger.info("Block scheduling completed")
+    
+    def cleanup(self):
+        """Clean up resources."""
+        self.data_transfer.cleanup()
+        self.logger.info("InferencePipelineManager cleanup completed")
+
+
+def main():
+    """Main function for the refactored inference pipeline."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_path", type=str)
+    parser.add_argument("--checkpoint_folder", type=str)
+    parser.add_argument("--output_folder", type=str)
+    parser.add_argument("--prompt_file_path", type=str)
+    parser.add_argument("--video_path", type=str)
+    parser.add_argument("--noise_scale", type=float, default=0.8)
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--width", type=int, default=832)
+    parser.add_argument("--fps", type=int, default=30)
+    parser.add_argument("--max_outstanding", type=int, default=1, help="max number of outstanding sends/recv to keep")
+    parser.add_argument("--dit_fsdp", action="store_true", default=False)
+    parser.add_argument("--t5_fsdp", action="store_true", default=False)
+    parser.add_argument("--ulysses_size", type=int, default=1)
+    parser.add_argument("--ring_size", type=int, default=1)
+    parser.add_argument("--step", type=int, default=2)
+    parser.add_argument("--seed", type=int, default=0, help="Random seed")
+    parser.add_argument("--schedule_block", action="store_true", default=False)
+    parser.add_argument("--profile", action="store_true", default=False, help="Enable synchronized throughput logging")
+    parser.add_argument("--t2v", action="store_true", default=False)
+    parser.add_argument("--model_type", type=str, default="T2V-1.3B", help="Model type (e.g., T2V-1.3B)")
+    parser.add_argument("--use_taehv", action="store_true", default=False, help="Use the lightweight TAEHV VAE for encode/decode")
+    parser.add_argument("--use_tensorrt", "--use_taehv_tensorrt", dest="use_tensorrt", action="store_true", default=False, help="Enable available TensorRT acceleration paths")
+    parser.add_argument("--fast", action="store_true", default=False, help="Enable the fast path: --use_taehv --use_tensorrt")
+    
+    args = parser.parse_args()
+    
+    torch.set_grad_enabled(False)
+    init_distributed()
+    
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    local_rank = int(os.environ.get("LOCAL_RANK", rank))
+    
+    assert world_size >= 2, "world_size must be at least 2"
+    
+    torch.cuda.set_device(local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+    
+    # Load configuration
+    config = merge_cli_config(args.config_path, args)
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+    LOGGER.info("Denoising Step List: %s", list(config.denoising_step_list))
+    
+    set_seed(args.seed)
+    
+    # Load input video
+    input_video_original = load_mp4_as_tensor(args.video_path, resize_hw=(args.height, args.width)).unsqueeze(0)
+    if input_video_original.dtype != torch.bfloat16:
+        input_video_original = input_video_original.to(dtype=torch.bfloat16).to(device)
+    
+    LOGGER.info("Input video tensor shape: %s", tuple(input_video_original.shape))
+    b, c, t, h, w = input_video_original.shape
+    
+    # Calculate number of chunks
+    chunk_size = 4 * config.num_frame_per_block
+    if rank == 0:
+        num_chunks = (t - 1) // chunk_size
+    else:
+        num_chunks = 0
+    num_chunks_tensor = torch.tensor([num_chunks], dtype=torch.int64, device=device)
+    dist.broadcast(num_chunks_tensor, src=0)
+    num_chunks = int(num_chunks_tensor.item())
+    
+    # Initialize pipeline manager
+    pipeline_manager = InferencePipelineManager(config, device, rank, world_size)
+    pipeline_manager.load_model(args.checkpoint_folder)
+
+    # Load prompts
+    dataset = TextDataset(args.prompt_file_path)
+    prompts = [dataset[0]]
+    num_steps = len(pipeline_manager.pipeline.denoising_step_list)
+    
+    # Determine block mode and setup block distribution
+    if rank == 0:
+        block_mode = 'input'
+    elif rank == world_size - 1:
+        block_mode = 'output'
+    else:
+        block_mode = 'middle'
+    
+    # Setup block distribution
+    total_blocks = pipeline_manager.pipeline.num_transformer_blocks
+    total_block_num = compute_default_block_distribution(total_blocks, world_size)
+    
+    block_num = torch.tensor(total_block_num, dtype=torch.int64, device=device)
+    
+    # Prepare pipeline
+    start_idx = 0
+    end_idx = 5
+    current_start = 0
+    current_end = pipeline_manager.pipeline.frame_seq_length * 2
+    
+    inp = input_video_original[:, :, start_idx:end_idx]
+    
+    # Only rank 0 performs VAE encoding operation
+    if rank == 0:
+        latents = pipeline_manager._timed_stream_encode(inp)
+        latents = latents.transpose(2, 1).contiguous().to(dtype=torch.bfloat16)
+        noise = torch.randn_like(latents)
+        noisy_latents = noise * args.noise_scale + latents * (1 - args.noise_scale)
+        
+        # First broadcast the shape information
+        latents_shape = torch.tensor(latents.shape, dtype=torch.int64, device=device)
+        pipeline_manager.communicator.broadcast_tensor(latents_shape, src=0)
+        # Then broadcast noisy_latents
+        pipeline_manager.communicator.broadcast_tensor(noisy_latents, src=0)
+    else:
+        # Other ranks receive shape info first
+        latents_shape = torch.zeros(5, dtype=torch.int64, device=device)
+        pipeline_manager.communicator.broadcast_tensor(latents_shape, src=0)
+        # Create tensor with same shape for receiving broadcast data
+        noisy_latents = torch.zeros(tuple(latents_shape.tolist()), dtype=torch.bfloat16, device=device)
+        # Receive the broadcasted noisy_latents
+        pipeline_manager.communicator.broadcast_tensor(noisy_latents, src=0)
+    
+    denoised_pred = pipeline_manager.prepare_pipeline(
+        text_prompts=prompts,
+        noise=noisy_latents,
+        block_mode=block_mode,
+        current_start=current_start,
+        current_end=current_end,
+        block_num=block_num[rank],
+    )
+
+    # Clear unused GPU memory
+    torch.cuda.empty_cache()
+
+    # Save initial result for final rank
+    if rank == world_size - 1:
+        results = {}
+        video = pipeline_manager._timed_stream_decode(denoised_pred)
+        video = (video * 0.5 + 0.5).clamp(0, 1)
+        video = video[0].permute(0, 2, 3, 1).contiguous()
+        results[0] = video.cpu().float().numpy()
+    
+    dist.barrier()
+    pipeline_manager.logger.info(f"Prepared, Block num: {block_num[rank].tolist()}")
+
+    used_mem = torch.cuda.memory_allocated(device) / 1024 / 1024 / 1024
+    total_mem = torch.cuda.get_device_properties(device).total_memory / 1024 / 1024 / 1024
+    pipeline_manager.logger.info(f"Current GPU memory usage: {used_mem:.2f} GB / {total_mem:.2f} GB")
+    
+    # Run appropriate loop based on rank
+    try:
+        if rank == 0:
+            pipeline_manager.run_rank_0_loop(
+                input_video_original, prompts, num_chunks, num_steps, chunk_size,
+                block_num, args.noise_scale, args.schedule_block, total_blocks
+            )
+        elif rank == world_size - 1:
+            pipeline_manager.run_final_rank_loop(
+                num_chunks, num_steps, chunk_size, block_num, args.output_folder,
+                args.fps, args.schedule_block, total_blocks, results
+            )
+        else:
+            pipeline_manager.run_middle_rank_loop(
+                num_chunks, num_steps, chunk_size, block_num, args.schedule_block, total_blocks
+            )
+    finally:
+        # Cleanup
+        pipeline_manager.cleanup()
+    
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

streamv2v/inference_wo_batch.py

@@ -0,0 +1,451 @@
+"""
+Single GPU Inference Pipeline - Refactored from inference_pipe.py
+
+This file extracts core logic from multi-GPU inference code to implement a complete 
+inference pipeline on a single GPU:
+1. VAE encode input video
+2. DiT inference (using input mode, processing all 30 blocks)
+3. VAE decode output video
+"""
+
+from models.wan.causal_stream_inference import CausalStreamInferencePipeline
+from models.util import set_seed
+from diffusers.utils import export_to_video
+from models.data import TextDataset
+import argparse
+import torch
+import os
+import time
+import numpy as np
+import logging
+from typing import List
+
+try:
+    from streamv2v.inference import compute_noise_scale_and_step, SingleGPUStreamSession
+    from streamv2v.inference_common import (
+        load_generator_state_dict,
+        load_mp4_as_tensor,
+        merge_cli_config,
+    )
+except ModuleNotFoundError:
+    from inference import compute_noise_scale_and_step, SingleGPUStreamSession
+    from inference_common import (
+        load_generator_state_dict,
+        load_mp4_as_tensor,
+        merge_cli_config,
+    )
+
+LOGGER = logging.getLogger(__name__)
+
+
+class SingleGPUInferencePipeline:
+    """
+    Single GPU Inference Pipeline Manager
+    
+    This class encapsulates the complete inference logic on a single GPU, 
+    including encoding, inference, and decoding.
+    """
+    
+    def __init__(self, config, device: torch.device):
+        """
+        Initialize the single GPU inference pipeline manager.
+        
+        Args:
+            config: Configuration object
+            device: GPU device
+        """
+        self.config = config
+        self.device = device
+        
+        # Setup logging
+        self.logger = logging.getLogger("SingleGPUInference")
+        self.logger.setLevel(logging.INFO)
+        # Prevent messages from propagating to the root logger (avoid double prints)
+        self.logger.propagate = False
+        
+        if not self.logger.handlers:
+            handler = logging.StreamHandler()
+            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+            handler.setFormatter(formatter)
+            self.logger.addHandler(handler)
+        
+        # Initialize pipeline
+        self.pipeline = CausalStreamInferencePipeline(config, device=str(device))
+        self.pipeline.to(device=str(device), dtype=torch.bfloat16)
+        
+        # Performance tracking
+        self.t_dit = 100.0
+        self.t_total = 100.0
+        self.processed = 0
+        self.processed_offset = 3
+        self.base_chunk_size = 4
+        self.t_refresh = 50
+        self.profile = bool(config.get("profile", False))
+        self.encode_fps_list: list[float] = []
+        self.decode_fps_list: list[float] = []
+        self._canonical_denoising_step_list = self.pipeline.denoising_step_list.clone()
+        
+        self.logger.info("Single GPU inference pipeline manager initialized")
+    
+    def load_model(self, checkpoint_folder: str):
+        """Load the model from checkpoint."""
+        ckpt_path, state_dict = load_generator_state_dict(checkpoint_folder)
+        try:
+            self.pipeline.generator.load_state_dict(state_dict, strict=True)
+        except RuntimeError as exc:
+            self.logger.warning(f"Strict load_state_dict failed: {exc}; retrying with strict=False")
+            self.pipeline.generator.load_state_dict(state_dict, strict=False)
+        self.logger.info(f"Model loaded successfully from {ckpt_path}")
+    
+    def prepare_pipeline(self, text_prompts: list, noise: torch.Tensor, 
+                        current_start: int, current_end: int, batch_denoise: bool=True):
+        """Prepare the pipeline for inference."""
+        # Use the original prepare method which now handles distributed environment gracefully
+        denoised_pred = self.pipeline.prepare(
+            text_prompts=text_prompts,
+            device=self.device,
+            dtype=torch.bfloat16,
+            block_mode='input',
+            noise=noise,
+            current_start=current_start,
+            current_end=current_end,
+            batch_denoise=batch_denoise,
+        )
+        return denoised_pred
+
+    def _sync_for_timing(self):
+        if self.profile:
+            torch.cuda.synchronize()
+
+    def _record_stage_fps(self, values: list[float], num_frames: int, elapsed: float) -> None:
+        if self.profile and elapsed > 0 and num_frames > 0:
+            values.append(num_frames / elapsed)
+
+    def _timed_stream_encode(self, images: torch.Tensor) -> torch.Tensor:
+        self._sync_for_timing()
+        start_time = time.time()
+        latents = self.pipeline.vae.stream_encode(images)
+        self._sync_for_timing()
+        self._record_stage_fps(self.encode_fps_list, int(images.shape[2]), time.time() - start_time)
+        return latents
+
+    def _timed_stream_decode(self, denoised_pred: torch.Tensor) -> torch.Tensor:
+        self._sync_for_timing()
+        start_time = time.time()
+        video = self.pipeline.vae.stream_decode_to_pixel(denoised_pred)
+        self._sync_for_timing()
+        self._record_stage_fps(self.decode_fps_list, int(video.shape[1]), time.time() - start_time)
+        return video
+
+    def reset_stream_state(self, reset_vae_flags: bool = True) -> None:
+        """Reset cached state before starting a new no-batch stream session."""
+        if reset_vae_flags:
+            self.pipeline.vae.model.first_encode = True
+            self.pipeline.vae.model.first_decode = True
+
+        self.pipeline.kv_cache1 = None
+        self.pipeline.crossattn_cache = None
+        self.pipeline.hidden_states = None
+        self.pipeline.block_x = None
+        self.processed = 0
+
+    def _encode_noisy_latents(self, images: torch.Tensor, noise_scale: float) -> torch.Tensor:
+        latents = self._timed_stream_encode(images)
+        latents = latents.transpose(2, 1).contiguous().to(dtype=torch.bfloat16)
+        noise = torch.randn_like(latents)
+        return noise * noise_scale + latents * (1 - noise_scale)
+
+    def _decode_video_array(self, denoised_pred: torch.Tensor, last_frame_only: bool = False) -> np.ndarray:
+        if last_frame_only:
+            denoised_pred = denoised_pred[[-1]]
+
+        video = self._timed_stream_decode(denoised_pred)
+        video = (video * 0.5 + 0.5).clamp(0, 1)
+        video = video[0].permute(0, 2, 3, 1).contiguous()
+        return video.detach().cpu().float().numpy()
+
+    def start_stream_session(self, prompt: str, images: torch.Tensor, noise_scale: float) -> tuple[SingleGPUStreamSession, np.ndarray]:
+        """Initialize a no-batch streaming session and return the first decoded frames."""
+        self.reset_stream_state(reset_vae_flags=True)
+        self.pipeline.denoising_step_list = self._canonical_denoising_step_list.clone()
+
+        chunk_size = self.base_chunk_size * self.pipeline.num_frame_per_block
+        current_start = 0
+        current_end = self.pipeline.frame_seq_length * (1 + chunk_size // self.base_chunk_size)
+
+        noisy_latents = self._encode_noisy_latents(images, noise_scale)
+        denoised_pred = self.prepare_pipeline(
+            text_prompts=[prompt],
+            noise=noisy_latents,
+            current_start=current_start,
+            current_end=current_end,
+            batch_denoise=False,
+        )
+        initial_video = self._decode_video_array(denoised_pred, last_frame_only=False)
+
+        session = SingleGPUStreamSession(
+            prompt=prompt,
+            noise_scale=noise_scale,
+            init_noise_scale=noise_scale,
+            chunk_size=chunk_size,
+            current_start=current_end,
+            current_end=current_end + (chunk_size // self.base_chunk_size) * self.pipeline.frame_seq_length,
+            last_image=images[:, :, [-1]],
+            processed=0,
+        )
+        return session, initial_video
+
+    def run_stream_batch(self, session: SingleGPUStreamSession, images: torch.Tensor, queue_wait_time: float | None = None) -> List[np.ndarray]:
+        """Process one or more chunk-aligned frame groups for an active no-batch stream session."""
+        num_frames = images.shape[2]
+        input_batch = num_frames // session.chunk_size
+        noise_scale, current_step = compute_noise_scale_and_step(
+            input_video_original=torch.cat([session.last_image, images], dim=2),
+            end_idx=num_frames + 1,
+            chunk_size=num_frames,
+            noise_scale=float(session.noise_scale),
+            init_noise_scale=float(session.init_noise_scale),
+        )
+        noisy_latents = self._encode_noisy_latents(images, noise_scale)
+
+        outputs: List[np.ndarray] = []
+
+        for batch_idx in range(input_batch):
+            if session.current_start // self.pipeline.frame_seq_length >= self.t_refresh:
+                session.current_start = self.pipeline.kv_cache_length - self.pipeline.frame_seq_length
+                session.current_end = session.current_start + (session.chunk_size // self.base_chunk_size) * self.pipeline.frame_seq_length
+
+            denoised_pred = self.pipeline.inference_wo_batch(
+                noise=noisy_latents[:, batch_idx].unsqueeze(1),
+                current_start=session.current_start,
+                current_end=session.current_end,
+                current_step=current_step,
+            )
+
+            session.processed += 1
+            self.processed = session.processed
+            outputs.append(self._decode_video_array(denoised_pred, last_frame_only=True))
+
+            session.current_start = session.current_end
+            session.current_end += (session.chunk_size // self.base_chunk_size) * self.pipeline.frame_seq_length
+
+        session.last_image = images[:, :, [-1]]
+        session.noise_scale = noise_scale
+        return outputs
+    
+    def run_inference(self, input_video_original: torch.Tensor, prompts: list, 
+                     num_chunks: int, chunk_size: int, noise_scale: float, 
+                     output_folder: str, fps: int, num_steps: int):
+        """
+        Run the complete single GPU inference pipeline.
+        
+        This method integrates the complete encoding, inference, and decoding pipeline.
+        """
+        self.logger.info("Starting single GPU inference pipeline")
+        
+        os.makedirs(output_folder, exist_ok=True)
+        results = {}
+        save_results = 0
+
+        fps_list = []
+        dit_fps_list = []
+        self.encode_fps_list = []
+        self.decode_fps_list = []
+        
+        # Initialize variables
+        start_idx = 0
+        end_idx = 1 + chunk_size
+        current_start = 0
+        current_end = self.pipeline.frame_seq_length * (1+chunk_size//4)
+        init_noise_scale = noise_scale
+        
+        self._sync_for_timing()
+        start_time = time.time()
+        
+        # Process first chunk (initialization)
+        if end_idx <= input_video_original.shape[2]:
+            inp = input_video_original[:, :, start_idx:end_idx]
+            
+            # VAE encoding
+            latents = self._timed_stream_encode(inp)
+            latents = latents.transpose(2, 1).contiguous().to(dtype=torch.bfloat16)
+            
+            noise = torch.randn_like(latents)
+            noisy_latents = noise * noise_scale + latents * (1 - noise_scale)
+            
+            # Prepare pipeline
+            denoised_pred = self.prepare_pipeline(
+                text_prompts=prompts,
+                noise=noisy_latents,
+                current_start=current_start,
+                current_end=current_end,
+                batch_denoise=False,
+            )
+            
+            # Save first result - only start decoding after num_steps
+            video = self._timed_stream_decode(denoised_pred)
+            video = (video * 0.5 + 0.5).clamp(0, 1)
+            video = video[0].permute(0, 2, 3, 1).contiguous()
+            results[save_results] = video.cpu().float().numpy()
+            save_results += 1
+        
+        # Process remaining chunks
+        while self.processed < num_chunks + num_steps - 1:
+            # Update indices
+            start_idx = end_idx
+            end_idx = end_idx + chunk_size
+            current_start = current_end
+            current_end = current_end + (chunk_size // 4) * self.pipeline.frame_seq_length
+            
+            if end_idx <= input_video_original.shape[2]:
+                inp = input_video_original[:, :, start_idx:end_idx]
+                
+                noise_scale, current_step = compute_noise_scale_and_step(
+                    input_video_original, end_idx, chunk_size, noise_scale, init_noise_scale
+                )
+                
+                # VAE encoding
+                latents = self._timed_stream_encode(inp)
+                latents = latents.transpose(2, 1).contiguous().to(dtype=torch.bfloat16)
+                
+                noise = torch.randn_like(latents)
+                noisy_latents = noise * noise_scale + latents * (1 - noise_scale)
+
+            if current_start//self.pipeline.frame_seq_length >= 50:
+                current_start = self.pipeline.kv_cache_length - self.pipeline.frame_seq_length
+                current_end = current_start + (chunk_size // 4) * self.pipeline.frame_seq_length
+                
+            self._sync_for_timing()
+            dit_start_time = time.time()
+            # DiT inference - using input mode to process all 30 blocks
+            denoised_pred = self.pipeline.inference_wo_batch(
+                noise=noisy_latents,
+                current_start=current_start,
+                current_end=current_end,
+                current_step=current_step,
+            )
+            
+            if self.processed >self.processed_offset:
+                self._sync_for_timing()
+                if self.profile:
+                    dit_fps_list.append(chunk_size / (time.time() - dit_start_time))
+            
+            self.processed += 1
+            
+            # VAE decoding - only start decoding after num_steps
+            video = self._timed_stream_decode(denoised_pred[[-1]])
+            video = (video * 0.5 + 0.5).clamp(0, 1)
+            video = video[0].permute(0, 2, 3, 1).contiguous()
+            
+            results[save_results] = video.cpu().float().numpy()
+            save_results += 1
+        
+            # Update timing
+            if self.profile:
+                self._sync_for_timing()
+                end_time = time.time()
+                t = end_time - start_time
+                fps_test = inp.shape[2] / t
+                fps_list.append(fps_test)
+                self.logger.info(f"Processed {self.processed}, time: {t:.4f} s, FPS: {fps_test:.4f}")
+                start_time = end_time
+        
+        # Save final video
+        video_list = [results[i] for i in range(num_chunks)]
+        video = np.concatenate(video_list, axis=0)
+        if self.profile and fps_list:
+            fps_avg = np.mean(np.array(fps_list))
+            dit_avg = np.mean(np.array(dit_fps_list)) if dit_fps_list else 0.0
+            encode_avg = np.mean(np.array(self.encode_fps_list)) if self.encode_fps_list else 0.0
+            decode_avg = np.mean(np.array(self.decode_fps_list)) if self.decode_fps_list else 0.0
+            self.logger.info(f"VAE Encode Average FPS: {encode_avg:.4f}")
+            self.logger.info(f"DiT Average FPS: {dit_avg:.4f}")
+            self.logger.info(f"VAE Decode Average FPS: {decode_avg:.4f}")
+            self.logger.info(f"Video shape: {video.shape}, Average FPS: {fps_avg:.4f}")
+        else:
+            self.logger.info(f"Video shape: {video.shape}")
+        
+        output_path = os.path.join(output_folder, f"output_{0:03d}.mp4")
+        export_to_video(video, output_path, fps=fps)
+        self.logger.info(f"Video saved to: {output_path}")
+        
+        self.logger.info("Single GPU inference pipeline completed")
+
+
+def main():
+    """Main function for the single GPU inference pipeline."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_path", type=str, required=True, help="Configuration file path")
+    parser.add_argument("--checkpoint_folder", type=str, required=True, help="Checkpoint folder path")
+    parser.add_argument("--output_folder", type=str, required=True, help="Output folder path")
+    parser.add_argument("--prompt_file_path", type=str, required=True, help="Prompt file path")
+    parser.add_argument("--video_path", type=str, required=True, help="Input video path")
+    parser.add_argument("--noise_scale", type=float, default=0.700, help="Noise scale")
+    parser.add_argument("--height", type=int, default=480, help="Video height")
+    parser.add_argument("--width", type=int, default=832, help="Video width")
+    parser.add_argument("--fps", type=int, default=16, help="Output video fps")
+    parser.add_argument("--step", type=int, default=2, help="Step")
+    parser.add_argument("--seed", type=int, default=0, help="Random seed")
+    parser.add_argument("--gpu_id", type=int, default=None, help="CUDA device index for single-GPU inference")
+    parser.add_argument("--t2v", action="store_true", default=False)
+    parser.add_argument("--model_type", type=str, default="T2V-1.3B", help="Model type (e.g., T2V-1.3B)")
+    parser.add_argument("--profile", action="store_true", default=False, help="Enable synchronized throughput logging")
+    parser.add_argument("--use_taehv", action="store_true", default=False, help="Use the lightweight TAEHV VAE for encode/decode")
+    parser.add_argument("--use_tensorrt", "--use_taehv_tensorrt", dest="use_tensorrt", action="store_true", default=False, help="Enable available TensorRT acceleration paths")
+    parser.add_argument("--fast", action="store_true", default=False, help="Enable the fast path: --use_taehv --use_tensorrt")
+    args = parser.parse_args()
+    
+    torch.set_grad_enabled(False)
+    
+    # Auto-detect device
+    if torch.cuda.is_available():
+        if args.gpu_id is not None:
+            torch.cuda.set_device(args.gpu_id)
+            device = torch.device(f"cuda:{args.gpu_id}")
+        else:
+            device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+    
+    # Load configuration
+    config = merge_cli_config(args.config_path, args)
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+    LOGGER.info("Denoising Step List: %s", list(config.denoising_step_list))
+
+    set_seed(args.seed)
+    
+    # Load input video
+    input_video_original = load_mp4_as_tensor(args.video_path, resize_hw=(args.height, args.width)).unsqueeze(0)
+    if input_video_original.dtype != torch.bfloat16:
+        input_video_original = input_video_original.to(dtype=torch.bfloat16).to(device)
+    
+    LOGGER.info("Input video tensor shape: %s", tuple(input_video_original.shape))
+    b, c, t, h, w = input_video_original.shape
+    
+    # Calculate number of chunks
+    chunk_size = 4 * config.num_frame_per_block
+    num_chunks = (t - 1) // chunk_size
+    
+    # Initialize pipeline manager
+    pipeline_manager = SingleGPUInferencePipeline(config, device)
+    pipeline_manager.load_model(args.checkpoint_folder)
+    
+    # Load prompts
+    dataset = TextDataset(args.prompt_file_path)
+    prompts = [dataset[0]]
+    num_steps = len(pipeline_manager.pipeline.denoising_step_list)
+    
+    # Run inference
+    try:
+        pipeline_manager.run_inference(
+            input_video_original, prompts, num_chunks, chunk_size, 
+            args.noise_scale, args.output_folder, args.fps, num_steps
+        )
+    except Exception as e:
+        LOGGER.exception("Error occurred during inference: %s", e)
+        raise
+
+
+if __name__ == "__main__":
+    main()