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