|
| 1 | +#!/bin/bash |
| 2 | +# Copyright 2026 Google LLC |
| 3 | +# |
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +# you may not use this file except in compliance with the License. |
| 6 | +# You may obtain a copy of the License at |
| 7 | +# |
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +# |
| 10 | +# Unless required by applicable law or agreed to in writing, software |
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +# See the License for the specific language governing permissions and |
| 14 | +# limitations under the License. |
| 15 | + |
| 16 | +# WAN T2V fast-serving example: AOT executable cache + converted-weights |
| 17 | +# cache + zero-exec warmup, with a tuned v7 2D-ring attention recipe. |
| 18 | +# |
| 19 | +# First run per (model, shape) pays one-time conversion + compile and |
| 20 | +# populates the caches; every later process start is ~25s to ready. |
| 21 | +# |
| 22 | +# Usage: |
| 23 | +# ./run_wan_fast_inference.sh [21|22] [steps] ["prompt..."] |
| 24 | +# Env overrides: |
| 25 | +# WAN_CACHE_ROOT cache root (default ~/.cache/maxdiffusion_wan) |
| 26 | +# OUTPUT_DIR video/metrics output (default /tmp/wan_out) |
| 27 | +# COMPILE_TE=true torch.compile the text encoder (adds ~30s to load, |
| 28 | +# saves ~10s/encode; worth it for long-lived processes) |
| 29 | +set -u |
| 30 | +MODEL=${1:-22} |
| 31 | +STEPS=${2:-40} |
| 32 | +PROMPT=${3:-""} |
| 33 | + |
| 34 | +PROJECT_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &> /dev/null && pwd)" |
| 35 | +cd "$PROJECT_ROOT" || exit 1 |
| 36 | +export PYTHONPATH="$PROJECT_ROOT/src:${PYTHONPATH:-}" |
| 37 | +export HF_HUB_ENABLE_HF_TRANSFER=1 |
| 38 | +export JAX_DEFAULT_MATMUL_PRECISION=bfloat16 |
| 39 | +export TORCHINDUCTOR_FX_GRAPH_CACHE=1 |
| 40 | + |
| 41 | +CACHE_ROOT=${WAN_CACHE_ROOT:-$HOME/.cache/maxdiffusion_wan} |
| 42 | +OUTPUT_DIR=${OUTPUT_DIR:-/tmp/wan_out} |
| 43 | +mkdir -p "$CACHE_ROOT/jax" "$CACHE_ROOT/aot_wan$MODEL" "$CACHE_ROOT/converted" "$OUTPUT_DIR" |
| 44 | + |
| 45 | +# Tuned collective/scheduler flag set for v7 (from the PR #430 2D-ring |
| 46 | +# baseline). One line: libtpu stops parsing at a literal backslash. |
| 47 | +export LIBTPU_INIT_ARGS="--xla_tpu_spmd_rng_bit_generator_unsafe=true --xla_tpu_enable_dot_strength_reduction=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_enable_async_collective_permute=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true --xla_tpu_scoped_vmem_limit_kib=65536 --xla_tpu_enable_async_all_to_all=true --xla_tpu_enable_all_experimental_scheduler_features=true --xla_tpu_enable_scheduler_memory_pressure_tracking=true --xla_tpu_host_transfer_overlap_limit=24 --xla_tpu_aggressive_opt_barrier_removal=ENABLED --xla_lhs_prioritize_async_depth_over_stall=ENABLED --xla_should_allow_loop_variant_parameter_in_chain=ENABLED --xla_should_add_loop_invariant_op_in_chain=ENABLED --xla_tpu_enable_ici_ag_pipelining=true --xla_max_concurrent_host_send_recv=100 --xla_tpu_scheduler_percent_shared_memory_limit=100 --xla_latency_hiding_scheduler_rerun=2 --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 --xla_tpu_enable_latency_hiding_scheduler=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_megacore_fusion=true --xla_tpu_megacore_fusion_allow_ags=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true --xla_tpu_sparse_core_all_gather_latency_multiplier=1 --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 --xla_tpu_enable_sparse_core_collective_aggregator=true --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true --xla_tpu_enable_sparse_core_reduce_scatter_v2=true --xla_tpu_enable_sparse_core_collective_offload_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true --xla_tpu_enable_concurrent_sparse_core_offloading=true --xla_tpu_assign_all_reduce_scatter_layout=true" |
| 48 | + |
| 49 | +if [ "$MODEL" = "21" ]; then |
| 50 | + CONFIG=src/maxdiffusion/configs/base_wan_14b.yml |
| 51 | + GUIDANCE_ARGS="" |
| 52 | +else |
| 53 | + CONFIG=src/maxdiffusion/configs/base_wan_27b.yml |
| 54 | + GUIDANCE_ARGS="guidance_scale_low=3.0 guidance_scale_high=4.0" |
| 55 | +fi |
| 56 | + |
| 57 | +PROMPT_ARG=() |
| 58 | +[ -n "$PROMPT" ] && PROMPT_ARG=("prompt=$PROMPT") |
| 59 | +RUN_NAME="wan${MODEL}_fast_$(date +%m%d-%H%M%S)" |
| 60 | + |
| 61 | +# libtpu's XLA:CPU AOT feature-mismatch log is cosmetic and ignores every |
| 62 | +# log-level env var; filter just that message from stderr. |
| 63 | +python src/maxdiffusion/generate_wan.py "$CONFIG" \ |
| 64 | + run_name="$RUN_NAME" \ |
| 65 | + output_dir="$OUTPUT_DIR" \ |
| 66 | + jax_cache_dir="$CACHE_ROOT/jax" \ |
| 67 | + aot_cache_dir="$CACHE_ROOT/aot_wan$MODEL" \ |
| 68 | + converted_weights_dir="$CACHE_ROOT/converted" \ |
| 69 | + attention=ulysses_ring_custom \ |
| 70 | + ulysses_shards=2 \ |
| 71 | + ici_data_parallelism=2 ici_fsdp_parallelism=1 \ |
| 72 | + ici_context_parallelism=4 ici_tensor_parallelism=1 \ |
| 73 | + per_device_batch_size=0.125 \ |
| 74 | + num_inference_steps="$STEPS" num_frames=81 width=1280 height=720 \ |
| 75 | + weights_dtype=bfloat16 activations_dtype=bfloat16 \ |
| 76 | + vae_spatial=4 vae_decode_chunk=-1 \ |
| 77 | + vae_weights_dtype=bfloat16 vae_dtype=bfloat16 \ |
| 78 | + text_encoder_dtype=bfloat16 compile_text_encoder="${COMPILE_TE:-false}" use_batched_text_encoder=false \ |
| 79 | + use_base2_exp=true use_experimental_scheduler=true \ |
| 80 | + fps=16 $GUIDANCE_ARGS \ |
| 81 | + flash_block_sizes='{"block_q":9472,"block_kv":1024,"block_kv_compute":1024,"block_kv_compute_in":1024,"heads_per_tile":1,"vmem_limit_bytes":67108864,"block_q_dkv":9472,"block_kv_dkv":1024,"block_kv_dkv_compute":1024}' \ |
| 82 | + "${PROMPT_ARG[@]}" \ |
| 83 | + 2> >(grep -vE --line-buffered 'cpu_aot_loader|machine type for execution' >&2) |
| 84 | + |
| 85 | +mp4=$(ls -t wan_output_*.mp4 2>/dev/null | head -1) |
| 86 | +if [ -n "$mp4" ]; then |
| 87 | + mv "$mp4" "$OUTPUT_DIR/${RUN_NAME}.mp4" |
| 88 | + echo "" |
| 89 | + echo "=== video saved: $OUTPUT_DIR/${RUN_NAME}.mp4 ===" |
| 90 | +fi |
0 commit comments