Update worker_runpod.py

worker_runpod.py

@@ -25,7 +25,7 @@ tokenxf = os.getenv("HF_API_TOKEN")
 # Low GPU memory mode
 low_gpu_memory_mode = False
 lora_path = "/content/shirtlift.safetensors"
-
+weight_dtype = torch.bfloat16
 def to_pil(image):
     if isinstance(image, Image.Image):
         return image
@@ -75,13 +75,13 @@ with torch.inference_mode():
     model_id = "/runpod-volume/model"
     transformer = CogVideoXTransformer3DModel.from_pretrained_2d(
         model_id, subfolder="transformer"
-    ).to(torch.bfloat16)
+    ).to(weight_dtype)
 
     vae = AutoencoderKLCogVideoX.from_pretrained(
         model_id, subfolder="vae"
-    ).to(torch.bfloat16)
+    ).to(weight_dtype)
 
-    text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder")
+    text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=weight_dtype)
 
     sampler_dict = {
         "Euler": EulerDiscreteScheduler,
@@ -102,7 +102,7 @@ with torch.inference_mode():
             text_encoder=text_encoder,
             transformer=transformer, 
             scheduler=scheduler,
-            torch_dtype=torch.bfloat16
+            torch_dtype=weight_dtype
         )
     else:
         pipeline = CogVideoX_Fun_Pipeline.from_pretrained(
@@ -111,16 +111,16 @@ with torch.inference_mode():
             text_encoder=text_encoder,
             transformer=transformer, 
             scheduler=scheduler,
-            torch_dtype=torch.bfloat16
+            torch_dtype=weight_dtype
         )
 
     pipeline = merge_lora(pipeline, lora_path, 1.00)
 
 
-# if low_gpu_memory_mode:
-#     pipeline.enable_sequential_cpu_offload()
-# else:
-#     pipeline.enable_model_cpu_offload()
+if low_gpu_memory_mode:
+    pipeline.enable_sequential_cpu_offload()
+else:
+    pipeline.enable_model_cpu_offload()
 
 
 
@@ -130,13 +130,13 @@ def generate(input):
     prompt = values["prompt"]
     print("starting Generate function")
     print(prompt)
-    negative_prompt = values.get("negative_prompt", "blurry, blurred, blurry face")
+    negative_prompt = values.get("negative_prompt", "The video is not of a high quality, it has a low resolution. Watermark present in each frame. Strange motion trajectory. blurry, blurred, grainy, distortion, blurry face")
     guidance_scale = values.get("guidance_scale", 6.0)
     seed = values.get("seed", 42)
     num_inference_steps = values.get("num_inference_steps", 18)
     base_resolution = values.get("base_resolution", 512)
     
-    video_length = values.get("video_length", 53)
+    video_length = values.get("video_length", 49)
     fps = values.get("fps", 10)
     
     save_path = "samples"