Update app.py

app.py

@@ -5,61 +5,68 @@ import requests
 import os
 from io import BytesIO
 
-from utils.planner import extract_scene_plan  # 🧠 Brain Layer (LLM)
+from utils.planner import extract_scene_plan  # 🧠 Brain Layer
 
-# Hugging Face ControlNet model
 CONTROLNET_MODEL = "lllyasviel/controlnet-sdxl-1.0-canny"
 HF_API = f"https://api-inference.huggingface.co/models/{CONTROLNET_MODEL}"
 API_KEY = os.getenv("HF_API_KEY")
-
 headers = {"Authorization": f"Bearer {API_KEY}"}
 
-# 🧠 Generate image + extract Brain Layer JSON
-def process_image(prompt, image):
+
+def process_image(prompt, image, num_variations):
     # Step 1: Brain Layer – extract structured JSON
     reasoning_json = extract_scene_plan(prompt)
 
-    # Step 2: Prepare image for HF API
+    # Step 2: Encode image once
     buffered = BytesIO()
     image.save(buffered, format="JPEG")
     img_bytes = buffered.getvalue()
+    encoded_image = base64.b64encode(img_bytes).decode("utf-8")
+
+    outputs = []
+
+    for i in range(num_variations):
+        payload = {
+            "inputs": {
+                "prompt": prompt,
+                "image": encoded_image,
+                "negative_prompt": "blurry, deformed, cropped"
+            },
+            "options": {"wait_for_model": True}
+        }
 
-    payload = {
-        "inputs": {
-            "prompt": prompt,
-            "image": base64.b64encode(img_bytes).decode("utf-8"),
-            "negative_prompt": "blurry, deformed, cropped"
-        },
-        "options": {"wait_for_model": True}
-    }
+        try:
+            response = requests.post(HF_API, headers=headers, json=payload)
+            if response.status_code == 200:
+                result_image = Image.open(BytesIO(response.content))
+                outputs.append(result_image)
+            else:
+                outputs.append(f"Error: {response.status_code} - {response.text}")
+        except Exception as e:
+            outputs.append(f"Exception: {e}")
 
-    # Step 3: Generate image
-    response = requests.post(HF_API, headers=headers, json=payload)
-    if response.status_code == 200:
-        result_image = Image.open(BytesIO(response.content))
-    else:
-        result_image = None
+    return outputs, reasoning_json
 
-    return result_image, reasoning_json
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 NewCrux AI Demo: Product → Lifestyle Image with Brain Layer")
+    gr.Markdown("# 🧠 NewCrux AI Demo: Product → Lifestyle Image Generator")
 
     with gr.Row():
         with gr.Column():
             prompt_input = gr.Textbox(label="Enter Prompt")
             image_input = gr.Image(type="pil", label="Upload Product Image")
+            variation_slider = gr.Slider(1, 4, step=1, label="Number of Variations", value=1)
             generate_btn = gr.Button("Generate")
 
         with gr.Column():
-            output_image = gr.Image(label="Generated Image")
+            output_gallery = gr.Gallery(label="Generated Images").style(grid=2, height="auto")
             json_output = gr.JSON(label="Brain Layer Reasoning (Scene Plan)")
 
     generate_btn.click(
         fn=process_image,
-        inputs=[prompt_input, image_input],
-        outputs=[output_image, json_output]
+        inputs=[prompt_input, image_input, variation_slider],
+        outputs=[output_gallery, json_output]
     )
 
 demo.launch()