Update app.py

app.py

@@ -1,19 +1,16 @@
 import gradio as gr
 from gradio_client import Client, handle_file
-from PIL import Image, ImageOps, ImageDraw, ImageFont
+from PIL import Image, ImageDraw, ImageFont
 import io, os, base64, requests, random
 
-# HF Client for Z-Image model
 client = Client("Tongyi-MAI/Z-Image-Turbo")
 
-# ---------------------------------------------
-# Core function to call HF Space
-# ---------------------------------------------
-def call_generate(prompt, resolution="1024x1536 ( 2:3 )", seed=None, steps=8, gallery_urls=None, random_seed=False):
+def call_generate(prompt, resolution="1024x1536 ( 2:3 )", seed=None, steps=8, gallery_files=None, random_seed=False):
+
     gallery_images = []
-    if gallery_urls:
-        for url in gallery_urls:
-            gallery_images.append({"image": handle_file(url), "caption": None})
+    if gallery_files:
+        for path in gallery_files:
+            gallery_images.append({"image": handle_file(path), "caption": None})
 
     payload = dict(
         prompt=prompt,
@@ -25,104 +22,93 @@ def call_generate(prompt, resolution="1024x1536 ( 2:3 )", seed=None, steps=8, ga
         gallery_images=gallery_images,
         api_name="/generate",
     )
-
     return client.predict(**payload)
 
-# ---------------------------------------------
-# Generate posters & overlay tagline + logo
-# ---------------------------------------------
 def generate_posters(product_name, tagline, product_img, logo_img, variants):
 
-    # Save uploaded files locally so we can pass URLs to the HF API
-    product_path = "product_temp.png"
-    logo_path = "logo_temp.png"
+    if product_img is None or logo_img is None:
+        return ["❌ Please upload both product image and logo."]
+
+    # Save uploaded files
+    product_path = "product.png"
+    logo_path = "logo.png"
 
     Image.open(product_img).save(product_path)
     Image.open(logo_img).save(logo_path)
 
-    product_url = "file://" + os.path.abspath(product_path)
-    logo_url = "file://" + os.path.abspath(logo_path)
-
     base_prompt = (
-        f"{product_name} on a sandy beach at sunset — composition: product foreground centered, "
-        f"low-angle shot, golden hour rim light, cinematic mood, glossy finish, photorealistic high detail."
+        f"{product_name} on a sandy beach at sunset, cinematic composition, golden rim light, "
+        f"glossy highlights, photorealistic 4K detail. Clean whitespace for text."
     )
 
     output_images = []
 
     for i in range(variants):
         seed = random.randint(1, 10**6)
-        prompt = base_prompt + f" Tagline: \"{tagline}\" Variation {i}."
+        full_prompt = base_prompt + f" Tagline: \"{tagline}\"."
 
         raw = call_generate(
-            prompt,
+            full_prompt,
             resolution="1024x1536 ( 2:3 )",
             seed=seed,
             steps=12,
-            gallery_urls=[product_url, logo_url],
+            gallery_files=[product_path, logo_path]
         )
 
-        # Try to decode base64 or URL response
+        # Parse response
         img_bytes = None
         if isinstance(raw, dict):
             for v in raw.values():
                 if isinstance(v, str) and v.startswith("data:"):
-                    img_bytes = base64.b64decode(v.split(",")[1])
+                    img_bytes = base64.b64decode(v.split(",", 1)[1])
                     break
                 if isinstance(v, list):
                     for item in v:
                         if isinstance(item, str) and item.startswith("data:"):
-                            img_bytes = base64.b64decode(item.split(",")[1])
+                            img_bytes = base64.b64decode(item.split(",", 1)[1])
                             break
 
-        if img_bytes is None and isinstance(raw, str) and raw.startswith("data:"):
-            img_bytes = base64.b64decode(raw.split(",", 1)[1])
-
-        if img_bytes is None and isinstance(raw, str) and raw.startswith("http"):
-            img_bytes = requests.get(raw).content
-
         if img_bytes is None:
-            return [f"Error decoding image: {raw}"]
+            return ["❌ Could not decode model output"]
 
         img = Image.open(io.BytesIO(img_bytes)).convert("RGBA")
 
-        # ---- Overlay Logo ---
+        # Overlay Logo
         try:
-            logo = Image.open(logo_img).convert("RGBA")
-            logo.thumbnail((int(img.width * 0.18), int(img.height * 0.08)), Image.LANCZOS)
-            logo_pos = (img.width - logo.width - 40, img.height - logo.height - 40)
-            img.alpha_composite(logo, logo_pos)
+            logo = Image.open(logo_path).convert("RGBA")
+            logo.thumbnail((int(img.width * 0.18), int(img.height * 0.08)))
+            img.alpha_composite(logo, (img.width - logo.width - 30, img.height - logo.height - 30))
         except:
             pass
 
-        # ---- Overlay Tagline ---
+        # Overlay Tagline
         draw = ImageDraw.Draw(img)
         try:
-            font = ImageFont.truetype("Arial.ttf", size=int(img.width*0.06))
+            font = ImageFont.truetype("Arial.ttf", size=int(img.width * 0.06))
         except:
             font = ImageFont.load_default()
 
-        text_w, text_h = draw.textsize(tagline, font=font)
-        x, y = 60, img.height - text_h - 60
-        draw.rectangle([x-15, y-15, x+text_w+15, y+text_h+15], fill=(0,0,0,160))
-        draw.text((x,y), tagline, font=font, fill=(255,255,255,255))
+        tw, th = draw.textsize(tagline, font=font)
+        x, y = 60, img.height - th - 60
+        draw.rectangle([x-12, y-12, x+tw+12, y+th+12], fill=(0,0,0,160))
+        draw.text((x, y), tagline, fill="white", font=font)
 
-        final_path = f"poster_{seed}_{i}.png"
-        img.convert("RGB").save(final_path)
-        output_images.append(final_path)
+        out_path = f"poster_{seed}_{i}.png"
+        img.convert("RGB").save(out_path)
+        output_images.append(out_path)
 
     return output_images
 
-# ---------------------------------------------
+
+# ---------------------------
 # GRADIO UI
-# ---------------------------------------------
+# ---------------------------
 with gr.Blocks(title="Rookus Poster Generator") as demo:
 
-    gr.Markdown("## 🖼️ Rookus AI Poster Generator\nCreate cinematic advertising posters automatically.")
+    gr.Markdown("## 🎨 Rookus AI Poster Generator")
 
-    with gr.Row():
-        product_name = gr.Textbox(label="Product Name", placeholder="Sneakers")
-        tagline = gr.Textbox(label="Tagline", placeholder="Purchase these beautiful sneakers now")
+    product_name = gr.Textbox(label="Product Name")
+    tagline = gr.Textbox(label="Tagline")
 
     with gr.Row():
         product_img = gr.Image(type="filepath", label="Product Image")
@@ -130,14 +116,13 @@ with gr.Blocks(title="Rookus Poster Generator") as demo:
 
     variants = gr.Slider(1, 6, value=3, step=1, label="Number of Variants")
 
-    generate_btn = gr.Button("Generate Posters 🎨")
-
-    gallery = gr.Gallery(label="Generated Posters", height="auto")
+    generate_btn = gr.Button("Generate Posters")
+    gallery = gr.Gallery(label="Output Posters")
 
     generate_btn.click(
-        fn=generate_posters,
+        generate_posters,
         inputs=[product_name, tagline, product_img, logo_img, variants],
-        outputs=[gallery]
+        outputs=gallery
     )
 
 demo.launch()