Update app.py

app.py

@@ -25,9 +25,13 @@ pipeline = StableDiffusion3Pipeline.from_pretrained(
 # Load the LoRA trained weights once at the start
 lora_path = "lora_trained_model.pt"  # Ensure this file is uploaded in the Space
 if os.path.exists(lora_path):
-    lora_state_dict = torch.load(lora_path, map_location=device, weights_only=True)
-    pipeline.load_lora_weights(lora_path)  # Load LoRA weights into the pipeline
-    print("✅ LoRA weights loaded successfully!")
+    lora_state_dict = torch.load(lora_path, map_location=device)
+    
+    try:
+        pipeline.load_lora_weights(lora_state_dict)  # Load LoRA weights into the pipeline
+        print("✅ LoRA weights loaded successfully!")
+    except ValueError as e:
+        print(f"❌ Error loading LoRA: {e}")
 else:
     print("⚠️ LoRA file not found! Running base model.")