app.py

import gradio as gr
from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO
import tempfile
import concurrent.futures
import os

api_key = os.environ.get("API_KEY")  
univin_model = os.environ.get("univin")  

client = genai.Client(api_key=api_key)

PROMPT_VARIATIONS = [
    "Chain-of-Thought: Step 1: Carefully analyze the input image and note every detail of the product including its text. Step 2: Generate a high-quality image with a minimalist white background and bright natural light while keeping the product and all text exactly as in the original image. Step 3: Verify that the generated image is perfect and upload-ready.",
    "Chain-of-Thought: Step 1: Examine the product in the input image with focus on preserving every detail and text. Step 2: Generate a high-quality image featuring a beautiful woman using the product in a natural setting, ensuring the product and its text remain exactly unchanged. Step 3: Confirm that the final image meets professional standards and is ready for upload.",
    "Chain-of-Thought: Step 1: Inspect the input image thoroughly, capturing all aspects of the product and its text. Step 2: Generate a high-quality image from a slightly elevated angle in a modern indoor setting with warm lighting, while keeping the product and its text intact. Step 3: Validate that the resulting image is flawless and upload-ready.",
    "Chain-of-Thought: Step 1: Analyze the original product image and note all details, ensuring the text is preserved. Step 2: Generate a high-quality image using a low angle shot with cool ambient lighting and a contemporary studio background, without altering any aspect of the product or text. Step 3: Ensure the image is perfect and ready to be uploaded.",
    "Chain-of-Thought: Step 1: Study the input image closely, ensuring every detail of the product and its text is noted. Step 2: Generate a high-quality image with a vibrant outdoor background and natural sunlight, while keeping the product and its text completely unchanged. Step 3: Double-check that the final image is flawless and upload-ready.",
    "Chain-of-Thought: Step 1: Review the product image carefully, preserving all details including the text. Step 2: Generate a high-quality image with a subtle gradient background and soft diffused lighting, ensuring that no changes are made to the product or its text. Step 3: Confirm that the generated image is perfect and ready for upload.",
    "Chain-of-Thought: Step 1: Examine the original product image and identify every detail and text element. Step 2: Generate a high-quality image featuring an artistic abstract background with high contrast lighting, strictly preserving the product and its text. Step 3: Verify that the output image is impeccable and upload-ready.",
    "Chain-of-Thought: Step 1: Analyze the provided product image, capturing all details including text. Step 2: Generate a high-quality image with a dynamic angle and a colorful background, while ensuring the product and its text are exactly as in the input. Step 3: Ensure the final image is flawless and ready to be uploaded.",
    "Chain-of-Thought: Step 1: Thoroughly inspect the input image to capture every detail of the product and its text. Step 2: Generate a high-quality image with a scenic nature background and golden hour lighting, keeping the product and its text intact. Step 3: Confirm that the generated image is perfect and upload-ready.",
    "Chain-of-Thought: Step 1: Review the original product image with a focus on preserving all details and text. Step 2: Generate a high-quality image where a beautiful woman is using the product, ensuring that the product and its text remain exactly unchanged. Step 3: Validate that the final image meets perfection and is ready for upload."
]

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
    )
    response = client.models.generate_content(
        model=univin_model,  
        contents=[text_input, input_image],
        config=types.GenerateContentConfig(response_modalities=['Text', '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
    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)
            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 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()