diff --git a/.github/scripts/runpod_api.py b/.github/scripts/runpod_api.py index 4d75a692b..4db552fe0 100644 --- a/.github/scripts/runpod_api.py +++ b/.github/scripts/runpod_api.py @@ -144,7 +144,8 @@ def execute_command(pod_id): # Copy the tarball to the pod scp_command = [ "scp", "-o", "StrictHostKeyChecking=no", "-o", - "UserKnownHostsFile=/dev/null", "-P", + "UserKnownHostsFile=/dev/null", "-o", "ServerAliveInterval=60", "-o", + "ServerAliveCountMax=10", "-P", str(ssh_port), "/tmp/repo.tar.gz", f"root@{ssh_ip}:/tmp/" ] subprocess.run(scp_command, check=True) @@ -163,7 +164,8 @@ def execute_command(pod_id): ssh_command = [ "ssh", "-o", "StrictHostKeyChecking=no", "-o", - "UserKnownHostsFile=/dev/null", "-p", + "UserKnownHostsFile=/dev/null", "-o", "ServerAliveInterval=60", "-o", + "ServerAliveCountMax=10", "-p", str(ssh_port), f"root@{ssh_ip}", remote_command ] diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 2ac34310e..00df66327 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -32,6 +32,8 @@ jobs: encoder-test: - 'fastvideo/v1/models/encoders/**' - 'fastvideo/v1/models/loaders/**' + - 'fastvideo/v1/tests/test_clip_encoder.py' + - 'fastvideo/v1/tests/test_llama_encoder.py' encoder-test: needs: change-filter @@ -71,18 +73,56 @@ jobs: --gpu-type "NVIDIA A40" --gpu-count 1 --volume-size 100 - --test-command "./env_setup.sh && - pip install vllm && - pip install pytest && - pytest ./fastvideo/v1/tests/ --ignore=./fastvideo/v1/tests/old_tests/ -s" + --test-command "pip install -e .[test] && + pip install flash-attn==2.7.0.post2 --no-build-isolation && + pytest ./fastvideo/v1/tests/test_clip_encoder.py ./fastvideo/v1/tests/test_llama_encoder.py -s" + + ssim-test: + if: ${{ github.event.pull_request.draft == false }} + runs-on: ubuntu-latest + environment: runpod-runners + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + + - name: Set up SSH key + run: | + mkdir -p ~/.ssh + echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa + chmod 600 ~/.ssh/id_rsa + ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub + + - name: Install dependencies + run: pip install requests + + - name: Run tests on RunPod + env: + JOB_ID: "ssim-test" + RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} + GITHUB_SHA: ${{ github.sha }} + GITHUB_REF: ${{ github.ref }} + GITHUB_REPOSITORY: ${{ github.repository }} + GITHUB_RUN_ID: ${{ github.run_id }} + timeout-minutes: 30 + run: >- + python .github/scripts/runpod_api.py + --gpu-type "NVIDIA A40" + --gpu-count 2 + --disk-size 100 + --volume-size 100 + --test-command "pip install -e .[test] && + pip install flash-attn==2.7.0.post2 --no-build-isolation && + python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan-diffusers --local_dir=data/FastHunyuan-diffusers --repo_type=model && + pytest ./fastvideo/v1/tests/ssim -vs" runpod-cleanup: - needs: [encoder-test] # Add other jobs to this list as you create them - if: >- - always() && - ${{ github.event.pull_request.draft == false && - (contains(needs.*.result, 'success') || - contains(needs.*.result, 'failure')) }} + needs: [encoder-test, ssim-test] # Add other jobs to this list as you create them + if: ${{ always() && github.event.pull_request.draft == false && (contains(needs.*.result, 'success') || contains(needs.*.result, 'failure'))}} runs-on: ubuntu-latest environment: runpod-runners steps: @@ -99,7 +139,7 @@ jobs: - name: Cleanup all RunPod instances env: - JOB_IDS: '["encoder-test"]' # JSON array of job IDs + JOB_IDS: '["encoder-test", "ssim-test"]' # JSON array of job IDs RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} GITHUB_RUN_ID: ${{ github.run_id }} run: python .github/scripts/runpod_cleanup.py diff --git a/.gitignore b/.gitignore index d90d7eb1e..adfc7fe9a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,4 @@ __pycache__ -*.mp4 .ipynb_checkpoints *.pth UCF-101/ @@ -11,6 +10,7 @@ wandb/ *.jpg *.safetensors *.mp4 +!fastvideo/v1/tests/ssim/reference_videos/*.mp4 *.png *.gif *.pth diff --git a/fastvideo/v1/entrypoints/cli/utils.py b/fastvideo/v1/entrypoints/cli/utils.py index 54a8f8143..86a4d22f4 100644 --- a/fastvideo/v1/entrypoints/cli/utils.py +++ b/fastvideo/v1/entrypoints/cli/utils.py @@ -40,12 +40,15 @@ def launch_distributed(num_gpus=None, args=None, master_port=None): logger.info("Running inference with %d GPU(s)", num_gpus) logger.info("Launching command: %s", " ".join(cmd)) + current_env["PYTHONIOENCODING"] = "utf-8" process = subprocess.Popen(cmd, env=current_env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True, - bufsize=1) + bufsize=1, + encoding='utf-8', + errors='replace') if process.stdout: for line in iter(process.stdout.readline, ''): diff --git a/fastvideo/v1/tests/ssim/README.md b/fastvideo/v1/tests/ssim/README.md new file mode 100644 index 000000000..a3cd2c305 --- /dev/null +++ b/fastvideo/v1/tests/ssim/README.md @@ -0,0 +1,35 @@ +The reference videos in the `reference_videos` directory are used as part of an e2e test to ensure consistency in video generation quality across code changes. `test_inference_similarity.py` compares newly generated videos against these references using Structural Similarity Index (SSIM) metrics to detect any regressions in visual quality across code changes. + +The reference videos were generated on commit `66107fd5b8469fed25972feb632cd48887dac451` of the FastVideo codebase. + +## Generation Details + +2 x NVIDIA A40 GPUs + +## Generation Parameters + +{ +"num_gpus": 2, +"model_path": "data/FastHunyuan-diffusers", +"height": 720, +"width": 1280, +"num_frames": 45, +"num_inference_steps": 6, +"guidance_scale": 1, +"embedded_cfg_scale": 6, +"flow_shift": 17, +"seed": 1024, +"sp_size": 2, +"tp_size": 2, +"vae_sp": true, +"use_v1_transformer": true, +"use_v1_vae": true, +"use_v1_text_encoder": true, +"fps": 24 +} + +#### Prompts + +1. Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting." + +2. "A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature." diff --git a/fastvideo/v1/tests/ssim/compute_ssim.py b/fastvideo/v1/tests/ssim/compute_ssim.py new file mode 100644 index 000000000..357322a00 --- /dev/null +++ b/fastvideo/v1/tests/ssim/compute_ssim.py @@ -0,0 +1,143 @@ +import torch +import numpy as np +import os +import argparse +from pytorch_msssim import ssim, ms_ssim +from torchvision.io import read_video + + +def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True): + print(f"Computing SSIM between {video1_path} and {video2_path}...") + + frames1, _, _ = read_video(video1_path, + pts_unit='sec', + output_format="TCHW") + frames2, _, _ = read_video(video2_path, + pts_unit='sec', + output_format="TCHW") + + # Ensure same number of frames + min_frames = min(frames1.shape[0], frames2.shape[0]) + frames1 = frames1[:min_frames] + frames2 = frames2[:min_frames] + + frames1 = frames1.float() / 255.0 + frames2 = frames2.float() / 255.0 + + if torch.cuda.is_available(): + frames1 = frames1.cuda() + frames2 = frames2.cuda() + + ssim_values = [] + + # Process each frame individually + for i in range(min_frames): + img1 = frames1[i:i + 1] + img2 = frames2[i:i + 1] + + with torch.no_grad(): + if use_ms_ssim: + value = ms_ssim(img1, img2, data_range=1.0) + else: + value = ssim(img1, img2, data_range=1.0) + + ssim_values.append(value.item()) + + if ssim_values: + mean_ssim = np.mean(ssim_values) + min_ssim = np.min(ssim_values) + max_ssim = np.max(ssim_values) + min_frame_idx = np.argmin(ssim_values) + max_frame_idx = np.argmax(ssim_values) + + print(f"Mean SSIM: {mean_ssim:.4f}") + print(f"Min SSIM: {min_ssim:.4f} (at frame {min_frame_idx})") + print(f"Max SSIM: {max_ssim:.4f} (at frame {max_frame_idx})") + + return mean_ssim, min_ssim, max_ssim + else: + print('No SSIM values calculated') + return 0, 0, 0 + + +def compare_folders(reference_folder, generated_folder, use_ms_ssim=True): + """ + Compare videos with the same filename between reference_folder and generated_folder + """ + reference_videos = [ + f for f in os.listdir(reference_folder) if f.endswith(('.mp4')) + ] + + results = {} + + for video_name in reference_videos: + ref_path = os.path.join(reference_folder, video_name) + gen_path = os.path.join(generated_folder, video_name) + + if os.path.exists(gen_path): + print(f"\nComparing {video_name}...") + try: + ssim_value = compute_video_ssim_torchvision( + ref_path, gen_path, use_ms_ssim) + results[video_name] = ssim_value + except Exception as e: + print(f"Error comparing {video_name}: {e}") + results[video_name] = None + else: + print( + f"\nSkipping {video_name} - no matching file in generated folder" + ) + + return results + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description='Compare videos using SSIM/MS-SSIM metrics') + parser.add_argument('--reference', + '-r', + type=str, + help='Path to reference videos directory') + parser.add_argument('--generated', + '-g', + type=str, + help='Path to generated videos directory') + parser.add_argument('--use-ms-ssim', + action='store_true', + help='Use MS-SSIM instead of SSIM') + args = parser.parse_args() + + script_dir = os.path.dirname(os.path.abspath(__file__)) + + reference_folder = args.reference if args.reference else os.path.join( + script_dir, 'reference_videos') + generated_folder = args.generated if args.generated else os.path.join( + script_dir, 'generated_videos') + + if not os.path.exists(reference_folder): + print(f"ERROR: Reference folder {reference_folder} does not exist!") + exit(1) + + if not os.path.exists(generated_folder): + print(f"ERROR: Generated folder {generated_folder} does not exist!") + exit(1) + + print(f"Comparing videos between {reference_folder} and {generated_folder}") + results = compare_folders(reference_folder, generated_folder, + args.use_ms_ssim) + + print("\n===== SSIM Results Summary =====") + for video_name, ssim_value in results.items(): + if ssim_value is not None: + print( + f"{video_name}: {ssim_value[0]:.4f}, Min SSIM: {ssim_value[1]:.4f}, Max SSIM: {ssim_value[2]:.4f}" + ) + else: + print(f"{video_name}: Error during comparison") + + valid_ssims = [v for v in results.values() if v is not None] + if valid_ssims: + avg_ssim = np.mean([v[0] for v in valid_ssims]) + print(f"\nAverage SSIM across all videos: {avg_ssim:.4f}") + else: + print("\nNo valid SSIM values to average") diff --git a/fastvideo/v1/tests/ssim/reference_videos/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 b/fastvideo/v1/tests/ssim/reference_videos/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 new file mode 100644 index 000000000..10cecbdb3 Binary files /dev/null and b/fastvideo/v1/tests/ssim/reference_videos/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 differ diff --git a/fastvideo/v1/tests/ssim/reference_videos/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 b/fastvideo/v1/tests/ssim/reference_videos/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 new file mode 100644 index 000000000..0f3ea33c5 Binary files /dev/null and b/fastvideo/v1/tests/ssim/reference_videos/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 differ diff --git a/fastvideo/v1/tests/ssim/test_inference_similarity.py b/fastvideo/v1/tests/ssim/test_inference_similarity.py new file mode 100644 index 000000000..53792e684 --- /dev/null +++ b/fastvideo/v1/tests/ssim/test_inference_similarity.py @@ -0,0 +1,177 @@ +import os +import pytest +import json +from fastvideo.v1.tests.ssim.compute_ssim import compute_video_ssim_torchvision +from fastvideo.v1.logger import init_logger +from fastvideo.v1.entrypoints.cli.utils import launch_distributed + +logger = init_logger(__name__) + +# Base parameters from the shell script +BASE_PARAMS = { + "num_gpus": 2, + "model_path": "data/FastHunyuan-diffusers", + "height": 720, + "width": 1280, + "num_frames": 45, + "num_inference_steps": 6, + "guidance_scale": 1, + "embedded_cfg_scale": 6, + "flow_shift": 17, + "seed": 1024, + "sp_size": 2, + "tp_size": 2, + "vae_sp": True, + "use_v1_transformer": True, + "use_v1_vae": True, + "use_v1_text_encoder": True, + "fps": 24, +} + +TEST_PROMPTS = [ + "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.", + "A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature." +] + + +def write_ssim_results(output_dir, ssim_values, reference_path, generated_path, + num_inference_steps, prompt): + """ + Write SSIM results to a JSON file in the same directory as the generated videos. + """ + try: + logger.info( + f"Attempting to write SSIM results to directory: {output_dir}") + + if not os.path.exists(output_dir): + os.makedirs(output_dir, exist_ok=True) + + mean_ssim, min_ssim, max_ssim = ssim_values + + result = { + "mean_ssim": mean_ssim, + "min_ssim": min_ssim, + "max_ssim": max_ssim, + "reference_video": reference_path, + "generated_video": generated_path, + "parameters": { + "num_inference_steps": num_inference_steps, + "prompt": prompt + } + } + + test_name = f"steps{num_inference_steps}_{prompt[:100]}" + result_file = os.path.join(output_dir, f"{test_name}_ssim.json") + logger.info(f"Writing JSON results to: {result_file}") + with open(result_file, 'w') as f: + json.dump(result, f, indent=2) + + logger.info(f"SSIM results written to {result_file}") + return True + except Exception as e: + logger.error(f"ERROR writing SSIM results: {str(e)}") + return False + + +@pytest.mark.parametrize("num_inference_steps", [6]) +@pytest.mark.parametrize("prompt", TEST_PROMPTS) +def test_inference_similarity(num_inference_steps, prompt): + """ + Test that runs inference with different parameters and compares the output + to reference videos using SSIM. + """ + script_dir = os.path.dirname(os.path.abspath(__file__)) + + base_output_dir = os.path.join(script_dir, 'generated_videos') + output_dir = os.path.join(base_output_dir, + f'num_inference_steps={num_inference_steps}') + output_video_name = f"{prompt[:100]}.mp4" + + os.makedirs(output_dir, exist_ok=True) + + launch_args = [ + "--num-inference-steps", + str(num_inference_steps), + "--prompt", + prompt, + "--output-path", + output_dir, + "--model-path", + BASE_PARAMS["model_path"], + "--height", + str(BASE_PARAMS["height"]), + "--width", + str(BASE_PARAMS["width"]), + "--num-frames", + str(BASE_PARAMS["num_frames"]), + "--guidance-scale", + str(BASE_PARAMS["guidance_scale"]), + "--embedded-cfg-scale", + str(BASE_PARAMS["embedded_cfg_scale"]), + "--flow-shift", + str(BASE_PARAMS["flow_shift"]), + "--seed", + str(BASE_PARAMS["seed"]), + "--sp-size", + str(BASE_PARAMS["sp_size"]), + "--tp-size", + str(BASE_PARAMS["tp_size"]), + "--fps", + str(BASE_PARAMS["fps"]), + ] + + if BASE_PARAMS["use_v1_transformer"]: + launch_args.append("--use-v1-transformer") + if BASE_PARAMS["use_v1_vae"]: + launch_args.append("--use-v1-vae") + if BASE_PARAMS["use_v1_text_encoder"]: + launch_args.append("--use-v1-text-encoder") + if BASE_PARAMS["vae_sp"]: + launch_args.append("--vae-sp") + + launch_distributed(num_gpus=BASE_PARAMS["num_gpus"], args=launch_args) + + assert os.path.exists( + output_dir), f"Output video was not generated at {output_dir}" + + reference_folder = os.path.join(script_dir, 'reference_videos') + + if not os.path.exists(reference_folder): + logger.error("Reference folder missing") + raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") + + # Find the matching reference video based on the prompt + reference_video_name = None + + for filename in os.listdir(reference_folder): + if filename.endswith('.mp4') and prompt[:100] in filename: + reference_video_name = filename + break + + if not reference_video_name: + logger.error(f"Reference video not found for prompt: {prompt}") + raise FileNotFoundError(f"Reference video missing") + + reference_video_path = os.path.join(reference_folder, reference_video_name) + generated_video_path = os.path.join(output_dir, output_video_name) + + logger.info( + f"Computing SSIM between {reference_video_path} and {generated_video_path}" + ) + ssim_values = compute_video_ssim_torchvision(reference_video_path, + generated_video_path, + use_ms_ssim=True) + + mean_ssim = ssim_values[0] + logger.info(f"SSIM mean value: {mean_ssim}") + logger.info(f"Writing SSIM results to directory: {output_dir}") + + success = write_ssim_results(output_dir, ssim_values, reference_video_path, + generated_video_path, num_inference_steps, + prompt) + + if not success: + logger.error("Failed to write SSIM results to file") + + min_acceptable_ssim = 1 + assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim}"