Update app.py

app.py

@@ -7,8 +7,9 @@ import tempfile
 import concurrent.futures
 import os
 
-api_key = os.environ.get("API_KEY")
-univin_model = os.environ.get("univin")
+api_key = os.environ.get("API_KEY")  
+univin_model = os.environ.get("univin")  
+
 client = genai.Client(api_key=api_key)
 
 PROMPT_VARIATIONS = [
@@ -21,31 +22,25 @@ PROMPT_VARIATIONS = [
 ]
 
 def process_variation(variation, input_image, product_name):
-    buffered = BytesIO()
-    input_image.save(buffered, format="PNG")
-    image_bytes = buffered.getvalue()
-
-    prompt_text = f"This is a product named '{product_name}'. {variation}"
-
+    text_input = (
+        f"Hi, this is a picture of a product. The name of the product is {product_name}.",
+        variation
+    )
     response = client.models.generate_content(
-        model=univin_model,
-        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'])
+        model=univin_model,  
+        contents=[text_input, input_image],
+        config=types.GenerateContentConfig(response_modalities=['Text', 'Image'])
     )
-
-    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
+    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
     return None
 
 def generate_images(input_image, product_name):
+    
     with concurrent.futures.ThreadPoolExecutor() as executor:
         futures = [
             executor.submit(process_variation, variation, input_image, product_name)
@@ -60,9 +55,9 @@ with gr.Blocks() as demo:
         input_image = gr.Image(type="pil", label="Upload Image")
         product_name = gr.Textbox(label="Product Name", placeholder="Enter the product name")
     generate_button = gr.Button("Generate Images")
-
+    
     gallery = gr.Gallery(label="Generated Images", elem_id="gallery", columns=4, height="auto", object_fit="contain", interactive=True)
-
+    
     generate_button.click(fn=generate_images, inputs=[input_image, product_name], outputs=gallery)
-
-demo.launch(debug=True)
+    
+demo.launch()