Update app.py

app.py

@@ -6,51 +6,59 @@ from PIL import Image
 import numpy as np
 import cv2
 
-# Load ControlNet (depth)
+# Detect GPU or fallback to CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load ControlNet model (depth conditioning)
 controlnet = ControlNetModel.from_pretrained(
-    "lllyasviel/sd-controlnet-depth", torch_dtype=torch.float16
-)
+    "lllyasviel/sd-controlnet-depth", torch_dtype=torch.float32
+).to(device)
 
-# Load SD pipeline with ControlNet and enable CUDA for faster generation
+# Load Stable Diffusion with ControlNet
 pipe = StableDiffusionControlNetPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5",
     controlnet=controlnet,
-    torch_dtype=torch.float16
-).to("cuda")
+    torch_dtype=torch.float32
+).to(device)
 
-# Depth model
-depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cuda")
+# Load depth estimation model
+depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to(device)
 depth_processor = DPTFeatureExtractor.from_pretrained("Intel/dpt-hybrid-midas")
 
 def generate(input_image, prompt):
+    # Convert image to RGB
     image = input_image.convert("RGB")
-    inputs = depth_processor(images=image, return_tensors="pt").to("cuda")
+
+    # Prepare depth inputs
+    inputs = depth_processor(images=image, return_tensors="pt").to(device)
 
     with torch.no_grad():
         outputs = depth_model(**inputs)
         depth = outputs.predicted_depth.squeeze().cpu().numpy()
 
+    # Normalize to 0–255 and convert to grayscale PIL
     depth = cv2.normalize(depth, None, 0, 255, norm_type=cv2.NORM_MINMAX)
     depth_image = Image.fromarray(depth.astype(np.uint8))
 
+    # Run image generation
     result = pipe(
         prompt=prompt,
         image=depth_image,
-        height=512,
-        width=768,
-        num_inference_steps=5  # 🔥 Faster generation
+        height=512,   # You can reduce this if slow
+        width=512,
+        num_inference_steps=10
     ).images[0]
 
     return result
 
-# Interface
+# Gradio interface
 gr.Interface(
     fn=generate,
     inputs=[
         gr.Image(type="pil", label="Upload Room Image"),
-        gr.Textbox(label="Enter Interior Style Prompt"),
+        gr.Textbox(label="Enter Interior Style Prompt", placeholder="e.g. modern Japanese living room"),
     ],
     outputs=gr.Image(type="pil", label="Generated Room"),
-    title="🛋️ AI Interior Designer",
-    description="Upload your room and get a styled redesign using ControlNet (Depth). Optimized for fast generation.",
+    title="🛋️ AI Room Redesign (ControlNet + Depth)",
+    description="Upload a room image and get a redesigned version based on your style prompt using ControlNet Depth.",
 ).launch()