Update app.py

app.py

@@ -8,7 +8,7 @@ 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"  # Replace to the model you would like to use
+model_id_default = "stable-diffusion-v1-5/stable-diffusion-v1-5"
 
 if torch.cuda.is_available():
     torch_dtype = torch.float16
@@ -20,19 +20,28 @@ MAX_IMAGE_SIZE = 1024
 
 
 def get_lora_sd_pipeline(
-    ckpt_dir='./output', 
+    ckpt_dir='./model_output', 
     base_model_name_or_path=model_id_default, 
     dtype=torch_dtype, 
-    device=device
+    device=device, 
+    adapter_name="pusheen"
 ):
     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, 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)
+    pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)
+
+    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
+        )
 
     if dtype in (torch.float16, torch.bfloat16):
         pipe.unet.half()
@@ -84,7 +93,8 @@ def infer(
     generator = torch.Generator(device).manual_seed(seed)
     pipe = get_lora_sd_pipeline(base_model_name_or_path=model_id)    
     pipe = pipe.to(device)
-    # pipe.fuse_lora(lora_scale=lora_scale)
+    pipe.fuse_lora(lora_scale=lora_scale)
+    pipe.safety_checker = None
     # prompt_embeds = encode_prompt(prompt, pipe.tokenizer, pipe.text_encoder)
     # negative_prompt_embeds = encode_prompt(negative_prompt, pipe.tokenizer, pipe.text_encoder)