Refactor and clean up Gradio UI and generation logic

Refactored the generate function to return a single image instead of a gallery, updated the Gradio UI to use an Image component instead of a Gallery, and improved code formatting for readability. Enhanced docstrings and argument descriptions, and made minor style and error handling improvements throughout the file.

app.py

@@ -57,7 +57,7 @@ RES_CHOICES = {
         "1536x1024 ( 3:2 )",
         "1024x1536 ( 2:3 )",
         "1600x896 ( 16:9 )",
-        "896x1600 ( 9:16 )", # not 900 coz divided by 16 needed
+        "896x1600 ( 9:16 )",  # not 900 coz divided by 16 needed
         "1680x720 ( 21:9 )",
         "720x1680 ( 9:21 )",
     ],
@@ -119,10 +119,14 @@ def load_models(model_path, enable_compile=False, attention_backend="native"):
             use_auth_token=use_auth_token,
         ).eval()
 
-        tokenizer = AutoTokenizer.from_pretrained(f"{model_path}", subfolder="tokenizer", use_auth_token=use_auth_token)
+        tokenizer = AutoTokenizer.from_pretrained(
+            f"{model_path}", subfolder="tokenizer", use_auth_token=use_auth_token
+        )
     else:
         vae = AutoencoderKL.from_pretrained(
-            os.path.join(model_path, "vae"), torch_dtype=torch.bfloat16, device_map="cuda"
+            os.path.join(model_path, "vae"),
+            torch_dtype=torch.bfloat16,
+            device_map="cuda",
         )
 
         text_encoder = AutoModel.from_pretrained(
@@ -145,7 +149,13 @@ def load_models(model_path, enable_compile=False, attention_backend="native"):
         torch._inductor.config.max_autotune_gemm_backends = "TRITON,ATEN"
         torch._inductor.config.triton.cudagraphs = False
 
-    pipe = ZImagePipeline(scheduler=None, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, transformer=None)
+    pipe = ZImagePipeline(
+        scheduler=None,
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        transformer=None,
+    )
 
     if enable_compile:
         pipe.vae.disable_tiling()
@@ -155,16 +165,18 @@ def load_models(model_path, enable_compile=False, attention_backend="native"):
             f"{model_path}", subfolder="transformer", use_auth_token=use_auth_token
         ).to("cuda", torch.bfloat16)
     else:
-        transformer = ZImageTransformer2DModel.from_pretrained(os.path.join(model_path, "transformer")).to(
-            "cuda", torch.bfloat16
-        )
+        transformer = ZImageTransformer2DModel.from_pretrained(
+            os.path.join(model_path, "transformer")
+        ).to("cuda", torch.bfloat16)
 
     pipe.transformer = transformer
     pipe.transformer.set_attention_backend(attention_backend)
 
     if enable_compile:
         print("Compiling transformer...")
-        pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune-no-cudagraphs", fullgraph=False)
+        pipe.transformer = torch.compile(
+            pipe.transformer, mode="max-autotune-no-cudagraphs", fullgraph=False
+        )
 
     pipe.to("cuda", torch.bfloat16)
 
@@ -254,7 +266,9 @@ class APIPromptExpander(PromptExpander):
             from openai import OpenAI
 
             api_key = self.api_config.get("api_key") or DASHSCOPE_API_KEY
-            base_url = self.api_config.get("base_url", "https://dashscope.aliyuncs.com/compatible-mode/v1")
+            base_url = self.api_config.get(
+                "base_url", "https://dashscope.aliyuncs.com/compatible-mode/v1"
+            )
 
             if not api_key:
                 print("Warning: DASHSCOPE_API_KEY not found.")
@@ -273,7 +287,9 @@ class APIPromptExpander(PromptExpander):
 
     def extend(self, prompt, system_prompt=None, seed=-1, **kwargs):
         if self.client is None:
-            return PromptOutput(False, "", seed, system_prompt, "API client not initialized")
+            return PromptOutput(
+                False, "", seed, system_prompt, "API client not initialized"
+            )
 
         if system_prompt is None:
             system_prompt = self.decide_system_prompt()
@@ -286,7 +302,10 @@ class APIPromptExpander(PromptExpander):
             model = self.api_config.get("model", "qwen3-max-preview")
             response = self.client.chat.completions.create(
                 model=model,
-                messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}],
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": prompt},
+                ],
                 temperature=0.7,
                 top_p=0.8,
             )
@@ -305,7 +324,11 @@ class APIPromptExpander(PromptExpander):
                 expanded_prompt = content
 
             return PromptOutput(
-                status=True, prompt=expanded_prompt, seed=seed, system_prompt=system_prompt, message=content
+                status=True,
+                prompt=expanded_prompt,
+                seed=seed,
+                system_prompt=system_prompt,
+                message=content,
             )
         except Exception as e:
             return PromptOutput(False, "", seed, system_prompt, str(e))
@@ -325,7 +348,11 @@ def init_app():
     global pipe, prompt_expander
 
     try:
-        pipe = load_models(MODEL_PATH, enable_compile=ENABLE_COMPILE, attention_backend=ATTENTION_BACKEND)
+        pipe = load_models(
+            MODEL_PATH,
+            enable_compile=ENABLE_COMPILE,
+            attention_backend=ATTENTION_BACKEND,
+        )
         print(f"Model loaded. Compile: {ENABLE_COMPILE}, Backend: {ATTENTION_BACKEND}")
 
         if ENABLE_WARMUP:
@@ -339,7 +366,9 @@ def init_app():
         pipe = None
 
     try:
-        prompt_expander = create_prompt_expander(backend="api", api_config={"model": "qwen3-max-preview"})
+        prompt_expander = create_prompt_expander(
+            backend="api", api_config={"model": "qwen3-max-preview"}
+        )
         print("Prompt expander initialized.")
     except Exception as e:
         print(f"Error initializing prompt expander: {e}")
@@ -365,55 +394,34 @@ def prompt_enhance(prompt, enable_enhance):
 
 @spaces.GPU
 def generate(
-    prompt, resolution="1024x1024 ( 1:1 )", seed=42, steps=9, shift=3.0, enhance=False, random_seed=True, gallery_images=None, progress=gr.Progress(track_tqdm=True)
+    prompt: str,
+    resolution: str = "1024x1024 ( 1:1 )",
+    seed: int = 42,
+    steps: int = 9,
+    shift: float = 3.0,
+    enhance: bool = False,
+    random_seed: bool = True,
+    progress: gr.Progress = gr.Progress(track_tqdm=True),
 ):
     """
-    Generate an image using the Z-Image model based on the provided prompt and settings.
+    Generate a single image using the Z-Image model based on the provided prompt and settings.
 
-    This function is triggered when the user clicks the "Generate" button. It processes
-    the input prompt (optionally enhancing it), configures generation parameters, and
-    produces an image using the Z-Image diffusion transformer pipeline.
+    This function is exposed as a Gradio/MCP tool via the main 'Generate' button.
+    It optionally enhances the prompt, configures generation parameters, and
+    returns exactly one image plus the seed used.
 
     Args:
-        prompt (str): Text prompt describing the desired image content
-        resolution (str): Output resolution in format "WIDTHxHEIGHT ( RATIO )" (e.g., "1024x1024 ( 1:1 )")
-            valid options, 1024 category:
-                    - "1024x1024 ( 1:1 )"
-                    - "1152x896 ( 9:7 )"
-                    - "896x1152 ( 7:9 )"
-                    - "1152x864 ( 4:3 )"
-                    - "864x1152 ( 3:4 )"
-                    - "1248x832 ( 3:2 )"
-                    - "832x1248 ( 2:3 )"
-                    - "1280x720 ( 16:9 )"
-                    - "720x1280 ( 9:16 )"
-                    - "1344x576 ( 21:9 )"
-                    - "576x1344 ( 9:21 )"
-                1280 category:
-                    - "1280x1280 ( 1:1 )"
-                    - "1440x1120 ( 9:7 )"
-                    - "1120x1440 ( 7:9 )"
-                    - "1472x1104 ( 4:3 )"
-                    - "1104x1472 ( 3:4 )"
-                    - "1536x1024 ( 3:2 )"
-                    - "1024x1536 ( 2:3 )"
-                    - "1600x896 ( 16:9 )"
-                    - "896x1600 ( 9:16 )"
-                    - "1680x720 ( 21:9 )"
-                    - "720x1680 ( 9:21 )"
-        seed (int): Seed for reproducible generation
-        steps (int): Number of inference steps for the diffusion process
-        shift (float): Time shift parameter for the flow matching scheduler
-        enhance (bool): This was Whether to enhance the prompt (DISABLED! Do not use)
-        random_seed (bool): Whether to generate a new random seed, if True will ignore the seed input
-        gallery_images (list): List of previously generated images to append to (only needed for the Gradio UI)
-        progress (gr.Progress): Gradio progress tracker for displaying generation progress (only needed for the Gradio UI)
+        prompt: Text prompt describing the desired image content.
+        resolution: Output resolution in format "WIDTHxHEIGHT ( RATIO )".
+        seed: Seed for reproducible generation. Ignored if random_seed is True.
+        steps: Number of inference steps for the diffusion process.
+        shift: Time shift parameter for the flow matching scheduler.
+        enhance: (Currently disabled in the UI) Whether to enhance the prompt.
+        random_seed: If True, a new random seed will be sampled.
+        progress: Gradio progress tracker (automatically provided by Gradio).
 
     Returns:
-        tuple: (gallery_images, seed_str, seed_int)
-            - gallery_images: Updated list of generated images including the new image
-            - seed_str: String representation of the seed used for generation
-            - seed_int: Integer representation of the seed used for generation
+        tuple[object, str, int]: (image, seed_str, seed_int)
     """
     if pipe is None:
         raise gr.Error("Model not loaded.")
@@ -431,7 +439,7 @@ def generate(
 
     try:
         resolution_str = resolution.split(" ")[0]
-    except:
+    except Exception:
         resolution_str = "1024x1024"
 
     image = generate_image(
@@ -444,11 +452,7 @@ def generate(
         shift=shift,
     )
 
-    if gallery_images is None:
-        gallery_images = []
-    gallery_images.append(image)
-
-    return gallery_images, str(new_seed), int(new_seed)
+    return image, str(new_seed), int(new_seed)
 
 
 init_app()
@@ -473,7 +477,9 @@ with gr.Blocks(title="Z-Image Demo") as demo:
 
     with gr.Row():
         with gr.Column(scale=1):
-            prompt_input = gr.Textbox(label="Prompt", lines=3, placeholder="Enter your prompt here...")
+            prompt_input = gr.Textbox(
+                label="Prompt", lines=3, placeholder="Enter your prompt here..."
+            )
             # PE components (Temporarily disabled)
             # with gr.Row():
             #     enable_enhance = gr.Checkbox(label="Enhance Prompt (DashScope)", value=False)
@@ -481,18 +487,33 @@ with gr.Blocks(title="Z-Image Demo") as demo:
 
             with gr.Row():
                 choices = [int(k) for k in RES_CHOICES.keys()]
-                res_cat = gr.Dropdown(value=1024, choices=choices, label="Resolution Category")
+                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=RESOLUTION_SET, label="Width x Height (Ratio)")
+                resolution = gr.Dropdown(
+                    value=initial_res_choices[0],
+                    choices=RESOLUTION_SET,
+                    label="Width x Height (Ratio)",
+                )
 
             with gr.Row():
                 seed = gr.Number(label="Seed", value=42, 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=False)
-                shift = gr.Slider(label="Time Shift", minimum=1.0, maximum=10.0, value=3.0, step=0.1)
+                steps = gr.Slider(
+                    label="Steps",
+                    minimum=1,
+                    maximum=100,
+                    value=8,
+                    step=1,
+                    interactive=False,
+                )
+                shift = gr.Slider(
+                    label="Time Shift", minimum=1.0, maximum=10.0, value=3.0, step=0.1
+                )
 
             generate_btn = gr.Button("Generate", variant="primary")
 
@@ -501,8 +522,12 @@ with gr.Blocks(title="Z-Image Demo") as demo:
             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
+            # Switched from Gallery -> single Image for MCP-friendly output
+            output_image = gr.Image(
+                label="Generated Image",
+                format="png",
+                height=600,
+                interactive=False,
             )
             used_seed = gr.Textbox(label="Seed Used", interactive=False)
 
@@ -513,27 +538,31 @@ with gr.Blocks(title="Z-Image Demo") as demo:
             res_choices = RES_CHOICES["1024"]
         return gr.update(value=res_choices[0], choices=res_choices)
 
-    res_cat.change(update_res_choices, inputs=res_cat, outputs=resolution, api_visibility="private")
-
-    # PE enhancement button (Temporarily disabled)
-    # enhance_btn.click(
-    #     prompt_enhance,
-    #     inputs=[prompt_input, enable_enhance],
-    #     outputs=[prompt_input, final_prompt_output]
-    # )
+    res_cat.change(
+        update_res_choices, inputs=res_cat, outputs=resolution, api_visibility="private"
+    )
 
     # Dummy enable_enhance variable set to False
     enable_enhance = gr.State(value=False)
 
     generate_btn.click(
         generate,
-        inputs=[prompt_input, resolution, seed, steps, shift, enable_enhance, random_seed, output_gallery],
-        outputs=[output_gallery, used_seed, seed],
-        api_visibility="public",
+        inputs=[
+            prompt_input,
+            resolution,
+            seed,
+            steps,
+            shift,
+            enable_enhance,
+            random_seed,
+        ],
+        outputs=[output_image, used_seed, seed],
+        api_visibility="public",  # exposed as MCP tool
+        api_name="generate_image",  # nice, stable name for tool clients
     )
 
-css='''
+css = """
 .fillable{max-width: 1230px !important}
-'''
+"""
 if __name__ == "__main__":
-    demo.launch(css=css, mcp_server=True)
+    demo.launch(css=css, mcp_server=True)