#!/usr/bin/env bash
set -euo pipefail

mode="${1:?usage: run_distributed_t2va.sh ulysses|tensor [WORLD_SIZE] [ATTENTION]}"
world_size="${2:-${H3_WORLD_SIZE:-}}"
if [[ -z "$world_size" ]]; then
  world_size="$(python -c 'import torch; print(torch.cuda.device_count())')"
fi
if [[ "$world_size" -lt 1 ]]; then
  echo "No CUDA devices are visible. Launch the container with --gpus all." >&2
  exit 1
fi
attention="${3:-sdpa}"
benchmark="${H3_DISTRIBUTED_BENCHMARK:-benchmarks/t2va-dialogue-quoted-864x480-141f-base12-sage2-seed440420.json}"
output_root="${H3_DISTRIBUTED_OUTPUT:-/output/h3-baselines}"
stem="distributed-${mode}-${world_size}gpu-${attention}"
model="${H3_MODEL_PATH:-/models/minimax_h3_fl2va_pruned_nvfp4.safetensors}"
text_encoder="${H3_TEXT_ENCODER_PATH:-/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors}"
extra_args=()
if [[ "${H3_SAVE_LATENTS:-1}" == "1" ]]; then
  extra_args+=(--save-latent "$output_root/$stem.latent.pt")
fi

torchrun \
  --standalone \
  --nnodes=1 \
  --nproc-per-node="$world_size" \
  tools/distributed_t2va.py \
  --benchmark "$benchmark" \
  --mode "$mode" \
  --attention "$attention" \
  --model "$model" \
  --text-encoder "$text_encoder" \
  --report "$output_root/$stem.json" \
  "${extra_args[@]}"