Update app.py

using base model for all products

app.py

@@ -2,97 +2,91 @@
 
 import gradio as gr
 from PIL import Image
-import base64
-import requests
 import os
-from io import BytesIO
-
+import torch
+from diffusers import StableDiffusionXLInpaintPipeline
 from utils.planner import extract_scene_plan  # 🧠 Brain Layer
 
-# 🔐 Hugging Face keys
-HF_API_KEY = os.getenv("HF_API_KEY")
-SDXL_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"  # ✅ Correct model for image-to-image
-SDXL_API_URL = f"https://api-inference.huggingface.co/models/{SDXL_MODEL_ID}"
-HEADERS = {"Authorization": f"Bearer {HF_API_KEY}"}
-
-# 🚀 Image generation (img2img)
-def process_image(prompt, image, num_variations):
+# ----------------------------
+# 🔧 Device Setup
+# ----------------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
+
+# ----------------------------
+# 📦 Load Inpainting Model
+# ----------------------------
+pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "diffusers/stable-diffusion-xl-1.0-inpainting",
+    torch_dtype=dtype,
+    variant="fp16" if device == "cuda" else None
+).to(device)
+
+pipe.enable_attention_slicing()
+pipe.enable_model_cpu_offload()
+
+# ----------------------------
+# 🎨 Image Generation Function
+# ----------------------------
+def process_image(prompt, image, mask, 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: Encode input image
-        buffered = BytesIO()
-        image.save(buffered, format="JPEG")
-        img_bytes = buffered.getvalue()
-        encoded_image = base64.b64encode(img_bytes).decode("utf-8")
+        # Resize inputs to 1024x1024 (required for SDXL)
+        image = image.resize((1024, 1024)).convert("RGB")
+        mask = mask.resize((1024, 1024)).convert("L")
 
-        # Step 3: Send image + prompt to HF API
-        outputs = []
+        results = []
         for i in range(num_variations):
-            payload = {
-                "image": encoded_image,
-                "prompt": prompt,
-                "negative_prompt": "blurry, deformed, cropped",
-                "strength": 25,
-                "guidance_scale": 7.5
-            }
-
-            print(f"📤 Sending request to HF (variation {i+1})")
-            response = requests.post(SDXL_API_URL, headers=HEADERS, json=payload)
-
-            if response.status_code == 200:
-                try:
-                    result_json = response.json()
-                    if "images" in result_json:
-                        base64_img = result_json["images"][0]
-                        result_image = Image.open(BytesIO(base64.b64decode(base64_img)))
-                        outputs.append(result_image)
-                        print(f"✅ Decoded image variation {i+1} successfully")
-                    else:
-                        print(f"⚠️ No 'images' key found in response")
-                        outputs.append("❌ No image in response.")
-                except Exception as decode_err:
-                    print("❌ Image decode error:", decode_err)
-                    outputs.append("❌ Failed to decode image.")
-            else:
-                print(f"❌ HF API error: {response.status_code} - {response.text}")
-                outputs.append(f"Error {response.status_code}: {response.text}")
-
-        return outputs, reasoning_json
+            print(f"🎨 Generating variation {i + 1}...")
+            output = pipe(
+                prompt=prompt,
+                image=image,
+                mask_image=mask,
+                strength=0.98,
+                guidance_scale=7.5,
+                num_inference_steps=40
+            ).images[0]
+            results.append(output)
+
+        return results, reasoning_json
 
     except Exception as e:
-        print("❌ General Exception in process_image:", e)
-        return ["Processing error occurred"], {"error": str(e)}
+        print("❌ Error during generation:", e)
+        return ["❌ Generation failed"], {"error": str(e)}
 
-# 🎨 Gradio UI
+# ----------------------------
+# 🖼️ Gradio UI
+# ----------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 NewCrux AI Demo: Image-to-Image using Fast SDXL + Brain Layer")
+    gr.Markdown("## 🧠 NewCrux Inpainting Demo (SDXL)\nUpload a product image, a mask, and a prompt to generate realistic content.")
 
     with gr.Row():
         with gr.Column():
-            prompt_input = gr.Textbox(label="Enter Prompt")
+            prompt_input = gr.Textbox(label="Prompt")
             image_input = gr.Image(type="pil", label="Upload Product Image")
+            mask_input = gr.Image(type="pil", label="Upload Mask (white = keep, black = replace)")
             variation_slider = gr.Slider(1, 4, step=1, value=1, label="Number of Variations")
             generate_btn = gr.Button("Generate")
 
         with gr.Column():
             output_gallery = gr.Gallery(
-                label="Generated Image Variations",
+                label="Generated Variations",
                 columns=2,
                 rows=2,
                 height="auto"
             )
-            json_output = gr.JSON(label="Brain Layer Reasoning (Scene Plan)")
+            json_output = gr.JSON(label="🧠 Brain Layer Reasoning")
 
     generate_btn.click(
         fn=process_image,
-        inputs=[prompt_input, image_input, variation_slider],
+        inputs=[prompt_input, image_input, mask_input, variation_slider],
         outputs=[output_gallery, json_output]
     )
 
-demo.launch(share=True)
+demo.launch()