@pytest.mark.skipif(
device_reference_folder is None,
reason=f"No reference videos for device {device_name}",
)
@pytest.mark.parametrize("prompt", TEST_PROMPTS)
@pytest.mark.parametrize("ATTENTION_BACKEND", ["TORCH_SDPA"])
@pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys()))
def test_gen3c_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
"""
Generate a GEN3C video and compare against the reference using MS-SSIM.
"""
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
script_dir = os.path.dirname(os.path.abspath(__file__))
base_output_dir = os.path.join(script_dir, "generated_videos", model_id)
output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND)
output_video_name = CANDIDATE_VIDEO_NAME
os.makedirs(output_dir, exist_ok=True)
BASE_PARAMS = MODEL_TO_PARAMS[model_id]
num_inference_steps = BASE_PARAMS["num_inference_steps"]
model_path = BASE_PARAMS["model_path"]
# Guard common misconfigurations to keep CI behavior explicit.
if model_path.lower() == "nvidia/gen3c-cosmos-7b":
pytest.skip(
"nvidia/GEN3C-Cosmos-7B is the official raw checkpoint repo, not Diffusers format. "
"Use GEN3C_MODEL_PATH=FastVideo/GEN3C-Cosmos-7B-Diffusers or a local converted path."
)
local_like = model_path.startswith(("/", "./", "../"))
if local_like and not os.path.exists(model_path):
pytest.skip(
f"Local GEN3C model path not found: {model_path}. "
"Set GEN3C_MODEL_PATH to a valid local path or HF Diffusers repo id."
)
if os.path.exists(model_path):
model_index_path = os.path.join(model_path, "model_index.json")
if not os.path.exists(model_index_path):
pytest.skip(
f"GEN3C_MODEL_PATH is not Diffusers-format (missing model_index.json): {model_path}"
)
init_kwargs = {
"num_gpus": BASE_PARAMS["num_gpus"],
"sp_size": BASE_PARAMS["sp_size"],
"tp_size": BASE_PARAMS["tp_size"],
}
if "flow_shift" in BASE_PARAMS:
init_kwargs["flow_shift"] = BASE_PARAMS["flow_shift"]
generation_kwargs = {
"num_inference_steps": num_inference_steps,
"output_path": os.path.join(output_dir, output_video_name),
"height": BASE_PARAMS["height"],
"width": BASE_PARAMS["width"],
"num_frames": BASE_PARAMS["num_frames"],
"guidance_scale": BASE_PARAMS["guidance_scale"],
"embedded_cfg_scale": BASE_PARAMS["embedded_cfg_scale"],
"seed": BASE_PARAMS["seed"],
"image_path": BASE_PARAMS["image_path"],
"fps": BASE_PARAMS["fps"],
}
if not os.path.exists(generation_kwargs["image_path"]):
pytest.skip(
f"GEN3C test image not found: {generation_kwargs['image_path']}. "
"Set GEN3C_TEST_IMAGE_PATH to a valid local image."
)
# Keep local reruns deterministic: remove prior candidate outputs so
# VideoGenerator does not auto-suffix (_1, _2, ...).
stale_pattern = os.path.join(output_dir, "gen3c_ssim_candidate*.mp4")
for stale_video in glob.glob(stale_pattern):
os.remove(stale_video)
generator = VideoGenerator.from_pretrained(
model_path=model_path, **init_kwargs
)
generator.generate_video(prompt, **generation_kwargs)
if isinstance(generator.executor, MultiprocExecutor):
generator.executor.shutdown()
assert os.path.exists(output_dir), f"Output not generated at {output_dir}"
reference_folder = os.path.join(
script_dir, device_reference_folder, model_id, ATTENTION_BACKEND
)
if not os.path.exists(reference_folder):
raise FileNotFoundError(
f"Reference video folder does not exist: {reference_folder}"
)
reference_video_path = os.path.join(reference_folder, BASELINE_VIDEO_NAME)
if not os.path.exists(reference_video_path):
raise FileNotFoundError(
f"Reference video not found: {reference_video_path}"
)
generated_video_path = os.path.join(output_dir, output_video_name)
logger.info(f"Computing SSIM: {reference_video_path} vs {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"GEN3C SSIM mean: {mean_ssim}")
write_ssim_results(
output_dir,
ssim_values,
reference_video_path,
generated_video_path,
num_inference_steps,
prompt,
)
# GEN3C SSIM threshold for stable L40S reference comparisons.
min_acceptable_ssim = 0.93
assert mean_ssim >= min_acceptable_ssim, (
f"SSIM {mean_ssim:.4f} < {min_acceptable_ssim} for {model_id} / {ATTENTION_BACKEND}"
)