Update app.py

app.py

@@ -18,59 +18,6 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 
-def get_lora_sd_pipeline(
-    ckpt_dir='./output', 
-    base_model_name_or_path=model_id_default, 
-    dtype=torch_dtype, 
-    device=device
-):
-    unet_sub_dir = os.path.join(ckpt_dir, "unet")
-    text_encoder_sub_dir = os.path.join(ckpt_dir, "text_encoder")
-
-    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.unet = PeftModel.from_pretrained(pipe.unet, unet_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()
-
-    pipe.to(device)
-    
-    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/
-                              
-#                               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]
 def infer(
     prompt,
@@ -85,10 +32,26 @@ def infer(
     progress=gr.Progress(track_tqdm=True),
 ):  
     generator = torch.Generator(device).manual_seed(seed)
-    pipe = get_lora_sd_pipeline(base_model_name_or_path=model_id)   
+
+    ckpt_dir='./output'
+    unet_sub_dir = os.path.join(ckpt_dir, "unet")
+    text_encoder_sub_dir = os.path.join(ckpt_dir, "text_encoder")
+
+    if model_id is None:
+        raise ValueError("Please specify the base model name or path")
+
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, 
+                                                   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 dtype in (torch.float16, torch.bfloat16):
+        pipe.unet.half()
+        pipe.text_encoder.half()
+
     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)
+    pipe.to(device)
     
     image = pipe(
         prompt=prompt,