Update app.py

app.py

@@ -6,7 +6,7 @@ import random
 from diffusers import StableDiffusion3Pipeline
 from diffusers.loaders import SD3LoraLoaderMixin
 
-# Device selection NOTE
+# Device selection
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
@@ -24,10 +24,10 @@ pipeline = StableDiffusion3Pipeline.from_pretrained(
 ).to(device)
 
 # Load the LoRA trained weights once at the start
-lora_path = "lora_trained_model.pt"  # Ensure this file is uploaded in the Space
+lora_path = "lora_trained_model.safetensors"  # Use the correct file name
 if os.path.exists(lora_path):
     try:
-        pipeline.load_lora_weights(lora_path)  # Use the correct method for loading LoRA weights
+        SD3LoraLoaderMixin.load_lora_into_model(pipeline, lora_path)  # Correct method
         print("✅ LoRA weights loaded successfully!")
     except Exception as e:
         print(f"❌ Error loading LoRA: {e}")
@@ -67,4 +67,4 @@ with gr.Blocks() as demo:
     generate_btn.click(generate_image, inputs=[prompt_input, seed_input], outputs=output_image)
 
 # Launch the Gradio app
-demo.launch()
+demo.launch()