Update worker_runpod.py

worker_runpod.py

@@ -13,6 +13,7 @@ with torch.inference_mode():
     vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float16)
     tokenizer = T5Tokenizer.from_pretrained(model_id, subfolder="tokenizer")
     pipe = CogVideoXImageToVideoPipeline.from_pretrained(model_id, tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, torch_dtype=torch.float16).to("cuda")
+    pipe = merge_lora(pipe, lora_path, lora_weight)
     # pipe.enable_model_cpu_offload()
 
 def download_file(url, save_dir, file_name):
@@ -31,7 +32,6 @@ def generate(input):
     values = input["input"]
     lora_path = "/content/shirtlift.safetensors"
     lora_weight = 1.0
-    pipe = merge_lora(pipe, lora_path, lora_weight)
     input_image = values['input_image_check']
     input_image = download_file(url=input_image, save_dir='/content/input', file_name='input_image_tost')
     prompt = values['prompt']