53 lines
2.2 KiB
Python
53 lines
2.2 KiB
Python
import os
|
|
import sys
|
|
import subprocess
|
|
from pathlib import Path
|
|
import pytest
|
|
|
|
NUM_NODES = "1"
|
|
NUM_GPUS_PER_NODE = "2"
|
|
|
|
# Set environment variables
|
|
os.environ["FASTVIDEO_ATTENTION_CONFIG"] = "assets/mask_strategy_wan.json"
|
|
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
|
|
|
|
def test_inference():
|
|
"""Test the inference functionality"""
|
|
# Create command as in wan_14B-STA.sh
|
|
cmd = [
|
|
"fastvideo", "generate",
|
|
"--model-path", "Wan-AI/Wan2.1-T2V-14B-Diffusers",
|
|
"--sp-size", "2",
|
|
"--tp-size", "2",
|
|
"--num-gpus", "2",
|
|
"--height", "768",
|
|
"--width", "1280",
|
|
"--num-frames", "69",
|
|
"--num-inference-steps", "2",
|
|
"--fps", "16",
|
|
"--guidance-scale", "5.0",
|
|
"--flow-shift", "5.0",
|
|
"--prompt", "A majestic lion strides across the golden savanna, its powerful frame glistening under the warm afternoon sun. The tall grass ripples gently in the breeze, enhancing the lion's commanding presence. The tone is vibrant, embodying the raw energy of the wild. Low angle, steady tracking shot, cinematic.",
|
|
"--negative-prompt", "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
|
|
"--seed", "1024",
|
|
"--output-path", "outputs_video/STA_1024/",
|
|
]
|
|
|
|
# Run the command
|
|
subprocess.run(cmd, check=True)
|
|
|
|
# Verify output directory exists
|
|
output_dir = Path("outputs_video/STA_1024/")
|
|
assert output_dir.exists(), f"Output directory {output_dir} does not exist"
|
|
|
|
# Verify that video files were generated
|
|
video_files = list(output_dir.glob("*.mp4"))
|
|
assert len(video_files) > 0, "No video files were generated"
|
|
|
|
# Verify the video file properties
|
|
for video_file in video_files:
|
|
assert video_file.stat().st_size > 0, f"Video file {video_file} is empty"
|
|
|
|
if __name__ == "__main__":
|
|
test_inference()
|