Update app.py

app.py

@@ -8,25 +8,23 @@ import cv2
 
 # Load ControlNet (depth)
 controlnet = ControlNetModel.from_pretrained(
-    "lllyasviel/sd-controlnet-depth", torch_dtype=torch.float32
+    "lllyasviel/sd-controlnet-depth", torch_dtype=torch.float16
 )
 
-# Load SD pipeline with ControlNet
+# Load SD pipeline with ControlNet and enable CUDA for faster generation
 pipe = StableDiffusionControlNetPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5",
     controlnet=controlnet,
-    torch_dtype=torch.float32
-).to("cpu")
-
-# ✅ Removed: pipe.enable_model_cpu_offload()
+    torch_dtype=torch.float16
+).to("cuda")
 
 # Depth model
-depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cpu")
+depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cuda")
 depth_processor = DPTFeatureExtractor.from_pretrained("Intel/dpt-hybrid-midas")
 
 def generate(input_image, prompt):
     image = input_image.convert("RGB")
-    inputs = depth_processor(images=image, return_tensors="pt").to("cpu")
+    inputs = depth_processor(images=image, return_tensors="pt").to("cuda")
 
     with torch.no_grad():
         outputs = depth_model(**inputs)
@@ -38,9 +36,9 @@ def generate(input_image, prompt):
     result = pipe(
         prompt=prompt,
         image=depth_image,
-        height=720,
-    width=1280,
-        num_inference_steps=15
+        height=512,
+        width=768,
+        num_inference_steps=5  # 🔥 Faster generation
     ).images[0]
 
     return result
@@ -54,5 +52,5 @@ gr.Interface(
     ],
     outputs=gr.Image(type="pil", label="Generated Room"),
     title="🛋️ AI Interior Designer",
-    description="Upload your room and get a styled redesign using ControlNet (Depth).",
+    description="Upload your room and get a styled redesign using ControlNet (Depth). Optimized for fast generation.",
 ).launch()