Update app.py

app.py

@@ -419,4 +419,138 @@ def generate(
     steps: int,
     shift: float,
     enhance: bool,
-    random_s_
+    random_seed: bool,
+    gallery_images,
+    progress=gr.Progress(track_tqdm=True),
+):
+    global pipe
+
+    # 进入该函数时，ZeroGPU 会尝试分配 GPU；若失败，torch.cuda.is_available 可能为 False
+    if pipe is None:
+        # 启动阶段不加载，第一次生成时再加载（CPU）
+        load_models_cpu_only(MODEL_PATH)
+
+    if torch.cuda.is_available():
+        move_pipe_to_device(torch.device("cuda"))
+        warmup_if_needed()
+    else:
+        # 在 ZeroGPU 排队/额度不足/未正确启用硬件时，给出友好报错
+        raise gr.Error(
+            "当前未获得 ZeroGPU 的 CUDA 资源（torch.cuda.is_available()==False）。"
+            "请确认 Space 硬件选择为 ZeroGPU，并在有额度/队列可用时重试。"
+        )
+
+    final_prompt = prompt
+    if enhance:
+        final_prompt, _ = prompt_enhance(prompt, True)
+
+    if random_seed:
+        seed = random.randint(1, 1_000_000)
+    else:
+        seed = int(seed) if int(seed) != -1 else random.randint(1, 1_000_000)
+
+    try:
+        resolution_str = resolution.split(" ")[0]  # "1024x1024 (1:1)" -> "1024x1024"
+    except Exception:
+        resolution_str = "1024x1024"
+
+    image, used_seed_str = generate_image(
+        prompt=final_prompt,
+        resolution=resolution_str,
+        seed=seed,
+        guidance_scale=0.0,
+        num_inference_steps=int(steps) + 1,
+        shift=float(shift),
+        progress=progress,
+    )
+
+    if gallery_images is None:
+        gallery_images = []
+    gallery_images.append(image)
+
+    return gallery_images, used_seed_str
+
+
+# -------------------- UI --------------------
+init_prompt_expander()
+
+with gr.Blocks(title="Z-Image Demo") as demo:
+    gr.Markdown(
+        """<div align="center">
+
+# Z-Image Generation Demo
+
+*An Efficient Image Generation Foundation Model with Single-Stream Diffusion Transformer*
+
+</div>"""
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            prompt_input = gr.Textbox(label="Prompt", lines=3, placeholder="Enter your prompt here...")
+
+            with gr.Row():
+                choices = [int(k) for k in RES_CHOICES.keys()]
+                res_cat = gr.Dropdown(value=1024, choices=choices, label="Resolution Category")
+                initial_res_choices = RES_CHOICES["1024"]
+                resolution = gr.Dropdown(
+                    value=initial_res_choices[0],
+                    choices=initial_res_choices,
+                    label="Width x Height (Ratio)",
+                )
+
+            with gr.Row():
+                seed = gr.Number(label="Seed", value=-1, precision=0)
+                random_seed = gr.Checkbox(label="Random Seed", value=True)
+
+            with gr.Row():
+                steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=8, step=1, interactive=True)
+                shift = gr.Slider(label="Time Shift", minimum=1.0, maximum=10.0, value=3.0, step=0.1)
+
+            enhance = gr.Checkbox(label="Enhance Prompt (DashScope)", value=False)
+
+            generate_btn = gr.Button("Generate", variant="primary")
+
+            gr.Markdown("### Example Prompts")
+            gr.Examples(examples=EXAMPLE_PROMPTS, inputs=prompt_input, label=None)
+
+        with gr.Column(scale=1):
+            output_gallery = gr.Gallery(
+                label="Generated Images",
+                columns=2,
+                rows=2,
+                height=600,
+                object_fit="contain",
+                format="png",
+                interactive=False,
+            )
+            used_seed = gr.Textbox(label="Seed Used", interactive=False)
+
+    def update_res_choices(_res_cat):
+        key = str(_res_cat)
+        res_choices = RES_CHOICES.get(key, RES_CHOICES["1024"])
+        return gr.update(value=res_choices[0], choices=res_choices)
+
+    res_cat.change(update_res_choices, inputs=res_cat, outputs=resolution)
+
+    def update_seed(current_seed, random_seed_enabled):
+        if random_seed_enabled:
+            new_seed = random.randint(1, 1_000_000)
+        else:
+            new_seed = int(current_seed) if int(current_seed) != -1 else random.randint(1, 1_000_000)
+        return gr.update(value=new_seed)
+
+    generate_btn.click(update_seed, inputs=[seed, random_seed], outputs=[seed])
+
+    generate_btn.click(
+        generate,
+        inputs=[prompt_input, resolution, seed, steps, shift, enhance, random_seed, output_gallery],
+        outputs=[output_gallery, used_seed],
+    )
+
+css = """
+.fillable { max-width: 1230px !important; }
+"""
+
+if __name__ == "__main__":
+    demo.queue().launch(css=css)