diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 000000000..41ecc35d4 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,56 @@ + + +# Benchmarks + +The `benchmarks/` directory contains runnable evaluation workflows and benchmark +metadata for measuring MaxDiffusion model outputs. Each subdirectory should own +one benchmark suite or evaluation family. + +Current support: + +* `vbench/`: [VBench](https://github.com/Vchitect/VBench) + video-generation evaluation support, based on the + [original VBench paper](https://arxiv.org/pdf/2311.17982). This currently + covers Wan text-to-video generation and VBench evaluation over a 110-prompt + downsampled subset. + +## Adding benchmark suites + +Add a new subdirectory under `benchmarks/` for each new benchmark or evaluation +family. Keep suite-specific scripts, prompt files, metadata, adapters, and +README instructions inside that subdirectory. + +When adding a new suite, document: + +* what the benchmark measures +* which MaxDiffusion models are supported +* required input files and generated outputs +* how to run generation and evaluation +* where results are written locally or in GCS + +## Current VBench flow + +For the current Wan VBench workflow, generate videos first and then evaluate +them. This compressed evaluation generates and evaluates exactly one video per +prompt: + +```bash +bash benchmarks/vbench/run_tpu_generation.sh --ssh GCS_BUCKET= TPU_NAME= RUN_NAME= +bash benchmarks/vbench/run_gpu_eval.sh --ssh GCS_BUCKET= GPU_NAME= RUN_NAME= +``` + +See `benchmarks/vbench/README.md` for the full script options and data layout. diff --git a/benchmarks/vbench/README.md b/benchmarks/vbench/README.md new file mode 100644 index 000000000..a218500fb --- /dev/null +++ b/benchmarks/vbench/README.md @@ -0,0 +1,141 @@ + + +# VBench + +`benchmarks/vbench/` contains the MaxDiffusion integration for +[VBench](https://github.com/Vchitect/VBench), a video-generation benchmark suite +introduced in the +[original VBench paper](https://arxiv.org/pdf/2311.17982). The current workflow +is intentionally scoped to Wan text-to-video evaluation. Other MaxDiffusion +video models can be added here later by adding model-specific generation +defaults and documenting their output contract. + +The current benchmark data is a 110-prompt downsampled subset of the full VBench +prompt set. This is a compressed workflow: it generates and evaluates exactly +one video per prompt. + +* `prompts_110.txt`: prompts passed to Wan generation +* `VBench_full_info_sub110.json`: VBench metadata for the same prompts and + dimensions + +The prompt file and JSON metadata must stay aligned by order. Prompt `N` in +`prompts_110.txt` should match entry `N` in `VBench_full_info_sub110.json`. + +## Files + +* `run_tpu_generation.sh`: generates Wan videos on a TPU VM, or locally on a + TPU host, and writes videos to GCS. +* `run_gpu_eval.sh`: downloads generated videos, prepares VBench-compatible + filenames, runs VBench evaluation on a GPU VM or local GPU host, and uploads + results to GCS by default. +* `gpu_eval_utils.py`: helper utilities for extracting VBench dimensions, + patching the upstream VBench checkout for compatibility, and preparing the + video manifest layout expected by VBench. +* `prompts_110.txt`: default Wan prompt file for this downsampled evaluation. +* `VBench_full_info_sub110.json`: default VBench metadata used by evaluation. + +## Workflow + +The scripts use a two-stage workflow: + +1. Generate videos with MaxDiffusion Wan on TPU. +2. Evaluate the generated videos with VBench on GPU. + +Use the same `GCS_BUCKET` and `RUN_NAME` for both stages. The generation script +stores videos at: + +```text +gs:////videos +``` + +The evaluation script reads from that location and writes results to: + +```text +gs:////vbench_results +``` + +## Generate Wan Videos + +Run on a TPU VM over SSH: + +```bash +bash benchmarks/vbench/run_tpu_generation.sh --ssh GCS_BUCKET= TPU_NAME= RUN_NAME= +``` + +Or run directly on a TPU host: + +```bash +bash benchmarks/vbench/run_tpu_generation.sh GCS_BUCKET= RUN_NAME= +``` + +Common options: + +* `PROMPT_FILE`: prompt file path. Defaults to + `./benchmarks/vbench/prompts_110.txt`. +* `CONFIG_FILE`: Wan config. Defaults to `src/maxdiffusion/configs/base_wan_27b.yml`. +* `EXTERNAL_DISK`: mounted TPU disk root for large local files. Defaults to + `/mnt/disks/external_disk`. +* `HF_CACHE_ROOT`: Hugging Face cache root. Defaults to + `$EXTERNAL_DISK/hf_cache`. + +The generation script keeps `GCS_VIDEO_DIR` fixed to `${RUN_NAME}/videos`, +because MaxDiffusion writes generated MP4s using +`base_output_directory=gs://${GCS_BUCKET}` and `run_name=${RUN_NAME}`. It uses +the fixed seed `12345`. + +## Run VBench Evaluation + +Run on a GPU VM over SSH: + +```bash +bash benchmarks/vbench/run_gpu_eval.sh --ssh GCS_BUCKET= GPU_NAME= RUN_NAME= +``` + +Or run directly on a GPU host: + +```bash +bash benchmarks/vbench/run_gpu_eval.sh GCS_BUCKET= RUN_NAME= +``` + +This is a compressed one-sample-per-prompt workflow. Evaluation rejects a +missing or duplicate video for any prompt. + +Common options: + +* `BENCHMARK_JSON`: VBench metadata file name. Defaults to + `VBench_full_info_sub110.json`. +* `DIMENSIONS`: space-separated VBench dimensions. Defaults to reading all + dimensions from `BENCHMARK_JSON`. +* `GCS_VIDEO_DIR`: generated video prefix. Defaults to `${RUN_NAME}/videos`. +* `GCS_RESULTS_DIR`: result prefix. Defaults to `${RUN_NAME}/vbench_results`. +* `UPLOAD_RESULTS`: whether to upload results to GCS. Defaults to `true`. +* `WORK_DIR`: working directory on the GPU host. Defaults to + `$HOME/vbench_evaluation`. + +In SSH mode, evaluation results are copied back from the GPU VM to the caller +machine and uploaded from the caller's credentials. This avoids relying on GPU +VM service-account scopes for GCS result uploads. + +## Extending Support + +Today this directory supports Wan VBench evals only. When adding another model, +keep the model-specific generation defaults explicit, document the command, and +ensure its output filenames can be mapped to the VBench prompt order. + +When adding another VBench prompt subset, keep the prompt text file and JSON +metadata together, use matching names, and validate that their entries are in +the same order. diff --git a/benchmarks/vbench/VBench_full_info_sub110.json b/benchmarks/vbench/VBench_full_info_sub110.json new file mode 100644 index 000000000..bc891f110 --- /dev/null +++ b/benchmarks/vbench/VBench_full_info_sub110.json @@ -0,0 +1,1072 @@ +[ + { + "prompt_en": "A tranquil tableau of kitchen", + "dimension": [ + "temporal_flickering" + ] + }, + { + "prompt_en": "A tranquil tableau of alley", + "dimension": [ + "temporal_flickering" + ] + }, + { + "prompt_en": "In a still frame, a park bench with a view of the lake", + "dimension": [ + "temporal_flickering" + ] + }, + { + "prompt_en": "A tranquil tableau of an antique bowl", + "dimension": [ + "temporal_flickering" + ] + }, + { + "prompt_en": "A tranquil tableau of a bunch of grapes", + "dimension": [ + "temporal_flickering" + ] + }, + { + "prompt_en": "In a still frame, phone booth", + "dimension": [ + "temporal_flickering" + ] + }, + { + "prompt_en": "A tranquil tableau of indoor library", + "dimension": [ + "temporal_flickering" + ] + }, + { + "prompt_en": "A tranquil tableau of house", + "dimension": [ + "temporal_flickering" + ] + }, + { + "prompt_en": "A tranquil tableau of in the heart of Plaka, the neoclassical architecture of the old city harmonizes with the ancient ruins", + "dimension": [ + "temporal_flickering" + ] + }, + { + "prompt_en": "A tranquil tableau of bar", + "dimension": [ + "temporal_flickering" + ] + }, + { + "prompt_en": "a horse and a sheep", + "dimension": [ + "multiple_objects" + ], + "auxiliary_info": { + "multiple_objects": { + "object": "horse and sheep" + } + } + }, + { + "prompt_en": "a couch and a potted plant", + "dimension": [ + "multiple_objects" + ], + "auxiliary_info": { + "multiple_objects": { + "object": "couch and potted plant" + } + } + }, + { + "prompt_en": "a teddy bear and a frisbee", + "dimension": [ + "multiple_objects" + ], + "auxiliary_info": { + "multiple_objects": { + "object": "teddy bear and frisbee" + } + } + }, + { + "prompt_en": "skis and a snowboard", + "dimension": [ + "multiple_objects" + ], + "auxiliary_info": { + "multiple_objects": { + "object": "skis and snowboard" + } + } + }, + { + "prompt_en": "an orange and a clock", + "dimension": [ + "multiple_objects" + ], + "auxiliary_info": { + "multiple_objects": { + "object": "orange and clock" + } + } + }, + { + "prompt_en": "a person and a sink", + "dimension": [ + "multiple_objects" + ], + "auxiliary_info": { + "multiple_objects": { + "object": "person and sink" + } + } + }, + { + "prompt_en": "an oven and scissors", + "dimension": [ + "multiple_objects" + ], + "auxiliary_info": { + "multiple_objects": { + "object": "oven and scissors" + } + } + }, + { + "prompt_en": "a vase and scissors", + "dimension": [ + "multiple_objects" + ], + "auxiliary_info": { + "multiple_objects": { + "object": "vase and scissors" + } + } + }, + { + "prompt_en": "a donut and a suitcase", + "dimension": [ + "multiple_objects" + ], + "auxiliary_info": { + "multiple_objects": { + "object": "donut and suitcase" + } + } + }, + { + "prompt_en": "a toothbrush and a sink", + "dimension": [ + "multiple_objects" + ], + "auxiliary_info": { + "multiple_objects": { + "object": "toothbrush and sink" + } + } + }, + { + "prompt_en": "A person is bandaging", + "dimension": [ + "human_action" + ] + }, + { + "prompt_en": "A person is cutting nails", + "dimension": [ + "human_action" + ] + }, + { + "prompt_en": "A person is sword fighting", + "dimension": [ + "human_action" + ] + }, + { + "prompt_en": "A person is folding clothes", + "dimension": [ + "human_action" + ] + }, + { + "prompt_en": "A person is riding a bike", + "dimension": [ + "human_action" + ] + }, + { + "prompt_en": "A person is shining shoes", + "dimension": [ + "human_action" + ] + }, + { + "prompt_en": "A person is finger snapping", + "dimension": [ + "human_action" + ] + }, + { + "prompt_en": "A person is dancing ballet", + "dimension": [ + "human_action" + ] + }, + { + "prompt_en": "A person is cheerleading", + "dimension": [ + "human_action" + ] + }, + { + "prompt_en": "A person is taking a shower", + "dimension": [ + "human_action" + ] + }, + { + "prompt_en": "a boat sailing smoothly on a calm lake", + "dimension": [ + "subject_consistency", + "dynamic_degree", + "motion_smoothness" + ] + }, + { + "prompt_en": "a motorcycle accelerating to gain speed", + "dimension": [ + "subject_consistency", + "dynamic_degree", + "motion_smoothness" + ] + }, + { + "prompt_en": "a train speeding down the tracks", + "dimension": [ + "subject_consistency", + "dynamic_degree", + "motion_smoothness" + ] + }, + { + "prompt_en": "a cat drinking water", + "dimension": [ + "subject_consistency", + "dynamic_degree", + "motion_smoothness" + ] + }, + { + "prompt_en": "a car slowing down to stop", + "dimension": [ + "subject_consistency", + "dynamic_degree", + "motion_smoothness" + ] + }, + { + "prompt_en": "a car stuck in traffic during rush hour", + "dimension": [ + "subject_consistency", + "dynamic_degree", + "motion_smoothness" + ] + }, + { + "prompt_en": "a dog running happily", + "dimension": [ + "subject_consistency", + "dynamic_degree", + "motion_smoothness" + ] + }, + { + "prompt_en": "a car turning a corner", + "dimension": [ + "subject_consistency", + "dynamic_degree", + "motion_smoothness" + ] + }, + { + "prompt_en": "a dog enjoying a peaceful walk", + "dimension": [ + "subject_consistency", + "dynamic_degree", + "motion_smoothness" + ] + }, + { + "prompt_en": "a sheep taking a peaceful walk", + "dimension": [ + "subject_consistency", + "dynamic_degree", + "motion_smoothness" + ] + }, + { + "prompt_en": "a spoon", + "dimension": [ + "object_class" + ], + "auxiliary_info": { + "object_class": { + "object": "spoon" + } + } + }, + { + "prompt_en": "a hair drier", + "dimension": [ + "object_class" + ], + "auxiliary_info": { + "object_class": { + "object": "hair drier" + } + } + }, + { + "prompt_en": "a kite", + "dimension": [ + "object_class" + ], + "auxiliary_info": { + "object_class": { + "object": "kite" + } + } + }, + { + "prompt_en": "a bus", + "dimension": [ + "object_class" + ], + "auxiliary_info": { + "object_class": { + "object": "bus" + } + } + }, + { + "prompt_en": "a potted plant", + "dimension": [ + "object_class" + ], + "auxiliary_info": { + "object_class": { + "object": "potted plant" + } + } + }, + { + "prompt_en": "an oven", + "dimension": [ + "object_class" + ], + "auxiliary_info": { + "object_class": { + "object": "oven" + } + } + }, + { + "prompt_en": "a cat", + "dimension": [ + "object_class" + ], + "auxiliary_info": { + "object_class": { + "object": "cat" + } + } + }, + { + "prompt_en": "a sandwich", + "dimension": [ + "object_class" + ], + "auxiliary_info": { + "object_class": { + "object": "sandwich" + } + } + }, + { + "prompt_en": "a fire hydrant", + "dimension": [ + "object_class" + ], + "auxiliary_info": { + "object_class": { + "object": "fire hydrant" + } + } + }, + { + "prompt_en": "a surfboard", + "dimension": [ + "object_class" + ], + "auxiliary_info": { + "object_class": { + "object": "surfboard" + } + } + }, + { + "prompt_en": "an orange vase", + "dimension": [ + "color" + ], + "auxiliary_info": { + "color": { + "color": "orange" + } + } + }, + { + "prompt_en": "a yellow vase", + "dimension": [ + "color" + ], + "auxiliary_info": { + "color": { + "color": "yellow" + } + } + }, + { + "prompt_en": "a pink suitcase", + "dimension": [ + "color" + ], + "auxiliary_info": { + "color": { + "color": "pink" + } + } + }, + { + "prompt_en": "a pink clock", + "dimension": [ + "color" + ], + "auxiliary_info": { + "color": { + "color": "pink" + } + } + }, + { + "prompt_en": "a pink bird", + "dimension": [ + "color" + ], + "auxiliary_info": { + "color": { + "color": "pink" + } + } + }, + { + "prompt_en": "a white bicycle", + "dimension": [ + "color" + ], + "auxiliary_info": { + "color": { + "color": "white" + } + } + }, + { + "prompt_en": "a purple bicycle", + "dimension": [ + "color" + ], + "auxiliary_info": { + "color": { + "color": "purple" + } + } + }, + { + "prompt_en": "an orange cat", + "dimension": [ + "color" + ], + "auxiliary_info": { + "color": { + "color": "orange" + } + } + }, + { + "prompt_en": "a pink umbrella", + "dimension": [ + "color" + ], + "auxiliary_info": { + "color": { + "color": "pink" + } + } + }, + { + "prompt_en": "a green car", + "dimension": [ + "color" + ], + "auxiliary_info": { + "color": { + "color": "green" + } + } + }, + { + "prompt_en": "A panda drinking coffee in a cafe in Paris by Hokusai, in the style of Ukiyo", + "dimension": [ + "appearance_style" + ], + "auxiliary_info": { + "appearance_style": { + "appearance_style": "by Hokusai, in the style of Ukiyo" + } + } + }, + { + "prompt_en": "The bund Shanghai, black and white", + "dimension": [ + "appearance_style" + ], + "auxiliary_info": { + "appearance_style": { + "appearance_style": "black and white" + } + } + }, + { + "prompt_en": "Gwen Stacy reading a book, black and white", + "dimension": [ + "appearance_style" + ], + "auxiliary_info": { + "appearance_style": { + "appearance_style": "black and white" + } + } + }, + { + "prompt_en": "A panda drinking coffee in a cafe in Paris, surrealism style", + "dimension": [ + "appearance_style" + ], + "auxiliary_info": { + "appearance_style": { + "appearance_style": "surrealism style" + } + } + }, + { + "prompt_en": "A boat sailing leisurely along the Seine River with the Eiffel Tower in background, pixel art", + "dimension": [ + "appearance_style" + ], + "auxiliary_info": { + "appearance_style": { + "appearance_style": "pixel art" + } + } + }, + { + "prompt_en": "Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, Van Gogh style", + "dimension": [ + "appearance_style" + ], + "auxiliary_info": { + "appearance_style": { + "appearance_style": "Van Gogh style" + } + } + }, + { + "prompt_en": "Gwen Stacy reading a book, oil painting", + "dimension": [ + "appearance_style" + ], + "auxiliary_info": { + "appearance_style": { + "appearance_style": "oil painting" + } + } + }, + { + "prompt_en": "a shark is swimming in the ocean by Hokusai, in the style of Ukiyo", + "dimension": [ + "appearance_style" + ], + "auxiliary_info": { + "appearance_style": { + "appearance_style": "by Hokusai, in the style of Ukiyo" + } + } + }, + { + "prompt_en": "Gwen Stacy reading a book by Hokusai, in the style of Ukiyo", + "dimension": [ + "appearance_style" + ], + "auxiliary_info": { + "appearance_style": { + "appearance_style": "by Hokusai, in the style of Ukiyo" + } + } + }, + { + "prompt_en": "Gwen Stacy reading a book, Van Gogh style", + "dimension": [ + "appearance_style" + ], + "auxiliary_info": { + "appearance_style": { + "appearance_style": "Van Gogh style" + } + } + }, + { + "prompt_en": "a shark is swimming in the ocean, tilt down", + "dimension": [ + "temporal_style" + ] + }, + { + "prompt_en": "An astronaut flying in space, tilt up", + "dimension": [ + "temporal_style" + ] + }, + { + "prompt_en": "A panda drinking coffee in a cafe in Paris, pan right", + "dimension": [ + "temporal_style" + ] + }, + { + "prompt_en": "An astronaut flying in space, racking focus", + "dimension": [ + "temporal_style" + ] + }, + { + "prompt_en": "An astronaut flying in space, with an intense shaking effect", + "dimension": [ + "temporal_style" + ] + }, + { + "prompt_en": "An astronaut flying in space, zoom out", + "dimension": [ + "temporal_style" + ] + }, + { + "prompt_en": "A beautiful coastal beach in spring, waves lapping on sand, racking focus", + "dimension": [ + "temporal_style" + ] + }, + { + "prompt_en": "A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, with an intense shaking effect", + "dimension": [ + "temporal_style" + ] + }, + { + "prompt_en": "An astronaut flying in space, zoom in", + "dimension": [ + "temporal_style" + ] + }, + { + "prompt_en": "a shark is swimming in the ocean, zoom in", + "dimension": [ + "temporal_style" + ] + }, + { + "prompt_en": "A car moving slowly on an empty street, rainy evening", + "dimension": [ + "overall_consistency", + "aesthetic_quality", + "imaging_quality" + ] + }, + { + "prompt_en": "A 3D model of a 1800s victorian house.", + "dimension": [ + "overall_consistency", + "aesthetic_quality", + "imaging_quality" + ] + }, + { + "prompt_en": "a shark is swimming in the ocean.", + "dimension": [ + "overall_consistency", + "aesthetic_quality", + "imaging_quality" + ] + }, + { + "prompt_en": "A tropical beach at sunrise, with palm trees and crystal-clear water in the foreground", + "dimension": [ + "overall_consistency", + "aesthetic_quality", + "imaging_quality" + ] + }, + { + "prompt_en": "A raccoon is playing the electronic guitar.", + "dimension": [ + "overall_consistency", + "aesthetic_quality", + "imaging_quality" + ] + }, + { + "prompt_en": "Robot dancing in Times Square.", + "dimension": [ + "overall_consistency", + "aesthetic_quality", + "imaging_quality" + ] + }, + { + "prompt_en": "A raccoon dressed in suit playing the trumpet, stage background", + "dimension": [ + "overall_consistency", + "aesthetic_quality", + "imaging_quality" + ] + }, + { + "prompt_en": "Clown fish swimming through the coral reef", + "dimension": [ + "overall_consistency", + "aesthetic_quality", + "imaging_quality" + ] + }, + { + "prompt_en": "A confused panda in calculus class", + "dimension": [ + "overall_consistency", + "aesthetic_quality", + "imaging_quality" + ] + }, + { + "prompt_en": "Origami dancers in white paper, 3D render, on white background, studio shot, dancing modern dance.", + "dimension": [ + "overall_consistency", + "aesthetic_quality", + "imaging_quality" + ] + }, + { + "prompt_en": "hospital", + "dimension": [ + "scene", + "background_consistency" + ], + "auxiliary_info": { + "scene": { + "scene": { + "scene": "hospital" + } + } + } + }, + { + "prompt_en": "ballroom", + "dimension": [ + "scene", + "background_consistency" + ], + "auxiliary_info": { + "scene": { + "scene": { + "scene": "ballroom" + } + } + } + }, + { + "prompt_en": "driveway", + "dimension": [ + "scene", + "background_consistency" + ], + "auxiliary_info": { + "scene": { + "scene": { + "scene": "driveway" + } + } + } + }, + { + "prompt_en": "art gallery", + "dimension": [ + "scene", + "background_consistency" + ], + "auxiliary_info": { + "scene": { + "scene": { + "scene": "art gallery" + } + } + } + }, + { + "prompt_en": "highway", + "dimension": [ + "scene", + "background_consistency" + ], + "auxiliary_info": { + "scene": { + "scene": { + "scene": "highway" + } + } + } + }, + { + "prompt_en": "mansion", + "dimension": [ + "scene", + "background_consistency" + ], + "auxiliary_info": { + "scene": { + "scene": { + "scene": "mansion" + } + } + } + }, + { + "prompt_en": "fountain", + "dimension": [ + "scene", + "background_consistency" + ], + "auxiliary_info": { + "scene": { + "scene": { + "scene": "fountain" + } + } + } + }, + { + "prompt_en": "bar", + "dimension": [ + "scene", + "background_consistency" + ], + "auxiliary_info": { + "scene": { + "scene": { + "scene": "bar" + } + } + } + }, + { + "prompt_en": "desert", + "dimension": [ + "scene", + "background_consistency" + ], + "auxiliary_info": { + "scene": { + "scene": { + "scene": "desert" + } + } + } + }, + { + "prompt_en": "baseball stadium", + "dimension": [ + "scene", + "background_consistency" + ], + "auxiliary_info": { + "scene": { + "scene": { + "scene": "baseball stadium" + } + } + } + }, + { + "prompt_en": "a toothbrush on the left of a sink, front view", + "dimension": [ + "spatial_relationship" + ], + "auxiliary_info": { + "spatial_relationship": { + "spatial_relationship": { + "object_a": "toothbrush", + "object_b": "sink", + "relationship": "on the left of" + } + } + } + }, + { + "prompt_en": "a potted plant on the left of a remote, front view", + "dimension": [ + "spatial_relationship" + ], + "auxiliary_info": { + "spatial_relationship": { + "spatial_relationship": { + "object_a": "potted plant", + "object_b": "remote", + "relationship": "on the left of" + } + } + } + }, + { + "prompt_en": "an orange on the bottom of a carrot, front view", + "dimension": [ + "spatial_relationship" + ], + "auxiliary_info": { + "spatial_relationship": { + "spatial_relationship": { + "object_a": "orange", + "object_b": "carrot", + "relationship": "on the bottom of" + } + } + } + }, + { + "prompt_en": "an oven on the top of a toaster, front view", + "dimension": [ + "spatial_relationship" + ], + "auxiliary_info": { + "spatial_relationship": { + "spatial_relationship": { + "object_a": "oven", + "object_b": "toaster", + "relationship": "on the top of" + } + } + } + }, + { + "prompt_en": "an apple on the top of a sandwich, front view", + "dimension": [ + "spatial_relationship" + ], + "auxiliary_info": { + "spatial_relationship": { + "spatial_relationship": { + "object_a": "apple", + "object_b": "sandwich", + "relationship": "on the top of" + } + } + } + }, + { + "prompt_en": "a zebra on the left of a giraffe, front view", + "dimension": [ + "spatial_relationship" + ], + "auxiliary_info": { + "spatial_relationship": { + "spatial_relationship": { + "object_a": "zebra", + "object_b": "giraffe", + "relationship": "on the left of" + } + } + } + }, + { + "prompt_en": "a frisbee on the left of a sports ball, front view", + "dimension": [ + "spatial_relationship" + ], + "auxiliary_info": { + "spatial_relationship": { + "spatial_relationship": { + "object_a": "frisbee", + "object_b": "sports ball", + "relationship": "on the left of" + } + } + } + }, + { + "prompt_en": "a bear on the right of a zebra, front view", + "dimension": [ + "spatial_relationship" + ], + "auxiliary_info": { + "spatial_relationship": { + "spatial_relationship": { + "object_a": "bear", + "object_b": "zebra", + "relationship": "on the right of" + } + } + } + }, + { + "prompt_en": "scissors on the left of a teddy bear, front view", + "dimension": [ + "spatial_relationship" + ], + "auxiliary_info": { + "spatial_relationship": { + "spatial_relationship": { + "object_a": "scissors", + "object_b": "teddy bear", + "relationship": "on the left of" + } + } + } + }, + { + "prompt_en": "a donut on the bottom of broccoli, front view", + "dimension": [ + "spatial_relationship" + ], + "auxiliary_info": { + "spatial_relationship": { + "spatial_relationship": { + "object_a": "donut", + "object_b": "broccoli", + "relationship": "on the bottom of" + } + } + } + } +] \ No newline at end of file diff --git a/benchmarks/vbench/gpu_eval_utils.py b/benchmarks/vbench/gpu_eval_utils.py new file mode 100644 index 000000000..a497bf404 --- /dev/null +++ b/benchmarks/vbench/gpu_eval_utils.py @@ -0,0 +1,174 @@ +#!/usr/bin/env python3 +# +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Small helpers for run_gpu_eval.sh.""" + +from __future__ import annotations + +import argparse +import glob +import json +import os +from pathlib import Path +import re +import shutil + + +def extract_dimensions(args: argparse.Namespace) -> None: + with open(args.json_file, encoding="utf-8") as f: + data = json.load(f) + + dimensions = dict.fromkeys( + dimension + for item in data + for dimension in item.get("dimension", []) + ) + print(" ".join(dimensions)) + + +def patch_file(path: Path, replacements: list[tuple[str, str]]) -> None: + if not path.exists(): + return + + text = path.read_text(encoding="utf-8") + for old, new in replacements: + if old not in text: + print(f"Warning: patch target {old!r} not found in {path}. Upstream may have changed.") + text = text.replace(old, new) + path.write_text(text, encoding="utf-8") + + +def patch_vbench(args: argparse.Namespace) -> None: + root = Path(args.vbench_dir) + patch_file( + root / "setup.py", + [("def check_torch_version():", "def check_torch_version():\n return")], + ) + patch_file( + root / "vbench/distributed.py", + [ + ( + "backend = 'gloo' if os.name == 'nt' else 'nccl'", + "backend = 'gloo' if (os.name == 'nt' or not torch.cuda.is_available()) else 'nccl'", + ), + ( + "torch.cuda.set_device(int(os.environ.get('LOCAL_RANK', '0')))", + "if torch.cuda.is_available():\n torch.cuda.set_device(int(os.environ.get('LOCAL_RANK', '0')))", + ), + ], + ) + patch_file( + root / "evaluate.py", + [ + ( + 'device = torch.device("cuda")', + 'assert torch.cuda.is_available(), "CUDA is not available, but is required for VBench evaluation."\n' + ' device = torch.device(f"cuda:{int(os.environ.get(\'LOCAL_RANK\', \'0\'))}")', + ), + ], + ) + patch_file( + root / "vbench/__init__.py", + [ + ("import os\n", "import os\nimport re\n"), + ( + ( + " for prompt_dict in full_info_list:\n" + " # if the prompt belongs to any dimension we want to evaluate\n" + ' if set(dimension_list) & set(prompt_dict["dimension"]): \n' + " prompt = prompt_dict['prompt_en']\n" + " prompt_dict['video_list'] = []\n" + " for i in range(5): # video index for the same prompt\n" + " intended_video_name = f'{prompt}{special_str}-{str(i)}{postfix}'" + ), + ( + " for idx, prompt_dict in enumerate(full_info_list):\n" + " # if the prompt belongs to any dimension we want to evaluate\n" + ' if set(dimension_list) & set(prompt_dict["dimension"]): \n' + " prompt = prompt_dict['prompt_en']\n" + " safe_prompt = re.sub(r'[^\\w\\s-]', '_', prompt).strip()[:120]\n" + " prompt_dict['video_list'] = []\n" + " for i in range(5): # video index for the same prompt\n" + " intended_video_name = f'{safe_prompt}_{idx}-{str(i)}{postfix}'\n" + " if intended_video_name not in video_names:\n" + " intended_video_name = f'{prompt}{special_str}-{str(i)}{postfix}'" + ), + ), + ], + ) + + +def prepare_videos(args: argparse.Namespace) -> None: + with open(args.json_file, encoding="utf-8") as f: + bench_data = json.load(f) + + downloaded = sorted(glob.glob(os.path.join(args.download_dir, "*.mp4"))) + print(f"Downloaded {len(downloaded)} videos from GCS.") + + prompt_video_map: dict[int, list[str]] = {} + for video_path in downloaded: + match = re.search(r"_(\d+)\.mp4$", os.path.basename(video_path)) + if match: + prompt_video_map.setdefault(int(match.group(1)), []).append(video_path) + + prepared_videos: list[tuple[str, int, str]] = [] + for idx, item in enumerate(bench_data): + prompt = item["prompt_en"] + safe_prompt = re.sub(r"[^\w\s-]", "_", prompt).strip()[:120] + candidates = prompt_video_map.get(idx, []) + if len(candidates) != 1: + raise ValueError( + f"Expected exactly one video for prompt index {idx} " + f"({prompt[:40]!r}), but found {len(candidates)}." + ) + prepared_videos.append((safe_prompt, idx, candidates[0])) + + shutil.rmtree(args.vbench_dir, ignore_errors=True) + os.makedirs(args.vbench_dir, exist_ok=True) + for safe_prompt, idx, source_path in prepared_videos: + target_path = os.path.join(args.vbench_dir, f"{safe_prompt}_{idx}-0.mp4") + try: + os.symlink(os.path.abspath(source_path), target_path) + except OSError: + shutil.copy2(source_path, target_path) + + print(f"Prepared {len(prepared_videos)} VBench video entries.") + + +def main() -> None: + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers(required=True) + + dimensions = subparsers.add_parser("dimensions") + dimensions.add_argument("json_file") + dimensions.set_defaults(func=extract_dimensions) + + patch = subparsers.add_parser("patch-vbench") + patch.add_argument("vbench_dir") + patch.set_defaults(func=patch_vbench) + + prepare = subparsers.add_parser("prepare-videos") + prepare.add_argument("json_file") + prepare.add_argument("download_dir") + prepare.add_argument("vbench_dir") + prepare.set_defaults(func=prepare_videos) + + args = parser.parse_args() + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/vbench/prompts_110.txt b/benchmarks/vbench/prompts_110.txt new file mode 100644 index 000000000..c548dc4d5 --- /dev/null +++ b/benchmarks/vbench/prompts_110.txt @@ -0,0 +1,110 @@ +A tranquil tableau of kitchen +A tranquil tableau of alley +In a still frame, a park bench with a view of the lake +A tranquil tableau of an antique bowl +A tranquil tableau of a bunch of grapes +In a still frame, phone booth +A tranquil tableau of indoor library +A tranquil tableau of house +A tranquil tableau of in the heart of Plaka, the neoclassical architecture of the old city harmonizes with the ancient ruins +A tranquil tableau of bar +a horse and a sheep +a couch and a potted plant +a teddy bear and a frisbee +skis and a snowboard +an orange and a clock +a person and a sink +an oven and scissors +a vase and scissors +a donut and a suitcase +a toothbrush and a sink +A person is bandaging +A person is cutting nails +A person is sword fighting +A person is folding clothes +A person is riding a bike +A person is shining shoes +A person is finger snapping +A person is dancing ballet +A person is cheerleading +A person is taking a shower +a boat sailing smoothly on a calm lake +a motorcycle accelerating to gain speed +a train speeding down the tracks +a cat drinking water +a car slowing down to stop +a car stuck in traffic during rush hour +a dog running happily +a car turning a corner +a dog enjoying a peaceful walk +a sheep taking a peaceful walk +a spoon +a hair drier +a kite +a bus +a potted plant +an oven +a cat +a sandwich +a fire hydrant +a surfboard +an orange vase +a yellow vase +a pink suitcase +a pink clock +a pink bird +a white bicycle +a purple bicycle +an orange cat +a pink umbrella +a green car +A panda drinking coffee in a cafe in Paris by Hokusai, in the style of Ukiyo +The bund Shanghai, black and white +Gwen Stacy reading a book, black and white +A panda drinking coffee in a cafe in Paris, surrealism style +A boat sailing leisurely along the Seine River with the Eiffel Tower in background, pixel art +Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, Van Gogh style +Gwen Stacy reading a book, oil painting +a shark is swimming in the ocean by Hokusai, in the style of Ukiyo +Gwen Stacy reading a book by Hokusai, in the style of Ukiyo +Gwen Stacy reading a book, Van Gogh style +a shark is swimming in the ocean, tilt down +An astronaut flying in space, tilt up +A panda drinking coffee in a cafe in Paris, pan right +An astronaut flying in space, racking focus +An astronaut flying in space, with an intense shaking effect +An astronaut flying in space, zoom out +A beautiful coastal beach in spring, waves lapping on sand, racking focus +A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, with an intense shaking effect +An astronaut flying in space, zoom in +a shark is swimming in the ocean, zoom in +A car moving slowly on an empty street, rainy evening +A 3D model of a 1800s victorian house. +a shark is swimming in the ocean. +A tropical beach at sunrise, with palm trees and crystal-clear water in the foreground +A raccoon is playing the electronic guitar. +Robot dancing in Times Square. +A raccoon dressed in suit playing the trumpet, stage background +Clown fish swimming through the coral reef +A confused panda in calculus class +Origami dancers in white paper, 3D render, on white background, studio shot, dancing modern dance. +hospital +ballroom +driveway +art gallery +highway +mansion +fountain +bar +desert +baseball stadium +a toothbrush on the left of a sink, front view +a potted plant on the left of a remote, front view +an orange on the bottom of a carrot, front view +an oven on the top of a toaster, front view +an apple on the top of a sandwich, front view +a zebra on the left of a giraffe, front view +a frisbee on the left of a sports ball, front view +a bear on the right of a zebra, front view +scissors on the left of a teddy bear, front view +a donut on the bottom of broccoli, front view diff --git a/benchmarks/vbench/run_gpu_eval.sh b/benchmarks/vbench/run_gpu_eval.sh new file mode 100644 index 000000000..696f58ba5 --- /dev/null +++ b/benchmarks/vbench/run_gpu_eval.sh @@ -0,0 +1,536 @@ +#!/usr/bin/env bash +# +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# ============================================================================== +# VBench GPU Evaluation Script +# +# Usage: +# bash benchmarks/vbench/run_gpu_eval.sh GCS_BUCKET= [KEY=VALUE ...] +# bash benchmarks/vbench/run_gpu_eval.sh --ssh GCS_BUCKET= GPU_NAME= [KEY=VALUE ...] +# ============================================================================== + +set -euo pipefail + +GCS_BUCKET="${GCS_BUCKET:-}" +RUN_NAME="${RUN_NAME:-wan-inference-aug-31-1}" +GCS_VIDEO_DIR="${GCS_VIDEO_DIR:-}" +GCS_RESULTS_DIR="${GCS_RESULTS_DIR:-}" +UPLOAD_RESULTS="${UPLOAD_RESULTS:-true}" +WORK_DIR_FROM_ARG="" +local_results_parent="" + +VBENCH_REPO="${VBENCH_REPO:-https://github.com/Vchitect/VBench.git}" +VBENCH_BRANCH="${VBENCH_BRANCH:-master}" +MAXDIFFUSION_REPO="${MAXDIFFUSION_REPO:-https://github.com/AI-Hypercomputer/maxdiffusion.git}" +MAXDIFFUSION_BRANCH="${MAXDIFFUSION_BRANCH:-main}" +BENCHMARK_JSON="${BENCHMARK_JSON:-VBench_full_info_sub110.json}" +BENCHMARK_JSON_URL="${BENCHMARK_JSON_URL:-}" +BENCHMARK_JSON_PATH="${BENCHMARK_JSON_PATH:-}" +VBENCH_UTIL_PATH="${VBENCH_UTIL_PATH:-}" + +GPU_NAME="${GPU_NAME:-}" +GPU_ZONE="${GPU_ZONE:-us-central1-a}" +GPU_PROJECT="${GPU_PROJECT:-}" +INTERNAL_IP="${INTERNAL_IP:-false}" +SSH_MODE=false + +DIMENSIONS_TEXT="${DIMENSIONS:-}" +DIMENSIONS=() +if [[ -n "${DIMENSIONS_TEXT}" ]]; then + # shellcheck disable=SC2206 + DIMENSIONS=(${DIMENSIONS_TEXT}) +fi + +usage() { + cat < [KEY=VALUE ...] + +Required: + GCS_BUCKET GCS bucket containing generated videos. + +Common options: + RUN_NAME Generation run name (default: ${RUN_NAME}) + BENCHMARK_JSON Benchmark JSON file name (default: ${BENCHMARK_JSON}) + DIMENSIONS Space-separated VBench dimensions (default: read from JSON) + GCS_VIDEO_DIR GCS video prefix (default: \${RUN_NAME}/videos) + GCS_RESULTS_DIR GCS results prefix (default: \${RUN_NAME}/vbench_results) + UPLOAD_RESULTS Upload results to GCS after evaluation (default: true; SSH mode uploads from the caller machine) + WORK_DIR Working directory on the GPU VM (default: \$HOME/vbench_evaluation) + +Metadata options: + BENCHMARK_JSON_PATH Local benchmark JSON path. + BENCHMARK_JSON_URL URL fallback for benchmark JSON. + MAXDIFFUSION_REPO Repo fallback for benchmark JSON. + MAXDIFFUSION_BRANCH Branch fallback for benchmark JSON. + +SSH options: + --ssh Copy this script to a GPU VM and run it there. + GPU_NAME GPU VM name. + GPU_ZONE GPU VM zone (default: ${GPU_ZONE}) + GPU_PROJECT Optional GCP project for the GPU VM. + INTERNAL_IP Connect to GPU VM using internal IP (default: false). +EOF +} + +die() { + echo "ERROR: $*" >&2 + exit 1 +} + +have() { + command -v "$1" >/dev/null 2>&1 +} + +step() { + echo "==> $*" +} + +parse_args() { + local arg key value + for arg in "$@"; do + case "${arg}" in + --ssh) + SSH_MODE=true + ;; + --help | -h) + usage + exit 0 + ;; + WORK_DIR=*) + WORK_DIR_FROM_ARG="${arg#*=}" + export WORK_DIR="${WORK_DIR_FROM_ARG}" + ;; + DIMENSIONS=*) + value="${arg#*=}" + # shellcheck disable=SC2206 + DIMENSIONS=(${value}) + ;; + *=*) + key="${arg%%=*}" + value="${arg#*=}" + export "${key}=${value}" + ;; + *) + if [[ -z "${GCS_BUCKET}" ]]; then + GCS_BUCKET="${arg}" + fi + ;; + esac + done +} + +absolute_file() { + local path="$1" + [[ -f "${path}" ]] || die "File not found: ${path}" + local dir base + dir="$(cd "$(dirname "${path}")" && pwd)" + base="$(basename "${path}")" + printf '%s/%s\n' "${dir}" "${base}" +} + +normalize_config() { + [[ -n "${GCS_BUCKET}" ]] || die "GCS_BUCKET is required. Example: bash benchmarks/vbench/run_gpu_eval.sh GCS_BUCKET=my-bucket" + + GCS_BUCKET="${GCS_BUCKET#gs://}" + GCS_BUCKET="${GCS_BUCKET%/}" + GCS_VIDEO_DIR="${GCS_VIDEO_DIR:-${RUN_NAME}/videos}" + GCS_RESULTS_DIR="${GCS_RESULTS_DIR:-${RUN_NAME}/vbench_results}" + [[ "${UPLOAD_RESULTS}" == "true" || "${UPLOAD_RESULTS}" == "false" ]] || die "UPLOAD_RESULTS must be true or false." + WORK_DIR="${WORK_DIR_FROM_ARG:-${WORK_DIR:-$HOME/vbench_evaluation}}" + BENCHMARK_JSON_URL="${BENCHMARK_JSON_URL:-https://raw.githubusercontent.com/AI-Hypercomputer/maxdiffusion/${MAXDIFFUSION_BRANCH}/benchmarks/vbench/${BENCHMARK_JSON}}" + + if [[ -n "${BENCHMARK_JSON_PATH}" ]]; then + BENCHMARK_JSON_PATH="$(absolute_file "${BENCHMARK_JSON_PATH}")" + fi + + SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" + UTIL_SCRIPT="${VBENCH_UTIL_PATH:-${SCRIPT_DIR}/gpu_eval_utils.py}" + UTIL_SCRIPT="$(absolute_file "${UTIL_SCRIPT}")" + + JSON_DEST="${WORK_DIR}/${BENCHMARK_JSON}" + DOWNLOAD_DIR="${WORK_DIR}/downloaded_videos" + VBENCH_VIDEO_DIR="${WORK_DIR}/vbench_videos" + RESULTS_DIR="${WORK_DIR}/evaluation_results" +} + +print_config() { + echo "==========================================================================" + echo "Starting VBench evaluation on GPU VM" + echo " GCS Bucket: gs://${GCS_BUCKET}" + echo " Video Source: gs://${GCS_BUCKET}/${GCS_VIDEO_DIR}" + echo " Results Target: gs://${GCS_BUCKET}/${GCS_RESULTS_DIR}" + echo " Working Dir: ${WORK_DIR}" + echo " Benchmark JSON: ${BENCHMARK_JSON}" + echo " Dimensions: ${DIMENSIONS[*]:-(auto-extracted from ${BENCHMARK_JSON})}" + echo "==========================================================================" +} + +shell_quote() { + printf "'%s'" "$(printf '%s' "$1" | sed "s/'/'\\\\''/g")" +} + +quote_command() { + local command="$1" + shift + local arg + for arg in "$@"; do + command="${command} $(shell_quote "${arg}")" + done + printf '%s\n' "${command}" +} + +run_over_ssh() { + [[ -n "${GPU_NAME}" ]] || die "GPU_NAME must be specified when using --ssh mode." + + local source_script remote_script remote_util remote_command remote_json + local gcloud_args=("--zone=${GPU_ZONE}") + source_script="$(absolute_file "${BASH_SOURCE[0]}")" + remote_script="/tmp/run_gpu_eval_${USER:-user}_$$.sh" + remote_util="/tmp/gpu_eval_utils_${USER:-user}_$$.py" + if [[ -n "${GPU_PROJECT}" ]]; then + gcloud_args+=("--project=${GPU_PROJECT}") + fi + if [[ "${INTERNAL_IP}" == "true" ]]; then + gcloud_args+=("--internal-ip") + fi + + echo "==========================================================================" + echo "Executing VBench evaluation remotely" + echo " GPU Name: ${GPU_NAME}" + echo " GPU Zone: ${GPU_ZONE}" + echo " GCS Bucket: gs://${GCS_BUCKET}" + echo " Video Dir: gs://${GCS_BUCKET}/${GCS_VIDEO_DIR}" + echo " Results Dir: gs://${GCS_BUCKET}/${GCS_RESULTS_DIR}" + echo "==========================================================================" + + local scp_script_cmd=("gcloud" "compute" "scp" "${source_script}" "${GPU_NAME}:${remote_script}" "${gcloud_args[@]}") + "${scp_script_cmd[@]}" + + local scp_util_cmd=("gcloud" "compute" "scp" "${UTIL_SCRIPT}" "${GPU_NAME}:${remote_util}" "${gcloud_args[@]}") + "${scp_util_cmd[@]}" + + local remote_args=( + "GCS_BUCKET=${GCS_BUCKET}" + "RUN_NAME=${RUN_NAME}" + "GCS_VIDEO_DIR=${GCS_VIDEO_DIR}" + "GCS_RESULTS_DIR=${GCS_RESULTS_DIR}" + "VBENCH_REPO=${VBENCH_REPO}" + "VBENCH_BRANCH=${VBENCH_BRANCH}" + "MAXDIFFUSION_REPO=${MAXDIFFUSION_REPO}" + "MAXDIFFUSION_BRANCH=${MAXDIFFUSION_BRANCH}" + "BENCHMARK_JSON=${BENCHMARK_JSON}" + "BENCHMARK_JSON_URL=${BENCHMARK_JSON_URL}" + "VBENCH_UTIL_PATH=${remote_util}" + "UPLOAD_RESULTS=false" + ) + + if [[ -n "${WORK_DIR_FROM_ARG}" ]]; then + remote_args+=("WORK_DIR=${WORK_DIR_FROM_ARG}") + fi + if [[ ${#DIMENSIONS[@]} -gt 0 ]]; then + remote_args+=("DIMENSIONS=${DIMENSIONS[*]}") + fi + if [[ -n "${BENCHMARK_JSON_PATH}" ]]; then + remote_json="/tmp/${BENCHMARK_JSON}" + local scp_json_cmd=("gcloud" "compute" "scp" "${BENCHMARK_JSON_PATH}" "${GPU_NAME}:${remote_json}" "${gcloud_args[@]}") + "${scp_json_cmd[@]}" + remote_args+=("BENCHMARK_JSON_PATH=${remote_json}") + fi + + remote_command="$(quote_command "bash $(shell_quote "${remote_script}")" "${remote_args[@]}")" + + local ssh_cmd=("gcloud" "compute" "ssh" "${GPU_NAME}" "${gcloud_args[@]}") + ssh_cmd+=("--command=${remote_command}") + "${ssh_cmd[@]}" + + if [[ "${UPLOAD_RESULTS}" == "true" ]]; then + local remote_results_dir previous_results_dir + if [[ -n "${WORK_DIR_FROM_ARG}" ]]; then + remote_results_dir="${WORK_DIR_FROM_ARG%/}/evaluation_results" + remote_results_dir="${remote_results_dir#~/}" + else + remote_results_dir="vbench_evaluation/evaluation_results" + fi + + local_results_parent="$(mktemp -d "${TMPDIR:-/tmp}/vbench_results.XXXXXX")" + trap 'rm -rf "${local_results_parent}"' EXIT + step "Copying evaluation results from ${GPU_NAME}:${remote_results_dir}..." + local scp_results_cmd=("gcloud" "compute" "scp" "--recurse" "${GPU_NAME}:${remote_results_dir}" "${local_results_parent}/" "${gcloud_args[@]}") + "${scp_results_cmd[@]}" + + previous_results_dir="${RESULTS_DIR}" + RESULTS_DIR="${local_results_parent}/evaluation_results" + upload_results + RESULTS_DIR="${previous_results_dir}" + rm -rf "${local_results_parent}" + trap - EXIT + fi +} + +python_supports_vbench() { + "$1" -c 'import sys; assert (3, 10) <= sys.version_info < (3, 13)' >/dev/null 2>&1 +} + +install_system_packages() { + step "Ensuring build, git, and video codec system packages..." + local packages=(git curl wget unzip zip python3-dev python3-venv build-essential pkg-config ffmpeg libsm6 libxext6 libgl1 libglib2.0-0) + + if have sudo; then + sudo apt-get update -y && sudo apt-get install -y "${packages[@]}" || true + elif have apt-get; then + apt-get update -y && apt-get install -y "${packages[@]}" || true + fi +} + +ensure_uv() { + export PATH="$HOME/.cargo/bin:$HOME/.local/bin:$PATH" + if ! have uv; then + curl -LsSf https://astral.sh/uv/install.sh | sh 2>/dev/null || python3 -m pip install --user uv 2>/dev/null || true + export PATH="$HOME/.cargo/bin:$HOME/.local/bin:$PATH" + fi +} + +create_venv() { + step "Setting up Python virtual environment (Python 3.10-3.12)..." + ensure_uv + + if [[ -d venv ]] && ! python_supports_vbench venv/bin/python3; then + echo "Recreating venv with compatible Python version." + rm -rf venv + fi + + if [[ ! -d venv ]]; then + local candidate + for candidate in python3 python3.12 python3.11 python3.10; do + if have "${candidate}" && python_supports_vbench "${candidate}"; then + "${candidate}" -m venv venv + break + fi + done + fi + + if [[ ! -d venv ]] && have uv; then + uv venv venv --python 3.11 --seed || uv venv venv --python 3.12 --seed || uv venv venv --python 3.10 --seed + fi + + [[ -d venv ]] || python3 -m venv venv + # shellcheck disable=SC1091 + source venv/bin/activate + python3 -m pip install --upgrade "pip<25" "setuptools<71" wheel packaging +} + +sync_vbench_repo() { + step "Syncing VBench repository..." + if [[ ! -d VBench/.git ]]; then + git clone -b "${VBENCH_BRANCH}" "${VBENCH_REPO}" VBench + else + git -C VBench fetch origin "${VBENCH_BRANCH}" || true + git -C VBench checkout "${VBENCH_BRANCH}" || true + git -C VBench pull origin "${VBENCH_BRANCH}" || true + fi +} + +show_gpu_status() { + step "Checking NVIDIA GPU and driver status..." + if have nvidia-smi; then + nvidia-smi || true + else + echo "WARNING: nvidia-smi not found. Ensure NVIDIA drivers are installed on GPU VMs." + fi +} + +install_cuda_torch() { + step "Installing PyTorch with CUDA support..." + if ! python3 -c 'import torch; assert torch.cuda.is_available()' 2>/dev/null; then + python3 -m pip uninstall -y torch torchvision 2>/dev/null || true + python3 -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 --extra-index-url https://pypi.org/simple || \ + python3 -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124 --extra-index-url https://pypi.org/simple || \ + python3 -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 --extra-index-url https://pypi.org/simple + fi + python3 -m pip install "numpy<2" +} + +patch_vbench_sources() { + step "Applying VBench compatibility patches..." + git -C VBench checkout -- setup.py vbench/distributed.py evaluate.py vbench/__init__.py >/dev/null 2>&1 || true + python3 "${UTIL_SCRIPT}" patch-vbench VBench +} + +install_vbench_dependencies() { + step "Installing VBench and evaluation dependencies..." + patch_vbench_sources + (cd VBench && python3 -m pip install --no-build-isolation -e . --extra-index-url https://download.pytorch.org/whl/cu121) + python3 -m pip install --no-build-isolation git+https://github.com/openai/CLIP.git --extra-index-url https://download.pytorch.org/whl/cu121 + python3 -m pip install --no-build-isolation 'git+https://github.com/facebookresearch/detectron2.git' --extra-index-url https://download.pytorch.org/whl/cu121 || true + python3 -m pip install pandas tabulate google-cloud-storage tqdm opencv-python decord + + if ! python3 -c 'import torch; assert torch.cuda.is_available()' 2>/dev/null && have nvidia-smi; then + echo "Reinstalling CUDA-enabled PyTorch after dependency resolution..." + python3 -m pip install --force-reinstall --no-deps torch torchvision --index-url https://download.pytorch.org/whl/cu121 --extra-index-url https://pypi.org/simple + fi + + python3 - <<'PY' +import torch +device_count = torch.cuda.device_count() if torch.cuda.is_available() else 0 +print(f"==> PyTorch {torch.__version__} | CUDA Available: {torch.cuda.is_available()} | Device Count: {device_count}") +PY +} + +gcs_cp() { + if have gcloud && gcloud storage cp "$@"; then + return 0 + fi + if have gsutil && gsutil -m cp "$@"; then + return 0 + fi + return 1 +} + +fetch_benchmark_json() { + step "Locating ${BENCHMARK_JSON}..." + mkdir -p "${WORK_DIR}" + + local repo_json + repo_json="${REPO_ROOT}/benchmarks/vbench/${BENCHMARK_JSON}" + + if [[ -n "${BENCHMARK_JSON_PATH}" ]]; then + cp "${BENCHMARK_JSON_PATH}" "${JSON_DEST}" + elif [[ -f "${repo_json}" ]]; then + cp "${repo_json}" "${JSON_DEST}" + elif [[ -f "${JSON_DEST}" ]]; then + echo "Using existing ${JSON_DEST}" + else + echo "Downloading ${BENCHMARK_JSON} from ${BENCHMARK_JSON_URL}..." + curl -fLsS "${BENCHMARK_JSON_URL}" -o "${JSON_DEST}" 2>/dev/null || \ + python3 -c "import urllib.request; urllib.request.urlretrieve('${BENCHMARK_JSON_URL}', '${JSON_DEST}')" 2>/dev/null || \ + gcs_cp "gs://${GCS_BUCKET}/${RUN_NAME}/${BENCHMARK_JSON}" "${JSON_DEST}" 2>/dev/null || \ + fetch_benchmark_json_from_repo + fi +} + +fetch_benchmark_json_from_repo() { + step "Cloning MaxDiffusion fallback metadata..." + if [[ ! -d maxdiffusion/.git ]]; then + git clone --depth 1 -b "${MAXDIFFUSION_BRANCH}" "${MAXDIFFUSION_REPO}" maxdiffusion + fi + cp "maxdiffusion/benchmarks/vbench/${BENCHMARK_JSON}" "${JSON_DEST}" +} + +extract_dimensions() { + if [[ ${#DIMENSIONS[@]} -eq 0 ]]; then + step "Extracting dimensions from ${JSON_DEST}..." + local extracted + extracted="$(python3 "${UTIL_SCRIPT}" dimensions "${JSON_DEST}")" + # shellcheck disable=SC2206 + DIMENSIONS=(${extracted}) + fi + + [[ ${#DIMENSIONS[@]} -gt 0 ]] || die "No VBench dimensions found. Pass DIMENSIONS=\"...\" or check ${JSON_DEST}." + echo "Dimensions to evaluate: ${DIMENSIONS[*]}" +} + +download_videos() { + step "Downloading generated videos from gs://${GCS_BUCKET}/${GCS_VIDEO_DIR}/..." + mkdir -p "${DOWNLOAD_DIR}" + gcs_cp "gs://${GCS_BUCKET}/${GCS_VIDEO_DIR}/*.mp4" "${DOWNLOAD_DIR}/" || die "Failed to download videos from GCS." +} + +prepare_videos() { + step "Aligning video filenames for VBench..." + mkdir -p "${VBENCH_VIDEO_DIR}" + python3 "${UTIL_SCRIPT}" prepare-videos "${JSON_DEST}" "${DOWNLOAD_DIR}" "${VBENCH_VIDEO_DIR}" +} + +run_vbench() { + step "Running VBench evaluate.py..." + mkdir -p "${RESULTS_DIR}" + ( + cd "${WORK_DIR}/VBench" + local num_gpus=1 + if have nvidia-smi; then + num_gpus=$( (nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || true) | wc -l) + fi + if [[ "${num_gpus}" -gt 1 ]]; then + torchrun --nproc_per_node="${num_gpus}" evaluate.py \ + --videos_path "${VBENCH_VIDEO_DIR}" \ + --full_json_dir "${JSON_DEST}" \ + --output_path "${RESULTS_DIR}" \ + --dimension "${DIMENSIONS[@]}" \ + --mode vbench_standard + else + python3 evaluate.py \ + --videos_path "${VBENCH_VIDEO_DIR}" \ + --full_json_dir "${JSON_DEST}" \ + --output_path "${RESULTS_DIR}" \ + --dimension "${DIMENSIONS[@]}" \ + --mode vbench_standard + fi + ) +} + +upload_results() { + step "Publishing results to gs://${GCS_BUCKET}/${GCS_RESULTS_DIR}/..." + shopt -s nullglob + local result_files=("${RESULTS_DIR}"/*) + shopt -u nullglob + + [[ ${#result_files[@]} -gt 0 ]] || die "No result files found in ${RESULTS_DIR}." + gcs_cp -r "${result_files[@]}" "gs://${GCS_BUCKET}/${GCS_RESULTS_DIR}/" || die "Failed to upload evaluation results to GCS." +} + +run_local() { + print_config + mkdir -p "${WORK_DIR}" + cd "${WORK_DIR}" + + install_system_packages + create_venv + sync_vbench_repo + show_gpu_status + install_cuda_torch + install_vbench_dependencies + fetch_benchmark_json + extract_dimensions + download_videos + prepare_videos + run_vbench + if [[ "${UPLOAD_RESULTS}" == "true" ]]; then + upload_results + else + step "Skipping GCS result upload." + fi + + echo "" + echo "==========================================================================" + echo "VBench evaluation finished successfully!" + if [[ "${UPLOAD_RESULTS}" == "true" ]]; then + echo "Results published to: gs://${GCS_BUCKET}/${GCS_RESULTS_DIR}/" + else + echo "Results upload skipped." + fi + echo "Local results stored in: ${RESULTS_DIR}/" + echo "==========================================================================" +} + +parse_args "$@" +normalize_config + +if [[ "${SSH_MODE}" == "true" ]]; then + run_over_ssh +else + run_local +fi diff --git a/benchmarks/vbench/run_tpu_generation.sh b/benchmarks/vbench/run_tpu_generation.sh new file mode 100644 index 000000000..378d0d58e --- /dev/null +++ b/benchmarks/vbench/run_tpu_generation.sh @@ -0,0 +1,422 @@ +#!/usr/bin/env bash +# +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# ============================================================================== +# VBench TPU Video Generation Script +# +# Usage: +# bash benchmarks/vbench/run_tpu_generation.sh GCS_BUCKET= [KEY=VALUE ...] +# bash benchmarks/vbench/run_tpu_generation.sh --ssh GCS_BUCKET= TPU_NAME= [KEY=VALUE ...] +# ============================================================================== + +set -euo pipefail + +GCS_BUCKET="${GCS_BUCKET:-}" +GCS_VIDEO_DIR="${GCS_VIDEO_DIR:-}" +TPU_NAME="${TPU_NAME:-}" +TPU_ZONE="${TPU_ZONE:-}" +TPU_PROJECT="${TPU_PROJECT:-}" +SSH_MODE=false +EXPLICIT_ARGS=() + +DEFAULT_FLASH_BLOCK_SIZES='{"block_q" : 3328, "block_kv_compute" : 256, "block_kv" : 2816, "block_kv_compute_in" : 256, "block_q_dkv": 3328, "block_kv_dkv" : 2816, "block_kv_dkv_compute" : 256, "block_q_dq" : 3328, "block_kv_dq" : 2816, "heads_per_tile" : 1}' + +WAN_OVERRIDES=( + "run_name|RUN_NAME|wan-inference" + "attention|ATTENTION|ulysses_custom" + "num_inference_steps|NUM_STEPS|40" + "num_frames|NUM_FRAMES|81" + "width|WIDTH|1280" + "height|HEIGHT|720" + "per_device_batch_size|PER_DEVICE_BATCH_SIZE|0.125" + "vae_spatial|VAE_SPATIAL|4" + "vae_decode_chunk|VAE_DECODE_CHUNK|4" + "vae_weights_dtype|VAE_WEIGHTS_DTYPE|bfloat16" + "vae_dtype|VAE_DTYPE|bfloat16" + "text_encoder_dtype|TEXT_ENCODER_DTYPE|bfloat16" + "compile_text_encoder|COMPILE_TEXT_ENCODER|true" + "ici_data_parallelism|ICI_DATA_PARALLELISM|2" + "ici_context_parallelism|ICI_CONTEXT_PARALLELISM|4" + "fps|FPS|16" + "use_kv_cache|USE_KV_CACHE|true" + "use_base2_exp|USE_BASE2_EXP|true" + "use_experimental_scheduler|USE_EXPERIMENTAL_SCHEDULER|true" + "use_batched_text_encoder|USE_BATCHED_TEXT_ENCODER|true" + "flash_block_sizes|FLASH_BLOCK_SIZES|" + "prompt_file|PROMPT_FILE|./benchmarks/vbench/prompts_110.txt" +) + +usage() { + cat < [KEY=VALUE ...] + +Required: + GCS_BUCKET GCS bucket for generated videos. + +Common options: + RUN_NAME Generation run name (default: wan-inference) + GCS_VIDEO_DIR GCS video prefix (default: \${RUN_NAME}/videos) + PROMPT_FILE Prompt file path (default: ./benchmarks/vbench/prompts_110.txt) + CONFIG_FILE WAN config file (default: src/maxdiffusion/configs/base_wan_27b.yml) + EXTERNAL_DISK Mounted disk root for large local files (default: /mnt/disks/external_disk) + HF_CACHE_ROOT Hugging Face cache root (default: \$EXTERNAL_DISK/hf_cache) + HF_HOME Hugging Face home directory (default: \$HF_CACHE_ROOT) + HF_HUB_CACHE Hugging Face Hub model cache (default: \$HF_HOME/hub) + HF_XET_CACHE Hugging Face Xet cache (default: \$HF_HOME/xet) + TMPDIR Temporary files directory (default: \$EXTERNAL_DISK/tmp) + +SSH options: + --ssh Copy this script to a TPU VM and run it there. + TPU_NAME TPU VM name (required when using --ssh). + TPU_ZONE TPU VM zone. + TPU_PROJECT TPU VM project. + REMOTE_DIR MaxDiffusion repo path on the TPU VM (default: \$HOME/maxdiffusion) + GIT_REPO Repo to clone in SSH mode (default: https://github.com/AI-Hypercomputer/maxdiffusion.git). + GIT_BRANCH Branch to sync in SSH mode (default: main) +EOF +} + +die() { + echo "ERROR: $*" >&2 + exit 1 +} + +have() { + command -v "$1" >/dev/null 2>&1 +} + +step() { + echo "==> $*" +} + +quote() { + printf "'%s'" "$(printf '%s' "$1" | sed "s/'/'\\\\''/g")" +} + +set_default() { + local var="$1" value="$2" + [[ -n "${!var-}" ]] || printf -v "${var}" '%s' "${value}" +} + +arg_was_explicit() { + local explicit_arg + for explicit_arg in "${EXPLICIT_ARGS[@]}"; do + [[ "${explicit_arg}" == "$1" ]] && return 0 + done + return 1 +} + +path_is_under_external_disk() { + local path="${1%/}" + [[ "${path}" == "${EXTERNAL_DISK}" || "${path}" == "${EXTERNAL_DISK}/"* ]] +} + +ignore_inherited_path_outside_external_disk() { + local var="$1" + if [[ -n "${!var-}" ]] && ! arg_was_explicit "${var}" && ! path_is_under_external_disk "${!var}"; then + unset "${var}" + fi +} + +parse_args() { + local arg key value + for arg in "$@"; do + case "${arg}" in + --ssh) + SSH_MODE=true + ;; + --help | -h) + usage + exit 0 + ;; + *=*) + key="${arg%%=*}" + value="${arg#*=}" + EXPLICIT_ARGS+=("${key}") + export "${key}=${value}" + ;; + *) + [[ -n "${GCS_BUCKET}" ]] || GCS_BUCKET="${arg}" + ;; + esac + done +} + +normalize_config() { + local key var value item + [[ -n "${GCS_BUCKET}" ]] || die "GCS_BUCKET is required. Example: bash benchmarks/vbench/run_tpu_generation.sh GCS_BUCKET=my-bucket" + + GCS_BUCKET="${GCS_BUCKET#gs://}" + GCS_BUCKET="${GCS_BUCKET%/}" + set_default CONFIG_FILE "src/maxdiffusion/configs/base_wan_27b.yml" + set_default FLASH_BLOCK_SIZES "${DEFAULT_FLASH_BLOCK_SIZES}" + set_default TPU_NAME "" + set_default TPU_ZONE "" + set_default TPU_PROJECT "" + set_default GIT_REPO "https://github.com/AI-Hypercomputer/maxdiffusion.git" + set_default GIT_BRANCH "main" + + if [[ "${SSH_MODE}" == "true" ]]; then + [[ -n "${TPU_NAME:-}" ]] || die "TPU_NAME is required when using --ssh." + fi + set_default EXTERNAL_DISK "/mnt/disks/external_disk" + EXTERNAL_DISK="${EXTERNAL_DISK%/}" + for var in HF_CACHE_ROOT HF_HOME HF_HUB_CACHE HF_XET_CACHE HF_ASSETS_CACHE HF_DATASETS_CACHE HF_MODULES_CACHE TRANSFORMERS_CACHE TMPDIR; do + ignore_inherited_path_outside_external_disk "${var}" + done + set_default HF_CACHE_ROOT "${EXTERNAL_DISK}/hf_cache" + set_default HF_HOME "${HF_CACHE_ROOT}" + set_default HF_HUB_CACHE "${HF_HOME}/hub" + set_default HF_XET_CACHE "${HF_HOME}/xet" + set_default HF_ASSETS_CACHE "${HF_HOME}/assets" + set_default HF_DATASETS_CACHE "${HF_HOME}/datasets" + set_default HF_MODULES_CACHE "${HF_HOME}/modules" + set_default TRANSFORMERS_CACHE "${HF_HUB_CACHE}" + set_default TMPDIR "${EXTERNAL_DISK}/tmp" + set_default VENV_DIR "${HOME}/maxdiffusion_venv" + + for item in "${WAN_OVERRIDES[@]}"; do + IFS='|' read -r key var value <<< "${item}" + set_default "${var}" "${value}" + done + set_default GCS_VIDEO_DIR "${RUN_NAME}/videos" + GCS_VIDEO_DIR="${GCS_VIDEO_DIR#/}" + GCS_VIDEO_DIR="${GCS_VIDEO_DIR%/}" + [[ "${GCS_VIDEO_DIR}" == "${RUN_NAME}/videos" ]] || die "GCS_VIDEO_DIR must be ${RUN_NAME}/videos for TPU generation; pass RUN_NAME to choose the output directory." + + SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + if [[ -n "${MAXDIFFUSION_ROOT:-}" ]]; then + REPO_ROOT="$(cd "${MAXDIFFUSION_ROOT}" && pwd)" + else + REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" + fi +} + +absolute_file() { + local path="$1" dir base + [[ -f "${path}" ]] || die "File not found: ${path}" + dir="$(cd "$(dirname "${path}")" && pwd)" + base="$(basename "${path}")" + printf '%s/%s\n' "${dir}" "${base}" +} + +print_config() { + echo "==========================================================================" + echo "Starting VBench video generation on TPU VM" + echo " Run Name: ${RUN_NAME}" + echo " GCS Bucket: gs://${GCS_BUCKET}" + echo " Video Target: gs://${GCS_BUCKET}/${GCS_VIDEO_DIR}" + echo " Prompt File: ${PROMPT_FILE}" + echo " Config File: ${CONFIG_FILE}" + echo " Seed: 12345" + echo " Repo Root: ${REPO_ROOT}" + echo " HF Cache: ${HF_CACHE_ROOT}" + echo " HF Hub: ${HF_HUB_CACHE}" + echo " HF Xet: ${HF_XET_CACHE}" + echo " Temp Dir: ${TMPDIR}" + echo "==========================================================================" +} + +emit_remote_arg() { + local var="$1" + printf ' %s\n' "$(quote "${var}=${!var-}")" +} + +emit_remote_arg_if_explicit() { + local var="$1" + if arg_was_explicit "${var}"; then + emit_remote_arg "${var}" + fi +} + +emit_remote_args() { + local item key var value + for var in GCS_BUCKET GCS_VIDEO_DIR CONFIG_FILE EXTERNAL_DISK HF_CACHE_ROOT HF_HOME HF_HUB_CACHE HF_XET_CACHE HF_ASSETS_CACHE HF_DATASETS_CACHE HF_MODULES_CACHE TRANSFORMERS_CACHE TMPDIR; do + emit_remote_arg "${var}" + done + emit_remote_arg_if_explicit VENV_DIR + for item in "${WAN_OVERRIDES[@]}"; do + IFS='|' read -r key var value <<< "${item}" + emit_remote_arg "${var}" + done +} + +run_over_ssh() { + local source_script remote_script remote_dir_label remote_dir_assignment remote_command + local gcloud_args=() + [[ -n "${TPU_ZONE}" ]] && gcloud_args+=("--zone=${TPU_ZONE}") + [[ -n "${TPU_PROJECT}" ]] && gcloud_args+=("--project=${TPU_PROJECT}") + + source_script="$(absolute_file "${BASH_SOURCE[0]}")" + remote_script="/tmp/run_tpu_generation_${USER:-user}_$$.sh" + remote_dir_label="${REMOTE_DIR:-\$HOME/maxdiffusion}" + if [[ -n "${REMOTE_DIR:-}" ]]; then + remote_dir_assignment="remote_dir=$(quote "${REMOTE_DIR}")" + else + remote_dir_assignment='remote_dir="${HOME}/maxdiffusion"' + fi + + echo "==========================================================================" + echo "Executing VBench generation remotely" + echo " TPU Name: ${TPU_NAME}" + echo " TPU Zone: ${TPU_ZONE:-}" + echo " TPU Project: ${TPU_PROJECT:-}" + echo " Remote Dir: ${remote_dir_label}" + echo " GCS Bucket: gs://${GCS_BUCKET}" + echo " Video Target: gs://${GCS_BUCKET}/${GCS_VIDEO_DIR}" + echo "==========================================================================" + + gcloud compute tpus tpu-vm scp "${source_script}" "${TPU_NAME}:${remote_script}" "${gcloud_args[@]}" + + remote_command=$(cat < [TPU VM] Syncing MaxDiffusion at \${remote_dir}..." +if [[ ! -d "\${remote_dir}/.git" ]]; then + mkdir -p "\$(dirname "\${remote_dir}")" + git clone "\${git_repo}" "\${remote_dir}" +fi +cd "\${remote_dir}" +git fetch origin "\${git_branch}" || true +git checkout "\${git_branch}" || true +git pull origin "\${git_branch}" || true + +args=( +$(emit_remote_args) + "MAXDIFFUSION_ROOT=\${remote_dir}" +) +bash $(quote "${remote_script}") "\${args[@]}" +EOF +) + + gcloud compute tpus tpu-vm ssh "${TPU_NAME}" "${gcloud_args[@]}" --command="${remote_command}" +} + +configure_cache_dirs() { + local cache_var external_device root_device + + [[ -d "${EXTERNAL_DISK}" ]] || die "External disk ${EXTERNAL_DISK} is not mounted. Mount it or pass EXTERNAL_DISK=." + [[ -w "${EXTERNAL_DISK}" ]] || die "External disk ${EXTERNAL_DISK} is not writable." + + root_device="$(df -P / | awk 'NR == 2 {print $1}')" + external_device="$(df -P "${EXTERNAL_DISK}" | awk 'NR == 2 {print $1}')" + if [[ "${external_device}" == "${root_device}" ]]; then + local free_kb free_gb + free_kb="$(df -P / | awk 'NR == 2 {print $4}')" + free_gb=$(( free_kb / 1024 / 1024 )) + if (( free_gb < 100 )); then + die "${EXTERNAL_DISK} is on root filesystem with only ${free_gb}GB free. Need at least 100GB." + else + echo "WARNING: HF cache is on root filesystem (${free_gb}GB free). Consider mounting a dedicated disk for large runs." + fi + fi + + for cache_var in HF_CACHE_ROOT HF_HOME HF_HUB_CACHE HF_XET_CACHE HF_ASSETS_CACHE HF_DATASETS_CACHE HF_MODULES_CACHE TRANSFORMERS_CACHE TMPDIR; do + case "${!cache_var%/}" in + "${EXTERNAL_DISK}" | "${EXTERNAL_DISK}/"*) ;; + *) die "${cache_var}=${!cache_var} must be under EXTERNAL_DISK=${EXTERNAL_DISK}." ;; + esac + done + + export HF_CACHE_ROOT HF_HOME HF_HUB_CACHE HF_XET_CACHE HF_ASSETS_CACHE HF_DATASETS_CACHE HF_MODULES_CACHE TRANSFORMERS_CACHE TMPDIR + mkdir -p "${HF_HOME}" "${HF_HUB_CACHE}" "${HF_XET_CACHE}" "${HF_ASSETS_CACHE}" "${HF_DATASETS_CACHE}" "${HF_MODULES_CACHE}" "${TRANSFORMERS_CACHE}" "${TMPDIR}" +} + +python_supports_wan() { + "$1" -c 'import sys; assert sys.version_info >= (3, 12)' >/dev/null 2>&1 +} + +create_python_env() { + export PATH="${HOME}/.cargo/bin:${HOME}/.local/bin:${PATH}" + if [[ -x "${VENV_DIR}/bin/python3" ]] && python_supports_wan "${VENV_DIR}/bin/python3"; then + # shellcheck disable=SC1091 + source "${VENV_DIR}/bin/activate" + return + fi + # Always create and use the virtual environment to avoid PEP 668 issues and system package conflicts. + + step "Creating Python 3.12 virtualenv..." + have uv || python3 -m pip install --user --upgrade uv 2>/dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh + [[ ! -d "${VENV_DIR}" ]] || rm -rf "${VENV_DIR}" + python3 -m uv venv "${VENV_DIR}" --python 3.12 --seed || uv venv "${VENV_DIR}" --python 3.12 --seed + # shellcheck disable=SC1091 + source "${VENV_DIR}/bin/activate" +} + +install_dependencies() { + step "Installing MaxDiffusion TPU dependencies..." + bash setup.sh MODE=stable DEVICE=tpu + python3 -m uv pip install -e . || uv pip install -e . || python3 -m pip install -e . +} + +run_generation() { + local item key var value + local -a args + step "Running WAN 2.2 27B inference with seed: 12345..." + args=(python3 src/maxdiffusion/generate_wan.py "${CONFIG_FILE}") + for item in "${WAN_OVERRIDES[@]}"; do + IFS='|' read -r key var value <<< "${item}" + args+=("${key}=${!var}") + done + args+=("seed=12345" "base_output_directory=gs://${GCS_BUCKET}") + "${args[@]}" +} + +sync_metadata() { + step "Syncing benchmark metadata to GCS..." + python3 - < 0 and prompt_word_count < prompt_enhancement_words_threshold - ) + enhance_prompt = prompt_enhancement_words_threshold > 0 and prompt_word_count < prompt_enhancement_words_threshold s0 = time.perf_counter() images = pipeline(