Update worker_runpod.py

worker_runpod.py

@@ -13,6 +13,8 @@ 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")
+    lora_path = "/content/shirtlift.safetensors"
+    lora_weight = 1.0
     pipe = merge_lora(pipe, lora_path, lora_weight)
     # pipe.enable_model_cpu_offload()
 
@@ -30,8 +32,7 @@ def download_file(url, save_dir, file_name):
 @torch.inference_mode()
 def generate(input):
     values = input["input"]
-    lora_path = "/content/shirtlift.safetensors"
-    lora_weight = 1.0
+
     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']