Update app.py

app.py

@@ -25,19 +25,17 @@ 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)
+    lora_checkpoint = torch.load(lora_path, map_location=device)
     
     try:
-        # Assuming `pipeline` has a method `load_lora_weights` to load LoRA weights directly
-        # If the pipeline does not support this method, we might need to apply the LoRA weights manually
-        pipeline.load_lora_weights(lora_state_dict)  # Load LoRA weights into the pipeline
+        # 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
+        
         print("✅ LoRA weights loaded successfully!")
-    except AttributeError:
-        print("❌ pipeline does not support load_lora_weights method. Attempting manual application.")
-        # Manual application of weights if load_lora_weights method does not exist
-        # This is just a placeholder; you'll need to update this part based on how your LoRA weights should be applied
-        # Example:
-        # pipeline.model.load_state_dict(lora_state_dict, strict=False)
+    except KeyError as e:
+        print(f"❌ Error: Missing key in checkpoint: {e}")
+    except Exception as e:
+        print(f"❌ Error loading LoRA: {e}")
 else:
     print("⚠️ LoRA file not found! Running base model.")