Update app.py

app.py

@@ -33,54 +33,41 @@ PROMPT_VARIATIONS = [
 ]
 
 def process_variation(variation, input_image, product_name):
-    prompt_text = f"Product Name: {product_name}\n{variation}"
-
-    request_content = [
-        types.Part(text=prompt_text),
-        types.Part(inline_data=types.Blob(mime_type="image/png", data=input_image.getvalue()))
-    ]
-
+    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=request_content,
-        config=types.GenerateContentConfig(response_modalities=['Image'])
+        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_pil_image, product_name):
-    input_buffer = BytesIO()
-    input_pil_image.save(input_buffer, format="PNG")
-
+def generate_images(input_image, product_name):
     with concurrent.futures.ThreadPoolExecutor() as executor:
         futures = [
-            executor.submit(process_variation, variation, input_buffer, product_name)
+            executor.submit(process_variation, variation, input_image, product_name)
             for variation in PROMPT_VARIATIONS
         ]
         output_files = [future.result() for future in futures if future.result() is not None]
-
     return output_files
 
 with gr.Blocks() as demo:
     gr.Markdown("# Uni-Imaginator")
     with gr.Row():
-        input_image = gr.Image(type="pil", label="Upload Product Image")
+        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", columns=3, height="auto", interactive=True)
 
-    generate_button.click(
-        fn=generate_images,
-        inputs=[input_image, product_name],
-        outputs=gallery
-    )
+    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)