algorithms/worldmem/models/dit.py

"""
References:
    - DiT: https://github.com/facebookresearch/DiT/blob/main/models.py
    - Diffusion Forcing: https://github.com/buoyancy99/diffusion-forcing/blob/main/algorithms/diffusion_forcing/models/unet3d.py
    - Latte: https://github.com/Vchitect/Latte/blob/main/models/latte.py
"""

from typing import Optional, Literal
import torch
from torch import nn
from .rotary_embedding_torch import RotaryEmbedding
from einops import rearrange
from .attention import SpatialAxialAttention, TemporalAxialAttention
from timm.models.vision_transformer import Mlp
from timm.layers.helpers import to_2tuple
import math
from collections import namedtuple
from typing import Optional, Callable

def modulate(x, shift, scale):
    fixed_dims = [1] * len(shift.shape[1:])
    shift = shift.repeat(x.shape[0] // shift.shape[0], *fixed_dims)
    scale = scale.repeat(x.shape[0] // scale.shape[0], *fixed_dims)
    while shift.dim() < x.dim():
        shift = shift.unsqueeze(-2)
        scale = scale.unsqueeze(-2)
    return x * (1 + scale) + shift

def gate(x, g):
    fixed_dims = [1] * len(g.shape[1:])
    g = g.repeat(x.shape[0] // g.shape[0], *fixed_dims)
    while g.dim() < x.dim():
        g = g.unsqueeze(-2)
    return g * x


class PatchEmbed(nn.Module):
    """2D Image to Patch Embedding"""

    def __init__(
        self,
        img_height=256,
        img_width=256,
        patch_size=16,
        in_chans=3,
        embed_dim=768,
        norm_layer=None,
        flatten=True,
    ):
        super().__init__()
        img_size = (img_height, img_width)
        patch_size = to_2tuple(patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
        self.num_patches = self.grid_size[0] * self.grid_size[1]
        self.flatten = flatten

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, x, random_sample=False):
        B, C, H, W = x.shape
        assert random_sample or (H == self.img_size[0] and W == self.img_size[1]), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        
        x = self.proj(x)
        if self.flatten:
            x = rearrange(x, "B C H W -> B (H W) C")
        else:
            x = rearrange(x, "B C H W -> B H W C")
        x = self.norm(x)
        return x


class TimestepEmbedder(nn.Module):
    """
    Embeds scalar timesteps into vector representations.
    """

    def __init__(self, hidden_size, frequency_embedding_size=256, freq_type='time_step'):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_size, bias=True),  # hidden_size is diffusion model hidden size
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size, bias=True),
        )
        self.frequency_embedding_size = frequency_embedding_size
        self.freq_type = freq_type

    @staticmethod
    def timestep_embedding(t, dim, max_period=10000, freq_type='time_step'):
        """
        Create sinusoidal timestep embeddings.
        :param t: a 1-D Tensor of N indices, one per batch element.
                          These may be fractional.
        :param dim: the dimension of the output.
        :param max_period: controls the minimum frequency of the embeddings.
        :return: an (N, D) Tensor of positional embeddings.
        """
        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
        half = dim // 2

        if freq_type == 'time_step':
            freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(device=t.device)
        elif freq_type == 'spatial': # ~(-5 5)
            freqs = torch.linspace(1.0, half, half).to(device=t.device) * torch.pi
        elif freq_type == 'angle': # 0-360
            freqs = torch.linspace(1.0, half, half).to(device=t.device) * torch.pi / 180


        args = t[:, None].float() * freqs[None]
        
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        if dim % 2:
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
        return embedding

    def forward(self, t):
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size, freq_type=self.freq_type)
        t_emb = self.mlp(t_freq)
        return t_emb


class FinalLayer(nn.Module):
    """
    The final layer of DiT.
    """

    def __init__(self, hidden_size, patch_size, out_channels):
        super().__init__()
        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))

    def forward(self, x, c):
        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
        x = modulate(self.norm_final(x), shift, scale)
        x = self.linear(x)
        return x


MEMORY_TYPE_NAMES = ("anchor", "dynamic", "revisit")
MEMORY_TYPE_ANCHOR = 0
MEMORY_TYPE_DYNAMIC = 1
MEMORY_TYPE_REVISIT = 2


class MemoryTokenCrossAttention(nn.Module):
    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, num_memory_types=3):
        super().__init__()
        mlp_hidden_dim = int(hidden_size * mlp_ratio)
        approx_gelu = lambda: nn.GELU(approximate="tanh")
        self.num_heads = num_heads
        self.num_memory_types = num_memory_types
        self.norm_q = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.norm_mem = nn.LayerNorm(hidden_size, eps=1e-6)
        self.attn = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True)
        self.norm_mlp = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.mlp = Mlp(
            in_features=hidden_size,
            hidden_features=mlp_hidden_dim,
            act_layer=approx_gelu,
            drop=0,
        )
        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
        self.memory_type_embed = nn.Embedding(num_memory_types, hidden_size)
        self.memory_type_scale = nn.Parameter(torch.ones(num_memory_types, hidden_size))
        self.memory_type_gate = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, num_memory_types, bias=True))
        self.last_gate_mean = None
        self.last_delta_ratio = None
        self.last_valid_fraction = None
        self.last_type_gate_mean = None
        for type_name in MEMORY_TYPE_NAMES[:num_memory_types]:
            setattr(self, f"last_type_gate_{type_name}_mean", None)
        nn.init.normal_(self.memory_type_embed.weight, std=0.02)
        self.reset_identity_init()

    def reset_identity_init(self):
        nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
        nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
        nn.init.constant_(self.memory_type_gate[-1].weight, 0)
        nn.init.constant_(self.memory_type_gate[-1].bias, 0)

    def _attend(self, query, memory_tokens, memory_token_mask=None, memory_token_gate=None):
        if memory_token_mask is None and memory_token_gate is None:
            out, _ = self.attn(query, memory_tokens, memory_tokens, need_weights=False)
            return out, None

        if memory_token_mask is None:
            memory_token_mask = torch.ones(
                memory_tokens.shape[:2],
                device=memory_tokens.device,
                dtype=torch.bool,
            )
        else:
            memory_token_mask = memory_token_mask.bool()
        gate_tensor = None
        if memory_token_gate is not None:
            if tuple(memory_token_gate.shape) != tuple(memory_tokens.shape[:2]):
                raise ValueError(
                    f"memory_token_gate must have shape {tuple(memory_tokens.shape[:2])}, "
                    f"got {tuple(memory_token_gate.shape)}"
                )
            gate_tensor = memory_token_gate.to(device=memory_tokens.device, dtype=query.dtype)
            memory_token_mask = memory_token_mask & (gate_tensor > 0)
        valid_rows = memory_token_mask.any(dim=1)
        out = torch.zeros_like(query)
        if valid_rows.any():
            attn_mask = None
            key_padding_mask = ~memory_token_mask[valid_rows]
            if gate_tensor is not None:
                gate_bias = torch.log(gate_tensor[valid_rows].clamp_min(1.0e-6))
                gate_bias = gate_bias[:, None, :].expand(-1, query.shape[1], -1)
                attn_mask = gate_bias.repeat_interleave(self.num_heads, dim=0)
                float_padding_mask = torch.zeros_like(gate_tensor[valid_rows], dtype=query.dtype)
                key_padding_mask = float_padding_mask.masked_fill(key_padding_mask, float("-inf"))
            attended, _ = self.attn(
                query[valid_rows],
                memory_tokens[valid_rows],
                memory_tokens[valid_rows],
                key_padding_mask=key_padding_mask,
                attn_mask=attn_mask,
                need_weights=False,
            )
            out[valid_rows] = attended.to(out.dtype)
        return out, valid_rows

    def _apply_memory_type(self, memory_tokens, memory_type_ids):
        if memory_type_ids is None:
            return memory_tokens
        memory_type_ids = memory_type_ids.to(device=memory_tokens.device, dtype=torch.long)
        type_embed = self.memory_type_embed(memory_type_ids).to(memory_tokens.dtype)
        type_scale = self.memory_type_scale[memory_type_ids].to(memory_tokens.dtype)
        while type_embed.dim() < memory_tokens.dim():
            type_embed = type_embed.unsqueeze(0)
            type_scale = type_scale.unsqueeze(0)
        return memory_tokens * type_scale + type_embed

    def _store_type_gate_diagnostics(self, stage_gate):
        with torch.no_grad():
            detached = stage_gate.detach().float()
            self.last_type_gate_mean = detached.mean()
            for type_idx, type_name in enumerate(MEMORY_TYPE_NAMES[: self.num_memory_types]):
                setattr(self, f"last_type_gate_{type_name}_mean", detached[..., type_idx].mean())

    def _type_stage_gate(self, c, memory_tokens, memory_type_ids):
        if memory_type_ids is None:
            return None
        memory_type_ids = memory_type_ids.to(device=memory_tokens.device, dtype=torch.long)
        stage_gate = torch.sigmoid(self.memory_type_gate(c)).to(memory_tokens.dtype)
        self._store_type_gate_diagnostics(stage_gate)
        if memory_tokens.dim() == 4:
            batch_size, num_frames, num_tokens = memory_tokens.shape[:3]
            if memory_type_ids.dim() == 1:
                gather_ids = memory_type_ids.view(1, 1, num_tokens).expand(batch_size, num_frames, num_tokens)
            elif tuple(memory_type_ids.shape) == (batch_size, num_frames, num_tokens):
                gather_ids = memory_type_ids
            else:
                raise ValueError(
                    "rank-4 memory_type_ids must have shape (M,) or (B,T,M), "
                    f"got {tuple(memory_type_ids.shape)}"
                )
            return torch.gather(stage_gate, dim=-1, index=gather_ids)
        if memory_tokens.dim() == 3:
            batch_size, num_tokens = memory_tokens.shape[:2]
            if memory_type_ids.dim() != 1:
                raise ValueError("rank-3 memory_type_ids must have shape (M,)")
            gather_ids = memory_type_ids.view(1, 1, num_tokens).expand(batch_size, stage_gate.shape[1], num_tokens)
            return torch.gather(stage_gate, dim=-1, index=gather_ids).mean(dim=1)
        raise ValueError(f"memory_tokens must be rank 3 or 4, got rank {memory_tokens.dim()}")

    def _combine_memory_gate(self, memory_tokens, memory_token_gate, type_stage_gate):
        combined_gate = type_stage_gate
        if memory_token_gate is not None:
            if tuple(memory_token_gate.shape) != tuple(memory_tokens.shape[:-1]):
                raise ValueError(
                    f"memory_token_gate must have shape {tuple(memory_tokens.shape[:-1])}, "
                    f"got {tuple(memory_token_gate.shape)}"
                )
            stream_gate = memory_token_gate.to(device=memory_tokens.device, dtype=memory_tokens.dtype)
            combined_gate = stream_gate if combined_gate is None else combined_gate * stream_gate
        return combined_gate

    def _valid_mask(self, valid_rows, batch_size, num_frames, dtype, device):
        if valid_rows is None:
            return None
        valid_rows = valid_rows.to(device=device, dtype=dtype)
        if valid_rows.numel() == batch_size:
            return valid_rows.view(batch_size, 1, 1, 1, 1)
        if valid_rows.numel() == batch_size * num_frames:
            return rearrange(valid_rows, "(b t) -> b t", b=batch_size, t=num_frames)[:, :, None, None, None]
        raise ValueError(f"valid_rows has incompatible shape: {tuple(valid_rows.shape)}")

    def _gate_valid_mask(self, valid_rows, batch_size, num_frames, dtype, device):
        if valid_rows is None:
            return None
        valid_rows = valid_rows.to(device=device, dtype=dtype)
        if valid_rows.numel() == batch_size:
            return valid_rows.view(batch_size, 1, 1)
        if valid_rows.numel() == batch_size * num_frames:
            return rearrange(valid_rows, "(b t) -> b t", b=batch_size, t=num_frames)[:, :, None]
        raise ValueError(f"valid_rows has incompatible shape: {tuple(valid_rows.shape)}")

    def _residual_gate(self, residual_gate, batch_size, num_frames, dtype, device):
        if residual_gate is None:
            return None
        if not torch.is_tensor(residual_gate):
            return torch.tensor(float(residual_gate), dtype=dtype, device=device).view(1, 1, 1, 1, 1)
        gate_tensor = residual_gate.to(device=device, dtype=dtype)
        if gate_tensor.dim() == 0:
            gate_tensor = gate_tensor.view(1, 1, 1, 1, 1)
        elif gate_tensor.dim() == 1:
            if gate_tensor.numel() == batch_size:
                gate_tensor = gate_tensor.view(batch_size, 1, 1, 1, 1)
            elif gate_tensor.numel() == batch_size * num_frames:
                gate_tensor = rearrange(gate_tensor, "(b t) -> b t", b=batch_size, t=num_frames)[:, :, None, None, None]
            else:
                raise ValueError(f"residual_gate has incompatible shape: {tuple(gate_tensor.shape)}")
        elif gate_tensor.dim() == 2:
            if tuple(gate_tensor.shape) != (batch_size, num_frames):
                raise ValueError(f"residual_gate must have shape (B,T), got {tuple(gate_tensor.shape)}")
            gate_tensor = gate_tensor[:, :, None, None, None]
        elif gate_tensor.dim() == 3:
            if tuple(gate_tensor.shape[:2]) != (batch_size, num_frames):
                raise ValueError(f"residual_gate must start with (B,T), got {tuple(gate_tensor.shape)}")
            gate_tensor = gate_tensor[:, :, :, None, None]
        else:
            while gate_tensor.dim() < 5:
                gate_tensor = gate_tensor.unsqueeze(-1)
        return gate_tensor

    def _store_diagnostics(self, output, base, gate_msa, gate_mlp, valid_rows):
        with torch.no_grad():
            batch_size, num_frames = base.shape[:2]
            gate_values = torch.cat(
                [gate_msa.detach().float().abs(), gate_mlp.detach().float().abs()],
                dim=-1,
            )
            gate_mask = self._gate_valid_mask(
                valid_rows,
                batch_size,
                num_frames,
                dtype=gate_values.dtype,
                device=gate_values.device,
            )
            if gate_mask is not None:
                gate_values = gate_values * gate_mask
                self.last_valid_fraction = valid_rows.detach().float().mean()
                valid_count = (gate_mask.sum() * gate_values.shape[-1]).clamp_min(1.0)
                self.last_gate_mean = gate_values.sum() / valid_count
            else:
                self.last_valid_fraction = base.detach().new_tensor(1.0, dtype=torch.float32)
                self.last_gate_mean = gate_values.mean()

            delta_norm = (output.detach().float() - base.detach().float()).norm()
            base_norm = base.detach().float().norm()
            self.last_delta_ratio = delta_norm / (base_norm + 1e-6)

    def forward(
        self,
        x,
        c,
        memory_tokens,
        memory_token_mask=None,
        residual_base=None,
        return_delta=False,
        residual_gate=None,
        memory_type_ids=None,
        memory_token_gate=None,
    ):
        B, T, H, W, D = x.shape
        if residual_base is None:
            residual_base = x
        m_shift_msa, m_scale_msa, m_gate_msa, m_shift_mlp, m_scale_mlp, m_gate_mlp = (
            self.adaLN_modulation(c).chunk(6, dim=-1)
        )
        query_source = modulate(self.norm_q(x), m_shift_msa, m_scale_msa)
        type_stage_gate = self._type_stage_gate(c, memory_tokens, memory_type_ids)
        effective_token_gate = self._combine_memory_gate(memory_tokens, memory_token_gate, type_stage_gate)
        if memory_tokens.dim() == 3:
            query = rearrange(query_source, "b t h w d -> b (t h w) d")
            memory_tokens = self._apply_memory_type(self.norm_mem(memory_tokens), memory_type_ids)
            valid_rows = None
            if memory_token_mask is not None:
                if tuple(memory_token_mask.shape) != tuple(memory_tokens.shape[:2]):
                    raise ValueError(
                        f"legacy memory mask must have shape {tuple(memory_tokens.shape[:2])}, "
                        f"got {tuple(memory_token_mask.shape)}"
                    )
            out, valid_rows = self._attend(
                query,
                memory_tokens,
                memory_token_mask=memory_token_mask,
                memory_token_gate=effective_token_gate,
            )
            out = rearrange(out, "b (t h w) d -> b t h w d", t=T, h=H, w=W)
        elif memory_tokens.dim() == 4:
            assert memory_tokens.shape[:2] == (B, T), (
                f"per-frame memory tokens must have shape (B, T, M, D), got {tuple(memory_tokens.shape)}"
            )
            query = rearrange(query_source, "b t h w d -> (b t) (h w) d")
            memory_tokens = self._apply_memory_type(self.norm_mem(memory_tokens), memory_type_ids)
            memory_tokens = rearrange(memory_tokens, "b t m d -> (b t) m d")
            if effective_token_gate is not None:
                effective_token_gate = rearrange(effective_token_gate, "b t m -> (b t) m")
            valid_rows = None
            if memory_token_mask is not None:
                expected_mask_shape = (B, T, memory_tokens.shape[1])
                if tuple(memory_token_mask.shape) != expected_mask_shape:
                    raise ValueError(
                        f"per-frame memory mask must have shape {expected_mask_shape}, "
                        f"got {tuple(memory_token_mask.shape)}"
                    )
                memory_token_mask = rearrange(memory_token_mask.bool(), "b t m -> (b t) m")
            out, valid_rows = self._attend(
                query,
                memory_tokens,
                memory_token_mask=memory_token_mask,
                memory_token_gate=effective_token_gate,
            )
            out = rearrange(out, "(b t) (h w) d -> b t h w d", b=B, t=T, h=H, w=W)
        else:
            raise ValueError(f"memory_tokens must be rank 3 or 4, got rank {memory_tokens.dim()}")

        valid_mask = self._valid_mask(valid_rows, B, T, dtype=out.dtype, device=out.device)
        residual_gate_tensor = self._residual_gate(residual_gate, B, T, dtype=out.dtype, device=out.device)
        attn_delta = gate(out, m_gate_msa)
        if valid_mask is not None:
            attn_delta = attn_delta * valid_mask
        if residual_gate_tensor is not None:
            attn_delta = attn_delta * residual_gate_tensor
        output = residual_base + attn_delta

        mlp_delta = gate(self.mlp(modulate(self.norm_mlp(output), m_shift_mlp, m_scale_mlp)), m_gate_mlp)
        if valid_mask is not None:
            mlp_delta = mlp_delta * valid_mask
        if residual_gate_tensor is not None:
            mlp_delta = mlp_delta * residual_gate_tensor
        output = output + mlp_delta
        self._store_diagnostics(output, residual_base, m_gate_msa, m_gate_mlp, valid_rows)
        if return_delta:
            return attn_delta + mlp_delta
        return output

class SpatioTemporalDiTBlock(nn.Module):
    def __init__(
        self,
        hidden_size,
        num_heads,
        reference_length,
        mlp_ratio=4.0,
        is_causal=True,
        spatial_rotary_emb: Optional[RotaryEmbedding] = None,
        temporal_rotary_emb: Optional[RotaryEmbedding] = None,
        use_memory_token_cross_attention=False,
        ref_mode='sequential'
    ):
        super().__init__()
        self.is_causal = is_causal
        mlp_hidden_dim = int(hidden_size * mlp_ratio)
        approx_gelu = lambda: nn.GELU(approximate="tanh")

        self.s_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.s_attn = SpatialAxialAttention(
            hidden_size,
            heads=num_heads,
            dim_head=hidden_size // num_heads,
            rotary_emb=spatial_rotary_emb
        )
        self.s_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.s_mlp = Mlp(
            in_features=hidden_size,
            hidden_features=mlp_hidden_dim,
            act_layer=approx_gelu,
            drop=0,
        )
        self.s_adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))

        self.t_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.t_attn = TemporalAxialAttention(
            hidden_size,
            heads=num_heads,
            dim_head=hidden_size // num_heads,
            is_causal=is_causal,
            rotary_emb=temporal_rotary_emb,
            reference_length=reference_length
        )
        self.t_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.t_mlp = Mlp(
            in_features=hidden_size,
            hidden_features=mlp_hidden_dim,
            act_layer=approx_gelu,
            drop=0,
        )
        self.t_adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))

        self.reference_length = reference_length
        self.use_memory_token_cross_attention = use_memory_token_cross_attention
        if self.use_memory_token_cross_attention:
            self.memory_token_cross_attn = MemoryTokenCrossAttention(hidden_size, num_heads, mlp_ratio=mlp_ratio)

        self.ref_mode = ref_mode

        if self.ref_mode == 'parallel':
            self.parallel_map = nn.Linear(hidden_size, hidden_size)

    def _expand_memory_stream(self, tokens, mask, stream_gate, type_idx, batch_size, num_frames):
        if tokens is None or tokens.shape[-2] == 0:
            return None
        if tokens.dim() == 3:
            if tokens.shape[0] != batch_size:
                raise ValueError(f"rank-3 memory tokens must start with B={batch_size}, got {tuple(tokens.shape)}")
            tokens = tokens[:, None].expand(-1, num_frames, -1, -1)
            if mask is None:
                mask = torch.ones(tokens.shape[:3], device=tokens.device, dtype=torch.bool)
            elif mask.dim() == 2:
                mask = mask[:, None].expand(-1, num_frames, -1)
            elif mask.dim() != 3:
                raise ValueError(f"rank-3 stream mask must have rank 2 or 3, got {tuple(mask.shape)}")
        elif tokens.dim() == 4:
            if tuple(tokens.shape[:2]) != (batch_size, num_frames):
                raise ValueError(
                    f"rank-4 memory tokens must start with (B,T)={(batch_size, num_frames)}, "
                    f"got {tuple(tokens.shape)}"
                )
            if mask is None:
                mask = torch.ones(tokens.shape[:3], device=tokens.device, dtype=torch.bool)
            elif mask.dim() != 3:
                raise ValueError(f"rank-4 stream mask must have rank 3, got {tuple(mask.shape)}")
        else:
            raise ValueError(f"memory stream tokens must be rank 3 or 4, got rank {tokens.dim()}")
        if tuple(mask.shape) != tuple(tokens.shape[:3]):
            raise ValueError(f"memory stream mask must have shape {tuple(tokens.shape[:3])}, got {tuple(mask.shape)}")
        gate_tensor = self._expand_memory_stream_gate(stream_gate, tokens)
        type_ids = torch.full((tokens.shape[2],), int(type_idx), device=tokens.device, dtype=torch.long)
        return tokens, mask.to(device=tokens.device, dtype=torch.bool), gate_tensor, type_ids

    def _expand_memory_stream_gate(self, stream_gate, tokens):
        batch_size, num_frames, num_tokens = tokens.shape[:3]
        if stream_gate is None:
            return torch.ones(tokens.shape[:3], device=tokens.device, dtype=tokens.dtype)
        if not torch.is_tensor(stream_gate):
            return torch.full(tokens.shape[:3], float(stream_gate), device=tokens.device, dtype=tokens.dtype)
        gate_tensor = stream_gate.to(device=tokens.device, dtype=tokens.dtype)
        if gate_tensor.dim() == 0:
            return gate_tensor.view(1, 1, 1).expand(batch_size, num_frames, num_tokens)
        if gate_tensor.dim() == 1:
            if gate_tensor.numel() != batch_size:
                raise ValueError(f"rank-1 memory gate must have B={batch_size} values, got {tuple(gate_tensor.shape)}")
            return gate_tensor.view(batch_size, 1, 1).expand(batch_size, num_frames, num_tokens)
        if gate_tensor.dim() == 2:
            if tuple(gate_tensor.shape) == (batch_size, num_frames):
                return gate_tensor[:, :, None].expand(batch_size, num_frames, num_tokens)
            if tuple(gate_tensor.shape) == (batch_size, num_tokens):
                return gate_tensor[:, None, :].expand(batch_size, num_frames, num_tokens)
            raise ValueError(
                f"rank-2 memory gate must have shape (B,T) or (B,M), got {tuple(gate_tensor.shape)}"
            )
        if gate_tensor.dim() == 3:
            if tuple(gate_tensor.shape) == (batch_size, num_frames, 1):
                return gate_tensor.expand(batch_size, num_frames, num_tokens)
            if tuple(gate_tensor.shape) == (batch_size, num_frames, num_tokens):
                return gate_tensor
            raise ValueError(
                f"rank-3 memory gate must have shape (B,T,1) or (B,T,M), got {tuple(gate_tensor.shape)}"
            )
        raise ValueError(f"memory gate rank must be <=3, got rank {gate_tensor.dim()}")

    def _pack_typed_memory_streams(
        self,
        batch_size,
        num_frames,
        memory_tokens=None,
        memory_token_mask=None,
        memory_dynamic_tokens=None,
        memory_dynamic_mask=None,
        memory_retrieval_tokens=None,
        memory_retrieval_mask=None,
        memory_anchor_gate=None,
        memory_dynamic_gate=None,
        memory_retrieval_gate=None,
    ):
        streams = []
        for tokens, mask, stream_gate, type_idx in (
            (memory_tokens, memory_token_mask, memory_anchor_gate, MEMORY_TYPE_ANCHOR),
            (memory_dynamic_tokens, memory_dynamic_mask, memory_dynamic_gate, MEMORY_TYPE_DYNAMIC),
            (memory_retrieval_tokens, memory_retrieval_mask, memory_retrieval_gate, MEMORY_TYPE_REVISIT),
        ):
            expanded = self._expand_memory_stream(tokens, mask, stream_gate, type_idx, batch_size, num_frames)
            if expanded is not None:
                streams.append(expanded)
        if not streams:
            return None
        packed_tokens = torch.cat([item[0] for item in streams], dim=2)
        packed_mask = torch.cat([item[1] for item in streams], dim=2)
        packed_gate = torch.cat([item[2] for item in streams], dim=2)
        packed_type_ids = torch.cat([item[3] for item in streams], dim=0)
        valid_gate = packed_gate.masked_fill(~packed_mask, 0)
        residual_gate = valid_gate.max(dim=2).values
        return packed_tokens, packed_mask, packed_gate, packed_type_ids, residual_gate

    def forward(self, x, c, current_frame=None, timestep=None, is_last_block=False, 
        pose_cond=None, mode="training", c_action_cond=None, reference_length=None,
        memory_tokens=None, memory_token_mask=None, memory_dynamic_tokens=None, memory_dynamic_mask=None,
        memory_retrieval_tokens=None, memory_retrieval_mask=None, memory_anchor_gate=None,
        memory_dynamic_gate=None, memory_retrieval_gate=None):
        B, T, H, W, D = x.shape

        # spatial block
        
        s_shift_msa, s_scale_msa, s_gate_msa, s_shift_mlp, s_scale_mlp, s_gate_mlp = self.s_adaLN_modulation(c).chunk(6, dim=-1)
        x = x + gate(self.s_attn(modulate(self.s_norm1(x), s_shift_msa, s_scale_msa)), s_gate_msa)
        x = x + gate(self.s_mlp(modulate(self.s_norm2(x), s_shift_mlp, s_scale_mlp)), s_gate_mlp)

        # temporal block
        if c_action_cond is not None:
            t_shift_msa, t_scale_msa, t_gate_msa, t_shift_mlp, t_scale_mlp, t_gate_mlp = self.t_adaLN_modulation(c_action_cond).chunk(6, dim=-1)
        else:
            t_shift_msa, t_scale_msa, t_gate_msa, t_shift_mlp, t_scale_mlp, t_gate_mlp = self.t_adaLN_modulation(c).chunk(6, dim=-1)
        
        x_t = x + gate(self.t_attn(modulate(self.t_norm1(x), t_shift_msa, t_scale_msa)), t_gate_msa)
        x_t = x_t + gate(self.t_mlp(modulate(self.t_norm2(x_t), t_shift_mlp, t_scale_mlp)), t_gate_mlp)

        if self.ref_mode == 'sequential':
            x = x_t

        if self.use_memory_token_cross_attention:
            memory_base = x
            packed_memory = self._pack_typed_memory_streams(
                B,
                T,
                memory_tokens=memory_tokens,
                memory_token_mask=memory_token_mask,
                memory_dynamic_tokens=memory_dynamic_tokens,
                memory_dynamic_mask=memory_dynamic_mask,
                memory_retrieval_tokens=memory_retrieval_tokens,
                memory_retrieval_mask=memory_retrieval_mask,
                memory_anchor_gate=memory_anchor_gate,
                memory_dynamic_gate=memory_dynamic_gate,
                memory_retrieval_gate=memory_retrieval_gate,
            )
            if packed_memory is not None:
                packed_tokens, packed_mask, packed_gate, packed_type_ids, residual_gate = packed_memory
                x = self.memory_token_cross_attn(
                    memory_base,
                    c,
                    packed_tokens,
                    packed_mask,
                    residual_gate=residual_gate,
                    memory_type_ids=packed_type_ids,
                    memory_token_gate=packed_gate,
                )

        if self.ref_mode == 'parallel':
            x = x_t + self.parallel_map(x)

        return x


class DiT(nn.Module):
    """
    Diffusion model with a Transformer backbone.
    """

    def __init__(
        self,
        input_h=18,
        input_w=32,
        patch_size=2,
        in_channels=16,
        hidden_size=1024,
        depth=12,
        num_heads=16,
        mlp_ratio=4.0,
        action_cond_dim=25,
        max_frames=32,
        reference_length=8,
        memory_token_cross_attention=False,
        memory_cross_attn_layers=None,
        ref_mode='sequential'
    ):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = in_channels
        self.patch_size = patch_size
        self.num_heads = num_heads
        self.max_frames = max_frames

        self.x_embedder = PatchEmbed(input_h, input_w, patch_size, in_channels, hidden_size, flatten=False)
        self.t_embedder = TimestepEmbedder(hidden_size)

        self.spatial_rotary_emb = RotaryEmbedding(dim=hidden_size // num_heads // 2, freqs_for="pixel", max_freq=256)
        self.temporal_rotary_emb = RotaryEmbedding(dim=hidden_size // num_heads)

        self.external_cond = nn.Linear(action_cond_dim, hidden_size) if action_cond_dim > 0 else nn.Identity()
        if memory_cross_attn_layers is None:
            memory_cross_attn_layer_set = None
        else:
            memory_cross_attn_layer_set = {int(layer_idx) for layer_idx in memory_cross_attn_layers}
            invalid_layers = sorted(
                layer_idx for layer_idx in memory_cross_attn_layer_set if layer_idx < 0 or layer_idx >= depth
            )
            if invalid_layers:
                raise ValueError(
                    f"memory_cross_attn_layers contains invalid indices {invalid_layers} for depth={depth}"
                )

        self.blocks = nn.ModuleList(
            [
                SpatioTemporalDiTBlock(
                    hidden_size,
                    num_heads,
                    mlp_ratio=mlp_ratio,
                    is_causal=True,
                    reference_length=reference_length,
                    spatial_rotary_emb=self.spatial_rotary_emb,
                    temporal_rotary_emb=self.temporal_rotary_emb,
                    use_memory_token_cross_attention=memory_token_cross_attention
                    and (memory_cross_attn_layer_set is None or block_idx in memory_cross_attn_layer_set),
                    ref_mode=ref_mode
                )
                for block_idx in range(depth)
            ]
        )
        self.memory_token_cross_attention = memory_token_cross_attention
        self.memory_cross_attn_layers = (
            None if memory_cross_attn_layer_set is None else tuple(sorted(memory_cross_attn_layer_set))
        )
        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
        self.initialize_weights()

    def initialize_weights(self):
        # Initialize transformer layers:
        def _basic_init(module):
            if isinstance(module, nn.Linear):
                torch.nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)

        self.apply(_basic_init)

        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
        w = self.x_embedder.proj.weight.data
        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
        nn.init.constant_(self.x_embedder.proj.bias, 0)

        # Initialize timestep embedding MLP:
        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)

        # Zero-out adaLN modulation layers in DiT blocks:
        for block in self.blocks:
            nn.init.constant_(block.s_adaLN_modulation[-1].weight, 0)
            nn.init.constant_(block.s_adaLN_modulation[-1].bias, 0)
            nn.init.constant_(block.t_adaLN_modulation[-1].weight, 0)
            nn.init.constant_(block.t_adaLN_modulation[-1].bias, 0)

        # Zero-out output layers:
        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
        nn.init.constant_(self.final_layer.linear.weight, 0)
        nn.init.constant_(self.final_layer.linear.bias, 0)

        if self.memory_token_cross_attention:
            for block in self.blocks:
                memory_adapter = getattr(block, "memory_token_cross_attn", None)
                if memory_adapter is not None:
                    memory_adapter.reset_identity_init()

    def memory_adapter_delta_diagnostics(self):
        diagnostics = {}
        ratios = []
        type_gate_values = {type_name: [] for type_name in MEMORY_TYPE_NAMES}
        shared_type_gate_values = []
        for block in self.blocks:
            adapter = getattr(block, "memory_token_cross_attn", None)
            if adapter is None:
                continue
            ratio = getattr(adapter, "last_delta_ratio", None)
            if ratio is not None:
                ratios.append(torch.as_tensor(ratio).detach().float())
            type_gate = getattr(adapter, "last_type_gate_mean", None)
            if type_gate is not None:
                shared_type_gate_values.append(torch.as_tensor(type_gate).detach().float())
            for type_name in MEMORY_TYPE_NAMES:
                value = getattr(adapter, f"last_type_gate_{type_name}_mean", None)
                if value is not None:
                    type_gate_values[type_name].append(torch.as_tensor(value).detach().float())
        if ratios:
            values = torch.stack(ratios)
            diagnostics["memory_adapter_delta_ratio_max"] = float(values.max().item())
            diagnostics["memory_adapter_delta_ratio_mean"] = float(values.mean().item())
        if shared_type_gate_values:
            values = torch.stack(shared_type_gate_values)
            diagnostics["memory_adapter_type_gate_mean"] = float(values.mean().item())
        for type_name, values_list in type_gate_values.items():
            if values_list:
                values = torch.stack(values_list)
                diagnostics[f"memory_adapter_type_gate_{type_name}_mean"] = float(values.mean().item())
        return diagnostics

    def unpatchify(self, x):
        """
        x: (N, H, W, patch_size**2 * C)
        imgs: (N, H, W, C)
        """
        c = self.out_channels
        p = self.x_embedder.patch_size[0]
        h = x.shape[1]
        w = x.shape[2]

        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
        x = torch.einsum("nhwpqc->nchpwq", x)
        imgs = x.reshape(shape=(x.shape[0], c, h * p, w * p))
        return imgs

    def forward(
        self,
        x,
        t,
        action_cond=None,
        pose_cond=None,
        current_frame=None,
        mode=None,
        reference_length=None,
        frame_idx=None,
        memory_tokens=None,
        memory_token_mask=None,
        memory_dynamic_tokens=None,
        memory_dynamic_mask=None,
        memory_retrieval_tokens=None,
        memory_retrieval_mask=None,
        memory_anchor_gate=None,
        memory_dynamic_gate=None,
        memory_retrieval_gate=None,
    ):
        """
        Forward pass of DiT.
        x: (B, T, C, H, W) tensor of spatial inputs (images or latent representations of images)
        t: (B, T,) tensor of diffusion timesteps
        """

        B, T, C, H, W = x.shape

        # add spatial embeddings
        x = rearrange(x, "b t c h w -> (b t) c h w")

        x = self.x_embedder(x)  # (B*T, C, H, W) -> (B*T, H/2, W/2, D) , C = 16, D = d_model
        # restore shape
        x = rearrange(x, "(b t) h w d -> b t h w d", t=T)
        # embed noise steps
        t = rearrange(t, "b t -> (b t)")

        c_t = self.t_embedder(t)  # (N, D)
        c = c_t.clone()
        c = rearrange(c, "(b t) d -> b t d", t=T)

        if torch.is_tensor(action_cond):
            c_action_cond = c + self.external_cond(action_cond)
        else:
            c_action_cond = None

        for i, block in enumerate(self.blocks):
            x = block(x, c, current_frame=current_frame, timestep=t, is_last_block= (i+1 == len(self.blocks)), 
                mode=mode, c_action_cond=c_action_cond, reference_length=reference_length,
                memory_tokens=memory_tokens, memory_token_mask=memory_token_mask,
                memory_dynamic_tokens=memory_dynamic_tokens, memory_dynamic_mask=memory_dynamic_mask,
                memory_retrieval_tokens=memory_retrieval_tokens, memory_retrieval_mask=memory_retrieval_mask,
                memory_anchor_gate=memory_anchor_gate, memory_dynamic_gate=memory_dynamic_gate,
                memory_retrieval_gate=memory_retrieval_gate)  # (N, T, H, W, D)
        x = self.final_layer(x, c) # (N, T, H, W, patch_size ** 2 * out_channels)
        # unpatchify
        x = rearrange(x, "b t h w d -> (b t) h w d")
        x = self.unpatchify(x)  # (N, out_channels, H, W)
        x = rearrange(x, "(b t) c h w -> b t c h w", t=T)
        return x


def DiT_S_2(
    action_cond_dim,
    reference_length,
    ref_mode,
    memory_token_cross_attention=False,
    memory_cross_attn_layers=None,
):
    return DiT(
        patch_size=2,
        hidden_size=1024,
        depth=16,
        num_heads=16,
        action_cond_dim=action_cond_dim,
        reference_length=reference_length,
        memory_token_cross_attention=memory_token_cross_attention,
        memory_cross_attn_layers=memory_cross_attn_layers,
        ref_mode=ref_mode
    )


DiT_models = {"DiT-S/2": DiT_S_2}