Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import gradio as gr
 import torch
-from diffusers import StableDiffusion3Pipeline
+from diffusers import StableDiffusion3Pipeline, SD3LoraLoaderMixin
 import spaces
 import random
 
@@ -25,16 +25,10 @@ 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_checkpoint = torch.load(lora_path, map_location=device)
-    
     try:
-        # Assuming the checkpoint contains 'model' key with the model state dict
-        pipeline.unet.load_state_dict(lora_checkpoint['model'], strict=False)  # Apply weights to the unet
-        
+        pipeline.load_lora_weights(lora_path)  # This automatically applies to the right components
         print("✅ LoRA weights loaded successfully!")
-    except KeyError as e:
-        print(f"❌ Error: Missing key in checkpoint: {e}")
-    except Exception as e:
+    except ValueError as e:
         print(f"❌ Error loading LoRA: {e}")
 else:
     print("⚠️ LoRA file not found! Running base model.")