# Streaming multi-view video-to-video generation for weather/style transfer

`custom/eval/vid2vid_eval/infer_multiview_video_transfer.py`: turn one or more
per-camera driving videos into a re-rendered version (e.g. snow) using an **image-to-image keyframe model** + 
the **Wan Fun-InP+canny interpolation model**.

The script is **streaming/causal**: it consumes the videos chunk-by-chunk (8 frames per chunk),
generating each chunk's last keyframe first (per-frame model) then interpolating the frames in
between. This is the same shape a live CARLA source would drive later. See the module docstring
for the full design.

- **First** `--videos` entry is always `front_wide`. A **single** video ⇒ single-view transfer
  (the cross-view code becomes a no-op).
- `--num-frames`: `1` ⇒ save a single **image**; `-1` ⇒ whole video; otherwise N frames
  (padded up to an 8-frame chunk boundary, then cropped back to N).

---

## 0. Prerequisites (host)

- NVIDIA GPU + driver + `nvidia-container-toolkit` (checked: RTX PRO 6000 Blackwell).
- Docker.
- AWS CLI with the `tri-e2e-dev` profile configured (`aws configure list-profiles` shows it).
- ~40 GB free disk for checkpoints + text encoder.

---

## 1. Docker image

Build the docker container following [the root README](https://github.com/TRI-ML/Any4D/blob/main/README.md)

## 2. Download assets from S3 / HF

Everything lands under `~/vid2vid_assets/`.

```bash
A=~/vid2vid_assets
mkdir -p $A/{ckpts,videos,wan/google/umt5-xxl,output}
PROF=tri-e2e-dev

# --- keyframe checkpoint for --kf-model anchor (anchor Fun-Control, single- AND multi-view) ---
aws --profile $PROF s3 cp \
  s3://gaia-e2e-wfm-datasets/any4d/a4d2/edge1frame/06-15-05-23_aw_canny_anchor/model/iter_000066000_001056000.pt \
  $A/ckpts/kf.pt

# --- keyframe checkpoint for --kf-model syncfz (pure Wan-T2V syncfz-K4); only needed for syncfz ---
aws --profile $PROF s3 cp \
  s3://gaia-e2e-wfm-datasets/any4d/a4d2/t2v_base/syncfzfix_500k/07-01-20-26_aw_t2i_mixphase_syncfz_k4/iter_290000/model/iter_000290000_009280000.pt \
  $A/ckpts/kf_syncfz.pt

# --- interpolation checkpoint (Wan Fun-InP+canny, latest EMA) ---
aws --profile $PROF s3 cp \
  s3://gaia-e2e-wfm-datasets/any4d/a4d2/inp/06-11-17-16_canny-iked1-41-ema/model/iter_000163000_001304000.pt \
  $A/ckpts/inp.pt

# --- VAE (from the S3 Wan mirror) ---
aws --profile $PROF s3 cp \
  s3://gaia-e2e-wfm-datasets/any4d/pretrained/wan/Wan2.1-Fun-1.3B-InP/Wan2.1_VAE.pth \
  $A/wan/Wan2.1_VAE.pth
```

Download the UMT5 tokenizer and encoder weights from HF repo `alibaba-pai/Wan2.1-Fun-1.3B-InP`:

```bash
HF=https://huggingface.co/alibaba-pai/Wan2.1-Fun-1.3B-InP/resolve/main
for f in tokenizer.json tokenizer_config.json special_tokens_map.json spiece.model; do
  curl -fsSL "$HF/google/umt5-xxl/$f" -o $A/wan/google/umt5-xxl/$f
done
curl -fsSL "$HF/models_t5_umt5-xxl-enc-bf16.pth" -o $A/wan/models_t5_umt5-xxl-enc-bf16.pth   # ~11 GB
```

Final `--wan-dir` layout (`$A/wan`):
```
Wan2.1_VAE.pth
models_t5_umt5-xxl-enc-bf16.pth
google/umt5-xxl/{tokenizer.json,tokenizer_config.json,special_tokens_map.json,spiece.model}
```

### Test videos

The GAIA validation set (`split_val_transfer.json`) stores per-camera mp4s under each sequence's
`rgb/` folder. Download the two cameras you want (front_wide first):

```bash
BASE=s3://gaia-e2e-wfm-datasets/anydata/cv_unified_adverse_weather/videos/parquet_all_weather
SEQ=chunk_00.parquet/ip602_1450029186244495000_1450029216244495000     # first val sequence, 299 frames @ 10 fps
for cam in front_wide front_tele; do
  aws --profile $PROF s3 cp "$BASE/$SEQ/rgb/$cam.mp4" $A/videos/$cam.mp4
done
```

(List the full split with
`aws --profile $PROF s3 cp $BASE/split_val_transfer.json - | python -m json.tool | less`.)

---

## 3. Run inference

10 seconds @ 10 fps = **100 frames**, two cameras, snow prompt (the script's default prompt is
already snow):

```bash
docker run --gpus all --shm-size=64g --ipc=host --rm \
  -v ~/Any4D:/workspace \
  -v ~/vid2vid_assets:/assets \
  -w /workspace \
  -e A4D_WAN_FORCE_SDPA=1 \
  cosmos-predict2-local:latest \
  bash -lc "PYTHONPATH=/workspace A4D_WAN_FORCE_SDPA=1 python custom/eval/vid2vid_eval/infer_multiview_video_transfer.py \
    --kf-ckpt  /assets/ckpts/kf.pt \
    --inp-ckpt /assets/ckpts/inp.pt \
    --wan-dir  /assets/wan \
    --vae-path /assets/wan/Wan2.1_VAE.pth \
    --videos   /assets/videos/front_wide.mp4 /assets/videos/front_tele.mp4 \
    --num-frames 100 \
    --out      /assets/output/snow_2cam"
```

Or just use the helper: `bash ~/vid2vid_assets/run_infer.sh 100 snow_2cam`.

> ⚠️ **Workaround — two env settings are required:**
> 1. **`A4D_WAN_FORCE_SDPA=1`** — the docker's FlashAttention-3 is compiled for Hopper (sm_90) and
>    has *no kernel image* for Blackwell (sm_120); without this you get
>    `CUDA error ... no kernel image is available for execution on the device`. This flag routes
>    attention through torch SDPA instead. (Skip it on an H100/H200.)
> 2. **`PYTHONPATH=/workspace`** — the script sets its module root to its own dir, so the repo root
>    must be on `PYTHONPATH` for `import custom.* / cosmos_predict2.* / imaginaire.*` to resolve.

### Variants
- **Keyframe model:** `--kf-model anchor` (default; 2-view cross-view Fun-Control, front_wide anchors
  the other cameras, `--kf-ckpt /assets/ckpts/kf.pt`) or `--kf-model syncfz` (pure Wan-T2V,
  per-camera, no cross-view anchor, `--kf-ckpt /assets/ckpts/kf_syncfz.pt`). Point `--kf-ckpt` at the
  matching checkpoint.
- **Single view:** pass one video → `--videos /assets/videos/front_wide.mp4`.
- **Single image:** `--num-frames 1` → writes `pred_00_front_wide.png` + `grid.png`.
- **Whole video:** `--num-frames -1`.
- **Different weather:** `--prompt "..." --neg-prompt "..."`.
- **Tone-match:** on by default (locks each keyframe's exposure to keyframe 0, removing
  inter-keyframe flicker). Add `--no-tone-match` to disable — do this for long clips with genuine
  lighting changes, where locking to the first frame is undesirable.

---

## 3b. Batch over a whole validation split

`batch_val_transfer.py` loads the pipes **once** and loops `run_sequence` over every sequence in a
split (resumable — a sequence whose `grid.mp4` exists is skipped; per-sequence failures are logged
and skipped). Pre-download all sequences' cameras to `<data-root>/<seq>/<camera>.mp4` first:

```bash
# host-side: front_wide + front_tele for every val sequence
A=~/vid2vid_assets ; PROF=tri-e2e-dev
BASE=s3://gaia-e2e-wfm-datasets/anydata/cv_unified_adverse_weather/videos/parquet_all_weather
aws --profile $PROF s3 cp $BASE/split_val_transfer.json - \
  | python3 -c 'import json,sys;[print(k) for k in json.load(sys.stdin)["sequences"]]' > $A/val_seqs.txt
while read seq; do
  mkdir -p "$A/val_videos/$seq"
  for cam in front_wide front_tele; do
    aws --profile $PROF s3 cp "$BASE/$seq/rgb/$cam.mp4" "$A/val_videos/$seq/$cam.mp4" --only-show-errors
  done
done < $A/val_seqs.txt
```

Then run the batch in the docker (note **`python -u`** — unbuffered, so progress is visible in the
log during the long run):

```bash
docker run --gpus all --shm-size=64g --ipc=host --rm \
  -v ~/Any4D:/workspace -v ~/vid2vid_assets:/assets -w /workspace \
  -e A4D_WAN_FORCE_SDPA=1 \
  cosmos-predict2-local:latest \
  bash -lc "PYTHONPATH=/workspace A4D_WAN_FORCE_SDPA=1 python -u \
    custom/eval/vid2vid_eval/batch_val_transfer.py \
    --kf-ckpt /assets/ckpts/kf.pt --inp-ckpt /assets/ckpts/inp.pt \
    --wan-dir /assets/wan --vae-path /assets/wan/Wan2.1_VAE.pth \
    --split-json /assets/split_val_transfer.json --data-root /assets/val_videos \
    --cameras front_wide front_tele --num-frames 100 \
    --out /assets/output/val_transfer"
```

Outputs land in `~/vid2vid_assets/output/val_transfer/<seq>/` (same file set as §4). ~96 sequences
at ~20 min each on one Blackwell GPU with SDPA. Short clips (<100 real frames) are padded to 100 by
duplicating the last frame (the `--num-frames` contract). Re-running resumes where it left off.

## 4. Outputs

Under `--out` (`~/vid2vid_assets/output/snow_2cam/`):
- `pred_00_front_wide.mp4`, `pred_01_front_tele.mp4` — per-camera predictions.
- `grid.mp4` — front_wide input (top) + output (bottom) large on the left; every other camera
  tiled as thumbnails to the right (5 per column, zero-padded to match height).
- `test_input/` — the exact preprocessed input frames fed to the model.

For `--num-frames 1` the same names are written as `.png` + `grid.png`.