fix save dir path

examples/hunyuanvideo/hyvideo/run_vae.py

@@ -1,3 +1,21 @@
+"""
+image input:
+    python run_vae.py --input-type image --image-path "path/to/input_image.jpg" --rec-path "reconstructed_image.jpg"
+video input:
+    python run_vae.py --input-type video --video-path "path/to/input_video.mp4" --rec-path "reconstructed_video.mp4"
+video folder input:
+    python run_vae.py \
+        --input-type folder \
+        --real-video-dir "path/to/real_videos" \
+        --generated-video-dir "path/to/generated_videos" \
+    use parallel inference:
+        msrun --bind_core=True --master_port=8000 --worker_num=8 --local_worker_num=8 --log_dir="parallel_logs" \
+        python run_vae.py \
+            --input-type folder \
+            --real-video-dir "path/to/real_videos" \
+            --generated-video-dir "path/to/generated_videos" \
+            --use-parallel True \
+"""
 import argparse
 import logging
 import os
@@ -20,7 +38,13 @@
 
 from hyvideo.constants import PRECISION_TO_TYPE, PRECISIONS, VAE_PATH
 from hyvideo.dataset.dataset_videobase import VideoDataset, create_dataloader
-from hyvideo.utils.data_utils import preprocess_image, preprocess_video, read_video, transform_to_rgb
+from hyvideo.utils.data_utils import (
+    preprocess_image,
+    preprocess_video,
+    read_video,
+    replace_folder_in_path,
+    transform_to_rgb,
+)
 from hyvideo.utils.ms_utils import init_env
 from hyvideo.vae import load_vae
 from hyvideo.vae.unet_causal_3d_blocks import GroupNorm, MSInterpolate, MSPad
@@ -140,8 +164,8 @@ def process_folder(args, vae, dtype, rank_id, device_num):
             file_path = file_paths[idx]
             if ".avi" in os.path.basename(file_path):
                 file_path = file_path.replace(".avi", ".mp4")
-            output_path = file_path.replace(
-                real_video_dir, generated_video_dir
+            output_path = replace_folder_in_path(
+                file_path, real_video_dir, generated_video_dir
             )  # the same folder structure as the real video folder
             if not os.path.exists(os.path.dirname(output_path)):
                 os.makedirs(os.path.dirname(output_path), exist_ok=True)

examples/hunyuanvideo/hyvideo/utils/data_utils.py

@@ -3,6 +3,7 @@
 import os
 import random
 import sys
+from pathlib import Path
 
 import numpy as np
 from decord import VideoReader, cpu
@@ -21,6 +22,29 @@
 logger = logging.getLogger(__name__)
 
 
+def replace_folder_in_path(original_path, old_folder, new_folder):
+    """
+    Replace a specific folder name in a given path.
+    This function takes an original path, an old folder name, and a new folder name.
+    It then searches for the first occurrence of the old folder name in the path parts
+    and replaces it with the new folder name. Finally, it reconstructs and returns the new path.
+    Args:
+        original_path (str or Path): The original path where the folder name replacement will occur.
+        old_folder (str): The name of the folder to be replaced.
+        new_folder (str): The new name to replace the old folder name with.
+    Returns:
+        Path: The new path after the folder name replacement.
+    """
+    assert old_folder in original_path, f"The folder '{old_folder}' does not exist in the path '{original_path}'."
+    original_path = Path(original_path)
+    path_parts = list(original_path.parts)
+    if old_folder in path_parts:
+        index = path_parts.index(old_folder)
+        path_parts[index] = new_folder
+    new_path = Path(*path_parts)
+    return new_path
+
+
 def crop(image, i, j, h, w):
     if len(image.shape) != 3:
         raise ValueError("image should be a 3D tensor")