models/dit.py

from __future__ import annotations

import copy
import math
from typing import Optional

import torch
import torch.nn as nn

from .attention import AttentionRuntime
from .blocks import DiTBlock, OutputHead
from .conditioning import LQProjection
from .config import FlashVSRModelConfig
from .embeddings import sinusoidal_embedding_1d


class FlashVSRDiT(nn.Module):
    def __init__(self, config: FlashVSRModelConfig):
        super().__init__()
        config.validate()
        self.config = config
        dit = config.dit
        self.dim = dit.dim
        self.freq_dim = dit.freq_dim
        self.patch_size = tuple(dit.patch_size)
        self.patch_embedding = nn.Conv3d(
            dit.in_channels, dit.dim, kernel_size=self.patch_size, stride=self.patch_size
        )
        self.text_embedding = nn.Sequential(
            nn.Linear(dit.text_dim, dit.dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(dit.dim, dit.dim),
        )
        self.time_embedding = nn.Sequential(
            nn.Linear(dit.freq_dim, dit.dim), nn.SiLU(), nn.Linear(dit.dim, dit.dim)
        )
        self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dit.dim, dit.dim * 6))

        base_attention = copy.deepcopy(config.attention)
        if base_attention.backend == "hybrid_sliding":
            base_attention.backend = (
                "sparse"
                if (
                    base_attention.is_sparse
                    or base_attention.sparse.local_range is not None
                )
                else "dense"
            )
        if base_attention.is_causal:
            patch_t = self.patch_size[0]
            if not base_attention.causal.pad_to_block and base_attention.causal.block_size_frames % patch_t:
                raise ValueError("causal block_size_frames must be divisible by temporal patch size")
            base_attention.causal.block_size_frames = math.ceil(
                base_attention.causal.block_size_frames / patch_t
            )

        replacement_layers = set(config.replacement.layers)
        attentions = []
        for layer_index in range(dit.num_layers):
            attention = copy.deepcopy(base_attention)
            if layer_index in replacement_layers:
                if attention.sliding.spatial_mode == "local":
                    attention.sparse.local_range = attention.sliding.local_range
                if attention.sliding.temporal_mode == "local":
                    attention.backend = "sliding_local"
                    # Temporal locality uses equality between fixed groups,
                    # not the original causal <= relation.
                    attention.is_causal = False
                    attention.causal.block_size_frames = 1
                elif attention.sliding.temporal_mode == "short_causal":
                    # Keep a bounded causal bridge into the previous
                    # chunk. The specialized layout matches the one-window
                    # stream cache used by this layer.
                    attention.backend = "sliding_local"
                else:
                    # Spatial-only ablation: change only local_range. Keep
                    # the original causal mask and normal stream KV cache.
                    attention.backend = "sparse"
            attentions.append(attention)
        self.blocks = nn.ModuleList(
            DiTBlock(dit.dim, dit.num_heads, dit.ffn_dim, attention, dit.eps)
            for attention in attentions
        )
        self.head = OutputHead(dit.dim, dit.out_channels, self.patch_size, dit.eps)
        self.LQ_proj_in = (
            LQProjection(config.lq_projection, dit.dim) if config.lq_projection.enabled else None
        )

    def patchify(self, x: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int, int]]:
        x = self.patch_embedding(x)
        shape = tuple(x.shape[2:])
        return x.permute(0, 2, 3, 4, 1).flatten(1, 3), shape

    def unpatchify(self, x: torch.Tensor, shape: tuple[int, int, int]) -> torch.Tensor:
        batch = x.shape[0]
        frames, height, width = shape
        pt, ph, pw = self.patch_size
        channels = self.config.dit.out_channels
        x = x.view(batch, frames, height, width, pt, ph, pw, channels)
        return x.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(
            batch, channels, frames * pt, height * ph, width * pw
        )

    def forward(
        self,
        x: torch.Tensor,
        timestep: torch.Tensor,
        context: torch.Tensor,
        *,
        lq_video: Optional[torch.Tensor] = None,
        layer_conditions: Optional[list[torch.Tensor]] = None,
        frame_valid_mask: Optional[torch.Tensor] = None,
        context_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if x.ndim != 5:
            raise ValueError(f"DiT input must be [B,C,F,H,W], got {x.shape}")
        if timestep.ndim == 0:
            timestep = timestep.expand(x.shape[0])
        time_embedding = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep).to(dtype=x.dtype))
        time_modulation = self.time_projection(time_embedding).unflatten(1, (6, self.dim))
        context = self.text_embedding(context)
        x, token_shape = self.patchify(x)
        if lq_video is not None:
            if self.LQ_proj_in is None:
                raise ValueError("lq_video was provided but LQ projection is disabled")
            if layer_conditions is not None:
                raise ValueError("Provide lq_video or layer_conditions, not both")
            layer_conditions = self.LQ_proj_in(lq_video)
        if frame_valid_mask is not None and frame_valid_mask.shape[-1] != token_shape[0]:
            expected_frames = token_shape[0] * self.patch_size[0]
            if frame_valid_mask.shape[-1] != expected_frames:
                raise ValueError(
                    f"frame_valid_mask has {frame_valid_mask.shape[-1]} frames, expected "
                    f"{token_shape[0]} token frames or {expected_frames} input frames"
                )
            frame_valid_mask = frame_valid_mask.view(
                frame_valid_mask.shape[0], token_shape[0], self.patch_size[0]
            ).all(dim=-1)
        runtime = AttentionRuntime(
            token_shape=token_shape,
            frame_valid_mask=frame_valid_mask,
        )
        for block_index, block in enumerate(self.blocks):
            if layer_conditions is not None and block_index < len(layer_conditions):
                condition = layer_conditions[block_index].to(device=x.device, dtype=x.dtype)
                if condition.shape != x.shape:
                    raise ValueError(
                        f"LQ condition {block_index} has shape {condition.shape}, expected {x.shape}"
                    )
                x = x + condition
            x = block(x, context, time_modulation, runtime, context_mask)
        return self.unpatchify(self.head(x, time_embedding), token_shape)

    def forward_stream(
        self,
        x: torch.Tensor,
        timestep: torch.Tensor,
        context: torch.Tensor,
        *,
        layer_conditions: list[torch.Tensor],
        state,
        context_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Numerically compatible FlashVSR 1.1 streaming DiT forward."""
        if timestep.ndim == 0:
            timestep = timestep.expand(x.shape[0])
        time_embedding = self.time_embedding(
            sinusoidal_embedding_1d(self.freq_dim, timestep).to(dtype=x.dtype)
        )
        time_modulation = self.time_projection(time_embedding).unflatten(1, (6, self.dim))
        context = self.text_embedding(context)
        x, token_shape = self.patchify(x)
        expected_frames = 4
        if token_shape[0] != expected_frames:
            raise ValueError(
                f"FlashVSR stream chunk {state.chunk_index} requires {expected_frames} "
                f"latent frames, got {token_shape[0]}"
            )
        if len(state.attention_cache) not in {0, len(self.blocks)}:
            raise ValueError("attention cache must be empty or contain one entry per DiT block")
        old_caches = state.attention_cache
        old_ready_events = getattr(state, "attention_cache_ready_events", [])
        # Move ownership out of the long-lived state before updating layers.
        # Each old layer cache is released immediately after that layer has
        # consumed it, instead of keeping complete old and new cache sets
        # alive until the end of the DiT forward.
        state.attention_cache = []
        state.attention_cache_ready_events = []
        next_caches = []
        next_ready_events = []
        offload_layers = getattr(state, "attention_cache_offload_layers", frozenset())
        if offload_layers and max(offload_layers) >= len(self.blocks):
            raise ValueError(
                f"stream cache offload layer index exceeds DiT depth {len(self.blocks)}"
            )
        transfer_stream = getattr(state, "attention_cache_transfer_stream", None)
        for block_index, block in enumerate(self.blocks):
            if block_index < len(layer_conditions):
                condition = layer_conditions[block_index].to(device=x.device, dtype=x.dtype)
                if condition.shape != x.shape:
                    raise ValueError(
                        f"stream LQ condition {block_index} has {condition.shape}, expected {x.shape}"
                    )
                x = x + condition
            cache = (
                None
                if not old_caches
                else old_caches[block_index]
            )
            if old_caches:
                old_caches[block_index] = None
                if old_ready_events:
                    ready_event = old_ready_events[block_index]
                    old_ready_events[block_index] = None
                else:
                    ready_event = None
            else:
                ready_event = None
            if ready_event is not None and x.is_cuda:
                # The previous chunk's D2H copy may still be in flight on
                # the transfer stream. This wait orders the host source before
                # the H2D copy performed by attention on the main stream.
                torch.cuda.current_stream(x.device).wait_event(ready_event)
            runtime = AttentionRuntime(
                token_shape=token_shape,
                temporal_offset=4 * state.chunk_index,
                stream_cache=cache,
                stream_topk=state.attention_topk,
                stream_cache_blocks=state.attention_cache_blocks,
                stream_block_size=state.attention_block_size,
                stream_ffn_chunk_tokens=state.ffn_chunk_tokens,
                stream_block_m=state.attention_block_m,
                stream_compact_cache=getattr(state, "compact_cache", False),
                stream_write_cache=getattr(state, "write_cache", True),
                stream_release_workspaces=getattr(state, "release_attention_workspaces", False),
            )
            trace = state.memory_trace
            if trace is not None and x.is_cuda:
                def cache_stats(value):
                    if value is None:
                        return {"tokens": 0, "bytes": 0, "device": None}
                    k, v = value
                    return {
                        "tokens": int(k.shape[2]),
                        "bytes": int(k.numel() * k.element_size() + v.numel() * v.element_size()),
                        "device": str(k.device),
                    }

                before_cache = cache_stats(cache)
                trace.append({
                    "chunk": int(state.chunk_index),
                    "layer": int(block_index),
                    "stage": "before",
                    "attention": type(block.self_attn).__name__,
                    "alloc_mib": torch.cuda.memory_allocated(x.device) / 1024**2,
                    "reserved_mib": torch.cuda.memory_reserved(x.device) / 1024**2,
                    "peak_alloc_mib": torch.cuda.max_memory_allocated(x.device) / 1024**2,
                    "x_tokens": int(x.shape[1]),
                    "x_bytes": int(x.numel() * x.element_size()),
                    "old_cache_tokens": before_cache["tokens"],
                    "old_cache_bytes": before_cache["bytes"],
                    "old_cache_device": before_cache["device"],
                })
            try:
                x, next_cache = block.forward_stream(
                    x, context, time_modulation, runtime, context_mask
                )
            except Exception as error:
                if trace is not None and x.is_cuda:
                    trace.append({
                        "chunk": int(state.chunk_index),
                        "layer": int(block_index),
                        "stage": "error",
                        "error_type": type(error).__name__,
                        "error": str(error).splitlines()[0],
                        "alloc_mib": torch.cuda.memory_allocated(x.device) / 1024**2,
                        "reserved_mib": torch.cuda.memory_reserved(x.device) / 1024**2,
                        "peak_alloc_mib": torch.cuda.max_memory_allocated(x.device) / 1024**2,
                    })
                raise
            # A cache is consumed only by the same layer on the next stream
            # chunk. Moving selected caches to host memory changes residency,
            # not mask construction or arithmetic. Keep BF16/layout intact;
            # this opt-in path preserves deterministic cache ordering.
            if block_index in offload_layers and next_cache is not None:
                # Page-locked destinations make CUDA D2H copies asynchronous
                # and avoid the very slow pageable staging path. Copies are
                # enqueued on the transfer stream after block computation;
                # the next chunk waits on its completion event before H2D.
                host_cache = []
                transfer_stream = getattr(
                    state, "attention_cache_transfer_stream", None
                )
                current_stream = (
                    torch.cuda.current_stream(x.device) if x.is_cuda else None
                )
                if transfer_stream is not None and current_stream is not None:
                    transfer_stream.wait_stream(current_stream)
                for cache_index, value in enumerate(next_cache):
                    value = value.detach()
                    if getattr(state, "attention_cache_pin_memory", True):
                        try:
                            previous = None if cache is None else cache[cache_index]
                            if (
                                getattr(state, "reuse_host_cache", False)
                                and previous is not None
                                and previous.device.type == "cpu"
                                and previous.is_pinned()
                                and previous.shape == value.shape
                                and previous.dtype == value.dtype
                            ):
                                # H2D consumption precedes block completion on
                                # the compute stream. The transfer stream wait
                                # above prevents overwriting a still-read source.
                                host_value = previous
                            else:
                                host_value = torch.empty_like(
                                    value, device="cpu", pin_memory=True
                                )
                            if transfer_stream is not None:
                                value.record_stream(transfer_stream)
                                with torch.cuda.stream(transfer_stream):
                                    host_value.copy_(value, non_blocking=True)
                            else:
                                host_value.copy_(value, non_blocking=True)
                        except RuntimeError:
                            # Gracefully fall back on systems with an
                            # insufficient pinned-memory limit. Correctness
                            # is unchanged; only transfer throughput degrades.
                            if transfer_stream is not None:
                                transfer_stream.synchronize()
                                transfer_stream = None
                            # A failed pin allocation cannot be safely queued
                            # on the transfer stream; fall back synchronously.
                            host_value = value.to(device="cpu")
                    else:
                        host_value = value.to(device="cpu")
                    host_cache.append(host_value)
                next_cache = tuple(host_cache)
                ready_event = None
                if transfer_stream is not None:
                    with torch.cuda.stream(transfer_stream):
                        ready_event = transfer_stream.record_event()
                next_ready_events.append(ready_event)
            else:
                next_ready_events.append(None)
            if trace is not None and x.is_cuda:
                next_stats = cache_stats(next_cache)
                trace.append({
                    "chunk": int(state.chunk_index),
                    "layer": int(block_index),
                    "stage": "after",
                    "attention": type(block.self_attn).__name__,
                    "alloc_mib": torch.cuda.memory_allocated(x.device) / 1024**2,
                    "reserved_mib": torch.cuda.memory_reserved(x.device) / 1024**2,
                    "peak_alloc_mib": torch.cuda.max_memory_allocated(x.device) / 1024**2,
                    "next_cache_tokens": next_stats["tokens"],
                    "next_cache_bytes": next_stats["bytes"],
                    "next_cache_device": next_stats["device"],
                })
            next_caches.append(next_cache)
            del runtime, cache
        state.attention_cache = next_caches
        state.attention_cache_ready_events = next_ready_events
        return self.unpatchify(self.head(x, time_embedding), token_shape)