Update app.py

app.py

@@ -7,13 +7,14 @@ import cv2
 from PIL import Image
 import time
 
+"""
 # Load models from Hugging Face
 sd_model_id = "bhoomikagp/sd2-interior-model-version2" ## test
 # sd_model_id = "bhoomikagp/sd3-interior-model" ## SD3 model issue loading
 controlnet_model_id = "lllyasviel/sd-controlnet-mlsd"
 # Load Stable Diffusion pipeline CUDA
 scheduler = EulerDiscreteScheduler.from_pretrained(sd_model_id, subfolder="scheduler")
-sd_pipeline = StableDiffusionPipeline.from_pretrained(sd_model_id, torch_dtype=torch.float16,scheduler=scheduler).to("cuda")
+#sd_pipeline = StableDiffusionPipeline.from_pretrained(sd_model_id, torch_dtype=torch.float16,scheduler=scheduler).to("cuda")
 # Load ControlNet and Stable Diffusion ControlNet pipeline
 controlnet = ControlNetModel.from_pretrained(controlnet_model_id, torch_dtype=torch.float16).to("cuda")
 controlnet_pipeline = StableDiffusionControlNetPipeline.from_pretrained(
@@ -22,6 +23,21 @@ controlnet_pipeline = StableDiffusionControlNetPipeline.from_pretrained(
     scheduler=scheduler,
     torch_dtype=torch.float16
 ).to("cuda")
+"""
+sd_model_id = "stabilityai/stable-diffusion-2-1"
+scheduler = DPMSolverMultistepScheduler.from_pretrained(sd_model_id, subfolder="scheduler")
+
+# Check if CUDA is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+# Initialize the pipeline with appropriate device and dtype
+torch_dtype = torch.float16 if device == "cuda" else torch.float32
+sd_pipeline = StableDiffusionPipeline.from_pretrained(
+    sd_model_id,
+    torch_dtype=torch_dtype,
+    scheduler=scheduler
+).to(device)
 
 # choices lists
 option_choices = ["living_room", "bedroom", "kitchen"]