Upload folder using huggingface_hub

README.md

@@ -1,13 +1,41 @@
 ---
-title: Pixelmoe
-emoji: 🌍
-colorFrom: pink
-colorTo: indigo
+title: PixelMoE
+emoji: 🧩
+colorFrom: green
+colorTo: red
 sdk: gradio
 sdk_version: 6.31.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Class-conditional 512x512 pixel-space MoE diffusion
+python_version: "3.12"
+startup_duration_timeout: 30m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# PixelMoE
+
+Interactive demo for **PixelMoE: Scaling Pixel-Space Diffusion with LLM-Style
+Mixture-of-Experts** ([arXiv 2610.18857](https://arxiv.org/abs/2610.18857)).
+
+Pick any of the 1000 ImageNet-1K classes and the Space samples a 512×512 image
+with the released **PixelMoE-B-E32A4/16** base model from
+[KlingTeam/PixelMoE](https://huggingface.co/KlingTeam/PixelMoE), using the
+authors' official sampling settings (Heun solver, 25 steps per hierarchical
+stage, CFG 3.7 applied while t < 0.8, second-EMA weights).
+
+PixelMoE runs diffusion entirely in pixel space: every image token activates
+4 of 32 routed experts plus a shared expert in each MoE block, and images are
+generated by a two-stage hierarchical flow — coarse structure first, then
+details at the target resolution.
+
+## Advanced settings
+
+- **Steps per stage** — denoising steps per hierarchical stage (25 in the
+  official `scripts/sample.sh`; each Heun step is two model evaluations).
+- **CFG scale / CFG interval max** — classifier-free guidance and the global
+  timestep below which it is active.
+- **Sigma min** — lower bound clamped into the velocity denominator.
+
+The model code is vendored from the MIT-licensed official repository
+[KlingAIResearch/PixelMoE](https://github.com/KlingAIResearch/PixelMoE)
+(`torch.compile` decorators and the deterministic-scatter toggle are adapted
+for the ZeroGPU runtime, otherwise inference is unchanged).

app.py

@@ -0,0 +1,252 @@
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import json
+import random
+import sys
+import time
+from pathlib import Path
+
+import spaces  # MUST come before torch / any CUDA-touching import
+import torch
+
+sys.path.insert(0, str(Path(__file__).parent / "pixelmoe_src"))
+
+from denoiser_inference import Denoiser
+
+import gradio as gr
+from huggingface_hub import hf_hub_download
+
+CHECKPOINT_REPO = "KlingTeam/PixelMoE"
+CHECKPOINT_FILE = "pixelmoe_b_e32a4_512.pth"
+MODEL_NAME = "PixelMoE-B-E32A4/16"
+IMG_SIZE = 512
+NUM_CLASSES = 1000
+# Generation defaults follow configs/pixelmoe_b_e32a4_512.sh + scripts/sample.sh
+DEFAULT_SIGMA_MIN = 0.03
+DEFAULT_CFG = 3.7
+DEFAULT_INTERVAL_MAX = 0.8
+DEFAULT_STEPS = 25  # per stage; scripts/sample.sh uses --num_sampling_steps 25 25
+
+_imagenet_classes = json.loads(
+    (Path(__file__).parent / "pixelmoe_src" / "imagenet_classes.json").read_text()
+)
+
+
+def _build_model() -> Denoiser:
+    model = Denoiser(
+        model=MODEL_NAME,
+        img_size=IMG_SIZE,
+        num_classes=NUM_CLASSES,
+        label_drop_prob=0.1,
+        sigma_min=DEFAULT_SIGMA_MIN,
+        noise_scale=2.0,
+        hierarchical_stages=(0, 1),
+        hierarchical_end_timesteps=(0.75, 0.0),
+        hierarchical_logit_mean=0.0,
+        hierarchical_logit_std=1.0,
+        hierarchical_noise_rho=-1.0 / 3.0,
+        hierarchical_min_stage_sigma=0.0,
+        sampling_method="heun",
+        num_sampling_steps=(DEFAULT_STEPS, DEFAULT_STEPS),
+        cfg=DEFAULT_CFG,
+        cfg_interval=(0.0, DEFAULT_INTERVAL_MAX),
+    )
+    ckpt_path = hf_hub_download(CHECKPOINT_REPO, CHECKPOINT_FILE)
+    checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    # sample.sh evaluates with the second EMA copy of the weights.
+    state_dict = checkpoint["model_ema2"]
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    if missing:
+        raise RuntimeError(f"Missing keys when loading checkpoint: {missing[:5]}...")
+    if unexpected:
+        raise RuntimeError(f"Unexpected keys when loading checkpoint: {unexpected[:5]}...")
+    del checkpoint
+    model.eval()
+    return model
+
+
+denoiser = _build_model().to("cuda")
+
+
+def _estimate_duration(class_name, num_steps, *args, **kwargs):
+    # Heun sampling: ~2 * steps NFE per stage on the two hierarchical stages.
+    return min(300, 30 + int(num_steps) * 2 * 2 * 2)
+
+
+@spaces.GPU(duration=_estimate_duration)
+def generate(
+    class_name: str,
+    num_steps: int = DEFAULT_STEPS,
+    cfg: float = DEFAULT_CFG,
+    interval_max: float = DEFAULT_INTERVAL_MAX,
+    sigma_min: float = DEFAULT_SIGMA_MIN,
+    seed: int = 0,
+    randomize_seed: bool = True,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Generate an ImageNet-1K class-conditional image with PixelMoE (512x512).
+
+    Args:
+        class_name: One of the 1000 ImageNet-1K class names.
+        num_steps: Denoising steps per hierarchical stage (25 in the paper's
+            sampling script; each Heun step costs two model evaluations).
+        cfg: Classifier-free guidance scale.
+        interval_max: Only apply CFG while the global timestep is below this
+            value (0.8 in the official config).
+        sigma_min: Lower bound clamped into the velocity denominator.
+        seed: RNG seed for reproducible sampling.
+        randomize_seed: Draw a fresh seed each run instead of the given one.
+    """
+    t0 = time.perf_counter()
+    try:
+        class_idx = _imagenet_classes.index(class_name)
+    except ValueError:
+        raise gr.Error(f"Unknown class name: {class_name!r}")
+
+    if randomize_seed:
+        seed = random.randint(0, 2**31 - 1)
+    seed = int(seed)
+
+    denoiser.cfg_scale = float(cfg)
+    denoiser.cfg_interval = (0.0, float(interval_max))
+    denoiser.sigma_min = float(sigma_min)
+    denoiser.sampling_steps = [int(num_steps), int(num_steps)]
+
+    generator_labels = torch.tensor([class_idx], dtype=torch.long, device="cuda")
+    g_state = torch.get_rng_state()
+    try:
+        torch.manual_seed(seed)
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            sampled = denoiser.generate(generator_labels)
+    finally:
+        torch.set_rng_state(g_state)
+
+    sampled = (
+        ((sampled + 1.0) / 2.0)
+        .detach()
+        .float()
+        .mul_(255)
+        .round_()
+        .clamp_(0, 255)
+        .to(dtype=torch.uint8)
+    )
+    image = sampled[0].permute(1, 2, 0).cpu().numpy()
+    elapsed = time.perf_counter() - t0
+    print(f"PixelMoE generated class={class_name!r} in {elapsed:.1f}s", flush=True)
+    return image, seed
+
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+main, .gradio-container, .fillable { max-width: 1100px !important; margin: 0 auto !important; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.HTML(
+            """
+            <div style="text-align: center; margin-bottom: 8px;">
+                <h1 style="margin-bottom: 4px;">🧩 PixelMoE</h1>
+                <p style="margin: 0; font-size: 1.05em; opacity: 0.85;">
+                    Pixel-space diffusion with LLM-style Mixture-of-Experts — pick an
+                    ImageNet-1K class and sample a 512×512 image.
+                </p>
+                <p style="margin: 6px 0 0; font-size: 0.9em; opacity: 0.75;">
+                    <a href="https://arxiv.org/abs/2610.18857">Paper</a> ·
+                    <a href="https://github.com/KlingAIResearch/PixelMoE">Code</a> ·
+                    <a href="https://huggingface.co/KlingTeam/PixelMoE">Weights</a> ·
+                    Base model, 512×512
+                </p>
+            </div>
+            """
+        )
+        with gr.Column():
+            class_name = gr.Dropdown(
+                choices=_imagenet_classes,
+                value="golden retriever",
+                label="ImageNet-1K class",
+                filterable=True,
+            )
+            with gr.Row():
+                run_button = gr.Button("Generate", variant="primary", scale=3)
+                seed_used = gr.Number(
+                    label="Seed used", value=0, precision=0, interactive=True, scale=1
+                )
+
+        output = gr.Image(label="Generated image", height=512, show_download_button=True)
+
+        with gr.Accordion("Advanced settings", open=False):
+            with gr.Row():
+                num_steps = gr.Slider(
+                    minimum=5, maximum=50, value=DEFAULT_STEPS, step=1,
+                    label="Steps per stage",
+                    info="The paper's sampling script uses 25 per stage (Heun ≈ 2 model calls each).",
+                )
+                cfg = gr.Slider(
+                    minimum=1.0, maximum=6.0, value=DEFAULT_CFG, step=0.1,
+                    label="CFG scale",
+                    info="3.7 in the official 512×512 Base config.",
+                )
+            with gr.Row():
+                interval_max = gr.Slider(
+                    minimum=0.1, maximum=1.0, value=DEFAULT_INTERVAL_MAX, step=0.01,
+                    label="CFG interval max",
+                    info="Guidance is applied only while t < this value.",
+                )
+                sigma_min = gr.Slider(
+                    minimum=0.0, maximum=0.1, value=DEFAULT_SIGMA_MIN, step=0.005,
+                    label="Sigma min",
+                    info="Clamp on the velocity denominator (0.03 for the Base model).",
+                )
+            with gr.Row():
+                seed = gr.Number(label="Seed", value=0, precision=0)
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+        gr.Examples(
+            examples=[
+                ["golden retriever"],
+                ["giant panda"],
+                ["yellow lady's slipper"],
+                ["tiger shark"],
+                ["school bus"],
+            ],
+            inputs=[class_name],
+            outputs=[output, seed_used],
+            fn=generate,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+        gr.Markdown(
+            """
+            PixelMoE runs diffusion directly in pixel space: each image token routes to 4 of
+            32 experts (plus a shared expert) in every MoE block, and images are generated in
+            two hierarchical flow-matching stages — coarse structure first, then detail at the
+            target resolution. This Space serves the released **PixelMoE-B-E32A4/16** base
+            model with its second-EMA weights and the official 512×512 sampling settings
+            (Heun solver, 25 steps per stage, CFG 3.7 with interval max 0.8).
+            """
+        )
+
+    inputs = [
+        class_name,
+        num_steps,
+        cfg,
+        interval_max,
+        sigma_min,
+        seed,
+        randomize_seed,
+    ]
+    outputs = [output, seed_used]
+    run_button.click(
+        generate, inputs=inputs, outputs=outputs, api_name="generate",
+        queue=True,
+    )
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)
+else:
+    demo.launch(mcp_server=True)

pixelmoe_src/deepseek_moe.py

@@ -0,0 +1,344 @@
+"""DeepSeek-style token-choice MoE using PyTorch grouped matrix multiply.
+"""
+
+import math
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def normalize_router_bias_update(imbalance, eps=1e-6):
+    """RMS-normalize the expert-load imbalance vector."""
+    return imbalance * torch.rsqrt(imbalance.square().mean() + float(eps))
+
+
+class DeepSeekMoEGate(nn.Module):
+    """Select top-k experts using FP32 sigmoid scores and balancing bias."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_experts: int,
+        top_k: int,
+        route_scale: float = 1.0,
+        router_bias_update_rate: float = 1e-3,
+    ):
+        super().__init__()
+        if num_experts <= 0:
+            raise ValueError("moe_num_experts must be positive")
+        if top_k <= 0 or top_k > num_experts:
+            raise ValueError("moe_top_k must be in [1, moe_num_experts]")
+
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.route_scale = float(route_scale)
+        self.router_bias_update_rate = float(router_bias_update_rate)
+
+        self.weight = nn.Parameter(torch.empty(num_experts, dim))
+        # Balancing bias is checkpointed and updated without gradients.
+        self.register_buffer(
+            "router_bias", 
+            torch.zeros(num_experts, dtype=torch.float32), 
+            persistent=True
+        )
+        # Cache this batch's expert loads until the optimizer update is accepted.
+        self.load_counts = None
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        self.router_bias.zero_()
+
+    def forward(self, hidden_states):
+        # Route packed tokens [num_tokens, dim] in FP32.
+        with torch.autocast(device_type=hidden_states.device.type, enabled=False):
+            logits = F.linear(
+                hidden_states.type(torch.float32),
+                self.weight.type(torch.float32),
+            )
+
+        scores = logits.sigmoid()
+
+        # Bias affects expert selection; aggregating weights use the unbiased scores.
+        selection_scores = scores + self.router_bias.to(
+            device=scores.device, dtype=scores.dtype
+        ).detach()
+        indices = torch.topk(selection_scores, self.top_k, dim=-1, sorted=False).indices
+        weights = scores.gather(1, indices)
+        # Normalize each token's selected weights, then apply the routing scale.
+        weights = weights / weights.sum(dim=-1, keepdim=True).clamp_min(1e-20)
+        weights = weights * self.route_scale
+        if self.training:
+            self.record_load_counts(indices)
+        return weights, indices
+
+    @torch.no_grad()
+    def record_load_counts(self, topk_idx):
+        # Each selected token contributes one count to each of its top-k experts.
+        counts = topk_idx.new_zeros(
+            (topk_idx.shape[0], self.num_experts),
+            dtype=torch.float32,
+        )
+        counts.scatter_(1, topk_idx, 1)
+        self.load_counts = counts.sum(dim=0)
+
+    @torch.no_grad()
+    def update_bias(self, update_rate=1e-3):
+        # Called after optimizer.step(); skipped batches clear their counts instead.
+        if self.load_counts is None:
+            return
+        counts = self.load_counts
+        if dist.is_initialized():
+            # Use the same averaged expert loads on every DDP rank.
+            dist.all_reduce(counts, op=dist.ReduceOp.SUM)
+            counts /= dist.get_world_size()
+        # RMS-normalize the deficit: underloaded experts gain bias, overloaded lose it.
+        delta = normalize_router_bias_update(counts.mean() - counts)
+        self.router_bias.add_(
+            delta.to(dtype=self.router_bias.dtype),
+            alpha=float(update_rate),
+        )
+        self.load_counts = None
+
+class DeepSeekTokenReorderer(nn.Module):
+    """Stably group token-expert assignments by expert index."""
+
+    def __init__(self, num_experts: int):
+        super().__init__()
+        self.num_experts = num_experts
+
+    def forward(self, top_scores, selected_experts_indices, token_ids):
+        """Group N * top_k token-expert assignments by expert.
+
+        Returns:
+            top_scores_experts_sorted: [N * top_k] mixing weights in expert
+                order, used to weight expert outputs.
+            token_indices_experts_sorted: [N * top_k] original token IDs in
+                the same order, used to gather inputs and scatter outputs.
+                Each token appears top_k times.
+            num_tokens_per_expert: [num_experts] route counts defining each
+                expert's contiguous input segment.
+        """
+        flat_experts = selected_experts_indices.reshape(-1)
+        if flat_experts.is_cuda:
+            num_tokens_per_expert = torch.histc(
+                flat_experts,
+                bins=self.num_experts,
+                min=0,
+                max=self.num_experts,
+            )
+        else:
+            # CPU histc does not support the int64 indices returned by topk.
+            num_tokens_per_expert = torch.bincount(
+                flat_experts,
+                minlength=self.num_experts,
+            )
+        sorted_idx = torch.argsort(flat_experts, stable=True)
+        top_scores_experts_sorted = top_scores.reshape(-1).index_select(0, sorted_idx)
+        token_indices_experts_sorted = token_ids.reshape(-1).index_select(0, sorted_idx)
+        return (
+            top_scores_experts_sorted,
+            token_indices_experts_sorted,
+            num_tokens_per_expert,
+        )
+
+
+def _run_experts_grouped_mm(
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w3: torch.Tensor,
+    x: torch.Tensor,
+    num_tokens_per_expert: torch.Tensor,
+) -> torch.Tensor:
+    """Run packed SwiGLU experts with grouped matrix multiplication."""
+    # Input rows are already grouped by expert; each row is one token-expert route.
+    if x.dim() != 2:
+        raise ValueError(f"grouped expert input must be 2D, got shape {tuple(x.shape)}")
+    if num_tokens_per_expert.numel() != w1.shape[0]:
+        raise ValueError("num_tokens_per_expert must contain one count per expert")
+
+    # Skip experts with no assigned tokens, preserving the order of active experts.
+    non_zero_mask = num_tokens_per_expert > 0
+    non_zero_indices = torch.nonzero(non_zero_mask, as_tuple=True)[0]
+    if len(non_zero_indices) == 0:
+        return x.new_empty((0, w2.shape[1]))
+    if len(non_zero_indices) < len(num_tokens_per_expert):
+        counts = num_tokens_per_expert[non_zero_mask]
+        active_w1 = w1[non_zero_mask]
+        active_w2 = w2[non_zero_mask]
+        active_w3 = w3[non_zero_mask]
+    else:
+        # Every expert has tokens, so use all expert weights directly.
+        counts = num_tokens_per_expert
+        active_w1 = w1
+        active_w2 = w2
+        active_w3 = w3
+    # Cumulative segment ends.
+    offsets = torch.cumsum(counts, dim=0, dtype=torch.int32)
+
+    # Compute each expert's gate (w1) and value (w3) projections in BF16.
+    x_bf16 = x.bfloat16()
+    gate = torch._grouped_mm(
+        x_bf16,
+        active_w1.bfloat16().transpose(-2, -1),
+        offs=offsets,
+    )
+    linear = torch._grouped_mm(
+        x_bf16,
+        active_w3.bfloat16().transpose(-2, -1),
+        offs=offsets,
+    )
+    # SwiGLU: multiply the SiLU-activated gate by the value projection.
+    activated_gate = F.silu(gate)
+    hidden = activated_gate * linear
+    # Project through w2, keeping the same expert-grouped row order.
+    output = torch._grouped_mm(
+        hidden,
+        active_w2.bfloat16().transpose(-2, -1),
+        offs=offsets,
+    )
+    return output.type_as(x)
+
+
+class DeepSeekGroupedExperts(nn.Module):
+    """Routed experts with matrices stacked along the expert dimension."""
+
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        num_experts: int,
+        num_layers: int,
+    ):
+        super().__init__()
+        self.w1 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim))
+        self.w2 = nn.Parameter(torch.empty(num_experts, output_dim, hidden_dim))
+        self.w3 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim))
+        self.num_layers = num_layers
+        self.reset_parameters()
+
+    def forward(self, x, num_tokens_per_expert):
+        return _run_experts_grouped_mm(
+            self.w1,
+            self.w2,
+            self.w3,
+            x,
+            num_tokens_per_expert,
+        )
+
+    def reset_parameters(self):
+        # Initialize expert projections with a truncated normal distribution.
+        nn.init.trunc_normal_(self.w1, mean=0.0, std=0.02)
+        nn.init.trunc_normal_(self.w2, mean=0.0, std=0.02)
+        nn.init.trunc_normal_(self.w3, mean=0.0, std=0.02)
+        # Scale the residual output projection by the transformer body depth.
+        self.w2.data /= (2 * self.num_layers) ** 0.5
+
+
+class DeepSeekSharedExpert(nn.Module):
+    """Dense SwiGLU shared expert used for every token."""
+
+    def __init__(self, dim: int, hidden_dim: int):
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x):
+        # Shared SwiGLU processes every image token independently of routed selection.
+        gate = self.w1(x)
+        linear = self.w3(x)
+        activated_gate = F.silu(gate)
+        product = activated_gate * linear
+        output = self.w2(product)
+        return output
+
+
+class DeepSeekMoEFFN(nn.Module):
+    """Standalone DeepSeek-style routed and shared expert FFN."""
+
+    def __init__(
+        self,
+        dim: int,
+        routed_inter_dim: int,
+        shared_inter_dim: int,
+        num_layers: int,
+        num_experts: int = 8,
+        top_k: int = 2,
+        route_scale: float = 1.0,
+        router_bias_update_rate: float = 1e-3,
+        drop: float = 0.0,
+    ):
+        super().__init__()
+        if routed_inter_dim is None or shared_inter_dim is None:
+            raise ValueError("MoE requires explicit routed_inter_dim and shared_inter_dim")
+        expert_hidden_dim = int(routed_inter_dim)
+        shared_hidden_dim = int(shared_inter_dim)
+        self.num_experts = int(num_experts)
+        self.top_k = top_k
+        self.gate = DeepSeekMoEGate(
+            dim=dim,
+            num_experts=num_experts,
+            top_k=top_k,
+            route_scale=route_scale,
+            router_bias_update_rate=router_bias_update_rate,
+        )
+        self.experts = DeepSeekGroupedExperts(
+            dim=dim,
+            hidden_dim=expert_hidden_dim,
+            output_dim=dim,
+            num_experts=self.num_experts,
+            num_layers=num_layers,
+        )
+        self.reorderer = DeepSeekTokenReorderer(num_experts=self.num_experts)
+        self.shared_experts = DeepSeekSharedExpert(dim, shared_hidden_dim)
+        self.ffn_dropout = nn.Dropout(drop)
+
+    @torch.no_grad()
+    def update_router_bias_from_latest(self):
+        self.gate.update_bias(self.gate.router_bias_update_rate)
+
+    @torch._dynamo.disable
+    def forward(self, hidden_states):
+        # Merge packed token dimensions, then select top-k experts per token.
+        original_shape = hidden_states.shape
+        hidden_states_flat = hidden_states.reshape(-1, original_shape[-1])
+        topk_weight, topk_idx = self.gate(hidden_states_flat)
+
+        # A token appears once per selected expert; token_ids allow restoring the original order.
+        token_ids = torch.arange(
+            hidden_states_flat.shape[0],
+            device=hidden_states.device,
+        ).unsqueeze(1).expand(-1, self.top_k)
+        (
+            top_scores_experts_sorted,
+            token_indices_experts_sorted,
+            num_tokens_per_expert,
+        ) = self.reorderer(topk_weight, topk_idx, token_ids.reshape(-1))
+
+        # Repeat each token ID across its features to gather complete rows.
+        gather_index = token_indices_experts_sorted.reshape(-1, 1).expand(
+            -1,
+            hidden_states_flat.shape[-1],
+        )
+        routed_input = torch.gather(hidden_states_flat, dim=0, index=gather_index)
+        routed_output = self.experts(routed_input, num_tokens_per_expert)
+        # Weight routed outputs in FP32, then cast to the input dtype.
+        routed_output = (
+            routed_output.float()
+            * top_scores_experts_sorted.float().unsqueeze(-1)
+        ).to(dtype=hidden_states.dtype)
+
+        # Start with the shared expert, then add weighted routed outputs back to each token.
+        output = self.shared_experts(hidden_states_flat)
+        # Inference note: the upstream repo toggles
+        # torch.use_deterministic_algorithms() around this scatter_add at
+        # inference time; the toggling is unsafe inside a forked ZeroGPU worker
+        # and the exact reduction order does not matter for sampling, so we
+        # always take the plain scatter_add path.
+        output = output.scatter_add(dim=0, index=gather_index, src=routed_output)
+        output = self.ffn_dropout(output)
+        return output.reshape(original_shape)

pixelmoe_src/denoiser_inference.py

@@ -0,0 +1,314 @@
+# Inference-only Denoiser for PixelMoE.
+# Mirrors denoiser.py from https://github.com/KlingAIResearch/PixelMoE
+# (MIT License, Copyright (c) 2026 KlingAI Research) with the training
+# logic (loss computation, EMA bookkeeping) removed and the constructor
+# taking plain values instead of argparse.Namespace.
+import math
+from typing import List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from model_pixelmoe import PixelMoE_models
+from hierarchical_image_noise import HierarchicalImageNoiseBuilder
+
+
+class Denoiser(nn.Module):
+    def __init__(
+        self,
+        model="PixelMoE-B-E32A4/16",
+        img_size=512,
+        num_classes=1000,
+        label_drop_prob=0.1,
+        sigma_min=0.03,
+        noise_scale=2.0,
+        hierarchical_stages=(0, 1),
+        hierarchical_end_timesteps=(0.75, 0.0),
+        hierarchical_logit_mean=0.0,
+        hierarchical_logit_std=1.0,
+        hierarchical_noise_rho=-1.0 / 3.0,
+        hierarchical_min_stage_sigma=0.0,
+        sampling_method="heun",
+        num_sampling_steps=(25, 25),
+        cfg=3.0,
+        cfg_interval=(0.0, 0.8),
+    ):
+        super().__init__()
+        self.net = PixelMoE_models[model](
+            input_size=img_size,
+            in_channels=3,
+            num_classes=num_classes,
+            proj_drop=0.0,
+            flux_attn_proj_dropout=False,
+            use_flash_attn=False,
+        )
+        self.img_size = img_size
+        self.num_classes = num_classes
+
+        self.label_drop_prob = label_drop_prob
+        self.sigma_min = sigma_min
+        self.noise_scale = noise_scale
+
+        # Hierarchical Flow training hyperparameters.
+        self.hierarchical_stages = list(hierarchical_stages)
+        self.hierarchical_logit_mean = hierarchical_logit_mean
+        self.hierarchical_logit_std = hierarchical_logit_std
+        self.hierarchical_noise_rho = hierarchical_noise_rho
+        self.hierarchical_min_stage_sigma = float(hierarchical_min_stage_sigma)
+        self.hierarchical_noise_builder = HierarchicalImageNoiseBuilder(
+            stages=self.hierarchical_stages,
+            end_timesteps=hierarchical_end_timesteps,
+            noise_rho=self.hierarchical_noise_rho,
+            noise_scale=self.noise_scale,
+            logit_mean=self.hierarchical_logit_mean,
+            logit_std=self.hierarchical_logit_std,
+            stage_logit_means=None,
+            stage_logit_stds=None,
+            min_stage_sigma=self.hierarchical_min_stage_sigma,
+        )
+
+        # generation hyper-parameters
+        self.method = sampling_method
+        self.sampling_steps = list(num_sampling_steps)
+        self.cfg_scale = cfg
+        self.cfg_interval = tuple(cfg_interval)
+
+    @torch.no_grad()
+    def generate(self, labels):
+        return self.generate_hierarchical(labels)
+
+    @torch.no_grad()
+    def generate_hierarchical(self, labels):
+        """Generate images from low to high resolution.
+
+        Each higher-resolution stage starts from the previous output,
+        upsampled and re-noised to match its starting distribution.
+        """
+        device = labels.device
+        dtype = next(self.net.parameters()).dtype
+        batch_size = labels.size(0)
+        num_stages = len(self.hierarchical_stages)
+
+        # Downsample full-resolution noise to the first stage. Multiplying by 2
+        # restores the noise variance after each 2x bilinear downsampling.
+        z = self.noise_scale * torch.randn(
+            batch_size,
+            3,
+            self.img_size,
+            self.img_size,
+            device=device,
+            dtype=dtype,
+        )
+        height = self.img_size
+        width = self.img_size
+        for _ in range(num_stages - 1):
+            height //= 2
+            width //= 2
+            z = F.interpolate(z, size=(height, width), mode="bilinear", align_corners=False) * 2
+
+        for stage_index in range(num_stages):
+            if stage_index > 0:
+                height *= 2
+                width *= 2
+                z = F.interpolate(z, size=(height, width), mode="nearest")
+                z = self._hierarchical_transition_renoise(z, stage_index)
+
+            z = self._sample_hierarchical_stage(
+                z,
+                stage_index,
+                labels,
+            )
+
+        return z
+
+    def _hierarchical_transition_renoise(self, latents, stage_index):
+        """Rescale the upsampled image and add noise for the next stage."""
+        previous_t_end = self.hierarchical_noise_builder.end_timesteps[stage_index - 1]
+        t_start = self.hierarchical_noise_builder.start_timesteps[stage_index]
+
+        rho = self.hierarchical_noise_rho
+        alpha = t_start / math.sqrt(1. - rho)
+        signal_scale = (1. - t_start) / (1. - previous_t_end)
+
+        # Rescale the previous output and add block-correlated noise with the
+        # same noise_scale used during training.
+        noise = self.sample_block_noise_like(latents) * self.noise_scale
+        return signal_scale * latents + alpha * noise
+
+    def sample_block_noise_like(self, latents):
+        """Sample correlated noise within each spatial 2x2 block.
+
+        Each block follows a 4D Gaussian with covariance:
+            cov = (1 - rho) * I + rho * 11^T
+
+        For rho=-1/3, the four pixels in a block are negatively correlated and
+        their sum has zero variance.
+        """
+        batch_size, channels, height, width = latents.shape
+        if height % 2 != 0 or width % 2 != 0:
+            raise ValueError("Hierarchical block noise requires even spatial sizes")
+
+        rho = self.hierarchical_noise_rho
+        covariance = (
+            torch.eye(4, device=latents.device, dtype=torch.float32) * (1 - rho)
+            + torch.ones(4, 4, device=latents.device, dtype=torch.float32) * rho
+        )
+        mean = torch.zeros(4, device=latents.device, dtype=torch.float32)
+        dist = torch.distributions.multivariate_normal.MultivariateNormal(mean, covariance)
+        block_number = batch_size * channels * (height // 2) * (width // 2)
+        noise = dist.sample((block_number,)).to(dtype=latents.dtype)
+        noise = noise.reshape(batch_size, channels, height // 2, width // 2, 2, 2)
+        noise = noise.permute(0, 1, 2, 4, 3, 5).reshape(batch_size, channels, height, width)
+        return noise
+
+    @torch.no_grad()
+    def _sample_hierarchical_stage(self, z, stage_index, labels):
+        """Denoise z through one stage using Euler or Heun steps."""
+        global_timesteps, sigmas = self._hierarchical_stage_timesteps(stage_index, z.device, z.dtype)
+
+        for i in range(len(sigmas) - 1):
+            t = global_timesteps[i].expand(z.shape[0]).to(device=z.device, dtype=z.dtype)
+            sigma = sigmas[i].expand(z.shape[0]).view(-1, 1, 1, 1)
+            sigma_next = sigmas[i + 1].expand(z.shape[0]).view(-1, 1, 1, 1)
+
+            if self.method == "euler":
+                z = self._hierarchical_euler_step(
+                    z, t, sigma, sigma_next, labels
+                )
+            elif self.method == "heun":
+                t_next = global_timesteps[i + 1].expand(z.shape[0]).to(device=z.device, dtype=z.dtype)
+                z = self._hierarchical_heun_step(
+                    z, t, t_next, sigma, sigma_next, labels
+                )
+            else:
+                raise NotImplementedError
+
+        return z
+
+    def _hierarchical_stage_timesteps(self, stage_index, device, dtype):
+        """Select this stage's grid points and return global timestep and stage-relative sigma."""
+        t_start = self.hierarchical_noise_builder.start_timesteps[stage_index]
+        t_end = self.hierarchical_noise_builder.end_timesteps[stage_index]
+        # A zero lower bound covers the full stage interval. A positive value
+        # excludes near-endpoint grid points, except for the endpoint added below.
+        min_sigma = self.hierarchical_min_stage_sigma
+        # Convert the stage-relative lower bound to global t. The true endpoint
+        # t_end is added below so sampling still finishes at sigma=0.
+        sampling_t_end = t_end + min_sigma * (t_start - t_end)
+        sampling_t_start = t_start
+        global_timesteps = self._hierarchical_global_sampling_timesteps(device, dtype)
+        global_timesteps = global_timesteps[
+            (global_timesteps <= sampling_t_start) & (global_timesteps >= sampling_t_end)
+        ]
+        global_timesteps = torch.cat([
+            global_timesteps,
+            torch.tensor([sampling_t_start, sampling_t_end, t_end], device=device, dtype=dtype),
+        ])
+        global_timesteps = torch.unique(global_timesteps, sorted=True).flip(0)
+        global_timesteps[0] = sampling_t_start
+        global_timesteps[-1] = t_end
+
+        denom = max(t_start - t_end, torch.finfo(dtype).eps)
+        sigmas = (global_timesteps - t_end) / denom
+        sigmas[-1] = 0.0
+        return global_timesteps, sigmas
+
+    def _hierarchical_global_sampling_timesteps(self, device, dtype):
+        """Build the shared global timestep grid from t=1 to t=0."""
+        if len(self.sampling_steps) == 1:
+            total_steps = self.sampling_steps[0] * len(self.hierarchical_stages)
+        else:
+            total_steps = sum(self.sampling_steps)
+
+        timesteps = torch.linspace(1.0, 0.0, total_steps + 1, device=device, dtype=dtype)
+        timesteps = self._logit_shift_sampling_timesteps(timesteps)
+        return timesteps
+
+    def _logit_shift_sampling_timesteps(self, timesteps):
+        """Adjust timestep spacing using the global Logit-Normal mean and std."""
+        eps = torch.finfo(timesteps.dtype).eps
+        t = timesteps.clamp(eps, 1.0 - eps)
+        exp_mean = math.exp(float(self.hierarchical_logit_mean))
+        shifted = exp_mean / (
+            exp_mean + torch.pow(1.0 / t - 1.0, float(self.hierarchical_logit_std))
+        )
+        shifted[0] = 1.0
+        shifted[-1] = 0.0
+        return shifted
+
+    @torch.no_grad()
+    def _hierarchical_euler_step(self, z, t, sigma, sigma_next, labels):
+        """Euler step: take one Euler step using the current velocity and change in sigma."""
+        velocity = self._forward_hierarchical_sample(
+            z,
+            t,
+            sigma,
+            labels,
+        )
+        return z + (sigma_next - sigma) * velocity
+
+    @torch.no_grad()
+    def _hierarchical_heun_step(self, z, t, t_next, sigma, sigma_next, labels):
+        """Heun step: average current and next velocities for a Heun step; use Euler at sigma=0."""
+        velocity = self._forward_hierarchical_sample(
+            z,
+            t,
+            sigma,
+            labels,
+        )
+        z_euler = z + (sigma_next - sigma) * velocity
+        if torch.all(sigma_next <= 0):
+            return z_euler
+
+        velocity_next = self._forward_hierarchical_sample(
+            z_euler,
+            t_next,
+            sigma_next,
+            labels,
+        )
+        return z + (sigma_next - sigma) * 0.5 * (velocity + velocity_next)
+
+    @torch.no_grad()
+    def _forward_hierarchical_sample(self, z, t, sigma, labels):
+        """Predict endpoint-derived velocity, applying CFG within its global-timestep interval."""
+        pred_cond = self.net(z, t, labels)[0]
+        v_cond = self._hierarchical_prediction_to_velocity(
+            pred_cond,
+            z,
+            sigma,
+        )
+
+        if self.cfg_scale == 1.0 or self.label_drop_prob <= 0.0:
+            return v_cond
+
+        uncond_labels = torch.full(
+            labels.shape,
+            self.num_classes,
+            device=labels.device,
+            dtype=labels.dtype,
+        )
+        pred_uncond = self.net(z, t, uncond_labels)[0]
+        v_uncond = self._hierarchical_prediction_to_velocity(
+            pred_uncond,
+            z,
+            sigma,
+        )
+
+        cfg_t = t.view(-1, 1, 1, 1)
+        low, high = self.cfg_interval
+        lower_ok = (low == 0) | (cfg_t > low)
+        upper_ok = (high == 1) | (cfg_t < high)
+        interval_mask = lower_ok & upper_ok
+        cfg_scale = torch.where(interval_mask, self.cfg_scale, 1.0)
+
+        v_guided = v_uncond + cfg_scale * (v_cond - v_uncond)
+        return v_guided
+
+    def _hierarchical_prediction_to_velocity(self, pred, noisy, sigma):
+        """Convert an endpoint prediction to velocity with a clamped sigma denominator."""
+        # v_theta = (x_sigma - predicted endpoint) / sigma.
+        dtype_eps = torch.finfo(sigma.dtype).eps
+        sampling_eps = float(self.sigma_min)
+        sigma_eps = max(sampling_eps, dtype_eps)
+        return (noisy - pred) / sigma.clamp_min(sigma_eps)

pixelmoe_src/hierarchical_image_noise.py

@@ -0,0 +1,475 @@
+# --------------------------------------------------------
+# References:
+# Pyramid Flow: https://github.com/jy0205/Pyramid-Flow
+# --------------------------------------------------------
+
+import math
+from typing import List, NamedTuple
+
+import torch
+import torch.nn.functional as F
+
+
+class HierarchicalImageNoiseOutput(NamedTuple):
+    # One entry per active stage in this batch.
+    noisy_images_list: List[torch.Tensor]
+    # Stage-relative timestep sigma; shape [stage_batch_size, 1, 1, 1].
+    sigmas_list: List[torch.Tensor]
+    # Global timestep t used to condition the model; shape [stage_batch_size].
+    timesteps_list: List[torch.Tensor]
+    # Stage endpoints; intermediate-stage endpoints can contain noise.
+    targets_list: List[torch.Tensor]
+    sample_indices_list: List[torch.Tensor]
+
+
+class HierarchicalImageNoiseBuilder:
+    """
+    Prepare noisy images and prediction targets for each resolution stage.
+
+    The input is a batch of images with shape [B, C, H, W]. Each image is
+    assigned to one stage for this training step; the number of images per
+    stage is determined by the fixed-density allocation. Sample a noise level
+    within the assigned stage using that stage's Logit-Normal distribution.
+
+    t is the global timestep and sigma is the timestep
+    within a stage: t = t_end + sigma * (t_start - t_end). Construct the noisy
+    image as x_sigma = sigma * x_start + (1 - sigma) * x_end.
+
+    Return noisy images, stage-relative sigmas, global timesteps, endpoint
+    targets, and original batch indices grouped by stage. Image tensors use
+    [stage_batch_size, C, stage_H, stage_W].
+    """
+
+    def __init__(
+        self,
+        stages=(0, 1),
+        end_timesteps=(0.75, 0.0),
+        noise_rho=-1.0 / 3.0,
+        noise_scale=1.0,
+        logit_mean=0.0,
+        logit_std=1.0,
+        stage_logit_means=None,
+        stage_logit_stds=None,
+        min_stage_sigma=0.0,
+    ):
+        self.stages = list(stages)
+        self.noise_scale = float(noise_scale)
+        self.logit_mean = logit_mean
+        self.logit_std = logit_std
+        self.stage_logit_means = [float(value) for value in (
+            [0.6] * len(self.stages) if stage_logit_means is None else stage_logit_means
+        )]
+        self.stage_logit_stds = [float(value) for value in (
+            [0.8] * len(self.stages) if stage_logit_stds is None else stage_logit_stds
+        )]
+        self.min_stage_sigma = float(min_stage_sigma)
+
+        if not 0.0 <= self.min_stage_sigma < 1.0:
+            raise ValueError("min_stage_sigma must be in [0, 1)")
+        if len(self.stage_logit_means) != len(self.stages):
+            raise ValueError("stage_logit_means must contain one value per hierarchical stage")
+        if len(self.stage_logit_stds) != len(self.stages):
+            raise ValueError("stage_logit_stds must contain one value per hierarchical stage")
+        if any(std <= 0.0 for std in self.stage_logit_stds):
+            raise ValueError("stage_logit_stds values must be positive")
+        self.end_timesteps = self._parse_end_timesteps(end_timesteps)
+        self.start_timesteps = self._compute_start_timesteps(self.end_timesteps, noise_rho)
+        self.training_end_timesteps = [
+            t_end + self.min_stage_sigma * (t_start - t_end)
+            for t_start, t_end in zip(self.start_timesteps, self.end_timesteps)
+        ]
+        self.fixed_density_bins = self._build_fixed_density_bins()
+
+    @torch.no_grad()
+    def __call__(self, x):
+        """
+        Build resolution stages and prepare their noisy training inputs.
+
+        Args
+            x (torch.Tensor):
+                Clean full-resolution images [B, C, H, W].
+        Returns
+            HierarchicalImageNoiseOutput: training tensors grouped by active stage.
+        """
+        images_list = self.get_hierarchical_images(x)
+        return self.add_continuous_hierarchical_noise(images_list)
+
+    # Stage intervals
+
+    def _parse_end_timesteps(self, end_timesteps):
+        # Convert stage endpoints to floats; each stage needs one end timestep.
+        values = [float(v) for v in end_timesteps]
+        if len(values) == len(self.stages):
+            end_timesteps = values
+        else:
+            raise ValueError("end_timesteps must contain one end t per hierarchical stage")
+
+        # Global timesteps lie in [0, 1].
+        if any(s < 0.0 or s > 1.0 for s in end_timesteps):
+            raise ValueError("end timesteps must be in [0, 1]")
+        # Later stages must end at an equal or smaller timestep, e.g. [0.75, 0.0].
+        if any(end_timesteps[i] < end_timesteps[i + 1] for i in range(len(end_timesteps) - 1)):
+            raise ValueError("end timesteps must be monotonically non-increasing")
+        return end_timesteps
+
+    @staticmethod
+    def _compute_start_timesteps(end_timesteps, rho):
+        start_timesteps = []
+        for i_s, _ in enumerate(end_timesteps):
+            # The first stage starts from pure noise at t=1.
+            if i_s == 0:
+                start_timesteps.append(1.0)
+                continue
+            # Match the next stage's start distribution after upsampling, rescaling, and re-noising.
+            previous_t_end = end_timesteps[i_s - 1]
+            t_start = (math.sqrt(1. - rho) * previous_t_end) / (math.sqrt(1. - rho) * previous_t_end + (1. - previous_t_end) * math.sqrt(-rho))
+            start_timesteps.append(t_start)
+        return start_timesteps
+
+    # Fixed-density allocation
+
+    @staticmethod
+    def _normal_cdf(value):
+        return 0.5 * (1.0 + math.erf(value / math.sqrt(2.0)))
+
+    @staticmethod
+    def _logit(value):
+        return math.log(value / (1.0 - value))
+
+    def _timestep_cdf(self, t):
+        t = min(max(float(t), 0.0), 1.0)
+        if t <= 0.0:
+            return 0.0
+        if t >= 1.0:
+            return 1.0
+        z = (self._logit(t) - float(self.logit_mean)) / float(self.logit_std)
+        return self._normal_cdf(z)
+
+    def _interval_mass(self, low, high):
+        if high <= low:
+            return 0.0
+        return max(0.0, self._timestep_cdf(high) - self._timestep_cdf(low))
+
+    def _build_fixed_density_bins(self):
+        """
+        Build stage-allocation probabilities from the global timestep distribution.
+        Overlapping intervals share their probability equally between covering stages.
+
+        Returns
+            List[dict]: entries with a stage index and normalized probability mass.
+                A stage can have multiple entries; all masses sum to 1.
+        """
+        # Split [0, 1] at each stage's training interval boundaries.
+        boundaries = [0.0, 1.0]
+        boundaries.extend(float(v) for v in self.start_timesteps)  # Add stage starts
+        boundaries.extend(float(v) for v in self.training_end_timesteps)  # Add training lower bounds.
+        boundaries = sorted(set(min(max(v, 0.0), 1.0) for v in boundaries))  # Clamp to [0, 1], remove duplicates, and sort.
+
+        bins = []  # One entry for each interval/stage pair.
+        for low, high in zip(boundaries[:-1], boundaries[1:]):
+            if high <= low:
+                continue
+
+            # Stage coverage is constant within each bin; check its midpoint.
+            midpoint = 0.5 * (low + high)  # A representative global t in this interval.
+            candidates = []  # Stages whose training intervals cover this bin.
+            for stage_index, (t_start, training_t_end) in enumerate(zip(self.start_timesteps, self.training_end_timesteps)):
+                stage_low = min(t_start, training_t_end)
+                stage_high = max(t_start, training_t_end)
+                if stage_low <= midpoint <= stage_high:
+                    candidates.append(stage_index)
+
+            if not candidates:
+                continue
+
+            # Divide this bin's global Logit-Normal probability among covering stages.
+            mass = self._interval_mass(low, high)  # P(low <= t <= high), not interval width.
+            if mass <= 0.0:
+                continue
+
+            split_mass = mass / len(candidates)  # Each covering stage receives an equal share.
+            for stage_index in candidates:
+                bins.append(
+                    {
+                        "stage": stage_index,
+                        "mass": split_mass,
+                    }
+                )
+
+        # Normalize over valid training intervals; uncovered bins have no quota.
+        total_mass = sum(item["mass"] for item in bins)
+        if total_mass <= 0.0:
+            raise ValueError("Fixed-density stage allocation produced zero probability mass")
+
+        for item in bins:
+            item["mass"] /= total_mass
+
+        return bins
+
+    @staticmethod
+    def _apportion_counts(batch_size, masses):
+        """
+        Convert bin probabilities to integer image quotas.
+
+        Args
+            batch_size (int):
+                Total number of images to assign.
+            masses (List[float]):
+                Normalized probability per bin entry.
+        Returns
+            List[int]: quotas per bin entry, summing to batch_size.
+        """
+        # Expected image counts may be fractional; first round each down.
+        expected = [mass * batch_size for mass in masses]
+        counts = [int(math.floor(value)) for value in expected]  # Initial integer quotas per bin entry.
+        remaining = batch_size - sum(counts)  # Number of images still unassigned.
+
+        # Sort bin indices by fractional remainder, largest first.
+        remainders = sorted(
+            range(len(expected)),
+            key=lambda i: expected[i] - counts[i],
+            reverse=True,
+        )
+
+        # Add one image to each of the top remaining entries.
+        for i in remainders[:remaining]:
+            counts[i] += 1
+
+        return counts
+
+    def _fixed_batch_bin_counts(self, batch_size):
+        """
+        Get integer image quotas in fixed_density_bins order.
+
+        Args
+            batch_size (int):
+                Number of images in the batch.
+        Returns
+            List[int]: one quota per bin entry, summing to batch_size.
+        """
+        # Keep separate bin entries even when they belong to the same stage.
+        masses = [item["mass"] for item in self.fixed_density_bins]  # Preserve bin order.
+        return self._apportion_counts(batch_size, masses)
+
+    # Stage-local sampling
+
+    def _sample_stage_local_timesteps(self, count, stage_index, device, dtype):
+        """
+        Sample noise levels within the selected stage's training interval.
+
+        Args
+            count (int):
+                Number of images assigned to this bin entry.
+            stage_index (int):
+                Stage whose interval and Logit-Normal parameters are used.
+            device (torch.device):
+                Device for sampled tensors.
+            dtype (torch.dtype):
+                Dtype for returned timesteps.
+        Returns
+            torch.Tensor: global timesteps [count], between the training lower bound and stage start.
+        """
+        low = float(self.training_end_timesteps[stage_index])  # Lower global timestep allowed during training.
+        high = float(self.start_timesteps[stage_index])  # Stage start in global timestep.
+
+        # Standard normal -> stage-specific normal -> sigmoid values in (0, 1).
+        normal = torch.randn(count, device=device, dtype=torch.float64)  # [count]
+        # u is relative to the truncated training interval. With a lower bound,
+        # the full-stage sigma is min_stage_sigma + (1 - min_stage_sigma) * u.
+        u = torch.sigmoid(
+            float(self.stage_logit_means[stage_index])
+            + float(self.stage_logit_stds[stage_index]) * normal
+        )
+
+        # Map the relative position u to the stage's global-timestep training interval.
+        timesteps = low + (high - low) * u  # [count]
+        return timesteps.clamp(min=low, max=high).to(dtype=dtype)
+
+    def _sample_fixed_batch_timesteps_and_stages(
+        self,
+        batch_size,
+        device,
+        dtype,
+    ):
+        """
+        Randomly assign images to fixed quotas, then sample their global timesteps.
+
+        Args
+            batch_size (int):
+                Number of images in the batch.
+            device (torch.device):
+                Device for sampled tensors.
+            dtype (torch.dtype):
+                Dtype for returned timesteps.
+        Returns
+            tuple: global timesteps [B] and stage indices [B], in original image order.
+        """
+        counts = self._fixed_batch_bin_counts(batch_size)  # Integer quota per bin entry; sum = B.
+        sampled_timesteps = torch.empty(batch_size, device=device, dtype=dtype)  # [B], global t per image.
+        chosen_stages = torch.empty(batch_size, device=device, dtype=torch.long)  # [B], stage index per image.
+
+        # Shuffle image indices so stage assignment does not depend on batch order.
+        permutation = torch.randperm(batch_size, device=device)
+
+        offset = 0  # Number of indices already taken from permutation.
+        for bin_info, count in zip(self.fixed_density_bins, counts):
+            if count == 0:
+                continue
+            indices = permutation[offset:offset + count]  # [count], original image indices for this quota.
+            stage_index = int(bin_info["stage"])
+
+            # Bin masses control counts; sample t over the stage's full training interval.
+            sampled_timesteps[indices] = self._sample_stage_local_timesteps(
+                count,
+                stage_index,
+                device,
+                dtype,
+            )
+            chosen_stages[indices] = stage_index  # Write results back in original batch order.
+            offset += count
+
+        return sampled_timesteps, chosen_stages
+
+    # Hierarchical images and noise
+
+    @torch.no_grad()
+    def get_hierarchical_images(self, x):
+        """
+        Build `[low_res, ..., high_res]` image tensors from a clean image batch.
+
+        Args
+            x (torch.Tensor):
+                Clean full-resolution images [B, C, H, W].
+        Returns
+            List[torch.Tensor]: clean images [B, C, H_stage, W_stage], from low to high resolution.
+        """
+        if x.ndim != 4:
+            raise ValueError(f"Expected image tensor [B, C, H, W], got shape {tuple(x.shape)}")
+
+        # Collect images from high to low resolution, then reverse to stage order.
+        image_list = [x]
+        height, width = x.shape[-2], x.shape[-1]
+        for _ in range(len(self.stages) - 1):
+            height //= 2
+            width //= 2
+            # Halve the image height and width for the next lower-resolution stage.
+            x = F.interpolate(x, size=(height, width), mode="bilinear", align_corners=False)
+            image_list.append(x)
+
+        image_list = list(reversed(image_list))
+        return image_list
+
+    @torch.no_grad()
+    def add_continuous_hierarchical_noise(self, images_list):
+        """
+        Construct noisy inputs and endpoint targets for the assigned images at each stage.
+
+        Args
+            images_list (List[torch.Tensor]):
+                Clean images [B, C, H_stage, W_stage], ordered from low to high resolution.
+                Every stage keeps the same original batch order.
+        Returns
+            HierarchicalImageNoiseOutput: per-active-stage noisy inputs, sigmas,
+                global timesteps, endpoint targets, and original image indices.
+        """
+        stages = len(self.stages)
+        if stages != len(images_list):
+            raise ValueError("Number of clean hierarchical images must match len(stages)")
+
+        base_image = images_list[-1]  # [B, C, H, W], the full-resolution clean images.
+        device = base_image.device
+        dtype = base_image.dtype
+        batch_size = base_image.shape[0]
+
+        # Scale full-resolution Gaussian noise, then downsample it for each stage.
+        # Multiplying by 2 after downsampling restores the noise variance.
+        noise = torch.randn_like(base_image) * self.noise_scale  # [B, C, H, W]
+        height, width = noise.shape[-2], noise.shape[-1]
+        noise_list = [noise]
+        cur_noise = noise
+        for _ in range(stages - 1):
+            height //= 2
+            width //= 2
+            cur_noise = F.interpolate(
+                cur_noise,
+                size=(height, width),
+                mode="bilinear",
+                align_corners=False,
+            ) * 2.0
+            noise_list.append(cur_noise)
+        noise_list = list(reversed(noise_list))  # Match images_list's low-to-high resolution order.
+
+        # Sample only where the stage-relative sigma is at least min_stage_sigma.
+        sampled_timesteps, chosen_stages = self._sample_fixed_batch_timesteps_and_stages(
+            batch_size,
+            device,
+            dtype,
+        )
+
+        stage_to_indices = [
+            torch.nonzero(chosen_stages == i_s, as_tuple=False).flatten()
+            for i_s in range(stages)
+        ]  # stage_to_indices[i] contains the original indices assigned to stage i.
+
+        # One tensor per active stage; stages with no assigned images are omitted.
+        noisy_images_list = []
+        sigmas_list = []
+        targets_list = []
+        timesteps_list = []
+        sample_indices_list = []
+
+        for i_s, indices in enumerate(stage_to_indices):
+            # Small batches can leave some stages empty.
+            if indices.numel() == 0:
+                continue
+
+            # Gather matching images, noise, and timesteps using original batch indices.
+            clean_image = images_list[i_s].index_select(0, indices)  # [B_stage, C, H_stage, W_stage]
+            stage_noise = noise_list[i_s].index_select(0, indices)  # Same shape as clean_image.
+            timesteps = sampled_timesteps.index_select(0, indices)  # [B_stage], global t.
+            t_start = torch.as_tensor(self.start_timesteps[i_s], device=device, dtype=dtype)
+            t_end = torch.as_tensor(self.end_timesteps[i_s], device=device, dtype=dtype)
+
+            # Stage-relative sigma is 1 at t_start and 0 at t_end.
+            denom = (t_start - t_end).clamp_min(torch.finfo(dtype).eps)
+            sigma = ((timesteps - t_end) / denom).clamp(0.0, 1.0)  # [B_stage]
+
+            # Construct the start image at global t_start.
+            if i_s == 0:
+                x_start = stage_noise  # First stage starts from pure noise.
+            else:
+                # Use the same images at the previous resolution, upsampled to this stage.
+                last_clean_image = images_list[i_s - 1].index_select(0, indices)
+                last_clean_image = F.interpolate(
+                    last_clean_image,
+                    size=(last_clean_image.shape[-2] * 2, last_clean_image.shape[-1] * 2),
+                    mode="nearest",
+                )
+                x_start = t_start * stage_noise + (1.0 - t_start) * last_clean_image
+
+            # Construct the endpoint the model must predict.
+            if i_s == stages - 1:
+                x_end = clean_image  # Final endpoint is clean, at t=0 in our configuration.
+            else:
+                x_end = t_end * stage_noise + (1.0 - t_end) * clean_image  # Intermediate endpoint retains noise.
+
+            # Broadcast sigma from [B_stage] to [B_stage, 1, 1, 1].
+            while sigma.ndim < x_start.ndim:
+                sigma = sigma.unsqueeze(-1)
+
+            # Interpolate from the stage's start image toward its endpoint target.
+            noisy_images = sigma * x_start + (1.0 - sigma) * x_end  # [B_stage, C, H_stage, W_stage]
+
+            noisy_images_list.append(noisy_images.to(dtype))
+            sigmas_list.append(sigma.to(dtype))  # [B_stage, 1, 1, 1], used by the loss.
+            timesteps_list.append(timesteps.to(dtype))  # [B_stage], used to condition the model.
+            targets_list.append(x_end.to(dtype))
+            sample_indices_list.append(indices)  # [B_stage], used to select matching class labels.
+
+        return HierarchicalImageNoiseOutput(
+            noisy_images_list=noisy_images_list,
+            sigmas_list=sigmas_list,
+            timesteps_list=timesteps_list,
+            targets_list=targets_list,
+            sample_indices_list=sample_indices_list,
+        )

pixelmoe_src/imagenet_classes.json

@@ -0,0 +1 @@
+["tench", "goldfish", "great white shark", "tiger shark", "hammerhead", "electric ray", "stingray", "cock", "hen", "ostrich", "brambling", "goldfinch", "house finch", "junco", "indigo bunting", "robin", "bulbul", "jay", "magpie", "chickadee", "water ouzel", "kite", "bald eagle", "vulture", "great grey owl", "European fire salamander", "common newt", "eft", "spotted salamander", "axolotl", "bullfrog", "tree frog", "tailed frog", "loggerhead", "leatherback turtle", "mud turtle", "terrapin", "box turtle", "banded gecko", "common iguana", "American chameleon", "whiptail", "agama", "frilled lizard", "alligator lizard", "Gila monster", "green lizard", "African chameleon", "Komodo dragon", "African crocodile", "American alligator", "triceratops", "thunder snake", "ringneck snake", "hognose snake", "green snake", "king snake", "garter snake", "water snake", "vine snake", "night snake", "boa constrictor", "rock python", "Indian cobra", "green mamba", "sea snake", "horned viper", "diamondback", "sidewinder", "trilobite", "harvestman", "scorpion", "black and gold garden spider", "barn spider", "garden spider", "black widow", "tarantula", "wolf spider", "tick", "centipede", "black grouse", "ptarmigan", "ruffed grouse", "prairie chicken", "peacock", "quail", "partridge", "African grey", "macaw", "sulphur-crested cockatoo", "lorikeet", "coucal", "bee eater", "hornbill", "hummingbird", "jacamar", "toucan", "drake", "red-breasted merganser", "goose", "black swan", "tusker", "echidna", "platypus", "wallaby", "koala", "wombat", "jellyfish", "sea anemone", "brain coral", "flatworm", "nematode", "conch", "snail", "slug", "sea slug", "chiton", "chambered nautilus", "Dungeness crab", "rock crab", "fiddler crab", "king crab", "American lobster", "spiny lobster", "crayfish", "hermit crab", "isopod", "white stork", "black stork", "spoonbill", "flamingo", "little blue heron", "American egret", "bittern", "crane", "limpkin", "European gallinule", "American coot", "bustard", "ruddy turnstone", "red-backed sandpiper", "redshank", "dowitcher", "oystercatcher", "pelican", "king penguin", "albatross", "grey whale", "killer whale", "dugong", "sea lion", "Chihuahua", "Japanese spaniel", "Maltese dog", "Pekinese", "Shih-Tzu", "Blenheim spaniel", "papillon", "toy terrier", "Rhodesian ridgeback", "Afghan hound", "basset", "beagle", "bloodhound", "bluetick", "black-and-tan coonhound", "Walker hound", "English foxhound", "redbone", "borzoi", "Irish wolfhound", "Italian greyhound", "whippet", "Ibizan hound", "Norwegian elkhound", "otterhound", "Saluki", "Scottish deerhound", "Weimaraner", "Staffordshire bullterrier", "American Staffordshire terrier", "Bedlington terrier", "Border terrier", "Kerry blue terrier", "Irish terrier", "Norfolk terrier", "Norwich terrier", "Yorkshire terrier", "wire-haired fox terrier", "Lakeland terrier", "Sealyham terrier", "Airedale", "cairn", "Australian terrier", "Dandie Dinmont", "Boston bull", "miniature schnauzer", "giant schnauzer", "standard schnauzer", "Scotch terrier", "Tibetan terrier", "silky terrier", "soft-coated wheaten terrier", "West Highland white terrier", "Lhasa", "flat-coated retriever", "curly-coated retriever", "golden retriever", "Labrador retriever", "Chesapeake Bay retriever", "German short-haired pointer", "vizsla", "English setter", "Irish setter", "Gordon setter", "Brittany spaniel", "clumber", "English springer", "Welsh springer spaniel", "cocker spaniel", "Sussex spaniel", "Irish water spaniel", "kuvasz", "schipperke", "groenendael", "malinois", "briard", "kelpie", "komondor", "Old English sheepdog", "Shetland sheepdog", "collie", "Border collie", "Bouvier des Flandres", "Rottweiler", "German shepherd", "Doberman", "miniature pinscher", "Greater Swiss Mountain dog", "Bernese mountain dog", "Appenzeller", "EntleBucher", "boxer", "bull mastiff", "Tibetan mastiff", "French bulldog", "Great Dane", "Saint Bernard", "Eskimo dog", "malamute", "Siberian husky", "dalmatian", "affenpinscher", "basenji", "pug", "Leonberg", "Newfoundland", "Great Pyrenees", "Samoyed", "Pomeranian", "chow", "keeshond", "Brabancon griffon", "Pembroke", "Cardigan", "toy poodle", "miniature poodle", "standard poodle", "Mexican hairless", "timber wolf", "white wolf", "red wolf", "coyote", "dingo", "dhole", "African hunting dog", "hyena", "red fox", "kit fox", "Arctic fox", "grey fox", "tabby", "tiger cat", "Persian cat", "Siamese cat", "Egyptian cat", "cougar", "lynx", "leopard", "snow leopard", "jaguar", "lion", "tiger", "cheetah", "brown bear", "American black bear", "ice bear", "sloth bear", "mongoose", "meerkat", "tiger beetle", "ladybug", "ground beetle", "long-horned beetle", "leaf beetle", "dung beetle", "rhinoceros beetle", "weevil", "fly", "bee", "ant", "grasshopper", "cricket", "walking stick", "cockroach", "mantis", "cicada", "leafhopper", "lacewing", "dragonfly", "damselfly", "admiral", "ringlet", "monarch", "cabbage butterfly", "sulphur butterfly", "lycaenid", "starfish", "sea urchin", "sea cucumber", "wood rabbit", "hare", "Angora", "hamster", "porcupine", "fox squirrel", "marmot", "beaver", "guinea pig", "sorrel", "zebra", "hog", "wild boar", "warthog", "hippopotamus", "ox", "water buffalo", "bison", "ram", "bighorn", "ibex", "hartebeest", "impala", "gazelle", "Arabian camel", "llama", "weasel", "mink", "polecat", "black-footed ferret", "otter", "skunk", "badger", "armadillo", "three-toed sloth", "orangutan", "gorilla", "chimpanzee", "gibbon", "siamang", "guenon", "patas", "baboon", "macaque", "langur", "colobus", "proboscis monkey", "marmoset", "capuchin", "howler monkey", "titi", "spider monkey", "squirrel monkey", "Madagascar cat", "indri", "Indian elephant", "African elephant", "lesser panda", "giant panda", "barracouta", "eel", "coho", "rock beauty", "anemone fish", "sturgeon", "gar", "lionfish", "puffer", "abacus", "abaya", "academic gown", "accordion", "acoustic guitar", "aircraft carrier", "airliner", "airship", "altar", "ambulance", "amphibian", "analog clock", "apiary", "apron", "ashcan", "assault rifle", "backpack", "bakery", "balance beam", "balloon", "ballpoint", "Band Aid", "banjo", "bannister", "barbell", "barber chair", "barbershop", "barn", "barometer", "barrel", "barrow", "baseball", "basketball", "bassinet", "bassoon", "bathing cap", "bath towel", "bathtub", "beach wagon", "beacon", "beaker", "bearskin", "beer bottle", "beer glass", "bell cote", "bib", "bicycle-built-for-two", "bikini", "binder", "binoculars", "birdhouse", "boathouse", "bobsled", "bolo tie", "bonnet", "bookcase", "bookshop", "bottlecap", "bow", "bow tie", "brass", "brassiere", "breakwater", "breastplate", "broom", "bucket", "buckle", "bulletproof vest", "bullet train", "butcher shop", "cab", "caldron", "candle", "cannon", "canoe", "can opener", "cardigan", "car mirror", "carousel", "carpenter's kit", "carton", "car wheel", "cash machine", "cassette", "cassette player", "castle", "catamaran", "CD player", "cello", "cellular telephone", "chain", "chainlink fence", "chain mail", "chain saw", "chest", "chiffonier", "chime", "china cabinet", "Christmas stocking", "church", "cinema", "cleaver", "cliff dwelling", "cloak", "clog", "cocktail shaker", "coffee mug", "coffeepot", "coil", "combination lock", "computer keyboard", "confectionery", "container ship", "convertible", "corkscrew", "cornet", "cowboy boot", "cowboy hat", "cradle", "crane", "crash helmet", "crate", "crib", "Crock Pot", "croquet ball", "crutch", "cuirass", "dam", "desk", "desktop computer", "dial telephone", "diaper", "digital clock", "digital watch", "dining table", "dishrag", "dishwasher", "disk brake", "dock", "dogsled", "dome", "doormat", "drilling platform", "drum", "drumstick", "dumbbell", "Dutch oven", "electric fan", "electric guitar", "electric locomotive", "entertainment center", "envelope", "espresso maker", "face powder", "feather boa", "file", "fireboat", "fire engine", "fire screen", "flagpole", "flute", "folding chair", "football helmet", "forklift", "fountain", "fountain pen", "four-poster", "freight car", "French horn", "frying pan", "fur coat", "garbage truck", "gasmask", "gas pump", "goblet", "go-kart", "golf ball", "golfcart", "gondola", "gong", "gown", "grand piano", "greenhouse", "grille", "grocery store", "guillotine", "hair slide", "hair spray", "half track", "hammer", "hamper", "hand blower", "hand-held computer", "handkerchief", "hard disc", "harmonica", "harp", "harvester", "hatchet", "holster", "home theater", "honeycomb", "hook", "hoopskirt", "horizontal bar", "horse cart", "hourglass", "iPod", "iron", "jack-o'-lantern", "jean", "jeep", "jersey", "jigsaw puzzle", "jinrikisha", "joystick", "kimono", "knee pad", "knot", "lab coat", "ladle", "lampshade", "laptop", "lawn mower", "lens cap", "letter opener", "library", "lifeboat", "lighter", "limousine", "liner", "lipstick", "Loafer", "lotion", "loudspeaker", "loupe", "lumbermill", "magnetic compass", "mailbag", "mailbox", "maillot", "maillot", "manhole cover", "maraca", "marimba", "mask", "matchstick", "maypole", "maze", "measuring cup", "medicine chest", "megalith", "microphone", "microwave", "military uniform", "milk can", "minibus", "miniskirt", "minivan", "missile", "mitten", "mixing bowl", "mobile home", "Model T", "modem", "monastery", "monitor", "moped", "mortar", "mortarboard", "mosque", "mosquito net", "motor scooter", "mountain bike", "mountain tent", "mouse", "mousetrap", "moving van", "muzzle", "nail", "neck brace", "necklace", "nipple", "notebook", "obelisk", "oboe", "ocarina", "odometer", "oil filter", "organ", "oscilloscope", "overskirt", "oxcart", "oxygen mask", "packet", "paddle", "paddlewheel", "padlock", "paintbrush", "pajama", "palace", "panpipe", "paper towel", "parachute", "parallel bars", "park bench", "parking meter", "passenger car", "patio", "pay-phone", "pedestal", "pencil box", "pencil sharpener", "perfume", "Petri dish", "photocopier", "pick", "pickelhaube", "picket fence", "pickup", "pier", "piggy bank", "pill bottle", "pillow", "ping-pong ball", "pinwheel", "pirate", "pitcher", "plane", "planetarium", "plastic bag", "plate rack", "plow", "plunger", "Polaroid camera", "pole", "police van", "poncho", "pool table", "pop bottle", "pot", "potter's wheel", "power drill", "prayer rug", "printer", "prison", "projectile", "projector", "puck", "punching bag", "purse", "quill", "quilt", "racer", "racket", "radiator", "radio", "radio telescope", "rain barrel", "recreational vehicle", "reel", "reflex camera", "refrigerator", "remote control", "restaurant", "revolver", "rifle", "rocking chair", "rotisserie", "rubber eraser", "rugby ball", "rule", "running shoe", "safe", "safety pin", "saltshaker", "sandal", "sarong", "sax", "scabbard", "scale", "school bus", "schooner", "scoreboard", "screen", "screw", "screwdriver", "seat belt", "sewing machine", "shield", "shoe shop", "shoji", "shopping basket", "shopping cart", "shovel", "shower cap", "shower curtain", "ski", "ski mask", "sleeping bag", "slide rule", "sliding door", "slot", "snorkel", "snowmobile", "snowplow", "soap dispenser", "soccer ball", "sock", "solar dish", "sombrero", "soup bowl", "space bar", "space heater", "space shuttle", "spatula", "speedboat", "spider web", "spindle", "sports car", "spotlight", "stage", "steam locomotive", "steel arch bridge", "steel drum", "stethoscope", "stole", "stone wall", "stopwatch", "stove", "strainer", "streetcar", "stretcher", "studio couch", "stupa", "submarine", "suit", "sundial", "sunglass", "sunglasses", "sunscreen", "suspension bridge", "swab", "sweatshirt", "swimming trunks", "swing", "switch", "syringe", "table lamp", "tank", "tape player", "teapot", "teddy", "television", "tennis ball", "thatch", "theater curtain", "thimble", "thresher", "throne", "tile roof", "toaster", "tobacco shop", "toilet seat", "torch", "totem pole", "tow truck", "toyshop", "tractor", "trailer truck", "tray", "trench coat", "tricycle", "trimaran", "tripod", "triumphal arch", "trolleybus", "trombone", "tub", "turnstile", "typewriter keyboard", "umbrella", "unicycle", "upright", "vacuum", "vase", "vault", "velvet", "vending machine", "vestment", "viaduct", "violin", "volleyball", "waffle iron", "wall clock", "wallet", "wardrobe", "warplane", "washbasin", "washer", "water bottle", "water jug", "water tower", "whiskey jug", "whistle", "wig", "window screen", "window shade", "Windsor tie", "wine bottle", "wing", "wok", "wooden spoon", "wool", "worm fence", "wreck", "yawl", "yurt", "web site", "comic book", "crossword puzzle", "street sign", "traffic light", "book jacket", "menu", "plate", "guacamole", "consomme", "hot pot", "trifle", "ice cream", "ice lolly", "French loaf", "bagel", "pretzel", "cheeseburger", "hotdog", "mashed potato", "head cabbage", "broccoli", "cauliflower", "zucchini", "spaghetti squash", "acorn squash", "butternut squash", "cucumber", "artichoke", "bell pepper", "cardoon", "mushroom", "Granny Smith", "strawberry", "orange", "lemon", "fig", "pineapple", "banana", "jackfruit", "custard apple", "pomegranate", "hay", "carbonara", "chocolate sauce", "dough", "meat loaf", "pizza", "potpie", "burrito", "red wine", "espresso", "cup", "eggnog", "alp", "bubble", "cliff", "coral reef", "geyser", "lakeside", "promontory", "sandbar", "seashore", "valley", "volcano", "ballplayer", "groom", "scuba diver", "rapeseed", "daisy", "yellow lady's slipper", "corn", "acorn", "hip", "buckeye", "coral fungus", "agaric", "gyromitra", "stinkhorn", "earthstar", "hen-of-the-woods", "bolete", "ear", "toilet tissue"]

pixelmoe_src/model_pixelmoe.py

@@ -0,0 +1,1222 @@
+# --------------------------------------------------------
+# 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,
+}

pixelmoe_src/util/model_util.py

@@ -0,0 +1,74 @@
+# --------------------------------------------------------
+# References:
+# Lightning-DiT: https://github.com/hustvl/LightningDiT
+# --------------------------------------------------------
+
+
+import torch
+from torch import nn
+import numpy as np
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        Normalize by RMS in FP32, apply a learned scale, and restore input dtype.
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return (self.weight * hidden_states).to(input_dtype)
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim]
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # Allocate half the embedding to each coordinate: width first, then height.
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_w, emb_h], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb

requirements.txt

@@ -0,0 +1,3 @@
+einops
+numpy
+https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt211-cu130-cp312/flash_attn-2.8.3-cp312-cp312-linux_x86_64.whl