Spaces:

kaupane
/

Inko

Running on Zero

Inko / src /models /dit_blocks.py
Restore caller SDPA priority after masked attention
444566f verified about 11 hours ago
37.2 kB
"""Attention, positional features and native Inko transformer blocks."""

from typing import Tuple, Optional, Sequence
import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from .varlen_attention import packed_attention

# Optional SDPA backend restriction for masked (padded-text) attention.
_SDPA_BACKENDS = None


class RMSNorm(nn.RMSNorm):
    """Fuse Ascend normalization while retaining FP32 arithmetic and gains."""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (
            x.device.type != "npu"
            or self.weight is None
            or len(self.normalized_shape) != 1
            or x.dtype not in (torch.float16, torch.bfloat16, torch.float32)
        ):
            return super().forward(x)
        import torch_npu

        # PyTorch's mixed-dtype RMSNorm promotes the input to FP32 and applies
        # the learned gain before casting back. Casting the gain to BF16 just
        # to use a fused kernel would change that computation.
        eps = self.eps if self.eps is not None else torch.finfo(x.dtype).eps
        y, _ = torch_npu.npu_rms_norm(x.float(), self.weight.float(), epsilon=eps)
        return y.to(x.dtype)


def sdpa_with_pad_mask(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    attn_mask: Optional[torch.Tensor],
) -> torch.Tensor:
    """SDPA with an optional bool key-padding mask of shape [B, 1, 1, S_k].
    A bool mask forces PyTorch off the flash path; the math fallback then
    materializes B*H*S^2 attention weights per layer (saved for backward),
    which at seq ~2K is tens of GiB per layer and OOMs long-caption batches.
    We convert the mask to an additive bias in q.dtype and use fused kernels:
    cuDNN first in inference mode, memory-efficient attention otherwise.
    mask=None keeps the default (flash) path untouched.
    """
    if attn_mask is None:
        return F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)

    return sdpa_with_bias(q, k, v, pad_bias_from_mask(attn_mask, q.dtype))


def set_sdpa_backends(names: Optional[Sequence[str]]) -> None:
    """Restrict SDPA backends for masked (padded-text) attention.
    ``None`` uses cuDNN first in inference mode and memory-efficient attention
    otherwise. An explicit list restricts attention to those backends, in order.
    Training keeps memory-efficient attention: on a
    single 4090 at 256p (micro-batch 16, 128-192 text tokens) cuDNN attention
    is ~30% slower than memory-efficient on the forward+backward training
    step, despite a faster forward-only kernel with an additive pad bias.
    """
    global _SDPA_BACKENDS
    if names is None:
        _SDPA_BACKENDS = None
        return
    from torch.nn.attention import SDPBackend

    resolved = []
    for name in names:
        member = getattr(SDPBackend, name.strip().upper(), None)
        if member is None:
            raise ValueError(f"unknown SDPA backend {name!r}")
        resolved.append(member)
    _SDPA_BACKENDS = tuple(resolved)


def _sdpa_backends():
    if _SDPA_BACKENDS is not None:
        return list(_SDPA_BACKENDS)
    from torch.nn.attention import SDPBackend

    if torch.is_inference_mode_enabled():
        return [
            SDPBackend.CUDNN_ATTENTION,
            SDPBackend.EFFICIENT_ATTENTION,
            SDPBackend.MATH,
        ]
    return [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]


def sdpa_with_bias(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    bias: torch.Tensor,
) -> torch.Tensor:
    """SDPA with a precomputed additive bias of shape [B, 1, 1, S_k]."""
    if q.device.type == "npu":
        # torch_npu's F.sdpa dispatch with a float additive bias materializes
        # per-head S^2 bias + P matrix per layer (measured: +2.8 GiB at
        # B8/S2304), which OOMs long-caption batches. npu_fusion_attention with
        # a bool (B,1,S,S) drop-mask (True = drop) is true flash attention:
        # +0.06 GiB saved for backward, maxdiff 2e-3 vs the float-bias path at
        # S=2304 bf16. Two gotchas, both measured: the scale argument defaults
        # to 1.0 (NOT 1/sqrt(D)) so it must be passed explicitly, and the mask
        # must be full (B,1,S,S) - aclnn rejects broadcastable (B,1,1,S).
        import torch_npu

        b, h, s_q, d = q.shape
        s_k = k.shape[2]
        # bias == -inf: identical to torch.isneginf for float tensors, but
        # aten.eq.Scalar has a torchair GE converter while aten.isneginf does
        # not (ascend-ta-bench3 NotImplementedError).
        drop = bias == float("-inf")
        mask = drop.expand(b, 1, s_q, s_k).contiguous()
        return torch_npu.npu_fusion_attention(
            q, k, v, h, "BNSD", atten_mask=mask, keep_prob=1.0, scale=d**-0.5
        )[0]
    from torch.nn.attention import sdpa_kernel

    backends = _sdpa_backends()
    if q.device.type != "cuda":
        from torch.nn.attention import SDPBackend

        backends = [SDPBackend.MATH]
    # sdpa_kernel only saves enabled backends. Preserve disabled backends'
    # positions too, so nested contexts do not reorder the caller's policy.
    priority = torch._C._get_sdp_priority_order()
    try:
        with sdpa_kernel(backends, set_priority=True):
            return F.scaled_dot_product_attention(q, k, v, attn_mask=bias, dropout_p=0.0)
    finally:
        torch._C._set_sdp_priority_order(priority)


def pad_bias_from_mask(attn_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
    """Convert a bool keep-mask [B, 1, 1, S_k] into an additive SDPA bias."""
    bias = torch.zeros(attn_mask.shape, dtype=dtype, device=attn_mask.device)
    return bias.masked_fill(~attn_mask, float("-inf"))


class TimestepEmbeddings(nn.Module):
    """Sinusoidal features; scaling applies here, never to the flow path."""

    def __init__(self, hidden_size: int):
        super().__init__()
        self.hidden_size = hidden_size
        self.register_buffer("factor", torch.tensor(1000, dtype=torch.int64))

    def _load_from_state_dict(
        self,
        state_dict,
        prefix,
        local_metadata,
        strict,
        missing_keys,
        unexpected_keys,
        error_msgs,
    ):
        factor = state_dict.get(prefix + "factor")
        if factor is None or factor.item() != 1000:
            error_msgs.append(f"{prefix}expected timestep feature factor 1000")
        super()._load_from_state_dict(
            state_dict,
            prefix,
            local_metadata,
            strict,
            missing_keys,
            unexpected_keys,
            error_msgs,
        )

    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
        """
        Args:
            timesteps: (B,)
        Returns:
            embedding: (B, N)
        """
        half = self.hidden_size // 2
        exponent = (
            -math.log(10000)
            * torch.arange(
                start=0, end=half, dtype=torch.float32, device=timesteps.device
            )
            / half
        )

        emb = torch.exp(exponent)
        times = timesteps.float() * 1000
        emb = times.unsqueeze(1) * emb.unsqueeze(0)

        return torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)


def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
    """
    Apply rotary embeddings
    Args:
        x: Input tensor [B, S, H, D]
        freqs_cis: Complex frequencies [S, D/2] or [B, S, D/2]
    Returns:
        Tensor with rotary embeddings applied [B, S, H, D]
    """
    # Reshape x to [B, S, H, D/2, 2] and view as complex
    x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))

    # Reshape for broadcasting: [S, D/2] -> [1, S, 1,  D/2]
    if freqs_cis.ndim == 2:
        freqs_cis = freqs_cis.unsqueeze(0)
    freqs_cis = freqs_cis.unsqueeze(2)

    # Apply rotation
    x_rotated = x_complex * freqs_cis

    # View as real and flatten the last two dimensions
    x_out = torch.view_as_real(x_rotated).flatten(3)

    return x_out.type_as(x)


def apply_rotary_emb_real(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
    """Equivalent FP32 rotation expressed without complex multiplication.
    This lets Inductor fuse surrounding casts/layout operations. Compiled
    rounding may differ; the complex implementation remains the default.
    """
    pairs = x.float().reshape(*x.shape[:-1], -1, 2)
    real, imag = pairs.unbind(-1)
    freqs_real = torch.view_as_real(freqs_cis)
    if freqs_real.ndim == 3:
        freqs_real = freqs_real.unsqueeze(0)
    cosine, sine = freqs_real.unsqueeze(2).unbind(-1)
    rotated = torch.stack(
        (real * cosine - imag * sine, real * sine + imag * cosine), dim=-1
    )
    return rotated.flatten(3).type_as(x)


def apply_rotary_emb_realfreq(
    x: torch.Tensor, freqs_real: torch.Tensor
) -> torch.Tensor:
    """apply_rotary_emb_real taking an already-real frequency tensor.
    freqs_real is [S, D/2, 2] or [B, S, D/2, 2] float32 (cos, sin pairs), the view_as_real of
    the complex table, materialized OUTSIDE compiled regions so no
    complex-typed tensor crosses the graph boundary (torchair/GE has no
    complex dtype support).
    """
    pairs = x.float().reshape(*x.shape[:-1], -1, 2)
    real, imag = pairs.unbind(-1)
    if freqs_real.ndim == 3:
        freqs_real = freqs_real.unsqueeze(0)
    cosine, sine = freqs_real.unsqueeze(2).unbind(-1)
    rotated = torch.stack(
        (real * cosine - imag * sine, real * sine + imag * cosine), dim=-1
    )
    return rotated.flatten(3).type_as(x)


def set_real_rope(model: nn.Module, enabled: bool) -> None:
    """Set an instance-local execution policy before compiling the model.
    This is not checkpoint state: loading weights must not override the launch
    policy. Other live models and the complex frequency tables are untouched.
    """
    if type(enabled) is not bool:
        raise ValueError("real RoPE policy must be boolean")
    for module in model.modules():
        if isinstance(module, (DoubleStreamAttention, SingleStreamAttention)):
            module.real_rope = enabled


class MSRoPE(nn.Module):
    """Image/text RoPE with automatic Dy-YaRN inference scaling.
    The 896p training buckets reach 1184 pixels on either axis, or 74 patches
    at 16 pixels per patch. Training always uses the original frequencies.
    DyPE follows https://arxiv.org/html/2510.20766v2, equations 9 and 13–17.
    Both axes share one scaling factor to preserve the canvas aspect ratio
    in phase space.
    """

    def __init__(
        self,
        theta: int = 10000,
        axes_dim: list = [64, 64],
        reference_hw: tuple[int, int] = (74, 74),
        centered: bool = False,
        decay_exponent: float = 2.0,
        yarn_alpha: float = 1.0,
        yarn_beta: float = 32.0,
    ):
        super().__init__()
        self.theta = theta
        self.axes_dim = axes_dim
        self.reference_hw = reference_hw
        self.decay_exponent = decay_exponent
        self.yarn_alpha = yarn_alpha
        self.yarn_beta = yarn_beta
        # centered: image coords symmetric around 0 ([-(n-1)/2, (n-1)/2]) instead of
        # 0-based from the corner. Text starts at max(height, width) on both axes.
        self.centered = centered

        # Precompute frequency tables for both spatial axes
        pos_index = torch.arange(2560)  # Initial table capacity; expands as needed.
        # Store as real to avoid safetensors issues with complex numbers
        self.register_buffer(
            "pos_freqs", torch.view_as_real(self._build_frequency_table(pos_index))
        )

        # Retain only the latest geometry/mode on each instance.
        self._cache = {}

    def needs_time(self, img_hw: tuple[int, int]) -> bool:
        return not self.training and any(
            size > reference for size, reference in zip(img_hw, self.reference_hw)
        )

    def _apply(self, fn, recurse=True):
        super()._apply(fn, recurse=recurse)
        if self.pos_freqs.dtype != torch.float32:
            # Rebuild instead of promoting already-rounded BF16/FP16 rotations.
            self.pos_freqs = torch.view_as_real(
                self._build_frequency_table(torch.arange(self.pos_freqs.shape[0]))
            ).to(self.pos_freqs.device)
        self._cache.clear()
        return self

    def rope_params(
        self,
        index: torch.Tensor,
        dim: int,
        scaling_factor: float = 1.0,
        *,
        reference: int,
        kappa: float = 1.0,
    ) -> torch.Tensor:
        """
        Generate RoPE parameters for given dimensions
        Args:
            index: Position indices [S]
            dim: Embedding dimension for this axis
            scaling_factor: Spatial extension along this axis
        Returns:
            Complex frequency tensor [S, dim/2]
        """
        assert dim % 2 == 0

        # Compute frequency components: 1/ theta^(2i/dim)
        powers = torch.arange(0, dim, 2, dtype=torch.float32, device=index.device).div(
            dim
        )
        inverse_freqs = 1.0 / torch.pow(self.theta, powers)
        if scaling_factor > 1.0 and kappa > 0.0:
            rotations = reference * inverse_freqs / (2 * math.pi)
            ramp = (
                (rotations - self.yarn_alpha * kappa)
                / ((self.yarn_beta - self.yarn_alpha) * kappa)
            ).clamp(0, 1)
            inverse_freqs = inverse_freqs * (ramp + (1 - ramp) / scaling_factor)
        freqs = torch.outer(index.float(), inverse_freqs)

        # Convert to complex polar form: e^(i*freqs)
        return torch.polar(torch.ones_like(freqs), freqs)

    def _build_frequency_table(
        self,
        pos_index: torch.Tensor,
        scaling_factors: tuple[float, float] = (1.0, 1.0),
        *,
        kappa: float = 1.0,
    ) -> torch.Tensor:
        """Build frequency table for both height and width axes"""
        height_freqs = self.rope_params(
            pos_index,
            self.axes_dim[0],
            scaling_factors[0],
            kappa=kappa,
            reference=self.reference_hw[0],
        )
        width_freqs = self.rope_params(
            pos_index,
            self.axes_dim[1],
            scaling_factors[1],
            kappa=kappa,
            reference=self.reference_hw[1],
        )

        # Concatenate height and width frequencies: [S, (height_dim+width_dim)/2]
        return torch.cat([height_freqs, width_freqs], dim=1)

    def forward(
        self,
        img_hw: tuple[int, int],
        txt_seq_len: int,
        device,
        *,
        time: tuple[float, ...] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Compute frequency tensors for image and text sequences
        Args:
            img_hw: Image patch-grid dimensions (height, width)
            txt_seq_len: Length of text sequence
            device: Device for tensor placement
            time: Flow time per image (0 noise, 1 image); required for active DyPE
        Returns:
            Image/text frequencies, with a batch dimension when times differ
        """
        height, width = img_hw
        device = torch.device(device)
        scaling_factors = (1.0, 1.0)
        if not self.training:
            shared = max(
                max(1.0, size / reference)
                for size, reference in zip(img_hw, self.reference_hw)
            )
            scaling_factors = (shared, shared)
        scaled = scaling_factors != (1.0, 1.0)
        kappas = (1.0,)
        if self.needs_time(img_hw):
            if time is None:
                raise ValueError("dynamic RoPE requires the current flow time")
            # Inko runs from t=0 noise to t=1 image, opposite to the paper.
            kappas = tuple((1 - t) ** self.decay_exponent for t in time)
            if len(set(kappas)) == 1:
                kappas = kappas[:1]

        # Keep the checkpoint's base table unscaled.
        max_needed = max(height, width) + txt_seq_len
        current_max = self.pos_freqs.shape[0]
        if max_needed > current_max:
            new_max = max(max_needed, current_max * 2)
            self.pos_freqs = torch.view_as_real(
                self._build_frequency_table(torch.arange(new_max))
            )
        if self.pos_freqs.device != device:
            self.pos_freqs = self.pos_freqs.to(device)

        cache_key = (
            height,
            width,
            txt_seq_len,
            scaling_factors,
            kappas,
            self.reference_hw,
            self.yarn_alpha,
            self.yarn_beta,
            self.centered,
            device,
        )
        if cache_key in self._cache:
            return self._cache[cache_key]

        images, texts = [], []
        max_img_pos = max(height, width)
        for kappa in kappas:
            images.append(
                self._compute_image_freqs(
                    height, width, scaling_factors, kappa=kappa
                ).to(device)
            )
            # Share the image basis without relocating the text diagonal.
            texts.append(
                torch.view_as_complex(self.pos_freqs)[max_img_pos:max_needed]
                if not scaled
                else self._build_frequency_table(
                    torch.arange(max_img_pos, max_needed),
                    scaling_factors,
                    kappa=kappa,
                ).to(device)
            )
        img_freqs = images[0] if len(images) == 1 else torch.stack(images)
        txt_freqs = texts[0] if len(texts) == 1 else torch.stack(texts)
        if scaled:
            # Eq. 9 multiplies logits by tau. Multiplying BOTH Q/K rotations by
            # sqrt(tau) applies it in every attention backend and RoPE layout.
            amplitude = math.sqrt(1 + 0.1 * math.log(max(scaling_factors)))
            img_freqs = img_freqs * amplitude
            txt_freqs = txt_freqs * amplitude

        self._cache.clear()
        self._cache[cache_key] = (img_freqs, txt_freqs)
        return img_freqs, txt_freqs

    def prepare_freqs(
        self,
        img_hw: tuple[int, int],
        txt_seq_len: int,
        device,
        *,
        time: tuple[float, ...] | None = None,
    ) -> torch.Tensor:
        """Concatenated [S_img + S_txt, D/2] freqs, optionally batched.
        Identical to applying ``forward()`` and concatenating the two halves;
        hoisting it out of the block loop removes one table lookup and two
        concatenations per layer.
        """
        img_freqs, txt_freqs = self.forward(img_hw, txt_seq_len, device, time=time)
        return torch.cat([img_freqs, txt_freqs], dim=-2)

    def _compute_image_freqs(
        self,
        height: int,
        width: int,
        scaling_factors: tuple[float, float],
        *,
        kappa: float = 1.0,
    ) -> torch.Tensor:
        """
        Compute frequency tensor for 2D image grid
        Args:
            height: Height of image
            width: Width of image
        Returns:
            Frequency tensor [height*width, total_dim] (complex)
        """
        h_dim, w_dim = self.axes_dim
        h_idx = torch.arange(height, dtype=torch.float32)
        w_idx = torch.arange(width, dtype=torch.float32)
        if self.centered:
            # Symmetric (possibly half-integer) coordinates around the origin.
            h_idx = h_idx - (height - 1) / 2.0
            w_idx = w_idx - (width - 1) / 2.0
        h_freqs = torch.view_as_real(
            self.rope_params(
                h_idx,
                h_dim,
                scaling_factors[0],
                kappa=kappa,
                reference=self.reference_hw[0],
            )
        )
        w_freqs = torch.view_as_real(
            self.rope_params(
                w_idx,
                w_dim,
                scaling_factors[1],
                kappa=kappa,
                reference=self.reference_hw[1],
            )
        )

        # Broadcast to create the grid
        # Expand to [H, W, dim, 2]
        h_freqs_grid = h_freqs.unsqueeze(1).expand(
            height, width, -1, -1
        )  # [height, width, h_dim//2, 2]
        w_freqs_grid = w_freqs.unsqueeze(0).expand(
            height, width, -1, -1
        )  # [height, width, w_dim//2, 2]

        # Concatenate
        freqs = torch.cat(
            [h_freqs_grid, w_freqs_grid], dim=2
        )  # [height, width, (h_dim+w_dim)//2, 2]

        # Flatten spatial dimensions
        freqs = freqs.flatten(0, 1)  # [height*width, (h_dim+w_dim)//2, 2]

        # Convert back to complex
        freqs = torch.view_as_complex(freqs)

        return freqs


class GatedFeedForward(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, dropout: float = 0.0):
        super().__init__()
        self.up_proj = nn.Linear(dim, hidden_dim * 2)
        self.down_proj = nn.Linear(hidden_dim, dim)
        self.dropout = nn.Dropout(dropout)
        self.dropout_out = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.up_proj(x)
        if x.device.type == "npu":
            import torch_npu

            x = torch_npu.npu_swiglu(x, dim=-1)
        else:
            x_gated, x_linear = x.chunk(2, dim=-1)
            x = F.silu(x_gated) * x_linear
        x = self.dropout(x)
        x = self.down_proj(x)
        return self.dropout_out(x)


def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)


class DoubleStreamAttention(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        qkv_bias: bool = True,
        rope_theta: int = 10000,
        rope_axes_dim: list = [64, 64],
        rope_centered: bool = False,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim**-0.5
        self.real_rope = False

        self.qkv_img = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.qkv_txt = nn.Linear(dim, dim * 3, bias=qkv_bias)

        self.q_norm_img = RMSNorm(self.head_dim, eps=1e-6)
        self.k_norm_img = RMSNorm(self.head_dim, eps=1e-6)
        self.q_norm_txt = RMSNorm(self.head_dim, eps=1e-6)
        self.k_norm_txt = RMSNorm(self.head_dim, eps=1e-6)

        self.rope = MSRoPE(
            theta=rope_theta,
            axes_dim=rope_axes_dim,
            centered=rope_centered,
        )

        self.proj_img = nn.Linear(dim, dim)
        self.proj_txt = nn.Linear(dim, dim)

    def forward(
        self,
        img_tokens: torch.Tensor,
        txt_tokens: torch.Tensor,
        img_hw: Tuple[int, int],
        txt_seq_len: int,
        txt_attention_mask: Optional[torch.Tensor] = None,
        attn_metadata=None,
        rope_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        rope_time: tuple[float, ...] | None = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        B, S_img, C = img_tokens.shape
        _, S_txt, _ = txt_tokens.shape

        # Image QKV
        qkv_img = (
            self.qkv_img(img_tokens)
            .reshape(B, S_img, 3, self.num_heads, self.head_dim)
            .permute(2, 0, 3, 1, 4)
        )
        q_img, k_img, v_img = qkv_img.unbind(0)  # [B, H, S, D]

        # Text QKV
        qkv_txt = (
            self.qkv_txt(txt_tokens)
            .reshape(B, S_txt, 3, self.num_heads, self.head_dim)
            .permute(2, 0, 3, 1, 4)
        )
        q_txt, k_txt, v_txt = qkv_txt.unbind(0)  # [B, H, S, D]

        # QK Norm
        q_img = self.q_norm_img(q_img)
        k_img = self.k_norm_img(k_img)
        q_txt = self.q_norm_txt(q_txt)
        k_txt = self.k_norm_txt(k_txt)

        # RoPE
        img_freqs, txt_freqs = (
            self.rope(img_hw, txt_seq_len, img_tokens.device, time=rope_time)
            if rope_freqs is None
            else rope_freqs
        )

        # Apply RoPE (need to transpose to [B, S, H, D] for apply_rotary_emb)
        rotate = apply_rotary_emb_real if self.real_rope else apply_rotary_emb
        q_img = rotate(q_img.transpose(1, 2), img_freqs).transpose(1, 2)
        k_img = rotate(k_img.transpose(1, 2), img_freqs).transpose(1, 2)

        q_txt = rotate(q_txt.transpose(1, 2), txt_freqs).transpose(1, 2)
        k_txt = rotate(k_txt.transpose(1, 2), txt_freqs).transpose(1, 2)

        # Concat
        q = torch.cat([q_img, q_txt], dim=2)
        k = torch.cat([k_img, k_txt], dim=2)
        v = torch.cat([v_img, v_txt], dim=2)

        # Prepare attention mask
        attn_mask = None
        if txt_attention_mask is not None and attn_metadata is None:
            # Image tokens always have attention
            img_mask = torch.ones(
                B,
                S_img,
                device=txt_attention_mask.device,
                dtype=txt_attention_mask.dtype,
            )
            full_mask = torch.cat([img_mask, txt_attention_mask], dim=1)

            # Convert to attention mask format [B, 1, 1, S_total]
            # Mask: 1 (valid) -> True (attend), 0 (padding) -> False (ignore)
            attn_mask = (full_mask > 0).unsqueeze(1).unsqueeze(2)
            attn_mask = attn_mask.to(dtype=torch.bool)

        # Attention with mask
        x = (
            packed_attention(q, k, v, attn_metadata)
            if attn_metadata is not None
            else sdpa_with_pad_mask(q, k, v, attn_mask)
        )

        # Split
        x_img = x[:, :, :S_img, :]
        x_txt = x[:, :, S_img:, :]

        # Reshape and Project
        x_img = x_img.transpose(1, 2).reshape(B, S_img, C)
        x_txt = x_txt.transpose(1, 2).reshape(B, S_txt, C)

        x_img = self.proj_img(x_img)
        x_txt = self.proj_txt(x_txt)

        return x_img, x_txt


class DoubleStreamDiTBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        c_dim: int,
        mlp_ratio: float = 4.0,
        rope_theta: int = 10000,
        rope_axes_dim: list = [64, 64],
    ):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.mlp_ratio = mlp_ratio

        # Normalize branch outputs before learned residual gates.
        self.norm_msa_img_out = RMSNorm(dim, eps=1e-6)
        self.norm_msa_txt_out = RMSNorm(dim, eps=1e-6)
        self.norm_mlp_img_out = RMSNorm(dim, eps=1e-6)
        self.norm_mlp_txt_out = RMSNorm(dim, eps=1e-6)

        # Modulation
        self.modulation_img = nn.Sequential(
            nn.SiLU(), nn.Linear(c_dim, 6 * dim, bias=True)
        )
        self.modulation_txt = nn.Sequential(
            nn.SiLU(), nn.Linear(c_dim, 6 * dim, bias=True)
        )

        # Attention
        self.norm1_img = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.norm1_txt = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.attn = DoubleStreamAttention(
            dim,
            num_heads,
            True,
            rope_theta,
            rope_axes_dim,
            rope_centered=True,
        )

        # MLP
        self.norm2_img = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.norm2_txt = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        mlp_hidden_dim = int(dim * mlp_ratio)

        self.mlp_img = GatedFeedForward(dim, mlp_hidden_dim)
        self.mlp_txt = GatedFeedForward(dim, mlp_hidden_dim)

    def forward(
        self,
        img_tokens: torch.Tensor,
        txt_tokens: torch.Tensor,
        c: torch.Tensor,
        img_hw: Tuple[int, int],
        txt_seq_len: int,
        txt_attention_mask: Optional[torch.Tensor] = None,
        attn_metadata=None,
        rope_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        rope_time: tuple[float, ...] | None = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Modulation
        # c: [B, c_dim]
        (
            shift_msa_img,
            scale_msa_img,
            gate_msa_img,
            shift_mlp_img,
            scale_mlp_img,
            gate_mlp_img,
        ) = self.modulation_img(c).chunk(6, dim=1)
        (
            shift_msa_txt,
            scale_msa_txt,
            gate_msa_txt,
            shift_mlp_txt,
            scale_mlp_txt,
            gate_mlp_txt,
        ) = self.modulation_txt(c).chunk(6, dim=1)

        # 1. Attention Block
        img_norm = modulate(self.norm1_img(img_tokens), shift_msa_img, scale_msa_img)
        txt_norm = modulate(self.norm1_txt(txt_tokens), shift_msa_txt, scale_msa_txt)

        img_attn, txt_attn = self.attn(
            img_norm,
            txt_norm,
            img_hw,
            txt_seq_len,
            txt_attention_mask,
            attn_metadata=attn_metadata,
            rope_freqs=rope_freqs,
            rope_time=rope_time,
        )

        img_attn = self.norm_msa_img_out(img_attn)
        txt_attn = self.norm_msa_txt_out(txt_attn)
        img_tokens = img_tokens + gate_msa_img.unsqueeze(1) * img_attn
        txt_tokens = txt_tokens + gate_msa_txt.unsqueeze(1) * txt_attn

        # 2. MLP Block
        img_norm = modulate(self.norm2_img(img_tokens), shift_mlp_img, scale_mlp_img)
        txt_norm = modulate(self.norm2_txt(txt_tokens), shift_mlp_txt, scale_mlp_txt)

        mlp_img_out = self.mlp_img(img_norm)
        mlp_txt_out = self.mlp_txt(txt_norm)
        mlp_img_out = self.norm_mlp_img_out(mlp_img_out)
        mlp_txt_out = self.norm_mlp_txt_out(mlp_txt_out)
        img_tokens = img_tokens + gate_mlp_img.unsqueeze(1) * mlp_img_out
        txt_tokens = txt_tokens + gate_mlp_txt.unsqueeze(1) * mlp_txt_out

        return img_tokens, txt_tokens

    def initialize_weights(self):
        # AdaLN-zero: Initialize modulation MLP to zero
        nn.init.constant_(self.modulation_img[-1].weight, 0)
        nn.init.constant_(self.modulation_img[-1].bias, 0)
        nn.init.constant_(self.modulation_txt[-1].weight, 0)
        nn.init.constant_(self.modulation_txt[-1].bias, 0)


class SingleStreamAttention(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        qkv_bias: bool = True,
        rope_theta: int = 10000,
        rope_axes_dim: list = [64, 64],
        rope_centered: bool = False,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim**-0.5
        self.real_rope = False

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.q_norm = RMSNorm(self.head_dim, eps=1e-6)
        self.k_norm = RMSNorm(self.head_dim, eps=1e-6)

        self.rope = MSRoPE(
            theta=rope_theta,
            axes_dim=rope_axes_dim,
            centered=rope_centered,
        )

        self.proj = nn.Linear(dim, dim)

    def forward(
        self,
        x: torch.Tensor,
        img_hw: Tuple[int, int],
        txt_seq_len: int,
        txt_attention_mask: Optional[torch.Tensor] = None,
        rope_freqs: Optional[torch.Tensor] = None,
        attn_bias: Optional[torch.Tensor] = None,
        attn_metadata=None,
        rope_time: tuple[float, ...] | None = None,
    ) -> torch.Tensor:
        B, S, C = x.shape

        qkv = (
            self.qkv(x)
            .reshape(B, S, 3, self.num_heads, self.head_dim)
            .permute(2, 0, 3, 1, 4)
        )
        q, k, v = qkv.unbind(0)

        q = self.q_norm(q)
        k = self.k_norm(k)

        if rope_freqs is not None:
            # Frequencies already cover [img, txt] in that order, so RoPE is a
            # single elementwise pass over the concatenated sequence. Under
            # the real-rope policy the hoisted table arrives in real
            # [S, D/2, 2] layout (converted once at model level, outside any
            # compiled block); otherwise it is complex.
            rotate = apply_rotary_emb_realfreq if self.real_rope else apply_rotary_emb
            q = rotate(q.transpose(1, 2), rope_freqs).transpose(1, 2)
            k = rotate(k.transpose(1, 2), rope_freqs).transpose(1, 2)
        else:
            rotate = apply_rotary_emb_real if self.real_rope else apply_rotary_emb
            # Need to apply different RoPE to img and txt parts
            # Assuming x is [img, txt]
            S_img = img_hw[0] * img_hw[1]

            q_img = q[:, :, :S_img, :]
            k_img = k[:, :, :S_img, :]
            q_txt = q[:, :, S_img:, :]
            k_txt = k[:, :, S_img:, :]

            img_freqs, txt_freqs = self.rope(
                img_hw, txt_seq_len, x.device, time=rope_time
            )

            q_img = rotate(q_img.transpose(1, 2), img_freqs).transpose(1, 2)
            k_img = rotate(k_img.transpose(1, 2), img_freqs).transpose(1, 2)
            q_txt = rotate(q_txt.transpose(1, 2), txt_freqs).transpose(1, 2)
            k_txt = rotate(k_txt.transpose(1, 2), txt_freqs).transpose(1, 2)

            q = torch.cat([q_img, q_txt], dim=2)
            k = torch.cat([k_img, k_txt], dim=2)

        if attn_metadata is not None:
            x = packed_attention(q, k, v, attn_metadata)
            return self.proj(x.transpose(1, 2).reshape(B, S, C))

        if attn_bias is not None:
            x = sdpa_with_bias(q, k, v, attn_bias)
            x = x.transpose(1, 2).reshape(B, S, C)
            return self.proj(x)

        # Prepare attention mask
        attn_mask = None
        if txt_attention_mask is not None:
            # Image tokens always have attention
            img_mask = torch.ones(
                B,
                S_img,
                device=txt_attention_mask.device,
                dtype=txt_attention_mask.dtype,
            )
            full_mask = torch.cat([img_mask, txt_attention_mask], dim=1)

            # Convert to attention mask format [B, 1, 1, S]
            # Mask : 1 (valid) -> True (attend), 0 (padding) -> False (ignore)
            attn_mask = (full_mask > 0).unsqueeze(1).unsqueeze(2)
            attn_mask = attn_mask.to(dtype=torch.bool)

        # Attention with mask
        x = sdpa_with_pad_mask(q, k, v, attn_mask)

        x = x.transpose(1, 2).reshape(B, S, C)

        x = self.proj(x)

        return x


class SingleStreamDiTBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        c_dim: int,
        mlp_ratio: float = 4.0,
        rope_theta: int = 10000,
        rope_axes_dim: list = [64, 64],
    ):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.mlp_ratio = mlp_ratio

        # Modulation
        self.modulation = nn.Sequential(nn.SiLU(), nn.Linear(c_dim, 3 * dim, bias=True))

        # Normalize branch outputs before the residual add. Learned gates and
        # affine norm gains remain outside this control of branch amplitude.
        self.norm_msa_out = RMSNorm(dim, eps=1e-6)
        self.norm_mlp_out = RMSNorm(dim, eps=1e-6)

        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.attn = SingleStreamAttention(
            dim,
            num_heads,
            True,
            rope_theta,
            rope_axes_dim,
            rope_centered=True,
        )

        # MLP
        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        mlp_hidden_dim = int(dim * mlp_ratio)

        self.mlp = GatedFeedForward(dim, mlp_hidden_dim)

    def forward(
        self,
        img_tokens: torch.Tensor,
        txt_tokens: torch.Tensor,
        c: torch.Tensor,
        img_hw: Tuple[int, int],
        txt_seq_len: int,
        txt_attention_mask: Optional[torch.Tensor] = None,
        rope_freqs: Optional[torch.Tensor] = None,
        attn_bias: Optional[torch.Tensor] = None,
        attn_metadata=None,
        rope_time: tuple[float, ...] | None = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Modulation
        # c: [B, c_dim]
        params = self.modulation(c)
        shift, scale, gate = params.chunk(3, dim=1)
        shift_msa = shift_mlp = shift
        scale_msa = scale_mlp = scale
        gate_msa = gate_mlp = gate

        # Concatenate for shared MSA/MLP processing
        S_img = img_tokens.shape[1]
        S_txt = txt_tokens.shape[1]
        x = torch.cat([img_tokens, txt_tokens], dim=1)

        # 1. Attention Block
        x_norm = modulate(self.norm1(x), shift_msa, scale_msa)

        x_attn = self.attn(
            x_norm,
            img_hw,
            txt_seq_len,
            txt_attention_mask,
            rope_freqs=rope_freqs,
            attn_bias=attn_bias,
            attn_metadata=attn_metadata,
            rope_time=rope_time,
        )

        x_attn = self.norm_msa_out(x_attn)
        x = x + gate_msa.unsqueeze(1) * x_attn

        # 2. MLP Block
        x_norm = modulate(self.norm2(x), shift_mlp, scale_mlp)
        mlp_out = self.mlp(x_norm)
        mlp_out = self.norm_mlp_out(mlp_out)
        x = x + gate_mlp.unsqueeze(1) * mlp_out

        img_tokens, txt_tokens = x.split([S_img, S_txt], dim=1)

        return img_tokens, txt_tokens

    def initialize_weights(self):
        # AdaLN-zero: Initialize modulation MLP to zero
        nn.init.constant_(self.modulation[-1].weight, 0)
        nn.init.constant_(self.modulation[-1].bias, 0)