v22: Add prompt polishing feature (AI-enhanced prompts)

app.py

@@ -1,5 +1,6 @@
-"""Z-Image-Turbo v21 - Stable baseline (no AoTI)"""
+"""Z-Image-Turbo v22 - With prompt polishing feature"""
 
+import os
 import torch
 import spaces
 import gradio as gr
@@ -7,12 +8,44 @@ import requests
 import io
 from PIL import Image
 from diffusers import DiffusionPipeline, ZImageImg2ImgPipeline
+from huggingface_hub import InferenceClient
 
 # Enable optimized backends
 torch.backends.cuda.enable_flash_sdp(True)
 torch.backends.cuda.enable_mem_efficient_sdp(True)
 torch.backends.cudnn.benchmark = True
 
+# Prompt polishing using HF Inference API
+def polish_prompt(original_prompt):
+    """Expand short prompts into detailed, high-quality prompts using AI."""
+    if not original_prompt or not original_prompt.strip():
+        return "Ultra HD, 4K, cinematic composition, highly detailed"
+    
+    api_key = os.environ.get("HF_TOKEN")
+    if not api_key:
+        return original_prompt  # Return original if no token
+    
+    system_prompt = """You are a prompt optimizer for AI image generation. 
+Rewrite the user's input into a detailed, expressive prompt that will produce stunning images.
+Keep it under 150 words. Be descriptive about lighting, atmosphere, style, and details.
+Do not explain - just output the improved prompt directly."""
+    
+    try:
+        client = InferenceClient(api_key=api_key)
+        completion = client.chat.completions.create(
+            model="Qwen/Qwen2.5-72B-Instruct",
+            max_tokens=200,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": original_prompt}
+            ],
+        )
+        polished = completion.choices[0].message.content
+        return polished.strip().replace("\n", " ")
+    except Exception as e:
+        print(f"Prompt polish error: {e}")
+        return original_prompt
+
 print("Loading Z-Image-Turbo pipeline...")
 
 # Load text-to-image pipeline
@@ -113,17 +146,35 @@ def upload_to_hf_cdn(image):
     except Exception as e:
         return f"Error: {str(e)}"
 
+def prepare_prompt(prompt, style, do_polish):
+    """Prepare the final prompt with optional polishing and style."""
+    if not prompt or not prompt.strip():
+        return "", ""
+    
+    base_prompt = prompt.strip()
+    
+    # Polish if enabled
+    if do_polish:
+        polished = polish_prompt(base_prompt)
+    else:
+        polished = base_prompt
+    
+    # Add style suffix
+    final_prompt = polished + STYLE_SUFFIXES.get(style, "")
+    
+    return final_prompt, polished
+
 @spaces.GPU
-def generate(prompt, style, ratio, steps, seed, randomize, progress=gr.Progress(track_tqdm=True)):
+def generate(prompt, style, ratio, steps, seed, randomize, do_polish, progress=gr.Progress(track_tqdm=True)):
     if randomize:
         seed = torch.randint(0, 2**32 - 1, (1,)).item()
     seed = int(seed)
     
     if not prompt or not prompt.strip():
-        return None, seed
+        return None, seed, ""
     
     w, h = RATIO_DIMS.get(ratio, (1024, 1024))
-    full_prompt = prompt.strip() + STYLE_SUFFIXES.get(style, "")
+    full_prompt, polished_prompt = prepare_prompt(prompt, style, do_polish)
     
     generator = torch.Generator("cuda").manual_seed(seed)
     image = pipe_t2i(
@@ -135,7 +186,7 @@ def generate(prompt, style, ratio, steps, seed, randomize, progress=gr.Progress(
         generator=generator,
     ).images[0]
     
-    return image, seed
+    return image, seed, polished_prompt if do_polish else ""
 
 @spaces.GPU
 def transform(input_image, prompt, style, strength, steps, seed, randomize, progress=gr.Progress(track_tqdm=True)):
@@ -265,6 +316,7 @@ with gr.Blocks(title="Z-Image Turbo", css=css) as demo:
             with gr.Row():
                 with gr.Column():
                     gen_prompt = gr.Textbox(label="Prompt", placeholder="Describe your image...", lines=3)
+                    gen_polish = gr.Checkbox(label="✨ Polish Prompt (AI-enhanced)", value=False)
                     gen_style = gr.Dropdown(choices=STYLES, value="None", label="Style")
                     gen_ratio = gr.Dropdown(choices=RATIOS, value="1:1 Square (1024x1024)", label="Aspect Ratio")
                     gen_steps = gr.Slider(minimum=4, maximum=16, value=8, step=1, label="Steps")
@@ -275,6 +327,7 @@ with gr.Blocks(title="Z-Image Turbo", css=css) as demo:
                 
                 with gr.Column():
                     gen_output = gr.Image(label="Generated Image", type="pil", format="png", interactive=False)
+                    gen_polished_prompt = gr.Textbox(label="Polished Prompt", interactive=False, visible=True, lines=2)
                     gen_seed_out = gr.Number(label="Seed Used", interactive=False)
                     with gr.Row():
                         gen_share_btn = gr.Button("📤 Share Image Link", variant="secondary")
@@ -282,8 +335,8 @@ with gr.Blocks(title="Z-Image Turbo", css=css) as demo:
             
             gr.Examples(examples=EXAMPLES_GENERATE, inputs=[gen_prompt, gen_style, gen_ratio, gen_steps, gen_seed, gen_randomize])
             
-            gen_btn.click(fn=generate, inputs=[gen_prompt, gen_style, gen_ratio, gen_steps, gen_seed, gen_randomize], outputs=[gen_output, gen_seed_out])
-            gen_prompt.submit(fn=generate, inputs=[gen_prompt, gen_style, gen_ratio, gen_steps, gen_seed, gen_randomize], outputs=[gen_output, gen_seed_out])
+            gen_btn.click(fn=generate, inputs=[gen_prompt, gen_style, gen_ratio, gen_steps, gen_seed, gen_randomize, gen_polish], outputs=[gen_output, gen_seed_out, gen_polished_prompt])
+            gen_prompt.submit(fn=generate, inputs=[gen_prompt, gen_style, gen_ratio, gen_steps, gen_seed, gen_randomize, gen_polish], outputs=[gen_output, gen_seed_out, gen_polished_prompt])
             gen_share_btn.click(fn=upload_to_hf_cdn, inputs=[gen_output], outputs=[gen_share_link])
         
         # TAB 2: Transform Image