Spaces:

hugging-apps
/

pixelmoe

Running on Zero

pixelmoe / pixelmoe_src /model_pixelmoe.py
Upload folder using huggingface_hub
d155c16 verified 1 day ago
47.6 kB
# --------------------------------------------------------
# References:
# JiT: https://github.com/LTH14/JiT
# FLUX: https://github.com/black-forest-labs/flux
# Megatron-LM (weight initialization): https://github.com/NVIDIA/Megatron-LM
# --------------------------------------------------------

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

from deepseek_moe import (
    DeepSeekMoEFFN,
    DeepSeekMoEGate,
    DeepSeekSharedExpert,
)

from typing import List, Optional
from einops import rearrange
from util.model_util import get_2d_sincos_pos_embed, RMSNorm

try:
    from flash_attn.flash_attn_interface import flash_attn_varlen_func
except Exception:
    flash_attn_varlen_func = None

# Spaces note: torch.compile is not supported inside ZeroGPU's forked GPU
# worker; the upstream decorators were plain torch.compile which would raise
# at import time there. The forward passes stay eager.


# Helpers

def modulate(x, shift: torch.Tensor, scale: torch.Tensor):
    # Apply scale and shift to packed tokens; inputs and output: [1, total_tokens, D].
    return x * (1 + scale) + shift


def expand_segment_embeddings(embed: torch.Tensor, hidden_length: List[int], batch_size: int):
    """
    Repeat each image's conditioning over its packed tokens.
    Args
        embed (torch.Tensor):
            Conditioning embeddings with shape [num_images, D].
        hidden_length (List[int]):
            Token count per image, including context tokens when present.
        batch_size (int):
            1 for packed tokens.
    Returns
        token embeddings with shape [1, sum(hidden_length), D].
    """
    num_segments = len(hidden_length)
    lengths = torch.as_tensor(hidden_length, device=embed.device, dtype=torch.long)
    embed = embed.reshape(batch_size, num_segments, -1).transpose(0, 1)  # [num_segments, 1, D]
    return torch.repeat_interleave(embed, lengths, dim=0).transpose(0, 1).contiguous()  # [1, sum(hidden_length), D]


@torch._dynamo.disable
def split_ragged_context_image_tokens(mlp_input, lengths, context_len):
    """
    Split packed tokens into context and image groups.
    Args
        mlp_input (torch.Tensor):
            Packed tokens with shape [1, sum(lengths), D].
        lengths (List[int]):
            Token count per image, including context tokens.
        context_len (int):
            Context token count per image.
    Returns
        context tokens [1, num_images * context_len, D],
        image tokens [1, sum(image_lengths), D],
        image token counts.
    """
    segments = torch.split(mlp_input, lengths, dim=1)  # Each segment: [1, lengths[i], D].
    context_input = torch.cat(
        [segment[:, :context_len] for segment in segments],
        dim=1,
    )  # [1, num_images * context_len, D]
    image_lengths = [length - context_len for length in lengths]  # Image token count per image.
    image_input = torch.cat(
        [segment[:, context_len:] for segment in segments],
        dim=1,
    )  # [1, sum(image_lengths), D]
    return context_input, image_input, image_lengths


@torch._dynamo.disable
def merge_ragged_context_image_tokens(
    context_output,
    image_output,
    image_lengths,
    context_len,
):
    """
    Reassemble packed tokens in each image's context-then-image order.
    Args
        context_output (torch.Tensor):
            Context tokens with shape [1, num_images * context_len, D].
        image_output (torch.Tensor):
            Image tokens with shape [1, sum(image_lengths), D].
        image_lengths (List[int]):
            Image token count per image, excluding context tokens.
        context_len (int):
            Context token count per image.
    Returns
        packed tokens [1, num_images * context_len + sum(image_lengths), D].
    """
    context_segments = torch.split(
        context_output,
        [context_len] * len(image_lengths),
        dim=1,
    )  # Each segment: [1, context_len, D].
    image_segments = torch.split(image_output, image_lengths, dim=1)  # Each segment: [1, image_lengths[i], D].
    return torch.cat(
        [
            torch.cat([context_segment, image_segment], dim=1)  # [1, context_len + image_lengths[i], D]
            for context_segment, image_segment in zip(
                context_segments,
                image_segments,
            )
        ],
        dim=1,
    )  # [1, num_images * context_len + sum(image_lengths), D]


# RoPE

def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
    """
    Build 2x2 rotation matrices for one position axis.
    Args
        pos (torch.Tensor):
            Positions with shape [B, L].
        dim (int):
            Even feature dimension for this axis.
        theta (int):
            Frequency base.
    Returns
        FP32 rotation matrices with shape [B, L, dim // 2, 2, 2].
    """
    assert dim % 2 == 0, "The dimension must be even."

    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
    omega = 1.0 / (theta**scale)

    batch_size = pos.shape[0]
    out = torch.einsum("...n,d->...nd", pos, omega)  # [B, L, dim // 2]
    cos_out = torch.cos(out)
    sin_out = torch.sin(out)

    # Rotation matrix: [[cos, -sin], [sin, cos]].
    stacked_out = torch.stack([cos_out, -sin_out, sin_out, cos_out], dim=-1)
    out = stacked_out.view(batch_size, -1, dim // 2, 2, 2)
    return out.float()  # [B, L, dim // 2, 2, 2]


def apply_rope(xq, xk, freqs_cis):
    """
    Rotate query and key channel pairs in FP32 (B=1 for packed tokens).
    Args
        xq (torch.Tensor):
            Query with shape [B, L, heads, head_dim].
        xk (torch.Tensor):
            Key with shape [B, L, heads, head_dim].
        freqs_cis (torch.Tensor):
            Rotation matrices [B, L, 1, head_dim // 2, 2, 2], shared across heads.
    Returns
        rotated query and key with their original shapes and dtypes.
    """
    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)  # [B, L, heads, head_dim // 2, 1, 2]
    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]  # [B, L, heads, head_dim // 2, 2]
    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)


class EmbedND(nn.Module):
    """
    Build RoPE matrices for each position axis and concatenate them.
    Args
        theta (int):
            Frequency base.
        axes_dim (List[int]):
            Feature dimension per axis; height and width in this model.
    """
    def __init__(self, theta: int, axes_dim: List[int]):
        super().__init__()
        self.theta = theta
        self.axes_dim = axes_dim

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        """
        Embed each token's position axes as RoPE matrices.
        Args
            ids (torch.Tensor):
                Positions [B, L, num_axes]; B=1 and num_axes=2 for image tokens.
        Returns
            torch.Tensor: FP32 matrices [B, L, 1, sum(axes_dim) // 2, 2, 2], shared across heads.
        """
        n_axes = ids.shape[-1]
        emb = torch.cat(
            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
            dim=-3,
        )  # [B, L, sum(axes_dim) // 2, 2, 2]
        return emb.unsqueeze(2)  # [B, L, 1, sum(axes_dim) // 2, 2, 2]


# Condition embedders

class TimestepEmbedder(nn.Module):
    """
    Embeds global timestep t into a vector; stage-relative sigma
    is used by the loss and solver, and is not passed to this embedder.
    """
    def __init__(self, hidden_size, frequency_embedding_size=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size, bias=True),
        )
        self.frequency_embedding_size = frequency_embedding_size

    @staticmethod
    def timestep_embedding(
        timesteps: torch.Tensor,
        embedding_dim: int,
        max_period: int = 10000,
    ):
        """
        Create sinusoidal embeddings of global timesteps.
        Args
            timesteps (torch.Tensor):
                a 1-D tensor of N global timestep values, one per image.
            embedding_dim (int):
                the dimension of the output.
            max_period (int):
                Sets the largest period in the frequency scale.
        Returns
            timestep embeddings with shape [N, embedding_dim].
        """
        assert len(timesteps.shape) == 1, "Timesteps should be a 1D tensor."

        half_dim = embedding_dim // 2
        exponent = -math.log(max_period) * torch.arange(
            start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
        ) # [Dim // 2]
        exponent = exponent / (half_dim - 1)

        emb = torch.exp(exponent)
        emb = timesteps[:, None].float() * emb[None, :] # [B, Dim // 2]

        # Concatenate cosine and sine embeddings.
        emb = torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1) # [B, Dim]

        # zero pad
        if embedding_dim % 2 == 1:
            emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
        return emb

    def forward(self, t):
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
        t_freq = t_freq.to(dtype=self.mlp[0].weight.dtype)
        t_emb = self.mlp(t_freq)
        return t_emb


class LabelEmbedder(nn.Module):
    """
    Embed class labels; index num_classes represents the unconditional label.
    """
    def __init__(self, num_classes, hidden_size):
        super().__init__()
        self.embedding_table = nn.Embedding(num_classes + 1, hidden_size)

    def forward(self, labels):
        embeddings = self.embedding_table(labels)
        return embeddings


# Attention

@torch._dynamo.disable
def build_flash_cu_seqlens(hidden_length, device):
    """
    Build cumulative token offsets for FlashAttention in eager mode.
    Args
        hidden_length (List[int]):
            Token count per image in packed order.
        device (torch.device):
            Device for the offsets tensor.
    Returns
        offsets with shape [num_images + 1], starting at zero.
    """
    offsets = [0]
    total = 0
    for length in hidden_length:
        length = int(length)
        total += length
        offsets.append(total)
    return torch.tensor(offsets, device=device, dtype=torch.int32)


class VarlenFlashAttn(nn.Module):
    @torch._dynamo.disable
    def forward(self, query, key, value, cu_seqlens, max_seqlen, scale=None):
        cu_seqlens = cu_seqlens.to(device=query.device, dtype=torch.int32).contiguous()
        return flash_attn_varlen_func(
            query,
            key,
            value,
            cu_seqlens_q=cu_seqlens,
            cu_seqlens_k=cu_seqlens,
            max_seqlen_q=max_seqlen,
            max_seqlen_k=max_seqlen,
            dropout_p=0.0,
            causal=False,
            softmax_scale=scale,
        )


class VarlenFlashSelfAttnSingle(nn.Module):
    def __init__(self):
        super().__init__()
        self.flash_attn = VarlenFlashAttn()

    def forward(
        self,
        query,
        key,
        value,
        scale,
        hidden_length=None,
        image_rotary_emb=None,
    ):
        if flash_attn_varlen_func is None:
            raise ImportError("flash-attn is required when use_flash_attn=True")

        # Stack Q/K/V and apply RoPE to Q/K only.
        qkv = torch.stack([query, key, value], dim=2)
        if image_rotary_emb is not None:
            qkv[:, :, 0], qkv[:, :, 1] = apply_rope(
                qkv[:, :, 0],
                qkv[:, :, 1],
                image_rotary_emb,
            )

        qkv = rearrange(qkv, "b s ... -> (b s) ...")  # [total_tokens, 3, heads, head_dim]
        query, key, value = qkv.unbind(1)

        # Token offsets keep attention separate for each image in the packed sequence.
        max_seqlen_q = max(hidden_length)
        cu_seqlens_q = build_flash_cu_seqlens(
            hidden_length,
            query.device,
        )
        output = self.flash_attn(query, key, value, cu_seqlens_q, max_seqlen_q, scale)
        return output.reshape(1, -1, output.shape[-2], output.shape[-1]).flatten(2, 3)  # [1, total_tokens, heads * head_dim]


class VarlenSelfAttnSingle:
    def __init__(self):
        pass

    def __call__(
        self,
        query,
        key,
        value,
        hidden_length=None,
        image_rotary_emb=None,
    ):

        # Stack Q/K/V and apply RoPE to Q/K only.
        qkv = torch.stack([query, key, value], dim=2)  # [1, total_tokens, 3, heads, head_dim]
        if image_rotary_emb is not None:
            qkv[:, :, 0], qkv[:, :, 1] = apply_rope(
                qkv[:, :, 0],
                qkv[:, :, 1],
                image_rotary_emb,
            )

        i_sum = 0
        output_hidden_list = []
        # Compute attention within each image's token segment.
        for length in hidden_length:
            qkv_tokens = qkv[:, i_sum:i_sum + length]  # [1, length, 3, heads, head_dim]

            query, key, value = qkv_tokens.unbind(2)
            # SDPA expects Q/K/V in [1, heads, length, head_dim].
            query = query.transpose(1, 2).contiguous()
            key = key.transpose(1, 2).contiguous()
            value = value.transpose(1, 2).contiguous()

            stage_hidden_states = F.scaled_dot_product_attention(
                query,
                key,
                value,
                dropout_p=0.0,
                is_causal=False,
            )
            stage_hidden_states = stage_hidden_states.transpose(1, 2).flatten(2, 3)  # [1, length, heads * head_dim]
            output_hidden_list.append(stage_hidden_states)
            i_sum += length

        return torch.cat(output_hidden_list, dim=1)  # [1, total_tokens, heads * head_dim]


class FLUXAttention(nn.Module):

    def __init__(
        self,
        query_dim: int,
        heads: int = 8,
        dim_head: int = 64,
        dropout: float = 0.0,
        eps: float = 1e-6,
        use_flash_attn=False,
    ):
        super().__init__()
        self.heads = heads
        self.scale = dim_head**-0.5
        inner_dim = heads * dim_head
        self.norm_q = RMSNorm(dim_head, eps=eps)
        self.norm_k = RMSNorm(dim_head, eps=eps)
        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
        self.to_k = nn.Linear(query_dim, inner_dim, bias=False)
        self.to_v = nn.Linear(query_dim, inner_dim, bias=False)
        self.to_out = nn.ModuleList([
            nn.Linear(inner_dim, query_dim, bias=False),
            nn.Dropout(dropout),
        ])
        self.use_flash_attn = use_flash_attn
        if self.use_flash_attn:
            self.varlen_flash_attn = VarlenFlashSelfAttnSingle()
        else:
            self.varlen_attn = VarlenSelfAttnSingle()

    def forward(
        self,
        hidden_states: torch.Tensor,
        hidden_length: List = None,
        image_rotary_emb: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:

        # Project packed tokens to Q/K/V: [1, total_tokens, inner_dim].
        query = self.to_q(hidden_states)
        key = self.to_k(hidden_states)
        value = self.to_v(hidden_states)

        inner_dim = key.shape[-1]
        head_dim = inner_dim // self.heads

        # Split attention heads: [1, total_tokens, heads, head_dim].
        query = query.view(query.shape[0], -1, self.heads, head_dim)
        key = key.view(key.shape[0], -1, self.heads, head_dim)
        value = value.view(value.shape[0], -1, self.heads, head_dim)

        # Apply RMSNorm to each head's Q/K features.
        query = self.norm_q(query)
        key = self.norm_k(key)

        # Compute per-image attention and merge heads: [1, total_tokens, inner_dim].
        if self.use_flash_attn:
            hidden_states = self.varlen_flash_attn(
                query, key, value, 
                self.scale, hidden_length, 
                image_rotary_emb,
            )
        else:
            hidden_states = self.varlen_attn(
                query, key, value, 
                hidden_length, 
                image_rotary_emb,
            )

        # Project back to the input width, then apply dropout.
        hidden_states = self.to_out[0](hidden_states)
        hidden_states = self.to_out[1](hidden_states)

        return hidden_states


# Model components

class BottleneckPatchEmbed(nn.Module):
    """
    Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, pca_dim=768, embed_dim=768):
        super().__init__()
        img_size = (img_size, img_size)
        patch_size = (patch_size, patch_size)
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
        self.num_patches = num_patches

        self.proj1 = nn.Conv2d(in_chans, pca_dim, kernel_size=patch_size, stride=patch_size, bias=False)
        self.proj2 = nn.Conv2d(pca_dim, embed_dim, kernel_size=1, stride=1, bias=True)

    def forward(self, x):
        x = self.proj2(self.proj1(x)).flatten(2).transpose(1, 2)
        return x


class SwiGLUFFN(nn.Module):
    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        drop=0.0,
    ) -> None:
        super().__init__()
        hidden_dim = int(hidden_dim * 2 / 3)
        self.w12 = nn.Linear(dim, 2 * hidden_dim, bias=False)
        self.w3 = nn.Linear(hidden_dim, dim, bias=False)
        self.ffn_dropout = nn.Dropout(drop)

    def forward(self, x):
        x12 = self.w12(x)
        x1, x2 = x12.chunk(2, dim=-1)
        hidden = F.silu(x1) * x2
        return self.w3(self.ffn_dropout(hidden))


class PixelMoEBlock(nn.Module):
    def __init__(
        self,
        hidden_size,
        num_heads,
        flux_attn_proj_dropout=False,
        use_flash_attn=False,
        mlp_ratio=4.0,
        proj_drop=0.0,
        use_moe=False,
        moe_num_experts=8,
        moe_top_k=2,
        moe_routed_inter_dim=None,
        moe_shared_inter_dim=None,
        moe_route_scale=1.0,
        context_token_count=0,
        moe_router_bias_update_rate=1e-3,
        moe_num_layers=None,
    ):
        super().__init__()
        self.norm1 = RMSNorm(hidden_size, eps=1e-6)
        attention_head_dim = hidden_size // num_heads
        self.attn = FLUXAttention(
            query_dim=hidden_size,
            dim_head=attention_head_dim,
            heads=num_heads,
            eps=1e-6,
            use_flash_attn=use_flash_attn,
            dropout=proj_drop if flux_attn_proj_dropout else 0.0,
        )

        self.norm2 = RMSNorm(hidden_size, eps=1e-6)
        mlp_hidden_dim = int(hidden_size * mlp_ratio)
        if use_moe:
            self.mlp = DeepSeekMoEFFN(
                hidden_size,
                num_experts=moe_num_experts,
                top_k=moe_top_k,
                routed_inter_dim=moe_routed_inter_dim,
                shared_inter_dim=moe_shared_inter_dim,
                route_scale=moe_route_scale,
                router_bias_update_rate=moe_router_bias_update_rate,
                drop=proj_drop,
                num_layers=moe_num_layers,
            )
        else:
            self.mlp = SwiGLUFFN(
                hidden_size,
                mlp_hidden_dim,
                drop=proj_drop,
            )
        self.class_token_mlp = (
            SwiGLUFFN(
                hidden_size,
                mlp_hidden_dim,
                drop=proj_drop,
            )
            if use_moe and context_token_count > 0
            else None
        )
        self.context_token_count = int(context_token_count)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
        )

    def _forward_mlp(
        self,
        mlp_input,
        hidden_length=None,
    ):
        # Use one FFN when no separate class-token FFN is needed.
        if self.class_token_mlp is None:
            return self.mlp(
                mlp_input,
            )

        context_len = self.context_token_count
        # Split packed context/image tokens.
        lengths = [int(length) for length in hidden_length]
        if any(length <= context_len for length in lengths):
            raise ValueError(
                "Every ragged segment must contain class-context and image tokens; "
                f"got lengths={lengths}, context_token_count={context_len}"
            )
        context_input, image_input, image_lengths = split_ragged_context_image_tokens(
            mlp_input,
            lengths,
            context_len,
        )

        context_output = self.class_token_mlp(context_input)
        image_output = self.mlp(
            image_input,
        )
        return merge_ragged_context_image_tokens(
            context_output,
            image_output,
            image_lengths,
            context_len,
        )

    def forward(
        self,
        x,
        c,
        hidden_length=None,
        image_rotary_emb=None,
    ):
        # Expand each image's AdaLN conditioning to its packed tokens.
        embed = self.adaLN_modulation(c)
        batch_emb = expand_segment_embeddings(embed, hidden_length, x.shape[0])  # [1, total_tokens, 6 * D]
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = batch_emb.chunk(6, dim=-1)

        # Attention with AdaLN modulation and a gated residual.
        attn_input = modulate(self.norm1(x), shift_msa, scale_msa)
        attn_output = self.attn(
            hidden_states=attn_input,
            hidden_length=hidden_length,
            image_rotary_emb=image_rotary_emb,
        )
        x = x + gate_msa * attn_output

        # FFN with AdaLN modulation and a gated residual.
        mlp_input = modulate(self.norm2(x), shift_mlp, scale_mlp)
        mlp_output = self._forward_mlp(
            mlp_input,
            hidden_length=hidden_length,
        )
        x = x + gate_mlp * mlp_output
        return x


class PixelMoEDDTBlock(nn.Module):
    """
    Wide DDT head block with token-wise AdaLN conditioning of shape [B, L, D].
    """
    def __init__(
        self,
        hidden_size,
        num_heads,
        use_flash_attn=False,
        mlp_ratio=4.0,
        proj_drop=0.0,
    ):
        super().__init__()
        self.norm1 = RMSNorm(hidden_size, eps=1e-6)
        attention_head_dim = hidden_size // num_heads
        self.attn = FLUXAttention(
            query_dim=hidden_size,
            dim_head=attention_head_dim,
            heads=num_heads,
            eps=1e-6,
            use_flash_attn=use_flash_attn,
        )

        self.norm2 = RMSNorm(hidden_size, eps=1e-6)
        mlp_hidden_dim = int(hidden_size * mlp_ratio)
        self.mlp = SwiGLUFFN(
            hidden_size,
            mlp_hidden_dim,
            drop=proj_drop,
        )
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
        )

    def forward(
        self,
        x,
        c,
        hidden_length=None,
        image_rotary_emb=None,
    ):
        # Per-token AdaLN shift, scale, and gates; c: [1, total_tokens, D].
        embed = self.adaLN_modulation(c)
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = embed.chunk(6, dim=-1)

        # Attention with AdaLN modulation and a gated residual.
        attn_input = modulate(self.norm1(x), shift_msa, scale_msa)
        attn_output = self.attn(
            hidden_states=attn_input,
            hidden_length=hidden_length,
            image_rotary_emb=image_rotary_emb,
        )
        x = x + attn_output * gate_msa

        # Dense SwiGLU FFN with AdaLN modulation and a gated residual.
        mlp_output = self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
        x = x + mlp_output * gate_mlp
        return x


class DDTFinalLayer(nn.Module):
    """
    Token-wise conditioned final layer used by PixelMoE.
    """
    def __init__(self, hidden_size, patch_size, out_channels):
        super().__init__()
        self.norm_final = RMSNorm(hidden_size)
        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):
        # Apply per-token AdaLN conditioning before the output projection.
        embed = self.adaLN_modulation(c)
        shift, scale = embed.chunk(2, dim=-1)
        x = modulate(self.norm_final(x), shift, scale)
        x = self.linear(x)  # [1, total_tokens, patch_size * patch_size * out_channels]
        return x


# PixelMoE

class PixelMoE(nn.Module):
    """PixelMoE body with a wide DDT head and FLUX variable-length attention."""
    def __init__(
        self,
        input_size=256,
        patch_size=16,
        in_channels=3,
        hidden_size=1024,
        depth=24,
        num_heads=16,
        mlp_ratio=4.0,
        proj_drop=0.0,
        num_classes=1000,
        bottleneck_dim=128,
        in_context_len=32,
        in_context_start=8,
        flux_attn_proj_dropout=False,
        use_flash_attn=False,
        ddt_head_dim=2048,
        ddt_head_depth=2,
        ddt_head_num_heads=16,
        moe_num_experts=8,
        moe_top_k=2,
        moe_routed_inter_dim=None,
        moe_shared_inter_dim=None,
        moe_route_scale=1.0,
        moe_n_dense_layers=1,
        moe_router_bias_update_rate=1e-3,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = in_channels
        self.patch_size = patch_size
        self.hidden_size = hidden_size
        self.input_size = input_size
        self.depth = depth
        self.in_context_len = in_context_len
        self.in_context_start = in_context_start
        self.use_flash_attn = use_flash_attn
        # Base initialization std is 0.02; residual outputs use body-depth scaling.
        self.megatron_init_std = 0.02
        self.megatron_init_num_layers = depth

        # time and class embed
        self.t_embedder = TimestepEmbedder(hidden_size)
        self.y_embedder = LabelEmbedder(num_classes, hidden_size)

        # linear embed
        self.x_embedder = BottleneckPatchEmbed(input_size, patch_size, in_channels, bottleneck_dim, hidden_size)

        # use fixed sin-cos embedding
        num_patches = self.x_embedder.num_patches
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False)

        # in-context cls token
        self.in_context_posemb = nn.Parameter(torch.zeros(1, self.in_context_len, hidden_size), requires_grad=True)
        torch.nn.init.normal_(self.in_context_posemb, std=.02)

        self.hierarchical_pos_embed_cache = {}

        attention_head_dim = hidden_size // num_heads
        self.hierarchical_rope_embed = EmbedND(
            theta=10000,
            axes_dim=[attention_head_dim // 2, attention_head_dim // 2],
        )

        # transformer
        self.blocks = nn.ModuleList([
            PixelMoEBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, use_flash_attn=use_flash_attn, 
                    flux_attn_proj_dropout=flux_attn_proj_dropout,
                    proj_drop=proj_drop if (depth // 4 * 3 > i >= depth // 4) else 0.0,
                    use_moe=i >= moe_n_dense_layers,
                    context_token_count=in_context_len if i >= in_context_start else 0,
                    moe_num_experts=moe_num_experts,
                    moe_top_k=moe_top_k,
                    moe_routed_inter_dim=moe_routed_inter_dim,
                    moe_shared_inter_dim=moe_shared_inter_dim,
                    moe_route_scale=moe_route_scale,
                    moe_router_bias_update_rate=moe_router_bias_update_rate,
                    moe_num_layers=self.megatron_init_num_layers,
                    )
            for i in range(depth)
        ])

        self.initialize_weights()

        self.ddt_x_embedder = BottleneckPatchEmbed(
            self.input_size,
            self.patch_size,
            self.in_channels,
            ddt_head_dim,
            ddt_head_dim,
        )
        self.ddt_s_projector = nn.Linear(self.hidden_size, ddt_head_dim)
        self.ddt_blocks = nn.ModuleList([
            PixelMoEDDTBlock(
                ddt_head_dim,
                ddt_head_num_heads,
                mlp_ratio=4.0,
                use_flash_attn=self.use_flash_attn,
            )
            for _ in range(ddt_head_depth)
        ])
        self.ddt_final_layer = DDTFinalLayer(ddt_head_dim, self.patch_size, self.out_channels)

        attention_head_dim = ddt_head_dim // ddt_head_num_heads
        self.ddt_hierarchical_rope_embed = EmbedND(
            theta=10000,
            axes_dim=[attention_head_dim // 2, attention_head_dim // 2],
        )
        self.initialize_ddt_weights()

    def initialize_weights(self):
        self._initialize_megatron_linear_modules(self)

        # Initialize (and freeze) pos_embed using sinusoidal embeddings:
        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches ** 0.5))
        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))

        self._initialize_megatron_patch_embed(self.x_embedder)
        nn.init.normal_(self.y_embedder.embedding_table.weight, mean=0.0, std=self.megatron_init_std)
        nn.init.normal_(self.in_context_posemb, mean=0.0, std=self.megatron_init_std)
        self._initialize_megatron_output_projections(self)

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

    def initialize_ddt_weights(self):
        self._initialize_megatron_linear_modules(self.ddt_s_projector)
        self._initialize_megatron_linear_modules(self.ddt_blocks)
        self._initialize_megatron_patch_embed(self.ddt_x_embedder)
        self._initialize_megatron_output_projections(self.ddt_blocks)

        for block in self.ddt_blocks:
            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)

        nn.init.constant_(self.ddt_final_layer.adaLN_modulation[-1].weight, 0)
        nn.init.constant_(self.ddt_final_layer.adaLN_modulation[-1].bias, 0)
        nn.init.constant_(self.ddt_final_layer.linear.weight, 0)
        nn.init.constant_(self.ddt_final_layer.linear.bias, 0)

    @property
    def megatron_output_init_std(self):
        return self.megatron_init_std / math.sqrt(2.0 * self.megatron_init_num_layers)

    def _initialize_megatron_linear_modules(self, root):
        """
        Initialize ordinary projections with the Megatron base std.
        """
        def _init_linear(module):
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=self.megatron_init_std)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)
            elif isinstance(module, DeepSeekMoEGate):
                nn.init.normal_(module.weight, mean=0.0, std=self.megatron_init_std)

        root.apply(_init_linear)

    def _initialize_megatron_patch_embed(self, patch_embed):
        for projection in (patch_embed.proj1, patch_embed.proj2):
            nn.init.normal_(projection.weight, mean=0.0, std=self.megatron_init_std)
            if projection.bias is not None:
                nn.init.constant_(projection.bias, 0)

    def _initialize_megatron_output_projections(self, root):
        """
        Initialize attention and FFN output projections with depth-scaled std.
        """
        for module in root.modules():
            if isinstance(module, FLUXAttention):
                output_projection = module.to_out[0]
                nn.init.normal_(output_projection.weight, mean=0.0, std=self.megatron_output_init_std)
            elif isinstance(module, SwiGLUFFN):
                nn.init.normal_(module.w3.weight, mean=0.0, std=self.megatron_output_init_std)
            elif isinstance(module, DeepSeekSharedExpert):
                nn.init.normal_(module.w2.weight, mean=0.0, std=self.megatron_output_init_std)

    def get_hierarchical_pos_embed(self, num_patches, device=None, dtype=None):
        """
        Get cached sin-cos position embeddings for a square patch grid.
        Args
            num_patches (int):
                Number of image patch tokens.
            device (torch.device, optional):
                Device for the returned embeddings.
            dtype (torch.dtype, optional):
                Dtype for the returned embeddings.
        Returns
            torch.Tensor: position embeddings [1, num_patches, hidden_size].
        """
        key = str(num_patches)
        if key not in self.hierarchical_pos_embed_cache:
            grid_size = int(num_patches ** 0.5)
            if grid_size * grid_size != num_patches:
                raise ValueError(f"Expected square patch grid, got {num_patches} patches")

            pos_embed = get_2d_sincos_pos_embed(self.hidden_size, grid_size)
            self.hierarchical_pos_embed_cache[key] = torch.from_numpy(pos_embed).float().unsqueeze(0)

        return self.hierarchical_pos_embed_cache[key].to(device=device, dtype=dtype)

    @torch.no_grad()
    def _prepare_image_ids(self, height, width, device):
        """
        Build height/width coordinates for image-token RoPE.
        Args
            height (int):
                Patch grid height.
            width (int):
                Patch grid width.
            device (torch.device):
                Device for the coordinates.
        Returns
            coordinates [1, height * width, 2], in row-major order.
        """
        height_pos = torch.arange(height, device=device, dtype=torch.float32)
        width_pos = torch.arange(width, device=device, dtype=torch.float32)
        image_ids = torch.stack(torch.meshgrid(height_pos, width_pos, indexing="ij"), dim=-1)

        return rearrange(image_ids.unsqueeze(0), "b h w c -> b (h w) c")


    def _ddt_condition(self, body_tokens, t_emb, hidden_length):
        """
        Combine body features and timestep embeddings to condition the DDT head.
        Args
            body_tokens (torch.Tensor):
                Packed body features [1, total_tokens, hidden_size].
            t_emb (torch.Tensor):
                Timestep embeddings [num_images, hidden_size].
            hidden_length (List[int]):
                Image token count per image, excluding context tokens.
        Returns
            torch.Tensor: DDT conditioning [1, total_tokens, ddt_head_dim].
        """
        t_tokens = expand_segment_embeddings(t_emb, hidden_length, body_tokens.shape[0])
        condition = F.silu(body_tokens + t_tokens)
        return self.ddt_s_projector(condition)

    def forward(self, sample, timesteps, labels):
        """
        Predict stage endpoints conditioned on global timesteps t.
        Args
            sample (torch.Tensor or List[torch.Tensor]):
                Images [stage_batch_size, C, H, W] for each active stage.
            timesteps (torch.Tensor or List[torch.Tensor]):
                Global timesteps [stage_batch_size] for each active stage.
            labels (torch.Tensor):
                Class labels [num_images], concatenated in stage order.
        Returns
            List[torch.Tensor]: endpoint predictions [stage_batch_size, C, H, W] per active stage.
        """
        samples = sample if isinstance(sample, list) else [sample]
        timestep_list = timesteps if isinstance(timesteps, list) else [timesteps]
        if len(samples) != len(timestep_list):
            raise ValueError("sample and timesteps must have the same number of active stages")

        stage_tokens = []
        stage_ddt_tokens = []
        stage_batch_sizes = []
        stage_token_counts = []
        hidden_length = []
        output_shapes = []
        image_rotary_emb = []
        concat_image_rotary_emb = []
        ddt_image_rotary_emb = []

        for stage_sample in samples:
            batch_size, _, image_height, image_width = stage_sample.shape  # [B_stage, C, H, W]
            token_height = image_height // self.patch_size  # Patch grid height.
            token_width = image_width // self.patch_size  # Patch grid width.
            # Embed image patches separately for the body and DDT head.
            tokens = self.x_embedder(stage_sample)  # [B_stage, num_patches, hidden_size]
            tokens = tokens + self.get_hierarchical_pos_embed(tokens.shape[1], tokens.device, tokens.dtype)  # Add [1, num_patches, hidden_size].
            ddt_tokens = self.ddt_x_embedder(stage_sample)  # [B_stage, num_patches, ddt_head_dim]
            token_count = tokens.shape[1]  # num_patches = token_height * token_width

            # Image tokens use grid coordinates; class context tokens use (0, 0).
            base_ids = self._prepare_image_ids(token_height, token_width, tokens.device)  # [1, num_patches, 2]
            label_ids = torch.zeros(
                1,
                self.in_context_len,
                base_ids.shape[-1],
                device=tokens.device,
                dtype=base_ids.dtype,
            )  # [1, in_context_len, 2]
            context_ids = torch.cat([label_ids, base_ids], dim=1)  # [1, in_context_len + num_patches, 2]

            # Pack all images into batch_size=1; hidden_length retains image boundaries.
            stage_tokens.append(rearrange(tokens, "b s d -> 1 (b s) d"))
            stage_ddt_tokens.append(rearrange(ddt_tokens, "b s d -> 1 (b s) d"))
            stage_batch_sizes.append(batch_size)
            stage_token_counts.append(token_count)
            hidden_length.extend([token_count] * batch_size)
            output_shapes.append((token_height, token_width))
            # Images in the same stage share the same positional grid.
            image_rotary_emb.extend([self.hierarchical_rope_embed(base_ids)] * batch_size)
            concat_image_rotary_emb.extend([self.hierarchical_rope_embed(context_ids)] * batch_size)
            ddt_image_rotary_emb.extend([self.ddt_hierarchical_rope_embed(base_ids)] * batch_size)

        # Concatenate stages in the same order for tokens, RoPE, and conditioning.
        image_rotary_emb = torch.cat(image_rotary_emb, dim=1)  # [1, total_tokens, 1, body_head_dim // 2, 2, 2]
        concat_image_rotary_emb = torch.cat(concat_image_rotary_emb, dim=1)  # Image + context tokens, in per-image order.
        ddt_image_rotary_emb = torch.cat(ddt_image_rotary_emb, dim=1)  # [1, total_tokens, 1, ddt_head_dim_per_head // 2, 2, 2]
        hidden_states = torch.cat(stage_tokens, dim=1)  # [1, total_tokens, hidden_size]
        ddt_hidden_states = torch.cat(stage_ddt_tokens, dim=1)  # [1, total_tokens, ddt_head_dim]
        timesteps = torch.cat(timestep_list, dim=0)  # [num_images], in stage order.
        if labels.shape[0] != timesteps.shape[0]:
            raise ValueError("labels must be concatenated in the same stage order as sample")
        t_emb = self.t_embedder(timesteps.reshape(-1))  # [num_images, hidden_size]
        y_emb = self.y_embedder(labels)  # [num_images, hidden_size]
        c = t_emb + y_emb  # [num_images, hidden_size]

        has_label_context = False  # Context tokens are inserted at in_context_start.
        image_hidden_length = list(hidden_length)  # Keep image-only lengths for removing context later.

        def add_label_context(states):
            # Prepend class embeddings plus learned context positions to each image.
            stage_segments = []
            token_offset = 0  # Start of the current stage in packed image tokens.
            label_offset = 0  # Start of the current stage in y_emb.
            next_hidden_length = []  # Token count per image, including context.

            for batch_size, token_count in zip(stage_batch_sizes, stage_token_counts):
                stage_state = states[:, token_offset:token_offset + batch_size * token_count]  # [1, B_stage * token_count, hidden_size]
                stage_state = stage_state.reshape(batch_size, token_count, -1)  # [B_stage, token_count, hidden_size]
                context = y_emb[label_offset:label_offset + batch_size].unsqueeze(1).repeat(1, self.in_context_len, 1)  # [B_stage, in_context_len, hidden_size]
                context = context + self.in_context_posemb  # Add learned positions [1, in_context_len, hidden_size].
                stage_state = torch.cat([context, stage_state], dim=1)  # [B_stage, in_context_len + token_count, hidden_size]
                stage_segments.append(rearrange(stage_state, "b s d -> 1 (b s) d"))  # Pack [context_1, image_1, context_2, image_2, ...].
                next_hidden_length.extend([token_count + self.in_context_len] * batch_size)
                token_offset += batch_size * token_count
                label_offset += batch_size

            return torch.cat(stage_segments, dim=1), next_hidden_length  # [1, total_tokens + num_images * in_context_len, hidden_size], lengths.

        def strip_label_context(states):
            # Remove each image's context prefix and restore image-only token lengths.
            stage_segments = []
            token_offset = 0  # Start of the current stage in packed context + image tokens.

            for batch_size, token_count in zip(stage_batch_sizes, stage_token_counts):
                full_count = token_count + self.in_context_len  # Context + image token count per image.
                stage_state = states[:, token_offset:token_offset + batch_size * full_count]  # [1, B_stage * full_count, hidden_size]
                stage_state = stage_state.reshape(batch_size, full_count, -1)  # [B_stage, full_count, hidden_size]
                stage_state = stage_state[:, self.in_context_len:]  # [B_stage, token_count, hidden_size]
                stage_segments.append(rearrange(stage_state, "b s d -> 1 (b s) d"))  # [1, B_stage * token_count, hidden_size]
                token_offset += batch_size * full_count

            return torch.cat(stage_segments, dim=1), list(image_hidden_length)  # [1, total_tokens, hidden_size], image-only lengths.

        # Insert class context at in_context_start, then process the remaining body blocks.
        for i, block in enumerate(self.blocks):
            if i == self.in_context_start:
                hidden_states, hidden_length = add_label_context(hidden_states)
                has_label_context = True

            hidden_states = block(
                hidden_states,
                c,
                hidden_length=hidden_length,
                image_rotary_emb=concat_image_rotary_emb if has_label_context else image_rotary_emb,
            )

        # The DDT head receives image tokens only.
        hidden_states, hidden_length = strip_label_context(hidden_states)

        ddt_condition = self._ddt_condition(hidden_states, t_emb, hidden_length)
        for block in self.ddt_blocks:
            ddt_hidden_states = block(
                ddt_hidden_states,
                ddt_condition,
                hidden_length=hidden_length,
                image_rotary_emb=ddt_image_rotary_emb,
            )

        pred = self.ddt_final_layer(ddt_hidden_states, ddt_condition)

        # Unpack each stage and reshape predicted patch pixels back to images.
        outputs = []
        token_offset = 0
        for batch_size, token_count, (height, width) in zip(stage_batch_sizes, stage_token_counts, output_shapes):
            stage_pred = pred[:, token_offset:token_offset + batch_size * token_count]
            stage_pred = stage_pred.reshape(
                shape=(batch_size, height, width, self.patch_size, self.patch_size, -1)
            )
            stage_pred = rearrange(stage_pred, "b h w p1 p2 c -> b (h p1) (w p2) c")
            outputs.append(rearrange(stage_pred, "b h w c -> b c h w"))
            token_offset += batch_size * token_count
        return outputs


# Model variants

def PixelMoE_B_E32A4_8(**kwargs):
    return PixelMoE(depth=12, hidden_size=768, num_heads=12,
                    bottleneck_dim=128, in_context_len=32, in_context_start=4,
                    patch_size=8,
                    ddt_head_dim=1536, ddt_head_depth=2, ddt_head_num_heads=12,
                    moe_num_experts=32, moe_top_k=4,
                    moe_routed_inter_dim=384, moe_shared_inter_dim=632,
                    moe_route_scale=2.0, moe_n_dense_layers=4, **kwargs)


def PixelMoE_B_E32A4_16(**kwargs):
    return PixelMoE(depth=12, hidden_size=768, num_heads=12,
                    bottleneck_dim=128, in_context_len=32, in_context_start=4,
                    patch_size=16,
                    ddt_head_dim=1536, ddt_head_depth=2, ddt_head_num_heads=12,
                    moe_num_experts=32, moe_top_k=4,
                    moe_routed_inter_dim=384, moe_shared_inter_dim=632,
                    moe_route_scale=2.0, moe_n_dense_layers=4, **kwargs)


def PixelMoE_L_E32A4_8(**kwargs):
    return PixelMoE(depth=24, hidden_size=1024, num_heads=16,
                    bottleneck_dim=128, in_context_len=32, in_context_start=8,
                    patch_size=8,
                    ddt_head_dim=2048, ddt_head_depth=2, ddt_head_num_heads=16,
                    moe_num_experts=32, moe_top_k=4,
                    moe_routed_inter_dim=512, moe_shared_inter_dim=836,
                    moe_route_scale=2.0, moe_n_dense_layers=4, **kwargs)


def PixelMoE_L_E32A4_16(**kwargs):
    return PixelMoE(depth=24, hidden_size=1024, num_heads=16,
                    bottleneck_dim=128, in_context_len=32, in_context_start=8,
                    patch_size=16,
                    ddt_head_dim=2048, ddt_head_depth=2, ddt_head_num_heads=16,
                    moe_num_experts=32, moe_top_k=4,
                    moe_routed_inter_dim=512, moe_shared_inter_dim=836,
                    moe_route_scale=2.0, moe_n_dense_layers=4, **kwargs)


def PixelMoE_H_E32A4_8(**kwargs):
    return PixelMoE(depth=32, hidden_size=1280, num_heads=16,
                    bottleneck_dim=256, in_context_len=32, in_context_start=10,
                    patch_size=8,
                    ddt_head_dim=2560, ddt_head_depth=4, ddt_head_num_heads=16,
                    moe_num_experts=32, moe_top_k=4,
                    moe_routed_inter_dim=640, moe_shared_inter_dim=1048,
                    moe_route_scale=2.0, moe_n_dense_layers=4, **kwargs)


def PixelMoE_H_E32A4_16(**kwargs):
    return PixelMoE(depth=32, hidden_size=1280, num_heads=16,
                    bottleneck_dim=256, in_context_len=32, in_context_start=10,
                    patch_size=16,
                    ddt_head_dim=2560, ddt_head_depth=4, ddt_head_num_heads=16,
                    moe_num_experts=32, moe_top_k=4,
                    moe_routed_inter_dim=640, moe_shared_inter_dim=1048,
                    moe_route_scale=2.0, moe_n_dense_layers=4, **kwargs)


PixelMoE_models = {
    "PixelMoE-B-E32A4/8": PixelMoE_B_E32A4_8,
    "PixelMoE-B-E32A4/16": PixelMoE_B_E32A4_16,
    "PixelMoE-L-E32A4/8": PixelMoE_L_E32A4_8,
    "PixelMoE-L-E32A4/16": PixelMoE_L_E32A4_16,
    "PixelMoE-H-E32A4/8": PixelMoE_H_E32A4_8,
    "PixelMoE-H-E32A4/16": PixelMoE_H_E32A4_16,
}