Blame for app.py

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):
    # Measured on ZeroGPU (sm_120): ~4.6s of model time at 10 steps, ~8.3s at
    # 25 steps, ~17.3s at 50 steps (Heun = 2 NFE per step, two hierarchical
    # stages). ~0.35s/step; budget 1.5x that to absorb cold-start weight
    # streaming and load variance, floored for the smallest runs.
    return min(60, max(10, 2 + int(num_steps * 0.5)))


@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); }
"""

# Gradio 6 moved `theme` and `css` from the Blocks constructor to launch().
with gr.Blocks() 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
                )

        # Gradio 6: toolbar buttons are configured via `buttons` (download is shown by default).
        output = gr.Image(label="Generated image", height=512, buttons=["download", "fullscreen"])

        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,
    )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)