Update app.py

app.py

@@ -26,21 +26,13 @@ def get_lora_sd_pipeline(
     device=device
 ):
     unet_sub_dir = os.path.join(ckpt_dir, "unet")
-    text_encoder_sub_dir = os.path.join(ckpt_dir, "text_encoder")
-    if os.path.exists(text_encoder_sub_dir) and base_model_name_or_path is None:
-        config = LoraConfig.from_pretrained(text_encoder_sub_dir)
-        base_model_name_or_path = config.base_model_name_or_path
-
+    
     if base_model_name_or_path is None:
         raise ValueError("Please specify the base model name or path")
 
     pipe = StableDiffusionPipeline.from_pretrained(base_model_name_or_path, torch_dtype=dtype).to(device)
-    pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir)
-
-    if os.path.exists(text_encoder_sub_dir):
-        pipe.text_encoder = PeftModel.from_pretrained(
-            pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name
-        )
+    pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, torch_dtype=torch_dtype)
+    pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, base_model_name_or_path, subfolder="text_encoder", torch_dtype=torch_dtype)
 
     if dtype in (torch.float16, torch.bfloat16):
         pipe.unet.half()
@@ -95,6 +87,11 @@ def infer(
     # pipe.fuse_lora(lora_scale=lora_scale)
     # prompt_embeds = encode_prompt(prompt, pipe.tokenizer, pipe.text_encoder)
     # negative_prompt_embeds = encode_prompt(negative_prompt, pipe.tokenizer, pipe.text_encoder)
+    if hasattr(pipe.unet, "set_lora_scale"):
+        pipe.unet.set_lora_scale(lora_scale)
+    else:
+        print("Warning: LoRA scale adjustment method not found on UNet.")
+
     
     image = pipe(
         prompt=prompt,