app.py

import os
import sys
import subprocess
import tempfile

import spaces


PID_REPO_URL = "https://github.com/nv-tlabs/PiD.git"
PID_REPO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "PiD")

if not os.path.exists(PID_REPO_DIR):
    print(f"[pid] cloning {PID_REPO_URL} -> {PID_REPO_DIR}", flush=True)
    subprocess.check_call(["git", "clone", "--depth", "1", PID_REPO_URL, PID_REPO_DIR])
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-e", PID_REPO_DIR])

# PiD's loader resolves paths relative to CWD, so chdir into the repo root.
os.chdir(PID_REPO_DIR)
sys.path.insert(0, PID_REPO_DIR)

import torch
import numpy as np
import gradio as gr
from PIL import Image
from types import SimpleNamespace
from huggingface_hub import snapshot_download

# Pull just the Flux-1 / Z-Image-compatible checkpoints from nvidia/PiD into the
# repo's expected checkpoints/ tree.
snapshot_download(
    repo_id="nvidia/PiD",
    local_dir=PID_REPO_DIR,
    allow_patterns=[
        "checkpoints/PiD_res2k_sr4x_official_flux_distill_4step/*",
        "checkpoints/ae.safetensors",
    ],
)

from pid._src.inference.checkpoint_registry import get_pid_checkpoint
from pid._src.inference.create_dataset import XtCaptureCallback
from pid._src.inference.pipeline_registry import (
    decode_with_pipeline_vae,
    extract_latent,
    load_pipeline,
)
from pid._src.utils.model_loader import load_model_from_checkpoint


DTYPE = torch.bfloat16
BACKBONE = "zimage"
CKPT_TYPE = "2k"
SR_SCALE = 4
PID_INFERENCE_STEPS = 4

print("[pid] loading Z-Image pipeline...", flush=True)
pipeline, pipe_cfg = load_pipeline(BACKBONE, dtype=DTYPE)
pipeline.to("cuda")

print("[pid] loading PiD decoder...", flush=True)
pid_meta = get_pid_checkpoint(BACKBONE, CKPT_TYPE)
pid_model, _pid_cfg = load_model_from_checkpoint(
    experiment_name=pid_meta.experiment,
    checkpoint_path=pid_meta.checkpoint_path,
    config_file="pid/_src/configs/pid/config.py",
    enable_fsdp=False,
    strict=False,
)
pid_model.eval()
print("[pid] ready", flush=True)


def _latent_to_pil(tensor: torch.Tensor) -> Image.Image:
    """[C, H, W] in [-1, 1] -> PIL.Image."""
    if tensor.dim() == 4:
        tensor = tensor.squeeze(0)
    arr = ((tensor.float().clamp(-1, 1) + 1) * 127.5).permute(1, 2, 0).cpu().numpy().astype(np.uint8)
    return Image.fromarray(arr)


def _pid_decode(latent: torch.Tensor, baseline_01: torch.Tensor, sigma: float, caption: str) -> Image.Image:
    baseline_neg1_1 = baseline_01 * 2.0 - 1.0
    lq_h, lq_w = baseline_01.shape[-2], baseline_01.shape[-1]
    data_batch = {
        pid_model.config.input_caption_key: [caption],
        "LQ_video_or_image": baseline_neg1_1.to(dtype=DTYPE, device="cuda"),
        "LQ_latent": latent.to(dtype=DTYPE, device="cuda"),
        "degrade_sigma": torch.tensor([sigma], device="cuda", dtype=torch.float32),
    }
    samples = pid_model.generate_samples_from_batch(
        data_batch,
        cfg_scale=1.0,
        num_steps=PID_INFERENCE_STEPS,
        seed=0,
        shift=None,
        image_size=(lq_h * SR_SCALE, lq_w * SR_SCALE),
    )
    return _latent_to_pil(samples[0])


def _evenly_spaced_capture_steps(total_steps: int, num_captures: int) -> list[int]:
    """Pick N capture indices spread across [1, total_steps-1]. The final x0 is always added separately."""
    if num_captures <= 0:
        return []
    # avoid 0 (no forward pass yet) and total_steps (== final clean, captured separately)
    raw = np.linspace(1, max(2, total_steps - 1), num_captures + 1)[1:]
    return sorted({int(round(x)) for x in raw})


@spaces.GPU(duration=240)
def generate(
    prompt: str,
    num_inference_steps: int = 28,
    num_captures: int = 4,
    guidance_scale: float = 5.0,
    seed: int = 0,
    resolution: int = 512,
    progress=gr.Progress(track_tqdm=True),
):
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    num_inference_steps = int(num_inference_steps)
    num_captures = int(num_captures)
    resolution = int(resolution)
    H = W = resolution

    capture_ks = set(_evenly_spaced_capture_steps(num_inference_steps, num_captures))
    progress(0.05, desc="Running Z-Image latent diffusion…")

    xt_cb = XtCaptureCallback(capture_ks) if capture_ks else None
    generator = torch.Generator(device="cuda").manual_seed(int(seed))
    gen_kwargs = dict(
        prompt=prompt,
        height=H,
        width=W,
        num_inference_steps=num_inference_steps,
        guidance_scale=float(guidance_scale),
        num_images_per_prompt=1,
        output_type="latent",
        generator=generator,
    )
    gen_kwargs.update(pipe_cfg.extra_generate_kwargs)
    if xt_cb is not None:
        gen_kwargs["callback_on_step_end"] = xt_cb
        gen_kwargs["callback_on_step_end_tensor_inputs"] = ["latents"]

    with torch.no_grad():
        raw_output = pipeline(**gen_kwargs)
    final_latent = extract_latent(pipeline, raw_output, pipe_cfg, H, W)

    progress(0.5, desc="Decoding each captured step with PiD…")
    outputs: list[tuple[Image.Image, str]] = []
    steps_iter = []
    if xt_cb is not None:
        for K in sorted(xt_cb.captured.keys()):
            xt_packed_cpu, sigma = xt_cb.captured[K]
            xt_packed = xt_packed_cpu.to(device="cuda", dtype=DTYPE)
            xt_latent = extract_latent(pipeline, SimpleNamespace(images=xt_packed), pipe_cfg, H, W)
            steps_iter.append((f"step {K:02d}/{num_inference_steps}", xt_latent, sigma))
    final_sigma = float(pipeline.scheduler.sigmas[-1].item())
    steps_iter.append((f"final x₀", final_latent, final_sigma))

    total = len(steps_iter)
    for i, (label, latent, sigma) in enumerate(steps_iter):
        progress(0.5 + 0.5 * (i / total), desc=f"PiD decoding {label}")
        with torch.no_grad():
            baseline_01 = decode_with_pipeline_vae(pipeline, latent, pipe_cfg)
            pid_img = _pid_decode(latent, baseline_01, sigma, prompt)
        outputs.append((pid_img, f"{label}  (σ={sigma:.3f})"))

    return outputs


DESCRIPTION = """
# 🪄 PiD — Pixel Diffusion Decoder for Z-Image

Each tile shows what NVIDIA's [PiD](https://github.com/nv-tlabs/PiD) (a 4-step
distilled pixel-space diffusion decoder) reconstructs from Z-Image's denoising
loop at progressive timesteps. The first few tiles come from noisy intermediate
latents (`xt`); the last tile is decoded from the final clean `x₀`.

PiD upsamples 4× during decode, so a 512² Z-Image latent track becomes a
2048² super-resolved image.
"""

with gr.Blocks() as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                value="A photorealistic close-up of a brown tabby cat sitting on a rustic wooden table, morning light, ultra-detailed fur",
                lines=3,
            )
            with gr.Row():
                resolution = gr.Slider(label="Z-Image resolution", minimum=256, maximum=1024, step=128, value=512)
                num_inference_steps = gr.Slider(label="Z-Image steps", minimum=8, maximum=50, step=1, value=28)
            with gr.Row():
                num_captures = gr.Slider(label="Intermediate captures", minimum=1, maximum=8, step=1, value=4)
                guidance_scale = gr.Slider(label="Guidance", minimum=1.0, maximum=10.0, step=0.5, value=5.0)
            seed = gr.Number(label="Seed", value=0, precision=0)
            run = gr.Button("Run", variant="primary")
        with gr.Column(scale=2):
            gallery = gr.Gallery(label="PiD-decoded denoising trajectory", columns=2, object_fit="contain")

    run.click(
        fn=generate,
        inputs=[prompt, num_inference_steps, num_captures, guidance_scale, seed, resolution],
        outputs=[gallery],
    )

if __name__ == "__main__":
    demo.queue().launch()