Update app.py

app.py

@@ -13,23 +13,23 @@ else:
 
 # Load the Stable Diffusion 3.5 model
 model_id = "stabilityai/stable-diffusion-3.5-medium"
-pipe = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipe.to("cpu")
+pipe = StableDiffusion3Pipeline.from_pretrained(model_id)  # Removed torch_dtype argument
+pipe.to("cpu")  # Ensuring it runs on CPU
+
+# Define the path to the LoRA model
+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 (assuming it's a PyTorch .pth file)
+    # Load the LoRA weights
     lora_weights = torch.load(lora_model_path, map_location="cpu")
-
-    # Apply LoRA weights to the parameters of the model
-    for name, param in pipe.named_parameters():
+    
+    # Apply weights to the UNet submodule
+    for name, param in pipe.unet.named_parameters():  # Accessing unet parameters
         if name in lora_weights:
-            param.data += lora_weights[name]  # Apply LoRA weights to the parameters
-
-    return pipe  # Return the updated model
+            param.data += lora_weights[name]
 
-# Define the path to the LoRA model (Local path in your Hugging Face Space)
-lora_model_path = "./lora_model.pth"  # Local path to LoRA model
+    return pipe
 
 # Load and apply the LoRA model weights
 pipe = load_lora_model(pipe, lora_model_path)