Update app.py

app.py

@@ -3,7 +3,6 @@
 import gradio as gr
 from PIL import Image
 import torch
-
 from diffusers import StableDiffusionXLImg2ImgPipeline
 from utils.planner import (
     extract_scene_plan,
@@ -12,13 +11,13 @@ from utils.planner import (
 )
 
 # ----------------------------
-# 🔧 Device Setup
+# 🖥️ Device Setup
 # ----------------------------
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if device == "cuda" else torch.float32
 
 # ----------------------------
-# ✅ Load SDXL Only Pipeline
+# 🧠 Load SDXL Img2Img Pipeline (no ControlNet)
 # ----------------------------
 pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
@@ -31,38 +30,38 @@ pipe.enable_model_cpu_offload()
 pipe.enable_attention_slicing()
 
 # ----------------------------
-# 🎨 Image Generation Function
+# 🎯 Image Generation Function
 # ----------------------------
 def process_image(prompt, image, num_variations):
     try:
-        print("🧠 User Prompt:", prompt)
+        print("🧠 Prompt received:", prompt)
         if image is None:
-            raise ValueError("🚫 Uploaded image is missing.")
+            raise ValueError("🚫 Please upload an image.")
 
-        # Step 1: Extract scene plan
+        # Step 1: Analyze image + prompt → structured plan
         scene_plan = extract_scene_plan(prompt, image)
         print("📋 Scene Plan:", scene_plan)
 
-        # Step 2: Generate enriched prompts
-        prompt_list = generate_prompt_variations_from_scene(scene_plan, prompt, num_variations)
-        print("✅ Enriched Prompts:", prompt_list)
+        # Step 2: Generate prompt variations
+        enriched_prompts = generate_prompt_variations_from_scene(scene_plan, prompt, num_variations)
+        print("✅ Enriched Prompts:", enriched_prompts)
 
         # Step 3: Generate negative prompt
         negative_prompt = generate_negative_prompt_from_scene(scene_plan)
         print("🚫 Negative Prompt:", negative_prompt)
 
-        # Step 4: Resize image to SDXL resolution
+        # Step 4: Resize image to 1024x1024
         image = image.resize((1024, 1024)).convert("RGB")
 
-        # Step 5: Generate outputs with SDXL only
+        # Step 5: Generate variations
         outputs = []
-        for i, enriched_prompt in enumerate(prompt_list):
-            print(f"🎨 Generating variation {i+1}...")
+        for i, enriched_prompt in enumerate(enriched_prompts):
+            print(f"✨ Generating Image {i + 1}...")
             result = pipe(
                 prompt=enriched_prompt,
                 negative_prompt=negative_prompt,
                 image=image,
-                strength=0.7,                # ← You can fine-tune this
+                strength=0.7,
                 guidance_scale=7.5,
                 num_inference_steps=30,
             )
@@ -71,7 +70,7 @@ def process_image(prompt, image, num_variations):
         return outputs
 
     except Exception as e:
-        print("❌ Generation Error:", e)
+        print("❌ Generation failed:", e)
         return [Image.new("RGB", (512, 512), color="red")]
 
 # ----------------------------
@@ -84,10 +83,14 @@ demo = gr.Interface(
         gr.Image(type="pil", label="Product Image"),
         gr.Slider(1, 5, value=3, step=1, label="Number of Variations")
     ],
-    output_gallery = gr.Gallery(label="Generated Variations", columns=2, rows=2, height="auto"),
+    outputs=gr.Gallery(label="Generated Images"),
     title="NewCrux Product Image Generator (SDXL Only)",
-    description="Upload a product image and enter a prompt. SDXL will generate enriched variations using AI."
+    description="Upload a product image and describe your vision. NewCrux will generate multiple enriched image variations using SDXL + AI prompt planning."
 )
 
+# ----------------------------
+# 🚀 Launch
+# ----------------------------
 if __name__ == "__main__":
     demo.launch()
+