Update app.py

app.py

@@ -14,17 +14,16 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if device == "cuda" else torch.float32
 
 # ----------------------------
-# 📦 Load Inpainting Model
+# 📦 Load Inpainting Model (SD 1.5-based for now)
 # ----------------------------
-
-
 pipe = StableDiffusionInpaintPipeline.from_pretrained(
     "runwayml/stable-diffusion-inpainting",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-).to("cuda" if torch.cuda.is_available() else "cpu")
+    torch_dtype=dtype
+).to(device)
 
-pipe.enable_attention_slicing()
-pipe.enable_model_cpu_offload()
+if device == "cuda":
+    pipe.enable_attention_slicing()
+    pipe.enable_model_cpu_offload()  # ✅ Only on GPU!
 
 # ----------------------------
 # 🎨 Image Generation Function
@@ -37,9 +36,9 @@ def process_image(prompt, image, mask, num_variations):
         reasoning_json = extract_scene_plan(prompt)
         print("🧠 Scene plan extracted:", reasoning_json)
 
-        # Resize inputs to 1024x1024 (required for SDXL)
-        image = image.resize((1024, 1024)).convert("RGB")
-        mask = mask.resize((1024, 1024)).convert("L")
+        # Step 2: Resize inputs (SD 1.5 = 512x512, SDXL = 1024x1024)
+        image = image.resize((512, 512)).convert("RGB")
+        mask = mask.resize((512, 512)).convert("L")
 
         results = []
         for i in range(num_variations):
@@ -64,7 +63,7 @@ def process_image(prompt, image, mask, num_variations):
 # 🖼️ Gradio UI
 # ----------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 NewCrux Inpainting Demo (SDXL)\nUpload a product image, a mask, and a prompt to generate realistic content.")
+    gr.Markdown("## 🧠 NewCrux Inpainting Demo (SD 1.5)\nUpload a product image, a mask, and a prompt to generate realistic content.")
 
     with gr.Row():
         with gr.Column():
@@ -93,4 +92,3 @@ demo.launch()
 
 
 
-