Add MCP support

Closes #6

app.py

@@ -365,8 +365,56 @@ def prompt_enhance(prompt, enable_enhance):
 
 @spaces.GPU
 def generate(
-    prompt, resolution, seed, steps, shift, enhance, random_seed, gallery_images, progress=gr.Progress(track_tqdm=True)
+    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)
 ):
+    """
+    Generate an 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.
+
+    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)
+
+    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
+    """
     if pipe is None:
         raise gr.Error("Model not loaded.")
 
@@ -465,7 +513,7 @@ 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)
+    res_cat.change(update_res_choices, inputs=res_cat, outputs=resolution, api_visibility="private")
 
     # PE enhancement button (Temporarily disabled)
     # enhance_btn.click(
@@ -481,10 +529,11 @@ with gr.Blocks(title="Z-Image Demo") as demo:
         generate,
         inputs=[prompt_input, resolution, seed, steps, shift, enable_enhance, random_seed, output_gallery],
         outputs=[output_gallery, used_seed, seed],
+        api_visibility="public",
     )
 
 css='''
 .fillable{max-width: 1230px !important}
 '''
 if __name__ == "__main__":
-    demo.launch(css=css)
+    demo.launch(css=css, mcp_server=True)