Update app.py

app.py

@@ -5,7 +5,6 @@ import os
 import torch
 from diffusers import StableDiffusionPipeline
 from peft import PeftModel, LoraConfig
-from diffusers import DiffusionPipeline
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_id_default = "stable-diffusion-v1-5/stable-diffusion-v1-5"
@@ -27,26 +26,22 @@ def get_lora_sd_pipeline(
 ):
     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, safety_checker=None).to(device)
+    pipe = StableDiffusionPipeline.from_pretrained(base_model_name_or_path, 
+                                                   torch_dtype=dtype, 
+                                                   safety_checker=None).to(device)
     pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir)
+    pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, text_encoder_sub_dir)
 
-    if os.path.exists(text_encoder_sub_dir):
-        pipe.text_encoder = PeftModel.from_pretrained(
-            pipe.text_encoder, text_encoder_sub_dir
-        )
-
-    # if dtype in (torch.float16, torch.bfloat16):
-    #     pipe.unet.half()
-    #     pipe.text_encoder.half()
+    if dtype in (torch.float16, torch.bfloat16):
+        pipe.unet.half()
+        pipe.text_encoder.half()
 
     pipe.to(device)
+    
     return pipe
 
     
@@ -91,7 +86,7 @@ def infer(
 ):  
     generator = torch.Generator(device).manual_seed(seed)
     pipe = get_lora_sd_pipeline(base_model_name_or_path=model_id)   
-    # pipe.fuse_lora(lora_scale=lora_scale)
+    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)
     
@@ -104,6 +99,7 @@ def infer(
         height=height,
         generator=generator,
     ).images[0]
+    print(device, torch_dtype)
     
     return image