Update worker_runpod.py

worker_runpod.py

@@ -17,13 +17,25 @@ from cogvideox.models.autoencoder_magvit import AutoencoderKLCogVideoX
 from cogvideox.pipeline.pipeline_cogvideox import CogVideoX_Fun_Pipeline
 from cogvideox.pipeline.pipeline_cogvideox_inpaint import CogVideoX_Fun_Pipeline_Inpaint
 from cogvideox.utils.lora_utils import merge_lora, unmerge_lora
-from cogvideox.utils.utils import get_image_to_video_latent, save_videos_grid, ASPECT_RATIO_512, get_closest_ratio, to_pil
+from cogvideox.utils.utils import get_image_to_video_latent, save_videos_grid
+from cogvideox.data.bucket_sampler import ASPECT_RATIO_512, get_closest_ratio
 from huggingface_hub import HfApi, HfFolder
 
 tokenxf = os.getenv("HF_API_TOKEN")
 # Low GPU memory mode
 low_gpu_memory_mode = False
 lora_path = "/content/shirtlift.safetensors"
+
+def to_pil(image):
+    if isinstance(image, Image.Image):
+        return image
+    if isinstance(image, torch.Tensor):
+        return tensor2pil(image)
+    if isinstance(image, np.ndarray):
+        return numpy2pil(image)
+    raise ValueError(f"Cannot convert {type(image)} to PIL.Image")
+
+
 def download_image(url, download_dir="asset"):
     # Ensure the download directory exists
     if not os.path.exists(download_dir):