Update worker_runpod.py

worker_runpod.py

@@ -72,7 +72,7 @@ def download_image(url, download_dir="/content"):
 # validation_image_start = values.get("validation_image_start", "https://example.com/path/to/image.png")
 # downloaded_image_path = download_image(validation_image_start)
 model_id = "/runpod-volume/model"
-transformer = CogVideoXTransformer3DModel.from_pretrained(model_id, subfolder="transformer").to(torch.bfloat16)
+transformer = CogVideoXTransformer3DModel.from_pretrained_2d(model_id, subfolder="transformer").to(torch.bfloat16)
 
 vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae").to(torch.bfloat16)