Update app.py

app.py

@@ -4,7 +4,6 @@ import gradio as gr
 from PIL import Image
 import os
 import torch
-from io import BytesIO
 from diffusers import StableDiffusionInpaintPipeline
 from rembg import remove
 from utils.planner import extract_scene_plan  # 🧠 Brain Layer
@@ -28,13 +27,12 @@ if device == "cuda":
     pipe.enable_model_cpu_offload()
 
 # ----------------------------
-# 🪄 Auto-Generate Mask (rembg)
+# 🪄 Auto-Generate Mask
 # ----------------------------
 def auto_generate_mask(image: Image.Image) -> Image.Image:
-    print("🪄 Generating mask using rembg")
-    no_bg = remove(image)
-    mask = Image.open(BytesIO(no_bg)).convert("L")
-    return mask
+    print("🪄 Generating mask using rembg...")
+    no_bg = remove(image)  # Already a PIL image
+    return no_bg.convert("L")  # Convert to grayscale for inpainting
 
 # ----------------------------
 # 🎨 Image Generation Function
@@ -43,38 +41,39 @@ def process_image(prompt, image, num_variations):
     try:
         print("🧠 Prompt received:", prompt)
 
-        # 🧠 Step 1: Brain Layer
+        # Step 1: Brain Layer
         reasoning_json = extract_scene_plan(prompt)
         print("🧠 Scene plan extracted:", reasoning_json)
 
-        # Step 2: Resize & Prepare Inputs
+        # Step 2: Prepare image + auto-mask
         image = image.resize((512, 512)).convert("RGB")
         mask = auto_generate_mask(image).resize((512, 512)).convert("L")
 
-        results = []
+        outputs = []
         for i in range(num_variations):
-            print(f"🎨 Generating variation {i + 1}...")
-            output = pipe(
+            print(f"🎨 Generating variation {i+1}")
+            result = pipe(
                 prompt=prompt,
                 image=image,
                 mask_image=mask,
                 strength=0.98,
                 guidance_scale=7.5,
                 num_inference_steps=40
-            ).images[0]
-            results.append(output)
+            )
+            generated_image = result.images[0]
+            outputs.append(generated_image)
 
-        return results, reasoning_json
+        return outputs, reasoning_json
 
     except Exception as e:
-        print("❌ Error during generation:", e)
-        return ["❌ Generation failed"], {"error": str(e)}
+        print("❌ Generation failed:", e)
+        return ["❌ Error during generation"], {"error": str(e)}
 
 # ----------------------------
-# 🖼️ Gradio UI
+# 🖼 Gradio UI
 # ----------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 NewCrux AI - Auto Inpainting Demo\nUpload a product image, describe your vision, and we’ll do the magic 🪄")
+    gr.Markdown("## 🧠 NewCrux AI — Auto Inpainting\nUpload a product image, enter a prompt, and let us generate magic!")
 
     with gr.Row():
         with gr.Column():
@@ -100,5 +99,3 @@ with gr.Blocks() as demo:
 
 demo.launch()
 
-
-