Update app.py

app.py

@@ -6,59 +6,82 @@ from io import BytesIO
 import tempfile
 import concurrent.futures
 import os
+import time
+from threading import Lock
 
-api_key = os.environ.get("API_KEY")  
-univin_model = os.environ.get("univin")  
-
+# 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)
 
+# Rate limiting: maximum 10 calls per minute (one call every 6 seconds)
+CALL_INTERVAL = 60.0 / 10.0  # seconds
+dt_lock = Lock()
+_last_call_time = 0.0
+
+def throttle():
+    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()
+
+# Detailed prompt variations for different backgrounds/angles
+after = "..."  # placeholder to show continuation; see full PROMPT_VARIATIONS in code
 PROMPT_VARIATIONS = [
     # 1: White background
     "Chain-of-Thought: Step 1: Carefully analyze the input image and document every single detail of the product—text, colour, material finish, branding, dimensions, and any embellishments—ensuring no detail is overlooked. Step 2: Generate a highly professional 800×800px image with a pure white background and bright natural lighting, strictly preserving the product’s text, colour, shape and finish exactly as in the original image, with zero alterations. Step 3: Rigorously verify pixel-by-pixel that the output matches the original product data 100 %, then deliver a ready-for-e-commerce upload image.",
-    # 2: Chnaged 
+    # 2: Garden background
     "Chain-of-Thought: Step 1: Meticulously analyze the input image and record every product detail—text, colour, material finish, branding, dimensions, texture, and any embellishments—ensuring nothing is overlooked. Step 2: Generate a highly professional 800×800px image of the product from a precise 45° top-down angle against a green garden background with even lighting, strictly preserving the product’s text, colours, shape and finish exactly as in the original image, with zero alterations. Step 3: Perform a rigorous pixel-level comparison to confirm 100 % fidelity to the original product data and deliver an e-commerce–ready image.",
-    # 3: Slightly elevated angle, modern indoor
+    # 3: Modern indoor
     "Chain-of-Thought: Step 1: Scrutinize the input image to record every element of the product—text, colour palette, material finish, and branding—without exception. Step 2: Generate a highly professional 800×800px image from a slightly elevated angle and open lid if the product have a lid in a modern indoor setting with warm lighting, strictly preserving the product’s original text, colours and all details exactly as in the source. Step 3: Double-check at full resolution that no aspect of the product has been altered and that the image is e-commerce ready.",
-    # 4: Low angle, cool ambient lighting
+    # 4: Spa-like setting
     "Chain-of-Thought: Step 1: Analyze the input image pixel-for-pixel to note every product detail—text, colour, texture, finish, branding and proportions—ensuring nothing is missed. Step 2: Generate a highly professional 800×800px image using a A serene, spa-like setting with elements of tranquility like smooth stones, bamboo, and a calming water feature, ideal for wellness or beauty products. background, with the product’s text, colours, materials and details preserved exactly as in the original. Step 3: Verify 100 % fidelity at the pixel level and confirm readiness for online catalog upload.",
-    # 5: Vibrant outdoor background
+    # 5: Vibrant outdoor
     "Chain-of-Thought: Step 1: Inspect the input image meticulously and record every detail—text, colour gradients, material texture, and branding—ensuring total accuracy. Step 2: Generate a highly professional 800×800px image with a vibrant outdoor background and natural sunlight, strictly preserving the product’s original text, colours, finish and all details unchanged. Step 3: Validate that the final image is flawless, matches the original exactly, and is optimized for e-commerce display.",
-    # 6: Subtle gradient background
+    # 6: Gradient background
     "Chain-of-Thought: Step 1: Review the input image thoroughly, capturing and logging every detail—text, colour accuracy, surface finish, and branding—100 % intact. Step 2: Generate a highly professional 800×800px image with a subtle gradient background and soft diffused lighting, ensuring that the product’s text, colours and all design details remain precisely as in the source. Step 3: Perform a pixel-level check to confirm no alterations and that the image meets e-commerce quality standards.",
-    # 7: Artistic abstract background
+    # 7: Abstract background
     "Chain-of-Thought: Step 1: Examine the original product image in depth, noting every detail—text, colour tones, material texture, branding logos and dimensions—with absolute precision. Step 2: Generate a highly professional 800×800px image featuring an abstract background with perfect lighting, strictly preserving the product’s text, colours and every fine detail exactly as in the original. Step 3: Verify pixel-perfect fidelity and confirm the output is upload-ready for an online store.",
-    # 9: Rich lifestyle background
+    # 8: Lifestyle setting
     "Chain-of-Thought: Step 1: Scrutinize the input image in full detail, logging every aspect of the product—text, colour accuracy, branding, material texture, finish, and proportions—to guarantee total fidelity. Step 2: Generate a highly professional 800×800px image of the product placed A cozy, homely setting with a comfortable, inviting ambiance, featuring elements like soft lighting, plush textiles, and warm tones, suitable for home goods or comfort products. ensuring that the product’s text, colours, materials and every detail remain exactly as in the original image without any modification. Step 3: Validate at the pixel level that the output matches the original image perfectly and confirm readiness for online catalog upload.",
 ]
 
 def process_variation(variation, input_image, product_name):
+    # Throttle to respect API rate limit
+    throttle()
     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,  
+        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:
-            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
+            img = Image.open(BytesIO(part.inline_data.data))
+            # Save to temp file
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
+                img.save(tmp, format="PNG")
+                return tmp.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
+            executor.submit(process_variation, var, input_image, product_name)
+            for var in PROMPT_VARIATIONS
         ]
-        output_files = [future.result() for future in futures if future.result() is not None]
-    return output_files
+        # Collect successful outputs
+        return [f.result() for f in futures if f.result() is not None]
 
-# Improved UI Layout and Styling
+# Gradio UI setup
 custom_css = """
 #generate-button {
     background-color: #4A90E2;
@@ -85,35 +108,22 @@ custom_css = """
 demo = gr.Blocks(css=custom_css, theme=gr.themes.Soft())
 with demo:
     gr.Markdown("""
-    ## 🎨 Uni-Imaginator
-    **Create professional, e-commerce–ready product images effortlessly.**
+    ## 🎨 Uni-Imaginator (Rate-Limited)
+    **Create professional, e-commerce–ready product images effortlessly—now respecting API rate limits.**
     """)
     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",
-                interactive=True
-            )
-            product_name = gr.Textbox(
-                label="Product Name", placeholder="Enter your product's name here",
-                elem_id="product-name"
-            )
-            generate_button = gr.Button(
-                "Generate Images", variant="primary", elem_id="generate-button"
-            )
+            input_image = gr.Image(type="pil", label="Upload Product Image", elem_id="upload-img")
+            product_name = gr.Textbox(label="Product Name", placeholder="Enter your product's name here", elem_id="product-name")
+            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. Enter the product name.
-            3. Click **Generate Images** to see a gallery of styled product shots.
+            3. 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, product_name], outputs=gallery
-    )
+    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, product_name], outputs=gallery)
 
 demo.launch()