app.py

try:
    import spaces
except ImportError:

    class _SpacesShim:
        @staticmethod
        def GPU(*args, **kwargs):
            def decorator(fn):
                return fn

            return decorator

    spaces = _SpacesShim()

import gc
import os
import random
import threading

import gradio as gr
import torch
from diffusers import FlowMatchEulerDiscreteScheduler, ZImagePipeline

MODEL_IDS = {
    "turbo": "Tongyi-MAI/Z-Image-Turbo",
    "normal": "Tongyi-MAI/Z-Image",
}

DEFAULT_MODEL = os.environ.get("DEFAULT_MODEL", "turbo")
if DEFAULT_MODEL not in MODEL_IDS:
    DEFAULT_MODEL = "turbo"

UNLOAD_ON_SWITCH = os.environ.get("UNLOAD_ON_SWITCH", "0") == "1"

MAX_SEED = 2_147_483_647

RESOLUTIONS = {
    "1024x1024 (1:1)": (1024, 1024),
    "1152x896 (9:7)": (1152, 896),
    "896x1152 (7:9)": (896, 1152),
    "1280x720 (16:9)": (1280, 720),
    "720x1280 (9:16)": (720, 1280),
    "1248x832 (3:2)": (1248, 832),
    "832x1248 (2:3)": (832, 1248),
}
DEFAULT_RESOLUTION = "1024x1024 (1:1)"

MODEL_DEFAULTS = {
    "turbo": dict(
        steps=9, steps_min=4, steps_max=12,  # 9 requested -> 8 real DiT forwards
        guidance=0.0, guidance_min=0.0, guidance_max=1.0,
    ),
    "normal": dict(
        steps=50, steps_min=20, steps_max=60,
        guidance=4.0, guidance_min=1.0, guidance_max=10.0,
    ),
}

PIPES = {}
_load_lock = threading.Lock()


def is_turbo(model_key):
    return model_key == "turbo"


def load_pipe(model_key):
    """Return the pipeline for model_key, loading and caching it on first use.

    Deliberately a plain function, NOT decorated with @spaces.GPU: HF's own
    ZeroGPU guidance is that pipelines should be built and moved to CUDA at
    the root of the process (like the eager load below), not inside a
    @spaces.GPU function - that's where the CUDA-call interception that
    lets `.to("cuda")` work without a real attached GPU is designed to run.
    Loading inside a @spaces.GPU function is supported but discouraged
    (slower transfers), so button clicks land here instead.
    https://huggingface.co/docs/hub/spaces-zerogpu
    """
    if model_key in PIPES:
        return PIPES[model_key]

    with _load_lock:
        if model_key in PIPES: 
            return PIPES[model_key]

        if UNLOAD_ON_SWITCH:
            for key in list(PIPES):
                del PIPES[key]
            gc.collect()
            try:
                torch.cuda.empty_cache()
            except Exception:
                pass

        model_id = MODEL_IDS[model_key]
        print(f"Loading {model_id} ...")
        pipe = ZImagePipeline.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16,
            low_cpu_mem_usage=False,
        )
        pipe.to("cuda")
        PIPES[model_key] = pipe
        return pipe

load_pipe(DEFAULT_MODEL)


def get_duration(model_key, *args, **kwargs):
    return 60 if is_turbo(model_key) else 120


@spaces.GPU(duration=get_duration)
def generate(
    model_key,
    prompt,
    negative_prompt,
    resolution,
    guidance_scale,
    num_inference_steps,
    shift,
    seed,
    randomize_seed,
    progress=gr.Progress(track_tqdm=True),
):
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    pipe = load_pipe(model_key)
    turbo = is_turbo(model_key)
    width, height = RESOLUTIONS[resolution]

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    generator = torch.Generator(device="cuda").manual_seed(seed)

    pipe.scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=float(shift))

    kwargs = dict(
        prompt=prompt,
        height=height,
        width=width,
        num_inference_steps=int(num_inference_steps),
        generator=generator,
        max_sequence_length=512,
    )

    if turbo:
        kwargs["guidance_scale"] = 0.0  # Turbo is distilled without CFG - must stay 0
    else:
        kwargs["guidance_scale"] = float(guidance_scale)
        kwargs["negative_prompt"] = negative_prompt or ""
        kwargs["cfg_normalization"] = False

    try:
        image = pipe(**kwargs).images[0]
    except Exception as e:
        raise gr.Error(f"Generation failed: {e}")

    return image, seed


def model_description(model_key):
    model_id = MODEL_IDS[model_key]
    name = model_id.split("/")[-1]
    blurb = (
        "This is the fast, distilled **Turbo** model — 8 steps, no negative prompt / CFG."
        if is_turbo(model_key)
        else "This is the full, undistilled base model — supports classifier-free guidance and negative prompts."
    )
    return f"""# {name}
Text-to-image generation with **[{model_id}](https://huggingface.co/{model_id})**.
{blurb}"""


def switch_model(model_key, progress):
    """Button handler: lazy-load model_key if needed, then refresh every
    piece of UI that depends on which model is currently active."""
    if model_key not in PIPES:
        gr.Info(f"Loading {MODEL_IDS[model_key]} for the first time — this can take a minute...")
        progress(0, desc="Loading model weights...")
        load_pipe(model_key)

    turbo = is_turbo(model_key)
    d = MODEL_DEFAULTS[model_key]

    return (
        model_key,                                                # model_state
        gr.update(value=model_description(model_key)),            # header
        gr.update(variant="primary" if turbo else "secondary"),   # turbo_btn
        gr.update(variant="secondary" if turbo else "primary"),   # normal_btn
        gr.update(visible=not turbo),                              # negative_prompt
        gr.update(
            minimum=d["guidance_min"], maximum=d["guidance_max"],
            value=d["guidance"], interactive=not turbo,
        ),                                                          # guidance_scale
        gr.update(
            minimum=d["steps_min"], maximum=d["steps_max"],
            value=d["steps"], interactive=not turbo,
        ),                                                          # num_inference_steps
    )


def switch_to_turbo(progress=gr.Progress(track_tqdm=True)):
    return switch_model("turbo", progress)


def switch_to_normal(progress=gr.Progress(track_tqdm=True)):
    return switch_model("normal", progress)


with gr.Blocks(title="Z-Image") as demo:
    model_state = gr.State(DEFAULT_MODEL)

    header = gr.Markdown(model_description(DEFAULT_MODEL))

    with gr.Row():
        turbo_btn = gr.Button(
            "⚡ Turbo (fast)",
            variant="primary" if is_turbo(DEFAULT_MODEL) else "secondary",
        )
        normal_btn = gr.Button(
            "🎨 Normal (full quality)",
            variant="secondary" if is_turbo(DEFAULT_MODEL) else "primary",
        )

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt", placeholder="Describe the image you want...", lines=3
            )
            negative_prompt = gr.Textbox(
                label="Negative prompt",
                placeholder="Things to avoid (optional)",
                lines=2,
                visible=not is_turbo(DEFAULT_MODEL),
            )
            resolution = gr.Dropdown(
                choices=list(RESOLUTIONS.keys()), value=DEFAULT_RESOLUTION, label="Resolution"
            )

            with gr.Accordion("Advanced settings", open=False):
                _d = MODEL_DEFAULTS[DEFAULT_MODEL]
                guidance_scale = gr.Slider(
                    minimum=_d["guidance_min"],
                    maximum=_d["guidance_max"],
                    value=_d["guidance"],
                    step=0.1,
                    label="Guidance scale (CFG)",
                    interactive=not is_turbo(DEFAULT_MODEL),
                )
                num_inference_steps = gr.Slider(
                    minimum=_d["steps_min"],
                    maximum=_d["steps_max"],
                    value=_d["steps"],
                    step=1,
                    label="Inference steps",
                    interactive=not is_turbo(DEFAULT_MODEL),
                )
                shift = gr.Slider(
                    minimum=1.0, maximum=10.0, value=3.0, step=0.1, label="Time shift"
                )
                with gr.Row():
                    seed = gr.Number(label="Seed", value=42, precision=0, minimum=0, maximum=MAX_SEED)
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

            generate_btn = gr.Button("Generate", variant="primary")

            gr.Examples(
                examples=[
                    ["Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"],
                    ["A cozy reading nook by a rainy window, warm lamp light, watercolor style"],
                    ["A wooden bakery signboard reading '新鲜出炉' in bold red characters, morning light"],
                ],
                inputs=prompt,
            )

        with gr.Column(scale=1):
            output_image = gr.Image(label="Result", format="png")
            used_seed = gr.Number(label="Seed used", interactive=False)

    switch_outputs = [
        model_state, header, turbo_btn, normal_btn,
        negative_prompt, guidance_scale, num_inference_steps,
    ]
    turbo_btn.click(fn=switch_to_turbo, inputs=None, outputs=switch_outputs)
    normal_btn.click(fn=switch_to_normal, inputs=None, outputs=switch_outputs)

    event_inputs = [
        model_state,
        prompt,
        negative_prompt,
        resolution,
        guidance_scale,
        num_inference_steps,
        shift,
        seed,
        randomize_seed,
    ]
    event_outputs = [output_image, used_seed]

    generate_btn.click(fn=generate, inputs=event_inputs, outputs=event_outputs)
    prompt.submit(fn=generate, inputs=event_inputs, outputs=event_outputs)

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