Update app.py

app.py

@@ -14,7 +14,10 @@ else:
 # Load the Stable Diffusion 3.5 model with lower precision (float16)
 model_id = "stabilityai/stable-diffusion-3.5-large"
 pipe = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype=torch.float16)  # Use float16 precision
-pipe.to("cuda")  # Ensuring it runs on GPU
+
+# Check for GPU availability and set device accordingly
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe.to(device)  # Use GPU if available, otherwise fallback to CPU
 
 # Define the path to the LoRA model
 lora_model_path = "./lora_model.pth"  # Assuming the file is saved locally
@@ -22,7 +25,7 @@ lora_model_path = "./lora_model.pth"  # Assuming the file is saved locally
 # Custom method to load and apply LoRA weights to the Stable Diffusion pipeline
 def load_lora_model(pipe, lora_model_path):
     # Load the LoRA weights
-    lora_weights = torch.load(lora_model_path, map_location="cpu")
+    lora_weights = torch.load(lora_model_path, map_location=device)  # Use correct device
     
     # Apply weights to the UNet submodule
     for name, param in pipe.unet.named_parameters():  # Accessing unet parameters
@@ -50,4 +53,4 @@ iface = gr.Interface(
     ],
     outputs="image"
 )
-iface.launch()
+iface.launch()