Update app.py

app.py

@@ -9,45 +9,40 @@ import os
 
 api_key = os.environ.get("API_KEY")
 univin_model = os.environ.get("univin")
-
 client = genai.Client(api_key=api_key)
 
 PROMPT_VARIATIONS = [
-    # 1. Professional White Background (Universal)
-    "A high-resolution professional image of the product perfectly centered on a pure white seamless background, preserving exact shape, color, texture, printed text, and logo precisely as in the input image. Ideal for professional e-commerce.",
-
-    # 2. Luxurious Spa Setting (Beauty, Skincare, Haircare)
+    "A high-resolution professional image of the product perfectly centered on a pure white seamless background, preserving exact shape, color, texture, printed text, and logo precisely as in the input image.",
     "A luxurious spa-themed image showcasing the original beauty or skincare product elegantly positioned on polished marble surrounded by spa accessories (stones, towels, plants). Preserve all original product details exactly as input.",
-
-    # 3. Professional Model Holding Product (Beauty/Skincare)
     "Professional image of a sophisticated woman model holding the original beauty or skincare product gently and elegantly, focusing sharply on product details while maintaining precise original shape, color, texture, text, and logos.",
-
-    # 4. Modern Tech Scene (Personal Electronics)
     "A sleek, modern technology-oriented image featuring the electronic product placed on a reflective black or dark-gray surface with subtle neon lighting, strictly maintaining original shape, color, texture, printed text, and logos.",
-
-    # 5. Lifestyle with Model (Candles/Soft Toys)
     "Warm lifestyle image featuring either a woman professionally holding a lit candle or a child warmly hugging a soft toy. Ensure exact preservation of the original product's shape, color, texture, printed text, and logos as provided.",
-
-    # 6. Natural Refreshing Composition (Haircare & Skincare)
     "Refreshing natural composition featuring the original skincare or haircare product among botanical elements and water droplets. Precisely preserve original product shape, color, texture, printed text, and logos from the input image."
 ]
 
 def process_variation(variation, input_image, product_name):
-    text_input = (
-        f"Hi, this is a picture of a product. The name of the product is {product_name}.",
-        variation
-    )
+    buffered = BytesIO()
+    input_image.save(buffered, format="PNG")
+    image_bytes = buffered.getvalue()
+
+    prompt_text = f"This is a product named '{product_name}'. {variation}"
+
     response = client.models.generate_content(
         model=univin_model,
-        contents=[text_input, input_image],
-        config=types.GenerateContentConfig(response_modalities=['Text', 'Image'])
+        contents=[
+            types.Part(text=prompt_text),
+            types.Part(inline_data=types.Blob(mime_type="image/png", data=image_bytes))
+        ],
+        config=types.GenerateContentConfig(response_modalities=['Image'])
     )
-    for part in response.candidates[0].content.parts:
-        if part.inline_data is not None:
-            generated_img = Image.open(BytesIO(part.inline_data.data))
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file:
-                generated_img.save(tmp_file, format="PNG")
-                return tmp_file.name
+
+    if response.candidates:
+        for part in response.candidates[0].content.parts:
+            if part.inline_data:
+                generated_img = Image.open(BytesIO(part.inline_data.data))
+                with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file:
+                    generated_img.save(tmp_file, format="PNG")
+                    return tmp_file.name
     return None
 
 def generate_images(input_image, product_name):
@@ -70,4 +65,4 @@ with gr.Blocks() as demo:
 
     generate_button.click(fn=generate_images, inputs=[input_image, product_name], outputs=gallery)
 
-demo.launch(debug=True)
+demo.launch(debug=True)