Spaces:
Running on Zero
Running on Zero
| 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))) | |
| 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) | |