replace z-image

app.py

@@ -3,15 +3,13 @@ import numpy as np
 import random
 import os
 from PIL import Image
-import requests
-import json
+from gradio_client import Client
 from dotenv import load_dotenv
 
 # Load environment variables
 load_dotenv()
 
-# MCP Server configuration
-MCP_SERVER_URL = "https://multimodalart-qwen-image-fast.hf.space/gradio_api/mcp/"
+# Client configuration
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 # --- UI Constants and Helpers ---
@@ -20,32 +18,48 @@ MAX_SEED = np.iinfo(np.int32).max
 
 def create_prompt(params):
     """
-    MCP サーバーに送信するためのプロンプト文字列を生成する
+    Generate a prompt string for sending to the client
 
     Args:
-        params (dict): プロンプト生成に使用するパラメータ辞書
-            - label (str): ボタンテキスト
-            - detail (str): デザイン詳細
-            - shape (str): ボタン形状
-            - layout (str): レイアウト配置
-            - background (str): 背景設定
+        params (dict): Dictionary of parameters used for prompt generation
+            - label (str): Comma-separated button texts, will be split into array of 4 labels
+            - detail (str): Design details
+            - shape (str): Button shape
+            - layout (str): Layout arrangement
+            - background (str): Background setting
 
     Returns:
-        str: 画像生成用の完全なプロンプト文字列
+        str: Complete prompt string for image generation
     """
-    label = params.get("label", "Button")
+    label_str = params.get("label", "Button")
     detail = params.get("detail", "")
     shape = params.get("shape", "rounded")
     layout = params.get("layout", "horizontal_3")
     background = params.get("background", "natural")
 
+    # Split labels by comma, take up to number of positions
+    labels = [l.strip() for l in label_str.split(",") if l.strip()]
+    if labels:
+        # If there are labels, repeat to up to 4
+        while len(labels) < 4:
+            labels.extend(labels)
+
     layout_config = {
-        "horizontal_3": {"count": 3, "desc": "3 horizontal rows"},
+        "horizontal_3": {
+            "positions": ["top", "middle", "bottom"],
+            "desc": "3 horizontal rows: top, middle, bottom",
+        },
         "vertical_2": {
-            "count": 2,
-            "desc": " 2 vertical tall buttons arranged side by side, vertical rectangular shape, 2 columns layout, slender buttons, portrait orientation",
+            "positions": [
+                "left middle-valign small text",
+                "right middle-valign small text",
+            ],
+            "desc": "2 vertical tall buttons arranged side by side, vertical rectangular shape, 2 columns layout, horizontal text, small text orientation is horizontal",
+        },
+        "box_2x2": {
+            "positions": ["left-top", "right-top", "left-bottom", "right-bottom"],
+            "desc": "2x2 grid: left-top, right-top, left-bottom, right-bottom",
         },
-        "box_2x2": {"count": 4, "desc": "2x2 grid arrangement"},
     }
     shape_descriptions = {
         "box": "box shape",
@@ -55,18 +69,23 @@ def create_prompt(params):
     }
 
     config = layout_config.get(layout, layout_config["horizontal_3"])
-    count = config["count"]
+    positions = config["positions"]
     layout_desc = config["desc"]
     shape_desc = shape_descriptions.get(shape, "rounded corners")
 
-    if label and label.strip():
-        base_prompt = f'Create {count} {detail} button design variations with "{label}" text in a 1024x1024 image.\n'
-    else:
-        base_prompt = f"Create {count} {detail} button design variations with no text label in a 1024x1024 image.\n"
+    base_prompt = (
+        f"Create {len(positions)} {detail} button designs in a 1024x1024 image.\n"
+    )
     base_prompt += f"Arranged in {layout_desc}.\n"
     base_prompt += f"{shape_desc.capitalize()}.\n"
     if detail:
         base_prompt += f"{detail.capitalize()} aesthetic: detailed visual elements and color palette.\n"
+    # Describe each button with position and label
+    for i, pos in enumerate(positions):
+        if i < len(labels):
+            base_prompt += f'{pos} button: "{labels[i]}".\n'
+        else:
+            base_prompt += f"{pos} button: empty text rectangle.\n"
     base_prompt += "Each button is a different design variation exploring the theme.\n"
     if background == "natural":
         base_prompt += "Natural background with subtle textures and ambient lighting.\n"
@@ -80,91 +99,50 @@ def create_prompt(params):
     return base_prompt
 
 
-def call_mcp_server_http(
+def call_client(
     prompt,
     seed,
     randomize_seed,
     aspect_ratio,
-    guidance_scale,
     num_inference_steps,
-    prompt_enhance=False,
 ):
     """
-    MCP (Model Context Protocol) サーバーに HTTP リクエストを送信して画像生成を行う
+    Call the gradio client for image generation
 
     Args:
-        prompt (str): 画像生成のためのプロンプトテキスト
-        seed (int): 乱数シード値
-        randomize_seed (bool): シードをランダム化するかどうか
-        aspect_ratio (str): 画像のアスペクト比（例: "1:1", "16:9", "4:3"）
-        guidance_scale (float): ガイダンススケール（1.0-5.0）
-        num_inference_steps (int): 推論ステップ数（4-28）
-        prompt_enhance (bool, optional): プロンプトエンハンスを使用するかどうか. Defaults to False.
+        prompt (str): Prompt text for image generation
+        seed (int): Random seed value
+        randomize_seed (bool): Whether to randomize the seed
+        aspect_ratio (str): Image aspect ratio (e.g., "1:1", "16:9", "4:3")
+        num_inference_steps (int): Number of inference steps (4-28)
 
     Returns:
-        tuple: (image, seed, error) - 画像オブジェクト、使用されたシード、エラーメッセージ
+        tuple: (image, seed, error) - Image object, used seed, error message
     """
-    import base64
-    import io
-
-    print(f"Calling MCP Endpoint: {MCP_SERVER_URL}")
     try:
-        payload = {
-            "jsonrpc": "2.0",
-            "method": "tools/call",
-            "params": {
-                "name": "Qwen_Image_Fast_infer",
-                "arguments": {
-                    "prompt": prompt,
-                    "seed": seed,
-                    "randomize_seed": randomize_seed,
-                    "aspect_ratio": aspect_ratio,
-                    "guidance_scale": guidance_scale,
-                    "num_inference_steps": num_inference_steps,
-                    "prompt_enhance": prompt_enhance,
-                },
-            },
-            "id": 1,
-        }
-        headers = {
-            "Content-Type": "application/json",
-            "Accept": "application/json, text/event-stream",
-            "User-Agent": "mcp-client-python/1.0",
+        # Map aspect_ratio to resolution
+        resolution_map = {
+            "1:1": "1024x1024 ( 1:1 )",
+            "16:9": "1280x720 ( 16:9 )",
+            "4:3": "1152x864 ( 4:3 )",
         }
-        if HF_TOKEN:
-            headers["Authorization"] = f"Bearer {HF_TOKEN}"
-
-        response = requests.post(
-            MCP_SERVER_URL, headers=headers, data=json.dumps(payload), timeout=120
+        resolution = resolution_map.get(aspect_ratio, "1024x1024 ( 1:1 )")
+
+        client = Client("Tongyi-MAI/Z-Image-Turbo")
+        result = client.predict(
+            prompt=prompt,
+            resolution=resolution,
+            seed=seed,
+            steps=num_inference_steps,
+            shift=3.0,
+            random_seed=randomize_seed,
+            gallery_images=[],
+            api_name="/generate",
         )
-
-        if response.status_code == 200:
-            try:
-                result_json = response.json()
-            except json.JSONDecodeError:
-                lines = response.text.strip().split("\n")
-                last_line = lines[-1]
-                if last_line.startswith("data: "):
-                    result_json = json.loads(last_line[6:])
-                else:
-                    result_json = json.loads(last_line)
-
-            if "error" in result_json:
-                return None, seed, str(result_json["error"])
-
-            if "result" in result_json:
-                content_list = result_json["result"].get("content", [])
-                for item in content_list:
-                    if item.get("type") == "image" and item.get("data"):
-                        image_bytes = base64.b64decode(item.get("data"))
-                        return Image.open(io.BytesIO(image_bytes)), seed, None
-            else:
-                return None, seed, "Unexpected JSON structure"
-        else:
-            return None, seed, f"Server Error: {response.status_code}"
+        # Assume result is PIL Image
+        return result, seed, None
     except Exception as e:
         return None, seed, str(e)
-    return Image.new("RGB", (1024, 1024), color="gray"), seed, "Unknown error"
 
 
 # --- Main Inference Logic ---
@@ -180,25 +158,27 @@ def run_inference_engine(
     aspect_ratio,
     guidance_scale,
     num_inference_steps,
+    hf_token: gr.Header = None,
 ):
     """
-    MCP サーバーを使用して UI ボタン画像を生成するメイン推論エンジン
+    generate UI button images
 
     Args:
-        label (str): ボタンに表示するテキスト
-        detail (str): デザインの詳細プロンプト
-        shape (str): ボタンの形状（"box", "rounded", "oval", "free"）
-        layout (str): レイアウト配置（"horizontal_3", "vertical_2", "box_2x2"）
-        background (str): 背景設定（"natural", "white", "black"）
-        seed (int): 乱数シード値
-        randomize_seed (bool): シードをランダム化するかどうか
-        aspect_ratio (str): 画像のアスペクト比 (1:1を使うこと)
-        guidance_scale (float): ガイダンススケール(常に1を使うこと)
-        num_inference_steps (int): 推論ステップ数 (常に8を使うこと)
+        label (str): Text to display on the button
+        detail (str): Detailed design prompt
+        shape (str): Button shape ("box", "rounded", "oval", "free")
+        layout (str): Layout arrangement ("horizontal_3", "vertical_2", "box_2x2")
+        background (str): Background setting ("natural", "white", "black")
+        seed (int): Random seed value
+        randomize_seed (bool): Whether to randomize the seed
+        aspect_ratio (str): Image aspect ratio (use 1:1)
+        guidance_scale (float): Guidance scale (use 1)
+        num_inference_steps (int): Number of inference steps (use 8)
 
     Yields:
-        tuple: (image, seed, status_message) - 生成された画像、使用されたシード、ステータスメッセージ
+        tuple: (image, seed, status_message) - Generated image, used seed, status message
     """
+    print(hf_token)
     yield None, seed, "Generating..."
     prompt_params = {
         "label": label,
@@ -209,18 +189,23 @@ def run_inference_engine(
     }
     prompt = create_prompt(prompt_params)
 
+    # Debug: Print the generated prompt
+    print(f"Generated prompt: {prompt}")
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    image, generated_seed, error = call_mcp_server_http(
-        prompt, seed, randomize_seed, aspect_ratio, guidance_scale, num_inference_steps
+    image, generated_seed, error = call_client(
+        prompt, seed, randomize_seed, aspect_ratio, num_inference_steps
     )
+    print(type(image))
+    print(image)
 
     if image is None:
         gr.Warning(f"Error: {error}")
         yield None, generated_seed, f"Error: {error}"
     else:
-        yield image, generated_seed, "Done!"
+        yield image[0][0]["image"], generated_seed, "Done!"
 
 
 # --- UI Customization ---
@@ -230,10 +215,10 @@ body { font-family: 'Helvetica Neue', Arial, sans-serif; }
 h1 { text-align: center; font-weight: 800; color: #333; margin-bottom: 0.5em; }
 .subtitle { text-align: center; color: #666; margin-bottom: 2em; }
 
-.generate-btn { 
-    background: linear-gradient(90deg, #6366f1 0%, #a855f7 100%) !important; 
-    border: none !important; 
-    color: white !important; 
+.generate-btn {
+    background: linear-gradient(90deg, #6366f1 0%, #a855f7 100%) !important;
+    border: none !important;
+    color: white !important;
     font-weight: bold !important;
     font-size: 1.2em !important;
     padding: 20px !important;
@@ -247,15 +232,15 @@ h1 { text-align: center; font-weight: 800; color: #333; margin-bottom: 0.5em; }
     box-shadow: 0 5px 15px rgba(99, 102, 241, 0.4);
 }
 
-.examples-container table { 
-    font-size: 0.85em !important; 
+.examples-container table {
+    font-size: 0.85em !important;
     margin-bottom: 0 !important;
 }
-.examples-container td { 
-    padding: 4px 8px !important; 
-    white-space: nowrap; 
-    overflow: hidden; 
-    text-overflow: ellipsis; 
+.examples-container td {
+    padding: 4px 8px !important;
+    white-space: nowrap;
+    overflow: hidden;
+    text-overflow: ellipsis;
     max-width: 150px;
 }
 .examples-container label {
@@ -265,9 +250,9 @@ h1 { text-align: center; font-weight: 800; color: #333; margin-bottom: 0.5em; }
     display: block;
 }
 
-#result-image { 
-    border-radius: 12px; 
-    box-shadow: 0 10px 30px rgba(0,0,0,0.1); 
+#result-gallery {
+    border-radius: 12px;
+    box-shadow: 0 10px 30px rgba(0,0,0,0.1);
     border: 1px solid #eee;
 }
 .input-group {
@@ -289,7 +274,7 @@ example_data = [
         "box",
         "horizontal_3",
         "black",
-        "examples/play.webp",
+        # "examples/play.webp",
     ],
     [
         "Buy",
@@ -297,7 +282,7 @@ example_data = [
         "rounded",
         "vertical_2",
         "white",
-        "examples/buy.webp",
+        # "examples/buy.webp",
     ],
     [
         "RPG",
@@ -305,7 +290,7 @@ example_data = [
         "free",
         "box_2x2",
         "natural",
-        "examples/rpg.webp",
+        # "examples/rpg.webp",
     ],
     [
         "Submit",
@@ -320,7 +305,7 @@ example_data = [
 with gr.Blocks(css=css, theme=theme, title="UI Button Generator") as demo:
     gr.Markdown("# 🎨 AI UI Button Generator")
     gr.Markdown(
-        "<div class='subtitle'>UI button material and design concept generation tool using Qwen Image Fast</div>"
+        "<div class='subtitle'>UI button material and design concept generation tool using Z-Image-Turbo</div>"
     )
 
     with gr.Column(elem_id="col-container"):
@@ -398,7 +383,7 @@ with gr.Blocks(css=css, theme=theme, title="UI Button Generator") as demo:
         # --- 2. Layout Construction Phase (.render()) ---
 
         with gr.Row(equal_height=False):
-            # 左カラム
+            # Left column
             with gr.Column(scale=1, min_width=400):
                 with gr.Group(elem_classes="input-group"):
                     gr.Markdown("### 📝 Basic Settings")
@@ -426,18 +411,18 @@ with gr.Blocks(css=css, theme=theme, title="UI Button Generator") as demo:
                         guidance_scale.render()
                         num_inference_steps.render()
 
-            # 右カラム
+            # Right column
             with gr.Column(scale=1):
                 gr.Markdown("### 🖼️ Generated Button")
                 result.render()
                 status_msg.render()
-                # ボタンとExamplesの行
+                # Row for button and Examples
                 # gr.Markdown("**Quick Presets (Click to try)**")
                 # Initialize here with fn, inputs, outputs
                 gr.Examples(
                     examples=example_data,
                     fn=run_inference_engine,
-                    inputs=[label, detail, shape, layout, background, result],
+                    inputs=[label, detail, shape, layout, background],
                     outputs=[result, seed, status_msg],
                     examples_per_page=3,
                     run_on_click=False,
@@ -465,4 +450,7 @@ with gr.Blocks(css=css, theme=theme, title="UI Button Generator") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True)
+    demo.launch(
+        show_error=True,
+        mcp_server=True,
+    )