nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add ssim test (#314)

Browse Source 
Co-authored-by: kevin314 <kevin.lin.cs1@gmail.com>
This commit is contained in:
Kevin Lin
2025-04-05 17:35:45 -07:00
committed by GitHub
co-authored by kevin314
parent 4324c1c21d
commit 8a3e7aa761
9 changed files with 415 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/scripts/runpod_api.py
+4 -2
View File
@@ -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
]
.github/workflows/pr-test.yml
+51 -11
View File
@@ -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
.gitignore
+1 -1
View File
@@ -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
fastvideo/v1/entrypoints/cli/utils.py
+4 -1
View File
@@ -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, ''):
fastvideo/v1/tests/ssim/README.md
+35
View File
@@ -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."
fastvideo/v1/tests/ssim/compute_ssim.py
+143
View File
@@ -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")
fastvideo/v1/tests/ssim/reference_videos/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4
BIN
View File
Binary file not shown.
fastvideo/v1/tests/ssim/reference_videos/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4
BIN
View File
Binary file not shown.
fastvideo/v1/tests/ssim/test_inference_similarity.py
+177
View File
@@ -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}"