Update worker_runpod.py

worker_runpod.py

@@ -71,17 +71,11 @@ def download_image(url, download_dir="asset"):
 # 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 = "/content/model"
-transformer = CogVideoXTransformer3DModel.from_pretrained_2d(
-    model_id, subfolder="transformer", torch_dtype=torch.bfloat16
-).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)
+vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae").to(torch.bfloat16)
 
-text_encoder = T5EncoderModel.from_pretrained(
-    model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16
-)
+text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder")
 
 sampler_dict = {
     "Euler": EulerDiscreteScheduler,