Update app.py

app.py

@@ -24,7 +24,8 @@ def get_lora_sd_pipeline(
     base_model_name_or_path=model_id_default, 
     dtype=torch_dtype, 
     device=device, 
-    adapter_name="default"
+    adapter_name="default",
+    lora_scale=1.0
 ):
     unet_sub_dir = os.path.join(ckpt_dir, "unet")
     text_encoder_sub_dir = os.path.join(ckpt_dir, "text_encoder")
@@ -38,6 +39,7 @@ def get_lora_sd_pipeline(
     pipe = StableDiffusionPipeline.from_pretrained(base_model_name_or_path, torch_dtype=dtype).to(device)
     pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)
     pipe.unet.set_adapter(adapter_name)
+    pipe.fuse_lora(lora_scale=lora_scale)
 
     if os.path.exists(text_encoder_sub_dir):
         pipe.text_encoder = PeftModel.from_pretrained(
@@ -52,30 +54,30 @@ def get_lora_sd_pipeline(
     return pipe
 
     
-def encode_prompt(prompt, tokenizer, text_encoder):
-    text_inputs = tokenizer(
-        prompt,
-        padding="max_length",
-        max_length=tokenizer.model_max_length,
-        return_tensors="pt",
-    )
-    with torch.no_grad():
-        if len(text_inputs.input_ids[0]) < tokenizer.model_max_length:
-            prompt_embeds = text_encoder(text_inputs.input_ids.to(text_encoder.device))[0]
-        else:
-            embeds = []
-            start = 0
-            while start < tokenizer.model_max_length:
-                end = start + tokenizer.model_max_length
-                part_of_text_inputs = text_inputs.input_ids[0][start:end]
-                if len(part_of_text_inputs) < tokenizer.model_max_length:
-                    part_of_text_inputs = torch.cat([part_of_text_inputs, torch.tensor([tokenizer.pad_token_id] * (tokenizer.model_max_length - len(part_of_text_inputs)))])
-                embeds.append(text_encoder(part_of_text_inputs.to(text_encoder.device).unsqueeze(0))[0])
-                start += int((8/
+# def encode_prompt(prompt, tokenizer, text_encoder):
+#     text_inputs = tokenizer(
+#         prompt,
+#         padding="max_length",
+#         max_length=tokenizer.model_max_length,
+#         return_tensors="pt",
+#     )
+#     with torch.no_grad():
+#         if len(text_inputs.input_ids[0]) < tokenizer.model_max_length:
+#             prompt_embeds = text_encoder(text_inputs.input_ids.to(text_encoder.device))[0]
+#         else:
+#             embeds = []
+#             start = 0
+#             while start < tokenizer.model_max_length:
+#                 end = start + tokenizer.model_max_length
+#                 part_of_text_inputs = text_inputs.input_ids[0][start:end]
+#                 if len(part_of_text_inputs) < tokenizer.model_max_length:
+#                     part_of_text_inputs = torch.cat([part_of_text_inputs, torch.tensor([tokenizer.pad_token_id] * (tokenizer.model_max_length - len(part_of_text_inputs)))])
+#                 embeds.append(text_encoder(part_of_text_inputs.to(text_encoder.device).unsqueeze(0))[0])
+#                 start += int((8/
                               
-                              11)*tokenizer.model_max_length)
-            prompt_embeds = torch.mean(torch.stack(embeds, dim=0), dim=0)
-    return prompt_embeds
+#                               11)*tokenizer.model_max_length)
+#             prompt_embeds = torch.mean(torch.stack(embeds, dim=0), dim=0)
+#     return prompt_embeds
 
 
 # @spaces.GPU #[uncomment to use ZeroGPU]
@@ -95,13 +97,12 @@ def infer(
     pipe = get_lora_sd_pipeline(base_model_name_or_path=model_id,
                                 adapter_name="sticker_of_funny_cat_Pusheen")    
     pipe = pipe.to(device)
-    prompt_embeds = encode_prompt(prompt, pipe.tokenizer, pipe.text_encoder)
-    negative_prompt_embeds = encode_prompt(negative_prompt, pipe.tokenizer, pipe.text_encoder)
-    # 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)
     
     image = pipe(
-        prompt_embeds=prompt_embeds,
-        negative_prompt_embeds=negative_prompt_embeds,
+        prompt=prompt,
+        negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
         width=width,