Update worker_runpod.py

worker_runpod.py

@@ -1,5 +1,4 @@
 import os, json, requests, random, runpod
-
 import torch
 from diffusers import AutoencoderKLCogVideoX, CogVideoXImageToVideoPipeline, CogVideoXTransformer3DModel
 from cogvideox.utils.lora_utils import merge_lora, unmerge_lora
@@ -8,14 +7,22 @@ from transformers import T5EncoderModel, T5Tokenizer
 
 with torch.inference_mode():
     model_id = "/content/model"
-    transformer = CogVideoXTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.float16)
-    text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.float16)
-    vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float16)
+    
+    # Load models and ensure they are placed on CUDA device
+    transformer = CogVideoXTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.float16).to("cuda")
+    text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.float16).to("cuda")
+    vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float16).to("cuda")
     tokenizer = T5Tokenizer.from_pretrained(model_id, subfolder="tokenizer")
+    
+    # Ensure the pipeline is on the same device (CUDA)
     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)
+    
+    # Merge Lora model and ensure it's on the same device
+    pipe = merge_lora(pipe, lora_path, lora_weight).to("cuda")
+
     # pipe.enable_model_cpu_offload()
 
 def download_file(url, save_dir, file_name):