Update app.py

app.py

@@ -22,10 +22,7 @@ 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
-
-# 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
+pipe.to(device)  # Ensuring the model is on the correct device (GPU or CPU)
 
 # Define the path to the LoRA model
 lora_model_path = "./lora_model.pth"  # Assuming the file is saved locally
@@ -33,7 +30,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=device)  # Use correct device
+    lora_weights = torch.load(lora_model_path, map_location=device)  # Load LoRA model to the correct device
     
     # Apply weights to the UNet submodule
     for name, param in pipe.unet.named_parameters():  # Accessing unet parameters