Upload app.py

app.py

@@ -40,8 +40,18 @@ def load_pipeline(model_id: str):
         base_model = "CompVis/stable-diffusion-v1-4"
         pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch_dtype)
         # Load the LoRA weights
-        pipe.unet = PeftModel.from_pretrained(pipe.unet, f"{model_id}/unet", torch_dtype=torch_dtype)
-        pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, f"{model_id}/text_encoder", torch_dtype=torch_dtype)
+        pipe.unet = PeftModel.from_pretrained(
+            pipe.unet, 
+            model_id, 
+            subfolder="unet", 
+            torch_dtype=torch_dtype
+        )
+        pipe.text_encoder = PeftModel.from_pretrained(
+            pipe.text_encoder, 
+            model_id, 
+            subfolder="text_encoder", 
+            torch_dtype=torch_dtype
+        )
     else:
         pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)