don't generate synthetic when load vbench fails

This commit is contained in:
Tara Ding 2026-04-27 22:28:42 -07:00
parent c39f7ea76c
commit 54ced2923b

View File

@ -144,13 +144,9 @@ def download_models(base_dir: Path, model: str, task: str) -> None:
def _try_download_vbench_i2v(input_dir: Path) -> list[str]: def _try_download_vbench_i2v(input_dir: Path) -> list[str]:
""" """
Download VBench I2V origin images from Google Drive via gdown (pip install gdown). Download VBench I2V origin images from Google Drive via gdown (pip install gdown).
Returns image basenames placed in *input_dir*, or [] on failure. Raises on any failure.
""" """
try: import gdown # type: ignore; raises ImportError if not installed
import gdown # type: ignore
except ImportError:
print("[setup] gdown not available; skipping VBench download. Install with: pip install gdown")
return []
import zipfile import zipfile
@ -163,11 +159,10 @@ def _try_download_vbench_i2v(input_dir: Path) -> list[str]:
with zipfile.ZipFile(zip_path, "r") as zf: with zipfile.ZipFile(zip_path, "r") as zf:
zf.extractall(str(input_dir)) zf.extractall(str(input_dir))
zip_path.unlink() zip_path.unlink()
except Exception as exc: except Exception:
print(f"[setup] VBench I2V download failed: {exc}")
if zip_path.exists(): if zip_path.exists():
zip_path.unlink() zip_path.unlink()
return [] raise
image_exts = {".png", ".jpg", ".jpeg", ".webp"} image_exts = {".png", ".jpg", ".jpeg", ".webp"}
filenames = sorted( filenames = sorted(
@ -206,20 +201,14 @@ def prepare_input_images(
) -> list[str]: ) -> list[str]:
""" """
Prepare benchmark input images in *input_dir*. Prepare benchmark input images in *input_dir*.
For "vbench_i2v", downloads from Google Drive and raises on failure.
Priority: Falls back to synthetic images only when image_source is not "vbench_i2v".
1. Reuse any images already present in the directory. Returns a list of image paths relative to *input_dir*.
2. Fetch from the source specified by *image_source* (e.g. "vbench_i2v").
3. Generate synthetic 720×720 white PNG placeholders with Pillow.
Returns a list of image basenames (not full paths).
""" """
input_dir.mkdir(parents=True, exist_ok=True) input_dir.mkdir(parents=True, exist_ok=True)
if image_source == "vbench_i2v": if image_source == "vbench_i2v":
filenames = _try_download_vbench_i2v(input_dir) return _try_download_vbench_i2v(input_dir)
if filenames:
return filenames
print(f"[setup] generating {num_images} synthetic 720×720 placeholder images ...") print(f"[setup] generating {num_images} synthetic 720×720 placeholder images ...")
return _generate_synthetic_images(input_dir, num_images) return _generate_synthetic_images(input_dir, num_images)