119 lines
3.7 KiB
Python
119 lines
3.7 KiB
Python
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pyarrow.parquet as pq
|
|
import pytest
|
|
import torch
|
|
|
|
from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
|
|
|
|
CAPTION = "a deterministic preprocessing smoke test"
|
|
MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
|
|
|
|
pytestmark = [
|
|
pytest.mark.skipif(os.environ.get("FASTVIDEO_PREPROCESSING_E2E") != "1",
|
|
reason="set FASTVIDEO_PREPROCESSING_E2E=1 to run the GPU integration test"),
|
|
pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"),
|
|
]
|
|
|
|
|
|
def test_t2v_preprocessing_writes_valid_parquet(tmp_path: Path) -> None:
|
|
raw_data_dir = tmp_path / "raw"
|
|
video_dir = raw_data_dir / "videos"
|
|
video_dir.mkdir(parents=True)
|
|
|
|
video_name = "sample.mp4"
|
|
source_video = Path(__file__).parents[1] / "nightly" / "reference_video_1_sample_v0.mp4"
|
|
shutil.copy2(source_video, video_dir / video_name)
|
|
(raw_data_dir / "videos2caption.json").write_text(
|
|
json.dumps([{
|
|
"path": video_name,
|
|
"cap": CAPTION,
|
|
"fps": 16.0,
|
|
"num_frames": 29,
|
|
"resolution": {
|
|
"height": 480,
|
|
"width": 832,
|
|
},
|
|
}]),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
output_dir = tmp_path / "preprocessed"
|
|
subprocess.run(
|
|
[
|
|
sys.executable,
|
|
"-m",
|
|
"torch.distributed.run",
|
|
"--standalone",
|
|
"--nproc-per-node=1",
|
|
"-m",
|
|
"fastvideo.pipelines.preprocess.v1_preprocessing_new",
|
|
"--model-path",
|
|
MODEL_PATH,
|
|
"--mode",
|
|
"preprocess",
|
|
"--workload-type",
|
|
"t2v",
|
|
"--preprocess.video-loader-type",
|
|
"torchvision",
|
|
"--preprocess.dataset-type",
|
|
"merged",
|
|
"--preprocess.dataset-path",
|
|
str(raw_data_dir),
|
|
"--preprocess.dataset-output-dir",
|
|
str(output_dir),
|
|
"--preprocess.preprocess-video-batch-size",
|
|
"1",
|
|
"--preprocess.dataloader-num-workers",
|
|
"0",
|
|
"--preprocess.max-height",
|
|
"64",
|
|
"--preprocess.max-width",
|
|
"64",
|
|
"--preprocess.num-frames",
|
|
"17",
|
|
"--preprocess.train-fps",
|
|
"16",
|
|
"--preprocess.samples-per-file",
|
|
"1",
|
|
"--preprocess.flush-frequency",
|
|
"1",
|
|
"--preprocess.video-length-tolerance-range",
|
|
"5",
|
|
],
|
|
check=True,
|
|
)
|
|
|
|
rank_output_dir = output_dir / "combined_parquet_dataset" / "worker_0"
|
|
parquet_file = rank_output_dir / "worker_0" / "data_chunk_0.parquet"
|
|
assert parquet_file.is_file()
|
|
assert not (rank_output_dir / "worker_0" / "data_chunk_0.parquet.tmp").exists()
|
|
|
|
table = pq.read_table(parquet_file)
|
|
assert table.schema.equals(pyarrow_schema_t2v, check_metadata=False)
|
|
assert table.num_rows == 1
|
|
|
|
row = table.to_pylist()[0]
|
|
assert row["id"] == video_name
|
|
assert row["file_name"] == video_name
|
|
assert row["caption"] == CAPTION
|
|
assert row["media_type"] == "video"
|
|
assert row["width"] == 64
|
|
assert row["height"] == 64
|
|
assert row["fps"] == 16.0
|
|
assert row["duration_sec"] == pytest.approx(17 / 16)
|
|
assert row["vae_latent_bytes"]
|
|
assert len(row["vae_latent_shape"]) == 4
|
|
assert all(size > 0 for size in row["vae_latent_shape"])
|
|
assert row["num_frames"] == row["vae_latent_shape"][1]
|
|
assert row["vae_latent_dtype"]
|
|
assert row["text_embedding_bytes"]
|
|
assert len(row["text_embedding_shape"]) == 2
|
|
assert all(size > 0 for size in row["text_embedding_shape"])
|
|
assert row["text_embedding_dtype"]
|