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
import time
from threading import Lock

# Load API credentials from environment
api_key = os.environ.get("API_KEY")
univin_model = os.environ.get("univin")
client = genai.Client(api_key=api_key)

CALL_INTERVAL = 2.0 / 10.0  # seconds
_dt_lock = Lock()
_last_call_time = 0.0

def throttle():
    """
    Ensures there's at least CALL_INTERVAL seconds between API calls.
    """
    global _last_call_time
    with _dt_lock:
        now = time.time()
        elapsed = now - _last_call_time
        if elapsed < CALL_INTERVAL:
            time.sleep(CALL_INTERVAL - elapsed)
        _last_call_time = time.time()


PROMPT_VARIATIONS = [
    # 1: Front view
    "Create a high-resolution 800×800px image of the EXACT SAME product from a front view. CRITICAL INSTRUCTIONS: The product must be a PIXEL-PERFECT match to the original in EVERY aspect - EXACT same text (including font, size, spacing, and positioning), EXACT same colors (including all shades, gradients, and color codes), EXACT same materials and textures, EXACT same branding elements, EXACT same dimensions and proportions, EXACT same markings and labels, EXACT same finish and surface details, EXACT same shadows and highlights, EXACT same reflections and gloss, EXACT same embossing or debossing if present. DO NOT modify ANY aspect of the product. ONLY change the viewing angle to front view. Maintain a pure white background with professional studio lighting. The output must be a high-quality, e-commerce ready image with perfect color accuracy and sharp details.",
    
    # 2: Left-side view
    "Create a high-resolution 800×800px image of the EXACT SAME product from a left-side three-quarter (45-degree) angle. CRITICAL INSTRUCTIONS: The product must be a PIXEL-PERFECT match to the original in EVERY aspect - EXACT same text (including font, size, spacing, and positioning), EXACT same colors (including all shades, gradients, and color codes), EXACT same materials and textures, EXACT same branding elements, EXACT same dimensions and proportions, EXACT same markings and labels, EXACT same finish and surface details, EXACT same shadows and highlights, EXACT same reflections and gloss, EXACT same embossing or debossing if present. DO NOT modify ANY aspect of the product. ONLY change the viewing angle to left-side view. Maintain a pure white background with professional studio lighting. The output must be a high-quality, e-commerce ready image with perfect color accuracy and sharp details.",
    
    # 3: Right-side view
    "Create a high-resolution 800×800px image of the EXACT SAME product from a right-side three-quarter (45-degree) angle. CRITICAL INSTRUCTIONS: The product must be a PIXEL-PERFECT match to the original in EVERY aspect - EXACT same text (including font, size, spacing, and positioning), EXACT same colors (including all shades, gradients, and color codes), EXACT same materials and textures, EXACT same branding elements, EXACT same dimensions and proportions, EXACT same markings and labels, EXACT same finish and surface details, EXACT same shadows and highlights, EXACT same reflections and gloss, EXACT same embossing or debossing if present. DO NOT modify ANY aspect of the product. ONLY change the viewing angle to right-side view. Maintain a pure white background with professional studio lighting. The output must be a high-quality, e-commerce ready image with perfect color accuracy and sharp details.",
    
    # 4: Top-down view
    "Create a high-resolution 800×800px image of the EXACT SAME product from a top-down (bird's-eye) angle. CRITICAL INSTRUCTIONS: The product must be a PIXEL-PERFECT match to the original in EVERY aspect - EXACT same text (including font, size, spacing, and positioning), EXACT same colors (including all shades, gradients, and color codes), EXACT same materials and textures, EXACT same branding elements, EXACT same dimensions and proportions, EXACT same markings and labels, EXACT same finish and surface details, EXACT same shadows and highlights, EXACT same reflections and gloss, EXACT same embossing or debossing if present. DO NOT modify ANY aspect of the product. ONLY change the viewing angle to top-down view. Maintain a pure white background with professional studio lighting. The output must be a high-quality, e-commerce ready image with perfect color accuracy and sharp details."
]
# Detailed prompt variations for different backgrounds/angles

def process_variation(variation, input_image):
    # Throttle to respect API rate limit
    throttle()

    # Build text + image input
    text_input = (
        "Hi, this is a picture of a product.",
        variation
    )

    # Call the GenAI API
    response = client.models.generate_content(
        model=univin_model,
        contents=[text_input, input_image],
        config=types.GenerateContentConfig(response_modalities=['Text', 'Image'])
    )

    # Extract generated image bytes
    for part in response.candidates[0].content.parts:
        if part.inline_data is not None:
            img = Image.open(BytesIO(part.inline_data.data))
            # Save to temp file and return its path
            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
                img.save(tmp, format="PNG")
                return tmp.name
    return None


def generate_images(input_image):
    """
    Generate one image per prompt variation in parallel, respecting rate limits.
    """
    with concurrent.futures.ThreadPoolExecutor() as executor:
        futures = [
            executor.submit(process_variation, var, input_image)
            for var in PROMPT_VARIATIONS
        ]
        return [f.result() for f in futures if f.result()]

# Gradio UI setup with custom styling
def build_interface():
    custom_css = """
    #generate-button {
        background-color: #4A90E2;
        color: white;
        border-radius: 8px;
        padding: 12px 24px;
        font-size: 16px;
        margin-top: 10px;
    }
    #gallery .gallery-item {
        border-radius: 12px;
        box-shadow: 0 4px 6px rgba(0,0,0,0.1);
        transition: transform 0.2s ease;
    }
    #gallery .gallery-item:hover {
        transform: scale(1.05);
    }
    #input-row {
        gap: 20px;
        align-items: start;
    }
    """

    demo = gr.Blocks(css=custom_css, theme=gr.themes.Soft())
    with demo:
        gr.Markdown(
            """
            ## 🎨 Uni-Imaginator
            **Create professional, e-commerce–ready product images effortlessly without worrying abot copyrights.**
            """
        )
        with gr.Row(elem_id="input-row"):
            with gr.Column(scale=1):
                input_image = gr.Image(
                    type="pil", label="Upload Product Image", elem_id="upload-img"
                )
                generate_button = gr.Button(
                    "Generate Images", variant="primary", elem_id="generate-button"
                )
            with gr.Column(scale=1):
                gr.Markdown(
                    """
                    **How to Use:**
                    1. Upload a clear photo of your product.
                    2. Click **Generate Images** and wait a few moments while images are processed within rate limits.
                    """
                )
        gallery = gr.Gallery(
            label="Generated Images",
            elem_id="gallery",
            columns=4,
            object_fit="contain",
            height="auto",
            show_label=False
        )
        generate_button.click(
            fn=generate_images,
            inputs=[input_image],
            outputs=gallery
        )
    return demo

# Launch the app
if __name__ == "__main__":
    demo_app = build_interface()
    demo_app.launch()