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( """

🧩 PixelMoE

Pixel-space diffusion with LLM-style Mixture-of-Experts — pick an ImageNet-1K class and sample a 512×512 image.

Paper · Code · Weights · Base model, 512×512

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