Update app.py

app.py

@@ -24,10 +24,12 @@ 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):
+    # Set device to 'cuda' if available, otherwise 'cpu'
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
     # When loading the LoRA weights
     lora_weights = torch.load(lora_model_path, map_location=device, weights_only=True)
 
-    
     # Print available attributes of the model to check access to `unet` (optional)
     print(dir(pipe))  # This will list all attributes and methods of the `pipe` object