Update app.py

app.py

@@ -1,19 +1,25 @@
 import gradio as gr
-import requests
 from PIL import Image
-from io import BytesIO
 import base64
+import requests
 import os
+from io import BytesIO
+
+from utils.planner import extract_scene_plan  # 🧠 Brain Layer (LLM)
+
+# 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")
 
-# Hugging Face ControlNet API (Canny version)
-HF_API = "https://api-inference.huggingface.co/models/lllyasviel/controlnet-sdxl-1.0-canny"
-API_KEY = os.getenv("HF_API_KEY")  # Secure: fetch from secret
+headers = {"Authorization": f"Bearer {API_KEY}"}
 
-headers = {
-    "Authorization": f"Bearer {API_KEY}"
-}
+# 🧠 Generate image + extract Brain Layer JSON
+def process_image(prompt, image):
+    # Step 1: Brain Layer – extract structured JSON
+    reasoning_json = extract_scene_plan(prompt)
 
-def generate_image(prompt, image):
+    # Step 2: Prepare image for HF API
     buffered = BytesIO()
     image.save(buffered, format="JPEG")
     img_bytes = buffered.getvalue()
@@ -27,23 +33,34 @@ def generate_image(prompt, image):
         "options": {"wait_for_model": True}
     }
 
+    # Step 3: Generate image
     response = requests.post(HF_API, headers=headers, json=payload)
-
     if response.status_code == 200:
-        img_out = Image.open(BytesIO(response.content))
-        return img_out
+        result_image = Image.open(BytesIO(response.content))
     else:
-        return f"Error: {response.status_code} - {response.text}"
+        result_image = None
+
+    return result_image, reasoning_json
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 NewCrux AI Demo: Product → Lifestyle Image")
+    gr.Markdown("# 🧠 NewCrux AI Demo: Product → Lifestyle Image with Brain Layer")
+
     with gr.Row():
-        input_image = gr.Image(type="pil", label="Upload Product Image")
-        prompt_text = gr.Textbox(label="Enter Prompt", placeholder="e.g., A runner on a beach wearing this shoe")
-    output_image = gr.Image(label="Generated Lifestyle Image")
-    generate_btn = gr.Button("Generate Image")
+        with gr.Column():
+            prompt_input = gr.Textbox(label="Enter Prompt")
+            image_input = gr.Image(type="pil", label="Upload Product Image")
+            generate_btn = gr.Button("Generate")
 
-    generate_btn.click(fn=generate_image, inputs=[prompt_text, input_image], outputs=output_image)
+        with gr.Column():
+            output_image = gr.Image(label="Generated Image")
+            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]
+    )
 
 demo.launch()
+