Add clip fix

Training:
- finetune.py: optional track/base param groups (WANTRACK_TRACK_GROUP) so the
  track pathway can take a different LR (WANTRACK_TRACK_LR_MULT); patch_embedding
  gradient masked to the track slot [:, 36:]. Raises if nothing matches.
- track_warmup.py callback: hold the pretrained DiT at lr 0 while the track
  pathway warms up, then ramp it back in.

14B configs + held-alloc launchers:
- finetune_wantrack_{synth_sparse_fixed,synth_sparse_random,openvid_sparse_merged_bias,
  synth_stage2}_14b + openvid_14b_{warmup,lwlr} + 1.3B stage2 configs.
- run_wan14b_held.sh (held alloc + auto-restart + rank cleanup, per-stage WANTRACK
  env), acquire_nodes.sh, chain_stepAB.sh, run_stepB_seed.sh, run_stepC_merge.sh,
  watch_gate.sh, run_synth_stage2_slurm.sh.

Init / diagnostics:
- convert_trackwan_init_v2.py: --pe-src to lift patch_embedding[:, 36:] and
  --track-src to lift track_encoder together (co-adapted merge); sharded-source support.
- check_gate_growth.py, gradient_flow_analysis.py, gradient_flow_summarize.py.

Data pipeline: segment_tracks.py sharding; decode/segment/tracks slurm scripts;
precompute_null_text, two_track_val, cfg_ablation, lwlr_experiment.

Inference: app_multi_trace.py updates.

.agents/memory/performance-dashboard.md

@@ -0,0 +1,331 @@
+# Performance Dashboard Memory
+
+Date: 2026-06-16
+Branch: `ci/dashboard`
+
+## Purpose
+
+This branch adds a local live dashboard for FastVideo performance benchmark
+history. It is intended for maintainer/operator use: inspect latest benchmark
+status, compare current values with recent baseline context, and view trends
+from the Hugging Face performance-tracking dataset.
+
+The dashboard is a FastAPI + React app. It is separate from the existing
+Svelte `ui/` app.
+
+## Main Files Added Or Changed
+
+Backend:
+
+- `fastvideo/performance_dashboard/__init__.py`
+- `fastvideo/performance_dashboard/__main__.py`
+- `fastvideo/performance_dashboard/api.py`
+- `fastvideo/performance_dashboard/metrics.py`
+- `fastvideo/performance_dashboard/service.py`
+
+Frontend:
+
+- `performance_dashboard/frontend/package.json`
+- `performance_dashboard/frontend/package-lock.json`
+- `performance_dashboard/frontend/tsconfig.json`
+- `performance_dashboard/frontend/vite.config.ts`
+- `performance_dashboard/frontend/index.html`
+- `performance_dashboard/frontend/scripts/build.mjs`
+- `performance_dashboard/frontend/src/main.tsx`
+- `performance_dashboard/frontend/src/api.ts`
+- `performance_dashboard/frontend/src/App.tsx`
+- `performance_dashboard/frontend/src/styles.css`
+
+Docs/tests:
+
+- `performance_dashboard/README.md`
+- `docs/contributing/performance_benchmarks.md`
+- `fastvideo/tests/performance/test_dashboard_service.py`
+- `fastvideo/tests/performance/test_dashboard_api.py`
+
+Shared HF utility change:
+
+- `fastvideo/tests/performance/hf_store.py`
+
+## Data Source
+
+The source of truth remains the Hugging Face dataset repo used by existing
+performance CI:
+
+```text
+HF_REPO_ID=FastVideo/performance-tracking
+```
+
+The dataset stores normalized JSON records emitted by
+`fastvideo/tests/performance/compare_baseline.py`. The current v1 normalized
+schema includes:
+
+- `model_id`
+- `timestamp`
+- `commit_sha`
+- `gpu_type`
+- `latency`
+- `throughput`
+- `memory`
+- `text_encoder_time_s`
+- `dit_time_s`
+- `vae_decode_time_s`
+- `success`
+
+Records are grouped by `(model_id, gpu_type)` for v1 dashboard behavior.
+
+## Local Cache
+
+The backend syncs the HF dataset to a local cache directory:
+
+```text
+PERFORMANCE_TRACKING_ROOT=/tmp/fastvideo-perf-dashboard
+```
+
+If `PERFORMANCE_TRACKING_ROOT` is not set, the dashboard defaults to:
+
+```text
+/tmp/fastvideo-perf-dashboard
+```
+
+The sync is performed through the existing helper:
+
+```python
+fastvideo.tests.performance.hf_store.sync_from_hf(...)
+```
+
+The dashboard then loads JSON files from the local cache through:
+
+```python
+fastvideo.tests.performance.hf_store.load_records(...)
+```
+
+## Authentication
+
+Originally `hf_store.py` only read `HF_API_KEY`. This caused local dashboard
+runs to fail when users had standard Hugging Face token variables set.
+
+`hf_store.py` now resolves tokens from the first available variable in:
+
+```text
+HF_API_KEY
+HUGGINGFACE_HUB_TOKEN
+HF_TOKEN
+```
+
+For local use:
+
+```bash
+export HF_TOKEN=hf_...
+```
+
+If the HF repo is private or gated, the token must have dataset read access.
+
+## Backend API
+
+The FastAPI app is created by:
+
+```python
+fastvideo.performance_dashboard.api:create_app
+```
+
+The module-level app is:
+
+```python
+fastvideo.performance_dashboard.api:app
+```
+
+Endpoints:
+
+- `GET /api/performance/health`
+- `POST /api/performance/refresh`
+- `GET /api/performance/records?days=90`
+- `GET /api/performance/summary?days=90`
+- `GET /api/performance/trends?days=90`
+
+`POST /api/performance/refresh` forces a fresh HF sync.
+
+## Status Semantics
+
+Important: the dashboard intentionally separates stored CI status from
+recomputed context.
+
+Stored status:
+
+- Comes directly from the latest JSON record's `success` field.
+- This is what the dashboard displays as `Stored Status`.
+- This is the primary latest status.
+
+Recomputed status:
+
+- Calculated locally from cached records for explanatory context.
+- Uses the latest record's metric values compared to the median of the latest
+  five previous successful records in the same `(model_id, gpu_type)` group.
+- Displayed separately as `Recomputed`.
+- Does not override the stored JSON `success` status.
+
+This distinction was added after observing that recomputing pass/fail from the
+local cache can disagree with the status originally uploaded by CI.
+
+## Time Window Behavior
+
+The default dashboard time window is 90 days.
+
+The selected `days` value affects:
+
+- trend charts
+- record browsing/filtering
+
+The selected `days` value does not affect:
+
+- latest stored status
+- latest summary baseline context
+
+Reason: latest status should not change when users widen or narrow the trend
+window. The API keeps `days` on `/summary` only for shared frontend filter
+state, but summary loading uses all cached records.
+
+This fixed a bug where changing from roughly 35 days to 42 days could change
+the latest status from pass to fail because older records entered the local
+baseline window.
+
+## Metric Logic
+
+Dashboard metric definitions live in:
+
+```text
+fastvideo/performance_dashboard/metrics.py
+```
+
+Tracked metrics:
+
+- `latency` lower is better
+- `throughput` higher is better
+- `memory` lower is better
+- `text_encoder_time_s` lower is better
+- `dit_time_s` lower is better
+- `vae_decode_time_s` lower is better
+
+Baseline context uses the median of up to five previous successful records for
+the same `(model_id, gpu_type)`.
+
+## Frontend Behavior
+
+The React app:
+
+- fetches `/api/performance/summary`
+- fetches `/api/performance/trends`
+- displays summary cards
+- displays latest rows by model/GPU
+- displays native SVG trend charts
+- has model/GPU/day filters
+- includes a refresh button
+- auto-refreshes every five minutes
+
+The UI is implemented without a charting library. Trend charts are native SVG
+in `performance_dashboard/frontend/src/App.tsx`.
+
+The production frontend build uses `esbuild` through
+`performance_dashboard/frontend/scripts/build.mjs`. Vite is still used for the
+dev server and `/api` proxy.
+
+Why esbuild for production build:
+
+- Vite/Rollup hit a local macOS native optional dependency code-signing issue
+  in this environment.
+- Direct esbuild worked reliably and is sufficient for this small dashboard.
+
+## Static Serving
+
+After frontend build, the FastAPI server serves:
+
+- static JS/CSS from `performance_dashboard/frontend/dist/assets`
+- `performance_dashboard/frontend/dist/index.html` for the dashboard page
+
+This allows a single local port to serve both the API and UI.
+
+## Local Run Workflow
+
+Build frontend:
+
+```bash
+cd performance_dashboard/frontend
+conda run -n fastvideo env PATH=/Applications/Codex.app/Contents/Resources/cua_node/bin:/usr/local/bin:/usr/bin:/bin \
+  /Applications/Codex.app/Contents/Resources/cua_node/bin/npm install
+conda run -n fastvideo env PATH=/Applications/Codex.app/Contents/Resources/cua_node/bin:/usr/local/bin:/usr/bin:/bin \
+  /Applications/Codex.app/Contents/Resources/cua_node/bin/npm run build
+```
+
+Run dashboard:
+
+```bash
+export HF_TOKEN=hf_...
+python -m fastvideo.performance_dashboard --host 0.0.0.0 --port 8000
+```
+
+Open locally:
+
+```text
+http://127.0.0.1:8000
+```
+
+## ngrok Workflow
+
+`python -m fastvideo.performance_dashboard --host 0.0.0.0 --port 8000`
+starts the actual local dashboard server.
+
+`ngrok http 8000` does not start the dashboard. It exposes the already-running
+local server through a temporary public URL.
+
+Typical flow:
+
+```bash
+python -m fastvideo.performance_dashboard --host 0.0.0.0 --port 8000
+ngrok http 8000
+```
+
+Use the HTTPS URL printed by ngrok to view the dashboard remotely.
+
+## Verification Commands
+
+Backend tests:
+
+```bash
+conda run -n fastvideo python -m pytest \
+  fastvideo/tests/performance/test_dashboard_service.py \
+  fastvideo/tests/performance/test_dashboard_api.py \
+  -q
+```
+
+Expected after latest changes:
+
+```text
+8 passed
+```
+
+Frontend build:
+
+```bash
+cd performance_dashboard/frontend
+conda run -n fastvideo env PATH=/Applications/Codex.app/Contents/Resources/cua_node/bin:/usr/local/bin:/usr/bin:/bin \
+  /Applications/Codex.app/Contents/Resources/cua_node/bin/npm run build
+```
+
+Expected:
+
+```text
+tsc && node scripts/build.mjs
+```
+
+with exit code 0.
+
+## Known Notes
+
+- `performance_dashboard/frontend/node_modules/` and
+  `performance_dashboard/frontend/dist/` are ignored by git.
+- `npm install` reported two high-severity audit findings in dependency tree.
+  `npm audit fix --force` was not run because it can introduce breaking
+  dependency upgrades.
+- Existing `fastvideo` package imports may emit platform warnings such as NPU
+  or macOS torch distributed messages. These are not dashboard-specific errors.
+

.agents/skills/add-reward-model/SKILL.md

@@ -0,0 +1,32 @@
+---
+name: add-reward-model
+description: Use when adding reusable reward models under fastvideo/train/methods/rl/rewards for RLHF or online RL training.
+---
+
+# Add Reward Model
+
+Use for reward models consumed by RL methods.
+
+## Placement
+
+- Put reusable reward code under `fastvideo/train/methods/rl/rewards/`.
+- Expose public builders from `fastvideo/train/methods/rl/rewards/__init__.py`.
+- Keep method-specific aggregation or advantage logic out of reward classes.
+
+## Media Inputs
+
+- Reward callables receive decoded media tensors.
+- Accept single-frame tensors as `[B, C, H, W]` and multi-frame tensors as `[B, C, T, H, W]` when practical.
+- Frame selection is reward-specific. Frame scorers such as PickScore and CLIPScore should explicitly select frame `0`; temporal rewards should inspect whichever frames they need.
+- Return one scalar reward per prompt/sample.
+
+## Attribution
+
+- If code is ported or closely adapted from another repo, add a short comment or docstring naming the source file/function.
+- Preserve SPDX headers used by FastVideo files.
+
+## Tests
+
+- Unit-test tensor layout handling without loading large reward checkpoints.
+- Allow fake scorer injection for multi-reward tests.
+- Test weighted reward aggregation and metric keys.

.agents/skills/add-rl-method/SKILL.md

@@ -0,0 +1,38 @@
+---
+name: add-rl-method
+description: Use when adding or modifying an RL/RLHF method under fastvideo/train/methods/rl, including DiffusionNFT-like methods.
+---
+
+# Add RL Method
+
+Use for new RL methods in the modular `fastvideo/train` stack.
+
+## Required Shape
+
+- Add the method under `fastvideo/train/methods/rl/`.
+- Subclass `TrainingMethod`.
+- Keep model-family logic in `ModelBase` wrappers.
+- Decode generated latents through `ModelBase.decode_latents`; add that hook to the new model wrapper instead of decoding inside the RL method.
+- Use `fastvideo/train/methods/rl/common/sampling.py` for generation unless the method has a documented reason to avoid sampling.
+- Use `fastvideo/train/methods/rl/common/prompt_sampling.py` for reusable grouped prompt sampling patterns such as DiffusionNFT K-repeat.
+- Use `fastvideo/train/methods/rl/rewards/` for reward models.
+
+## Optimization
+
+- Return `manages_optimization() == True` only when the method must own a nonstandard outer/inner loop.
+- If using managed optimization, implement `managed_train_step(data_stream, iteration)`.
+- Existing trainer callbacks, checkpointing, tracking, and validation should still work.
+
+## Config
+
+- Put method knobs under `method`.
+- Put sampler knobs under `method.sampling`.
+- Do not put scheduler or trajectory policy into model configs.
+- Do not split a diffusers-style scheduler from its built-in `step()` solver in YAML; use `trajectory` only for higher-level ODE vs re-noise behavior.
+- Avoid fixed timestep lists in examples unless reproducing a known baseline; prefer scheduler-generated defaults.
+
+## Tests
+
+- Add fake-model tests for sampler/method behavior.
+- Add config parse tests for the public YAML.
+- Confirm existing train methods stay on the default Trainer path.

.agents/skills/index.jsonl

@@ -8,3 +8,6 @@
 {"name": "decompose-pipeline-pr", "description": "Decompose an oversized FastVideo pipeline PR into a stack of independently-reviewable PRs. Tiers the diff by blast radius (invisible / dead code / cross-cutting infra / activation), produces a branch graph and worktree bootstrap, drafts the AGENTS.md manifest, flags missing tests on cross-cutting infra changes, and extracts lessons from the PR body. Worked example: PR #1280 daVinci-MagiHuman (9.8k LOC) decomposed into 10 stacked PRs.", "path": "decompose-pipeline-pr/SKILL.md", "status": "tested", "trust": "medium"}
 {"name": "reseed-performance-baseline", "description": "Re-seed the HF performance-tracking baseline for an intentional runtime, dependency, or environment-caused benchmark shift. Use when performance CI fails because metrics such as latency, throughput, component time, or peak memory changed for an accepted reason and the rolling median baseline must be advanced by replicating one reviewed shifted source result into three success=true records, or five records when explicitly requested", "path": "reseed-performance-baseline/SKILL.md", "status": "draft", "trust": "low"}
 {"name": "add-model", "description": "Add a new model (or variant) to FastVideo: DiT + configs + pipeline + presets + registry + tests. Walks through FastVideo's single stage-based pipeline architecture with exact file paths and registration hooks.", "path": "add-model/SKILL.md", "status": "draft", "trust": "low"}
+{"name": "rlhf-training-abstractions", "description": "Use when changing FastVideo RLHF/RL training infrastructure, especially sampler, reward, scheduler trajectory, or method boundaries under fastvideo/train.", "path": "rlhf-training-abstractions/SKILL.md", "status": "draft", "trust": "low"}
+{"name": "add-rl-method", "description": "Use when adding or modifying an RL/RLHF method under fastvideo/train/methods/rl, including DiffusionNFT-like methods.", "path": "add-rl-method/SKILL.md", "status": "draft", "trust": "low"}
+{"name": "add-reward-model", "description": "Use when adding reusable reward models under fastvideo/train/methods/rl/rewards for RLHF or online RL training.", "path": "add-reward-model/SKILL.md", "status": "draft", "trust": "low"}

.agents/skills/rlhf-training-abstractions/SKILL.md

@@ -0,0 +1,41 @@
+---
+name: rlhf-training-abstractions
+description: Use when changing FastVideo RLHF/RL training infrastructure, especially sampler, reward, scheduler trajectory, or method boundaries under fastvideo/train.
+---
+
+# RLHF Training Abstractions
+
+Use this skill before editing RLHF-style training code in `fastvideo/train`.
+
+## Boundaries
+
+- RL methods live under `fastvideo/train/methods/rl/` and own algorithm logic: reward collection, advantage computation, policy loss, KL/reference terms, and optimizer cadence.
+- Rewards live under `fastvideo/train/methods/rl/rewards/` and must be reusable across RL methods.
+- RL methods pass decoded media to rewards; each reward decides whether to use the first frame, sampled frames, or the full video.
+- Sampling lives under `fastvideo/train/methods/rl/common/` and must use `ModelBase` primitives plus scheduler math, not model-family inference pipelines.
+- Model wrappers under `fastvideo/train/models/` own model-specific forward details.
+- Model wrappers also own model-specific latent decoding via `ModelBase.decode_latents`; RL methods should not reach into VAE normalization internals.
+- Shared RL helpers such as K-repeat prompt sampling belong under `fastvideo/train/methods/rl/common/` when they are reusable across RL methods.
+
+## Anti-Patterns
+
+- Do not bind RL methods to inference pipeline classes such as `WanDMDPipeline`.
+- Do not hardcode timestep lists in a method when the scheduler can generate them.
+- Do not put reward-model code inside one RL method.
+- Do not make existing non-RL methods use method-managed optimization unless explicitly requested.
+
+## Sampling Policy
+
+- Prefer YAML-configured `method.sampling` with `scheduler`, `trajectory`, `num_steps`, `timesteps`, and `sigmas`.
+- Treat diffusers-style scheduler classes as owning both the timestep schedule and their `step()` update rule; avoid a separate `solver` field unless a new sampler truly implements solver math outside the scheduler object.
+- Missing `timesteps` means “ask the scheduler”; explicit `timesteps` or `sigmas` are overrides.
+- ODE-style trajectories should not re-noise between denoising steps.
+- SDE/re-noise behavior must be explicit in config.
+
+## Validation
+
+- Run focused local tests for sampler config and Trainer opt-in behavior.
+- Verify existing train methods still report `manages_optimization() == False`.
+- Keep fixed-prompt validation helpers in `fastvideo/train/methods/rl/common/validation.py` so new RL methods can reuse sharding and captions.
+- Test distributed prompt grouping helpers separately from heavyweight model loading.
+- Run `pre-commit run --files <changed paths>`; respect configured excludes.

.buildkite/pipeline.yml

@@ -217,6 +217,17 @@ steps:
             limit: 2
       agents:
         queue: "default"
+    - label: ":test_tube: Eval Metrics Tests"
+      if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "eval"
+      command: "timeout 90m .buildkite/scripts/pr_test.sh"
+      retry:
+        automatic:
+          - exit_status: 128
+            limit: 3
+          - exit_status: -1
+            limit: 2
+      agents:
+        queue: "default"
 
   # ============================================================
   # Fastcheck: Runs on every PR (~10-15 min parallel)
@@ -483,3 +494,15 @@ steps:
                   - TEST_TYPE=train_framework
                 agents:
                   queue: "default"
+            - path:
+                - "fastvideo/eval/**"
+                - "fastvideo/tests/eval/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 90m .buildkite/scripts/pr_test.sh"
+                label: ":test_tube: Eval Metrics Tests"
+                env:
+                  - TEST_TYPE=eval
+                agents:
+                  queue: "default"

.buildkite/scripts/pr_test.sh

@@ -76,7 +76,7 @@ EFFECTIVE_PR=${BUILDKITE_PULL_REQUEST:-false}
 if [ "$EFFECTIVE_PR" = "false" ] && [ -n "${PR_NUMBER:-}" ]; then
     EFFECTIVE_PR=$PR_NUMBER
 fi
-MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$EFFECTIVE_PR BUILDKITE_BRANCH=${BUILDKITE_BRANCH:-} TEST_SCOPE=${TEST_SCOPE:-} IMAGE_VERSION=$IMAGE_VERSION"
+MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$EFFECTIVE_PR BUILDKITE_BRANCH=${BUILDKITE_BRANCH:-} TEST_SCOPE=${TEST_SCOPE:-} BUILDKITE_BUILD_URL=${BUILDKITE_BUILD_URL:-} BUILDKITE_BUILD_ID=${BUILDKITE_BUILD_ID:-} BUILDKITE_JOB_ID=${BUILDKITE_JOB_ID:-} IMAGE_VERSION=$IMAGE_VERSION"
 
 POST_RUN_HOOK=""
 
@@ -219,6 +219,10 @@ case "$TEST_TYPE" in
         log "Running fastvideo.train framework tests..."
         MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_train_framework_tests"
         ;;
+    "eval")
+        log "Running eval metric tests..."
+        MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_eval_tests"
+        ;;
     "lora_extraction")
         log "Running LoRA extraction tests..."
         MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_lora_extraction_tests"

.github/mergify.yml

@@ -272,7 +272,7 @@ pull_request_rules:
         remove: [needs-rebase]
 
   # ============================================================
-  # Auto-merge and auto-rebase
+  # Auto-merge
   # ============================================================
 
   - name: auto-merge when ready and all checks pass
@@ -290,16 +290,6 @@ pull_request_rules:
       merge:
         method: squash
 
-  - name: auto-update when ready
-    conditions:
-      - label=ready
-      - "#approved-reviews-by>=1"
-      - -conflict
-      - -closed
-      - -draft
-    actions:
-      update: {}
-
   # ============================================================
   # PR title format help
   # ============================================================

.github/workflows/ci-slash-commands.yml

@@ -125,7 +125,7 @@ jobs:
           set -euo pipefail
           TEST_NAME=$(echo "$COMMENT" | grep -oP '(?<=/test\s)\S+' | head -1 || true)
 
-          VALID="encoder vae transformer kernel unit dreamverse ssim training lora-inference lora-training distillation self-forcing vsa vmoba performance api train-framework full fastcheck pre-commit"
+          VALID="encoder vae transformer kernel unit dreamverse ssim training lora-inference lora-training distillation self-forcing vsa vmoba performance api train-framework eval full fastcheck pre-commit"
           if [ -z "$TEST_NAME" ] || ! echo "$VALID" | grep -qw "$TEST_NAME"; then
             echo "Unknown test: '$TEST_NAME'. Valid: $VALID"
             exit 1
@@ -139,7 +139,7 @@ jobs:
             [distillation]=distillation_dmd [self-forcing]=self_forcing
             [vsa]=training_vsa [vmoba]=inference_vmoba
             [performance]=performance [api]=api_server
-            [train-framework]=train_framework
+            [train-framework]=train_framework [eval]=eval
           )
 
           if [ "$TEST_NAME" = "full" ]; then

.github/workflows/publish-kernel.yml

@@ -53,23 +53,15 @@ jobs:
       fail-fast: false
       matrix:
           os: [ubuntu-22.04]
-          python-version: ['3.10', '3.11', '3.12']
+          python-version: ['3.12']
           torch-cuda:
-            # - torch-version: '2.5.1'
-            #   cuda-version: '12.4.1'
-            #   torch-cuda-short: 'cu124'
-            # - torch-version: '2.6.0'
-            #   cuda-version: '12.6.3'
-            #   torch-cuda-short: 'cu126'
-            # - torch-version: '2.7.1'
-            #   cuda-version: '12.8.0'
-            #   torch-cuda-short: 'cu128'
-            # - torch-version: '2.9.1'
-            #   cuda-version: '12.8.0'
-            #   torch-cuda-short: 'cu128'
-            - torch-version: '2.10.0'
-              cuda-version: '12.8.0'
-              torch-cuda-short: 'cu128'
+            # torch 2.12 dropped cu128; cu130 (CUDA 13) is the default wheel, cu126 covers older drivers.
+            - torch-version: '2.12.0'
+              cuda-version: '12.6.3'
+              torch-cuda-short: 'cu126'
+            - torch-version: '2.12.0'
+              cuda-version: '13.0.0'
+              torch-cuda-short: 'cu130'
 
     steps:
       - name: Free up disk space
@@ -104,7 +96,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Install CUDA ${{ matrix.torch-cuda.cuda-version }}
-        uses: Jimver/cuda-toolkit@v0.2.21
+        uses: Jimver/cuda-toolkit@v0.2.35
         id: cuda-toolkit
         with:
           cuda: ${{ matrix.torch-cuda.cuda-version }}
@@ -152,9 +144,27 @@ jobs:
 
           cd fastvideo-kernel
           git submodule update --init --recursive  # Ensure ThunderKittens submodule is initialized
-          # Release builds are produced on GPU-less runners, so force-enable TK and target Hopper.
-          export TORCH_CUDA_ARCH_LIST="9.0a"
-          export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
+          # Release builds run on GPU-less runners, so force-enable TK and set the arch explicitly.
+          # The CUDA 13 (cu130) leg additionally targets Blackwell sm_120a and builds the
+          # attn_qat_infer FP4 inference kernels (they need CUDA Toolkit 12.8+, so the cu126 leg
+          # stays Hopper-only and ships without them).
+          # Cap Ninja build concurrency to fit the 16 GB standard runner. CUTLASS FP4
+          # and ThunderKittens template TUs each need several GB; Ninja's default
+          # (cores+2 ≈ 6 parallel jobs) exhausts RAM and the host OOM-kills the build
+          # (SIGTERM / exit 143, surfaced as "runner lost communication").
+          if [ "${{ matrix.torch-cuda.torch-cuda-short }}" = "cu130" ]; then
+            export TORCH_CUDA_ARCH_LIST="9.0a;12.0a"
+            export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
+            # A single FP4 TU (attn_qat_infer) can use ~8-12 GB on its own, so serialize.
+            export CMAKE_BUILD_PARALLEL_LEVEL=1
+          else
+            export TORCH_CUDA_ARCH_LIST="9.0a"
+            export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
+            # No FP4 here; only the two TK TUs are heavy (single-arch) and they fit
+            # side by side. Uncapped (cores+2=6) has built but occasionally OOMs when
+            # the light TUs pile on; -j4 overlaps them safely (~12-14 GB peak).
+            export CMAKE_BUILD_PARALLEL_LEVEL=4
+          fi
 
           # Build standard wheel (no local version suffix) for PyPI
           python -m build --wheel --outdir dist
@@ -183,8 +193,8 @@ jobs:
           mv fixed_dist/*.whl dist/
 
       - name: Upload wheel artifact
-        # Only upload if it's the "main" CUDA version we want on PyPI
-        # We upload all to artifacts for inspection/GH releases, but give them distinct artifact names
+        # Upload every matrix leg as a distinct artifact (for inspection / GitHub releases).
+        # The publish job below selects which CUDA build is pushed to PyPI.
         uses: actions/upload-artifact@v4
         with:
           name: fastvideo_kernel-py${{ matrix.python-version }}-${{ matrix.torch-cuda.torch-cuda-short }}-torch${{ matrix.torch-cuda.torch-version }}
@@ -204,13 +214,15 @@ jobs:
 
       - uses: actions/setup-python@v5
         with:
-          python-version: '3.10'
+          python-version: '3.12'
 
       - name: Download PyPI wheels
+        # Publish only the cu130 (CUDA 13) wheel to PyPI — it includes the Blackwell FP4 kernels.
+        # The cu126 wheel stays available as a build artifact / GitHub-release asset.
         uses: actions/download-artifact@v4
         with:
           path: fastvideo-kernel/dist/
-          pattern: 'fastvideo_kernel-py*'
+          pattern: 'fastvideo_kernel-py*-cu130-*'
           merge-multiple: true
 
       - name: Install uv

.gitignore

@@ -130,3 +130,11 @@ apps/dreamverse/web/.env.production.local
 .sisyphus/
 openspec/
 fastvideo/tests/ssim/reference_videos/**
+
+# Editor logs and local Python version pins (accidentally committed)
+*.nvimlog
+.nvimlog
+.python-version
+
+# Gradio runtime/flagging artifacts
+.gradio/

.nvimlog

@@ -1 +0,0 @@
-WRN 2026-03-26T13:46:33.469 ?.19646    server_start:193: Failed to start server: operation not permitted: /var/folders/z_/h_6myyk14d1b7z87z3vy4mjh0000gn/T/nvim.dsynkd/iSe0el/nvim.19646.0

.python-version

@@ -1 +0,0 @@
-3.12

README.md

@@ -9,8 +9,9 @@
 **FastVideo is a unified post-training and real-time inference framework for accelerated video generation.**
 
 ## NEWS
-- `2026/03/17`: Release Live demo: [Into the Dreamverse: Vibe Directing in FastVideo](https://dreamverse.fastvideo.org/), check out the [Blog](https://haoailab.com/blogs/dreamverse/).
-- `2026/03/13`: Release Live demo: [Create a 5s 1080p Video in 4.5s with FastVideo on a Single GPU](https://1080p.fastvideo.org/), check out the [Blog](https://haoailab.com/blogs/fastvideo_realtime_1080p/).
+- `2026/06/23`: Release FastWan-QAD: 5s of Video generated in 1.8s E2E. [FastWan-QAD models](https://huggingface.co/FastVideo/FastWan-QAD-FP8-1.3B), check out the [Blog](https://haoailab.com/blogs/fastwan-qad/).
+- `2026/03/17`: Release demo: [Into the Dreamverse: Vibe Directing in FastVideo](), check out the [Blog](https://haoailab.com/blogs/dreamverse/).
+- `2026/03/13`: Release demo: [Create a 5s 1080p Video in 4.5s with FastVideo on a Single GPU](), check out the [Blog](https://haoailab.com/blogs/fastvideo_realtime_1080p/).
 - `2025/11/19`: Release [CausalWan2.2 I2V A14B Preview](https://huggingface.co/FastVideo/CausalWan2.2-I2V-A14B-Preview-Diffusers) models, [Blog](https://hao-ai-lab.github.io/blogs/fastvideo_causalwan_preview/) and [Inference Code!](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py).
 - `2025/08/04`: Release [FastWan](https://hao-ai-lab.github.io/FastVideo/distillation/dmd) models and [Sparse-Distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).
 

apps/dreamverse/docker/README.md

@@ -32,7 +32,7 @@ checkout with the `dreamverse` extra, installs the FA4
 flash-attention fork, builds native FFmpeg, and installs FlashInfer for NVFP4
 quantization.
 
-FastVideo's pinned `fastvideo-kernel==0.2.6` package is installed by default.
+FastVideo's pinned `fastvideo-kernel==0.3.0` package is installed by default.
 To rebuild `fastvideo-kernel` from this checkout during the image build, set:
 
 ```bash

build_env.sh

@@ -0,0 +1,70 @@
+#!/usr/bin/env bash
+# FastVideo (trackwan_bidir) env — CORRECTED build. CUDA 12.9 container, aarch64 + GB200 (sm_100a).
+# Ordering fix: install repo-pinned torch (2.11.0 cu128) FIRST, build the local Blackwell kernel
+# against it, then core deps (relaxed kernel pin so the local 0.3.0 satisfies), then FA4.
+set +e
+say(){ echo; echo "==================== $* ===================="; }
+
+export DEBIAN_FRONTEND=noninteractive
+export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
+export PATH=$CUDA_HOME/bin:/mnt/lustre/vlm-s4duan/bin:$PATH
+export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
+export UV_CACHE_DIR=/mnt/lustre/vlm-s4duan/.uv/cache
+export UV_PYTHON_INSTALL_DIR=/mnt/lustre/vlm-s4duan/.uv/python
+export TORCH_CUDA_ARCH_LIST=10.0a
+export MAX_JOBS=${MAX_JOBS:-64}
+cd /mnt/lustre/vlm-s4duan/FastVideo || { echo NO_REPO; exit 1; }
+
+say "ENSURE toolchain"
+if ! command -v g++ >/dev/null || ! command -v cmake >/dev/null || ! command -v ninja >/dev/null; then
+  apt-get update -y >/tmp/apt.log 2>&1 && apt-get install -y --no-install-recommends g++ gcc make cmake ninja-build git ca-certificates >>/tmp/apt.log 2>&1 && echo "apt ok" || { echo APT_FAIL; tail -20 /tmp/apt.log; }
+fi
+nvcc --version | tail -2
+
+say "FRESH venv (python 3.12 on Lustre)"
+rm -rf .venv
+uv venv --python 3.12 --seed .venv || exit 1
+source .venv/bin/activate
+python --version
+
+say "STAGE 1: torch 2.11.0 cu128 (repo pin, Blackwell-capable)"
+uv pip install --torch-backend=cu128 torch==2.11.0 torchvision torchaudio 2>&1 | tail -15
+python -c "import torch;print('TORCH',torch.__version__,torch.version.cuda,'avail',torch.cuda.is_available())" || echo TORCH_FAIL
+
+say "STAGE 2: submodules + build local fastvideo-kernel (sm_100a) against torch 2.11"
+git submodule update --init --recursive fastvideo-kernel/include/cutlass fastvideo-kernel/include/tk 2>&1 | tail -5
+( cd fastvideo-kernel && TORCH_CUDA_ARCH_LIST=10.0a bash ./build.sh ) 2>&1 | tail -20
+echo "KERNEL_EXIT=$?"
+python -c "import fastvideo_kernel;print('kernel import OK')" || echo KERNEL_IMPORT_FAIL
+
+say "STAGE 3: core deps  uv pip install -e .[dev]  (local kernel already satisfies relaxed pin)"
+uv pip install -e ".[dev]" 2>&1 | tail -40
+echo "CORE_EXIT=$?"
+python -c "import fastvideo;print('fastvideo import OK')" || echo FASTVIDEO_IMPORT_FAIL
+
+say "STAGE 4: FA4 — OFFICIAL Dao-AILab flash_attn/cute (NOT the XOR-op fork; fork is stale/broken)"
+# Package name is flash-attn-4; pass the bare git URL (name-qualified spec is rejected).
+# --prerelease allow lets it pull the matching nvidia-cutlass-dsl==4.6.0.dev0 + quack>=0.5.3.
+uv pip install --prerelease allow \
+  "git+https://github.com/Dao-AILab/flash-attention.git#subdirectory=flash_attn/cute" 2>&1 | tail -25
+python -c "from flash_attn.cute.interface import _flash_attn_fwd,_flash_attn_bwd; print('FA4_CUTE_OK')" 2>&1
+
+say "VERIFY"
+python - <<'PY'
+import importlib, torch
+print("torch", torch.__version__, torch.version.cuda, "cuda_avail", torch.cuda.is_available())
+if torch.cuda.is_available():
+    print("device0", torch.cuda.get_device_name(0), "cc", torch.cuda.get_device_capability(0))
+    x = torch.randn(2048, 2048, device="cuda", dtype=torch.bfloat16)
+    y = (x @ x); torch.cuda.synchronize()
+    print("bf16 matmul on GPU OK", tuple(y.shape))
+for m in ["fastvideo", "fastvideo_kernel", "flash_attn"]:
+    try: importlib.import_module(m); print("import OK:", m)
+    except Exception as e: print("import FAIL:", m, repr(e)[:180])
+try:
+    from flash_attn.cute.interface import _flash_attn_fwd; print("FA4 flash_attn.cute.interface OK")
+except Exception as e:
+    print("FA4 cute interface:", repr(e)[:180])
+PY
+
+say "DONE build_env.sh (v2)"

build_trackwan_init.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+set +e
+export HF_HOME=/mnt/lustre/vlm-s4duan/.hf
+export TOKENIZERS_PARALLELISM=false
+cd /mnt/lustre/vlm-s4duan/FastVideo || exit 1
+source .venv/bin/activate
+
+BASE=/mnt/lustre/vlm-s4duan/models/Wan2.1-Fun-1.3B-InP-Diffusers
+OUT=/mnt/lustre/vlm-s4duan/models/trackwan_1.3b_init
+mkdir -p /mnt/lustre/vlm-s4duan/models
+
+echo "==================== download Fun-InP diffusers base ===================="
+python - <<PY
+from huggingface_hub import snapshot_download
+p = snapshot_download("weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers", local_dir="$BASE")
+print("downloaded to", p)
+PY
+echo "DL_EXIT=$?"
+
+echo "==================== base transformer sanity (I2V => in_channels 36) ===================="
+python - <<PY
+import json
+c = json.load(open("$BASE/transformer/config.json"))
+print("in_channels:", c.get("in_channels"), "image_dim:", c.get("image_dim"), "num_layers:", c.get("num_layers"), "hidden:", c.get("hidden_size") or c.get("dim"))
+PY
+
+echo "==================== convert_trackwan_init.py ===================="
+python data_pipeline/convert_trackwan_init.py --base "$BASE" --out "$OUT"
+echo "CONVERT_EXIT=$?"
+
+echo "==================== trackwan init result ===================="
+ls -la "$OUT" 2>&1 | head -25
+python - <<PY
+import json
+c = json.load(open("$OUT/transformer/config.json"))
+print("trackwan in_channels:", c.get("in_channels"), "| has track_config:", "track_config" in c)
+PY
+echo "==================== DONE build_trackwan_init ===================="

data_pipeline/720_stage_1/00_make_overfit_subset.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+# Phase 0, step 0 — carve a SMALL overfit dataset out of the 720p openvid parquets.
+# Symlinks a few data_chunk parquets into a dedicated dir; the map-style loader walks it for
+# *.parquet, so a couple of chunks (~32 clips each) is a good overfit set. Non-destructive
+# (symlinks only). Re-run to rebuild; it clears the stale map_style_cache.
+set -uo pipefail
+
+SRC_ROOT=${SRC_ROOT:-/home/hal-shared/motionstream/data/openvid-wantrack-parquets}
+SRC_SHARD=${SRC_SHARD:-shard000}
+N_CHUNKS=${N_CHUNKS:-2}                       # ~32 clips/chunk -> ~64 clips
+OUT=${OUT:-/home/hal-kevin/data/motion-stream-test/overfit_subset_720p/combined_parquet_dataset}
+
+src_worker="$SRC_ROOT/$SRC_SHARD/combined_parquet_dataset/worker_0"
+[ -d "$src_worker" ] || { echo "[subset] source not found: $src_worker" >&2; exit 1; }
+
+dst_worker="$OUT/worker_0"
+rm -rf "$OUT"                                 # drop old subset + its map_style_cache
+mkdir -p "$dst_worker"
+
+n=0
+for f in $(ls "$src_worker"/data_chunk_*.parquet | sort -V | head -n "$N_CHUNKS"); do
+  ln -s "$(readlink -f "$f")" "$dst_worker/$(basename "$f")"
+  n=$((n + 1))
+done
+
+echo "[subset] linked $n parquet chunk(s) from $SRC_SHARD -> $OUT"
+python - "$OUT" <<'PY'
+import glob, sys, pyarrow.parquet as pq
+fs = glob.glob(f"{sys.argv[1]}/**/*.parquet", recursive=True)
+rows = sum(pq.ParquetFile(f).metadata.num_rows for f in fs)
+print(f"[subset] {len(fs)} file(s), {rows} clips total")
+PY
+echo "[subset] point the overfit config data_path at: $OUT"

data_pipeline/720_stage_1/01_build_init.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+# Phase 0, step 1 — build the d64 + bias 14B WanTrack init for the overfit.
+# CPU only, needs ~62GB RAM (loads the 14B base) -> run on a COMPUTE node, not the login node.
+# pretrained channels are preserved; the added track slot is ZERO-init (--pe-init zero, matching
+# upstream's trackwan_14b_i2v_d64_zero_init_bias) and the track_encoder gets default init WITH bias
+# (--use-bias-defaults), matching TRACKWAN_TRACK_BIAS=1 at train time.
+set -uo pipefail
+cd ~/FastVideo
+
+BASE=${BASE:-/home/hal-kevin/models/Wan2.1-I2V-14B-720P-Diffusers}
+OUT=${OUT:-/home/hal-kevin/models/trackwan_14b_i2v_d64_bias_init}
+ID_DIM=${ID_DIM:-64}
+PE_INIT=${PE_INIT:-zero}
+
+python data_pipeline/convert_trackwan_init_v2.py \
+  --base "$BASE" \
+  --out  "$OUT" \
+  --id-dim "$ID_DIM" \
+  --pe-init "$PE_INIT" \
+  --use-bias-defaults
+
+echo "[init] built $OUT (id_dim=$ID_DIM, pe-init=$PE_INIT, bias=on)"
+echo "[init] expect: 'added 4 track_encoder tensors' (2 weights + 2 bias)"
+python -c "
+from safetensors import safe_open
+ks=[k for k in safe_open('$OUT/transformer/diffusion_pytorch_model.safetensors','pt').keys() if 'track_encoder' in k]
+import json; c=json.load(open('$OUT/transformer/config.json'))
+print('[init] in_channels', c['in_channels'], '| track_config.id_dim', c['track_config']['id_dim'])
+print('[init] track_encoder keys:', sorted(ks))
+"

data_pipeline/720_stage_1/02_run_overfit.sh

@@ -0,0 +1,56 @@
+#!/bin/bash
+# Phase 0, step 2 — overfit the track pathway. 14B/720p, d64+bias, sparse, HEAD TRAINABLE.
+# Run the SAME command on BOTH racks (2x4 GB200). MASTER_ADDR = rack0 host on both; NODE_RANK
+# differs (0 on rack0, 1 on rack1). STAGE=A uses fixed track IDs (lock onto the pattern), STAGE=B
+# uses random IDs; run A first, then B (B resumes A's checkpoint via resume_from_checkpoint: latest).
+#
+#   rack0: MASTER_ADDR=hpc-rack-1-6 NODE_RANK=0 STAGE=A bash data_pipeline/720_stage_1/02_run_overfit.sh
+#   rack1: MASTER_ADDR=hpc-rack-1-6 NODE_RANK=1 STAGE=A bash data_pipeline/720_stage_1/02_run_overfit.sh
+# Extra args pass through to the trainer, e.g. ... STAGE=B bash 02_run_overfit.sh --training.loop.max_train_steps 1600
+set -uo pipefail
+cd ~/FastVideo
+
+: "${MASTER_ADDR:?set MASTER_ADDR to rack-0 hostname (reachable from both racks)}"
+: "${NODE_RANK:?set NODE_RANK: 0 on the master rack, 1 on the other}"
+STAGE=${STAGE:-A}
+MASTER_PORT=${MASTER_PORT:-29502}
+CFG=data_pipeline/720_stage_1/finetune_wantrack_overfit_14b_720p_d64_bias.yaml
+
+# Stage A = deterministic track IDs (fixed sampling); Stage B = random IDs.
+# B writes to its OWN output dir (keeps A's dir pristine, A/B checkpoints separated) — mirrors
+# upstream chain_stepAB.sh + run_stepB_seed.sh. Seed B's dir once with A's final checkpoint via
+# 02b_seed_stageB.sh BEFORE launching B, so resume_from_checkpoint: latest picks up A's weights.
+OUT_A=/home/hal-kevin/data/motion-stream-test/overfit_14b_720p_d64_bias_out
+OUT_B=/home/hal-kevin/data/motion-stream-test/overfit_14b_720p_d64_bias_stageB_out
+case "$STAGE" in
+  A) FIXED=1; OUT_DIR=$OUT_A ;;
+  B) FIXED=0; OUT_DIR=$OUT_B ;;
+  *) echo "STAGE must be A or B (got '$STAGE')" >&2; exit 1 ;;
+esac
+if [ "$STAGE" = B ] && [ ! -d "$OUT_B" ]; then
+  echo "[overfit] ERROR: STAGE=B but $OUT_B does not exist." >&2
+  echo "          Seed it first (run ONCE): bash data_pipeline/720_stage_1/02b_seed_stageB.sh" >&2
+  exit 1
+fi
+
+# Pin NCCL to the four active 400G InfiniBand HCAs (keep off the 200G Ethernet port).
+export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_3,mlx5_4
+# export NCCL_DEBUG=INFO   # uncomment on the FIRST launch to confirm NET/IB, then re-comment.
+
+FASTVIDEO_FA4=1 \
+  TRACKWAN_TRACK_BIAS=1 WANTRACK_FREEZE_HEAD=0 \
+  WANTRACK_AUG=1 WANTRACK_SPARSE=1 WANTRACK_EXTRA_RANDOM=20 \
+  WANTRACK_PMASK=0 WANTRACK_MASK_CHUNK=0 \
+  WANTRACK_IMAGE_COND=0 \
+  WANTRACK_FIXED_SAMPLE=${FIXED} \
+  torchrun \
+  --nnodes=2 --nproc_per_node=4 --node_rank=${NODE_RANK} \
+  --rdzv_id=overfit_14b_720p --rdzv_backend=c10d --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
+  --log-dir ~/FastVideo/torchrun_logs \
+  -m fastvideo.train.entrypoint.train \
+  --config ${CFG} \
+  --training.checkpoint.output_dir ${OUT_DIR} \
+  --training.tracker.run_name overfit-14b-720p-d64-bias-stage${STAGE} \
+  --training.checkpoint.resume_from_checkpoint latest \
+  "$@" \
+  2>&1 | tee -a data_pipeline/720_stage_1/overfit_node${NODE_RANK}_stage${STAGE}.log

data_pipeline/720_stage_1/02b_seed_stageB.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+# Seed Stage B's output dir with Stage A's FINAL checkpoint, so B runs in its own dir while its
+# config still uses resume_from_checkpoint: latest (which is what makes crash-restarts safe).
+# Mirrors upstream examples/train/run_stepB_seed.sh.
+#
+# Run this ONCE (not per-rack) on the shared filesystem, AFTER Stage A finishes, BEFORE launching
+# Stage B. Hardlinked (cp -al), not copied: a 14B training-state checkpoint is 100s of GB, and both
+# dirs are on the same filesystem (/home), so hardlinks are instant and use no extra space.
+#
+#   bash data_pipeline/720_stage_1/02b_seed_stageB.sh
+#   # then, on BOTH racks:
+#   MASTER_ADDR=hpc-rack-1-6 NODE_RANK=0 STAGE=B bash data_pipeline/720_stage_1/02_run_overfit.sh \
+#     --training.loop.max_train_steps 3000
+set -uo pipefail
+
+OUT_A=${OUT_A:-/home/hal-kevin/data/motion-stream-test/overfit_14b_720p_d64_bias_out}
+OUT_B=${OUT_B:-/home/hal-kevin/data/motion-stream-test/overfit_14b_720p_d64_bias_stageB_out}
+
+# Pick A's latest COMPLETE checkpoint (dcp/.metadata present) — refuse to seed off a partial save.
+STEP=${STEP:-}
+if [ -z "$STEP" ]; then
+  for d in $(for c in "$OUT_A"/checkpoint-*; do n=${c##*checkpoint-}; echo "$n"; done | sort -rn); do
+    [ -f "$OUT_A/checkpoint-$d/dcp/.metadata" ] && { STEP=$d; break; }
+  done
+fi
+[ -n "$STEP" ] || { echo "[B-seed] no complete checkpoint found under $OUT_A" >&2; exit 1; }
+
+SRC="$OUT_A/checkpoint-$STEP"
+DST="$OUT_B/checkpoint-$STEP"
+[ -f "$SRC/dcp/.metadata" ] || { echo "[B-seed] $SRC incomplete (no dcp/.metadata) — refusing" >&2; exit 1; }
+
+mkdir -p "$OUT_B"
+if [ -e "$DST" ]; then
+  echo "[B-seed] $DST already exists — leaving it alone"
+else
+  echo "[B-seed] hardlinking $SRC -> $DST"
+  cp -al "$SRC" "$DST"
+fi
+echo "[B-seed] Stage B dir now seeded at step $STEP:"
+for c in "$OUT_B"/checkpoint-*; do echo "  $c"; done
+echo "[B-seed] Launch B with --training.loop.max_train_steps > $STEP (e.g. $((STEP + 1000)))."

data_pipeline/720_stage_1/03_export.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+# Phase 0, step 3 — export the overfit DCP checkpoint to a diffusers model dir.
+# The Phase-1 merge (convert_trackwan_init_v2.py --track-src/--pe-src) reads this diffusers dir.
+# Only 1 GPU needed (DCP reshards automatically). Run on a compute node.
+set -uo pipefail
+cd ~/FastVideo
+
+# --checkpoint accepts an output_dir (auto-picks the latest checkpoint-<step>), a specific
+# checkpoint-<step> dir, or its dcp/ subdir.
+# Default = the Stage-B dir (the FINAL overfit; B refines A and is what the merge lifts). Override
+# CKPT to the Stage-A dir only if you deliberately want to export A.
+CKPT=${CKPT:-/home/hal-kevin/data/motion-stream-test/overfit_14b_720p_d64_bias_stageB_out}
+OUT=${OUT:-/home/hal-kevin/models/overfit14b_export}
+CFG=${CFG:-data_pipeline/720_stage_1/finetune_wantrack_overfit_14b_720p_d64_bias.yaml}
+
+# TRACKWAN_TRACK_BIAS=1 MUST match training: it toggles bias on track_encoder.{proj,temporal_conv}
+# at build time (track_encoder.py:73). The overfit trained with bias=1, so the checkpoint carries
+# those bias params; building bias-less here fails to load them ("track_encoder.proj.bias not found").
+TRACKWAN_TRACK_BIAS=1 \
+python -m fastvideo.train.entrypoint.dcp_to_diffusers \
+  --checkpoint "$CKPT" \
+  --output-dir "$OUT" \
+  --config "$CFG" \
+  --role student \
+  --overwrite \
+  || { echo "[export] FAILED (see traceback above) -- nothing written" >&2; exit 1; }
+
+# Guard against a silent partial write (the entrypoint can exit 0 yet write no weights).
+ls "$OUT"/transformer/*.safetensors >/dev/null 2>&1 \
+  || { echo "[export] ERROR: no transformer/*.safetensors under $OUT -- export did not complete" >&2; exit 1; }
+
+echo "[export] wrote diffusers model -> $OUT"
+echo "[export] next (Phase 1 merge):"
+echo "  python data_pipeline/convert_trackwan_init_v2.py \\"
+echo "    --base /home/hal-kevin/models/Wan2.1-I2V-14B-720P-Diffusers \\"
+echo "    --out  /home/hal-kevin/models/trackwan_14b_i2v_d64_merged_from_overfit_bias \\"
+echo "    --id-dim 64 --pe-init random \\"
+echo "    --track-src $OUT/transformer --pe-src $OUT/transformer"

data_pipeline/720_stage_1/04_merge.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+# Step 04 — MERGE. Graft the overfit's co-adapted track pathway (track_encoder + patch-embed track
+# slot [36:52]) onto a PRISTINE 14B base, discarding the overfit's base degradation. CPU only,
+# ~62GB RAM -> run on a compute node. Needs 03_export.sh to have produced the overfit diffusers dir.
+set -uo pipefail
+cd ~/FastVideo
+
+BASE=${BASE:-/home/hal-kevin/models/Wan2.1-I2V-14B-720P-Diffusers}
+SRC=${SRC:-/home/hal-kevin/models/overfit14b_export}     # produced by 03_export.sh
+OUT=${OUT:-/home/hal-kevin/models/trackwan_14b_i2v_d64_merged_from_overfit_bias}
+ID_DIM=${ID_DIM:-64}
+
+[ -d "$SRC/transformer" ] || { echo "[merge] $SRC/transformer not found — run 03_export.sh first" >&2; exit 1; }
+
+# --track-src + --pe-src from the SAME export = the merge: encoder AND its co-adapted track slot
+# lifted together (bias copied through); everything else comes from the pristine --base.
+python data_pipeline/convert_trackwan_init_v2.py \
+  --base "$BASE" \
+  --out  "$OUT" \
+  --id-dim "$ID_DIM" --pe-init random \
+  --track-src "$SRC/transformer" \
+  --pe-src    "$SRC/transformer"
+
+echo "[merge] built merged init -> $OUT"
+python -c "
+import json; c=json.load(open('$OUT/transformer/config.json'))
+print('[merge] in_channels', c['in_channels'], '| id_dim', c['track_config']['id_dim'])
+"
+echo "[merge] next: 05_run_openvid_stage1.sh (init_from defaults to this dir)"

data_pipeline/720_stage_1/05_run_openvid_stage1.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+# Step 05 — OpenVid Stage 1 (the big run, ~10 days). 14B/720p, merged init, HEAD FROZEN, sparse.
+# Run the SAME command on BOTH racks. MASTER_ADDR = rack0 host on both; NODE_RANK 0 on rack0, 1 on rack1.
+#   rack0: MASTER_ADDR=hpc-rack-1-6 NODE_RANK=0 bash data_pipeline/720_stage_1/05_run_openvid_stage1.sh
+#   rack1: MASTER_ADDR=hpc-rack-1-6 NODE_RANK=1 bash data_pipeline/720_stage_1/05_run_openvid_stage1.sh
+# resume_from_checkpoint: latest -> relaunch both to continue after any interruption.
+set -uo pipefail
+cd ~/FastVideo
+
+: "${MASTER_ADDR:?set MASTER_ADDR to rack-0 hostname (reachable from both racks)}"
+: "${NODE_RANK:?set NODE_RANK: 0 on the master rack, 1 on the other}"
+MASTER_PORT=${MASTER_PORT:-29503}
+NNODES=${NNODES:-8}                 # 8 nodes x 4 GB200 = 32 GPUs (matches num_gpus:32 / replicate 8 in the config)
+GPUS_PER_NODE=${GPUS_PER_NODE:-4}
+CFG=data_pipeline/720_stage_1/finetune_wantrack_openvid_stage1_14b_720p_d64_bias.yaml
+
+export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_3,mlx5_4
+# export NCCL_DEBUG=INFO   # uncomment on the FIRST launch to confirm NET/IB, then re-comment.
+
+# Stage 1: sparse conditioning, HEAD FROZEN (train the DiT to use the merged track pathway), no masking.
+# Every WANTRACK_/TRACKWAN_ knob is set EXPLICITLY (no reliance on code defaults) to match upstream
+# stage-1 (D) and to avoid inheriting the overfit launcher's opposite settings (IMAGE_COND=0, FIXED_SAMPLE=1).
+FASTVIDEO_FA4=1 \
+  TRACKWAN_TRACK_BIAS=1 \
+  WANTRACK_FREEZE_HEAD=1 \
+  WANTRACK_IMAGE_COND=1 \
+  WANTRACK_AUG=1 \
+  WANTRACK_SPARSE=1 WANTRACK_EXTRA_RANDOM=20 WANTRACK_EXTRA_MODE=random \
+  WANTRACK_FIXED_SAMPLE=0 \
+  WANTRACK_PMASK=0 WANTRACK_MASK_CHUNK=0 \
+  WANTRACK_TRACK_DROP=0 WANTRACK_MOTION_DROP=0 WANTRACK_TEXT_DROP=0 \
+  torchrun \
+  --nnodes=${NNODES} --nproc_per_node=${GPUS_PER_NODE} --node_rank=${NODE_RANK} \
+  --rdzv_id=openvid_stage1_14b_720p --rdzv_backend=c10d --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
+  --log-dir ~/FastVideo/torchrun_logs \
+  -m fastvideo.train.entrypoint.train \
+  --config ${CFG} \
+  --training.distributed.num_gpus $((NNODES * GPUS_PER_NODE)) \
+  --training.checkpoint.resume_from_checkpoint latest \
+  "$@" \
+  2>&1 | tee -a data_pipeline/720_stage_1/openvid_stage1_node${NODE_RANK}.log

data_pipeline/720_stage_1/05b_run_openvid_stage1_2rack.sh

@@ -0,0 +1,66 @@
+#!/bin/bash
+# Step 05 (2-RACK variant) — OpenVid Stage 1 on 2 racks (2x4 = 8 GB200) instead of 8 nodes.
+#
+# SAME effective global batch as the 8-node 05: grad_accum is bumped 4x (4 -> 16) to compensate for
+# 4x fewer GPUs, so the optimization is equivalent (see note below) — only the wall-clock is ~4x
+# (~48 days vs ~12). Run the SAME command on BOTH racks; MASTER_ADDR = rack0 host on both,
+# NODE_RANK 0 on rack0, 1 on rack1.
+#   rack0: MASTER_ADDR=hpc-rack-1-6 NODE_RANK=0 bash data_pipeline/720_stage_1/05b_run_openvid_stage1_2rack.sh
+#   rack1: MASTER_ADDR=hpc-rack-1-6 NODE_RANK=1 bash data_pipeline/720_stage_1/05b_run_openvid_stage1_2rack.sh
+# resume_from_checkpoint: latest -> relaunch both to continue after any interruption.
+set -uo pipefail
+cd ~/FastVideo
+
+: "${MASTER_ADDR:?set MASTER_ADDR to rack-0 hostname (reachable from both racks)}"
+: "${NODE_RANK:?set NODE_RANK: 0 on the master rack, 1 on the other}"
+MASTER_PORT=${MASTER_PORT:-29503}
+NNODES=${NNODES:-2}                 # 2 racks x 4 GB200 = 8 GPUs
+GPUS_PER_NODE=${GPUS_PER_NODE:-4}
+CFG=data_pipeline/720_stage_1/finetune_wantrack_openvid_stage1_14b_720p_d64_bias.yaml
+# bf16 parquets (complete 259k set, ~4.7TB). Loader honors the per-field _dtype; same clips/order
+# as fp32 (so val_sample_indices are unchanged). Override to the fp32 set if ever needed:
+#   DATA_PATH=/home/hal-shared/motionstream/data/openvid-wantrack-parquets
+DATA_PATH=${DATA_PATH:-/home/hal-shared/motionstream/data/openvid-wantrack-parquets-bf16}
+
+# --- batch math: hold the effective global batch EQUAL to the 8-node 05 ------------------
+# 8-node config: replicate 8 x shard 4 = 32 GPUs, grad_accum 4.
+# 2 racks:       replicate 2 x shard 4 =  8 GPUs, grad_accum 16.
+# 4x fewer GPUs x 4x grad_accum = same effective batch under EITHER counting convention
+# (num_gpus- or replicate_dim-based both scale by 4). Model still shards across 4 GPUs
+# (shard_dim 4) as in the 8-node run, so per-GPU memory is unchanged (no OOM risk from this).
+HSDP_REPLICATE=${HSDP_REPLICATE:-2}
+HSDP_SHARD=${HSDP_SHARD:-4}
+GRAD_ACCUM=${GRAD_ACCUM:-16}
+
+export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_3,mlx5_4
+# export NCCL_DEBUG=INFO   # uncomment on the FIRST launch to confirm NET/IB, then re-comment.
+
+# Stage 1 env — identical to 05 (every knob explicit): sparse, HEAD FROZEN, CLIP on, no masking.
+FASTVIDEO_FA4=1 \
+  TRACKWAN_TRACK_BIAS=1 \
+  WANTRACK_FREEZE_HEAD=1 \
+  WANTRACK_IMAGE_COND=1 \
+  WANTRACK_AUG=1 \
+  WANTRACK_SPARSE=1 WANTRACK_EXTRA_RANDOM=20 WANTRACK_EXTRA_MODE=random \
+  WANTRACK_FIXED_SAMPLE=0 \
+  WANTRACK_PMASK=0 WANTRACK_MASK_CHUNK=0 \
+  WANTRACK_TRACK_DROP=0 WANTRACK_MOTION_DROP=0 WANTRACK_TEXT_DROP=0 \
+  torchrun \
+  --nnodes=${NNODES} --nproc_per_node=${GPUS_PER_NODE} --node_rank=${NODE_RANK} \
+  --rdzv_id=openvid_stage1_14b_720p_2rack --rdzv_backend=c10d --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
+  --log-dir ~/FastVideo/torchrun_logs \
+  -m fastvideo.train.entrypoint.train \
+  --config ${CFG} \
+  --training.data.data_path ${DATA_PATH} \
+  --training.distributed.num_gpus $((NNODES * GPUS_PER_NODE)) \
+  --training.distributed.hsdp_replicate_dim ${HSDP_REPLICATE} \
+  --training.distributed.hsdp_shard_dim ${HSDP_SHARD} \
+  --training.loop.gradient_accumulation_steps ${GRAD_ACCUM} \
+  --training.checkpoint.resume_from_checkpoint latest \
+  --training.checkpoint.training_state_checkpointing_steps 20 \
+  --training.checkpoint.checkpoints_total_limit 50 \
+  --callbacks.track_validation.validate_at_start false \
+  --callbacks.track_validation.every_steps 100 \
+  --callbacks.track_validation.val_sample_indices "[1660, 1888]" \
+  "$@" \
+  2>&1 | tee -a data_pipeline/720_stage_1/openvid_stage1_2rack_node${NODE_RANK}.log

data_pipeline/720_stage_1/06_export_stage1.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+# Step 06 — export the OpenVid Stage-1 DCP checkpoint to a diffusers dir. Stage 2 (07) inits its
+# WEIGHTS from this export (init_from, fresh optimizer/step), so it must be a diffusers model dir,
+# not a DCP resume. Only 1 GPU needed. Run on a compute node.
+set -uo pipefail
+cd ~/FastVideo
+
+CKPT=${CKPT:-/home/hal-kevin/data/motion-stream-test/openvid_stage1_14b_720p_out}   # auto-picks latest
+OUT=${OUT:-/home/hal-kevin/models/openvid_stage1_14b_export}
+CFG=${CFG:-data_pipeline/720_stage_1/finetune_wantrack_openvid_stage1_14b_720p_d64_bias.yaml}
+
+python -m fastvideo.train.entrypoint.dcp_to_diffusers \
+  --checkpoint "$CKPT" \
+  --output-dir "$OUT" \
+  --config "$CFG" \
+  --role student \
+  --overwrite
+
+echo "[export] stage-1 teacher -> $OUT"
+echo "[export] next: 07_run_synth_stage2.sh (init_from defaults to this dir)"

data_pipeline/720_stage_1/07_run_synth_stage2.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+# Step 07 — Stage 2 (robustness, ~1 day). 14B/720p, inits from exported Stage-1 teacher, HEAD FROZEN,
+# sparse + heavy masking/dropout. Produces the FINAL bidir teacher. Run on BOTH racks (same command,
+# NODE_RANK 0/1, MASTER_ADDR = rack0).
+#   rack0: MASTER_ADDR=hpc-rack-1-6 NODE_RANK=0 bash data_pipeline/720_stage_1/07_run_synth_stage2.sh
+#   rack1: MASTER_ADDR=hpc-rack-1-6 NODE_RANK=1 bash data_pipeline/720_stage_1/07_run_synth_stage2.sh
+set -uo pipefail
+cd ~/FastVideo
+
+: "${MASTER_ADDR:?set MASTER_ADDR to rack-0 hostname (reachable from both racks)}"
+: "${NODE_RANK:?set NODE_RANK: 0 on the master rack, 1 on the other}"
+MASTER_PORT=${MASTER_PORT:-29504}
+NNODES=${NNODES:-8}                 # 8 nodes x 4 GB200 = 32 GPUs (matches num_gpus:32 / replicate 8 in the config)
+GPUS_PER_NODE=${GPUS_PER_NODE:-4}
+CFG=data_pipeline/720_stage_1/finetune_wantrack_synth_stage2_14b_720p_d64_bias.yaml
+
+export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_3,mlx5_4
+# export NCCL_DEBUG=INFO   # first launch only.
+
+# Stage 2: sparse + robustness masking/dropout, HEAD still FROZEN.
+# Every WANTRACK_/TRACKWAN_ knob is set EXPLICITLY (no code defaults) to match upstream stage-2 (E):
+# same as stage-1 but with masking (PMASK=0.2/CHUNK=8) and dropout (TRACK=0.5, MOTION=0.3) ON.
+FASTVIDEO_FA4=1 \
+  TRACKWAN_TRACK_BIAS=1 \
+  WANTRACK_FREEZE_HEAD=1 \
+  WANTRACK_IMAGE_COND=1 \
+  WANTRACK_AUG=1 \
+  WANTRACK_SPARSE=1 WANTRACK_EXTRA_RANDOM=20 WANTRACK_EXTRA_MODE=random \
+  WANTRACK_FIXED_SAMPLE=0 \
+  WANTRACK_PMASK=0.2 WANTRACK_MASK_CHUNK=8 \
+  WANTRACK_TRACK_DROP=0.5 WANTRACK_MOTION_DROP=0.3 WANTRACK_TEXT_DROP=0 \
+  torchrun \
+  --nnodes=${NNODES} --nproc_per_node=${GPUS_PER_NODE} --node_rank=${NODE_RANK} \
+  --rdzv_id=synth_stage2_14b_720p --rdzv_backend=c10d --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
+  --log-dir ~/FastVideo/torchrun_logs \
+  -m fastvideo.train.entrypoint.train \
+  --config ${CFG} \
+  --training.distributed.num_gpus $((NNODES * GPUS_PER_NODE)) \
+  --training.checkpoint.resume_from_checkpoint latest \
+  "$@" \
+  2>&1 | tee -a data_pipeline/720_stage_1/synth_stage2_node${NODE_RANK}.log

data_pipeline/720_stage_1/README.md

@@ -0,0 +1,83 @@
+# Stage 1 — Bidirectional 14B/720p track teacher (full pipeline)
+
+The complete bidirectional teacher: overfit the track pathway → merge it into a pristine 14B
+base → OpenVid stage-1 (frozen head) → synth stage-2 (robustness). The final teacher then seeds
+the causal (Self-Forcing) student. Full rationale: `../notes/trackwan_14b_720p_teacher_merge_plan.md`.
+
+Recipe: **d64 + bias, sparse conditioning, flow_shift 6, 720p.** Hardware: 2×4 GB200, 400G IB,
+manual 2-node `torchrun` (rack0 = `hpc-rack-1-6` = NODE_RANK 0; rack1 = `hpc-rack-1-8` = NODE_RANK 1;
+`MASTER_ADDR` = rack0 on both).
+
+## Files / run order
+
+| Step | Script | What | Where |
+|---|---|---|---|
+| 0 | `00_make_overfit_subset.sh` | carve a small overfit set (symlinks ~2 parquet chunks) | login/compute |
+| 1 | `01_build_init.sh` | build `trackwan_14b_i2v_d64_bias_init` | compute (CPU, ~62GB RAM) |
+| 2 | `02_run_overfit.sh` | overfit track pathway, head trainable (Stage A→B) | **both racks** |
+| 3 | `03_export.sh` | overfit DCP → diffusers | compute (1 GPU) |
+| 4 | `04_merge.sh` | graft pathway into pristine 14B base → merged init | compute (CPU, ~62GB) |
+| 5 | `05_run_openvid_stage1.sh` | OpenVid stage-1, **head frozen** (~10 days) | **both racks** |
+| 6 | `06_export_stage1.sh` | stage-1 DCP → diffusers | compute (1 GPU) |
+| 7 | `07_run_synth_stage2.sh` | stage-2 robustness, head frozen (~1 day) → **final teacher** | **both racks** |
+
+Configs (referenced by the scripts): `finetune_wantrack_overfit_14b_720p_d64_bias.yaml`,
+`finetune_wantrack_openvid_stage1_14b_720p_d64_bias.yaml`,
+`finetune_wantrack_synth_stage2_14b_720p_d64_bias.yaml`.
+
+## Commands
+
+```bash
+# 0-1: data + init
+bash data_pipeline/720_stage_1/00_make_overfit_subset.sh
+bash data_pipeline/720_stage_1/01_build_init.sh                    # compute node
+
+# 2: overfit — both racks, Stage A (fixed IDs) then Stage B (random IDs, resumes A)
+MASTER_ADDR=hpc-rack-1-6 NODE_RANK=0 STAGE=A bash data_pipeline/720_stage_1/02_run_overfit.sh   # rack0
+MASTER_ADDR=hpc-rack-1-6 NODE_RANK=1 STAGE=A bash data_pipeline/720_stage_1/02_run_overfit.sh   # rack1
+#   ...when track-following is clear in validation, stop and run Stage B (bump steps as needed):
+MASTER_ADDR=hpc-rack-1-6 NODE_RANK=0 STAGE=B bash data_pipeline/720_stage_1/02_run_overfit.sh --training.loop.max_train_steps 1600
+MASTER_ADDR=hpc-rack-1-6 NODE_RANK=1 STAGE=B bash data_pipeline/720_stage_1/02_run_overfit.sh --training.loop.max_train_steps 1600
+
+# 3-4: export + merge
+bash data_pipeline/720_stage_1/03_export.sh                        # compute node
+bash data_pipeline/720_stage_1/04_merge.sh                         # compute node
+
+# 5: OpenVid stage-1 — both racks (~10 days; resume by relaunching)
+MASTER_ADDR=hpc-rack-1-6 NODE_RANK=0 bash data_pipeline/720_stage_1/05_run_openvid_stage1.sh   # rack0
+MASTER_ADDR=hpc-rack-1-6 NODE_RANK=1 bash data_pipeline/720_stage_1/05_run_openvid_stage1.sh   # rack1
+
+# 6-7: export stage-1 + synth stage-2 -> final teacher
+bash data_pipeline/720_stage_1/06_export_stage1.sh                 # compute node
+MASTER_ADDR=hpc-rack-1-6 NODE_RANK=0 bash data_pipeline/720_stage_1/07_run_synth_stage2.sh      # rack0
+MASTER_ADDR=hpc-rack-1-6 NODE_RANK=1 bash data_pipeline/720_stage_1/07_run_synth_stage2.sh      # rack1
+```
+
+## Notes & decisions baked in
+
+- **Topology / batch:** stage-1 and stage-2 target the **upstream recipe on 8 nodes × 4 GB200 = 32
+  GPUs** — `num_gpus 32, hsdp_replicate_dim 8, hsdp_shard_dim 4, grad_accum 4 → global bs 128`.
+  Stage-1 = 4,800 steps = 2.4 epochs over 259k ≈ **~12 days**; stage-2 = 600 steps ≈ **~1.5 days**.
+  (On only 8 GPUs / 2 nodes this recipe is ~49 days — if you drop back, set `NNODES=2`,
+  `hsdp_replicate_dim 2`, `grad_accum 2` → global bs 16, and use ~8000 stage-1 steps for a ~10-day
+  half-epoch.) The overfit (steps 0–3) stays on **2 nodes** — it's ~64 clips, so 32 GPUs / bs 128
+  would exceed the dataset; leave it at `NNODES=2` there.
+- **Launching the 8-node stages (05, 07):** the scripts take `NNODES` (default 8) and pass a matching
+  `--training.distributed.num_gpus`. But manually running one command on each of 8 nodes (NODE_RANK
+  0…7) is impractical — **use a SLURM launcher** (`srun` spans all nodes with `--node-rank=$SLURM_PROCID`,
+  as in `examples/train/run_slurm.sh` / `run_wan14b_held.sh`). The overfit (02) is fine to launch
+  manually on 2 nodes.
+- **Env knobs** live in the launch scripts: overfit = `FREEZE_HEAD=0`; stage-1/2 = `FREEZE_HEAD=1`;
+  stage-2 adds `TRACK_DROP=0.5 MOTION_DROP=0.3 PMASK=0.2 MASK_CHUNK=8`. `TRACKWAN_TRACK_BIAS=1` and
+  `WANTRACK_SPARSE=1 EXTRA_RANDOM=20` throughout.
+- **Stage-2 data** defaults to the openvid parquets + masking (you have no synth generated). Swap
+  `data_path` to a synth set if you build one; the robustness comes from the masking either way.
+- **`dit_precision`:** fp32 master (default) for the real teacher stages; upstream uses bf16 — flip
+  only if memory-blocked.
+- **Monitor track-following, not loss** — the `track_validation` with-track vs no-track/adversarial
+  deltas. Gate the overfit (step 2) on these before merging, and watch stage-1's first validation
+  (step 2000) before committing the full ~10 days. Check efficiency with `python scripts/mfu_estimate.py`.
+- **Not runnable end-to-end yet:** steps 4–7 chain via default paths (04 reads 03's export, 05 reads
+  04's merge, etc.), so they're correct now but only *run* once the prior step's output exists.
+- **Two things to smoke-test once:** `dcp_to_diffusers` on a real checkpoint (steps 3 & 6), and that
+  the overfit actually learns track-following on ~64 clips (if weak: `N_CHUNKS=4` in step 0, or more steps).

data_pipeline/720_stage_1/finetune_wantrack_openvid_stage1_14b_720p_d64_bias.yaml

@@ -0,0 +1,86 @@
+# Step 05 config — OpenVid Stage 1. 14B/720p, merged d64+bias init, HEAD FROZEN, sparse.
+# The big run: the DiT learns to USE the frozen (merged) track pathway on the full 720p openvid set.
+#
+# BATCH: upstream recipe = global bs 128 on 32 GPUs (8 nodes x 4). This config targets that: 4800
+# steps x 128 = 614k samples = 2.4 epochs over 259k, ~12 days on 32 GPUs. (On only 8 GPUs, bs 128
+# would be grad_accum 16 = ~49 days; if you drop back to 8 nodes/GPUs, use grad_accum 2 = global
+# bs 16 and ~8000 steps for a ~10-day half-epoch instead.)
+models:
+  student:
+    _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel
+    init_from: /home/hal-kevin/models/trackwan_14b_i2v_d64_merged_from_overfit_bias   # from 04_merge.sh
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  model_path: /home/hal-kevin/models/trackwan_14b_i2v_d64_merged_from_overfit_bias
+  # dit_precision: fp32 master (default) for the real teacher (standard MFU / stable updates at lr 1e-5).
+  # upstream sets bf16; flip only if memory-blocked (see notes/trackwan_14b_720p_teacher_merge_plan.md).
+
+  distributed:
+    num_gpus: 32
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 8   # across the 8 nodes
+    hsdp_shard_dim: 4       # within each node (over NVLink)
+
+  data:
+    data_path: /home/hal-shared/motionstream/data/openvid-wantrack-parquets   # full 259k, 720p
+    dataloader_num_workers: 2
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 31
+    num_height: 720
+    num_width: 1280
+    num_frames: 121
+
+  optimizer:
+    learning_rate: 1.0e-5    # MotionStream teacher stage-1
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4800            # upstream stage-1: 2.4 epochs over 259k at bs128; ~12 days on 32 GPUs; resumable
+    gradient_accumulation_steps: 4   # 1 x 32 GPU x 4 = global batch 128
+
+  checkpoint:
+    output_dir: /home/hal-kevin/data/motion-stream-test/openvid_stage1_14b_720p_out
+    training_state_checkpointing_steps: 200
+    checkpoints_total_limit: 5
+    resume_from_checkpoint: latest
+
+  tracker:
+    trackers: [wandb]
+    project_name: wantrack-bidir
+    run_name: openvid-stage1-14b-720p-d64-bias
+    entity: s4duan-uc-san-diego
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  track_validation:
+    _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback
+    every_steps: 2000            # 14B/720p validation is a multi-hour job; keep it rare
+    num_val_samples: 1
+    val_sample_indices: [28]
+    num_inference_steps: 20
+    guidance_scale: 3.0
+    motion_guidance_scale: 1.5
+    fps: 24
+    grid_stride: 3
+    tail: 12
+    include_heldout: false
+    validate_at_start: false
+    seed: 1000
+
+pipeline:
+  flow_shift: 6

data_pipeline/720_stage_1/finetune_wantrack_overfit_14b_720p_d64_bias.yaml

@@ -0,0 +1,85 @@
+# Phase 0 overfit — 14B/720p, d64 + bias, sparse conditioning, HEAD TRAINABLE.
+# Goal: co-adapt the track pathway (track_encoder + patch-embed track slot) on a small subset so
+# the model follows tracks. The base is DISCARDED by the Phase-1 merge (only the track pathway is
+# lifted), so dit_precision: bf16 here is fine + faster. Sparse/bias/freeze knobs come from the
+# launch script (02_run_overfit.sh), not this file.
+models:
+  student:
+    _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel
+    init_from: /home/hal-kevin/models/trackwan_14b_i2v_d64_bias_init
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  model_path: /home/hal-kevin/models/trackwan_14b_i2v_d64_bias_init
+  dit_precision: bf16     # overfit base is thrown away by the merge; bf16 for speed
+
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 2   # across the 2 racks
+    hsdp_shard_dim: 4       # within each rack
+
+  data:
+    data_path: /home/hal-kevin/data/motion-stream-synth/preprocessed_i2v_track/combined_parquet_dataset
+    dataloader_num_workers: 2
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 31
+    num_height: 720
+    num_width: 1280
+    num_frames: 121
+
+  optimizer:
+    learning_rate: 1.0e-4   # upstream overfit LR (synth_sparse_*_14b), constant
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 2000            # upstream overfit steps (per stage); stop earlier if tracks lock in
+    gradient_accumulation_steps: 1   # 1 x 8 GPU x 1 = global batch 8 (fine on a tiny subset)
+
+  checkpoint:
+    output_dir: /home/hal-kevin/data/motion-stream-test/overfit_14b_720p_d64_bias_out
+    training_state_checkpointing_steps: 100
+    checkpoints_total_limit: 30    # keep effectively all (20 ckpts @ every-100 for 2000 steps) so you can export the BEST, not just the latest
+    resume_from_checkpoint: latest
+
+  tracker:
+    trackers: [wandb]
+    project_name: wantrack-bidir
+    run_name: overfit-14b-720p-d64-bias
+    entity: s4duan-uc-san-diego
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  track_validation:
+    _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback
+    # Validate on the overfit clips themselves — you want to SEE track-following emerge. The
+    # with-track vs no-track/adversarial deltas are the signal to stop the overfit.
+    every_steps: 200
+    num_val_samples: 2
+    val_sample_indices: [0, 10]
+    num_inference_steps: 20
+    guidance_scale: 3.0
+    motion_guidance_scale: 1.5
+    fps: 24
+    grid_stride: 3
+    tail: 12
+    include_heldout: false
+    validate_at_start: false
+    seed: 1000
+
+pipeline:
+  flow_shift: 6

data_pipeline/720_stage_1/finetune_wantrack_synth_stage2_14b_720p_d64_bias.yaml

@@ -0,0 +1,84 @@
+# Step 07 config — Stage 2 (robustness). 14B/720p, inits from the exported Stage-1 teacher, HEAD
+# still FROZEN, sparse + heavy masking/dropout. Short finetune (lr 1e-6) -> the FINAL teacher.
+#
+# DATA: upstream uses a synthetic set here for controlled motion. You don't have synth generated,
+# so this defaults to the SAME 720p openvid parquets + the stage-2 masking env (the robustness comes
+# from the masking, applied on whatever data). Swap data_path to a synth set if you generate one
+# (data_pipeline/gen_synth_worker.py). Masking knobs are set in 07_run_synth_stage2.sh, not here.
+models:
+  student:
+    _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel
+    init_from: /home/hal-kevin/models/openvid_stage1_14b_export   # from 06_export_stage1.sh
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  model_path: /home/hal-kevin/models/openvid_stage1_14b_export
+
+  distributed:
+    num_gpus: 32
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 8   # across the 8 nodes
+    hsdp_shard_dim: 4       # within each node
+
+  data:
+    data_path: /home/hal-shared/motionstream/data/openvid-wantrack-parquets   # swap to synth if available
+    dataloader_num_workers: 2
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 31
+    num_height: 720
+    num_width: 1280
+    num_frames: 121
+
+  optimizer:
+    learning_rate: 1.0e-6    # low LR, short robustness finetune
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 600             # upstream stage-2; ~1.5 days on 32 GPUs at global bs 128
+    gradient_accumulation_steps: 4   # 1 x 32 GPU x 4 = global batch 128
+
+  checkpoint:
+    output_dir: /home/hal-kevin/data/motion-stream-test/synth_stage2_14b_720p_out
+    training_state_checkpointing_steps: 100
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: latest
+
+  tracker:
+    trackers: [wandb]
+    project_name: wantrack-bidir
+    run_name: synth-stage2-14b-720p-d64-bias
+    entity: s4duan-uc-san-diego
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  track_validation:
+    _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback
+    every_steps: 300
+    num_val_samples: 1
+    val_sample_indices: [28]
+    num_inference_steps: 20
+    guidance_scale: 3.0
+    motion_guidance_scale: 1.5
+    fps: 24
+    grid_stride: 3
+    tail: 12
+    include_heldout: false
+    validate_at_start: false
+    seed: 1000
+
+pipeline:
+  flow_shift: 6

data_pipeline/add_captions.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fill in real captions on a videos2caption.json built from OpenVid clips.
+
+The OpenVid-WanTrack shards ship mp4s only -- no metadata -- so a manifest built by
+scanning the clip directory has empty ``cap`` fields. Stage 5 would then encode empty
+strings through T5, giving every clip an identical null text embedding: text conditioning
+(and the joint text+motion CFG that depends on a meaningful conditional/null contrast)
+would be silently dead.
+
+OpenVid-1M's caption CSV keys on exactly the same filenames (``---_iRTHryQ_13_0to241.mp4``),
+so the join is 1:1 on basename -- no id parsing required.
+
+    python data_pipeline/add_captions.py --manifest <root>/videos2caption.json
+
+The CSV is downloaded once from the Hub (~300 MB) and cached; pass --captions-csv to use
+a local copy. Clips with no caption keep "" and are counted, never silently dropped.
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import sys
+from pathlib import Path
+
+CAPTION_REPO = "nkp37/OpenVid-1M"
+CAPTION_FILES = ["data/train/OpenVid-1M.csv", "data/train/OpenVidHD.csv"]
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--manifest", type=Path, required=True, help="videos2caption.json to patch in place.")
+    p.add_argument("--captions-csv", type=Path, nargs="*", default=None,
+                   help="Local caption CSV(s). Default: download+cache from the Hub.")
+    p.add_argument("--cache-dir", type=Path, default=Path.home() / ".cache/openvid_captions",
+                   help="Where downloaded caption CSVs are cached.")
+    p.add_argument("--min-coverage", type=float, default=0.9,
+                   help="Fail if fewer than this fraction of clips get a caption (0 = never fail).")
+    p.add_argument("--dry-run", action="store_true", help="Report coverage without writing.")
+    return p.parse_args()
+
+
+def caption_paths(args: argparse.Namespace) -> list[Path]:
+    if args.captions_csv:
+        return list(args.captions_csv)
+    from huggingface_hub import hf_hub_download
+    args.cache_dir.mkdir(parents=True, exist_ok=True)
+    out = []
+    for f in CAPTION_FILES:
+        try:
+            out.append(Path(hf_hub_download(CAPTION_REPO, f, repo_type="dataset",
+                                            cache_dir=str(args.cache_dir))))
+        except Exception as e:  # noqa: BLE001
+            print(f"[caption] warn: could not fetch {f} ({e})", flush=True)
+    if not out:
+        sys.exit("[caption] ERROR: no caption CSV available")
+    return out
+
+
+def load_captions(paths: list[Path]) -> dict[str, str]:
+    """basename -> caption. Later files do not overwrite earlier hits."""
+    caps: dict[str, str] = {}
+    csv.field_size_limit(10 * 1024 * 1024)      # captions can be long
+    for p in paths:
+        n0 = len(caps)
+        with p.open(newline="", encoding="utf-8", errors="replace") as fh:
+            for row in csv.DictReader(fh):
+                key, cap = row.get("video"), row.get("caption")
+                if key and cap and key not in caps:
+                    caps[key] = cap.strip()
+        print(f"[caption] {p.name}: +{len(caps) - n0} captions (total {len(caps)})", flush=True)
+    return caps
+
+
+def main() -> None:
+    args = parse_args()
+    items = json.loads(args.manifest.read_text())
+    caps = load_captions(caption_paths(args))
+
+    hit = 0
+    missing: list[str] = []
+    for it in items:
+        name = Path(it.get("path", "")).name
+        cap = caps.get(name)
+        if cap:
+            it["cap"] = [cap]
+            hit += 1
+        else:
+            it.setdefault("cap", [""])
+            missing.append(name)
+
+    cov = hit / max(len(items), 1)
+    print(f"[caption] matched {hit}/{len(items)} clips ({cov*100:.1f}%)", flush=True)
+    if missing:
+        print(f"[caption] first few unmatched: {missing[:3]}", flush=True)
+
+    if cov < args.min_coverage:
+        sys.exit(f"[caption] ERROR: coverage {cov*100:.1f}% < required {args.min_coverage*100:.0f}%. "
+                 "Refusing to write -- training on empty captions silently breaks text conditioning.")
+
+    if args.dry_run:
+        print("[caption] dry run, manifest not written", flush=True)
+        return
+    tmp = args.manifest.with_suffix(".json.tmp")
+    tmp.write_text(json.dumps(items, indent=2))
+    tmp.replace(args.manifest)
+    print(f"[caption] wrote {args.manifest}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/benchmark_tracks.sh

@@ -0,0 +1,148 @@
+#!/bin/bash
+# Benchmark Stage 3 (extract_tracks) + Stage 4 (segment_tracks) end-to-end on the
+# motion-stream-test source videos, writing ALL outputs to a separate bench dir so
+# the real dataset is never touched.
+#
+# Usage:
+#   bash data_pipeline/benchmark_tracks.sh                 # 50 videos, 4 GPUs
+#   GPUS=0 bash data_pipeline/benchmark_tracks.sh          # single GPU
+#   LIMIT=5 bash data_pipeline/benchmark_tracks.sh         # smoke run
+#   VIZ=1 bash data_pipeline/benchmark_tracks.sh           # include viz mp4 rendering
+#   FUSED=1 bash data_pipeline/benchmark_tracks.sh         # fused stage 3+4 (extract --segment)
+#
+# Results append to $OUT_DIR/benchmark_results.txt (tagged with git commit); each
+# entry records gpus/videos/viz/fused so runs stay comparable.
+set -euo pipefail
+
+SRC_VIDEOS=${SRC_VIDEOS:-/home/hal-kevin/data/motion-stream-test/videos}
+SRC_MANIFEST=${SRC_MANIFEST:-/home/hal-kevin/data/motion-stream-test/videos2caption.json}
+OUT_DIR=${OUT_DIR:-/home/hal-kevin/data/motion-stream-qtest}
+GPUS=${GPUS:-0,1,2,3}
+LIMIT=${LIMIT:-}
+VIZ=${VIZ:-0}    # 0 = lean run (production-like for large-scale), 1 = render overlay mp4s
+FUSED=${FUSED:-0}  # 1 = single fused pass (extract_tracks --segment); stage 4 not run
+AMP=${AMP:-0}      # 1 = bf16 autocast for CoTracker (--amp)
+COMPILE=${COMPILE:-0}  # 1 = torch.compile the main CoTracker pass (--compile)
+OVERRIDE_EVERY=${OVERRIDE_EVERY:-3}  # --vis-override-every (2 rides the entry-sweep mask cache in fused mode)
+
+cd "$(dirname "$0")/.."
+IFS=',' read -ra GPU_ARR <<< "$GPUS"
+WORLD_SIZE=${#GPU_ARR[@]}
+MANIFEST=bench_manifest.json
+LOG="$OUT_DIR/benchmark.log"
+RESULTS="$OUT_DIR/benchmark_results.txt"
+LIMIT_ARGS=()
+[[ -n "$LIMIT" ]] && LIMIT_ARGS=(--limit "$LIMIT")
+VIZ_ARGS=()
+[[ "$VIZ" == "1" ]] && VIZ_ARGS=(--viz --viz-dir "$OUT_DIR/bench_viz")
+FUSED_ARGS=()
+if [[ "$FUSED" == "1" ]]; then
+    FUSED_ARGS=(--segment --vis-override-every "$OVERRIDE_EVERY")
+    [[ "$VIZ" == "1" ]] && FUSED_ARGS+=("${VIZ_ARGS[@]}")
+fi
+SPEED_ARGS=()
+[[ "$AMP" == "1" ]] && SPEED_ARGS+=(--amp)
+if [[ "$COMPILE" == "1" ]]; then
+    SPEED_ARGS+=(--compile)
+    # Persist compile artifacts in $HOME (default /tmp/torchinductor_* is wiped by
+    # reboots/tmp-cleaners), so repeat runs start warm instead of recompiling.
+    export TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR:-$HOME/.cache/torchinductor}
+    export TRITON_CACHE_DIR=${TRITON_CACHE_DIR:-$HOME/.cache/triton}
+fi
+
+# --- setup: symlink source videos, copy manifest without points_path -------------
+mkdir -p "$OUT_DIR"
+ln -sfn "$SRC_VIDEOS" "$OUT_DIR/bench_videos"
+python - "$SRC_MANIFEST" "$OUT_DIR/$MANIFEST" <<'PY'
+import json, sys
+items = json.load(open(sys.argv[1]))
+for it in items:
+    it.pop("points_path", None)  # stage 3 re-patches these to the bench tracks dir
+json.dump(items, open(sys.argv[2], "w"), indent=2)
+print(f"[bench] manifest: {len(items)} items -> {sys.argv[2]}")
+PY
+rm -rf "$OUT_DIR/bench_tracks"
+[[ "$VIZ" == "1" ]] && rm -rf "$OUT_DIR/bench_viz"
+: > "$LOG"
+
+N_VIDEOS=$(ls "$OUT_DIR"/bench_videos/*.mp4 | wc -l)
+[[ -n "$LIMIT" ]] && N_VIDEOS=$LIMIT
+COMMIT=$(git rev-parse --short HEAD 2>/dev/null || echo unknown)$(git diff --quiet 2>/dev/null || echo -dirty)
+echo "[bench] commit=$COMMIT  gpus=$GPUS  videos=$N_VIDEOS  log=$LOG"
+
+wait_workers() {  # wait_workers <stage-name> <pid...>
+    local stage=$1 fail=0 pid
+    shift
+    for pid in "$@"; do
+        wait "$pid" || fail=$((fail + 1))
+    done
+    [[ $fail -gt 0 ]] && echo "[bench] WARNING: $fail $stage worker(s) failed — check $LOG"
+    return 0
+}
+
+# --- stage 3: extract tracks ------------------------------------------------------
+t0=$(date +%s)
+PIDS=()
+for i in "${!GPU_ARR[@]}"; do
+    CUDA_VISIBLE_DEVICES=${GPU_ARR[$i]} python -u data_pipeline/extract_tracks.py \
+        --data-dir "$OUT_DIR" \
+        --videos-subdir bench_videos \
+        --out-subdir bench_tracks \
+        --manifest "$MANIFEST" \
+        --grid-size 50 \
+        --device cuda \
+        --detect-entries \
+        --sam-conf 0.75 --sam-iou 0.9 --sam-imgsz 1024 \
+        --entry-sample-every 2 --entry-min-area 0.001 --entry-new-area 0.5 \
+        "${FUSED_ARGS[@]}" \
+        "${SPEED_ARGS[@]}" \
+        --force \
+        --rank "$i" --world-size "$WORLD_SIZE" \
+        "${LIMIT_ARGS[@]}" \
+        >> "$LOG" 2>&1 &
+    PIDS+=($!)
+done
+wait_workers "stage-3" "${PIDS[@]}"
+t1=$(date +%s)
+S3=$((t1 - t0)); [[ $S3 -eq 0 ]] && S3=1
+N_NPZ=$(ls "$OUT_DIR"/bench_tracks/*.npz 2>/dev/null | wc -l || true)
+echo "[bench] stage 3: ${S3}s for $N_NPZ npz"
+[[ "$N_NPZ" -eq "$N_VIDEOS" ]] || echo "[bench] WARNING: expected $N_VIDEOS npz — check $LOG"
+
+# --- stage 4: segment tracks (skipped when FUSED=1: stage 3 already segmented) -----
+S4=0
+if [[ "$FUSED" != "1" ]]; then
+t2=$(date +%s)
+PIDS=()
+for i in "${!GPU_ARR[@]}"; do
+    CUDA_VISIBLE_DEVICES=${GPU_ARR[$i]} python -u data_pipeline/segment_tracks.py \
+        --data-dir "$OUT_DIR" \
+        --videos-subdir bench_videos \
+        --manifest "$MANIFEST" \
+        --conf 0.75 --iou 0.9 --imgsz 1024 \
+        --vis-override-every "$OVERRIDE_EVERY" \
+        "${VIZ_ARGS[@]}" \
+        --force \
+        --rank "$i" --world-size "$WORLD_SIZE" \
+        "${LIMIT_ARGS[@]}" \
+        >> "$LOG" 2>&1 &
+    PIDS+=($!)
+done
+wait_workers "stage-4" "${PIDS[@]}"
+t3=$(date +%s)
+S4=$((t3 - t2)); [[ $S4 -eq 0 ]] && S4=1
+echo "[bench] stage 4: ${S4}s"
+fi
+
+# --- summary ----------------------------------------------------------------------
+{
+    echo "=== $(date -u '+%Y-%m-%d %H:%M:%S') UTC  commit=$COMMIT  gpus=$GPUS  videos=$N_VIDEOS  viz=$VIZ  fused=$FUSED  amp=$AMP  compile=$COMPILE  override=$OVERRIDE_EVERY ==="
+    awk -v s3="$S3" -v s4="$S4" -v w="$WORLD_SIZE" -v n="$N_VIDEOS" -v fused="$FUSED" 'BEGIN {
+        label = (fused == "1") ? "stage 3+4 (fused):" : "stage 3 (extract):"
+        printf "%s %5ds total  %6.1fs/video/worker  %5.1f videos/min\n", label, s3, s3*w/n, 60*n/s3
+        if (fused != "1")
+            printf "stage 4 (segment): %5ds total  %6.1fs/video/worker  %5.1f videos/min\n", s4, s4*w/n, 60*n/s4
+        printf "end-to-end:        %5ds total\n", s3+s4
+    }'
+} | tee -a "$RESULTS"
+echo "[bench] results appended to $RESULTS"

data_pipeline/build_manifest.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Build videos2caption.json + merge.txt for the final track dataset.
+
+Pairs each 121f/720p clip (--clips-dir) with its CoTracker npz (--tracks-dir) and
+caption (--captions JSON {video.mp4: caption}). Output is exactly the manifest
+`fastvideo/pipelines/preprocess/v1_preprocess.py --preprocess_task i2v_track` consumes
+(same format used for the overfit run), so the dataset is training-ready.
+"""
+from __future__ import annotations
+import argparse, json
+from pathlib import Path
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--clips-dir", required=True)
+    ap.add_argument("--tracks-dir", required=True)
+    ap.add_argument("--captions", required=True, help="JSON {video.mp4: caption}")
+    ap.add_argument("--out-dir", required=True)
+    ap.add_argument("--fps", type=int, default=24)
+    ap.add_argument("--num-frames", type=int, default=121)
+    ap.add_argument("--height", type=int, default=720)
+    ap.add_argument("--width", type=int, default=1280)
+    a = ap.parse_args()
+
+    caps = json.load(open(a.captions))
+    tracks = {p.stem: p for p in Path(a.tracks_dir).glob("*.npz")}
+    items, missing_track, missing_cap = [], 0, 0
+    for clip in sorted(Path(a.clips_dir).glob("*.mp4")):
+        tr = tracks.get(clip.stem)
+        if tr is None:
+            missing_track += 1; continue
+        cap = caps.get(clip.name, "")
+        if not cap:
+            missing_cap += 1
+        items.append({
+            "path": clip.name,
+            "cap": [cap],
+            "points_path": str(tr.resolve()),
+            "fps": float(a.fps),
+            "num_frames": a.num_frames,
+            "resolution": {"width": a.width, "height": a.height},
+        })
+    outd = Path(a.out_dir); outd.mkdir(parents=True, exist_ok=True)
+    (outd / "videos2caption.json").write_text(json.dumps(items, indent=2))
+    (outd / "merge.txt").write_text(f"{a.clips_dir},{outd / 'videos2caption.json'}")
+    print(f"[manifest] {len(items)} clips paired (clip+track+caption); "
+          f"skipped {missing_track} clips w/o track, {missing_cap} w/o caption")
+    print(f"[manifest] wrote {outd}/videos2caption.json + merge.txt")
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/cfg_ablation.py

@@ -0,0 +1,241 @@
+#!/usr/bin/env python3
+"""CFG-ablation eval on a trackwan checkpoint.
+
+Reproduces ``fastvideo/train/callbacks/track_validation.py::_sample`` denoise
+formula EXACTLY, but with configurable ``w_text`` and ``w_motion`` so we can
+isolate which CFG branch is causing artifacts.
+
+Runs each val sample under N configs; uploads all videos to one wandb run for
+side-by-side visual comparison.
+
+Usage (single-GPU, on any held alloc):
+  python data_pipeline/cfg_ablation.py \
+    --model-dir /mnt/lustre/vlm-s4duan/exports/synth_stage2_paperLR_ckpt400 \
+    --yaml examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_paperLR.yaml \
+    --val-parquet /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset \
+    --wandb-run-name cfg_abl_paperLR_ckpt400
+"""
+from __future__ import annotations
+
+import argparse
+import glob
+import os
+import sys
+import time
+from pathlib import Path
+
+import imageio.v2 as imageio
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__)))
+import trackwan_infer as twi  # noqa: E402
+from fastvideo.forward_context import set_forward_context  # noqa: E402
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (  # noqa: E402
+    FlowMatchEulerDiscreteScheduler,
+)
+
+
+def load_val_samples(parquet_dir: str, text_len: int) -> list:
+    from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v_track
+    from fastvideo.dataset.utils import collate_rows_from_parquet_schema
+    fs = sorted(glob.glob(os.path.join(parquet_dir, "**", "*.parquet"), recursive=True))
+    if not fs:
+        raise FileNotFoundError(parquet_dir)
+    rows = []
+    for f in fs:
+        rows.extend(pq.read_table(f).to_pylist())
+    batch = collate_rows_from_parquet_schema(rows, pyarrow_schema_i2v_track,
+                                             text_padding_length=int(text_len), cfg_rate=0.0)
+    infos = batch.get("info_list") or [{} for _ in rows]
+    return [{
+        "id": str(infos[i].get("id", f"clip{i:03d}")) if i < len(infos) else f"clip{i:03d}",
+        "text_embedding": batch["text_embedding"][i:i + 1].clone(),
+        "text_attention_mask": batch["text_attention_mask"][i:i + 1].clone(),
+        "first_frame_latent": batch["first_frame_latent"][i:i + 1].clone(),
+        "clip_feature": batch["clip_feature"][i:i + 1].clone(),
+        "vae_latent": batch["vae_latent"][i:i + 1].clone(),
+        "track_points": batch["track_points"][i:i + 1].clone(),
+        "track_visibility": batch["track_visibility"][i:i + 1].clone(),
+        "object_ids": batch["object_ids"][i:i + 1].clone() if "object_ids" in batch else None,
+        "track_weights": batch["track_weights"][i:i + 1].clone() if "track_weights" in batch else None,
+        "caption": str(infos[i].get("caption", "") if i < len(infos) else ""),
+    } for i in range(len(rows))]
+
+
+@torch.no_grad()
+def generate_with_cfg(model, sample: dict, *, w_text: float, w_motion: float,
+                       num_steps: int, seed: int,
+                       null_text: torch.Tensor | None = None,
+                       null_text_mask: torch.Tensor | None = None,
+                       no_track: bool = False) -> np.ndarray:
+    """Reproduce track_validation._sample denoise EXACTLY, w/ configurable CFG.
+
+    If ``null_text`` is provided, use it as the ∅ text (canonical Wan/T5 convention
+    = UMT5-encoded empty string). Otherwise fall back to zeros_like(txt).
+    If ``no_track`` is True, pass tp=None/tv=None to every forward (matches the
+    val callback's `track_val/no_track` panel).
+    """
+    device = model.device
+    dtype = torch.bfloat16
+
+    ff = sample["first_frame_latent"].to(device, dtype)
+    txt = sample["text_embedding"].to(device, dtype)
+    mask = sample["text_attention_mask"].to(device, dtype)
+    clip = sample["clip_feature"].to(device, dtype)
+    if no_track:
+        tp = None
+        tv = None
+    else:
+        # CRITICAL: sparse-sample tracks the same way training + track_validation do.
+        # Feeding the raw 2500 tracks is out-of-distribution and produces melted output.
+        tp_raw = sample["track_points"].float()
+        tv_raw = sample["track_visibility"].float()
+        oid = sample.get("object_ids")
+        tw = sample.get("track_weights")
+        aug_gen = torch.Generator(device=device).manual_seed(int(seed))
+        tp_s, tv_s = model._augment_tracks(tp_raw, tv_raw, aug_gen, object_ids=oid, track_weights=tw)
+        if tp_s is None or tv_s is None:  # motion_drop fired -> fall back to raw
+            tp_s, tv_s = tp_raw, tv_raw
+        tp = tp_s.to(device, dtype)
+        tv = tv_s.to(device, dtype)
+    cond20 = model._build_i2v_cond_concat(ff)
+
+    _, _, T, H, W = ff.shape
+    gen = torch.Generator(device="cpu").manual_seed(int(seed))
+    latents = torch.randn((1, 16, T, H, W), generator=gen, dtype=torch.float32).to(device)
+    sched = FlowMatchEulerDiscreteScheduler(shift=float(model.timestep_shift))
+    sched.set_timesteps(int(num_steps), device=device)
+
+    do_text_cfg = (w_text != 1.0)
+    do_motion_cfg = (w_motion != 1.0) and not no_track  # no tracks -> motion CFG is a no-op
+
+    def _fwd(text_e, tp_e, tv_e, mi, ts_):
+        with torch.autocast(device.type, dtype=dtype), \
+             set_forward_context(current_timestep=ts_, attn_metadata=None):
+            return model.transformer(hidden_states=mi, encoder_hidden_states=text_e,
+                                     encoder_attention_mask=mask, timestep=ts_,
+                                     encoder_hidden_states_image=clip,
+                                     track_points=tp_e, track_visibility=tv_e,
+                                     return_dict=False)
+
+    if do_text_cfg:
+        if null_text is not None:
+            # Broadcast the [seq_len, dim] null embedding to match txt shape [B, seq_len, dim]
+            nl = null_text.to(device, dtype)
+            if nl.dim() == 2:
+                nl = nl.unsqueeze(0).expand_as(txt).contiguous()
+            txt_null = nl
+        else:
+            txt_null = torch.zeros_like(txt)
+    else:
+        txt_null = None
+
+    for tt in sched.timesteps:
+        mi = torch.cat([latents.to(dtype), cond20], dim=1)
+        ts_ = tt.reshape(1).to(device, dtype)
+        v_full = _fwd(txt, tp, tv, mi, ts_)
+        if not do_text_cfg and not do_motion_cfg:
+            v = v_full
+        elif do_text_cfg and not do_motion_cfg:
+            v_no_text = _fwd(txt_null, tp, tv, mi, ts_)
+            v = v_no_text + w_text * (v_full - v_no_text)
+        elif do_motion_cfg and not do_text_cfg:
+            v_no_motion = _fwd(txt, None, None, mi, ts_)
+            v = v_no_motion + w_motion * (v_full - v_no_motion)
+        else:
+            # joint text+motion CFG (matches track_validation._sample)
+            v_no_text = _fwd(txt_null, tp, tv, mi, ts_)
+            v_no_motion = _fwd(txt, None, None, mi, ts_)
+            alpha = w_text / max(w_text + w_motion, 1e-6)
+            v_base = alpha * v_no_text + (1.0 - alpha) * v_no_motion
+            v = v_base + w_text * (v_full - v_no_text) + w_motion * (v_full - v_no_motion)
+        latents = sched.step(v.float(), tt, latents.float(), return_dict=False)[0]
+
+    return twi.decode_to_pixels(model, latents)
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--model-dir", required=True)
+    p.add_argument("--yaml", required=True)
+    p.add_argument("--val-parquet", required=True)
+    p.add_argument("--wandb-project", default="wantrack-bidir")
+    p.add_argument("--wandb-run-name", required=True)
+    p.add_argument("--out-dir", default="/mnt/lustre/vlm-s4duan/cfg_ablation")
+    p.add_argument("--steps", type=int, default=30)
+    p.add_argument("--seed", type=int, default=1000)
+    p.add_argument("--fps", type=int, default=24)
+    p.add_argument("--configs", default="3.0,1.5;1.0,1.5;1.0,1.0",
+                   help="';'-sep list of 'w_text,w_motion' pairs")
+    p.add_argument("--null-text-pt",
+                   help="Path to precomputed UMT5(``'') embedding .pt file. If given, this is "
+                        "used as the ∅ text branch instead of zeros_like(txt).")
+    p.add_argument("--no-track", action="store_true",
+                   help="Force track_points=None throughout denoise. Reproduces the val "
+                        "callback's `track_val/no_track` panel.")
+    args = p.parse_args()
+
+    null_text = None
+    null_text_mask = None
+    if args.null_text_pt:
+        blob = torch.load(args.null_text_pt, map_location="cpu")
+        null_text = blob["embedding"]
+        null_text_mask = blob.get("attention_mask")
+        print(f"[abl] loaded null-text embedding: shape={list(null_text.shape)} "
+              f"norm={null_text.norm().item():.3f}", flush=True)
+
+    import wandb
+    Path(args.out_dir).mkdir(parents=True, exist_ok=True)
+
+    print(f"[abl] loading model {args.model_dir}", flush=True)
+    model, tc = twi.load_trackwan(args.model_dir, args.yaml)
+    text_len = int(getattr(tc.data, "text_padding_length", 256))
+
+    print(f"[abl] loading val samples from {args.val_parquet}", flush=True)
+    samples = load_val_samples(args.val_parquet, text_len)
+    print(f"[abl] {len(samples)} val samples", flush=True)
+
+    configs = [tuple(float(x) for x in pair.split(","))
+               for pair in args.configs.split(";")]
+    print(f"[abl] configs: {configs}", flush=True)
+
+    run = wandb.init(project=args.wandb_project, name=args.wandb_run_name,
+                     config={"model_dir": args.model_dir, "steps": args.steps,
+                             "configs": [list(c) for c in configs]})
+
+    for i, s in enumerate(samples):
+        sid = s["id"] or f"clip{i:03d}"
+        print(f"[abl] === sample {i}: {sid} ===", flush=True)
+        wandb_log = {"sample": i, "caption": s["caption"][:120]}
+        for wt, wm in configs:
+            t0 = time.time()
+            frames = generate_with_cfg(model, s, w_text=wt, w_motion=wm,
+                                        num_steps=args.steps, seed=args.seed + i,
+                                        null_text=null_text, null_text_mask=null_text_mask,
+                                        no_track=args.no_track)
+            dt = time.time() - t0
+            tag = f"wt{wt}_wm{wm}".replace(".", "p")
+            fn = Path(args.out_dir) / f"{sid}_{tag}.mp4"
+            imageio.mimsave(str(fn), frames, fps=args.fps, macro_block_size=1)
+            wandb_log[f"gen_{tag}"] = wandb.Video(str(fn), fps=args.fps, format="mp4")
+            print(f"[abl]   {tag}: {dt:.1f}s -> {fn.name}", flush=True)
+        wandb.log(wandb_log)
+
+    # also render GT reference once
+    for i, s in enumerate(samples):
+        try:
+            ref = twi.decode_reference(model, s["vae_latent"].to(model.device, torch.bfloat16))
+            fn = Path(args.out_dir) / f"{s['id']}_gt.mp4"
+            imageio.mimsave(str(fn), ref, fps=args.fps, macro_block_size=1)
+            wandb.log({"sample": i, "reference_gt": wandb.Video(str(fn), fps=args.fps, format="mp4")})
+        except Exception as e:
+            print(f"[abl] gt render failed for sample {i}: {e}", flush=True)
+
+    wandb.finish()
+    print("[abl] DONE", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/check_gate_growth.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Report the WanTrack "gate" — patch_embedding.weight[:, 36:] — from a DCP checkpoint.
+
+The overfit stages (A/B) exist to grow this slot from EXACTLY ZERO up to roughly the scale the
+successful 1.3B merged init reached (std ~0.011), co-adapting it with track_encoder. If it is
+still ~0 there is no bootstrap happening and step C would merge nothing; if it explodes past the
+pretrained channel scale (~0.0124) the track pathway is overwhelming the image prior.
+
+Loads ONLY the tensors of interest (DCP resharding is per-tensor), so it costs ~MBs, not the
+~200GB a full checkpoint load would.
+
+Usage: python data_pipeline/check_gate_growth.py <checkpoint-dir> [more dirs...]
+"""
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+import torch
+import torch.distributed.checkpoint as dcp
+from torch.distributed.checkpoint import FileSystemReader
+
+BASE_IN = 36
+# NOTE: the live module nests the conv, so a TRAINED checkpoint stores
+# 'patch_embedding.proj.weight', while the converter writes 'patch_embedding.weight' into the
+# init safetensors. Try both — silently matching neither is how the earlier LWLR experiment
+# ended up measuring nothing.
+PE_KEYS = ["patch_embedding.proj.weight", "patch_embedding.weight"]
+KEYS = PE_KEYS + [
+    "track_encoder.proj.weight",
+    "track_encoder.temporal_conv.weight",
+]
+
+
+def _find(md_keys, want: str) -> str | None:
+    """DCP nests model tensors under a role prefix (e.g. 'student.'); match on the suffix."""
+    exact = [k for k in md_keys if k == want]
+    if exact:
+        return exact[0]
+    cands = [k for k in md_keys if k.endswith(want)]
+    return cands[0] if cands else None
+
+
+def report(ckpt: Path) -> None:
+    dcp_dir = ckpt / "dcp" if (ckpt / "dcp").is_dir() else ckpt
+    reader = FileSystemReader(str(dcp_dir))
+    md = reader.read_metadata()
+    planned = md.state_dict_metadata
+    sd = {}
+    resolved = {}
+    for want in KEYS:
+        k = _find(list(planned.keys()), want)
+        if k is None:
+            continue
+        meta = planned[k]
+        sd[k] = torch.empty(tuple(meta.size), dtype=meta.properties.dtype)
+        resolved[want] = k
+    if not sd:
+        print(f"{ckpt.name}: none of {KEYS} found in checkpoint metadata")
+        return
+    dcp.load(sd, checkpoint_id=str(dcp_dir))
+
+    out = [f"{ckpt.name:>18s}"]
+    pe_k = next((resolved[k] for k in PE_KEYS if k in resolved), None)
+    if pe_k is None:
+        out.append("!! patch_embedding NOT FOUND — cannot report gate")
+    else:
+        pe = sd[pe_k].float()
+        out.append(f"gate[:, {BASE_IN}:] std={pe[:, BASE_IN:].std():.6f}")
+        out.append(f"pretrained std={pe[:, :BASE_IN].std():.6f}")
+    for w in ("track_encoder.proj.weight", "track_encoder.temporal_conv.weight"):
+        if w in resolved:
+            out.append(f"{w.split('.')[1]}_norm={sd[resolved[w]].float().norm():.4f}")
+    print("  ".join(out))
+
+
+def main() -> None:
+    if len(sys.argv) < 2:
+        print(__doc__)
+        sys.exit(1)
+    for key, default in [("RANK", "0"), ("LOCAL_RANK", "0"), ("WORLD_SIZE", "1"),
+                         ("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", "29555")]:
+        os.environ.setdefault(key, default)
+    print(f"reference: 1.3B merged-from-overfit gate std = 0.011418 (grown from 0.000000)")
+    for p in sys.argv[1:]:
+        try:
+            report(Path(p))
+        except Exception as e:  # keep going across checkpoints
+            print(f"{Path(p).name}: ERROR {type(e).__name__}: {e}")
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/convert_droid.py

@@ -0,0 +1,237 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stage 0b (DROID): build a track-conditioning training set from DROID.
+
+Goal (per the project decision): a set where the **first frame is similar** across
+clips but the **motion is diverse**, with a single generic prompt, so the model is
+forced to read the point tracks to reduce loss (content cannot determine the output).
+
+DROID has no scene grouping in its LeRobot metadata, so "similar first frame" is
+obtained by *clustering frame-0*: we download a candidate pool, embed each first
+frame (downscaled RGB), and pick the ``--num-clips`` episodes whose first frames are
+mutually most similar (the tightest dense cluster -- the centroid that minimises the
+radius to its k-th nearest neighbour, then its k nearest).
+
+Output mirrors ``generate_videos.py`` exactly so the existing downstream
+(``extract_tracks.py`` -> ``v1_preprocess.py --preprocess_task i2v_track``) ingests it
+unchanged:
+
+    <out>/videos/vid_000000.mp4 ...        libx264, ``--num-frames`` frames @ ``--fps``
+    <out>/videos2caption.json              FastVideo manifest (one generic caption)
+    <out>/merge.txt                        "<videos_dir>,<json_path>"
+
+Source: HuggingFace LeRobot mirror ``IPEC-COMMUNITY/droid_lerobot`` (v2.0): per-episode
+av1 mp4 at ``videos/chunk-{c:03d}/{video_key}/episode_{i:06d}.mp4`` (av1 -> decoded
+with PyAV; re-encoded to libx264 so downstream decord reads it). 180x320, 15 fps.
+
+Run on a compute node (network + a little CPU; no GPU)::
+
+    srun --jobid=<job> --overlap --ntasks=1 .venv/bin/python data_pipeline/convert_droid.py \
+        --out-dir /mnt/weka/home/hao.zhang/shao/data/motion_pipeline/droid_track_200 \
+        --num-clips 200 --candidate-scan 1500
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+
+import numpy as np
+
+REPO = "IPEC-COMMUNITY/droid_lerobot"
+CHUNK_SIZE = 1000  # meta/info.json chunks_size
+DEFAULT_VIDEO_KEY = "observation.images.exterior_image_1_left"
+DEFAULT_PROMPT = "a robot arm manipulating objects on a tabletop"
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--out-dir", type=Path, required=True)
+    p.add_argument("--num-clips", type=int, default=200, help="How many episodes to keep (the cluster size).")
+    p.add_argument("--candidate-scan",
+                   type=int,
+                   default=1500,
+                   help="How many episodes to download+embed before clustering. Larger => tighter cluster.")
+    p.add_argument("--video-key", type=str, default=DEFAULT_VIDEO_KEY)
+    p.add_argument("--num-frames", type=int, default=121)
+    p.add_argument("--fps", type=int, default=24, help="Written fps == train_fps so preprocess does NOT resample.")
+    p.add_argument("--frame-start", type=int, default=0, help="First source-frame index of the kept window.")
+    p.add_argument("--prompt", type=str, default=DEFAULT_PROMPT)
+    p.add_argument("--embed-size", type=int, default=64, help="Frame-0 is resized to NxN for the similarity embedding.")
+    p.add_argument("--no-cluster", action="store_true", help="Skip clustering; keep the first --num-clips candidates.")
+    p.add_argument("--download-workers", type=int, default=16)
+    p.add_argument("--seed", type=int, default=0)
+    return p.parse_args()
+
+
+def list_candidate_episodes(num_frames: int, scan: int) -> list[int]:
+    """Episode indices (length >= num_frames), first ``scan`` in index order."""
+    from huggingface_hub import hf_hub_download
+    ep_path = hf_hub_download(REPO, "meta/episodes.jsonl", repo_type="dataset")
+    out: list[int] = []
+    with open(ep_path) as f:
+        for line in f:
+            d = json.loads(line)
+            if int(d.get("length", 0)) >= num_frames:
+                out.append(int(d["episode_index"]))
+                if len(out) >= scan:
+                    break
+    return out
+
+
+def episode_video_path(ep_idx: int, video_key: str) -> str:
+    chunk = ep_idx // CHUNK_SIZE
+    return f"videos/chunk-{chunk:03d}/{video_key}/episode_{ep_idx:06d}.mp4"
+
+
+def download_one(ep_idx: int, video_key: str) -> tuple[int, str | None]:
+    from huggingface_hub import hf_hub_download
+    try:
+        p = hf_hub_download(REPO, episode_video_path(ep_idx, video_key), repo_type="dataset")
+        return ep_idx, p
+    except Exception:  # noqa: BLE001
+        return ep_idx, None
+
+
+def decode_frames(path: str, start: int, count: int) -> np.ndarray | None:
+    """Decode ``count`` frames starting at ``start`` (RGB uint8 [T,H,W,3]) via PyAV (handles av1)."""
+    import av
+    try:
+        container = av.open(path)
+        frames = []
+        for idx, frame in enumerate(container.decode(video=0)):
+            if idx >= start:
+                frames.append(frame.to_ndarray(format="rgb24"))
+                if len(frames) >= count:
+                    break
+        container.close()
+    except Exception:  # noqa: BLE001
+        return None
+    if len(frames) < count:
+        return None
+    return np.stack(frames, axis=0)
+
+
+def embed_frame0(path: str, size: int) -> np.ndarray | None:
+    from PIL import Image
+    fr = decode_frames(path, 0, 1)
+    if fr is None:
+        return None
+    img = Image.fromarray(fr[0]).resize((size, size), Image.BILINEAR)
+    v = np.asarray(img, dtype=np.float32).reshape(-1)
+    n = np.linalg.norm(v) + 1e-8
+    return v / n
+
+
+def pick_tightest_cluster(embs: np.ndarray, k: int) -> np.ndarray:
+    """Return indices of the k mutually-most-similar embeddings.
+
+    Centroid = the point minimising the distance to its k-th nearest neighbour
+    (densest region); then take that point's k nearest neighbours.
+    """
+    # cosine distance (embs are L2-normalised) -> 1 - sim
+    sim = embs @ embs.T
+    dist = 1.0 - sim
+    np.fill_diagonal(dist, 0.0)
+    part = np.partition(dist, kth=min(k, dist.shape[0] - 1), axis=1)
+    kth = part[:, min(k, dist.shape[0] - 1)]  # radius to k-th NN per row
+    centroid = int(np.argmin(kth))
+    order = np.argsort(dist[centroid])[:k]
+    return order
+
+
+def main() -> None:
+    args = parse_args()
+    out_dir = args.out_dir
+    videos_dir = out_dir / "videos"
+    videos_dir.mkdir(parents=True, exist_ok=True)
+    json_path = out_dir / "videos2caption.json"
+    merge_path = out_dir / "merge.txt"
+
+    print(f"[droid] listing candidates (length>={args.num_frames}, scan={args.candidate_scan})...", flush=True)
+    candidates = list_candidate_episodes(args.num_frames, args.candidate_scan)
+    print(f"[droid] {len(candidates)} candidate episodes", flush=True)
+
+    # Download (cached) + embed frame-0 in parallel.
+    print(f"[droid] downloading + embedding frame-0 ({args.download_workers} workers)...", flush=True)
+    ep_to_path: dict[int, str] = {}
+    with ThreadPoolExecutor(max_workers=args.download_workers) as ex:
+        futs = {ex.submit(download_one, ep, args.video_key): ep for ep in candidates}
+        for i, fut in enumerate(as_completed(futs), 1):
+            ep, path = fut.result()
+            if path is not None:
+                ep_to_path[ep] = path
+            if i % 100 == 0:
+                print(f"[droid]   downloaded {i}/{len(candidates)}", flush=True)
+    print(f"[droid] downloaded {len(ep_to_path)} mp4s; embedding...", flush=True)
+
+    eps_ok: list[int] = []
+    embs: list[np.ndarray] = []
+    for ep in candidates:
+        if ep not in ep_to_path:
+            continue
+        e = embed_frame0(ep_to_path[ep], args.embed_size)
+        if e is not None:
+            eps_ok.append(ep)
+            embs.append(e)
+    embs_arr = np.stack(embs, axis=0)
+    print(f"[droid] embedded {len(eps_ok)} frame-0s", flush=True)
+
+    if args.no_cluster or len(eps_ok) <= args.num_clips:
+        sel_local = list(range(min(args.num_clips, len(eps_ok))))
+        print(f"[droid] no clustering -> first {len(sel_local)} candidates", flush=True)
+    else:
+        sel_local = pick_tightest_cluster(embs_arr, args.num_clips).tolist()
+        # report cluster tightness
+        sub = embs_arr[sel_local]
+        c = sub.mean(0, keepdims=True)
+        c /= np.linalg.norm(c) + 1e-8
+        mean_sim = float((sub @ c.T).mean())
+        print(f"[droid] tightest-cluster of {len(sel_local)} clips; mean cos-sim to centroid={mean_sim:.3f}",
+              flush=True)
+
+    selected_eps = [eps_ok[i] for i in sel_local]
+
+    # Decode + write libx264 mp4s + manifest.
+    import imageio.v2 as imageio
+    records = []
+    seq = 0
+    for ep in selected_eps:
+        frames = decode_frames(ep_to_path[ep], args.frame_start, args.num_frames)
+        if frames is None:
+            print(f"[droid]   skip ep{ep}: decode<{args.num_frames} frames", flush=True)
+            continue
+        H, W = int(frames.shape[1]), int(frames.shape[2])
+        name = f"vid_{seq:06d}.mp4"
+        imageio.mimsave(str(videos_dir / name),
+                        list(frames),
+                        fps=args.fps,
+                        codec="libx264",
+                        macro_block_size=1,
+                        output_params=["-pix_fmt", "yuv420p"])
+        records.append({
+            "idx": seq,
+            "path": name,
+            "cap": [args.prompt],
+            "fps": float(args.fps),
+            "duration": float(args.num_frames) / float(args.fps),
+            "num_frames": int(args.num_frames),
+            "resolution": {
+                "width": W,
+                "height": H
+            },
+            "source_episode": int(ep),
+        })
+        seq += 1
+        if seq % 25 == 0:
+            print(f"[droid]   wrote {seq}/{len(selected_eps)}", flush=True)
+
+    json_path.write_text(json.dumps(records, indent=2))
+    merge_path.write_text(f"{videos_dir.resolve()},{json_path.resolve()}\n")
+    print(f"[droid] DONE: wrote {len(records)} clips -> {videos_dir}", flush=True)
+    print(f"[droid] manifest -> {json_path}", flush=True)
+    print(f"[droid] next: extract_tracks.py --data-dir {out_dir} ; then v1_preprocess i2v_track", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/convert_parquets_to_bf16.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Convert openvid-wantrack parquets to bf16 (a COPY; never mutates the source).
+
+Casts the big float tensor fields (vae_latent, first_frame_latent, text_embedding,
+clip_feature, track_points, track_visibility) from float32 to bfloat16 and tags their
+``*_dtype`` column ``"bfloat16"``. Integer / tiny fields (object_ids, track_weights) and all
+metadata are copied unchanged. Training already downcasts these fields to bf16
+(``wan.py`` / ``wantrack.py``), so the quality loss is negligible while the files ~halve.
+
+Requires the loader change in ``fastvideo/dataset/utils.py`` that honors the ``*_dtype``
+column (both decoders). Without it, the bf16 bytes would be misread as float32.
+
+The output mirrors the source tree (shardNNN/combined_parquet_dataset/worker_N/*.parquet), so
+point training ``data_path`` at ``--dst`` once you've converted what you want. Resumable:
+already-written files are skipped. CPU only; run on a compute node (I/O + RAM heavy).
+
+Examples:
+  # first half of all parquets -> a bf16 sibling dir
+  python data_pipeline/convert_parquets_to_bf16.py --fraction 0.5
+
+  # a fixed number of files, dry-run first
+  python data_pipeline/convert_parquets_to_bf16.py --limit 100 --dry-run
+"""
+from __future__ import annotations
+
+import argparse
+import os
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+import torch
+
+DEFAULT_SRC = "/home/hal-shared/motionstream/data/openvid-wantrack-parquets"
+DEFAULT_DST = "/home/hal-shared/motionstream/data/openvid-wantrack-parquets-bf16"
+
+# Big float fields consumed at bf16 by training -> safe to store bf16.
+DEFAULT_BF16_FIELDS = [
+    "vae_latent", "first_frame_latent", "text_embedding", "clip_feature",
+    "track_points", "track_visibility",
+]
+# Left untouched (integer labels / tiny): object_ids, track_weights, all metadata.
+
+_SRC_STR_TO_TORCH = {
+    "float32": torch.float32, "float16": torch.float16, "float64": torch.float64,
+}
+
+
+def _to_bf16_bytes(b: bytes, dtype_str: str) -> tuple[bytes, str]:
+    """Re-encode a raw float tensor blob as bf16. Returns (bytes, dtype_label)."""
+    if not b:                         # empty optional field -> leave as-is
+        return b, (dtype_str or "")
+    if dtype_str == "bfloat16":       # already converted
+        return b, dtype_str
+    src = _SRC_STR_TO_TORCH.get(dtype_str or "float32")
+    if src is None:
+        raise ValueError(f"cannot convert stored dtype {dtype_str!r} to bf16")
+    t = torch.frombuffer(bytearray(b), dtype=src).to(torch.bfloat16)
+    return t.view(torch.uint8).numpy().tobytes(), "bfloat16"
+
+
+def convert_file(src_path: str, dst_path: str, fields: list[str]) -> tuple[int, int]:
+    """Convert one parquet file. Returns (src_bytes, dst_bytes)."""
+    tbl = pq.read_table(src_path)
+    names = list(tbl.schema.names)
+    cols: dict[str, object] = {n: tbl.column(n) for n in names}
+
+    for fld in fields:
+        bkey, dkey = f"{fld}_bytes", f"{fld}_dtype"
+        if bkey not in cols or dkey not in cols:
+            continue
+        b_list = tbl.column(bkey).to_pylist()
+        d_list = tbl.column(dkey).to_pylist()
+        new_b, new_d = [], []
+        for b, d in zip(b_list, d_list, strict=True):
+            nb, nd = _to_bf16_bytes(b, d)
+            new_b.append(nb)
+            new_d.append(nd)
+        cols[bkey] = pa.array(new_b, type=pa.binary())
+        cols[dkey] = pa.array(new_d, type=pa.string())
+
+    out = pa.table([cols[n] for n in names], schema=tbl.schema)
+    os.makedirs(os.path.dirname(dst_path), exist_ok=True)
+    tmp = dst_path + ".tmp"
+    pq.write_table(out, tmp)
+    os.replace(tmp, dst_path)
+    return os.path.getsize(src_path), os.path.getsize(dst_path)
+
+
+def verify_file(src_path: str, dst_path: str, fields: list[str]) -> None:
+    """Round-trip check: one bf16 field on row 0 must equal src fp32 -> bf16."""
+    s = pq.ParquetFile(src_path).read_row_group(0).slice(0, 1).to_pylist()[0]
+    d = pq.ParquetFile(dst_path).read_row_group(0).slice(0, 1).to_pylist()[0]
+    for fld in fields:
+        sb, db = s.get(f"{fld}_bytes"), d.get(f"{fld}_bytes")
+        if not sb:
+            continue
+        assert d.get(f"{fld}_dtype") == "bfloat16", f"{fld}: dtype not tagged bfloat16"
+        src_t = torch.frombuffer(bytearray(sb), dtype=_SRC_STR_TO_TORCH[s[f"{fld}_dtype"]]).to(torch.bfloat16)
+        dst_t = torch.frombuffer(bytearray(db), dtype=torch.bfloat16)
+        assert torch.equal(dst_t, src_t), f"{fld}: bf16 round-trip mismatch"
+        return  # one field is enough
+    return
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--src", default=DEFAULT_SRC)
+    p.add_argument("--dst", default=DEFAULT_DST)
+    p.add_argument("--fraction", type=float, default=0.5, help="fraction of the sorted parquet list to convert (default first 0.5)")
+    p.add_argument("--limit", type=int, default=None, help="convert at most N files (overrides --fraction)")
+    p.add_argument("--offset", type=int, default=0, help="skip the first N files of the sorted list (parallelize by running disjoint --offset/--limit ranges)")
+    p.add_argument("--fields", default=",".join(DEFAULT_BF16_FIELDS), help="comma-separated fields to cast to bf16")
+    p.add_argument("--overwrite", action="store_true", help="re-convert files already present in --dst")
+    p.add_argument("--no-verify", action="store_true", help="skip the per-file round-trip check")
+    p.add_argument("--dry-run", action="store_true", help="list what would be converted; write nothing")
+    args = p.parse_args()
+
+    src_root = os.path.realpath(args.src)
+    dst_root = os.path.realpath(args.dst)
+    fields = [f.strip() for f in args.fields.split(",") if f.strip()]
+    if os.path.commonpath([src_root, dst_root]) == src_root and dst_root != src_root:
+        raise SystemExit(f"--dst {dst_root} is inside --src; choose a separate directory")
+    if src_root == dst_root:
+        raise SystemExit("refusing to convert in place; --dst must differ from --src")
+
+    all_files = []
+    for root, _, files in os.walk(src_root):
+        for f in files:
+            if f.endswith(".parquet"):
+                all_files.append(os.path.join(root, f))
+    all_files.sort()
+    n_total = len(all_files)
+    n_take = args.limit if args.limit is not None else int(args.fraction * n_total)
+    selected = all_files[args.offset:args.offset + n_take]
+    print(f"[bf16] {n_total} parquet(s) found; converting {len(selected)} "
+          f"[offset {args.offset}, {'limit ' + str(args.limit) if args.limit is not None else f'fraction {args.fraction}'}]")
+    print(f"[bf16] fields -> bf16: {fields}")
+    print(f"[bf16] src={src_root}\n[bf16] dst={dst_root}")
+
+    if args.dry_run:
+        for f in selected[:5]:
+            print("  would convert:", os.path.relpath(f, src_root))
+        if len(selected) > 5:
+            print(f"  ... and {len(selected) - 5} more")
+        return
+
+    src_tot = dst_tot = done = skipped = 0
+    for i, sp in enumerate(selected, 1):
+        rel = os.path.relpath(sp, src_root)
+        dp = os.path.join(dst_root, rel)
+        if os.path.exists(dp) and not args.overwrite:
+            skipped += 1
+            continue
+        sb, db = convert_file(sp, dp, fields)
+        if not args.no_verify:
+            verify_file(sp, dp, fields)
+        src_tot += sb
+        dst_tot += db
+        done += 1
+        if done % 20 == 0 or i == len(selected):
+            gb = 1024 ** 3
+            print(f"[bf16] {i}/{len(selected)} | converted {done}, skipped {skipped} | "
+                  f"{src_tot/gb:.1f}GB -> {dst_tot/gb:.1f}GB"
+                  f"{f' ({dst_tot/src_tot*100:.0f}%)' if src_tot else ''}")
+
+    print(f"[bf16] done: converted {done}, skipped {skipped}. Output tree at {dst_root}")
+    print(f"[bf16] point training data_path at {dst_root} once you've converted enough shards.")
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/convert_trackwan_init.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Build a WanTrack init checkpoint from a base Wan diffusers model.
+
+The WanTrack DiT widens the patch-embed input to 52 channels and adds a
+``track_encoder``. Works from either base:
+  - a Wan **I2V** model (e.g. Wan2.1-Fun-1.3B-InP: in_channels=36 = 16 noisy +
+    4 mask + 16 first-frame, with CLIP image cross-attention): the 36 pretrained
+    image-conditioning channels are kept and only the 16 track channels are
+    zero-init, so first-frame conditioning works from step 0.  **Recommended.**
+  - a Wan **T2V** model (in_channels=16): the 20 I2V + 16 track channels are all
+    zero-init, so I2V must be learned from scratch (slower, weaker).
+
+It produces a diffusers transformer dir whose weights load *strictly* into
+``TrackWanTransformer3DModel``:
+  - ``patch_embedding.weight`` zero-padded base_in -> 52 input channels (pretrained
+    weights occupy the first base_in channels; the new track channels start at
+    zero, so a freshly converted model reproduces the base at step 0),
+  - ``track_encoder.*`` added (proj normal-init; the patch-embed track channels are the
+    single zero-conv -> zero track contribution at step 0, but gradient still flows so the
+    track pathway can learn -- see build_track_encoder_state for the deadlock rationale),
+  - ``config.json`` gets ``in_channels=52`` + ``track_config`` (CLIP ``image_dim``
+    is inherited from the base config when present).
+
+Other pipeline components (vae / text_encoder / tokenizer / scheduler /
+model_index.json) are symlinked so the output is a complete, loadable model dir.
+
+Usage (no GPU / no fastvideo import needed):
+    python data_pipeline/convert_trackwan_init.py \
+        --base <base diffusers model dir> --out <output dir>
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file, save_file
+from torch import nn
+
+NEW_IN_CHANNELS = 52
+TRACK_CHANNELS = 16
+ID_DIM = 128
+VAE_T_COMP = 4
+TRACK_CONFIG = {
+    "id_dim": ID_DIM,
+    "track_channels": TRACK_CHANNELS,
+    "vae_spatial_compression": 8,
+    "vae_temporal_compression": VAE_T_COMP,
+    "max_track_id": 100_000,
+    # The single zero-conv is the patch-embed track channels (zero-padded below), NOT the
+    # track head -> proj is normal-init so the track pathway actually receives gradient.
+    "zero_init_head": False,
+}
+
+
+def build_track_encoder_state() -> dict[str, torch.Tensor]:
+    """Match TrackEncoder's params: temporal_conv + proj, BOTH normal-initialized.
+
+    IMPORTANT (deadlock fix): the track signal passes through two layers in series --
+    ``track_encoder.proj`` then the patch-embed track channels [36:52]. The single
+    ControlNet-style zero-conv is the *patch-embed track channels* (kept at 0 by the
+    zero-pad in main()), which already guarantees zero track contribution at step 0
+    (teacher behavior). ``proj`` must therefore be NON-zero so gradient can reach the
+    patch-embed track channels (grad ∝ proj output): if BOTH were zero, each layer's
+    gradient is gated by the other being nonzero -> both stay exactly 0 forever and the
+    track pathway never learns (observed: proj & patch-embed[36:52] frozen at 0.0 after
+    4000 steps). So leave proj at its default Conv init here."""
+    # bias=False MUST match the model (fastvideo/models/dits/trackwan/track_encoder.py): TrackEncoder
+    # omits the conv bias (a bias broadcasts to every latent cell and densifies the sparse track
+    # signal -- load-bearing). Emitting bias tensors breaks strict loading with
+    # "track_encoder.proj.bias not found in custom model state dict", so build bias-free and write
+    # only the two weight tensors the model actually has.
+    temporal_conv = nn.Conv3d(ID_DIM, TRACK_CHANNELS, kernel_size=(VAE_T_COMP, 1, 1),
+                              stride=(VAE_T_COMP, 1, 1), bias=False)
+    proj = nn.Conv3d(TRACK_CHANNELS, TRACK_CHANNELS, kernel_size=1, bias=False)  # default init (NOT zero) -> breaks the deadlock
+    return {
+        "track_encoder.temporal_conv.weight": temporal_conv.weight.detach().clone(),
+        "track_encoder.proj.weight": proj.weight.detach().clone(),
+    }
+
+
+def main() -> None:
+    global ID_DIM
+    p = argparse.ArgumentParser()
+    p.add_argument("--base",
+                   required=True,
+                   help="Base Wan I2V (e.g. Wan2.1-Fun-1.3B-InP) or T2V diffusers model dir (with transformer/).")
+    p.add_argument("--out", required=True, help="Output model dir for the WanTrack init.")
+    p.add_argument("--id-dim", type=int, default=ID_DIM,
+                   help="sinusoidal track-id posemb dim (MotionStream d; 64 for d64 init, default 128)")
+    args = p.parse_args()
+
+    ID_DIM = args.id_dim
+    TRACK_CONFIG["id_dim"] = args.id_dim
+
+    base = Path(args.base)
+    out = Path(args.out)
+    if not (base / "transformer" / "config.json").exists():
+        raise FileNotFoundError(f"{base}/transformer/config.json not found")
+    out.mkdir(parents=True, exist_ok=True)
+
+    # 1) Symlink every top-level entry except transformer/ (vae, text_encoder, ...).
+    for entry in os.listdir(base):
+        if entry == "transformer":
+            continue
+        src = (base / entry).resolve()
+        dst = out / entry
+        if dst.is_symlink() or dst.exists():
+            if dst.is_dir() and not dst.is_symlink():
+                shutil.rmtree(dst)
+            else:
+                dst.unlink()
+        os.symlink(src, dst)
+
+    # 2) Convert transformer/.
+    tdir = out / "transformer"
+    tdir.mkdir(exist_ok=True)
+    cfg = json.loads((base / "transformer" / "config.json").read_text())
+    base_in = int(cfg["in_channels"])
+    cfg["in_channels"] = NEW_IN_CHANNELS
+    cfg["track_config"] = TRACK_CONFIG
+    (tdir / "config.json").write_text(json.dumps(cfg, indent=2))
+
+    sf_files = sorted((base / "transformer").glob("*.safetensors"))
+    if not sf_files:
+        raise FileNotFoundError(f"No transformer safetensors under {base}/transformer")
+    state: dict[str, torch.Tensor] = {}
+    for sf in sf_files:
+        state.update(load_file(str(sf)))  # merge shards if the base is sharded
+
+    pe_key = "patch_embedding.weight"
+    if pe_key not in state:
+        cands = [k for k in state if "patch_embed" in k and k.endswith(".weight")]
+        if len(cands) != 1:
+            raise KeyError(f"Could not find patch_embedding weight; candidates={cands}")
+        pe_key = cands[0]
+    w = state[pe_key]  # [out, base_in, 1, 2, 2]
+    if w.shape[1] != base_in:
+        raise ValueError(f"{pe_key} in_ch {w.shape[1]} != config in_channels {base_in}")
+    new_w = torch.zeros((w.shape[0], NEW_IN_CHANNELS, *w.shape[2:]), dtype=w.dtype)
+    new_w[:, :base_in] = w  # pretrained channels first; new channels zero
+    state[pe_key] = new_w
+    print(f"[convert] padded {pe_key}: {tuple(w.shape)} -> {tuple(new_w.shape)}")
+
+    te_state = build_track_encoder_state()
+    for k, v in te_state.items():
+        state[k] = v.to(w.dtype)
+    print(f"[convert] added {len(te_state)} track_encoder tensors")
+
+    save_file(state, str(tdir / "diffusion_pytorch_model.safetensors"), metadata={"format": "pt"})
+    print(f"[convert] wrote {tdir/'diffusion_pytorch_model.safetensors'} ({len(state)} keys)")
+    print(f"[convert] done -> {out}")
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/convert_trackwan_init_v2.py

@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Build a WanTrack init checkpoint from a base Wan diffusers model, with several
+strategies for the patch_embedding track-channel slot and the track_encoder.
+
+Modes for patch_embedding[:, base_in:] (the new 16 in-channels):
+  --pe-init zero    : zero-pad (original stage-1 recipe)
+  --pe-init random  : normal-init with std matched to base first-N channels
+
+Modes for track_encoder.{proj, temporal_conv}:
+  --track-src default            : default Conv init (used by original convert)
+  --track-src <ckpt.safetensors> : lift both convs from an existing WanTrack safetensors
+                                   (e.g. a trained 1.3B ckpt). Bias is copied when present
+                                   and its shape matches.
+
+Usage:
+  python data_pipeline/convert_trackwan_init_v2.py \
+      --base /mnt/lustre/vlm-s4duan/models/Wan2.1-I2V-14B-720P-Diffusers \
+      --out  /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_random_init \
+      --id-dim 64 --pe-init random
+
+  python data_pipeline/convert_trackwan_init_v2.py \
+      --base /mnt/lustre/vlm-s4duan/models/Wan2.1-I2V-14B-720P-Diffusers \
+      --out  /mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_partial_merged \
+      --id-dim 64 --pe-init random \
+      --track-src /mnt/lustre/vlm-s4duan/exports/synth_stage2_paperLR_ckpt600/transformer/model.safetensors
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file, save_file
+from torch import nn
+
+NEW_IN_CHANNELS = 52
+TRACK_CHANNELS = 16
+VAE_T_COMP = 4
+DEFAULT_ID_DIM = 128
+
+
+def build_track_encoder_default(id_dim: int, dtype: torch.dtype, use_bias: bool = False) -> dict[str, torch.Tensor]:
+    """Default (non-zero) Conv init for track_encoder.proj + temporal_conv."""
+    temporal_conv = nn.Conv3d(id_dim, TRACK_CHANNELS,
+                              kernel_size=(VAE_T_COMP, 1, 1), stride=(VAE_T_COMP, 1, 1),
+                              bias=use_bias)
+    proj = nn.Conv3d(TRACK_CHANNELS, TRACK_CHANNELS, kernel_size=1, bias=use_bias)
+    d = {
+        "track_encoder.temporal_conv.weight": temporal_conv.weight.detach().to(dtype).clone(),
+        "track_encoder.proj.weight": proj.weight.detach().to(dtype).clone(),
+    }
+    if use_bias:
+        d["track_encoder.temporal_conv.bias"] = temporal_conv.bias.detach().to(dtype).clone()
+        d["track_encoder.proj.bias"] = proj.bias.detach().to(dtype).clone()
+    return d
+
+
+def load_src_state(src_path: str) -> dict[str, torch.Tensor]:
+    """Load a source state dict from a .safetensors FILE or a diffusers model DIR.
+
+    A 14B export is sharded across several *.safetensors, so accept a directory (either the
+    model root or its transformer/ subdir) and merge every shard.
+    """
+    p = Path(src_path)
+    if p.is_file():
+        return load_file(str(p))
+    tdir = p / "transformer" if (p / "transformer").is_dir() else p
+    shards = sorted(tdir.glob("*.safetensors"))
+    if not shards:
+        raise FileNotFoundError(f"no *.safetensors under {tdir}")
+    state: dict[str, torch.Tensor] = {}
+    for s in shards:
+        state.update(load_file(str(s)))
+    print(f"[src] loaded {len(state)} tensors from {len(shards)} shard(s) in {tdir}")
+    return state
+
+
+def lift_track_encoder_from_src(src_path: str, id_dim: int, dtype: torch.dtype) -> dict[str, torch.Tensor]:
+    """Load track_encoder.{proj, temporal_conv}.{weight, bias} from an existing WanTrack ckpt.
+
+    Verifies shape compatibility. If bias is missing in the source we skip it; if the
+    source shapes don't match the expected shapes we fall back to default init for that
+    key and print a warning.
+    """
+    src_state = load_src_state(src_path)
+    expected_shapes = {
+        "track_encoder.temporal_conv.weight": (TRACK_CHANNELS, id_dim, VAE_T_COMP, 1, 1),
+        "track_encoder.proj.weight":          (TRACK_CHANNELS, TRACK_CHANNELS, 1, 1, 1),
+        "track_encoder.temporal_conv.bias":   (TRACK_CHANNELS,),
+        "track_encoder.proj.bias":            (TRACK_CHANNELS,),
+    }
+    picked: dict[str, torch.Tensor] = {}
+    for k, want in expected_shapes.items():
+        if k not in src_state:
+            continue
+        got = tuple(src_state[k].shape)
+        if got != want:
+            print(f"[track-src] SHAPE MISMATCH on {k}: src={got} vs expected={want}; SKIPPING (will need fallback)")
+            continue
+        picked[k] = src_state[k].detach().to(dtype).clone()
+        print(f"[track-src] lifted {k} shape={got}")
+    # If either weight is missing, we still need to keep the model loadable.
+    fallback = build_track_encoder_default(id_dim, dtype, use_bias=False)
+    for k, v in fallback.items():
+        if k not in picked:
+            print(f"[track-src] {k} not in source -> default init")
+            picked[k] = v
+    return picked
+
+
+def std_of_first_n(w: torch.Tensor, n: int) -> float:
+    """Std of the pretrained first-N in-channels slice of patch_embedding.weight."""
+    with torch.no_grad():
+        s = w[:, :n].float().std().item()
+    return float(s)
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--base", required=True, help="Base Wan diffusers model dir.")
+    p.add_argument("--out",  required=True, help="Output diffusers model dir.")
+    p.add_argument("--id-dim", type=int, default=DEFAULT_ID_DIM)
+    p.add_argument("--pe-init", choices=("zero", "random"), default="zero",
+                   help="patch_embedding init strategy for the ADDED track channels [:, base_in:]. "
+                        "Base pretrained channels [:, :base_in] are always preserved.")
+    p.add_argument("--pe-random-seed", type=int, default=1234)
+    p.add_argument("--track-src", default=None,
+                   help="Path to a .safetensors (or diffusers model dir) with track_encoder.* to lift.")
+    p.add_argument("--pe-src", default=None,
+                   help="Path to a .safetensors (or diffusers model dir) to lift patch_embedding.weight[:, base_in:] "
+                        "from — i.e. the TRACK SLOT only; the pretrained [:, :base_in] channels always come from "
+                        "--base. Use together with --track-src pointing at the SAME source so the encoder and the "
+                        "track slot stay CO-ADAPTED (lifting the encoder alone is worse than random init).")
+    p.add_argument("--use-bias-defaults", action="store_true",
+                   help="When --track-src is not set, build track_encoder convs WITH bias (matches TRACKWAN_TRACK_BIAS=1 training).")
+    args = p.parse_args()
+
+    id_dim = int(args.id_dim)
+    TRACK_CONFIG = {
+        "id_dim": id_dim,
+        "track_channels": TRACK_CHANNELS,
+        "vae_spatial_compression": 8,
+        "vae_temporal_compression": VAE_T_COMP,
+        "max_track_id": 100_000,
+        "zero_init_head": False,
+    }
+
+    base = Path(args.base)
+    out = Path(args.out)
+    if not (base / "transformer" / "config.json").exists():
+        raise FileNotFoundError(f"{base}/transformer/config.json not found")
+    out.mkdir(parents=True, exist_ok=True)
+
+    # 1) Symlink every top-level entry except transformer/
+    for entry in os.listdir(base):
+        if entry == "transformer":
+            continue
+        src_p = (base / entry).resolve()
+        dst = out / entry
+        if dst.is_symlink() or dst.exists():
+            if dst.is_dir() and not dst.is_symlink():
+                shutil.rmtree(dst)
+            else:
+                dst.unlink()
+        os.symlink(src_p, dst)
+
+    # 2) Rewrite transformer/config.json
+    tdir = out / "transformer"
+    tdir.mkdir(exist_ok=True)
+    cfg = json.loads((base / "transformer" / "config.json").read_text())
+    base_in = int(cfg["in_channels"])
+    cfg["in_channels"] = NEW_IN_CHANNELS
+    cfg["track_config"] = TRACK_CONFIG
+    (tdir / "config.json").write_text(json.dumps(cfg, indent=2))
+
+    # 3) Load full base transformer state
+    sf_files = sorted((base / "transformer").glob("*.safetensors"))
+    if not sf_files:
+        raise FileNotFoundError(f"No safetensors under {base}/transformer")
+    state: dict[str, torch.Tensor] = {}
+    for sf in sf_files:
+        state.update(load_file(str(sf)))
+    dtype = next(iter(state.values())).dtype
+
+    # 4) Grow patch_embedding.weight from base_in -> 52 input channels
+    pe_key = "patch_embedding.weight"
+    if pe_key not in state:
+        raise KeyError(f"{pe_key} not found in base transformer")
+    w = state[pe_key]  # [hidden, base_in, 1, 2, 2]
+    if w.shape[1] != base_in:
+        raise ValueError(f"{pe_key} in_ch {w.shape[1]} != config in_channels {base_in}")
+    new_w = torch.empty((w.shape[0], NEW_IN_CHANNELS, *w.shape[2:]), dtype=w.dtype)
+    new_w[:, :base_in] = w  # pretrained channels first
+    if args.pe_init == "zero":
+        new_w[:, base_in:] = 0
+        print(f"[convert] pe-init=zero -> new channels are zero")
+    else:
+        std = std_of_first_n(w, base_in)
+        g = torch.Generator().manual_seed(int(args.pe_random_seed))
+        noise = torch.empty_like(new_w[:, base_in:], dtype=torch.float32)
+        noise.normal_(mean=0.0, std=std, generator=g)
+        new_w[:, base_in:] = noise.to(w.dtype)
+        print(f"[convert] pe-init=random N(0, {std:.5f}) matching first-{base_in} std")
+    # 4b) Optionally overwrite ONLY the track slot [:, base_in:] from a trained source.
+    # This is the other half of the "merged" recipe: the track slot and the track_encoder were
+    # co-adapted during the overfit, so they must be lifted TOGETHER (--pe-src + --track-src from
+    # the same ckpt). Lifting the encoder alone leaves it talking to a random projection.
+    if args.pe_src:
+        src_pe_state = load_src_state(args.pe_src)
+        if pe_key not in src_pe_state:
+            raise KeyError(f"{pe_key} not found in --pe-src {args.pe_src}")
+        src_pe = src_pe_state[pe_key]
+        want = (w.shape[0], NEW_IN_CHANNELS, *w.shape[2:])
+        if tuple(src_pe.shape) != want:
+            raise ValueError(f"--pe-src {pe_key} shape {tuple(src_pe.shape)} != expected {want}")
+        slot = src_pe[:, base_in:].to(w.dtype)
+        new_w[:, base_in:] = slot
+        print(f"[pe-src] lifted {pe_key}[:, {base_in}:] from source "
+              f"(std={slot.float().std().item():.5f}, absmax={slot.float().abs().max().item():.5f})")
+
+    state[pe_key] = new_w
+    print(f"[convert] patch_embedding.weight: {tuple(w.shape)} -> {tuple(new_w.shape)}")
+    with torch.no_grad():
+        print(f"[convert] pe[:, :{base_in}] std={new_w[:, :base_in].float().std().item():.5f} (pretrained) | "
+              f"pe[:, {base_in}:] std={new_w[:, base_in:].float().std().item():.5f} (track slot)")
+
+    # 5) Add track_encoder
+    if args.track_src:
+        te = lift_track_encoder_from_src(args.track_src, id_dim, dtype)
+    else:
+        te = build_track_encoder_default(id_dim, dtype, use_bias=args.use_bias_defaults)
+    for k, v in te.items():
+        state[k] = v
+    print(f"[convert] added {len(te)} track_encoder tensors")
+
+    save_file(state, str(tdir / "diffusion_pytorch_model.safetensors"), metadata={"format": "pt"})
+    print(f"[convert] wrote {tdir/'diffusion_pytorch_model.safetensors'} ({len(state)} keys)")
+    print(f"[convert] done -> {out}")
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/decode_roundtrip_videos.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stage 2: produce VAE round-trip videos for CoTracker.
+
+Encode each source video through the FastVideo WanVAE (use_feature_cache=False,
+with causal-boundary fix) then immediately decode it back. The resulting videos
+differ slightly from the originals (compression artifacts, mild color shift) but
+are exactly what Stage 5 (preprocess_to_parquet) will store as latents and what
+the validation callback will decode for reference. CoTracker tracks extracted from
+these round-trip videos therefore align with training latents and the validation
+reference display.
+
+Usage:
+    python data_pipeline/decode_roundtrip_videos.py \\
+        --data-dir /home/hal-kevin/data/motion-physics \\
+        --vae-path /home/hal-kevin/models/trackwan_1.3b_i2v_control_init/vae
+
+    # Re-run specific indices only
+    python data_pipeline/decode_roundtrip_videos.py \\
+        --data-dir /home/hal-kevin/data/motion-physics \\
+        --vae-path /home/hal-kevin/models/trackwan_1.3b_i2v_control_init/vae \\
+        --index 4 7 12
+"""
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import imageio
+import numpy as np
+import torch
+from safetensors.torch import load_file as safetensors_load_file
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from fastvideo.configs.models.vaes.wanvae import WanVAEConfig
+from fastvideo.dataset.transform import center_crop_th_tw, resize
+from fastvideo.models.vaes.wanvae import AutoencoderKLWan
+
+TARGET_H, TARGET_W = 480, 832
+NUM_FRAMES = 121
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--data-dir", type=Path, required=True, help="Dataset root (contains videos/, etc.).")
+    p.add_argument("--vae-path", type=Path, required=True, help="Path to the VAE directory (contains diffusion_pytorch_model.safetensors).")
+    p.add_argument("--video-subdir", type=str, default="videos", help="Input video subdirectory.")
+    p.add_argument("--out-subdir", type=str, default="roundtrip_videos", help="Output subdirectory.")
+    p.add_argument("--num-frames", type=int, default=NUM_FRAMES, help="Number of frames per video.")
+    p.add_argument("--height", type=int, default=TARGET_H)
+    p.add_argument("--width", type=int, default=TARGET_W)
+    p.add_argument("--fps", type=int, default=24)
+    p.add_argument("--device", type=str, default="cuda")
+    p.add_argument("--index", type=int, nargs="+", default=None, metavar="IDX",
+                   help="Process only these video indices (e.g. --index 4 7 12).")
+    p.add_argument("--limit", type=int, default=None, help="Process only first N videos (smoke test).")
+    p.add_argument("--force", action="store_true", help="Re-encode even if output already exists.")
+    return p.parse_args()
+
+
+def load_vae(vae_path: Path, device: str) -> AutoencoderKLWan:
+    config = WanVAEConfig(use_feature_cache=False)
+    vae = AutoencoderKLWan(config).to(device).eval()
+    weights = safetensors_load_file(str(vae_path / "diffusion_pytorch_model.safetensors"))
+    vae.load_state_dict(weights, strict=True)
+    return vae
+
+
+def load_video(path: Path, num_frames: int, height: int, width: int) -> torch.Tensor:
+    """Return pixel tensor [1, C, T, H, W] in [-1, 1]."""
+    reader = imageio.get_reader(str(path))
+    frames = [reader.get_data(i) for i in range(num_frames)]
+    reader.close()
+    clip = torch.from_numpy(np.stack(frames)).permute(0, 3, 1, 2).float() / 255.0
+    clip = center_crop_th_tw(clip, height, width, top_crop=False)
+    clip = resize(clip, (height, width), interpolation_mode="bilinear")
+    pixel = (clip * 2.0 - 1.0).permute(1, 0, 2, 3).unsqueeze(0)  # [1,C,T,H,W]
+    return pixel
+
+
+@torch.no_grad()
+def roundtrip(vae: AutoencoderKLWan, pixel: torch.Tensor, device: str) -> np.ndarray:
+    """Return decoded frames as uint8 numpy [T, H, W, C]."""
+    pixel = pixel.to(device)
+    with torch.autocast(device, dtype=torch.float32):
+        latent = vae.encode(pixel).mean
+    decoded = vae.decode(latent)
+    frames = decoded[0].permute(1, 2, 3, 0).float().cpu()
+    frames = ((frames / 2 + 0.5).clamp(0, 1) * 255).byte().numpy()
+    return frames
+
+
+def main() -> None:
+    args = parse_args()
+    device = args.device
+
+    videos_dir = args.data_dir / args.video_subdir
+    out_dir = args.data_dir / args.out_subdir
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    videos = sorted(videos_dir.glob("*.mp4"))
+    if args.index is not None:
+        wanted = {f"vid_{i:06d}.mp4" for i in args.index}
+        videos = [v for v in videos if v.name in wanted]
+    if args.limit is not None:
+        videos = videos[:args.limit]
+    if not videos:
+        print(f"[roundtrip] no videos found in {videos_dir}", flush=True)
+        return
+
+    print(f"[roundtrip] loading VAE from {args.vae_path} ...", flush=True)
+    vae = load_vae(args.vae_path, device)
+
+    print(f"[roundtrip] {len(videos)} videos → {out_dir}", flush=True)
+    for k, vpath in enumerate(videos, 1):
+        out_path = out_dir / vpath.name
+        if out_path.exists() and not args.force:
+            print(f"[roundtrip] [{k}/{len(videos)}] {vpath.name} already exists, skipping", flush=True)
+            continue
+
+        pixel = load_video(vpath, args.num_frames, args.height, args.width)
+        frames = roundtrip(vae, pixel, device)
+        imageio.mimsave(str(out_path), frames, fps=args.fps, macro_block_size=1)
+        print(f"[roundtrip] [{k}/{len(videos)}] {vpath.name} → {out_path.name}  "
+              f"shape={frames.shape}", flush=True)
+
+    print("[roundtrip] done.", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/decode_wansyn32k.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python3
+"""Decode FastVideo Wan2.2-Syn-121x704x1280_32k parquet latents to 480x832 mp4s.
+
+The HF dataset stores Wan 2.2 5B VAE latents (48-channel, 16x spatial, 4x
+temporal). We can't feed these to our Wan 2.1 WanTrack model directly, so we
+decode → pixels → downsize to our training resolution 480x832 → save mp4.
+Downstream (SAM + CoTracker + i2v_track preprocess) then treats these as if they
+were raw synth mp4s.
+
+Idempotent + resumable: skips samples whose mp4 already exists.
+
+Slurm sharded via ``--num-shards``/``--shard``. Each worker processes
+``parquets[shard::num_shards]``.
+"""
+from __future__ import annotations
+
+import argparse
+import glob
+import json
+import os
+import time
+from pathlib import Path
+
+import imageio.v2 as imageio
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+from diffusers import AutoencoderKLWan
+from PIL import Image
+
+
+def _resize_frames(frames: np.ndarray, out_h: int, out_w: int) -> np.ndarray:
+    """Center-crop-then-resize a (T,H,W,3) uint8 stack to (T,out_h,out_w,3)."""
+    T, H, W, _ = frames.shape
+    src_ar = W / H
+    dst_ar = out_w / out_h
+    if src_ar > dst_ar:
+        new_w = int(round(H * dst_ar))
+        x0 = (W - new_w) // 2
+        cropped = frames[:, :, x0:x0 + new_w, :]
+    else:
+        new_h = int(round(W / dst_ar))
+        y0 = (H - new_h) // 2
+        cropped = frames[:, y0:y0 + new_h, :, :]
+    out = np.empty((T, out_h, out_w, 3), dtype=np.uint8)
+    for t in range(T):
+        out[t] = np.asarray(Image.fromarray(cropped[t]).resize((out_w, out_h), Image.BILINEAR))
+    return out
+
+
+@torch.no_grad()
+def decode_row(vae: AutoencoderKLWan, latent_bytes: bytes, latent_shape: list[int],
+               device: torch.device, dtype: torch.dtype) -> np.ndarray:
+    """Decode one row's Wan 2.2 5B latent to a [T,704,1280,3] uint8 pixel stack."""
+    lat = np.frombuffer(latent_bytes, dtype=np.float32).reshape(latent_shape)
+    lat = torch.from_numpy(lat).to(device=device, dtype=dtype).unsqueeze(0)  # [1,C,T,H,W]
+    # unnormalize per-channel
+    m = torch.tensor(vae.config.latents_mean, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
+    s = torch.tensor(vae.config.latents_std, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
+    lat = lat * s + m
+    pix = vae.decode(lat, return_dict=False)[0]  # [1,3,T_out,H*16,W*16] in [-1,1]
+    pix = pix.clamp(-1, 1).float()
+    pix = ((pix + 1.0) * 127.5).round().to(torch.uint8)
+    pix = pix[0].permute(1, 2, 3, 0).cpu().numpy()  # [T,H,W,3]
+    return pix
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--parquet-dir", required=True, help="dir containing train/Part_*/latents_chunk_*.parquet")
+    p.add_argument("--vae-dir", required=True, help="Wan 2.2 5B VAE dir")
+    p.add_argument("--out-dir", required=True, help="output root; writes videos/vid_*.mp4 + meta/*.json")
+    p.add_argument("--out-h", type=int, default=480)
+    p.add_argument("--out-w", type=int, default=832)
+    p.add_argument("--fps", type=int, default=24)
+    p.add_argument("--shard", type=int, default=0)
+    p.add_argument("--num-shards", type=int, default=1)
+    p.add_argument("--dtype", default="bfloat16", choices=["float32", "bfloat16", "float16"])
+    p.add_argument("--limit-rows", type=int, default=None, help="cap rows for smoke test")
+    args = p.parse_args()
+
+    device = torch.device("cuda")
+    dtype = getattr(torch, args.dtype)
+
+    vids_dir = Path(args.out_dir) / "videos"
+    meta_dir = Path(args.out_dir) / "meta"
+    vids_dir.mkdir(parents=True, exist_ok=True)
+    meta_dir.mkdir(parents=True, exist_ok=True)
+
+    files = sorted(glob.glob(os.path.join(args.parquet_dir, "**", "*.parquet"), recursive=True))
+    if not files:
+        raise FileNotFoundError(f"no *.parquet under {args.parquet_dir}")
+    my_files = files[args.shard::args.num_shards]
+    print(f"[dec {args.shard}/{args.num_shards}] {len(my_files)} parquets", flush=True)
+
+    vae = AutoencoderKLWan.from_pretrained(args.vae_dir, torch_dtype=dtype).to(device).eval()
+
+    n_ok = 0
+    n_skip = 0
+    n_err = 0
+    t0 = time.time()
+    row_budget = args.limit_rows or 10**12
+    processed_rows = 0
+    for fk, f in enumerate(my_files):
+        try:
+            table = pq.read_table(f)
+        except Exception as e:  # skip unreadable parquets
+            print(f"[dec {args.shard}] parquet read fail {f}: {e}", flush=True)
+            continue
+        # Dataset's `id` field only encodes chunk index (not Part), so the same id
+        # appears in many Parts. Prefix Part directory so vid_ids are unique.
+        part_name = Path(f).parent.name  # e.g. "Part_57"
+        rows = table.to_pylist()
+        for r in rows:
+            if processed_rows >= row_budget:
+                break
+            processed_rows += 1
+            row_key = str(r.get("id") or r.get("file_name") or f"{fk:05d}_{n_ok:04d}")
+            vid_id = f"{part_name}_{row_key}"
+            vid_path = vids_dir / f"{vid_id}.mp4"
+            meta_path = meta_dir / f"{vid_id}.json"
+            if vid_path.exists() and meta_path.exists():
+                n_skip += 1
+                continue
+            try:
+                pix = decode_row(vae, r["vae_latent_bytes"], list(r["vae_latent_shape"]),
+                                 device=device, dtype=dtype)
+                resized = _resize_frames(pix, args.out_h, args.out_w)  # [T, out_h, out_w, 3]
+                # atomic mp4 write (imageio infers format from extension → keep .mp4)
+                tmp = vid_path.parent / f".{vid_path.name}.tmp.mp4"
+                imageio.mimsave(str(tmp), resized, fps=args.fps, codec="libx264",
+                                quality=7, macro_block_size=1)
+                tmp.rename(vid_path)
+                meta = {
+                    "path": vid_path.name,
+                    "id": vid_id,
+                    "cap": [str(r.get("caption") or "")],
+                    "fps": float(args.fps),
+                    "num_frames": int(resized.shape[0]),
+                    "resolution": [int(args.out_h), int(args.out_w)],
+                    "src_resolution": [int(r.get("height", 704)), int(r.get("width", 1280))],
+                    "src_dataset": "FastVideo/Wan2.2-Syn-121x704x1280_32k",
+                }
+                meta_path.write_text(json.dumps(meta))
+                n_ok += 1
+                if n_ok % 20 == 0:
+                    dt = time.time() - t0
+                    rate = n_ok / max(dt, 1e-9)
+                    print(f"[dec {args.shard}] ok={n_ok} skip={n_skip} err={n_err} "
+                          f"rate={rate:.2f}/s", flush=True)
+            except Exception as e:  # skip a bad row, keep going
+                n_err += 1
+                print(f"[dec {args.shard}] row {vid_id} FAILED: {e}", flush=True)
+        if processed_rows >= row_budget:
+            break
+
+    print(f"[dec {args.shard}] DONE ok={n_ok} skip={n_skip} err={n_err} "
+          f"in {(time.time()-t0)/60:.1f}min", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/diag_first_frame.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Diagnose the first-frame I2V conditioning + visualize synthetic tracks.
+
+(1) First-frame conditioning bug: the preprocessing stored ``first_frame_latent``
+    using vae.scaling_factor/shift_factor (absent for Wan -> effectively RAW),
+    while the model normalizes latents with per-channel latents_mean/std. This
+    script decodes the stored conditioning under both interpretations and compares
+    to the GT first frame to localize the space mismatch, and shows the current
+    model's generated first frame.
+
+(2) Synthetic-track previews: overlays authored controls (pan/zoom/drag/...) on the
+    real first frame so you can see what the control signals look like.
+
+Outputs PNGs to --out (default research_log/figures/).
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+import imageio.v2 as imageio
+import numpy as np
+import torch
+
+sys.path.insert(0, os.path.dirname(__file__))
+import synthetic_tracks as st  # noqa: E402
+import trackwan_infer as twi  # noqa: E402
+
+
+def _mse(a, b):
+    a = a[:min(len(a), len(b))].astype(np.float32)
+    b = b[:len(a)].astype(np.float32)
+    return float(((a - b) ** 2).mean())
+
+
+@torch.no_grad()
+def _decode_frame0(model, latent_bcthw, *, already_normalized: bool) -> np.ndarray:
+    """latent [1,16,T,H,W] -> frame0 uint8 (H,W,3). If raw, normalize first."""
+    from fastvideo.training.training_utils import normalize_dit_input
+    x = latent_bcthw.to(model.device, torch.bfloat16)
+    if not already_normalized:
+        x = normalize_dit_input("wan", x, model.vae)
+    px = model.decode_latents(x.permute(0, 2, 1, 3, 4))[0]  # [3,T,H,W] in [0,1]
+    f0 = (px[:, 0].clamp(0, 1).float().cpu().numpy() * 255).astype(np.uint8)
+    return np.transpose(f0, (1, 2, 0))  # H,W,3
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--export", required=True)
+    p.add_argument("--yaml", default="examples/train/scenario/worldmodel/finetune_wantrack_i2v.yaml")
+    p.add_argument("--data", default="/mnt/weka/home/hao.zhang/shao/data/motion_pipeline/"
+                   "wan22_a14b_720p_24fps/preprocessed_i2v_track/combined_parquet_dataset")
+    p.add_argument("--out", default="research_log/figures")
+    p.add_argument("--clip", type=int, default=0)
+    p.add_argument("--steps", type=int, default=30)
+    args = p.parse_args()
+    os.makedirs(args.out, exist_ok=True)
+
+    model, tc = twi.load_trackwan(args.export, args.yaml)
+    text_len = int(tc.pipeline_config.text_encoder_configs[0].arch_config.text_len)
+    s = twi.load_conditioning_from_parquet(args.data, [args.clip], text_len)[0]
+
+    vae_latent = s["vae_latent"]          # raw GT latents
+    ff = s["first_frame_latent"]          # stored conditioning latent
+
+    # GT first frame: vae_latent is raw -> normalize -> decode.
+    gt0 = _decode_frame0(model, vae_latent, already_normalized=False)
+
+    # Stored conditioning decoded TWO ways to find its true space:
+    #  (a) "model's view": treat stored as already-normalized (decode_latents denorms it).
+    #      This is effectively what the model was conditioned on.
+    cond_as_norm0 = _decode_frame0(model, ff, already_normalized=True)
+    #  (b) "as raw": treat stored as raw -> normalize -> decode.
+    cond_as_raw0 = _decode_frame0(model, ff, already_normalized=False)
+
+    # Current model's generated first frame (GT tracks).
+    Tpx = (vae_latent.shape[2] - 1) * int(tc.pipeline_config.vae_config.arch_config.temporal_compression_ratio) + 1
+    lat = twi.generate(model, first_frame_latent=ff, text_embedding=s["text_embedding"],
+                       text_attention_mask=s["text_attention_mask"],
+                       track_points=s["track_points"][:, :Tpx], track_visibility=s["track_visibility"][:, :Tpx],
+                       num_steps=args.steps, seed=1000)
+    gen0 = twi.decode_to_pixels(model, lat)[0]
+
+    # Save individual + side-by-side.
+    panels = {"1_gt_first_frame": gt0,
+              "2_cond_as_normalized_MODELVIEW": cond_as_norm0,
+              "3_cond_as_raw_then_normalized": cond_as_raw0,
+              "4_generated_first_frame": gen0}
+    for name, img in panels.items():
+        imageio.imwrite(os.path.join(args.out, f"firstframe_clip{args.clip}_{name}.png"), img)
+    strip = np.concatenate([gt0, cond_as_norm0, cond_as_raw0, gen0], axis=1)
+    imageio.imwrite(os.path.join(args.out, f"firstframe_clip{args.clip}_compare.png"), strip)
+
+    print("=== first-frame diagnosis (clip %d) ===" % args.clip)
+    print(f"  MSE(GT, cond-as-normalized / MODEL'S VIEW) = {_mse(gt0, cond_as_norm0):8.1f}")
+    print(f"  MSE(GT, cond-as-raw-then-normalized)       = {_mse(gt0, cond_as_raw0):8.1f}")
+    print(f"  MSE(GT, generated first frame)             = {_mse(gt0, gen0):8.1f}")
+    print("  -> if 'as-raw' MSE << 'model's view' MSE, the stored conditioning is RAW")
+    print("     and the model was fed a mis-normalized first frame (the bug).")
+
+    # ---- synthetic track previews on the real first frame ----
+    for name in ["pan_right", "zoom_in", "rotate_cw", "drag_center_right", "swirl"]:
+        tr, vis = st.preset(name, Tpx, 50, strength=0.25)
+        H, W, _ = gt0.shape
+        ov = st._overlay_preview(gt0, st.to_pixel(tr, H, W), vis, stride=3)
+        imageio.imwrite(os.path.join(args.out, f"trackpreview_{name}.png"), ov)
+    print(f"\n[viz] wrote first-frame panels + synthetic track previews to {args.out}/")
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/download_resilient.sh

@@ -0,0 +1,47 @@
+#!/usr/bin/env bash
+# Self-healing OpenVidHD download: HF unauthenticated throttles ~per-200GB then
+# resets after a cooldown. Run N idempotent download shards; watchdog detects a
+# stall (no zip growth) and kills+cooldowns+relaunches (resume) until all 98 parts
+# done. Runs as ONE srun on the node so local kill works. Honors $HF_TOKEN if set.
+set +e
+WORK=/mnt/lustre/vlm-s4duan
+export HF_HOME=$WORK/.hf
+cd "$WORK/FastVideo"; source .venv/bin/activate
+NSHARD=${NSHARD:-4}
+DC(){ ls "$WORK"/openvid_1m/_extracted/*.done 2>/dev/null | wc -l; }
+ZB(){ find "$WORK"/openvid_1m/_zips -type f -printf '%s\n' 2>/dev/null | awk '{s+=$1}END{print s+0}'; }
+
+pkill -f openvid_download_hd.py 2>/dev/null; sleep 3
+round=0
+while [ "$(DC)" -lt 98 ]; do
+  round=$((round+1))
+  echo "[$(date +%H:%M)] ROUND $round start parts=$(DC)/98 ${HF_TOKEN:+(token set)}"
+  pids=()
+  for s in $(seq 0 $((NSHARD-1))); do
+    python data_pipeline/openvid_download_hd.py --shard "$s" --num-shards "$NSHARD" \
+      --videos-dir "$WORK"/openvid_1m/videos --zip-dir "$WORK"/openvid_1m/_zips/"$s" \
+      --only-list "$WORK"/openvid/OpenVidHD_filtered.txt &
+    pids+=($!)
+  done
+  last=$(ZB); stall=0
+  while :; do
+    sleep 120
+    [ "$(DC)" -ge 98 ] && break
+    alive=0; for p in "${pids[@]}"; do kill -0 "$p" 2>/dev/null && alive=1; done
+    [ "$alive" -eq 0 ] && { echo "[$(date +%H:%M)] round done naturally parts=$(DC)/98"; break; }
+    now=$(ZB)
+    if [ "$now" -le "$last" ]; then stall=$((stall+1)); else stall=0; fi
+    last=$now
+    echo "[$(date +%H:%M)] parts=$(DC)/98 zip=$((now/1000000000))GB stall=$stall"
+    if [ "$stall" -ge 2 ]; then
+      echo "[$(date +%H:%M)] STALL -> kill + cooldown"
+      kill "${pids[@]}" 2>/dev/null; sleep 5; kill -9 "${pids[@]}" 2>/dev/null
+      pkill -9 -f openvid_download_hd.py 2>/dev/null
+      break
+    fi
+  done
+  wait 2>/dev/null
+  [ "$(DC)" -ge 98 ] && break
+  echo "[$(date +%H:%M)] cooldown 180s (parts $(DC)/98)"; sleep 180
+done
+echo "ALL_PARTS_DONE $(DC)/98 videos=$(ls "$WORK"/openvid_1m/videos/*.mp4 2>/dev/null | wc -l)"

data_pipeline/dp_sweep.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+# DP concurrency sweep on ONE GPU: K workers sharing GPU 3, aggregate throughput.
+set +e
+WORK=/mnt/lustre/vlm-s4duan
+source "$WORK/FastVideo/.venv/bin/activate"
+export CUDA_HOME=/usr/local/cuda
+export PATH=$CUDA_HOME/bin:$PATH
+export TORCH_HOME=$WORK/.torch
+export TOKENIZERS_PARALLELISM=false
+V=$WORK/dp_test/videos; T=$WORK/dp_test/tracks; L=$WORK/dp_test/videos.txt
+mkdir -p "$V" "$T"
+if [ ! -f "$V/clip_24.mp4" ]; then
+  for i in $(seq -w 1 24); do cp -f "$WORK/wan_overfit/data/videos/vid_000000.mp4" "$V/clip_$i.mp4"; done
+fi
+ls "$V"/*.mp4 > "$L"
+echo "test set: $(wc -l < "$L") videos on GPU 3 (121f 720p, grid 50)"
+for K in 1 2 3 4; do
+  rm -f "$T"/*.npz
+  echo "===================== K=$K procs/GPU ====================="
+  pids=()
+  for s in $(seq 0 $((K-1))); do
+    CUDA_VISIBLE_DEVICES=3 python data_pipeline/extract_tracks_mp.py \
+      --video-list "$L" --out-dir "$T" --gpus-per-node 1 \
+      --shard "$s" --num-shards "$K" --fps 24 --num-frames 121 --grid-size 50 &
+    pids+=($!)
+  done
+  wait "${pids[@]}"
+done
+echo "SWEEP_DONE"

data_pipeline/droid_dataset_dashboard.py

@@ -0,0 +1,224 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Inspect a track-conditioning dataset: are the clips diverse, and do the
+CoTracker tracks actually capture the motion?
+
+Two concerns this surfaces:
+  1. Near-duplicate clips (the frame-0 cluster can over-select the same scene).
+  2. Weak control signal: CoTracker uses a frame-0 grid, so anything that ENTERS
+     the frame later (or that the tracker loses) is never queried -> its motion is
+     invisible to the tracks. A clip can look dynamic yet have near-static tracks.
+
+``render`` (CPU; run on a node) writes, per clip:
+  - <id>__overlay.mp4   the video with the CoTracker tracks drawn (dots + tails),
+  - <id>__thumb.png     first frame,
+and a ``stats.json`` with per-clip motion-coverage + a near-duplicate grouping.
+
+``serve`` (no GPU) is a gradio gallery: clips sorted by motion (least first, so
+dead/duplicate clips float to the top), each with its overlay + stats.
+
+    # 1) render artifacts (node)
+    srun --jobid=<job> --overlap --ntasks=1 .venv/bin/python data_pipeline/droid_dataset_dashboard.py render \
+        --data-dir /mnt/weka/home/hao.zhang/shao/data/motion_pipeline/droid_track_200 --out <dash_dir>
+    # 2) serve (login node ok)
+    .venv/bin/python data_pipeline/droid_dataset_dashboard.py serve --dash-dir <dash_dir> --share
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(__file__))
+
+MOVE_THRESH_PX = 8.0  # a point "moves" if it travels > this from frame 0 (input px)
+VIS_THRESH = 0.5
+
+
+def _load_frames(path: str) -> np.ndarray:
+    try:
+        from decord import VideoReader, cpu
+        vr = VideoReader(path, ctx=cpu(0))
+        return vr.get_batch(list(range(len(vr)))).asnumpy()
+    except Exception:  # noqa: BLE001
+        import imageio.v2 as imageio
+        rd = imageio.get_reader(path, format="ffmpeg")
+        fr = np.stack([np.asarray(f) for f in rd], axis=0)
+        rd.close()
+        return fr
+
+
+def _clip_stats(tracks: np.ndarray, vis: np.ndarray) -> dict[str, Any]:
+    """Motion coverage from tracks [T,N,2] (px) + vis [T,N]."""
+    T, N, _ = tracks.shape
+    v = vis > VIS_THRESH
+    # displacement of each point from its frame-0 position, only where visible
+    disp = np.sqrt(((tracks - tracks[0:1])**2).sum(-1))  # [T,N]
+    max_disp = np.where(v, disp, 0.0).max(0)  # [N] per-point peak travel
+    moving = max_disp > MOVE_THRESH_PX
+    return {
+        "n_points": int(N),
+        "frac_visible": float(v.mean()),
+        "frac_moving": float(moving.mean()),  # fraction of grid points that ever move
+        "n_moving": int(moving.sum()),
+        "mean_motion_px": float(max_disp.mean()),
+        "p95_motion_px": float(np.percentile(max_disp, 95)),
+        "max_motion_px": float(max_disp.max()),
+    }
+
+
+def cmd_render(args: argparse.Namespace) -> None:
+    from fastvideo.train.callbacks.track_validation import (_draw_overlay, _grid_colors, _subsample)
+    data_dir = Path(args.data_dir)
+    manifest = json.loads((data_dir / "videos2caption.json").read_text())
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    import imageio.v2 as imageio
+    from PIL import Image
+
+    embs: list[np.ndarray] = []
+    entries: list[dict[str, Any]] = []
+    for rec in manifest:
+        cid = Path(rec["path"]).stem
+        vpath = str(data_dir / "videos" / rec["path"])
+        ppath = rec.get("points_path")
+        frames = _load_frames(vpath)
+        d = np.load(ppath)
+        tracks = d["tracks"].astype(np.float32)[:frames.shape[0]]  # [T,N,2] px
+        vis = d["visibility"].astype(np.float32)[:frames.shape[0]]
+        H, W = int(frames.shape[1]), int(frames.shape[2])
+
+        stats = _clip_stats(tracks, vis)
+
+        # overlay (subsampled grid for clarity)
+        grid = int(round(tracks.shape[1]**0.5))
+        tn = tracks / np.array([W, H], np.float32)  # normalize for _subsample
+        tr_s, vs_s = _subsample(tn, vis, grid, args.stride)
+        colors = _grid_colors(grid, args.stride)[:tr_s.shape[1]]
+        trpx = tr_s.copy()
+        trpx[..., 0] *= W
+        trpx[..., 1] *= H
+        ov = _draw_overlay(frames, trpx, vs_s, colors, args.tail, 2, VIS_THRESH)
+        imageio.mimsave(str(out / f"{cid}__overlay.mp4"), ov, fps=args.fps, macro_block_size=1)
+        Image.fromarray(frames[0]).save(str(out / f"{cid}__thumb.png"))
+
+        # frame-0 embedding for near-duplicate grouping (32x32 gray, L2-norm)
+        g = np.asarray(Image.fromarray(frames[0]).convert("L").resize((32, 32)), np.float32).reshape(-1)
+        embs.append(g / (np.linalg.norm(g) + 1e-8))
+
+        entries.append({"id": cid, "episode": rec.get("source_episode"), "caption": rec["cap"][0], **stats})
+        print(
+            f"[dash] {cid} move%={stats['frac_moving']*100:4.1f} mean={stats['mean_motion_px']:5.1f}px "
+            f"max={stats['max_motion_px']:5.1f}px",
+            flush=True)
+
+    # near-duplicate grouping: union-find on cosine-sim >= threshold
+    E = np.stack(embs, 0)
+    sim = E @ E.T
+    n = len(entries)
+    parent = list(range(n))
+
+    def find(i: int) -> int:
+        while parent[i] != i:
+            parent[i] = parent[parent[i]]
+            i = parent[i]
+        return i
+
+    for i in range(n):
+        for j in range(i + 1, n):
+            if sim[i, j] >= args.dup_thresh:
+                parent[find(i)] = find(j)
+    groups: dict[int, list[int]] = {}
+    for i in range(n):
+        groups.setdefault(find(i), []).append(i)
+    group_id = {}
+    for gi, (_, members) in enumerate(sorted(groups.items(), key=lambda kv: -len(kv[1]))):
+        for m in members:
+            group_id[m] = gi
+    for i, e in enumerate(entries):
+        e["dup_group"] = int(group_id[i])
+
+    n_groups = len(groups)
+    n_dead = sum(1 for e in entries if e["frac_moving"] < 0.02)
+    summary = {
+        "n_clips": n,
+        "n_dup_groups": n_groups,
+        "largest_dup_group": max(len(v) for v in groups.values()),
+        "n_low_motion_clips(<2%)": n_dead,
+        "median_frac_moving": float(np.median([e["frac_moving"] for e in entries])),
+        "median_mean_motion_px": float(np.median([e["mean_motion_px"] for e in entries])),
+        "dup_thresh": args.dup_thresh,
+        "move_thresh_px": MOVE_THRESH_PX,
+    }
+    (out / "stats.json").write_text(json.dumps({"summary": summary, "clips": entries}, indent=2))
+    print(f"[dash] SUMMARY: {json.dumps(summary, indent=2)}", flush=True)
+    print(f"[dash] wrote dashboard artifacts -> {out}", flush=True)
+
+
+def cmd_serve(args: argparse.Namespace) -> None:
+    import gradio as gr
+    dash = Path(args.dash_dir)
+    data = json.loads((dash / "stats.json").read_text())
+    summary, clips = data["summary"], data["clips"]
+    # sort: least motion first (dead/duplicate clips surface), then by dup group
+    order = sorted(range(len(clips)), key=lambda i: (clips[i]["frac_moving"], clips[i]["dup_group"]))
+
+    def label(i: int) -> str:
+        c = clips[i]
+        return (f"#{i} ep{c['episode']} | move {c['frac_moving']*100:.0f}% | "
+                f"mean {c['mean_motion_px']:.0f}px | dupgrp {c['dup_group']}")
+
+    gallery_items = [(str(dash / f"{clips[i]['id']}__thumb.png"), label(i)) for i in order]
+
+    def show(evt):  # gr.SelectData; unannotated so gradio's get_type_hints doesn't eval it
+        i = order[evt.index]
+        c = clips[i]
+        return (str(dash / f"{c['id']}__overlay.mp4"), json.dumps(c, indent=2))
+
+    head = (f"### DROID dataset inspection — {summary['n_clips']} clips\n"
+            f"- near-duplicate groups: **{summary['n_dup_groups']}** "
+            f"(largest group **{summary['largest_dup_group']}** clips)\n"
+            f"- low-motion clips (<2% points move): **{summary['n_low_motion_clips(<2%)']}**\n"
+            f"- median fraction of points moving: **{summary['median_frac_moving']*100:.1f}%**, "
+            f"median mean-motion **{summary['median_mean_motion_px']:.1f}px**\n\n"
+            f"Sorted least-motion first. Overlay = CoTracker tracks (dots + tails); "
+            f"watch for moving objects with NO dots on them (off-frame entries / lost tracks).")
+
+    with gr.Blocks(title="DROID dataset dashboard") as demo:
+        gr.Markdown(head)
+        with gr.Row():
+            gal = gr.Gallery(value=gallery_items, columns=6, height=560, label="clips (least motion first)")
+            with gr.Column():
+                vid = gr.Video(label="track overlay")
+                meta = gr.Code(label="clip stats", language="json")
+        gal.select(show, None, [vid, meta])
+    demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share)
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub = p.add_subparsers(dest="cmd", required=True)
+    r = sub.add_parser("render")
+    r.add_argument("--data-dir", required=True, help="dataset root (videos/ + videos2caption.json + tracks/)")
+    r.add_argument("--out", required=True)
+    r.add_argument("--stride", type=int, default=3, help="subsample the NxN grid for overlay clarity")
+    r.add_argument("--tail", type=int, default=12)
+    r.add_argument("--fps", type=int, default=12)
+    r.add_argument("--dup-thresh", type=float, default=0.985, help="cosine-sim on frame-0 to call clips duplicates")
+    r.set_defaults(func=cmd_render)
+    s = sub.add_parser("serve")
+    s.add_argument("--dash-dir", required=True)
+    s.add_argument("--host", default="0.0.0.0")
+    s.add_argument("--port", type=int, default=7872)
+    s.add_argument("--share", action="store_true")
+    s.set_defaults(func=cmd_serve)
+    args = p.parse_args()
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/eval_sparse_track_epe.py

@@ -0,0 +1,416 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""Evaluate super-sparse track adherence via CoTracker EPE.
+
+For each val clip, we build the exact adversarial case the user described:
+  * Pick ONE active foreground trace from the GT parquet — the highest-motion
+    foreground track (object_ids >= 0, visible in frame 0).
+  * Pick N static background anchors — random background tracks (object_ids == -1)
+    with position frozen at frame-0 (visibility=1 across all frames).
+  * Generate video conditioned on (active + anchors) — this is what training
+    was supposed to teach the model to handle.
+  * Also generate an ablation with track_points=None (unconditional motion).
+  * Run CoTracker on both generated videos, initialised at frame 0 with the
+    conditioning points. Compare CoTracker-extracted tracks against the GT
+    trace to compute End-Point-Error in pixel space.
+
+Metrics logged per clip and aggregated to wandb:
+  epe_active_full   EPE(GT_active, CoTracker(gen_full))
+  epe_active_notrack EPE(GT_active, CoTracker(gen_no_track))
+  epe_bg            EPE(static_anchor, CoTracker(gen_full)) — should be small
+  DELTA = epe_active_notrack - epe_active_full — positive means training helped.
+
+Usage:
+  srun --overlap --jobid=529 --ntasks=1 -w hpc-rack-2-13 --chdir=$PWD bash -lc '
+    source .venv/bin/activate
+    export HOME=/mnt/lustre/vlm-s4duan HF_HOME=/mnt/lustre/vlm-s4duan/.hf \
+      TORCH_HOME=/mnt/lustre/vlm-s4duan/.torch MPLCONFIGDIR=/mnt/lustre/vlm-s4duan/.mpl \
+      TRITON_CACHE_DIR=/tmp/triton_eval TOKENIZERS_PARALLELISM=false NCCL_CUMEM_ENABLE=0 \
+      PYTHONPATH=$PWD TRACKWAN_TRACK_BIAS=1 CUDA_VISIBLE_DEVICES=2 \
+      WANDB_API_KEY=<key> WANDB_MODE=online
+    python data_pipeline/eval_sparse_track_epe.py \
+      --model-dir /mnt/lustre/vlm-s4duan/exports/merged_bias_ckpt4800 \
+      --yaml examples/train/scenario/worldmodel/finetune_wantrack_openvid_sparse_1p3b_merged_bias.yaml \
+      --data-path /mnt/lustre/vlm-s4duan/openvid_1m/combined_parquet_dataset \
+      --num-clips 5 --steps 30 --w-motion 1.5 \
+      --wandb-run-name eval_mb4800_sparse
+  '
+"""
+from __future__ import annotations
+
+import argparse
+import glob
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+# make trackwan_infer importable
+sys.path.insert(0, os.path.join(os.path.dirname(__file__)))
+
+
+# ==========================================================================
+# CoTracker
+# ==========================================================================
+def load_cotracker(device: str) -> torch.nn.Module:
+    """Load CoTracker v3 offline. Uses the warm hub cache under $TORCH_HOME/hub."""
+    hub_dir = Path(torch.hub.get_dir())
+    local = hub_dir / "facebookresearch_co-tracker_main"
+    if local.exists():
+        model = torch.hub.load(str(local), "cotracker3_offline", source="local", trust_repo=True)
+    else:
+        model = torch.hub.load("facebookresearch/co-tracker", "cotracker3_offline", trust_repo=True)
+    return model.to(device).eval()
+
+
+@torch.no_grad()
+def cotracker_query(cotracker: torch.nn.Module, video: np.ndarray, queries_xy: np.ndarray,
+                    device: str) -> tuple[np.ndarray, np.ndarray]:
+    """Track `queries_xy` (N, 2 pixel coords) forward through `video` [T,H,W,3] uint8.
+
+    Returns tracks [T,N,2] (pixel coords) and visibility [T,N] (bool).
+    """
+    T, H, W, C = video.shape
+    v = torch.from_numpy(video).permute(0, 3, 1, 2).float()[None].to(device)  # [1,T,3,H,W]
+    N = queries_xy.shape[0]
+    q = np.zeros((1, N, 3), dtype=np.float32)
+    q[0, :, 0] = 0  # start-frame index
+    q[0, :, 1] = queries_xy[:, 0]  # x
+    q[0, :, 2] = queries_xy[:, 1]  # y
+    tracks, visibility = cotracker(v, queries=torch.from_numpy(q).to(device))
+    return tracks[0].cpu().numpy(), visibility[0].cpu().numpy()
+
+
+# ==========================================================================
+# Sparse-conditioning sampler
+# ==========================================================================
+def pick_sparse_conditioning(row: dict, *, num_active: int, num_bg: int,
+                             seed: int) -> dict:
+    """From a preprocessed parquet row, pick the sparse case: `num_active` moving fg
+    traces (highest-displacement fg tracks visible in frame 0) + `num_bg` static bg
+    anchors (random bg tracks, position frozen at frame-0).
+
+    Everything is normalised to [0,1] to match training. Pixel-space is only used at
+    the metric-computation step.
+    """
+    tp = row["track_points"]  # [T,N,2] normalised
+    tv = row["track_visibility"]  # [T,N]
+    oid = row["object_ids"]  # [N] int, -1 for bg
+    T, N, _ = tp.shape
+    rng = np.random.RandomState(seed)
+
+    fg = np.where((oid >= 0) & (tv[0] > 0.5))[0]
+    bg = np.where((oid == -1) & (tv[0] > 0.5))[0]
+    if fg.size == 0:
+        raise RuntimeError("no visible foreground tracks in frame 0")
+    # displacement of each fg over the full clip (only over visible frames)
+    disp = np.linalg.norm(tp[-1, fg] - tp[0, fg], axis=-1)
+    active_idx = fg[np.argsort(-disp)[:num_active]]  # top-`num_active` moving fg
+    n_bg = min(num_bg, bg.size)
+    bg_idx = rng.choice(bg, size=n_bg, replace=False) if n_bg > 0 else np.zeros((0, ), dtype=int)
+
+    active_gt = tp[:, active_idx]  # [T, num_active, 2] — GT motion
+    active_vis = tv[:, active_idx]
+    bg_static_pos = np.broadcast_to(tp[0:1, bg_idx], (T, n_bg, 2)).copy()  # position frozen at f0
+    bg_vis = np.ones((T, n_bg), dtype=np.float32)
+
+    cond_tracks = np.concatenate([active_gt, bg_static_pos], axis=1)  # [T, num_active+n_bg, 2]
+    cond_vis = np.concatenate([active_vis, bg_vis], axis=1)
+    return {
+        "active_idx": active_idx,
+        "active_gt": active_gt,  # [T,num_active,2] normalised
+        "active_vis": active_vis,
+        "bg_static_pos": bg_static_pos,  # [T,n_bg,2] normalised (const in t)
+        "cond_tracks": cond_tracks,  # [T, num_active+n_bg, 2] normalised
+        "cond_vis": cond_vis,
+        "n_active": len(active_idx),
+        "n_bg": n_bg,
+    }
+
+
+# ==========================================================================
+# Metrics
+# ==========================================================================
+def epe_norm_to_px(pred_norm: np.ndarray, gt_norm: np.ndarray, vis: np.ndarray,
+                   res_h: int, res_w: int) -> float:
+    """L2 distance between predicted/GT normalised tracks, evaluated in pixel space,
+    averaged only over visible frames × visible points."""
+    diff = (pred_norm - gt_norm) * np.array([res_w, res_h])[None, None, :]
+    l2 = np.linalg.norm(diff, axis=-1)  # [T, N]
+    mask = vis > 0.5
+    if not mask.any():
+        return float("nan")
+    return float(l2[mask].mean())
+
+
+def epe_px(pred_px: np.ndarray, gt_px: np.ndarray, vis: np.ndarray) -> float:
+    """L2 in pixel space — both inputs already in pixels."""
+    l2 = np.linalg.norm(pred_px - gt_px, axis=-1)
+    mask = vis > 0.5
+    if not mask.any():
+        return float("nan")
+    return float(l2[mask].mean())
+
+
+# ==========================================================================
+# Trace overlay for wandb visualisation
+# ==========================================================================
+def overlay_tracks(video: np.ndarray, tracks_norm: np.ndarray, vis: np.ndarray,
+                   colors: list[tuple[int, int, int]], radius: int = 3,
+                   tail: int = 12) -> np.ndarray:
+    """Draw normalised tracks on a video [T,H,W,3] uint8.  Small tail + current dot."""
+    from PIL import Image, ImageDraw
+    T, H, W, C = video.shape
+    N = tracks_norm.shape[1]
+    out = []
+    for t in range(T):
+        img = Image.fromarray(np.ascontiguousarray(video[t]))
+        dr = ImageDraw.Draw(img)
+        t0 = max(0, t - tail)
+        for i in range(N):
+            col = colors[i % len(colors)]
+            # tail
+            if t - t0 >= 1:
+                pts = [(float(tracks_norm[k, i, 0] * W), float(tracks_norm[k, i, 1] * H))
+                       for k in range(t0, t + 1)]
+                dr.line(pts, fill=col, width=1)
+            if vis[t, i] > 0.5:
+                x = float(tracks_norm[t, i, 0] * W)
+                y = float(tracks_norm[t, i, 1] * H)
+                dr.ellipse([x - radius, y - radius, x + radius, y + radius], fill=col)
+        out.append(np.asarray(img))
+    return np.stack(out)
+
+
+# ==========================================================================
+# Parquet loader (small subset — avoid loading the whole 259k row set)
+# ==========================================================================
+def load_first_n_from_parquet(data_path: str, n: int, text_len: int) -> list:
+    import pyarrow.parquet as pq
+    from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v_track
+    from fastvideo.dataset.utils import collate_rows_from_parquet_schema
+    files = sorted(glob.glob(os.path.join(data_path, "**", "*.parquet"), recursive=True))
+    if not files:
+        raise FileNotFoundError(data_path)
+    rows: list = []
+    for f in files:
+        rows.extend(pq.read_table(f).to_pylist())
+        if len(rows) >= n:
+            break
+    sel = rows[:n]
+    batch = collate_rows_from_parquet_schema(sel, pyarrow_schema_i2v_track,
+                                             text_padding_length=int(text_len), cfg_rate=0.0)
+    infos = batch.get("info_list") or [{} for _ in sel]
+    out = []
+    for i in range(len(sel)):
+        out.append({
+            "text_embedding": batch["text_embedding"][i:i + 1].clone(),
+            "text_attention_mask": batch["text_attention_mask"][i:i + 1].clone(),
+            "first_frame_latent": batch["first_frame_latent"][i:i + 1].clone(),
+            "clip_feature": batch["clip_feature"][i:i + 1].clone(),
+            "vae_latent": batch["vae_latent"][i:i + 1].clone(),
+            "track_points": batch["track_points"][i].numpy(),  # [T,N,2] normalised
+            "track_visibility": batch["track_visibility"][i].numpy(),  # [T,N]
+            "object_ids": batch["object_ids"][i].numpy(),  # [N]
+            "caption": str(infos[i].get("caption", "") if i < len(infos) else ""),
+        })
+    return out
+
+
+# ==========================================================================
+# Main
+# ==========================================================================
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--model-dir", required=True)
+    p.add_argument("--yaml", required=True)
+    p.add_argument("--data-path", required=True)
+    p.add_argument("--num-clips", type=int, default=5)
+    p.add_argument("--num-active", type=int, default=1)
+    p.add_argument("--num-bg", type=int, default=20)
+    p.add_argument("--steps", type=int, default=30)
+    p.add_argument("--w-motion", type=float, default=1.5)
+    p.add_argument("--seed", type=int, default=1234)
+    p.add_argument("--fps", type=int, default=24)
+    p.add_argument("--wandb-project", default="wantrack-bidir")
+    p.add_argument("--wandb-run-name", required=True)
+    p.add_argument("--out-dir", default="/mnt/lustre/vlm-s4duan/eval_sparse_epe")
+    args = p.parse_args()
+
+    Path(args.out_dir).mkdir(parents=True, exist_ok=True)
+
+    import wandb
+    import imageio.v2 as imageio
+    import trackwan_infer as twi
+
+    print(f"[eval] loading model {args.model_dir}", flush=True)
+    model, tc = twi.load_trackwan(args.model_dir, args.yaml)
+    text_len = int(tc.data.text_padding_length) if hasattr(tc.data, "text_padding_length") else 256
+
+    print(f"[eval] loading {args.num_clips} parquet clips", flush=True)
+    samples = load_first_n_from_parquet(args.data_path, args.num_clips, text_len)
+
+    print("[eval] loading CoTracker", flush=True)
+    device = model.device
+    cotracker = load_cotracker(device.type + ":" + str(device.index) if device.index is not None else device.type)
+
+    res_h = int(tc.data.num_height)
+    res_w = int(tc.data.num_width)
+    palette = [(255, 60, 60), (60, 255, 60), (60, 90, 255), (255, 200, 60), (200, 60, 255), (60, 255, 220)]
+
+    run = wandb.init(project=args.wandb_project, name=args.wandb_run_name,
+                     config={"model_dir": args.model_dir, "num_clips": args.num_clips,
+                             "num_active": args.num_active, "num_bg": args.num_bg,
+                             "steps": args.steps, "w_motion": args.w_motion, "res": [res_h, res_w]})
+
+    per_clip: list = []
+
+    for idx, s in enumerate(samples):
+        try:
+            sc = pick_sparse_conditioning(s, num_active=args.num_active, num_bg=args.num_bg,
+                                          seed=args.seed + idx)
+        except RuntimeError as exc:
+            print(f"[eval] clip {idx}: {exc}; skipping", flush=True)
+            continue
+
+        # ── Generate two videos: with sparse tracks (motion CFG) + without any tracks ──
+        tp_t = torch.from_numpy(sc["cond_tracks"])[None].float()  # [1,T,N,2]
+        tv_t = torch.from_numpy(sc["cond_vis"])[None].float()  # [1,T,N]
+
+        seed = args.seed + idx
+        common = dict(first_frame_latent=s["first_frame_latent"],
+                      text_embedding=s["text_embedding"],
+                      text_attention_mask=s["text_attention_mask"],
+                      clip_feature=s["clip_feature"],
+                      num_steps=args.steps, seed=seed)
+
+        print(f"[eval] clip {idx}: generating full ({sc['n_active']} active + {sc['n_bg']} anchors)", flush=True)
+        t0 = time.time()
+        lat_full = twi.generate(model, track_points=tp_t, track_visibility=tv_t,
+                                guidance_scale=args.w_motion, **common)
+        gen_full = twi.decode_to_pixels(model, lat_full)  # [T,H,W,3] uint8
+        t_full = time.time() - t0
+        print(f"[eval] clip {idx}: full took {t_full:.1f}s", flush=True)
+
+        print(f"[eval] clip {idx}: generating no-track", flush=True)
+        t0 = time.time()
+        lat_no = twi.generate(model, track_points=None, track_visibility=None,
+                              guidance_scale=1.0, **common)
+        gen_no = twi.decode_to_pixels(model, lat_no)
+        t_no = time.time() - t0
+        print(f"[eval] clip {idx}: no-track took {t_no:.1f}s", flush=True)
+
+        T, H, W, _ = gen_full.shape
+
+        # ── CoTracker on both generations, queried at frame-0 with same points ──
+        query_norm = sc["cond_tracks"][0]  # [N_total, 2] in [0,1]
+        queries_px = query_norm * np.array([W, H])[None]  # pixel space
+        print(f"[eval] clip {idx}: CoTracker on full", flush=True)
+        tk_full_px, tk_full_vis = cotracker_query(cotracker, gen_full, queries_px, device=device.type)
+        print(f"[eval] clip {idx}: CoTracker on no-track", flush=True)
+        tk_no_px, tk_no_vis = cotracker_query(cotracker, gen_no, queries_px, device=device.type)
+
+        # GT in pixel space, at the SAME resolution as the generation
+        gt_full_px = sc["cond_tracks"] * np.array([W, H])[None, None, :]  # [T,N_total,2]
+        active_slice = slice(0, sc["n_active"])
+        bg_slice = slice(sc["n_active"], sc["n_active"] + sc["n_bg"])
+
+        # Active EPE: CoTracker(gen) vs the GT trace we specified.
+        # Filter by GT visibility (active_vis) AND CoTracker's own visibility.
+        active_vis_gt = sc["active_vis"]  # [T,n_active]
+        both_vis_full = (active_vis_gt > 0.5) & (tk_full_vis[:, active_slice] > 0.5)
+        both_vis_no = (active_vis_gt > 0.5) & (tk_no_vis[:, active_slice] > 0.5)
+
+        epe_active_full = epe_px(tk_full_px[:, active_slice],
+                                 gt_full_px[:, active_slice],
+                                 both_vis_full.astype(np.float32))
+        epe_active_no = epe_px(tk_no_px[:, active_slice],
+                               gt_full_px[:, active_slice],
+                               both_vis_no.astype(np.float32))
+
+        # BG EPE (only for full): CoTracker(gen_full) vs static anchor position
+        bg_static_px = sc["bg_static_pos"] * np.array([W, H])[None, None, :]
+        bg_vis_mask = (tk_full_vis[:, bg_slice] > 0.5).astype(np.float32)
+        epe_bg = epe_px(tk_full_px[:, bg_slice], bg_static_px, bg_vis_mask) if sc["n_bg"] > 0 else float("nan")
+
+        delta = epe_active_no - epe_active_full  # positive => tracks helped
+
+        # ── Save + log videos ──
+        overlay_full = overlay_tracks(gen_full, sc["cond_tracks"], sc["cond_vis"], palette)
+        overlay_no = overlay_tracks(gen_no, sc["cond_tracks"], sc["cond_vis"], palette)
+
+        # Also show the CoTracker-extracted tracks over the full generation — this is
+        # what the model *actually* produced motion-wise.
+        cotr_tracks_norm = tk_full_px / np.array([W, H])[None, None, :]
+        overlay_cotr = overlay_tracks(gen_full, cotr_tracks_norm, tk_full_vis.astype(np.float32), palette)
+
+        clip_dir = Path(args.out_dir) / f"clip_{idx:03d}"
+        clip_dir.mkdir(parents=True, exist_ok=True)
+        p_full = clip_dir / "gen_full_overlay.mp4"
+        p_no = clip_dir / "gen_notrack_overlay.mp4"
+        p_cotr = clip_dir / "gen_full_cotracker.mp4"
+        imageio.mimsave(p_full, overlay_full, fps=args.fps, macro_block_size=1)
+        imageio.mimsave(p_no, overlay_no, fps=args.fps, macro_block_size=1)
+        imageio.mimsave(p_cotr, overlay_cotr, fps=args.fps, macro_block_size=1)
+
+        print(f"[eval] clip {idx}: EPE_active_full={epe_active_full:.2f}  "
+              f"EPE_active_notrack={epe_active_no:.2f}  DELTA={delta:.2f}  "
+              f"EPE_bg={epe_bg:.2f}", flush=True)
+
+        # per-clip logs
+        entry = {
+            "clip": idx,
+            "epe_active_full_px": epe_active_full,
+            "epe_active_notrack_px": epe_active_no,
+            "delta_px": delta,
+            "epe_bg_px": epe_bg,
+            "n_active": sc["n_active"],
+            "n_bg": sc["n_bg"],
+            "sec_gen_full": t_full,
+            "sec_gen_notrack": t_no,
+            "caption": s["caption"][:120],
+        }
+        per_clip.append(entry)
+        wandb.log({
+            f"clip_{idx:03d}/gen_full": wandb.Video(str(p_full), fps=args.fps, format="mp4"),
+            f"clip_{idx:03d}/gen_notrack": wandb.Video(str(p_no), fps=args.fps, format="mp4"),
+            f"clip_{idx:03d}/gen_full_cotracker": wandb.Video(str(p_cotr), fps=args.fps, format="mp4"),
+            "clip": idx,
+            "epe_active_full_px": epe_active_full,
+            "epe_active_notrack_px": epe_active_no,
+            "delta_px": delta,
+            "epe_bg_px": epe_bg,
+        })
+
+    # ── Aggregate + summary ──
+    if per_clip:
+        arr = lambda k: np.array([r[k] for r in per_clip if not np.isnan(r[k])])
+        summary = {
+            "n_clips": len(per_clip),
+            "mean_epe_active_full_px": float(np.mean(arr("epe_active_full_px"))),
+            "median_epe_active_full_px": float(np.median(arr("epe_active_full_px"))),
+            "mean_epe_active_notrack_px": float(np.mean(arr("epe_active_notrack_px"))),
+            "median_epe_active_notrack_px": float(np.median(arr("epe_active_notrack_px"))),
+            "mean_delta_px": float(np.mean(arr("delta_px"))),
+            "median_delta_px": float(np.median(arr("delta_px"))),
+            "mean_epe_bg_px": float(np.mean(arr("epe_bg_px"))),
+        }
+        (Path(args.out_dir) / f"{args.wandb_run_name}_summary.json").write_text(json.dumps({
+            "summary": summary, "per_clip": per_clip}, indent=2))
+        wandb.log({f"summary/{k}": v for k, v in summary.items()})
+        wandb.summary.update(summary)
+        print("=" * 70)
+        print(f"[eval] summary ({len(per_clip)} clips):")
+        for k, v in summary.items():
+            print(f"  {k}: {v:.3f}" if isinstance(v, float) else f"  {k}: {v}")
+        print("=" * 70)
+
+    wandb.finish()
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/extract_tracks.py

@@ -0,0 +1,536 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stage 0c: extract dense point tracks from generated videos with CoTracker v3.
+
+For each .mp4 produced by ``generate_videos.py`` we run CoTracker v3 (``cotracker3_offline``)
+with a ``grid_size``x``grid_size`` regular query grid (default 50x50 = 2500 points) and save
+the per-frame tracks + visibility. We then patch ``points_path`` (absolute) into the
+manifest so the future points-aware preprocess task can find them (mirrors how MatrixGame2
+references ``action_path``).
+
+Tracks are stored in ORIGINAL video pixel coordinates. The full 2500-point grid + visibility
+are kept; the trainer samples 1-200 points per step.
+
+If ``--detect-entries`` is set, FastSAM detects objects entering after frame 0 (frames are
+segmented in batched forwards). Grid points landing on each new object are tracked from its
+entry frame T_entry — all entry events share a single extra CoTracker pass (chunked if the
+combined query count exceeds grid_size^2) — and replace dead background slots (those
+permanently occluded from T_entry onwards), keeping N = grid_size^2 throughout.
+
+Run on a GPU node (never the login node), e.g.:
+
+    srun --jobid=<shao_wm jobid> --overlap --ntasks=1 \\
+        .venv/bin/python data_pipeline/extract_tracks.py \\
+        --data-dir /mnt/weka/home/hao.zhang/shao/data/motion_pipeline/wan22_t2v_720p \\
+        --grid-size 50 --detect-entries
+
+torch.hub note: prefetch once on the login node (internet) so the shared cache is warm:
+    .venv/bin/python -c "import torch; torch.hub.load('facebookresearch/co-tracker','cotracker3_offline')"
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import queue
+import threading
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+HUB_REPO = "facebookresearch/co-tracker"
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--data-dir", type=Path, required=True, help="Dataset root from generate_videos.py.")
+    p.add_argument("--videos-subdir", type=str, default="videos")
+    p.add_argument("--out-subdir", type=str, default="tracks")
+    p.add_argument("--manifest", type=str, default="videos2caption.json")
+    p.add_argument("--grid-size", type=int, default=50, help="NxN query grid (N*N points).")
+    p.add_argument("--model", type=str, default="cotracker3_offline")
+    p.add_argument("--device", type=str, default="cuda")
+    p.add_argument("--downscale", type=float, default=1.0,
+                   help="Run tracking at this spatial scale (coords rescaled back to original px). "
+                        "Use <1.0 (e.g. 0.5) if full-res OOMs.")
+    p.add_argument("--limit", type=int, default=None, help="Only process the first N videos (for smoke tests).")
+    p.add_argument("--index", type=int, nargs="+", default=None, metavar="IDX",
+                   help="Only process videos with these indices (e.g. --index 4 7 12).")
+    p.add_argument("--rank", type=int, default=0, help="GPU rank for sharding (0-indexed).")
+    p.add_argument("--world-size", type=int, default=1, help="Total number of parallel processes.")
+    p.add_argument("--force", action="store_true", help="Re-extract even if .npz already exists.")
+    p.add_argument("--verbose", action="store_true", help="Print per-mask debug info for entry detection.")
+    # Entry-frame detection
+    p.add_argument("--detect-entries", action="store_true",
+                   help="Detect objects entering after frame 0 and replace dead background slots.")
+    p.add_argument("--sam-model", type=str, default="FastSAM-s.pt")
+    p.add_argument("--sam-conf", type=float, default=0.75)
+    p.add_argument("--sam-iou", type=float, default=0.9)
+    p.add_argument("--sam-imgsz", type=int, default=1024)
+    p.add_argument("--sam-batch", type=int, default=16,
+                   help="Frames per batched FastSAM forward during entry detection.")
+    p.add_argument("--amp", action="store_true",
+                   help="Run CoTracker under bf16 autocast (~1.5-2x faster, slightly different coords).")
+    p.add_argument("--compile", action="store_true",
+                   help="torch.compile the main CoTracker pass (fixed shape; compiled once per worker). "
+                        "The variable-size entry pass stays eager to avoid recompiles.")
+    p.add_argument("--prefetch", type=int, default=2,
+                   help="Decode up to N videos ahead on a background thread so the GPU never waits "
+                        "on video decode (0 = disable).")
+    # Fused Stage 4
+    p.add_argument("--segment", action="store_true",
+                   help="Fused Stage 4: also compute object_ids/n_objects/track_weights and the "
+                        "vis-override sweep in this pass, reusing the decoded video and the "
+                        "entry-detection FastSAM masks. Uses --sam-conf/--sam-iou/--sam-imgsz; "
+                        "masks are unfiltered (no min-area-frac/max-masks).")
+    p.add_argument("--vis-override-every", type=int, default=3,
+                   help="(with --segment) run FastSAM every N frames and set vis=True for object "
+                        "points inside masks; 0 disables.")
+    p.add_argument("--viz", action="store_true",
+                   help="(with --segment) render a track-overlay mp4 after each video.")
+    p.add_argument("--viz-dir", type=str, default=None,
+                   help="Output directory for viz mp4s (default: <data-dir>/viz).")
+    p.add_argument("--entry-sample-every", type=int, default=5,
+                   help="Check for new objects every N frames.")
+    p.add_argument("--entry-new-area", type=float, default=0.3,
+                   help="A mask triggers entry detection if at least this fraction of its area "
+                        "is not covered by any frame-0 mask (catches partially-entering objects).")
+    p.add_argument("--entry-min-area", type=float, default=0.005,
+                   help="Min area fraction for a new-object mask to be considered.")
+    args = p.parse_args()
+    if args.viz and not args.segment:
+        p.error("--viz requires --segment (object IDs are needed for the overlay)")
+    return args
+
+
+def load_cotracker(model_name: str, device: str):
+    """Load CoTracker via torch.hub, preferring the warm local cache (offline-safe)."""
+    hub_dir = Path(torch.hub.get_dir())
+    local = hub_dir / (HUB_REPO.replace("/", "_") + "_main")
+    try:
+        if local.exists():
+            model = torch.hub.load(str(local), model_name, source="local", trust_repo=True)
+        else:
+            model = torch.hub.load(HUB_REPO, model_name, trust_repo=True)
+    except Exception as e:  # noqa: BLE001
+        raise RuntimeError(
+            f"Failed to load CoTracker ({e!r}). On a node without internet, prefetch on the "
+            "login node first: .venv/bin/python -c \"import torch; "
+            "torch.hub.load('facebookresearch/co-tracker','cotracker3_offline')\""
+        ) from e
+    return model.to(device).eval()
+
+
+def read_video(path: Path) -> tuple[torch.Tensor, int, int]:
+    """Return (video[1,T,C,H,W] float 0-255, H, W)."""
+    try:
+        from decord import VideoReader, cpu
+        vr = VideoReader(str(path), ctx=cpu(0))
+        frames = vr.get_batch(list(range(len(vr)))).asnumpy()  # (T, H, W, C) uint8
+    except Exception:  # noqa: BLE001
+        import imageio
+        reader = imageio.get_reader(str(path), format="ffmpeg")
+        frames = np.stack([np.asarray(f) for f in reader], axis=0)
+        reader.close()
+    vid = torch.from_numpy(np.ascontiguousarray(frames))
+    h, w = int(vid.shape[1]), int(vid.shape[2])
+    video = vid.permute(0, 3, 1, 2).unsqueeze(0).float()  # (1,T,C,H,W)
+    return video, h, w
+
+
+@torch.no_grad()
+def track_one(model, video: torch.Tensor, grid_size: int, downscale: float, device: str,
+              amp: bool = False) -> tuple[np.ndarray, np.ndarray]:
+    """Return tracks (T,N,2) in original pixel coords and visibility (T,N)."""
+    _, t, c, h, w = video.shape
+    track_video = video
+    if downscale != 1.0:
+        sh, sw = max(1, int(round(h * downscale))), max(1, int(round(w * downscale)))
+        track_video = F.interpolate(video[0], size=(sh, sw), mode="bilinear", align_corners=False).unsqueeze(0)
+    else:
+        sh, sw = h, w
+
+    with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp and device.startswith("cuda")):
+        pred_tracks, pred_vis = model(track_video.to(device), grid_size=grid_size)
+    tracks = pred_tracks[0].float().cpu().numpy()
+    vis = pred_vis[0].cpu().numpy()
+
+    if downscale != 1.0:
+        tracks[..., 0] *= w / float(sw)
+        tracks[..., 1] *= h / float(sh)
+    return tracks.astype(np.float32), vis
+
+
+@torch.no_grad()
+def track_with_queries(model, video: torch.Tensor, queries_txy: np.ndarray, downscale: float,
+                       device: str, H: int, W: int, amp: bool = False) -> tuple[np.ndarray, np.ndarray]:
+    """Track explicit query points. queries_txy: [K,3] as (t, x, y) in original pixel coords."""
+    track_video = video
+    sh, sw = H, W
+    q = queries_txy.astype(np.float32).copy()
+    if downscale != 1.0:
+        sh = max(1, int(round(H * downscale)))
+        sw = max(1, int(round(W * downscale)))
+        track_video = F.interpolate(video[0], size=(sh, sw), mode="bilinear", align_corners=False).unsqueeze(0)
+        q[:, 1] *= sw / float(W)
+        q[:, 2] *= sh / float(H)
+
+    q_tensor = torch.from_numpy(q).unsqueeze(0).to(device)  # [1, K, 3]
+    with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp and device.startswith("cuda")):
+        pred_tracks, pred_vis = model(track_video.to(device), queries=q_tensor)
+    tracks = pred_tracks[0].float().cpu().numpy()
+    vis = pred_vis[0].cpu().numpy()
+
+    if downscale != 1.0:
+        tracks[..., 0] *= W / float(sw)
+        tracks[..., 1] *= H / float(sh)
+    return tracks.astype(np.float32), vis
+
+
+def make_grid_queries(grid_size: int, H: int, W: int, frame_t: int) -> np.ndarray:
+    """Generate a grid_size x grid_size uniform query grid at frame_t. Returns [N,3] (t,x,y)."""
+    ys = np.linspace(0, H - 1, grid_size)
+    xs = np.linspace(0, W - 1, grid_size)
+    xx, yy = np.meshgrid(xs, ys)
+    pts_xy = np.column_stack([xx.ravel(), yy.ravel()])  # [N, 2]
+    return np.column_stack([np.full(len(pts_xy), float(frame_t)), pts_xy]).astype(np.float32)
+
+
+
+def fastsam_masks_batch(sam_model, frames_rgb: list[np.ndarray], conf: float, iou: float,
+                        imgsz: int, H: int, W: int, batch: int = 16) -> list[np.ndarray]:
+    """Run FastSAM on uint8 HxWx3 RGB frames in batched forwards, return bool masks [M,H,W] per frame."""
+    out: list[np.ndarray] = []
+    for s in range(0, len(frames_rgb), max(1, batch)):
+        res = sam_model(frames_rgb[s:s + max(1, batch)], device="cuda", retina_masks=True,
+                        imgsz=imgsz, conf=conf, iou=iou, verbose=False)
+        for r in res:
+            if r.masks is None:
+                out.append(np.zeros((0, H, W), bool))
+                continue
+            masks = r.masks.data.cpu().numpy().astype(bool)
+            if masks.shape[0] and masks.shape[1:] != (H, W):
+                import cv2
+                masks = np.stack([
+                    cv2.resize(m.astype(np.uint8), (W, H), interpolation=cv2.INTER_NEAREST).astype(bool)
+                    for m in masks
+                ])
+            out.append(masks)
+    return out
+
+
+def max_iou_with_set(mask: np.ndarray, others: np.ndarray) -> float:
+    """Max IOU of a single [H,W] bool mask against a set [M,H,W]."""
+    if others.shape[0] == 0:
+        return 0.0
+    inter = (mask & others).reshape(others.shape[0], -1).sum(1)
+    union = (mask | others).reshape(others.shape[0], -1).sum(1)
+    return float(np.where(union > 0, inter / union, 0.0).max())
+
+
+def detect_entry_events(video: torch.Tensor, sam_model, H: int, W: int, conf: float, iou: float,
+                        imgsz: int, sample_every: int, new_area_thresh: float,
+                        min_area_frac: float, sam_batch: int = 16
+                        ) -> tuple[dict[int, np.ndarray], np.ndarray, dict[int, np.ndarray]]:
+    """Return ({frame_t: new_masks [M,H,W]}, masks0_union [H,W], {frame_t: all masks}) for
+    frames where new regions appear that weren't covered by any frame-0 mask. The third
+    element caches every segmented frame's raw masks for reuse by the fused --segment pass.
+
+    Uses new-area fraction rather than IOU so partially-entering objects (partially visible
+    at frame 0, more visible later) are correctly detected: only the newly visible region
+    triggers detection and receives replacement tracks.
+    """
+    T = video.shape[1]
+    min_area_px = int(min_area_frac * H * W)
+
+    # One batched FastSAM sweep over frame 0 + all sampled frames (the per-frame
+    # mask filtering below is sequential, but the model forwards are independent).
+    sample_ts = list(range(sample_every, T, sample_every))
+    frames = [video[0, t].permute(1, 2, 0).numpy().astype(np.uint8) for t in [0, *sample_ts]]
+    all_masks = fastsam_masks_batch(sam_model, frames, conf, iou, imgsz, H, W, batch=sam_batch)
+    masks_by_frame = dict(zip([0, *sample_ts], all_masks))
+    masks0 = all_masks[0]
+    # Exclude large background masks (table surface, floor, walls) from masks0_union.
+    # Only object-sized masks define "known territory" — background covers the whole frame
+    # and would suppress detection of the ball moving to a new position on that surface.
+    if masks0.shape[0] > 0:
+        areas = masks0.reshape(masks0.shape[0], -1).sum(1)
+        object_masks0 = masks0[areas < 0.4 * H * W]
+        masks0_union = object_masks0.any(axis=0) if object_masks0.shape[0] > 0 else np.zeros((H, W), bool)
+    else:
+        masks0_union = np.zeros((H, W), bool)
+
+    # Tracks newly-covered regions across entry events to avoid re-detecting the same
+    # entering object at multiple sample frames.
+    claimed_new = np.zeros((H, W), bool)
+
+    entry_events: dict[int, np.ndarray] = {}
+    for frame_t, masks_t in zip(sample_ts, all_masks[1:]):
+        if masks_t.shape[0] == 0:
+            continue
+
+        new_masks = []
+        for m in masks_t:
+            area = int(m.sum())
+            if area < min_area_px:
+                continue
+            # New area: part of this mask not present in ANY frame-0 mask.
+            new_region = m & ~masks0_union
+            new_area_frac = float(new_region.sum()) / max(area, 1)
+            if new_area_frac < new_area_thresh:
+                continue
+            # Skip if this new region was already claimed by a prior entry event.
+            if float((new_region & claimed_new).sum()) / max(int(new_region.sum()), 1) > 0.5:
+                continue
+            new_masks.append(m)
+            # Dilate claimed region by ~2x the object radius so a fast-moving
+            # object doesn't re-trigger entry detection on subsequent sample frames.
+            radius = int(np.sqrt(int(new_region.sum()) / np.pi))
+            dil = max(1, radius * 2)
+            ys, xs = np.where(new_region)
+            y0, y1 = max(0, ys.min() - dil), min(H, ys.max() + dil + 1)
+            x0, x1 = max(0, xs.min() - dil), min(W, xs.max() + dil + 1)
+            claimed_new[y0:y1, x0:x1] = True
+
+        if new_masks:
+            entry_events[frame_t] = np.stack(new_masks)
+
+    return entry_events, masks0_union, masks_by_frame
+
+
+def decoded_videos(todo: list[tuple[int, Path, Path]], prefetch: int):
+    """Yield (k, vpath, out_path, video, h, w), decoding up to `prefetch` videos ahead
+    on a background thread (decord/av release the GIL, so decode overlaps GPU compute)."""
+    if prefetch <= 0:
+        for k, vpath, out_path in todo:
+            yield (k, vpath, out_path, *read_video(vpath))
+        return
+
+    q: queue.Queue = queue.Queue(maxsize=prefetch)
+
+    def producer() -> None:
+        for k, vpath, out_path in todo:
+            try:
+                payload = read_video(vpath)
+            except Exception as e:  # noqa: BLE001
+                payload = e
+            q.put((k, vpath, out_path, payload))
+        q.put(None)
+
+    threading.Thread(target=producer, daemon=True).start()
+    while (item := q.get()) is not None:
+        k, vpath, out_path, payload = item
+        if isinstance(payload, Exception):
+            # One unreadable clip must not take down the worker and everything it has
+            # left to process -- real-world shards contain corrupt/truncated files.
+            print(f"[track] [{k}] {vpath.name}: DECODE FAILED ({payload}), skipping", flush=True)
+            continue
+        yield (k, vpath, out_path, *payload)
+
+
+def patch_manifest(manifest_path: Path, stem_to_points: dict[str, Path]) -> int:
+    if not manifest_path.exists():
+        return 0
+    items = json.loads(manifest_path.read_text())
+    patched = 0
+    for item in items:
+        stem = Path(item.get("path", "")).stem
+        pts = stem_to_points.get(stem)
+        if pts is not None:
+            item["points_path"] = str(pts.resolve())
+            patched += 1
+    tmp = manifest_path.with_suffix(f".json.tmp{os.getpid()}")   # pid-unique: never shared
+    tmp.write_text(json.dumps(items, indent=2))
+    tmp.replace(manifest_path)
+    return patched
+
+
+def main() -> None:
+    args = parse_args()
+    videos_dir = args.data_dir / args.videos_subdir
+    out_dir = args.data_dir / args.out_subdir
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    videos = sorted(videos_dir.glob("*.mp4"))
+    if args.index is not None:
+        wanted = {f"vid_{i:06d}.mp4" for i in args.index}
+        videos = [v for v in videos if v.name in wanted]
+    if args.limit is not None:
+        videos = videos[:args.limit]
+    all_videos = videos                      # pre-shard list: only rank 0 patches the manifest,
+    if args.world_size > 1:                  # and it must cover every rank's videos
+        videos = videos[args.rank::args.world_size]
+    if not videos:
+        print(f"[track] no videos found in {videos_dir}", flush=True)
+        return
+
+    cotracker = load_cotracker(args.model, args.device)
+    # Compile only the fixed-shape main pass; entry-pass query counts vary per video and
+    # would recompile every time, so track_with_queries keeps the eager model.
+    cotracker_main = torch.compile(cotracker) if args.compile else cotracker
+    print(f"[track] loaded {args.model}; {len(videos)} videos, grid={args.grid_size}x{args.grid_size}"
+          f"{' (compiled main pass)' if args.compile else ''}", flush=True)
+
+    sam_model = None
+    if args.detect_entries or args.segment:
+        from ultralytics import FastSAM
+        sam_model = FastSAM(args.sam_model)
+        print(f"[track] entries={args.detect_entries} segment={args.segment}: sam={args.sam_model}, "
+              f"conf={args.sam_conf}, every={args.entry_sample_every}f", flush=True)
+
+    stem_to_fps: dict[str, float] = {}
+    if args.viz:
+        mpath = args.data_dir / args.manifest
+        items = json.loads(mpath.read_text()) if mpath.exists() else []
+        stem_to_fps = {Path(it.get("path", "")).stem: it.get("fps", 24) for it in items}
+
+    # points_path is deterministic, so rank 0 can record it for every video without
+    # doing their work -- avoids all ranks rewriting the manifest at once.
+    stem_to_points = {v.stem: out_dir / f"{v.stem}.npz" for v in all_videos}
+    todo: list[tuple[int, Path, Path]] = []
+    for k, vpath in enumerate(videos, 1):
+        out_path = out_dir / f"{vpath.stem}.npz"
+        if out_path.exists() and not args.force:
+            continue
+        todo.append((k, vpath, out_path))
+
+    for k, vpath, out_path, video, h, w in decoded_videos(todo, args.prefetch):
+        tracks, vis = track_one(cotracker_main, video, args.grid_size, args.downscale, args.device, amp=args.amp)
+
+        n_replaced = 0
+        masks_cache: dict[int, np.ndarray] = {}
+        if args.detect_entries:
+            entry_events, masks0_union, masks_cache = detect_entry_events(
+                video, sam_model, h, w,
+                conf=args.sam_conf, iou=args.sam_iou, imgsz=args.sam_imgsz,
+                sample_every=args.entry_sample_every,
+                new_area_thresh=args.entry_new_area,
+                min_area_frac=args.entry_min_area,
+                sam_batch=args.sam_batch,
+            )
+
+            # Build queries for all (entry frame, mask) pairs at once: only grid points
+            # landing inside each new region are tracked, instead of a full grid_size^2
+            # pass per mask. Groups keep disjoint column ranges [qa, qb) in the batch.
+            groups: list[tuple[int, int, np.ndarray, np.ndarray, int, int]] = []
+            q_parts: list[np.ndarray] = []
+            n_q = 0
+            for frame_t, new_masks in sorted(entry_events.items()):
+                grid_q = make_grid_queries(args.grid_size, h, w, frame_t)
+                gxi = np.clip(grid_q[:, 1].round().astype(int), 0, w - 1)
+                gyi = np.clip(grid_q[:, 2].round().astype(int), 0, h - 1)
+                for mi, new_mask in enumerate(new_masks):
+                    new_region = new_mask & ~masks0_union
+                    q = grid_q[new_region[gyi, gxi]]
+                    groups.append((frame_t, mi, new_mask, new_region, n_q, n_q + len(q)))
+                    q_parts.append(q)
+                    n_q += len(q)
+
+            e_tracks = e_vis = None
+            if n_q > 0:
+                all_q = np.concatenate(q_parts, axis=0)
+                max_q = args.grid_size * args.grid_size  # bound memory to the main pass
+                et_parts, ev_parts = [], []
+                for s in range(0, n_q, max_q):
+                    ct, cv = track_with_queries(
+                        cotracker, video, all_q[s:s + max_q], args.downscale,
+                        args.device, h, w, amp=args.amp)
+                    et_parts.append(ct)
+                    ev_parts.append(cv)
+                e_tracks = np.concatenate(et_parts, axis=1)
+                e_vis = np.concatenate(ev_parts, axis=1)
+
+            for frame_t, mi, new_mask, new_region, qa, qb in groups:
+                # Covered slots: original frame-0 grid points whose position at
+                # T_entry falls inside the new object's region.
+                orig_xi = np.clip(tracks[frame_t, :, 0].round().astype(int), 0, w - 1)
+                orig_yi = np.clip(tracks[frame_t, :, 1].round().astype(int), 0, h - 1)
+                covered_dst = np.where(new_region[orig_yi, orig_xi])[0]
+                dead_dst = covered_dst if covered_dst.size > 0 else \
+                    np.where((~vis[frame_t:].astype(bool)).all(axis=0))[0]
+                if dead_dst.size == 0:
+                    continue
+
+                n_src = qb - qa
+                n = min(n_src, len(dead_dst))
+                if n > 0:
+                    e_vis[:frame_t, qa:qb] = False  # object didn't exist before entry frame
+                    dst = dead_dst[:n]
+                    tracks[:, dst] = e_tracks[:, qa:qa + n]
+                    vis[:, dst] = e_vis[:, qa:qa + n]
+                    # CoTracker can mark points invisible even at their query frame
+                    # when the object is entering from the edge. Force vis=True at
+                    # T_entry so segment_tracks sees these as first-visible there.
+                    vis[frame_t, dst] = True
+                    n_replaced += n
+                if args.verbose:
+                    mask_area = int(new_mask.sum())
+                    new_region_area = int(new_region.sum())
+                    vis_at_entry = int(vis[frame_t, dst].sum()) if n > 0 else 0
+                    print(f"  [entry] t={frame_t} mask#{mi}: area={mask_area}px "
+                          f"new_region={new_region_area}px "
+                          f"object_src={n_src} covered_dst={len(covered_dst)} "
+                          f"dead_dst={len(dead_dst)} -> replacing {n}, "
+                          f"vis[{frame_t}, dst].sum()={vis_at_entry}", flush=True)
+
+            frames_str = ",".join(str(t) for t in sorted(entry_events)) if entry_events else "none"
+            print(f"[track] [{k}/{len(videos)}] {vpath.name}: "
+                  f"{len(entry_events)} entry event(s) at frames [{frames_str}], "
+                  f"{n_replaced} slots replaced", flush=True)
+
+        # Fused Stage 4: object IDs + vis override + track weights, reusing the decoded
+        # video and any masks already computed by entry detection.
+        seg_extra: dict[str, np.ndarray] = {}
+        if args.segment:
+            import segment_tracks as seg
+
+            def get_masks(frame_ts: list[int], _video=video, _cache=masks_cache,
+                          _h=h, _w=w) -> dict[int, np.ndarray]:
+                missing = [t for t in frame_ts if t not in _cache]
+                if missing:
+                    frames_np = [_video[0, t].permute(1, 2, 0).numpy().astype(np.uint8) for t in missing]
+                    new_masks = fastsam_masks_batch(sam_model, frames_np, args.sam_conf, args.sam_iou,
+                                                    args.sam_imgsz, _h, _w, batch=args.sam_batch)
+                    _cache.update(zip(missing, new_masks))
+                return {t: _cache[t] for t in frame_ts}
+
+            vis = vis.astype(bool)
+            oid, n_objects, vis, weights, n_overrides, uframes = seg.segment_tracks_arrays(
+                tracks, vis, h, w, get_masks, args.vis_override_every, verbose=args.verbose)
+            seg_extra = dict(object_ids=oid, n_objects=np.int64(n_objects), track_weights=weights)
+            print(f"[track] [{k}/{len(videos)}] {vpath.name} seg: {n_objects} objs across "
+                  f"{len(uframes)} frames, {int((oid >= 0).sum())}/{oid.shape[0]} pts labeled, "
+                  f"{n_overrides} vis overrides", flush=True)
+
+            if args.viz:
+                viz_dir = Path(args.viz_dir) if args.viz_dir else (args.data_dir / "viz")
+                stem_dir = viz_dir / vpath.stem
+                stem_dir.mkdir(parents=True, exist_ok=True)
+                frames_np = video[0].permute(0, 2, 3, 1).numpy().astype(np.uint8)  # [T,H,W,3]
+                seg.render_viz(frames_np[:tracks.shape[0]], tracks, vis, oid, stem_dir / "tracks.mp4",
+                               fps=int(stem_to_fps.get(vpath.stem, 24)))
+                print(f"  viz -> {stem_dir}/", flush=True)
+
+        tmp = out_path.with_name(out_path.stem + ".tmp.npz")
+        np.savez(
+            tmp, tracks=tracks, visibility=vis,
+            grid_size=args.grid_size, height=h, width=w,
+            num_frames=tracks.shape[0],
+            **seg_extra,
+        )
+        tmp.replace(out_path)
+        print(f"[track] [{k}/{len(videos)}] {vpath.name} -> {out_path.name} "
+              f"tracks={tracks.shape} vis={vis.shape}", flush=True)
+
+    # Only rank 0 writes: concurrent read-modify-write from every rank corrupted the
+    # manifest (interleaved writes to a shared temp file -> invalid JSON).
+    if args.rank == 0:
+        n = patch_manifest(args.data_dir / args.manifest, stem_to_points)
+        print(f"[track] done; patched points_path into {n} manifest entries", flush=True)
+    else:
+        print(f"[track] done (rank {args.rank}; manifest patched by rank 0)", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/extract_tracks_mp.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Data-parallel CoTracker v3 track extraction for large video sets (OpenVid-1M).
+
+CoTracker batching does NOT help (s/video is flat, memory linear), so we scale by
+DATA PARALLELISM: launch many single-video workers, several per GPU, across nodes.
+Each worker processes video_list[shard::num_shards] on one pinned GPU and is idempotent.
+
+Sharding + GPU pinning come from Slurm env by default:
+  shard      = SLURM_PROCID     (0..num_shards-1)
+  num_shards = SLURM_NTASKS
+  gpu        = SLURM_LOCALID % gpus_per_node
+so `srun --ntasks-per-node=(gpus*procs_per_gpu)` gives procs_per_gpu workers per GPU.
+
+Resamples each clip to --fps / --num-frames and resizes to --height x --width (720p)
+before tracking, matching the training spec.
+"""
+from __future__ import annotations
+import argparse, os, time
+from pathlib import Path
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+HUB = "facebookresearch/co-tracker"
+
+
+def load_cotracker(device: str):
+    local = Path(torch.hub.get_dir()) / (HUB.replace("/", "_") + "_main")
+    if local.exists():
+        m = torch.hub.load(str(local), "cotracker3_offline", source="local", trust_repo=True)
+    else:
+        m = torch.hub.load(HUB, "cotracker3_offline", trust_repo=True)
+    return m.to(device).eval()
+
+
+def read_clip(path: str, fps: int, num_frames: int, H: int, W: int,
+              min_height: int = 0, aspect: float | None = None, aspect_tol: float = 0.1):
+    """Decode (PyAV — decord has no aarch64 wheels, opencv needs GUI libs), resample
+    to `fps`, take `num_frames`, resize to HxW. Sequential decode so it is codec-robust
+    for OpenVid's varied encodings. Enforces resolution/aspect by probing stream metadata
+    (fast, no decode). Returns (1,T,C,H,W) on success, or a reason string
+    ('lowres' | 'aspect' | 'short') on reject."""
+    import av
+    container = av.open(str(path))
+    try:
+        stream = container.streams.video[0]
+        nh, nw = int(stream.height or 0), int(stream.width or 0)
+        if min_height and nh and nh < min_height:
+            return "lowres"
+        if aspect is not None and nw and nh and abs((nw / nh) - aspect) > aspect_tol:
+            return "aspect"
+        native = float(stream.average_rate) if stream.average_rate else float(fps)
+        step = native / float(fps)
+        idxs = [int(round(i * step)) for i in range(num_frames)]
+        picked, wi, cur = [], 0, 0
+        for frame in container.decode(video=0):
+            img = None
+            while wi < len(idxs) and idxs[wi] == cur:   # handles repeated idxs (fps up-sample)
+                if img is None:
+                    img = frame.to_ndarray(format="rgb24")  # (H,W,3) uint8
+                picked.append(img); wi += 1
+            cur += 1
+            if wi >= len(idxs):
+                break
+    finally:
+        container.close()
+    if wi < len(idxs):
+        return "short"  # video too short at target fps
+    arr = np.ascontiguousarray(np.stack(picked))  # (T,H,W,C) RGB uint8
+    vid = torch.from_numpy(arr).permute(0, 3, 1, 2).unsqueeze(0).float()  # 1,T,C,h,w
+    vid = F.interpolate(vid[0], size=(H, W), mode="bilinear", align_corners=False).unsqueeze(0)
+    return vid
+
+
+@torch.no_grad()
+def track(model, video, grid, device):
+    tr, vis = model(video.to(device), grid_size=grid)
+    return tr[0].float().cpu().numpy(), vis[0].cpu().numpy()
+
+
+def save_clip(video, path, fps):
+    """Write the resampled 121f/720p clip so the training video IS the tracked frames."""
+    import imageio
+    arr = video[0].permute(0, 2, 3, 1).clamp(0, 255).byte().cpu().numpy()  # (T,C,H,W)->(T,H,W,C)
+    imageio.mimwrite(path, arr, fps=fps, codec="libx264", quality=7, macro_block_size=1)
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--video-list", required=True, help="Text file: one video path per line.")
+    ap.add_argument("--out-dir", required=True, help="Where <stem>.npz tracks are written.")
+    ap.add_argument("--grid-size", type=int, default=50)
+    ap.add_argument("--num-shards", type=int, default=int(os.environ.get("SLURM_NTASKS", "1")))
+    ap.add_argument("--shard", type=int, default=int(os.environ.get("SLURM_PROCID", "0")))
+    ap.add_argument("--gpus-per-node", type=int, default=4)
+    ap.add_argument("--fps", type=int, default=24)
+    ap.add_argument("--num-frames", type=int, default=121)
+    ap.add_argument("--height", type=int, default=720)
+    ap.add_argument("--width", type=int, default=1280)
+    ap.add_argument("--min-height", type=int, default=720, help="skip clips whose native height < this (0=off)")
+    ap.add_argument("--aspect", type=float, default=16.0 / 9.0, help="target native W/H aspect")
+    ap.add_argument("--aspect-tol", type=float, default=0.12, help="allowed |native_aspect - target| (large=off)")
+    ap.add_argument("--clips-dir", default=None, help="also save the resampled 121f/720p clip here (aligned to tracks)")
+    a = ap.parse_args()
+
+    local = int(os.environ.get("SLURM_LOCALID", str(a.shard)))
+    gpu = local % a.gpus_per_node
+    torch.cuda.set_device(gpu)
+    device = f"cuda:{gpu}"
+    Path(a.out_dir).mkdir(parents=True, exist_ok=True)
+    if a.clips_dir:
+        Path(a.clips_dir).mkdir(parents=True, exist_ok=True)
+
+    vids = [l.strip() for l in open(a.video_list) if l.strip()]
+    mine = vids[a.shard::a.num_shards]
+    model = load_cotracker(device)
+    print(f"[shard {a.shard}/{a.num_shards}] gpu={gpu} localid={local} videos={len(mine)}", flush=True)
+
+    t0 = time.time(); done = 0; err = 0; rej = {"short": 0, "lowres": 0, "aspect": 0}
+    for vp in mine:
+        out = Path(a.out_dir) / f"{Path(vp).stem}.npz"
+        if out.exists():
+            continue
+        try:
+            res = read_clip(vp, a.fps, a.num_frames, a.height, a.width,
+                            a.min_height, a.aspect, a.aspect_tol)
+            if isinstance(res, str):
+                rej[res] = rej.get(res, 0) + 1; continue
+            tr, vis = track(model, res, a.grid_size, device)
+            if a.clips_dir:
+                save_clip(res, str(Path(a.clips_dir) / f"{Path(vp).stem}.mp4"), a.fps)
+            tmp = out.with_suffix(".tmp.npz")
+            np.savez(tmp, tracks=tr.astype(np.float32), visibility=vis, grid_size=a.grid_size,
+                     height=a.height, width=a.width, num_frames=tr.shape[0], fps=a.fps)
+            tmp.replace(out); done += 1
+        except Exception as e:  # noqa: BLE001
+            err += 1
+            print(f"[shard {a.shard}] ERR {Path(vp).name}: {repr(e)[:110]}", flush=True)
+    dt = time.time() - t0
+    rate = done / dt if dt > 0 else 0.0
+    print(f"[shard {a.shard}] DONE done={done} rej={rej} err={err} in {dt:.1f}s -> {rate:.3f} vid/s", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/filter_clips.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Select clips that already match the target geometry; symlink them for downstream stages.
+
+Real shards are usually uniform but not guaranteed. Rather than re-encoding every clip to
+force conformity (a lossy no-op when the clip already matches -- measured 40 dB on the
+OpenVid shard, worse than the VAE round-trip's own distortion), this scans container
+metadata (fast: no frame decode) and links through only the clips that conform.
+
+Non-conforming clips are reported and listed in ``skipped_clips.json`` so nothing is
+silently dropped -- re-run them through ``resize_videos.py`` if you want them included.
+
+    python data_pipeline/filter_clips.py \\
+        --src-dir <root>/raw_videos --out-dir <root>/videos \\
+        --height 720 --width 1280 --num-frames 121
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import cv2
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--src-dir", type=Path, required=True)
+    p.add_argument("--out-dir", type=Path, required=True, help="Symlinks to conforming clips land here.")
+    p.add_argument("--height", type=int, required=True)
+    p.add_argument("--width", type=int, required=True)
+    p.add_argument("--num-frames", type=int, default=121,
+                   help="Required exact frame count (0 = don't check).")
+    p.add_argument("--report", type=str, default="skipped_clips.json",
+                   help="Written next to --out-dir; lists every skipped clip and why.")
+    p.add_argument("--needs-resize-list", type=str, default="needs_resize.txt",
+                   help="Written next to --out-dir; names of clips that are readable but at the "
+                        "wrong geometry, i.e. rescuable by resize_videos.py --include-list. "
+                        "Unreadable clips are excluded (nothing can rescue those).")
+    p.add_argument("--clean", action="store_true",
+                   help="Empty --out-dir of *.mp4 first (links AND regular files -- a stale real "
+                        "file would otherwise shadow the link and be used silently). --out-dir is "
+                        "a derived directory; never point it at original footage.")
+    return p.parse_args()
+
+
+def probe(path: Path) -> tuple[int, int, int]:
+    """Return (width, height, n_frames) from container metadata; (0,0,0) if unreadable."""
+    cap = cv2.VideoCapture(str(path))
+    if not cap.isOpened():
+        return (0, 0, 0)
+    wh = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
+          int(cap.get(cv2.CAP_PROP_FRAME_COUNT)))
+    cap.release()
+    return wh
+
+
+def main() -> None:
+    args = parse_args()
+    args.out_dir.mkdir(parents=True, exist_ok=True)
+    if args.clean:
+        for old in args.out_dir.glob("*.mp4"):
+            old.unlink()
+
+    clips = sorted(args.src_dir.glob("*.mp4"))
+    ok = 0
+    skipped: list[dict] = []
+    for c in clips:
+        w, h, n = probe(c)
+        if (w, h) == (0, 0):
+            skipped.append({"clip": c.name, "reason": "unreadable"})
+            continue
+        if (w, h) != (args.width, args.height):
+            skipped.append({"clip": c.name, "reason": "resolution", "got": f"{w}x{h}"})
+            continue
+        if args.num_frames and n != args.num_frames:
+            skipped.append({"clip": c.name, "reason": "frames", "got": n})
+            continue
+        link = args.out_dir / c.name
+        target = c.resolve()
+        if link.is_symlink() and link.readlink() == target:
+            pass                       # already correct
+        else:
+            if link.exists() or link.is_symlink():
+                link.unlink()          # replace a stale file/link rather than trusting it
+            link.symlink_to(target)
+        ok += 1
+
+    report_path = args.out_dir.parent / args.report
+    report_path.write_text(json.dumps(
+        {"src": str(args.src_dir), "required": f"{args.width}x{args.height}@{args.num_frames}f",
+         "total": len(clips), "kept": ok, "skipped": skipped}, indent=2))
+
+    rescuable = [s["clip"] for s in skipped if s["reason"] in ("resolution", "frames")]
+    list_path = args.out_dir.parent / args.needs_resize_list
+    list_path.write_text("\n".join(rescuable) + ("\n" if rescuable else ""))
+
+    by_reason: dict[str, int] = {}
+    for s in skipped:
+        by_reason[s["reason"]] = by_reason.get(s["reason"], 0) + 1
+    detail = ", ".join(f"{k}={v}" for k, v in sorted(by_reason.items())) or "none"
+    print(f"[filter] {ok}/{len(clips)} clips match {args.width}x{args.height}"
+          f"@{args.num_frames}f -> {args.out_dir}", flush=True)
+    print(f"[filter] skipped: {detail}  (details in {report_path})", flush=True)
+    if rescuable:
+        print(f"[filter] {len(rescuable)} rescuable by resize -> {list_path}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/find_static_windows.py

@@ -0,0 +1,353 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Find relatively-static WINDOWS inside (long) videos, before tracking.
+
+The camera in egocentric video turns and the whole view sweeps off-frame; a filter that
+seeds points once and tracks the whole clip would die at the first turn and mis-score
+everything after. Instead we compute a PER-FRAME global-motion signal that RE-SEEDS every
+frame (so it recovers after a turn: motion spikes during the turn, drops when static), then
+slide a window and locate the calm stretches.
+
+Per-frame motion m[t] = median Lucas-Kanade optical-flow magnitude between frames t-1 and t
+(fresh goodFeaturesToTrack each step), normalized by the frame diagonal. Camera pan/turn ->
+large; static -> small. A window is static iff m[t] stays low across ALL its frames (we score
+by a high percentile so one calm-but-not-perfect frame is fine but a turn is not). We also
+report per-window SURVIVAL: seed a grid at the window START, LK-track to the window END, and
+measure the fraction still tracked & in-frame -> "would CoTracker keep traces here?".
+
+    .venv/bin/python data_pipeline/find_static_windows.py scan  --videos-dir <dir> --out <dir> --window-sec 5 --limit 12
+    .venv/bin/python data_pipeline/find_static_windows.py serve --viz-dir <dir> --share
+"""
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(__file__))
+
+
+def read_gray_lowres(path: str, max_frames: int = 900, target_w: int = 320) -> tuple[list, float, int]:
+    """Decode DIRECTLY at low res (fast + tiny memory) -> gray frames + fps + total_frames.
+    Only used for the motion analysis; full-res frames are read on demand for extraction."""
+    import cv2
+    from decord import VideoReader, cpu
+    cap = cv2.VideoCapture(path)  # fast metadata read (no full decode)
+    W0 = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) or 640
+    H0 = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) or 360
+    fps = float(cap.get(cv2.CAP_PROP_FPS)) or 30.0
+    cap.release()
+    tw = min(target_w, W0)
+    th = int(round(H0 * tw / max(W0, 1)))
+    try:
+        vr = VideoReader(path, ctx=cpu(0), width=tw, height=th)
+    except Exception:  # noqa: BLE001  (older decord without width/height)
+        vr = VideoReader(path, ctx=cpu(0))
+    n_total = len(vr)
+    n = min(n_total, max_frames)
+    batch = vr.get_batch(list(range(n))).asnumpy()
+    gray = [cv2.cvtColor(f, cv2.COLOR_RGB2GRAY) for f in batch]
+    return gray, fps, n_total
+
+
+def read_rgb_frames(path: str, indices: Any) -> np.ndarray:
+    """Full-res RGB for a specific set of frame indices (decord random access)."""
+    from decord import VideoReader, cpu
+    vr = VideoReader(path, ctx=cpu(0))
+    return vr.get_batch(list(indices)).asnumpy()
+
+
+def motion_series(gray: list) -> tuple[np.ndarray, np.ndarray, float]:
+    """Per-frame normalized global camera motion via phase correlation (FFT-based global
+    shift between consecutive frames, ~1ms/frame). Recovers after turns: spikes during a
+    turn, drops when static. Returns (m, shifts, diag): m[t] = |shift|/diag, shifts[t] =
+    signed (dx,dy)/diag (used to accumulate net view drift over a window)."""
+    import cv2
+    H, W = gray[0].shape
+    diag = float(np.hypot(H, W))
+    hann = cv2.createHanningWindow((W, H), cv2.CV_32F)
+    m = np.zeros(len(gray), np.float32)
+    shifts = np.zeros((len(gray), 2), np.float32)
+    prev = gray[0].astype(np.float32)
+    for t in range(1, len(gray)):
+        cur = gray[t].astype(np.float32)
+        (dx, dy), _ = cv2.phaseCorrelate(prev, cur, hann)
+        shifts[t] = (dx / diag, dy / diag)
+        m[t] = float(np.hypot(dx, dy)) / diag
+        prev = cur
+    m[0] = m[1] if len(m) > 1 else 0.0
+    return m, shifts, diag
+
+
+def window_drift(shifts: np.ndarray, s: int, L: int) -> float:
+    """Net view drift over a window: max cumulative excursion of the accumulated per-frame
+    shifts, normalized by diagonal. Large => the camera slowly panned the view off-screen
+    even if no single frame moved much. Free (from the phase-correlation shifts)."""
+    c = np.cumsum(shifts[s:s + L], axis=0)
+    return float(np.max(np.hypot(c[:, 0], c[:, 1]))) if len(c) else 0.0
+
+
+def window_survival(gray: list, s: int, L: int, grid: int = 24) -> float:
+    """Seed a grid at frame s, LK-track to s+L-1; fraction still tracked & in-frame."""
+    import cv2
+    H, W = gray[0].shape
+    xs = np.linspace(W * 0.05, W * 0.95, grid)
+    ys = np.linspace(H * 0.05, H * 0.95, grid)
+    gx, gy = np.meshgrid(xs, ys)
+    pts = np.stack([gx.ravel(), gy.ravel()], 1).astype(np.float32).reshape(-1, 1, 2)
+    alive = np.ones(pts.shape[0], bool)
+    cur = pts.copy()
+    lk = dict(winSize=(21, 21), maxLevel=3)
+    for t in range(s + 1, min(s + L, len(gray))):
+        nxt, stt, _ = cv2.calcOpticalFlowPyrLK(gray[t - 1], gray[t], cur, None, **lk)
+        stt = stt.ravel().astype(bool)
+        x, y = nxt[:, 0, 0], nxt[:, 0, 1]
+        inb = (x >= 0) & (x < W) & (y >= 0) & (y < H)
+        alive &= stt & inb
+        cur = nxt
+    return float(alive.mean())
+
+
+def best_windows(m: np.ndarray, L: int, pct: float = 95, top: int = 1, min_gap: int | None = None) -> list:
+    """Return up to `top` low-motion windows (start index) minimizing the pct-percentile of
+    per-frame motion, non-overlapping by min_gap (default L)."""
+    if len(m) < L:
+        return []
+    min_gap = min_gap or L
+    scores = np.array([np.percentile(m[s:s + L], pct) for s in range(len(m) - L + 1)])
+    order = np.argsort(scores)
+    picks: list[int] = []
+    for s in order:
+        if all(abs(s - p) >= min_gap for p in picks):
+            picks.append(int(s))
+            if len(picks) >= top:
+                break
+    return picks
+
+
+def greedy_valid_windows(m: np.ndarray,
+                         shifts: np.ndarray,
+                         L: int,
+                         pct: float,
+                         p95_thresh: float,
+                         drift_thresh: float,
+                         cand_step: int | None = None) -> list:
+    """Greedily take ALL non-overlapping windows passing p95_motion<=p95_thresh AND
+    net-drift<=drift_thresh, calmest-first. Returns list of (start, p95, drift). All metrics
+    come from the phase-correlation signal (no LK), so this is cheap."""
+    if len(m) < L:
+        return []
+    cand_step = cand_step or max(1, L // 6)
+    cand = [(s, float(np.percentile(m[s:s + L], pct))) for s in range(0, len(m) - L + 1, cand_step)]
+    cand.sort(key=lambda x: x[1])  # calmest first
+    accepted: list[tuple[int, float, float]] = []
+    for s, p in cand:
+        if p > p95_thresh:
+            break  # sorted -> everything after is worse
+        if any(abs(s - a[0]) < L for a in accepted):
+            continue  # overlaps an accepted window
+        drift = window_drift(shifts, s, L)
+        if drift > drift_thresh:
+            continue
+        accepted.append((int(s), round(p, 4), round(drift, 3)))
+    return sorted(accepted)
+
+
+def cmd_scan(args: argparse.Namespace) -> None:
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    import imageio.v2 as imageio
+
+    vids = sorted(Path(args.videos_dir).glob("**/*.mp4"))[:args.limit]
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    entries = []
+    for i, vp in enumerate(vids, 1):
+        gray, fps, _ = read_gray_lowres(str(vp), args.max_frames)
+        L = int(round(args.window_sec * fps))
+        if len(gray) < L:
+            print(f"[win] [{i}/{len(vids)}] {vp.name}: only {len(gray)}f < {L}, skip", flush=True)
+            continue
+        m, shifts, diag = motion_series(gray)
+        picks = best_windows(m, L, args.pct, top=args.top)
+        stem = vp.stem
+        # motion curve with chosen windows shaded
+        fig, ax = plt.subplots(figsize=(12, 3.2))
+        ax.plot(np.arange(len(m)) / fps, m, lw=0.8, color="steelblue")
+        ax.axhline(args.calm_thresh, color="red", ls="--", lw=0.8, label=f"calm thresh {args.calm_thresh}")
+        wstats = []
+        for j, s in enumerate(picks):
+            drift = window_drift(shifts, s, L)
+            p95 = float(np.percentile(m[s:s + L], args.pct))
+            ax.axvspan(s / fps, (s + L) / fps, color="limegreen", alpha=0.3)
+            ax.text((s + L / 2) / fps, m.max() * 0.9, f"drift {drift:.3f}", ha="center", fontsize=8)
+            wstats.append({
+                "start_frame": s,
+                "start_sec": round(s / fps, 2),
+                "len_frames": L,
+                "p95_motion": round(p95, 4),
+                "drift": round(drift, 3)
+            })
+        ax.set_xlabel("time (s)")
+        ax.set_ylabel("per-frame motion (norm)")
+        ax.set_title(f"{stem}  ({len(m)}f @ {fps:.0f}fps)  green = static {args.window_sec}s window")
+        ax.legend(fontsize=8)
+        fig.tight_layout()
+        curve = f"{stem}_motion.png"
+        fig.savefig(str(out / curve), dpi=80)
+        plt.close(fig)
+        # preview: the best window as a short clip (full-res, read on demand)
+        prev = ""
+        if picks:
+            s = picks[0]
+            rgbw = read_rgb_frames(str(vp), range(s, s + L))
+            imageio.mimsave(str(out / f"{stem}_win.mp4"), rgbw, fps=int(round(fps)), macro_block_size=1)
+            prev = f"{stem}_win.mp4"
+        entries.append({
+            "id": stem,
+            "src": str(vp),
+            "fps": round(fps, 2),
+            "n_frames": len(m),
+            "curve": curve,
+            "preview": prev,
+            "windows": wstats,
+            "median_motion": round(float(np.median(m)), 4)
+        })
+        best = wstats[0] if wstats else {}
+        print(
+            f"[win] [{i}/{len(vids)}] {stem}: {len(m)}f, best window @{best.get('start_sec')}s "
+            f"p95={best.get('p95_motion')} surv={best.get('survival')}",
+            flush=True)
+    (out / "manifest.json").write_text(json.dumps(entries, indent=2))
+    print(f"[win] scanned {len(entries)} videos -> {out}", flush=True)
+
+
+def cmd_build(args: argparse.Namespace) -> None:
+    """Scan untrimmed source clips, greedily extract all non-overlapping static windows, and
+    write them as target_frames-frame clips (subsampled to span the whole window) + manifest."""
+    import imageio.v2 as imageio
+    vids = sorted(Path(args.videos_dir).glob("**/*.mp4"))
+    out = Path(args.out)
+    (out / "videos").mkdir(parents=True, exist_ok=True)
+    GENERIC = "a first-person egocentric view of a person performing a kitchen task with their hands"
+    man = []
+    idx = 0
+    for vi, vp in enumerate(vids, 1):
+        if idx >= args.target_count:
+            break
+        try:
+            gray, fps, _ = read_gray_lowres(str(vp), args.max_frames)
+        except Exception as e:  # noqa: BLE001
+            print(f"[build] {vp.name}: read fail {e}", flush=True)
+            continue
+        L = int(round(args.window_sec * fps))
+        if len(gray) < L:
+            continue
+        m, shifts, _ = motion_series(gray)
+        wins = greedy_valid_windows(m, shifts, L, args.pct, args.p95_thresh, args.drift_thresh)
+        stride = max(1, round(L / args.target_frames))
+        out_fps = round(fps / stride, 3)
+        part = vp.stem.split("_")[0]
+        n_from_clip = 0
+        for (s, p95, drift) in wins:
+            if idx >= args.target_count:
+                break
+            sel = np.arange(s, s + L, stride)[:args.target_frames]
+            if sel.size < args.target_frames:
+                continue
+            f = read_rgb_frames(str(vp), sel)  # full-res, only the window frames
+            f = f[:, :f.shape[1] // 2 * 2, :f.shape[2] // 2 * 2]  # even dims for libx264
+            name = f"vid_{idx:06d}.mp4"
+            imageio.mimsave(str(out / "videos" / name), f, fps=int(round(out_fps)), macro_block_size=1)
+            man.append({
+                "idx": idx,
+                "path": name,
+                "cap": [GENERIC],
+                "fps": out_fps,
+                "num_frames": int(args.target_frames),
+                "duration": round(args.target_frames / out_fps, 3),
+                "resolution": [int(f.shape[1]), int(f.shape[2])],
+                "participant": part,
+                "orig": vp.stem,
+                "win_start_sec": round(s / fps, 2),
+                "p95_motion": p95,
+                "drift": drift
+            })
+            idx += 1
+            n_from_clip += 1
+        if n_from_clip:
+            print(f"[build] [{vi}/{len(vids)}] {vp.stem}: +{n_from_clip} windows (total {idx})", flush=True)
+    (out / "videos2caption.json").write_text(json.dumps(man, indent=2))
+    (out / "merge.txt").write_text(f"{out}/videos,{out}/videos2caption.json\n")
+    from collections import Counter
+    print(
+        f"[build] wrote {len(man)} static clips -> {out} | per-participant {dict(Counter(x['participant'] for x in man))}",
+        flush=True)
+
+
+def cmd_serve(args: argparse.Namespace) -> None:
+    import gradio as gr
+    viz = Path(args.viz_dir).resolve()
+    entries = json.loads((viz / "manifest.json").read_text())
+    gal = [(str(viz / e["curve"]), e["id"]) for e in entries]
+
+    def show(evt: gr.SelectData):
+        e = entries[evt.index]
+        vp = str(viz / e["preview"]) if e.get("preview") else None
+        return str(viz / e["curve"]), vp, json.dumps(e["windows"], indent=2)
+
+    with gr.Blocks(title="Static-window finder") as demo:
+        gr.Markdown("### Per-frame camera motion + static-window finder\n"
+                    "The curve is per-frame global motion (re-seeded each frame, so it **recovers after a turn** "
+                    "— spikes = turns, valleys = static). Green span = the chosen static window; `surv` = grid "
+                    "survival seeded at the window start. Click a clip -> motion curve + a preview of the best "
+                    "window + per-window stats. We keep windows with low `p95_motion` and high `survival`.")
+        with gr.Row():
+            g = gr.Gallery(value=gal, columns=2, height=560, label="videos (click)")
+            with gr.Column():
+                curve = gr.Image(label="per-frame motion (green = static window)")
+                vid = gr.Video(label="preview of best static window")
+                meta = gr.Code(label="windows", language="json")
+        g.select(show, None, [curve, vid, meta])
+    demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share, allowed_paths=[str(viz)])
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub = p.add_subparsers(dest="cmd", required=True)
+    r = sub.add_parser("scan")
+    r.add_argument("--videos-dir", required=True)
+    r.add_argument("--out", required=True)
+    r.add_argument("--window-sec", type=float, default=5.0)
+    r.add_argument("--pct", type=float, default=95.0, help="percentile of per-frame motion used to score a window")
+    r.add_argument("--calm-thresh", type=float, default=0.01, help="reference 'calm' motion line for the plot")
+    r.add_argument("--top", type=int, default=2, help="static windows to find per video")
+    r.add_argument("--max-frames", type=int, default=900, help="cap frames scanned per video (bounds cost)")
+    r.add_argument("--limit", type=int, default=12)
+    r.set_defaults(func=cmd_scan)
+    b = sub.add_parser("build")
+    b.add_argument("--videos-dir", required=True)
+    b.add_argument("--out", required=True)
+    b.add_argument("--window-sec", type=float, default=5.0)
+    b.add_argument("--target-frames", type=int, default=121, help="frames per output clip (window subsampled to span)")
+    b.add_argument("--target-count", type=int, default=200)
+    b.add_argument("--pct", type=float, default=95.0)
+    b.add_argument("--p95-thresh", type=float, default=0.004, help="max p95 per-frame motion for a valid window")
+    b.add_argument("--drift-thresh", type=float, default=0.35, help="max net view drift (frac of diagonal)")
+    b.add_argument("--max-frames", type=int, default=1500, help="cap frames scanned per source clip")
+    b.set_defaults(func=cmd_build)
+    s = sub.add_parser("serve")
+    s.add_argument("--viz-dir", required=True)
+    s.add_argument("--host", default="0.0.0.0")
+    s.add_argument("--port", type=int, default=7890)
+    s.add_argument("--share", action="store_true")
+    s.set_defaults(func=cmd_serve)
+    a = p.parse_args()
+    a.func(a)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/gen_eval_artifacts.py

@@ -0,0 +1,239 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pre-generate controllability-eval artifacts for the viewer app.
+
+For one checkpoint, over clips x counterfactuals, generate the video and save:
+  - <clip>_<ctrl>__gen.mp4      raw generation
+  - <clip>_<ctrl>__input.mp4    gen + the INPUT control tracks (what we asked for)
+  - <clip>_<ctrl>__tracked.mp4  gen + tracks RE-EXTRACTED from the gen (CoTracker3)
+  - <clip>_<ctrl>__heat.mp4     gen + points colored by per-point EPE (green=followed,
+                                red=ignored) -- the EPE heatmap (when tracks apply)
+plus per clip: <clip>__original.mp4 (real clip) and <clip>__original_tracks.mp4.
+
+A manifest.json records EPE / n_points / coverage and the file names so the viewer
+(examples/inference/gradio/trackwan/app_viewer.py) can browse without a GPU.
+
+Run one process per checkpoint (parallel across GPUs)::
+
+    srun ... env CUDA_VISIBLE_DEVICES=5 .venv/bin/python data_pipeline/gen_eval_artifacts.py \
+      --export <export_3000> --name "step 3000" --out <artifacts_root>/step3000 \
+      --data <funinp_now parquet> --clips 0 1 2 3
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from typing import Any
+
+import imageio.v2 as imageio
+import numpy as np
+import torch
+
+sys.path.insert(0, os.path.dirname(__file__))
+import synthetic_tracks as st  # noqa: E402
+import trackwan_infer as twi  # noqa: E402
+from fastvideo.eval.metrics.motion.control_sensitivity.metric import (  # noqa: E402
+    compute_control_sensitivity, render_diff_heatmap)
+
+HEAT_SCALE = 40.0  # px error mapped green(0)->red(>=40)
+
+
+def _draw(frames: Any, trpx: Any, vis: Any, colors: Any) -> Any:
+    from fastvideo.train.callbacks.track_validation import _draw_overlay
+    return _draw_overlay(frames, trpx, vis, colors, 9, 2, 0.65)
+
+
+def _solid(m: Any, rgb: Any) -> np.ndarray:
+    return np.tile(np.array([rgb], np.uint8), (m, 1))
+
+
+def _heat_colors(perpt: Any) -> np.ndarray:
+    e = np.clip(np.nan_to_num(perpt, nan=HEAT_SCALE) / HEAT_SCALE, 0.0, 1.0)
+    r = (255 * e).astype(np.uint8)
+    g = (255 * (1 - e)).astype(np.uint8)
+    b = np.full_like(r, 40)
+    return np.stack([r, g, b], 1)
+
+
+def _retrack_core(frames: Any,
+                  tr_px: Any,
+                  vs: Any,
+                  ct: Any,
+                  device: Any,
+                  max_points: int = 600,
+                  seed: int = 0,
+                  qf: int = 0,
+                  vt: float = 0.5) -> dict[str, Any] | None:
+    """Mirror compute_epe but also return the extracted tracks + per-point error."""
+    from fastvideo.eval.metrics.motion.cotracker_epe.metric import _retrack
+    T = min(frames.shape[0], tr_px.shape[0])
+    frames, tr_px, vs = frames[:T], tr_px[:T], vs[:T]
+    idx = np.nonzero(vs[qf] > vt)[0]
+    if idx.size == 0:
+        return None
+    if max_points and idx.size > max_points:
+        rng = np.random.default_rng(seed)
+        idx = np.sort(rng.choice(idx, size=max_points, replace=False))
+    q_xy = tr_px[qf, idx]
+    queries = np.concatenate([np.full((idx.size, 1), qf, np.float32), q_xy], axis=1)
+    rt, _ = _retrack(ct, frames, queries, device)  # [T,M,2] px
+    Tr = min(T, rt.shape[0])
+    tgt = tr_px[:Tr][:, idx]
+    pred = rt[:Tr]
+    mask = vs[:Tr][:, idx] > vt
+    d = np.sqrt(((pred - tgt)**2).sum(-1))  # [Tr,M]
+    valid = d[mask]
+    epe = float(valid.mean()) if valid.size else None
+    perpt = np.array([d[:, m][mask[:, m]].mean() if mask[:, m].any() else np.nan for m in range(idx.size)])
+    disp = np.sqrt(((tgt - tgt[0:1])**2).sum(-1))  # input travel from frame 0
+    moving = disp.max(0) > 8.0
+    mm = mask & moving[None, :]
+    epe_moving = float(d[mm].mean()) if d[mm].size else None
+    return dict(tgt=tgt,
+                pred=pred,
+                mask=mask,
+                epe=epe,
+                epe_moving=epe_moving,
+                perpt=perpt,
+                n=int(idx.size),
+                n_moving=int(moving.sum()),
+                cov=float(mask.mean()))
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--export", required=True)
+    p.add_argument("--name", required=True, help="checkpoint label e.g. 'step 3000'")
+    p.add_argument("--out", required=True)
+    p.add_argument("--yaml", default="examples/train/scenario/worldmodel/finetune_wantrack_i2v.yaml")
+    p.add_argument("--data", required=True)
+    p.add_argument("--clips", type=int, nargs="+", default=[0, 1, 2, 3])
+    p.add_argument("--steps", type=int, default=30)
+    p.add_argument("--seed", type=int, default=1000)
+    p.add_argument("--fps", type=int, default=24)
+    args = p.parse_args()
+
+    os.makedirs(args.out, exist_ok=True)
+    model, tc = twi.load_trackwan(args.export, args.yaml)
+    text_len = int(tc.pipeline_config.text_encoder_configs[0].arch_config.text_len)
+    samples = twi.load_conditioning_from_parquet(args.data, args.clips, text_len)
+    from fastvideo.eval.metrics.motion.cotracker_epe.metric import load_cotracker
+    ct = load_cotracker(model.device)
+    num_lat_t = samples[0]["first_frame_latent"].shape[2]
+    ratio = int(tc.pipeline_config.vae_config.arch_config.temporal_compression_ratio)
+    Tpx = (num_lat_t - 1) * ratio + 1
+    dev = model.device
+
+    def save(frames: Any, name: str) -> str:
+        path = os.path.join(args.out, name)
+        imageio.mimsave(path, frames, fps=args.fps, macro_block_size=1)
+        return name
+
+    manifest = {"name": args.name, "clips": {}}
+    for ci, s in zip(args.clips, samples, strict=False):
+        gt_t = s["track_points"][0].numpy()[:Tpx]
+        gt_v = s["track_visibility"][0].numpy()[:Tpx]
+        swap = samples[(args.clips.index(ci) + 1) % len(samples)]
+        g = st.make_grid(50)
+        pan_t, pan_v = st.pan(g, Tpx, 0.25, 0.0)
+        zoom_t, zoom_v = st.zoom(g, Tpx, 1.25)
+        drag_t, drag_v = st.drag(g, Tpx, center=(0.5, 0.5), dx=0.3, dy=0.0, radius=0.2)
+        drag_vs = st.select_radius(drag_v, g, center=(0.5, 0.5), radius=0.2)
+        controls = {
+            "gt": (gt_t, gt_v),
+            "none": (None, None),
+            "pan_right": (pan_t, pan_v),
+            "zoom_in": (zoom_t, zoom_v),
+            "drag_dense": (drag_t, drag_v),
+            "drag_sparse": (drag_t, drag_vs),
+            "swap": (swap["track_points"][0].numpy()[:Tpx], swap["track_visibility"][0].numpy()[:Tpx]),
+        }
+
+        ref = twi.decode_reference(model, s["vae_latent"])
+        H, W = ref.shape[1], ref.shape[2]
+        gtpx = gt_t.copy()
+        gtpx[..., 0] *= W
+        gtpx[..., 1] *= H
+        entry = {
+            "caption":
+            s["caption"][:240],
+            "H":
+            int(H),
+            "W":
+            int(W),
+            "original":
+            save(ref, f"clip{ci}__original.mp4"),
+            "original_tracks":
+            save(_draw(ref, gtpx, gt_v, _solid(gt_t.shape[1], [0, 220, 255])), f"clip{ci}__original_tracks.mp4"),
+            "controls": {}
+        }
+
+        frames_gt = None
+        for cname, (tr, vs) in controls.items():
+            tp = torch.from_numpy(tr)[None].float() if tr is not None else None
+            tv = torch.from_numpy(vs)[None].float() if vs is not None else None
+            lat = twi.generate(model,
+                               first_frame_latent=s["first_frame_latent"],
+                               text_embedding=s["text_embedding"],
+                               text_attention_mask=s["text_attention_mask"],
+                               track_points=tp,
+                               track_visibility=tv,
+                               clip_feature=s["clip_feature"],
+                               num_steps=args.steps,
+                               seed=args.seed)
+            frames = twi.decode_to_pixels(model, lat)
+            if cname == "gt":
+                frames_gt = frames
+            trpx = None
+            if tr is not None:
+                trpx = tr.copy()
+                trpx[..., 0] *= W
+                trpx[..., 1] *= H
+            rec = {
+                "gen": save(frames, f"clip{ci}_{cname}__gen.mp4"),
+                "epe": None,
+                "epe_moving": None,
+                "n_points": 0,
+                "coverage": 0.0
+            }
+            if trpx is not None:
+                core = _retrack_core(frames, trpx, vs, ct, dev)
+                if core is not None:
+                    m = core["tgt"].shape[1]
+                    rec["input"] = save(_draw(frames, core["tgt"], core["mask"], _solid(m, [0, 220, 255])),
+                                        f"clip{ci}_{cname}__input.mp4")
+                    rec["tracked"] = save(_draw(frames, core["pred"], core["mask"], _solid(m, [255, 220, 0])),
+                                          f"clip{ci}_{cname}__tracked.mp4")
+                    rec["heat"] = save(_draw(frames, core["pred"], core["mask"], _heat_colors(core["perpt"])),
+                                       f"clip{ci}_{cname}__heat.mp4")
+                    rec.update(epe=core["epe"],
+                               epe_moving=core["epe_moving"],
+                               n_points=core["n"],
+                               n_moving=core["n_moving"],
+                               coverage=core["cov"])
+            # intervention diff vs the GT-track generation (counterfactual control sensitivity)
+            if cname != "gt" and frames_gt is not None:
+                sens = compute_control_sensitivity(frames_gt, frames, trpx, vs)
+                rec["diff"] = save(render_diff_heatmap(frames_gt, frames), f"clip{ci}_{cname}__diff.mp4")
+                sr = sens["sensitivity_roi"]
+                rec.update(sensitivity_roi=(round(sr, 4) if sr == sr else None),
+                           bg_leakage=round(sens["bg_leakage"], 4),
+                           localization_iou=round(sens["localization_iou"], 4),
+                           mean_diff=round(sens["mean_diff"], 4))
+            entry["controls"][cname] = rec
+            print(
+                f"[{args.name}] clip{ci} {cname:12s} "
+                f"EPE={rec['epe'] if rec['epe'] is None else round(float(rec['epe']),2)} "
+                f"EPEmv={rec['epe_moving'] if rec['epe_moving'] is None else round(float(rec['epe_moving']),2)} "
+                f"sens={rec.get('sensitivity_roi')} iou={rec.get('localization_iou')}",
+                flush=True)
+        manifest["clips"][str(ci)] = entry
+
+    with open(os.path.join(args.out, "manifest.json"), "w") as f:
+        json.dump(manifest, f, indent=2)
+    print("ARTIFACTS_DONE", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/gen_synth_i2v_worker.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Single-GPU I2V synthetic gen: one SHARED seed image + a fixed caption list -> videos.
+
+Reproduces the wantrack_synth_toy setup (all clips start from the same seed frame; each caption
+varies only the motion) but at 720p/24fps. Uses Wan2.1-I2V-14B-720P: for every caption it
+generates a clip conditioned on --seed-image at --height x --width.
+
+Unlike the T2V worker there is NO first-frame drop: I2V's frame 0 IS the (clean) seed, so we
+keep all num_frames. Output layout matches gen_synth_worker.py (videos/, meta/, manifest_shards/)
+so merge_synth_manifests.py + the tracks/preprocess stages work unchanged.
+
+Idempotent/resumable (skips finished mp4s). Parallelize across GPUs with --worker-id/--num-workers.
+
+  CUDA_VISIBLE_DEVICES=0 python data_pipeline/gen_synth_i2v_worker.py \
+    --seed-image data_pipeline/synth_toy_720p/synthetic_seed.png \
+    --captions   data_pipeline/synth_toy_720p/captions.txt \
+    --output-dir /home/hal-kevin/data/motion-stream-synth
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+from pathlib import Path
+
+import numpy as np
+
+MODEL_DEFAULT = "/home/hal-kevin/models/Wan2.1-I2V-14B-720P-Diffusers"
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--seed-image", type=Path, required=True, help="shared first-frame conditioning image")
+    p.add_argument("--captions", type=Path, required=True, help="one caption per line")
+    p.add_argument("--output-dir", type=Path, required=True)
+    p.add_argument("--model", default=MODEL_DEFAULT)
+    p.add_argument("--worker-id", type=int, default=0)
+    p.add_argument("--num-workers", type=int, default=1)
+    p.add_argument("--height", type=int, default=720)
+    p.add_argument("--width", type=int, default=1280)
+    p.add_argument("--num-frames", type=int, default=121, help="kept 1:1 (I2V, no drop); must be 4k+1")
+    p.add_argument("--fps", type=int, default=24)
+    p.add_argument("--steps", type=int, default=40)
+    p.add_argument("--guidance-scale", type=float, default=5.0)
+    p.add_argument("--seed-base", type=int, default=1024, help="per-clip seed = seed_base + caption_idx")
+    return p.parse_args()
+
+
+def main():
+    a = parse_args()
+    assert (a.num_frames - 1) % 4 == 0, f"num_frames {a.num_frames} must be 4k+1"
+    assert a.seed_image.exists(), f"seed image not found: {a.seed_image}"
+
+    videos_dir = a.output_dir / "videos"
+    meta_dir = a.output_dir / "meta"
+    shard_dir = a.output_dir / "manifest_shards"
+    for d in (videos_dir, meta_dir, shard_dir):
+        d.mkdir(parents=True, exist_ok=True)
+    shard_manifest = shard_dir / f"worker_{a.worker_id:04d}.jsonl"
+    fail_log = a.output_dir / f"failures_worker_{a.worker_id:04d}.log"
+
+    captions = [ln.strip() for ln in a.captions.read_text().splitlines() if ln.strip()]
+    # Fixed order (caption line -> vid index); this worker owns a stride slice.
+    my_idx = list(range(a.worker_id, len(captions), a.num_workers))
+    print(f"[w{a.worker_id}/{a.num_workers}] {len(my_idx)} caption(s) "
+          f"seed={a.seed_image.name} {a.width}x{a.height}@{a.fps}fps x{a.num_frames}f", flush=True)
+
+    import imageio.v2 as imageio
+    from fastvideo import VideoGenerator
+    t0 = time.time()
+    g = VideoGenerator.from_pretrained(
+        a.model, num_gpus=1, use_fsdp_inference=False,
+        dit_cpu_offload=False, vae_cpu_offload=False,
+        text_encoder_cpu_offload=True, pin_cpu_memory=True,
+    )
+    print(f"[w{a.worker_id}] model ready in {time.time()-t0:.1f}s", flush=True)
+
+    seed_path = str(a.seed_image.resolve())
+    n_ok = 0
+    for gi in my_idx:
+        fp = videos_dir / f"vid_{gi:06d}.mp4"
+        if fp.exists():
+            continue
+        prompt = captions[gi]
+        tmp = videos_dir / f".tmp_w{a.worker_id}_{gi:06d}.mp4"
+        t = time.time()
+        try:
+            res = g.generate_video(
+                prompt, image_path=seed_path, save_video=False, return_frames=True,
+                height=a.height, width=a.width, num_frames=a.num_frames, fps=a.fps,
+                seed=a.seed_base + gi, num_inference_steps=a.steps,
+                guidance_scale=a.guidance_scale,
+            )
+            if isinstance(res, list):
+                res = res[0]
+            frames = np.asarray(res["frames"])          # I2V: keep all frames, no drop
+            if frames.shape[0] != a.num_frames:
+                raise RuntimeError(f"got {frames.shape[0]} frames, want {a.num_frames}")
+            imageio.mimsave(tmp, list(frames), fps=a.fps, format="mp4")
+            os.replace(tmp, fp)
+        except Exception as e:                          # keep the worker alive
+            with fail_log.open("a") as f:
+                f.write(json.dumps({"idx": gi, "err": repr(e)[:500]}) + "\n")
+            print(f"[w{a.worker_id}] FAIL idx={gi}: {e!r}", flush=True)
+            if tmp.exists():
+                tmp.unlink()
+            continue
+        dt = time.time() - t
+        rec = {"idx": gi, "path": fp.name, "cap": [prompt], "fps": float(a.fps),
+               "num_frames": a.num_frames, "duration": a.num_frames / float(a.fps),
+               "resolution": {"width": a.width, "height": a.height},
+               "gen_seconds": round(dt, 1)}
+        with shard_manifest.open("a") as f:
+            f.write(json.dumps(rec) + "\n")
+        (meta_dir / f"vid_{gi:06d}.json").write_text(json.dumps(rec))
+        n_ok += 1
+        print(f"[w{a.worker_id}] {n_ok}/{len(my_idx)} done, {dt:.0f}s idx={gi}", flush=True)
+    print(f"[w{a.worker_id}] DONE_WORKER made {n_ok} new videos", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/gen_synth_worker.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Single-GPU data-parallel worker for large-scale Wan2.2-T2V-A14B synthetic gen.
+
+One VideoGenerator per GPU (num_gpus=1, no FSDP/SP/TP). Each worker owns a stride
+slice of the prompt list: prompts[worker_id :: num_workers]. Idempotent/resumable:
+a video whose final mp4 exists is skipped, so a requeue after a cordon just continues.
+
+First-frame saturation fix baked in: generate (num_frames + drop) frames, drop the
+first `drop` decoded frames, keep `num_frames`. (See research_log/08-first-frame-saturation.md.)
+
+Per-worker manifest shard (manifest_shards/worker_<id>.jsonl) avoids 48 procs racing on
+one file; merge_manifests.py compiles them into the FastVideo videos2caption.json.
+"""
+from __future__ import annotations
+import argparse, json, os, time, traceback
+from pathlib import Path
+import numpy as np
+
+MODEL = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--prompts", type=Path, required=True)
+    p.add_argument("--output-dir", type=Path, required=True)
+    p.add_argument("--worker-id", type=int, required=True)
+    p.add_argument("--num-workers", type=int, required=True)
+    p.add_argument("--max-videos", type=int, default=None, help="global cap on #prompts consumed (across all workers)")
+    p.add_argument("--height", type=int, default=720)
+    p.add_argument("--width", type=int, default=1280)
+    p.add_argument("--num-frames", type=int, default=121, help="frames to KEEP")
+    p.add_argument("--drop", type=int, default=8, help="leading frames dropped (gen = keep+drop, must be 4k+1)")
+    p.add_argument("--fps", type=int, default=16)
+    p.add_argument("--steps", type=int, default=40)
+    p.add_argument("--guidance-scale", type=float, default=4.0)
+    p.add_argument("--guidance-scale-2", type=float, default=3.0)
+    p.add_argument("--seed-base", type=int, default=1024, help="per-video seed = seed_base + global_prompt_idx")
+    p.add_argument("--shuffle-seed", type=int, default=1234, help="deterministic prompt shuffle (same across workers)")
+    return p.parse_args()
+
+def load_prompts(path: Path):
+    lines = [ln.strip() for ln in path.read_text(encoding="utf-8").splitlines() if ln.strip()]
+    return lines
+
+def main():
+    a = parse_args()
+    gen_frames = a.num_frames + a.drop
+    assert (gen_frames - 1) % 4 == 0, f"gen frames {gen_frames} must be 4k+1"
+
+    videos_dir = a.output_dir / "videos"
+    meta_dir = a.output_dir / "meta"
+    shard_dir = a.output_dir / "manifest_shards"
+    for d in (videos_dir, meta_dir, shard_dir):
+        d.mkdir(parents=True, exist_ok=True)
+    shard_manifest = shard_dir / f"worker_{a.worker_id:04d}.jsonl"
+    fail_log = a.output_dir / f"failures_worker_{a.worker_id:04d}.log"
+
+    prompts = load_prompts(a.prompts)
+    # Deterministic global shuffle so prompt order is stable across restarts and workers.
+    order = np.random.RandomState(a.shuffle_seed).permutation(len(prompts))
+    if a.max_videos is not None:
+        order = order[:a.max_videos]
+    # This worker's stride slice.
+    my_positions = list(range(a.worker_id, len(order), a.num_workers))
+    my_global_idx = [int(order[pos]) for pos in my_positions]
+    print(f"[w{a.worker_id}/{a.num_workers}] assigned {len(my_global_idx)} prompts "
+          f"(gen {gen_frames}->keep {a.num_frames}, {a.width}x{a.height}@{a.fps}fps)", flush=True)
+
+    # already-done set from existing mp4s
+    def final_path(idx): return videos_dir / f"vid_{idx:06d}.mp4"
+
+    import imageio.v2 as imageio
+    from fastvideo import VideoGenerator
+    t0 = time.time()
+    g = VideoGenerator.from_pretrained(
+        MODEL, num_gpus=1, use_fsdp_inference=False,
+        dit_cpu_offload=False, vae_cpu_offload=False,
+        text_encoder_cpu_offload=True, pin_cpu_memory=True,
+    )
+    print(f"[w{a.worker_id}] model ready in {time.time()-t0:.1f}s", flush=True)
+
+    n_ok = 0
+    for gi in my_global_idx:
+        fp = final_path(gi)
+        if fp.exists():
+            continue
+        prompt = prompts[gi]
+        tmp = videos_dir / f".tmp_w{a.worker_id}_{gi:06d}.mp4"
+        t = time.time()
+        try:
+            res = g.generate_video(
+                prompt, save_video=False, return_frames=True,
+                height=a.height, width=a.width, num_frames=gen_frames, fps=a.fps,
+                seed=a.seed_base + gi, num_inference_steps=a.steps,
+                guidance_scale=a.guidance_scale, guidance_scale_2=a.guidance_scale_2,
+            )
+            if isinstance(res, list):
+                res = res[0]
+            frames = np.asarray(res["frames"])[a.drop:a.drop + a.num_frames]
+            if frames.shape[0] != a.num_frames:
+                raise RuntimeError(f"got {frames.shape[0]} frames after drop, want {a.num_frames}")
+            imageio.mimsave(tmp, list(frames), fps=a.fps, format="mp4")
+            os.replace(tmp, fp)  # atomic publish
+        except Exception as e:  # keep the worker alive
+            with fail_log.open("a") as f:
+                f.write(json.dumps({"idx": gi, "err": repr(e)[:500]}) + "\n")
+            print(f"[w{a.worker_id}] FAIL idx={gi}: {e!r}", flush=True)
+            if tmp.exists(): tmp.unlink()
+            continue
+        dt = time.time() - t
+        rec = {"idx": gi, "path": fp.name, "cap": [prompt], "fps": float(a.fps),
+               "num_frames": a.num_frames, "resolution": {"width": a.width, "height": a.height},
+               "gen_seconds": round(dt, 1)}
+        with shard_manifest.open("a") as f:
+            f.write(json.dumps(rec) + "\n")
+        (meta_dir / f"vid_{gi:06d}.json").write_text(json.dumps(rec))
+        n_ok += 1
+        if n_ok % 10 == 0:
+            print(f"[w{a.worker_id}] {n_ok} done, last {dt:.0f}s idx={gi}", flush=True)
+    print(f"[w{a.worker_id}] DONE_WORKER made {n_ok} new videos", flush=True)
+
+if __name__ == "__main__":
+    main()

data_pipeline/generate_videos.py

@@ -0,0 +1,235 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stage 0b: generate a synthetic (video, prompt) dataset with Wan2.2-T2V-A14B.
+
+Text prompts -> .mp4 videos + a FastVideo-compatible manifest, so the existing
+preprocess pipeline can ingest the result directly. CoTracker point extraction is a
+separate stage (``extract_tracks.py``).
+
+Design (see notes/DECISIONS.md):
+- Generate at the *training* fps/length (default 16 fps, 81 frames ~= 5 s) so per-frame
+  point tracks align 1:1 with frames and no resampling is needed downstream.
+- T2V only. The eventual I2V+points model uses frame 0 as the conditioning image at
+  training time, so no input image is needed here.
+- Idempotent/resumable: each finished video is appended to ``manifest.jsonl`` and skipped
+  on re-run.
+
+Run on a GPU node (never the login node), e.g.:
+
+    srun --jobid=<shao_wm jobid> --overlap --ntasks=1 \
+        .venv/bin/python data_pipeline/generate_videos.py \
+        --prompts examples/dataset/vidprom/prompts/vidprom_filtered_extended.txt \
+        --output-dir /mnt/weka/home/hao.zhang/shao/data/motion_pipeline/wan22_t2v_720p \
+        --num-videos 50 --num-gpus 8
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import shutil
+import time
+import warnings
+from pathlib import Path
+
+DEFAULT_MODEL = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--prompts", type=Path, required=True, help="Text file, one prompt per line.")
+    p.add_argument("--output-dir", type=Path, required=True, help="Dataset root (videos/, manifest, ...).")
+    p.add_argument("--model", type=str, default=DEFAULT_MODEL)
+    p.add_argument("--num-videos", type=int, default=50, help="How many prompts to generate (after --start).")
+    p.add_argument("--start", type=int, default=0, help="Offset into the prompt list (for sharding).")
+    p.add_argument("--shuffle", action="store_true", help="Deterministically shuffle prompts before slicing.")
+    p.add_argument("--num-gpus", type=int, default=8)
+    p.add_argument("--height", type=int, default=720)
+    p.add_argument("--width", type=int, default=1280)
+    p.add_argument("--num-frames", type=int, default=121)
+    p.add_argument("--fps", type=int, default=24)
+    p.add_argument("--trim-start-frames", type=int, default=0,
+                   help="Drop this many frames from the start of each generated video (VAE warm-up artifact). "
+                        "Generation runs for num_frames+trim_start_frames and the head is discarded.")
+    p.add_argument("--seed", type=int, default=1024, help="Base seed; per-video seed = seed + global index.")
+    p.add_argument("--num-inference-steps", type=int, default=None, help="Override model default if set.")
+    p.add_argument("--negative-prompt", type=str, default=None)
+    p.add_argument("--image",
+                   type=str,
+                   default=None,
+                   help="If set, do I2V from this image (same for every prompt); else T2V.")
+    # Offload controls. Default OFF: Wan2.2-A14B is ~56GB bf16 and fits on a single H200 (143GB),
+    # so offloading (esp. layerwise) only makes generation ~10x slower. Enable on small GPUs.
+    p.add_argument("--dit-cpu-offload", action="store_true", help="Offload DiT to CPU (slow).")
+    p.add_argument("--dit-layerwise-offload",
+                   action="store_true",
+                   help="Stream DiT layers from CPU per step (very slow; only for tiny GPUs).")
+    p.add_argument("--text-encoder-cpu-offload", action="store_true", help="Offload text encoder to CPU.")
+    p.add_argument("--vae-cpu-offload", action="store_true", help="Offload VAE to CPU.")
+    return p.parse_args()
+
+
+def load_prompts(path: Path, start: int, num: int, shuffle: bool, seed: int) -> list[tuple[int, str]]:
+    if not path.exists():
+        raise FileNotFoundError(f"Prompts file not found: {path}\n"
+                                "Download it first (login node): cd examples/dataset/vidprom && ./download_dataset.sh")
+    lines = [ln.strip() for ln in path.read_text(encoding="utf-8").splitlines() if ln.strip()]
+    indexed = list(enumerate(lines))  # global index is stable w.r.t. the raw file order
+    if shuffle:
+        random.Random(seed).shuffle(indexed)
+    return indexed[start:start + num]
+
+
+def read_done_indices(manifest_jsonl: Path) -> set[int]:
+    done: set[int] = set()
+    if manifest_jsonl.exists():
+        for ln in manifest_jsonl.read_text().splitlines():
+            ln = ln.strip()
+            if not ln:
+                continue
+            try:
+                done.add(int(json.loads(ln)["idx"]))
+            except (json.JSONDecodeError, KeyError, ValueError):
+                continue
+    return done
+
+
+def rebuild_manifest(manifest_jsonl: Path, videos_dir: Path, json_path: Path, merge_path: Path) -> int:
+    """Compile manifest.jsonl -> videos2caption.json + merge.txt (FastVideo format)."""
+    records: dict[int, dict] = {}
+    if manifest_jsonl.exists():
+        for ln in manifest_jsonl.read_text().splitlines():
+            ln = ln.strip()
+            if not ln:
+                continue
+            try:
+                rec = json.loads(ln)
+            except json.JSONDecodeError:
+                continue
+            records[int(rec["idx"])] = rec
+    ordered = [records[k] for k in sorted(records)]
+
+    tmp = json_path.with_suffix(".json.tmp")
+    tmp.write_text(json.dumps(ordered, indent=2))
+    tmp.replace(json_path)
+    merge_path.write_text(f"{videos_dir.resolve()},{json_path.resolve()}\n")
+    return len(ordered)
+
+
+def main() -> None:
+    args = parse_args()
+    warnings.filterwarnings("ignore", category=DeprecationWarning, module="fastvideo.*")
+
+    output_dir = args.output_dir
+    videos_dir = output_dir / "videos"
+    videos_dir.mkdir(parents=True, exist_ok=True)
+    manifest_jsonl = output_dir / "manifest.jsonl"
+    json_path = output_dir / "videos2caption.json"
+    merge_path = output_dir / "merge.txt"
+    fail_log = output_dir / "failures.log"
+
+    selected = load_prompts(args.prompts, args.start, args.num_videos, args.shuffle, args.seed)
+    done = read_done_indices(manifest_jsonl)
+    todo = [(i, pr) for (i, pr) in selected if i not in done]
+    print(f"[gen] {len(selected)} selected, {len(done)} already done, {len(todo)} to generate", flush=True)
+
+    if not todo:
+        n = rebuild_manifest(manifest_jsonl, videos_dir, json_path, merge_path)
+        print(f"[gen] nothing to do; manifest has {n} entries -> {json_path}", flush=True)
+        return
+
+    # Import here so --help works without loading torch/fastvideo.
+    from fastvideo import VideoGenerator
+
+    generator = VideoGenerator.from_pretrained(
+        args.model,
+        num_gpus=args.num_gpus,
+        use_fsdp_inference=False,
+        dit_cpu_offload=args.dit_cpu_offload,
+        dit_layerwise_offload=args.dit_layerwise_offload,
+        vae_cpu_offload=args.vae_cpu_offload,
+        text_encoder_cpu_offload=args.text_encoder_cpu_offload,
+        pin_cpu_memory=True,
+    )
+
+    extra: dict = {}
+    if args.num_inference_steps is not None:
+        extra["num_inference_steps"] = args.num_inference_steps
+    if args.negative_prompt is not None:
+        extra["negative_prompt"] = args.negative_prompt
+
+    # Wan VAE requires num_frames = 4k+1. Round up gen_frames to satisfy this,
+    # then trim the actual excess (may be more than trim_start_frames).
+    _raw = args.num_frames + args.trim_start_frames
+    gen_frames = _raw if (_raw - 1) % 4 == 0 else _raw + (4 - (_raw - 1) % 4)
+    actual_trim = gen_frames - args.num_frames
+    duration = float(args.num_frames) / float(args.fps)
+    for n_done, (idx, prompt) in enumerate(todo, 1):
+        final_path = videos_dir / f"vid_{idx:06d}.mp4"
+        if final_path.exists():  # belt-and-suspenders vs manifest
+            continue
+        tmp_dir = videos_dir / f".tmp_{idx:06d}"
+        if tmp_dir.exists():
+            shutil.rmtree(tmp_dir, ignore_errors=True)
+        tmp_dir.mkdir(parents=True, exist_ok=True)
+
+        t0 = time.time()
+        try:
+            i2v_kwargs = {"image_path": args.image} if args.image else {}
+            generator.generate_video(
+                prompt,
+                output_path=str(tmp_dir),
+                save_video=True,
+                height=args.height,
+                width=args.width,
+                num_frames=gen_frames,
+                fps=args.fps,
+                seed=args.seed + idx,
+                **i2v_kwargs,
+                **extra,
+            )
+            produced = sorted(tmp_dir.glob("*.mp4"))
+            if not produced:
+                raise RuntimeError("no .mp4 produced by generate_video")
+            src = str(produced[0])
+            if actual_trim > 0:
+                import subprocess as _sp
+                tmp_trim = str(final_path) + ".trim.mp4"
+                _sp.run(
+                    ["ffmpeg", "-y", "-i", src,
+                     "-vf", f"trim=start_frame={actual_trim},setpts=PTS-STARTPTS",
+                     "-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "18", "-an", tmp_trim],
+                    check=True, capture_output=True,
+                )
+                Path(tmp_trim).replace(final_path)
+            else:
+                shutil.move(src, str(final_path))
+        except Exception as e:  # noqa: BLE001 - keep the batch alive, log and move on
+            with fail_log.open("a") as f:
+                f.write(json.dumps({"idx": idx, "error": repr(e), "prompt": prompt}) + "\n")
+            print(f"[gen] FAILED idx={idx}: {e!r}", flush=True)
+            shutil.rmtree(tmp_dir, ignore_errors=True)
+            continue
+        shutil.rmtree(tmp_dir, ignore_errors=True)
+
+        record = {
+            "idx": idx,
+            "path": final_path.name,  # basename; folder in merge.txt is videos_dir
+            "cap": [prompt],
+            "fps": float(args.fps),
+            "duration": duration,
+            "num_frames": int(args.num_frames),
+            "resolution": {
+                "width": args.width,
+                "height": args.height
+            },
+        }
+        with manifest_jsonl.open("a") as f:
+            f.write(json.dumps(record) + "\n")
+        print(f"[gen] [{n_done}/{len(todo)}] idx={idx} {time.time()-t0:.1f}s -> {final_path.name}", flush=True)
+
+    n = rebuild_manifest(manifest_jsonl, videos_dir, json_path, merge_path)
+    print(f"[gen] done; manifest has {n} entries -> {json_path}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/lwlr_experiment.py

@@ -0,0 +1,246 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""Layer-wise learning-rate (LWLR) bootstrap experiment.
+
+Question: on a 14B random-init WanTrack model, does giving the track pathway
+(track_encoder.* + patch_embedding.weight[:, 36:]) a 100x higher LR than the
+rest actually let it grow into functional gradient magnitudes within a small
+number of steps? Or is the encoder stuck regardless?
+
+Runs a short training loop on real OpenVid samples, tracking weight+grad norms
+of the track pathway at each step. Compares against the "known good" magnitude
+observed in 1.3B_merged (per-sample track_encoder.proj grad ~0.5).
+"""
+from __future__ import annotations
+
+import argparse
+import glob
+import json
+import os
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+import torch.nn.functional as F
+from safetensors.torch import load_file as st_load
+
+# Distributed init BEFORE importing fastvideo (same pattern as gradient_flow_analysis)
+def _init_distributed_1gpu() -> None:
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("LOCAL_RANK", "0")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29513")
+    os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+    from fastvideo.distributed.parallel_state import (
+        maybe_init_distributed_environment_and_model_parallel,
+    )
+    maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1)
+
+
+_init_distributed_1gpu()
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from gradient_flow_analysis import (  # noqa: E402
+    build_model_from_diffusers,
+    load_samples,
+    build_i2v_cond_concat,
+    flow_matching_noisy,
+    Sample,
+)
+
+
+TRACK_PATH_PARAMS = (
+    "track_encoder.proj.weight",
+    "track_encoder.temporal_conv.weight",
+    "track_encoder.proj.bias",           # may not exist depending on TRACKWAN_TRACK_BIAS
+    "track_encoder.temporal_conv.bias",
+)
+
+
+def split_param_groups(model, high_lr: float, base_lr: float) -> list[dict]:
+    """Give track pathway (track_encoder.* + patch_embedding[:, 36:]) a high LR.
+
+    patch_embedding.weight is a SINGLE tensor of shape [hidden, 52, 1, 2, 2] — we can't
+    give sub-slice a different LR through standard param groups. Workaround: put the
+    whole patch_embedding in a "hybrid" group with high LR (the base 36 channels will
+    also move faster, but by only 100x for a short warmup that's small — 100 steps of
+    high LR is still << stage-1 4800 steps of normal LR).
+    """
+    high_group, base_group = [], []
+    named = dict(model.named_parameters())
+    for n, p in named.items():
+        if not p.requires_grad:
+            continue
+        # High-LR pathway: encoder convs (+ bias if present) + the whole patch_embedding.weight
+        is_encoder = n.startswith("track_encoder.")
+        is_patch_embed = (n == "patch_embedding.weight")
+        if is_encoder or is_patch_embed:
+            high_group.append(p)
+        else:
+            base_group.append(p)
+    return [
+        {"params": high_group, "lr": high_lr, "name": "track_pathway"},
+        {"params": base_group, "lr": base_lr, "name": "dit_body"},
+    ]
+
+
+def snapshot_norms(model, sample_grad_i: int) -> dict:
+    """Return weight/grad norms for the diagnostic params."""
+    out = {}
+    for n, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        # Track params of interest
+        keep = (
+            n.startswith("track_encoder.")
+            or n == "patch_embedding.weight"
+        )
+        if not keep:
+            continue
+        w_norm = p.detach().float().norm().item()
+        g_norm = p.grad.detach().float().norm().item() if p.grad is not None else 0.0
+        # For patch_embedding, split into base/track slices
+        if n == "patch_embedding.weight":
+            w_base = p.detach()[:, :36].float()
+            w_track = p.detach()[:, 36:].float()
+            out["patch_embedding.weight[:, :36]"] = {
+                "w_norm": w_base.norm().item(), "g_norm": 0.0,
+            }
+            out["patch_embedding.weight[:, 36:]"] = {
+                "w_norm": w_track.norm().item(), "g_norm": 0.0,
+            }
+            if p.grad is not None:
+                g_base = p.grad.detach()[:, :36].float()
+                g_track = p.grad.detach()[:, 36:].float()
+                out["patch_embedding.weight[:, :36]"]["g_norm"] = g_base.norm().item()
+                out["patch_embedding.weight[:, 36:]"]["g_norm"] = g_track.norm().item()
+        else:
+            out[n] = {"w_norm": w_norm, "g_norm": g_norm}
+    return out
+
+
+def one_step(model, s: Sample, device: torch.device, flow_shift: float,
+              from_ctx_ns) -> float:
+    """Forward+backward on one sample; grads accumulate."""
+    vae = s.vae_latent.unsqueeze(0).to(device, dtype=torch.bfloat16)
+    ff = s.first_frame_latent.unsqueeze(0).to(device, dtype=torch.bfloat16)
+    clip = s.clip_feature.unsqueeze(0).to(device, dtype=torch.bfloat16)
+    text = s.text_embedding.unsqueeze(0).to(device, dtype=torch.bfloat16)
+    mask = torch.ones(1, s.text_embedding.shape[0], device=device, dtype=torch.bfloat16)
+    tp = s.track_points.unsqueeze(0).to(device, dtype=torch.bfloat16)
+    tv = s.track_visibility.unsqueeze(0).to(device, dtype=torch.bfloat16)
+
+    num_latent_t = vae.shape[2]
+    expected_T = (num_latent_t - 1) * 4 + 1
+    tp = tp[:, :expected_T]
+    tv = tv[:, :expected_T]
+
+    sigma_u = float(torch.rand(1).item())
+    noise = torch.randn_like(vae, dtype=torch.float32)
+    clean_f = vae.float()
+    noisy_f, _, ts_val = flow_matching_noisy(clean_f, noise, sigma_u,
+                                              flow_shift=flow_shift)
+    noisy = noisy_f.to(torch.bfloat16)
+    cond20 = build_i2v_cond_concat(ff, vae_temporal_compression=4)
+    hs = torch.cat([noisy, cond20], dim=1)
+    ts = torch.tensor([ts_val], device=device, dtype=torch.bfloat16)
+
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16), from_ctx_ns(current_timestep=ts, attn_metadata=None):
+        pred = model(
+            hidden_states=hs, encoder_hidden_states=text, encoder_attention_mask=mask,
+            timestep=ts, encoder_hidden_states_image=clip,
+            track_points=tp, track_visibility=tv, return_dict=False,
+        )
+        target = noise - clean_f
+        loss = F.mse_loss(pred.float(), target.float())
+        loss.backward()
+    return float(loss.item())
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", default="/mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_random_init")
+    ap.add_argument("--n-samples", type=int, default=16)
+    ap.add_argument("--n-steps", type=int, default=60)
+    ap.add_argument("--batch-accum", type=int, default=4, help="gradient accumulation steps per optim step")
+    ap.add_argument("--high-lr", type=float, default=1e-3, help="LR for track pathway")
+    ap.add_argument("--base-lr", type=float, default=1e-5, help="LR for the rest of DiT")
+    ap.add_argument("--flow-shift", type=float, default=6.0)
+    ap.add_argument("--parquet-glob", type=str,
+                    default="/mnt/lustre/vlm-s4duan/openvid_1m/combined_parquet_dataset/shard_*/**/*.parquet")
+    ap.add_argument("--out", type=str,
+                    default="/mnt/lustre/vlm-s4duan/gradient_analysis_14b/lwlr_trace.json")
+    ap.add_argument("--seed", type=int, default=42)
+    args = ap.parse_args()
+
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+
+    from fastvideo.forward_context import set_forward_context
+    device = torch.device("cuda")
+
+    samples = load_samples(args.parquet_glob, args.n_samples, seed=args.seed)
+    print(f"loaded {len(samples)} samples")
+    assert len(samples) >= args.batch_accum, "need enough samples for accumulation"
+
+    print(f"loading {args.model} ...")
+    model = build_model_from_diffusers(args.model, device)
+
+    # Set up param groups
+    param_groups = split_param_groups(model, high_lr=args.high_lr, base_lr=args.base_lr)
+    n_high = sum(p.numel() for p in param_groups[0]["params"])
+    n_base = sum(p.numel() for p in param_groups[1]["params"])
+    print(f"track_pathway params: {n_high/1e6:.2f}M  (LR={args.high_lr:g})")
+    print(f"dit_body params:      {n_base/1e6:.2f}M  (LR={args.base_lr:g})")
+    optim = torch.optim.AdamW(param_groups, weight_decay=0.0)  # no wd so we see raw drift
+
+    trace: list[dict] = []
+    # step 0 (before any update)
+    print("\n=== step 0 (init) ===")
+    norms0 = snapshot_norms(model, -1)
+    for k, v in norms0.items():
+        print(f"  {k:<45s}  w={v['w_norm']:.4f}  g={v['g_norm']:.4f}")
+    trace.append({"step": 0, "loss": None, "norms": norms0})
+
+    step = 0
+    accum = 0
+    losses_accum = []
+    for outer in range(args.n_steps * args.batch_accum):
+        s = samples[outer % len(samples)]
+        loss = one_step(model, s, device, args.flow_shift, set_forward_context)
+        losses_accum.append(loss)
+        accum += 1
+        if accum < args.batch_accum:
+            continue
+        step += 1
+        # snapshot BEFORE step (grads are accumulated over the batch)
+        gs = snapshot_norms(model, -1)
+        optim.step()
+        optim.zero_grad(set_to_none=True)
+        avg_loss = float(np.mean(losses_accum))
+        losses_accum = []
+        accum = 0
+        if step % 5 == 0 or step <= 5:
+            print(f"\n=== step {step} (loss={avg_loss:.4f}) ===")
+            for k, v in gs.items():
+                marker = ""
+                if k == "track_encoder.proj.weight":
+                    marker = "  <-- 1.3B_merged saw per_sample_norm ~0.50"
+                elif k == "patch_embedding.weight[:, 36:]":
+                    marker = "  <-- 1.3B_merged saw per_sample_norm ~0.73"
+                print(f"  {k:<45s}  w={v['w_norm']:.4f}  g={v['g_norm']:.4f}{marker}")
+        trace.append({"step": step, "loss": avg_loss, "norms": gs})
+
+    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
+    with open(args.out, "w") as f:
+        json.dump({"config": vars(args), "trace": trace}, f, indent=2)
+    print(f"\nwrote trace -> {args.out}")
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/merge_synth_manifests.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Compile per-worker manifest shards -> FastVideo videos2caption.json + merge.txt.
+Idempotent; run anytime (progress check) or at the end. Dedups by idx, verifies mp4 exists."""
+from __future__ import annotations
+import argparse, json
+from pathlib import Path
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--output-dir", type=Path, required=True)
+    a = p.parse_args()
+    videos_dir = a.output_dir / "videos"
+    shard_dir = a.output_dir / "manifest_shards"
+    recs: dict[int, dict] = {}
+    for sh in sorted(shard_dir.glob("worker_*.jsonl")):
+        for ln in sh.read_text().splitlines():
+            ln = ln.strip()
+            if not ln:
+                continue
+            try:
+                r = json.loads(ln)
+            except json.JSONDecodeError:
+                continue
+            if (videos_dir / r["path"]).exists():
+                recs[int(r["idx"])] = r
+    ordered = [recs[k] for k in sorted(recs)]
+    j = a.output_dir / "videos2caption.json"
+    j.write_text(json.dumps(ordered, indent=2))
+    (a.output_dir / "merge.txt").write_text(f"{videos_dir.resolve()},{j.resolve()}\n")
+    # also count mp4s on disk (may exceed manifest if a worker died mid-write)
+    on_disk = len(list(videos_dir.glob("vid_*.mp4")))
+    print(f"manifest: {len(ordered)} entries; {on_disk} mp4s on disk -> {j}")
+
+if __name__ == "__main__":
+    main()

data_pipeline/notes/DECISIONS.md

@@ -0,0 +1,65 @@
+# Data Pipeline — Decisions & Rationale
+
+## Why generation + tracking live OUTSIDE `fastvideo/pipelines/preprocess/`
+FastVideo's preprocess pipeline is structurally a *consumer of existing .mp4 files*
+(`VideoCaptionMergedDataset` reads a `merge.txt`→`annotation.json` over a folder and
+asserts each video path exists). There is no seam to *produce* a video inside it. So
+generation (Wan2.2) and tracking (CoTracker) are standalone Stage-0 scripts that
+produce the files preprocess later ingests — matching the repo convention (everything
+upstream of preprocess is plain scripts, not pipeline stages).
+
+## Two scripts, not one fused process
+- The `.mp4` on disk is the checkpoint; generation is the expensive part. Tracking
+  config will be tweaked many times — must not re-run Wan each time.
+- Wan2.2-14B is huge (multi-GPU + offload); CoTracker is tiny. "Generate all → free
+  model → track all" is the natural ordering and is exactly two phases.
+- Independent restart/parallelism. Both scripts are idempotent (skip done work).
+
+## Generate at training fps/length (no temporal alignment headache)
+CoTracker tracks per *video* frame, but the VAE compresses time (~4× for Wan). If we
+generated at arbitrary fps we'd have to re-index tracks by the preprocess frame-sampler
+(`sample_frame_index`) and then fold onto latent time. Since we *control* generation,
+we emit videos already at the train fps/length (default 16 fps, 81 frames ≈ 5 s) so
+source frames == sampled frames, 1:1. Tracks then align to frames trivially; the only
+remaining fold (frames→latent-frames) happens in the model/trainer.
+
+## Generation is T2V even though the model is I2V+points
+Synthetic videos come from text→video (Wan2.2-T2V-A14B). At *training* time the model
+uses frame 0 as the I2V conditioning image + the point tracks as motion control. So we
+never need an input image during generation.
+
+## Manifest format (FastVideo-compatible)
+`generate_videos.py` emits the same shape the existing loader expects:
+- `videos/vid_NNNNNN.mp4`
+- `videos2caption.json`: list of `{path(basename), cap:[...], fps, duration, num_frames, resolution}`
+- `merge.txt`: one line `<videos_dir>,<videos2caption.json>`
+Loader actually consumes only `path, cap, fps, duration` (+ optional conditioning path);
+`resolution`/`num_frames` are informational. We know all of these at generation time, so
+we write the manifest directly and SKIP `scripts/dataset_preparation/prepare_json_file.py`
+(it exists only to *recover* fps/duration by re-probing videos of unknown provenance).
+
+## points_path stored as ABSOLUTE path
+The future preprocess loader joins `folder + relative_path`. Tracks live in a sibling
+`tracks/` dir, not under `videos/`. Storing `points_path` absolute makes `os.path.join`
+return it unchanged, so it resolves regardless of the manifest folder. Mirrors how
+MatrixGame2 references `action_path`, which is the pattern the points preprocess task
+will copy.
+
+## CoTracker v3 via torch.hub (`cotracker3_offline`)
+The repo only vendors CoTracker **v2** under `third_party/eval/vbench` (eval-only). For
+v3 we pull `facebookresearch/co-tracker` `cotracker3_offline` via torch.hub. Prefetch on
+the login node (has internet); the hub cache lives on shared weka so compute nodes reuse
+it offline (`extract_tracks.py` falls back to `source="local"` from the cache dir).
+
+## Output location OUTSIDE the git repo
+Generated videos are large (≈2–5 MB each → ~25 GB at 5k). Default output dir is
+`/mnt/weka/home/hao.zhang/shao/data/motion_pipeline/...` (outside `FastVideo/`) to avoid
+bloating the working tree / accidental commits.
+
+## Cluster
+- `shao_wm` = SLURM job **1788946** on node **fs-mbz-gpu-538**: 8×GPU (~140 GB each),
+  128 CPU, 768 GB RAM, 5-day limit. As of setup all 8 GPUs were idle.
+- Attach with `srun --jobid=1788946 --overlap ...`. Never run heavy work on the login
+  node (`fs-mbz-login-big-001`) — it lags the shared connection.
+- `sqs` is a `.bashrc` alias (`squeue -u hao.zhang`); not available in non-interactive
+  shells — use `squeue -u $USER` directly.

data_pipeline/notes/PROGRESS.md

@@ -0,0 +1,40 @@
+# Data Pipeline — Progress Log
+
+Project: build a `(video, prompt, point-tracks)` dataset to train a motion-controlled
+real-time interactive video model (bidirectional finetune stage first).
+
+Branch: `shao/realtime-bidir`. This `data_pipeline/` dir holds the **Stage-0** data
+generation scripts (everything *upstream* of FastVideo's `preprocess` pipeline).
+
+## Pipeline shape (where we are)
+
+```
+vidprom prompts ──> generate_videos.py ──> videos/*.mp4 + videos2caption.json + merge.txt
+                                              │
+                    extract_tracks.py  ──> tracks/*.npz  (+ patch points_path into the json)
+                                              │
+                    [TODO] v1_preprocess.py --preprocess_task i2v_points ──> parquet
+                                              │
+                    [TODO] bidir I2V+points trainer
+```
+
+## Status
+
+- [x] Located existing FastVideo preprocess stack + MatrixGame2 control-signal pattern (the template).
+- [x] Confirmed env = repo `.venv` (`/mnt/weka/home/hao.zhang/shao/FastVideo/.venv/bin/python`).
+- [x] Confirmed Wan2.2-T2V-A14B weights cached (`~/.cache/huggingface/hub/models--Wan-AI--Wan2.2-T2V-A14B-Diffusers`).
+- [x] Wrote `generate_videos.py` (Wan2.2-14B T2V → mp4 + manifest).
+- [x] Wrote `extract_tracks.py` (CoTracker v3 50×50 grid → npz, patch manifest).
+- [~] Smoke-test launched on fs-mbz-gpu-538 (`run.sh --smoke`, background). Node validated:
+      internet=yes (downloads OK on the node), 8 GPUs idle, `.venv` imports fastvideo + torch
+      2.11.0/CUDA, cotracker3_offline loads (25.4M params). Prompts downloaded: 248,221 (140 MB).
+- [ ] Scale to 50, then 5k.
+- [ ] Extend preprocess with an `i2v_points` task (mirror MatrixGame2).
+
+## Open questions / decisions pending
+
+- Training model for the bidir stage: 1.3B/480p for plumbing bring-up vs straight to 14B (data gen is fixed at 14B/720p regardless — user confirmed).
+- Confirm conditioning shape: I2V + points (input image + per-frame points → video). Assumed yes.
+- Point sampling at train time: store full 2500-pt grid + visibility; sample 1–200 in the trainer.
+
+See `DECISIONS.md` for rationale, `RUNBOOK.md` for exact commands.

data_pipeline/notes/RUNBOOK.md

@@ -0,0 +1,70 @@
+# Data Pipeline — Runbook
+
+All heavy work runs on the `shao_wm` allocation via `srun --overlap`. Do NOT run model
+code on the login node.
+
+## 0. Env + node handles
+```bash
+PY=/mnt/weka/home/hao.zhang/shao/FastVideo/.venv/bin/python
+JOBID=$(squeue -u "$USER" -n shao_wm -h -o %i | head -1)   # shao_wm allocation id
+echo "jobid=$JOBID"
+# live GPU usage on the node (lightweight):
+srun --jobid="$JOBID" --overlap nvidia-smi \
+  --query-gpu=index,memory.used,memory.total,utilization.gpu --format=csv,noheader
+```
+
+## 1. One-time setup (login node — has internet)
+```bash
+cd /mnt/weka/home/hao.zhang/shao/FastVideo
+# prompts (self-forcing / vidprom):
+[ -f examples/dataset/vidprom/prompts/vidprom_filtered_extended.txt ] || \
+  ( cd examples/dataset/vidprom && ./download_dataset.sh )
+# prefetch CoTracker v3 into the shared torch.hub cache so compute nodes work offline:
+$PY -c "import torch; torch.hub.load('facebookresearch/co-tracker','cotracker3_offline'); print('cotracker ok')"
+```
+
+**Pick free GPUs yourself** from step 0 and pass them via `CUDA_VISIBLE_DEVICES` — the node
+is shared and the scripts do NOT auto-select. Pinning to busy GPUs OOMs.
+
+## 2. Generate videos (Wan2.2-T2V-A14B, 720p, 16fps, 81f)
+```bash
+GPUS=2,3   # <- two currently-free GPUs from step 0
+srun --jobid="$JOBID" --overlap --ntasks=1 env CUDA_VISIBLE_DEVICES=$GPUS \
+  $PY data_pipeline/generate_videos.py \
+  --prompts examples/dataset/vidprom/prompts/vidprom_filtered_extended.txt \
+  --output-dir /mnt/weka/home/hao.zhang/shao/data/motion_pipeline/wan22_t2v_720p \
+  --num-videos 50 --num-gpus 2          # num-gpus must match the device count
+```
+- Smoke first: `--num-videos 2`.
+- Idempotent: re-run resumes from `manifest.jsonl`.
+
+## 3. Extract point tracks (CoTracker v3, 50×50 grid)
+```bash
+srun --jobid="$JOBID" --overlap --ntasks=1 env CUDA_VISIBLE_DEVICES=$GPUS \
+  $PY data_pipeline/extract_tracks.py \
+  --data-dir /mnt/weka/home/hao.zhang/shao/data/motion_pipeline/wan22_t2v_720p \
+  --grid-size 50
+```
+- Idempotent: skips videos whose `tracks/<stem>.npz` exists.
+- If 720p OOMs CoTracker, add `--downscale 0.5` (coords are rescaled back to original px).
+
+## 4. One-shot wrapper (derives num_gpus from CUDA_VISIBLE_DEVICES)
+```bash
+CUDA_VISIBLE_DEVICES=2,3 bash data_pipeline/run.sh            # full run (defaults)
+CUDA_VISIBLE_DEVICES=2,3 bash data_pipeline/run.sh --smoke    # 2 videos
+```
+
+## Outputs
+```
+<output-dir>/
+  videos/vid_000000.mp4 ...
+  tracks/vid_000000.npz ...        # keys: tracks (T,N,2 px), visibility (T,N), grid_size, height, width, fps, num_frames
+  videos2caption.json              # FastVideo manifest (+ points_path after step 3)
+  merge.txt
+  manifest.jsonl                   # incremental resume log (generation)
+```
+
+## Gotchas
+- Check GPU availability before launching (step 0) — the node may not be fully free.
+- `generate_video(...)` logs a Deprecation warning (use of legacy API) — harmless.
+- First Wan2.2 run loads ~14B MoE weights; expect a few min of startup before frame 1.

data_pipeline/notes/pipeline.md

@@ -0,0 +1,148 @@
+# Data Pipeline
+
+All commands run from `/home/hal-kevin/FastVideo`.
+
+---
+
+## Stage 1 — Generate Videos
+
+```bash
+python data_pipeline/generate_videos.py \
+  --prompts examples/dataset/motion-test/prompts.txt \
+  --output-dir /home/hal-kevin/data/motion-stream-test \
+  --num-videos 100 \
+  --num-gpus 4 \
+  --trim-start-frames 10 \
+  --num-inference-steps 40
+```
+
+Output: `motion-physics/videos/vid_000000.mp4 ...`
+
+---
+
+## Stage 2 — VAE Round-Trip Videos
+
+Encode then decode each video through the FastVideo WanVAE so CoTracker runs on
+the same frames that training will see (Stage 5 re-encodes these, so tracks must
+be extracted from the decoded version, not the raw source).
+
+```bash
+python data_pipeline/decode_roundtrip_videos.py \
+  --data-dir /home/hal-kevin/data/motion-stream-test \
+  --vae-path /home/hal-kevin/models/trackwan_1.3b_i2v_control_init/vae
+```
+
+Output: `motion-physics/roundtrip_videos/vid_000000.mp4 ...`
+
+---
+
+## Stage 3 — Extract Tracks
+
+Run CoTracker on the round-trip videos, parallelized across 4 GPUs:
+
+```bash
+bash data_pipeline/run_extract_tracks.sh
+```
+
+Pass extra args (e.g. `--force`, `--limit 5`) directly — they are forwarded to each worker.
+The script no longer hardcodes `--force`; existing `.npz` are skipped unless you pass it:
+
+```bash
+bash data_pipeline/run_extract_tracks.sh --force
+```
+
+Speed knobs: `--sam-batch 16` (frames per batched FastSAM forward, default 16) and `--amp`
+(bf16 autocast for CoTracker, ~1.5-2x faster but slightly different coords — validate before
+adopting). Entry events now share one extra CoTracker pass (queries filtered to the new-object
+regions) instead of a full 2500-point pass per mask.
+
+**Fused mode:** `--segment` (with `--vis-override-every 2`) runs Stage 4 inside this pass —
+object IDs, vis override, and track weights — reusing the decoded video and the entry-detection
+FastSAM masks, so Stage 4 does not need to run at all. Same results as the standalone stage
+(shared implementation). `--viz`/`--viz-dir` render the same overlay mp4s as standalone Stage 4
+(slow — skip for large-scale runs); `--min-area-frac`/`--max-masks` remain standalone-only.
+Benchmark it with `FUSED=1 bash data_pipeline/benchmark_tracks.sh` (add `VIZ=1` for renders).
+
+Single-GPU alternative:
+
+```bash
+python data_pipeline/extract_tracks.py \
+  --data-dir /home/hal-kevin/data/motion-stream-test \
+  --videos-subdir roundtrip_videos \
+  --grid-size 50 \
+  --device cuda \
+  --detect-entries \
+  --sam-conf 0.75 \
+  --sam-iou 0.9 \
+  --sam-imgsz 1024
+```
+
+Output: `motion-stream-test/tracks/vid_000000.npz ...`
+
+---
+
+## Stage 4 — Segment Tracks
+
+Assign object IDs and compute motion weights, parallelized across 4 GPUs:
+
+```bash
+bash data_pipeline/run_segment_tracks.sh
+```
+
+The script no longer hardcodes `--force` (pass it to re-process npz that already have
+`object_ids`). Each video is now decoded once and FastSAM runs in batched forwards
+(`--sam-batch 16`), shared between object-ID assignment and the vis override sweep.
+
+Single-GPU alternative:
+
+```bash
+python data_pipeline/segment_tracks.py \
+  --data-dir /home/hal-kevin/data/motion-stream-test \
+  --videos-subdir roundtrip_videos \
+  --conf 0.75 --iou 0.9 --imgsz 1024 \
+  --vis-override-every 2 \
+  --force \
+  --viz
+```
+
+Adds `object_ids`, `n_objects`, `track_weights` to each `.npz`.
+
+---
+
+## Stage 5 — Preprocess to Parquet
+
+Reads raw `videos/` for VAE encoding and `tracks/` for track data.
+
+```bash
+torchrun --nproc_per_node=1 -m fastvideo.pipelines.preprocess.v1_preprocess \
+  --model_path /home/hal-kevin/models/trackwan_1.3b_i2v_control_init \
+  --data_merge_path /home/hal-kevin/data/motion-stream-test/data_merge.txt \
+  --output_dir /home/hal-kevin/data/motion-stream-test/preprocessed_i2v_track \
+  --preprocess_task i2v_track \
+  --num_frames 121 \
+  --num_latent_t 31 \
+  --train_fps 24 \
+  --max_height 480 \
+  --max_width 832 \
+  --preprocess_video_batch_size 1 \
+  --samples_per_file 64
+```
+
+`--train_fps` must match the source video fps (24 here). If omitted it defaults to 30,
+and `FrameSamplingStage` resamples with interval `fps/train_fps = 0.8` — duplicating
+every 5th frame and covering only the first ~97 of 121 frames. The stored latents then
+encode slowed, stuttering motion that no longer aligns with the tracks (extracted at
+native fps), which shows up as drifting motion in validation reference videos.
+
+Output: `motion-physics/preprocessed_i2v_track/combined_parquet_dataset/`
+
+---
+
+## Stage 6 — Training
+
+```bash
+python -m fastvideo.train.train \
+  --config examples/train/scenario/worldmodel/finetune_wantrack_golf_overfit.yaml
+```
+
+Update the yaml to point at your data and checkpoint directories.

data_pipeline/notes/processed_dataset_README_bf16.md

@@ -0,0 +1,97 @@
+---
+license: apache-2.0
+task_categories:
+- text-to-video
+- image-to-video
+tags:
+- video-generation
+- point-tracking
+- motionstream
+- wantrack
+- fastvideo
+---
+
+# OpenVid-WanTrack Processed (v2, 720p, **bf16**)
+
+FastVideo preprocessing parquets for training the TrackWan point-track-conditioned I2V model on
+the OpenVid-derived WanTrack set. Each row is one 121-frame clip with its VAE latents, text and
+image conditioning, and dense CoTracker3 tracks — everything the trainer memory-maps, so no video
+decoding happens at train time.
+
+**This is the `bfloat16` variant** of `…/openvid-wantrack-processed` (v2, 720p): the large float
+tensor fields are stored in **bf16** instead of float32, so the dataset is roughly **half the size**.
+Everything else (clips, ids, shapes, layout) is identical.
+
+## ⚠️ Precision — read before loading
+
+- The big tensors — `vae_latent`, `first_frame_latent`, `clip_feature`, `text_embedding`,
+  `track_points`, `track_visibility` — are **`bfloat16`**. Each field's `_dtype` column says so.
+- `object_ids` and `track_weights` are kept **`float32`** (small integer/label fields).
+- **You must honor the per-field `_dtype` when decoding.** numpy has **no** bfloat16, so
+  `np.frombuffer(bytes, "bfloat16")` fails — decode via `torch.frombuffer` (see Loading below).
+- **Quality is unaffected for training:** the TrackWan trainer already downcasts these fields to
+  bf16 before use, so storing bf16 just pre-applies the exact rounding the model does anyway.
+- The FastVideo trainer's loader honors `_dtype`, so pointing `data_path` at this set "just works".
+
+## Layout
+
+```
+shard000/combined_parquet_dataset/worker_*/data_chunk_*.parquet
+shard001/combined_parquet_dataset/worker_*/data_chunk_*.parquet
+...
+shard259/...                      # shard259 is a 110-clip remainder; all others are 1000
+```
+
+~259,110 clips across 260 shards. Clip ids join 1:1 with `openvid-wantrack-clips` (videos),
+`openvid-wantrack-tracks-v2` (raw npz tracks), and OpenVid-1M captions.
+
+## Row schema (`pyarrow_schema_i2v_track`, 33 columns)
+
+Scalars: `id, file_name, caption, media_type, width, height, num_frames, duration_sec, fps`.
+
+Tensors — each stored as a triplet `<name>_bytes` (raw buffer), `_shape` (list<int64>), `_dtype`:
+
+| tensor | shape (720p) | dtype | description |
+|--------|--------------|-------|-------------|
+| `vae_latent` | `[16, 31, 90, 160]` | **bfloat16** | WanVAE latent of the clip (training target) |
+| `first_frame_latent` | `[16, 31, 90, 160]` | **bfloat16** | I2V conditioning: VAE-encode of `[frame0, zeros...]` |
+| `clip_feature` | `[257, 1280]` | **bfloat16** | CLIP image embedding of frame 0 |
+| `text_embedding` | `[L, 4096]` | **bfloat16** | T5 caption embedding (variable length `L`, padding stripped) |
+| `track_points` | `[121, 2500, 2]` | **bfloat16** | CoTracker tracks, **normalized [0,1]** |
+| `track_visibility` | `[121, 2500]` | **bfloat16** | per-frame visibility |
+| `object_ids` | `[2500]` | float32 | FastSAM object id per track (-1 = background) |
+| `track_weights` | `[2500]` | float32 | low-rank motion weight in [0,1] |
+
+`num_frames=31` for the latents (VAE 4x temporal compression: `(121-1)/4+1`); `track_points`
+stay at native `121`. Text embedding length varies per row (padding removed), so read the
+per-row `_shape`.
+
+## Config
+
+- Video: 1280x720, 121 frames, 24 fps
+- VAE: FastVideo WanVAE (latents encoded in fp32, **stored as bf16**), `use_feature_cache=True`
+- CLIP: frame-0 image embedding; T5: caption text embedding
+- Tracks: CoTracker3, 50x50 grid (2500 points), FastSAM segmentation
+
+## Loading
+
+numpy cannot represent bfloat16, so decode through torch, honoring each field's `_dtype`:
+
+```python
+import glob, torch, pyarrow.parquet as pq
+
+_STR2T = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
+
+def decode(row, name):
+    dt = _STR2T[row[f"{name}_dtype"]]
+    # bytearray() -> writable buffer that doesn't alias the parquet row
+    return torch.frombuffer(bytearray(row[f"{name}_bytes"]), dtype=dt).reshape(row[f"{name}_shape"])
+
+files = glob.glob("**/*.parquet", recursive=True)          # all shards
+row = pq.read_table(files[0]).slice(0, 1).to_pylist()[0]
+lat = decode(row, "vae_latent")        # torch.bfloat16, shape [16, 31, 90, 160]
+tracks = decode(row, "track_points")   # torch.bfloat16, normalized [0,1]
+```
+
+The FastVideo trainer discovers all parquets under the dataset root via `os.walk` and its loader
+honors the `_dtype` column, so point `data_path` at the directory containing the `shard*/` folders.

data_pipeline/openvid_download.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Download + extract OpenVid / OpenVidHD video shards from HuggingFace.
+
+Videos are only distributed as ~30-50 GB zip shards (no per-file fetch), so this
+pulls whole shards, extracts the .mp4s into --videos-dir, and (unless --keep-zip)
+deletes the zip. Pair with openvid_filter.py (which clips to keep) + extract_tracks_mp.py
+(resolution/aspect probe + tracking). Idempotent: skips shards already extracted.
+
+Some parts are HF-split (e.g. OpenVid_part102_partaa/ab) — those must be
+concatenated before unzip; this handles single-file parts (all OpenVidHD parts).
+"""
+from __future__ import annotations
+import argparse, zipfile
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+
+REPO = "nkp37/OpenVid-1M"
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--shards", nargs="+", required=True,
+                    help="repo paths, e.g. OpenVidHD/OpenVidHD_part_1.zip OpenVid_part0.zip")
+    ap.add_argument("--videos-dir", required=True, help="where .mp4s are extracted")
+    ap.add_argument("--zip-dir", default=None, help="scratch for downloaded zips (default: <videos-dir>/../_zips)")
+    ap.add_argument("--keep-zip", action="store_true", help="don't delete the zip after extract")
+    ap.add_argument("--limit", type=int, default=0, help="extract at most N mp4s per shard (0=all; for smoke tests)")
+    ap.add_argument("--only-list", default=None, help="file of basenames; extract only these (e.g. the ≥5s filtered list)")
+    a = ap.parse_args()
+    only = None
+    if a.only_list:
+        only = {l.strip() for l in open(a.only_list) if l.strip()}
+
+    vdir = Path(a.videos_dir); vdir.mkdir(parents=True, exist_ok=True)
+    zdir = Path(a.zip_dir) if a.zip_dir else vdir.parent / "_zips"
+    zdir.mkdir(parents=True, exist_ok=True)
+    marker_dir = vdir.parent / "_extracted"; marker_dir.mkdir(parents=True, exist_ok=True)
+
+    for shard in a.shards:
+        done_marker = marker_dir / (Path(shard).name + ".done")
+        if done_marker.exists():
+            print(f"[dl] skip {shard} (already extracted)", flush=True); continue
+        print(f"[dl] downloading {shard} ...", flush=True)
+        zp = hf_hub_download(REPO, shard, repo_type="dataset", local_dir=str(zdir))
+        print(f"[dl] extracting {zp} -> {vdir}", flush=True)
+        n = 0
+        with zipfile.ZipFile(zp) as z:
+            for m in z.namelist():
+                if m.lower().endswith(".mp4"):
+                    if only is not None and Path(m).name not in only:
+                        continue
+                    # flatten: write basename directly into videos-dir
+                    (vdir / Path(m).name).write_bytes(z.read(m))
+                    n += 1
+                    if a.limit and n >= a.limit:
+                        break
+        done_marker.write_text(str(n))
+        if not a.keep_zip:
+            Path(zp).unlink(missing_ok=True)
+        print(f"[dl] {shard}: extracted {n} mp4s", flush=True)
+    print("[dl] done", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/openvid_download_hd.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Robust OpenVidHD (1080p, 16:9) full downloader for the trackwan bidir dataset.
+
+OpenVidHD parts are a mix of single-file `.zip` (parts 1-14) and SPLIT parts
+(`OpenVidHD_part_<i>_part_aa` + `_part_ab`) that must be concatenated before unzip
+(mirrors the official download_scripts/download_OpenVid.py). This enumerates the HF
+repo, plans each part, downloads (single or split+cat), extracts only the mp4s in
+--only-list (the >=5s filtered set), writes them into --videos-dir, and deletes the
+zip. Idempotent via per-part .done markers, so it resumes after interruption.
+
+Shard by --shard/--num-shards (e.g. SLURM env) to spread the download across nodes.
+"""
+from __future__ import annotations
+import argparse, os, re, subprocess, zipfile, collections
+from pathlib import Path
+from huggingface_hub import HfApi, hf_hub_download
+
+REPO = "nkp37/OpenVid-1M"
+
+
+def plan_parts():
+    """Return {part_num: [repo_files]} for OpenVidHD (single .zip or split _part_aa/ab)."""
+    files = HfApi().list_repo_files(REPO, repo_type="dataset")
+    parts = collections.defaultdict(list)
+    for f in files:
+        if not f.startswith("OpenVidHD/"):
+            continue
+        m = re.search(r"OpenVidHD_part_(\d+)", f)
+        if m:
+            parts[int(m.group(1))].append(f)
+    return {k: sorted(v) for k, v in parts.items()}
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--videos-dir", required=True)
+    ap.add_argument("--zip-dir", required=True, help="scratch for zips (deleted after extract unless --keep-zip)")
+    ap.add_argument("--only-list", required=True, help="basenames to keep (the >=5s filtered list)")
+    ap.add_argument("--keep-zip", action="store_true")
+    ap.add_argument("--num-shards", type=int, default=int(os.environ.get("SLURM_NTASKS", "1")))
+    ap.add_argument("--shard", type=int, default=int(os.environ.get("SLURM_PROCID", "0")))
+    ap.add_argument("--parts", default="all", help="'all' or comma range like 1-14 / 1,2,3")
+    a = ap.parse_args()
+
+    vdir = Path(a.videos_dir); vdir.mkdir(parents=True, exist_ok=True)
+    zdir = Path(a.zip_dir); zdir.mkdir(parents=True, exist_ok=True)
+    mdir = vdir.parent / "_extracted"; mdir.mkdir(parents=True, exist_ok=True)
+    only = {l.strip() for l in open(a.only_list) if l.strip()}
+
+    parts = plan_parts()
+    keys = sorted(parts)
+    if a.parts != "all":
+        want = set()
+        for tok in a.parts.split(","):
+            if "-" in tok:
+                lo, hi = tok.split("-"); want |= set(range(int(lo), int(hi) + 1))
+            else:
+                want.add(int(tok))
+        keys = [k for k in keys if k in want]
+    keys = keys[a.shard::a.num_shards]
+    print(f"[hd-dl shard {a.shard}/{a.num_shards}] {len(keys)} parts: {keys[:6]}{'...' if len(keys)>6 else ''}", flush=True)
+
+    for i in keys:
+        marker = mdir / f"OpenVidHD_part_{i}.done"
+        if marker.exists():
+            continue
+        fs = parts[i]
+        zip_path = zdir / f"OpenVidHD_part_{i}.zip"
+        try:
+            if len(fs) == 1 and fs[0].endswith(".zip"):
+                p = hf_hub_download(REPO, fs[0], repo_type="dataset", local_dir=str(zdir))
+                zip_path = Path(p)
+            else:  # split: download _part_aa/_part_ab, concat
+                locals_ = []
+                for f in fs:
+                    locals_.append(hf_hub_download(REPO, f, repo_type="dataset", local_dir=str(zdir)))
+                with open(zip_path, "wb") as out:
+                    for lp in sorted(locals_):
+                        with open(lp, "rb") as src:
+                            while (chunk := src.read(1 << 24)):
+                                out.write(chunk)
+                for lp in locals_:
+                    Path(lp).unlink(missing_ok=True)
+            n = 0
+            with zipfile.ZipFile(zip_path) as z:
+                for m in z.namelist():
+                    if m.lower().endswith(".mp4") and Path(m).name in only:
+                        (vdir / Path(m).name).write_bytes(z.read(m))
+                        n += 1
+            marker.write_text(str(n))
+            if not a.keep_zip:
+                zip_path.unlink(missing_ok=True)
+            print(f"[hd-dl] part_{i}: extracted {n} clips", flush=True)
+        except Exception as e:  # noqa: BLE001
+            print(f"[hd-dl] part_{i} FAILED: {repr(e)[:150]}", flush=True)
+    print(f"[hd-dl shard {a.shard}] done", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/openvid_filter.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stage-1 (metadata) filter for OpenVid — from the dataset CSV.
+
+The OpenVid CSV columns are: video, caption, aesthetic score, motion score,
+temporal consistency score, camera motion, frame, fps, seconds.
+There is NO resolution / aspect-ratio column, so 720p + 16:9 CANNOT be filtered
+here — those are enforced downstream by probing (extract_tracks_mp.py
+--min-height / --aspect-tol). This stage keeps clips that are long enough to yield
+`num_frames` at the target fps, with optional motion/quality gates (tracking wants
+motion; static/low-motion clips give degenerate tracks).
+
+Output: newline-delimited video basenames to keep (feeds the shard extraction +
+tracking stages). Also writes a captions sidecar (video -> caption) for later use.
+"""
+from __future__ import annotations
+import argparse, csv, json, sys
+
+csv.field_size_limit(sys.maxsize)  # captions are long
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--csv", required=True, help="OpenVid-1M.csv or OpenVidHD.csv")
+    ap.add_argument("--out", required=True, help="output: one video basename per line")
+    ap.add_argument("--captions-out", default=None, help="optional JSON {video: caption}")
+    ap.add_argument("--num-frames", type=int, default=121)
+    ap.add_argument("--target-fps", type=int, default=24)
+    ap.add_argument("--min-seconds", type=float, default=None,
+                    help="override; default = num_frames/target_fps (+2%% margin)")
+    ap.add_argument("--min-motion", type=float, default=0.0, help="drop clips below this motion score")
+    ap.add_argument("--min-aesthetic", type=float, default=0.0)
+    ap.add_argument("--exclude-static", action="store_true", help="drop camera motion == 'static'")
+    a = ap.parse_args()
+
+    min_secs = a.min_seconds if a.min_seconds is not None else (a.num_frames / a.target_fps) * 1.02
+    kept = 0; total = 0; caps = {}
+    with open(a.csv, newline="") as f, open(a.out, "w") as o:
+        r = csv.DictReader(f)
+        for row in r:
+            total += 1
+            try:
+                secs = float(row["seconds"]); motion = float(row["motion score"]); aes = float(row["aesthetic score"])
+            except (KeyError, ValueError):
+                continue
+            if secs < min_secs:            continue
+            if motion < a.min_motion:      continue
+            if aes < a.min_aesthetic:      continue
+            if a.exclude_static and row.get("camera motion", "").strip().lower() == "static":
+                continue
+            o.write(row["video"] + "\n"); kept += 1
+            if a.captions_out is not None:
+                caps[row["video"]] = row.get("caption", "")
+    if a.captions_out is not None:
+        json.dump(caps, open(a.captions_out, "w"))
+    print(f"[filter] min_seconds={min_secs:.3f} min_motion={a.min_motion} min_aes={a.min_aesthetic} "
+          f"exclude_static={a.exclude_static}")
+    print(f"[filter] kept {kept}/{total} ({100*kept/max(total,1):.1f}%) -> {a.out}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/precompute_null_text.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python3
+"""Precompute the UMT5 embedding of an empty string.
+
+This is the canonical Wan/T5 null-text convention: feed "" to the text encoder
+and use the resulting non-zero embedding as the ∅ in classifier-free guidance.
+Contrast with `torch.zeros_like(text_embedding)`, which is a different tensor
+distribution that the base model was never asked to handle.
+
+Saves a dict:
+  {"embedding": [seq_len, dim], "attention_mask": [seq_len]}
+so the val callback can just index into these when it needs a null text.
+
+Usage:
+  python data_pipeline/precompute_null_text.py \
+    --model-dir /mnt/lustre/vlm-s4duan/exports/synth_stage2_paperLR_ckpt400 \
+    --output /mnt/lustre/vlm-s4duan/exports/null_text_umt5.pt
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import torch
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--model-dir", required=True,
+                   help="Diffusers export dir containing text_encoder/ and tokenizer/")
+    p.add_argument("--output", required=True)
+    p.add_argument("--text-len", type=int, default=256,
+                   help="Pad/truncate to this length (matches training text_padding_length)")
+    p.add_argument("--device", default="cuda")
+    args = p.parse_args()
+
+    from transformers import AutoTokenizer
+    try:
+        from transformers import UMT5EncoderModel  # type: ignore
+    except ImportError:
+        UMT5EncoderModel = None  # type: ignore
+    from transformers import AutoModel
+
+    model_dir = Path(args.model_dir)
+    tok = AutoTokenizer.from_pretrained(model_dir / "tokenizer")
+    if UMT5EncoderModel is not None:
+        try:
+            enc = UMT5EncoderModel.from_pretrained(model_dir / "text_encoder", torch_dtype=torch.bfloat16)
+        except Exception:
+            enc = AutoModel.from_pretrained(model_dir / "text_encoder", torch_dtype=torch.bfloat16)
+    else:
+        enc = AutoModel.from_pretrained(model_dir / "text_encoder", torch_dtype=torch.bfloat16)
+    enc = enc.to(args.device).eval()
+
+    tokens = tok(
+        [""],
+        max_length=args.text_len,
+        padding="max_length",
+        truncation=True,
+        return_tensors="pt",
+        return_attention_mask=True,
+    )
+    input_ids = tokens["input_ids"].to(args.device)
+    attn_mask = tokens["attention_mask"].to(args.device)
+    print(f"tokenized '' -> {input_ids.shape}, non-pad tokens = {int(attn_mask.sum().item())}")
+    print(f"  first 5 token ids: {input_ids[0, :5].tolist()}")
+
+    with torch.no_grad():
+        out = enc(input_ids=input_ids, attention_mask=attn_mask)
+    # take last_hidden_state; strip the batch dim
+    hs = out.last_hidden_state[0].float().cpu()  # [seq_len, dim]
+    am = attn_mask[0].cpu()
+    print(f"embedding shape={list(hs.shape)}  norm={hs.norm().item():.3f}  mean={hs.mean().item():.5f}")
+    print(f"attn_mask shape={list(am.shape)}  sum={int(am.sum().item())}")
+
+    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+    torch.save({"embedding": hs, "attention_mask": am, "text_len": args.text_len}, args.output)
+    print(f"[null-text] saved -> {args.output}")
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/resize_videos.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stage 2 (VAE-free variant): crop+resize source videos to the training geometry.
+
+Applies exactly the same ``center_crop_th_tw`` + ``resize`` transform as
+``decode_roundtrip_videos.py`` (and as Stage 5's ``CenterCropResizeVideo``), but skips
+the VAE encode/decode. Tracks extracted from these videos land in the same coordinate
+frame as the training latents; they just don't carry the VAE's reconstruction artifacts.
+
+Purpose: the geometry is what track alignment *requires*; the VAE round-trip is what it
+*may* require. This script exists so the two can be A/B'd -- extract tracks from
+``resized_videos/`` and from ``roundtrip_videos/``, then diff the npz. If the track delta
+is small, large-scale runs can skip the VAE pass entirely (it is pure GPU cost per clip).
+
+Usage:
+    python data_pipeline/resize_videos.py \\
+        --data-dir /home/hal-kevin/data/motion-stream-test
+
+CPU-only; parallelize with --index / --limit sharding if needed.
+"""
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import imageio
+import numpy as np
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from fastvideo.dataset.transform import center_crop_th_tw, resize
+
+TARGET_H, TARGET_W = 480, 832
+NUM_FRAMES = 121
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--data-dir", type=Path, required=True, help="Dataset root (contains videos/, etc.).")
+    p.add_argument("--video-subdir", type=str, default="videos", help="Input video subdirectory.")
+    p.add_argument("--out-subdir", type=str, default="resized_videos", help="Output subdirectory.")
+    p.add_argument("--num-frames", type=int, default=NUM_FRAMES)
+    p.add_argument("--height", type=int, default=TARGET_H)
+    p.add_argument("--width", type=int, default=TARGET_W)
+    p.add_argument("--fps", type=int, default=24)
+    p.add_argument("--index", type=int, nargs="+", default=None, metavar="IDX",
+                   help="Process only these video indices (e.g. --index 4 7 12). Assumes vid_%06d naming.")
+    p.add_argument("--include-list", type=Path, default=None,
+                   help="Text file of clip filenames (one per line) to process; everything else in "
+                        "--video-subdir is ignored. Pairs with filter_clips.py's needs_resize.txt "
+                        "so only off-spec clips are re-encoded.")
+    p.add_argument("--limit", type=int, default=None, help="Process only first N videos (smoke test).")
+    p.add_argument("--rank", type=int, default=0, help="Shard index for CPU-parallel runs (0-indexed).")
+    p.add_argument("--world-size", type=int, default=1, help="Total number of parallel processes.")
+    p.add_argument("--min-frames", type=int, default=None,
+                   help="Skip clips with fewer than this many frames (default: --num-frames). "
+                        "Set 0 to keep short clips (output T then varies per clip).")
+    p.add_argument("--force", action="store_true", help="Re-write even if output already exists.")
+    return p.parse_args()
+
+
+def resize_video(path: Path, num_frames: int, height: int, width: int) -> np.ndarray:
+    """Crop+resize to the training geometry. Returns uint8 frames [T, H, W, C].
+
+    Reads sequentially and stops at num_frames or end-of-file, so clips shorter than
+    num_frames yield what they have rather than raising (real-world shards are ragged).
+    """
+    reader = imageio.get_reader(str(path))
+    frames = []
+    for i, frame in enumerate(reader):
+        if i >= num_frames:
+            break
+        frames.append(np.asarray(frame))
+    reader.close()
+    if not frames:
+        raise ValueError(f"no frames decoded from {path}")
+    clip = torch.from_numpy(np.stack(frames)).permute(0, 3, 1, 2).float() / 255.0
+    clip = center_crop_th_tw(clip, height, width, top_crop=False)
+    clip = resize(clip, (height, width), interpolation_mode="bilinear")
+    return (clip.clamp(0, 1) * 255).byte().permute(0, 2, 3, 1).numpy()
+
+
+def main() -> None:
+    args = parse_args()
+    videos_dir = args.data_dir / args.video_subdir
+    out_dir = args.data_dir / args.out_subdir
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    videos = sorted(videos_dir.glob("*.mp4"))
+    if args.include_list is not None:
+        wanted = {ln.strip() for ln in args.include_list.read_text().splitlines() if ln.strip()}
+        videos = [v for v in videos if v.name in wanted]
+    if args.index is not None:
+        wanted = {f"vid_{i:06d}.mp4" for i in args.index}
+        videos = [v for v in videos if v.name in wanted]
+    if args.limit is not None:
+        videos = videos[:args.limit]
+    if args.world_size > 1:
+        videos = videos[args.rank::args.world_size]
+    if not videos:
+        print(f"[resize] no videos found in {videos_dir}", flush=True)
+        return
+
+    min_frames = args.num_frames if args.min_frames is None else args.min_frames
+    print(f"[resize] {len(videos)} videos → {out_dir} ({args.height}x{args.width}, no VAE)"
+          f"{f' [shard {args.rank}/{args.world_size}]' if args.world_size > 1 else ''}", flush=True)
+    n_ok = n_short = n_err = 0
+    for k, vpath in enumerate(videos, 1):
+        out_path = out_dir / vpath.name
+        if out_path.exists() and not args.force:
+            continue
+        try:
+            frames = resize_video(vpath, args.num_frames, args.height, args.width)
+        except Exception as e:  # noqa: BLE001
+            n_err += 1
+            print(f"[resize] [{k}/{len(videos)}] {vpath.name}: DECODE FAILED ({e}), skipping", flush=True)
+            continue
+        if frames.shape[0] < min_frames:
+            n_short += 1
+            print(f"[resize] [{k}/{len(videos)}] {vpath.name}: only {frames.shape[0]} frames "
+                  f"(< {min_frames}), skipping", flush=True)
+            continue
+        # Dot-prefixed so a leftover temp is NOT picked up by downstream `*.mp4` globs
+        # (a stale "<name>.tmp.mp4" once got fed to the tracker and killed the worker).
+        tmp = out_path.with_name(f".{out_path.stem}.tmp.mp4")
+        imageio.mimsave(str(tmp), frames, fps=args.fps, macro_block_size=1)
+        tmp.replace(out_path)
+        n_ok += 1
+        if k % 50 == 0 or k == len(videos):
+            print(f"[resize] [{k}/{len(videos)}] ok={n_ok} short={n_short} err={n_err}", flush=True)
+
+    print(f"[resize] done. ok={n_ok} short={n_short} err={n_err}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/run.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+# Stage-0 data pipeline: generate videos (Wan2.2-14B T2V) -> extract tracks (CoTracker v3).
+# Runs both steps on the shao_wm allocation via `srun --overlap`.
+#
+# You pick the GPUs (the node is shared); pass them explicitly:
+#   CUDA_VISIBLE_DEVICES=2,3 bash data_pipeline/run.sh            # full run (defaults below)
+#   CUDA_VISIBLE_DEVICES=2,3 bash data_pipeline/run.sh --smoke    # 2 videos
+#   CUDA_VISIBLE_DEVICES=2,3 NUM_VIDEOS=50 OUTPUT_DIR=... bash data_pipeline/run.sh
+#
+# num_gpus is derived from CUDA_VISIBLE_DEVICES. Never run model code on the login node.
+set -euo pipefail
+
+REPO=/mnt/weka/home/hao.zhang/shao/FastVideo
+PY="$REPO/.venv/bin/python"
+PROMPTS="${PROMPTS:-$REPO/examples/dataset/vidprom/prompts/vidprom_filtered_extended.txt}"
+OUTPUT_DIR="${OUTPUT_DIR:-/mnt/weka/home/hao.zhang/shao/data/motion_pipeline/wan22_t2v_720p}"
+JOB_NAME="${JOB_NAME:-shao_wm}"
+NUM_VIDEOS="${NUM_VIDEOS:-50}"
+GRID_SIZE="${GRID_SIZE:-50}"
+
+[[ "${1:-}" == "--smoke" ]] && NUM_VIDEOS=2
+
+: "${CUDA_VISIBLE_DEVICES:?set CUDA_VISIBLE_DEVICES to the GPUs to use, e.g. CUDA_VISIBLE_DEVICES=2,3}"
+IFS=',' read -ra _gpus <<< "$CUDA_VISIBLE_DEVICES"
+NUM_GPUS="${NUM_GPUS:-${#_gpus[@]}}"
+
+JOBID="${JOBID:-$(squeue -u "$USER" -n "$JOB_NAME" -h -o %i | head -1)}"
+[[ -n "$JOBID" ]] || { echo "[run] ERROR: no running '$JOB_NAME' allocation" >&2; exit 1; }
+echo "[run] jobid=$JOBID  CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES  num_gpus=$NUM_GPUS  num_videos=$NUM_VIDEOS"
+
+# /usr/bin/env (absolute: bare `env` may resolve to a broken ~/.local/bin/env) pins the
+# device list inside the step. Not SLURM --export: its comma parsing splits "2,3".
+SRUN=(srun --jobid="$JOBID" --overlap --ntasks=1 /usr/bin/env "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES")
+
+echo "[run] === generate videos ==="
+"${SRUN[@]}" "$PY" "$REPO/data_pipeline/generate_videos.py" \
+  --prompts "$PROMPTS" --output-dir "$OUTPUT_DIR" --num-videos "$NUM_VIDEOS" --num-gpus "$NUM_GPUS"
+
+echo "[run] === extract tracks ==="
+"${SRUN[@]}" "$PY" "$REPO/data_pipeline/extract_tracks.py" \
+  --data-dir "$OUTPUT_DIR" --grid-size "$GRID_SIZE"
+
+echo "[run] done -> $OUTPUT_DIR"

data_pipeline/run_decode_wansyn32k_slurm.sh

@@ -0,0 +1,25 @@
+#!/usr/bin/env bash
+# Wan 2.2 5B VAE decode of the 32k HF latent dataset -> 480x832 mp4s.
+# Sharded across NODES * GPUS * PROCS_PER_GPU workers via SLURM_PROCID.
+set -euo pipefail
+WORK=/mnt/lustre/vlm-s4duan
+: "${PARQUET_DIR:?}"; : "${OUT_DIR:?}"; : "${VAE_DIR:?}"
+NODES=${NODES:-8}; GPUS=${GPUS:-4}; PROCS_PER_GPU=${PROCS_PER_GPU:-2}
+TASKS_PER_NODE=$(( GPUS * PROCS_PER_GPU ))
+CPUS_PER_TASK=$(( 128 / TASKS_PER_NODE )); [ "$CPUS_PER_TASK" -lt 1 ] && CPUS_PER_TASK=1
+NUM_SHARDS=$(( NODES * TASKS_PER_NODE ))
+
+mkdir -p "$OUT_DIR" "$WORK/logs"
+echo "nodes=$NODES gpus/node=$GPUS procs/gpu=$PROCS_PER_GPU -> $NUM_SHARDS workers"
+
+sbatch -N "$NODES" --gres=gpu:$GPUS --ntasks-per-node=$TASKS_PER_NODE --exclusive \
+  --cpus-per-task=$CPUS_PER_TASK --mem=0 -t 12:00:00 -J decode_wansyn32k \
+  --chdir="$WORK/FastVideo" -o "$WORK/logs/decode_wansyn32k_%j.out" -e "$WORK/logs/decode_wansyn32k_%j.out" \
+  --wrap "srun --chdir=$WORK/FastVideo \
+    bash -lc 'source .venv/bin/activate && \
+      export HOME=$WORK HF_HOME=$WORK/.hf TORCH_HOME=$WORK/.torch \
+        MPLCONFIGDIR=$WORK/.mpl TOKENIZERS_PARALLELISM=false && \
+      export CUDA_VISIBLE_DEVICES=\$(( SLURM_LOCALID % $GPUS )) && \
+      python data_pipeline/decode_wansyn32k.py \
+        --parquet-dir $PARQUET_DIR --vae-dir $VAE_DIR --out-dir $OUT_DIR \
+        --num-shards $NUM_SHARDS --shard \$SLURM_PROCID'"

data_pipeline/run_download_hd_slurm.sh

@@ -0,0 +1,21 @@
+#!/usr/bin/env bash
+# Multi-node OpenVidHD download+extract (>=5s clips only), split-shard aware.
+# Parts are sharded across tasks (SLURM_PROCID/NTASKS) so the ~8.8TB download runs
+# in parallel across nodes. CPU/network only — no GPU. Idempotent (resumes).
+#
+# Usage:
+#   NODES=8 TASKS_PER_NODE=3 bash data_pipeline/run_download_hd_slurm.sh
+set -euo pipefail
+WORK=/mnt/lustre/vlm-s4duan
+NODES=${NODES:-8}; TASKS_PER_NODE=${TASKS_PER_NODE:-3}   # concurrent HF downloads per node
+OUT=${OUT:-$WORK/openvid_1m}                             # goal folder
+FILTER=${FILTER:-$WORK/openvid/OpenVidHD_filtered.txt}
+mkdir -p "$OUT/videos" "$OUT/_zips" "$WORK/logs"
+echo "download OpenVidHD -> $OUT/videos  ($NODES nodes x $TASKS_PER_NODE = $((NODES*TASKS_PER_NODE)) parallel downloads)"
+
+sbatch -N "$NODES" --ntasks-per-node=$TASKS_PER_NODE --cpus-per-task=16 --mem=0 --exclusive \
+  -t 48:00:00 -J openvid_dl --chdir="$WORK/FastVideo" \
+  -o "$WORK/logs/openvid_dl_%j.out" -e "$WORK/logs/openvid_dl_%j.out" \
+  --wrap "srun --chdir=$WORK/FastVideo bash -lc 'source .venv/bin/activate && export HF_HOME=$WORK/.hf && \
+    python data_pipeline/openvid_download_hd.py \
+      --videos-dir $OUT/videos --zip-dir $OUT/_zips --only-list $FILTER'"

data_pipeline/run_extract_tracks.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+# Run extract_tracks.py in parallel across 4 GPUs.
+# Usage: bash data_pipeline/run_extract_tracks.sh [extra args]
+
+DATA_DIR=/home/hal-kevin/data/motion-stream-test
+WORLD_SIZE=4
+LOG_FILE=data_pipeline/extract_tracks.log
+
+> $LOG_FILE  # truncate on each run
+
+echo "[track] launching $WORLD_SIZE workers... logging to $LOG_FILE"
+
+for RANK in $(seq 0 $((WORLD_SIZE - 1))); do
+    CUDA_VISIBLE_DEVICES=$RANK python -u data_pipeline/extract_tracks.py \
+        --data-dir $DATA_DIR \
+        --videos-subdir roundtrip_videos \
+        --grid-size 50 \
+        --device cuda \
+        --detect-entries \
+        --sam-conf 0.75 \
+        --sam-iou 0.9 \
+        --sam-imgsz 1024 \
+        --entry-sample-every 2 \
+        --entry-min-area 0.001 \
+        --entry-new-area 0.5 \
+        --rank $RANK --world-size $WORLD_SIZE \
+        "$@" \
+        >> $LOG_FILE 2>&1 &
+done
+
+wait
+echo "[track] all done. log at $LOG_FILE"

data_pipeline/run_openvid_shard.sh

@@ -0,0 +1,405 @@
+#!/bin/bash
+# End-to-end preprocessing of one OpenVid-WanTrack shard, WITHOUT the VAE round-trip:
+#
+#   download shard -> extract -> crop+resize to 720p (CPU, parallel) -> fused tracks (GPU)
+#
+# NOTE (deliberate, per request): this skips Stage 2's VAE encode/decode and tracks the
+# resized source frames directly. The 50-clip A/B (ab_vae_roundtrip.sh) found tracks then
+# differ from round-trip tracks by ~5.5px on shared grid points, plus different entry-object
+# sets (n_objects differed on 23/50 clips). Fine for a throughput measurement or a training
+# A/B; see notes before adopting for a production set.
+#
+# NOTE (geometry): --height/--width define the coordinate frame the tracks live in. They
+# must match what Stage 5 crops/resizes to, or tracks won't align with the latents. 720p
+# here is NOT the current training geometry (480x832) -- set HEIGHT/WIDTH accordingly if
+# these tracks are meant to feed the existing training config.
+#
+# Usage:
+#   bash data_pipeline/run_openvid_shard.sh                    # shard 0, 720p, 4 GPUs
+#   LIMIT=50 bash data_pipeline/run_openvid_shard.sh           # quick smoke run
+#   SHARD=3 HEIGHT=480 WIDTH=832 bash data_pipeline/run_openvid_shard.sh
+#   SKIP_DOWNLOAD=1 bash data_pipeline/run_openvid_shard.sh    # shard already on disk
+set -euo pipefail
+
+REPO_ID=${REPO_ID:-noctuashap/openvid-wantrack-clips}
+SHARD=${SHARD:-0}
+SHARD_NAME=$(printf "clips-%05d.tar" "$SHARD")
+DATA_ROOT=${DATA_ROOT:-$(printf "/home/hal-shared/motionstream/data/openvid-wantrack/shard%03d" "$SHARD")}
+HEIGHT=${HEIGHT:-720}
+WIDTH=${WIDTH:-1280}
+NUM_FRAMES=${NUM_FRAMES:-121}
+FPS=${FPS:-24}
+GPUS=${GPUS:-0,1,2,3}
+CPU_WORKERS=${CPU_WORKERS:-$(( $(nproc) > 16 ? 16 : $(nproc) ))}
+LIMIT=${LIMIT:-}
+AMP=${AMP:-1}          # bf16 CoTracker (validated: ~1.1px delta, ~1.5x faster)
+COMPILE=${COMPILE:-0}  # torch.compile main pass (~10-15%; pays a per-worker warmup)
+VIZ=${VIZ:-0}
+SKIP_DOWNLOAD=${SKIP_DOWNLOAD:-0}
+TRACKS=${TRACKS:-1}              # 0 = skip tracking (parquet-only pass over existing npz)
+PARQUET=${PARQUET:-0}            # 1 = also run Stage 5 (v1_preprocess) -> training parquets
+# Stage 5's dataloader defaults to 1 worker, so video decode blocks the GPU between clips
+# (the same bubble --prefetch removes in tracking). More workers overlap decode with encode,
+# but each holds a decoded 720p clip (~334MB of raw frames), so this also drives host RAM.
+PARQUET_WORKERS=${PARQUET_WORKERS:-2}
+# How many samples the parquet writer buffers before flushing to disk. At 720p each sample's
+# latents are ~2.4x the 480p reference, so the default (256) can OOM a 127GB node. Flushing
+# every samples_per_file keeps the in-RAM buffer small.
+PARQUET_FLUSH=${PARQUET_FLUSH:-64}
+MODEL_PATH=${MODEL_PATH:-/home/hal-kevin/models/trackwan_1.3b_i2v_control_init}
+RESUME=${RESUME:-0}              # 1 = continue an interrupted run: skip finished phases AND
+                                 #     already-tracked clips (implies FORCE_TRACKS=0)
+FORCE_TRACKS=${FORCE_TRACKS:-$([[ "$RESUME" == "1" ]] && echo 0 || echo 1)}
+DRY_RUN=${DRY_RUN:-0}            # 1 = print the worker command lines and exit (no work done)
+
+cd "$(dirname "$0")/.."
+IFS=',' read -ra GPU_ARR <<< "$GPUS"
+WORLD_SIZE=${#GPU_ARR[@]}
+DL_DIR="$DATA_ROOT/download"
+RAW_DIR="$DATA_ROOT/raw_videos"
+LOG_DIR="$DATA_ROOT/logs"
+mkdir -p "$DATA_ROOT" "$LOG_DIR"
+
+# Two concurrent runs would interleave resize writes with the tracker's video glob, so the
+# tracker would see a partially-populated dir and silently process a subset. Refuse to overlap.
+LOCK="$DATA_ROOT/.run.lock"
+if ! mkdir "$LOCK" 2>/dev/null; then
+    echo "[openvid] ERROR: another run holds $LOCK (remove it if stale)" >&2; exit 1
+fi
+trap 'rmdir "$LOCK" 2>/dev/null || true' EXIT
+
+LIMIT_ARGS=(); [[ -n "$LIMIT" ]] && LIMIT_ARGS=(--limit "$LIMIT")
+secs() { date +%s; }
+hms() { awk -v s="$1" 'BEGIN{printf "%dm%02ds", s/60, s%60}'; }
+
+# --- progress tracking: survives SIGKILL (cluster reapers), enables RESUME=1 --------
+PROGRESS="$DATA_ROOT/progress.json"
+prog_set() {  # prog_set <phase> <json-object>
+    python - "$PROGRESS" "$1" "$2" "$SHARD_NAME" "${HEIGHT}x${WIDTH}@${NUM_FRAMES}f" <<'PY'
+import datetime, json, sys
+from pathlib import Path
+p, phase, payload, shard, target = Path(sys.argv[1]), sys.argv[2], json.loads(sys.argv[3]), sys.argv[4], sys.argv[5]
+d = json.loads(p.read_text()) if p.exists() else {}
+d.update(shard=shard, target=target)   # informational: the most recent run's target
+now = datetime.datetime.now().isoformat(timespec="seconds")
+# `target` is recorded PER PHASE: a later run at a different geometry must not be able to
+# reuse filter/track outputs produced for the old one.
+d.setdefault("phases", {})[phase] = {**payload, "target": target, "ts": now}
+d["updated"] = now
+tmp = p.with_suffix(".json.tmp")            # atomic: a kill mid-write must not corrupt state
+tmp.write_text(json.dumps(d, indent=2))
+tmp.replace(p)
+PY
+}
+prog_done() {  # prog_done <phase> -> 0 if that phase completed for this target
+    python - "$PROGRESS" "$1" "${HEIGHT}x${WIDTH}@${NUM_FRAMES}f" <<'PY'
+import json, sys
+from pathlib import Path
+p = Path(sys.argv[1])
+if not p.exists():
+    sys.exit(1)
+ph = json.loads(p.read_text()).get("phases", {}).get(sys.argv[2], {})
+# a different target geometry invalidates that phase's outputs
+sys.exit(0 if (ph.get("done") and ph.get("target") == sys.argv[3]) else 1)
+PY
+}
+[[ "$RESUME" == "1" ]] && echo "[openvid] RESUME=1 -- finished phases and already-tracked clips will be skipped"
+
+echo "[openvid] shard=$SHARD_NAME  target=${HEIGHT}x${WIDTH}  gpus=$GPUS  cpu_workers=$CPU_WORKERS"
+echo "[openvid] data root: $DATA_ROOT"
+
+# --- 1. download -------------------------------------------------------------------
+if [[ "$SKIP_DOWNLOAD" != "1" && ! -f "$DL_DIR/$SHARD_NAME" ]]; then
+    echo "[openvid] downloading $SHARD_NAME (~3.3 GB) ..."
+    t=$(secs)
+    if command -v hf >/dev/null 2>&1; then
+        hf download "$REPO_ID" "$SHARD_NAME" --repo-type dataset --local-dir "$DL_DIR"
+    elif command -v huggingface-cli >/dev/null 2>&1; then
+        huggingface-cli download "$REPO_ID" "$SHARD_NAME" --repo-type dataset --local-dir "$DL_DIR"
+    else
+        echo "[openvid] ERROR: neither 'hf' nor 'huggingface-cli' found (pip install -U huggingface_hub)" >&2
+        exit 1
+    fi
+    echo "[openvid] download: $(hms $(( $(secs) - t )))"
+    prog_set download "{\"done\": true, \"secs\": $(( $(secs) - t ))}"
+else
+    echo "[openvid] download: skipped (have $DL_DIR/$SHARD_NAME)"
+    prog_set download '{"done": true, "note": "pre-existing"}'
+fi
+
+# --- 2. extract --------------------------------------------------------------------
+if [[ -z "$(ls -A "$RAW_DIR" 2>/dev/null)" ]]; then
+    echo "[openvid] extracting ..."
+    t=$(secs)
+    mkdir -p "$RAW_DIR"
+    tar -xf "$DL_DIR/$SHARD_NAME" -C "$RAW_DIR"
+    # flatten any nested layout so *.mp4 all sit directly in RAW_DIR
+    find "$RAW_DIR" -mindepth 2 -name '*.mp4' -exec mv -t "$RAW_DIR" {} + 2>/dev/null || true
+    find "$RAW_DIR" -mindepth 1 -type d -empty -delete 2>/dev/null || true
+    echo "[openvid] extract: $(hms $(( $(secs) - t )))"
+else
+    echo "[openvid] extract: skipped (raw_videos/ non-empty)"
+fi
+N_RAW=$(ls "$RAW_DIR"/*.mp4 2>/dev/null | wc -l || true)
+echo "[openvid] raw clips: $N_RAW"
+[[ "$N_RAW" -gt 0 ]] || { echo "[openvid] ERROR: no mp4s extracted" >&2; exit 1; }
+prog_set extract "{\"done\": true, \"raw_clips\": $N_RAW}"
+
+# --- 3. crop+resize to target geometry (CPU, parallel) ------------------------------
+# If the clips are already at the target geometry, resizing is a pure lossy re-encode
+# (measured 40 dB / 1.9-per-255 on this shard -- the same order as the VAE round-trip's
+# distortion, for zero benefit) plus a duplicate copy on disk. Track the raw clips instead.
+VID_SUBDIR=videos
+if [[ "$RESUME" == "1" ]] && prog_done filter && [[ -n "$(ls -A "$DATA_ROOT/videos" 2>/dev/null)" ]]; then
+    T_RESIZE=0
+    N_VID=$(ls "$DATA_ROOT"/videos/*.mp4 2>/dev/null | wc -l || true)
+    echo "[openvid] filter: skipped (progress.json says done; $N_VID clips staged)"
+elif [[ "${FORCE_RESIZE:-0}" != "1" ]]; then
+    # Scan every clip's metadata (no decode) and symlink through the conforming ones.
+    # Clips at other resolutions/lengths are skipped and listed in skipped_clips.json.
+    t=$(secs)
+    python -u data_pipeline/filter_clips.py \
+        --src-dir "$RAW_DIR" --out-dir "$DATA_ROOT/videos" \
+        --height "$HEIGHT" --width "$WIDTH" --num-frames "$NUM_FRAMES" --clean \
+        2>&1 | tee -a "$LOG_DIR/filter.log"
+    # Rescue pass: clips that are readable but off-spec get re-encoded to the target
+    # (only these -- the conforming majority stays symlinked, never re-encoded).
+    NEEDS="$DATA_ROOT/needs_resize.txt"
+    N_RESCUE=$(wc -l < "$NEEDS" 2>/dev/null || echo 0)
+    if [[ "$N_RESCUE" -gt 0 ]]; then
+        echo "[openvid] rescuing $N_RESCUE off-spec clip(s) by resize -> ${HEIGHT}x${WIDTH} ..."
+        pids=()
+        for i in $(seq 0 $((CPU_WORKERS - 1))); do
+            python -u data_pipeline/resize_videos.py \
+                --data-dir "$DATA_ROOT" \
+                --video-subdir raw_videos \
+                --out-subdir videos \
+                --include-list "$NEEDS" \
+                --height "$HEIGHT" --width "$WIDTH" \
+                --num-frames "$NUM_FRAMES" --fps "$FPS" \
+                --rank "$i" --world-size "$CPU_WORKERS" \
+                >> "$LOG_DIR/resize.log" 2>&1 &
+            pids+=($!)
+        done
+        rfail=0; for p in "${pids[@]}"; do wait "$p" || rfail=$((rfail + 1)); done
+        [[ $rfail -gt 0 ]] && echo "[openvid] WARNING: $rfail rescue worker(s) failed -- see $LOG_DIR/resize.log"
+    fi
+    T_RESIZE=$(( $(secs) - t ))
+    N_VID=$(ls "$DATA_ROOT"/videos/*.mp4 2>/dev/null | wc -l || true)
+    N_SKIP=$(( N_RAW - N_VID ))
+    echo "[openvid] filter: $(hms $T_RESIZE)   conforming: $N_VID / $N_RAW   skipped: $N_SKIP"
+    # only record success if something was actually staged
+    [[ "$N_VID" -gt 0 ]] && prog_set filter "{\"done\": true, \"kept\": $N_VID, \"skipped\": $N_SKIP, \"secs\": $T_RESIZE}"
+    if [[ "$N_VID" -eq 0 ]]; then
+        echo "[openvid] ERROR: no clips match ${HEIGHT}x${WIDTH}@${NUM_FRAMES}f." >&2
+        echo "[openvid]        Run with FORCE_RESIZE=1 to re-encode them to the target instead." >&2
+        exit 1
+    fi
+else
+echo "[openvid] resizing $NATIVE -> ${HEIGHT}x${WIDTH} across $CPU_WORKERS CPU workers ..."
+t=$(secs)
+pids=()
+for i in $(seq 0 $((CPU_WORKERS - 1))); do
+    python -u data_pipeline/resize_videos.py \
+        --data-dir "$DATA_ROOT" \
+        --video-subdir raw_videos \
+        --out-subdir videos \
+        --height "$HEIGHT" --width "$WIDTH" \
+        --num-frames "$NUM_FRAMES" --fps "$FPS" \
+        --rank "$i" --world-size "$CPU_WORKERS" \
+        "${LIMIT_ARGS[@]}" \
+        >> "$LOG_DIR/resize.log" 2>&1 &
+    pids+=($!)
+done
+rfail=0; for p in "${pids[@]}"; do wait "$p" || rfail=$((rfail + 1)); done
+[[ $rfail -gt 0 ]] && echo "[openvid] WARNING: $rfail resize worker(s) failed -- see $LOG_DIR/resize.log"
+T_RESIZE=$(( $(secs) - t ))
+find "$DATA_ROOT/$VID_SUBDIR" -name '.*.tmp.mp4' -delete 2>/dev/null || true   # drop any stale temps
+N_VID=$(ls "$DATA_ROOT"/$VID_SUBDIR/*.mp4 2>/dev/null | wc -l || true)
+echo "[openvid] resize: $(hms $T_RESIZE)   usable clips: $N_VID / $N_RAW"
+fi
+[[ "$N_VID" -gt 0 ]] || { echo "[openvid] ERROR: no clips available to track" >&2; exit 1; }
+
+# --- 4. manifest (so tracks get points_path patched + Stage 5 has an entry point) ---
+python - "$DATA_ROOT" "$FPS" "$NUM_FRAMES" "$HEIGHT" "$WIDTH" "$VID_SUBDIR" <<'PY'
+import json, sys
+from pathlib import Path
+root, fps, nf, h, w = Path(sys.argv[1]), float(sys.argv[2]), int(sys.argv[3]), int(sys.argv[4]), int(sys.argv[5])
+tracks_dir = root / "tracks"
+items = []
+for i, p in enumerate(sorted((root / sys.argv[6]).glob("*.mp4"))):
+    it = {"idx": i, "path": p.name, "cap": [""], "fps": fps, "num_frames": nf,
+          "duration": nf / fps, "resolution": {"width": w, "height": h}}
+    # This manifest gets rewritten every run, so re-attach points_path here whenever the
+    # npz exists. In PHASE=parquet, tracking is skipped and never patches it back, so Stage 5
+    # would otherwise see no track sidecar (PreprocessPipeline_I2V_Track then errors).
+    npz = tracks_dir / f"{p.stem}.npz"
+    if npz.exists():
+        it["points_path"] = str(npz.resolve())
+    items.append(it)
+(root / "videos2caption.json").write_text(json.dumps(items, indent=2))
+n_pts = sum(1 for it in items if "points_path" in it)
+print(f"[openvid] manifest: {len(items)} entries ({n_pts} with points_path) -> {root/'videos2caption.json'}")
+PY
+
+# Real captions from OpenVid-1M (joins 1:1 on clip filename). Without this every clip
+# carries an empty prompt and Stage 5 bakes identical null T5 embeddings into the parquets.
+if [[ "${CAPTIONS:-1}" == "1" ]]; then
+    python -u data_pipeline/add_captions.py \
+        --manifest "$DATA_ROOT/videos2caption.json" \
+        --min-coverage "${MIN_CAPTION_COVERAGE:-0.9}" \
+        2>&1 | tee -a "$LOG_DIR/captions.log" | tail -3
+    cap_rc=${PIPESTATUS[0]}
+    if [[ "$cap_rc" -ne 0 ]]; then
+        echo "[openvid] ERROR: caption join failed -- see $LOG_DIR/captions.log" >&2
+        echo "[openvid]        set CAPTIONS=0 to proceed with empty captions (tracks are still valid;" >&2
+        echo "[openvid]        parquets built from them would have dead text conditioning)." >&2
+        exit 1
+    fi
+fi
+
+# --- 5. fused tracks (stages 3+4 in one pass, no VAE round-trip) --------------------
+if [[ "$TRACKS" != "1" ]]; then
+    T_TRACK=0; tfail=0
+    N_NPZ=0
+    N_NPZ_ALL=$(ls "$DATA_ROOT"/tracks/*.npz 2>/dev/null | wc -l || true)
+    echo "[openvid] tracks: SKIPPED (TRACKS=0)   existing npz: $N_NPZ_ALL / $N_VID"
+else
+echo "[openvid] extracting tracks across $WORLD_SIZE GPUs ..."
+SPEED=(); [[ "$AMP" == "1" ]] && SPEED+=(--amp); [[ "$COMPILE" == "1" ]] && SPEED+=(--compile)
+VIZ_ARGS=(); [[ "$VIZ" == "1" ]] && VIZ_ARGS=(--viz --viz-dir "$DATA_ROOT/viz")
+if [[ "$COMPILE" == "1" ]]; then
+    export TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR:-$HOME/.cache/torchinductor}
+    export TRITON_CACHE_DIR=${TRITON_CACHE_DIR:-$HOME/.cache/triton}
+fi
+FORCE_ARGS=(); [[ "$FORCE_TRACKS" == "1" ]] && FORCE_ARGS=(--force)
+if [[ "$DRY_RUN" == "1" ]]; then
+    echo "[openvid] DRY RUN -- rank 0 track worker would be:"
+    echo "  CUDA_VISIBLE_DEVICES=${GPU_ARR[0]} python -u data_pipeline/extract_tracks.py" \
+         "--data-dir $DATA_ROOT --videos-subdir $VID_SUBDIR --out-subdir tracks" \
+         "--grid-size 50 --device cuda --detect-entries --sam-conf 0.75 --sam-iou 0.9 --sam-imgsz 1024" \
+         "--entry-sample-every 2 --entry-min-area 0.001 --entry-new-area 0.5" \
+         "--segment --vis-override-every 3 ${SPEED[*]} ${VIZ_ARGS[*]} ${FORCE_ARGS[*]} ${LIMIT_ARGS[*]}" \
+         "--rank 0 --world-size $WORLD_SIZE"
+    exit 0
+fi
+t=$(secs)
+pids=()
+for i in "${!GPU_ARR[@]}"; do
+    CUDA_VISIBLE_DEVICES=${GPU_ARR[$i]} python -u data_pipeline/extract_tracks.py \
+        --data-dir "$DATA_ROOT" \
+        --videos-subdir "$VID_SUBDIR" \
+        --out-subdir tracks \
+        --grid-size 50 --device cuda \
+        --detect-entries --sam-conf 0.75 --sam-iou 0.9 --sam-imgsz 1024 \
+        --entry-sample-every 2 --entry-min-area 0.001 --entry-new-area 0.5 \
+        --segment --vis-override-every 3 \
+        "${SPEED[@]}" "${VIZ_ARGS[@]}" "${FORCE_ARGS[@]}" "${LIMIT_ARGS[@]}" \
+        --rank "$i" --world-size "$WORLD_SIZE" \
+        >> "$LOG_DIR/tracks.log" 2>&1 &
+    pids+=($!)
+done
+tfail=0; for p in "${pids[@]}"; do wait "$p" || tfail=$((tfail + 1)); done
+[[ $tfail -gt 0 ]] && echo "[openvid] WARNING: $tfail track worker(s) failed -- see $LOG_DIR/tracks.log"
+T_TRACK=$(( $(secs) - t ))
+# Count only npz written by THIS run -- counting the whole dir would fold in earlier runs
+# and (with FORCE_TRACKS=0) report a rate for work that was skipped.
+N_NPZ=$(find "$DATA_ROOT/tracks" -name '*.npz' -newermt "@$t" 2>/dev/null | wc -l || true)
+N_NPZ_ALL=$(ls "$DATA_ROOT"/tracks/*.npz 2>/dev/null | wc -l || true)
+echo "[openvid] tracks: $(hms $T_TRACK)   npz this run: $N_NPZ   total in dir: $N_NPZ_ALL / $N_VID"
+if [[ "$N_NPZ_ALL" -ge "$N_VID" && "$tfail" -eq 0 && -z "$LIMIT" ]]; then
+    prog_set tracks "{\"done\": true, \"npz\": $N_NPZ_ALL, \"clips\": $N_VID, \"secs\": $T_TRACK}"
+    echo "[openvid] shard COMPLETE"
+else
+    prog_set tracks "{\"done\": false, \"npz\": $N_NPZ_ALL, \"clips\": $N_VID, \"secs\": $T_TRACK}"
+    [[ "$N_NPZ_ALL" -lt "$N_VID" ]] && \
+        echo "[openvid] INCOMPLETE: $(( N_VID - N_NPZ_ALL )) clips remain -- resume with: RESUME=1 SKIP_DOWNLOAD=1 bash $0"
+fi
+fi
+
+# --- 6. Stage 5: parquets (opt-in; the tracks above are already usable without this) --
+T_PARQUET=0
+if [[ "$PARQUET" == "1" ]]; then
+    if [[ "$N_NPZ_ALL" -lt "$N_VID" ]]; then
+        echo "[openvid] SKIPPING parquets: tracks incomplete ($N_NPZ_ALL/$N_VID)" >&2
+    else
+        # num_latent_t = (num_frames - 1)/4 + 1 for WanVAE's 4x temporal compression.
+        NLT=$(( (NUM_FRAMES - 1) / 4 + 1 ))
+        echo "$DATA_ROOT/$VID_SUBDIR,$DATA_ROOT/videos2caption.json" > "$DATA_ROOT/data_merge.txt"
+        # Parquet output location. Default keeps it beside the shard's other data; set
+        # PARQUET_ROOT to collect all shards' parquets under one tree (their own dir),
+        # each in a per-shard subdir so shards stay independent (wipe/verify are per-shard).
+        if [[ -n "${PARQUET_ROOT:-}" ]]; then
+            PQ_OUT="$PARQUET_ROOT/$(printf 'shard%03d' "$SHARD")"
+        else
+            PQ_OUT="$DATA_ROOT/preprocessed_i2v_track"
+        fi
+        # Stage 5 now resumes by clip id (preprocess_pipeline_base.py): a re-run skips clips
+        # already written and appends only the rest, so a reaper kill costs minutes (the
+        # unflushed buffer), not the whole shard. Do NOT wipe -- that would discard progress.
+        echo "[openvid] Stage 5: parquets (${HEIGHT}x${WIDTH}, ${NUM_FRAMES}f, num_latent_t=$NLT, train_fps=$FPS) ..."
+        t=$(secs)
+        # --train_fps MUST equal the source fps: a mismatch makes FrameSamplingStage resample
+        # (duplicating/dropping frames) so latents no longer align with the tracks in the same row.
+        # v1_preprocess asserts num_gpus == 1 (fastvideo/pipelines/preprocess/v1_preprocess.py:27),
+        # so Stage 5 is single-GPU per shard. Scale it by running shards concurrently
+        # (one GPU each) rather than by raising nproc_per_node.
+        CUDA_VISIBLE_DEVICES="${PARQUET_GPU:-${GPU_ARR[0]}}" \
+        torchrun --nproc_per_node=1 -m fastvideo.pipelines.preprocess.v1_preprocess \
+            --model_path "$MODEL_PATH" \
+            --data_merge_path "$DATA_ROOT/data_merge.txt" \
+            --output_dir "$PQ_OUT" \
+            --preprocess_task i2v_track \
+            --num_frames "$NUM_FRAMES" \
+            --num_latent_t "$NLT" \
+            --train_fps "$FPS" \
+            --max_height "$HEIGHT" \
+            --max_width "$WIDTH" \
+            --preprocess_video_batch_size 1 \
+            --dataloader_num_workers "$PARQUET_WORKERS" \
+            --samples_per_file "${PARQUET_SAMPLES:-64}" \
+            --flush_frequency "$PARQUET_FLUSH" \
+            >> "$LOG_DIR/parquet.log" 2>&1
+        prc=$?
+        T_PARQUET=$(( $(secs) - t ))
+        N_PQ=$(find "$PQ_OUT" -name '*.parquet' 2>/dev/null | wc -l || true)
+        # Verify integrity: total rows must equal clip count AND be free of duplicate ids.
+        # A reaper kill mid-run leaves a partial (rows < N_VID) which the next attempt wipes+redoes.
+        read -r N_ROWS N_UNIQ < <(python - "$PQ_OUT" <<'PY'
+import glob, sys
+import pyarrow.parquet as pq
+ids = []
+for f in glob.glob(f"{sys.argv[1]}/**/*.parquet", recursive=True):
+    ids += pq.read_table(f, columns=["id"]).column("id").to_pylist()
+print(len(ids), len(set(ids)))
+PY
+)
+        if [[ $prc -eq 0 && "$N_ROWS" == "$N_VID" && "$N_UNIQ" == "$N_VID" ]]; then
+            echo "[openvid] parquets: $(hms $T_PARQUET)   rows=$N_ROWS unique=$N_UNIQ files=$N_PQ"
+            prog_set parquet "{\"done\": true, \"rows\": $N_ROWS, \"files\": $N_PQ, \"secs\": $T_PARQUET}"
+        else
+            echo "[openvid] Stage 5 INCOMPLETE (rc=$prc, rows=$N_ROWS unique=$N_UNIQ, want $N_VID) -- see $LOG_DIR/parquet.log" >&2
+            prog_set parquet "{\"done\": false, \"rows\": $N_ROWS, \"unique\": $N_UNIQ, \"secs\": $T_PARQUET}"
+        fi
+    fi
+fi
+
+# --- 7. summary --------------------------------------------------------------------
+RESULTS="$DATA_ROOT/shard_results.txt"
+{
+    echo "=== $(date -u '+%Y-%m-%d %H:%M:%S') UTC  shard=$SHARD_NAME  ${HEIGHT}x${WIDTH}  gpus=$GPUS  cpu=$CPU_WORKERS  amp=$AMP compile=$COMPILE viz=$VIZ  (no VAE round-trip) ==="
+    awk -v r="$T_RESIZE" -v tk="$T_TRACK" -v n="$N_NPZ" -v w="$WORLD_SIZE" -v c="$CPU_WORKERS" 'BEGIN {
+        if (n == 0) { print "no npz produced"; exit }
+        printf "resize (CPU): %6ds  %5.2fs/clip/worker\n", r, r*c/n
+        printf "tracks (GPU): %6ds  %5.2fs/clip/worker  %5.1f clips/min\n", tk, tk*w/n, 60*n/tk
+        printf "total:        %6ds for %d clips\n", r+tk, n
+        printf "  -> 259k clips on %d GPUs (tracks only): %.1f h\n", w, 259000*(tk*w/n)/w/3600
+        printf "  -> full shard (1000 clips) at this rate: %.1f min\n", (tk/n)*1000/60
+    }'
+    # `|| true`: a false [[ ]] would make this block (and so the piped tee) exit non-zero
+    # under `set -e -o pipefail`, failing the whole shard after the work already succeeded.
+    { [[ "$VIZ" == "1" ]] && echo "  NOTE: viz=1 -- overlay rendering dominates the GPU phase (~4x); projections are pessimistic, not a throughput measurement."; } || true
+    { [[ -n "$LIMIT" ]] && echo "  NOTE: limit=$LIMIT -- per-worker startup is a large share at this size; totals understate steady-state throughput."; } || true
+} | tee -a "$RESULTS"
+echo "[openvid] appended to $RESULTS"
+echo "[openvid] outputs: $DATA_ROOT/{videos,tracks}$([[ "$VIZ" == "1" ]] && echo ",viz")   logs: $LOG_DIR"

data_pipeline/run_openvid_shards.sh

@@ -0,0 +1,215 @@
+#!/bin/bash
+# Drive run_openvid_shard.sh over a range of shards, with resume and per-shard accounting.
+#
+# Every shard runs with RESUME=1, so re-invoking after an interruption (cluster reaper,
+# node loss, Ctrl-C) picks up exactly where it stopped: finished shards are skipped via
+# their progress.json, and a half-finished shard resumes at the first untracked clip.
+#
+# Usage:
+#   SHARDS=0-130 bash data_pipeline/run_openvid_shards.sh
+#   SHARDS=0-9,20,30-35 bash data_pipeline/run_openvid_shards.sh
+#   SHARDS=0-130 CLEANUP=1 bash data_pipeline/run_openvid_shards.sh   # drop tar+raw after each
+#   SHARDS=0-3 LIMIT=20 bash data_pipeline/run_openvid_shards.sh      # smoke run
+#
+# Passes through the per-shard knobs (GPUS, HEIGHT/WIDTH, AMP, COMPILE, VIZ, CPU_WORKERS,
+# LIMIT, DATA_ROOT_BASE); see run_openvid_shard.sh for their meanings.
+set -uo pipefail          # NOT -e: one bad shard must not kill a 130-shard run
+
+SHARDS=${SHARDS:-0}
+CLEANUP=${CLEANUP:-0}       # 1 = delete download/ and raw_videos/ once a shard completes
+STOP_ON_FAIL=${STOP_ON_FAIL:-0}
+# PARALLEL=1 runs one shard pipeline per GPU concurrently instead of one shard at a time
+# across all GPUs. Stage 5 (v1_preprocess) asserts a single GPU, so sequential mode leaves
+# 3 of 4 GPUs idle for the whole parquet phase; shard-level parallelism keeps them all busy.
+PARALLEL=${PARALLEL:-0}
+GPUS=${GPUS:-0,1,2,3}
+# PHASE picks what this invocation produces:
+#   tracks  -- tracks only, one shard at a time across all GPUs (fastest per-shard: ~13 min)
+#   parquet -- Stage 5 only, over shards that already have tracks; forces PARALLEL=1 because
+#              v1_preprocess is single-GPU, so concurrency has to come from running shards
+#   both    -- everything per shard (honours PARALLEL as set)
+PHASE=${PHASE:-both}
+case "$PHASE" in
+    tracks)  export TRACKS=1 PARQUET=0 ;;
+    parquet) export TRACKS=0 PARQUET=1; PARALLEL=1 ;;
+    both)    export TRACKS=1 ;;
+    *) echo "[shards] ERROR: PHASE must be tracks|parquet|both (got '$PHASE')" >&2; exit 1 ;;
+esac
+DATA_ROOT_BASE=${DATA_ROOT_BASE:-/home/hal-shared/motionstream/data/openvid-wantrack/shard}
+shard_root() { printf "%s%03d" "$DATA_ROOT_BASE" "$1"; }   # zero-padded: shard000 .. shard259
+
+cd "$(dirname "$0")/.."
+SHARD_SCRIPT=data_pipeline/run_openvid_shard.sh
+
+# --- expand "0-9,20,30-35" into a list ---------------------------------------------
+expand() {
+    local spec=$1 tok lo hi out=()
+    IFS=',' read -ra toks <<< "$spec"
+    for tok in "${toks[@]}"; do
+        if [[ "$tok" =~ ^([0-9]+)-([0-9]+)$ ]]; then
+            lo=${BASH_REMATCH[1]}; hi=${BASH_REMATCH[2]}
+            (( lo <= hi )) || { echo "[shards] ERROR: bad range '$tok'" >&2; exit 1; }
+            for ((i = lo; i <= hi; i++)); do out+=("$i"); done
+        elif [[ "$tok" =~ ^[0-9]+$ ]]; then
+            out+=("$tok")
+        else
+            echo "[shards] ERROR: cannot parse '$tok' (want N or A-B, comma-separated)" >&2; exit 1
+        fi
+    done
+    printf '%s\n' "${out[@]}"
+}
+mapfile -t SHARD_LIST < <(expand "$SHARDS")
+N_SHARDS=${#SHARD_LIST[@]}
+
+# 260 shards exist: clips-00000.tar .. clips-00259.tar
+for s in "${SHARD_LIST[@]}"; do
+    (( s <= 259 )) || { echo "[shards] ERROR: shard $s out of range (max 259)" >&2; exit 1; }
+done
+
+# A shard counts as complete only once every phase this run is producing has finished --
+# with PARQUET=1 that includes Stage 5, so a shard whose tracks landed but whose parquets
+# failed is retried rather than skipped.
+case "$PHASE" in
+    tracks)  REQUIRED_PHASES="tracks" ;;
+    parquet) REQUIRED_PHASES="parquet" ;;          # tracks were the previous pass's job
+    both)    REQUIRED_PHASES="tracks"; [[ "${PARQUET:-0}" == "1" ]] && REQUIRED_PHASES="tracks parquet" ;;
+esac
+shard_complete() {  # shard_complete <n> -> 0 if all required phases are done
+    python - "$(shard_root "$1")/progress.json" $REQUIRED_PHASES <<'PY'
+import json, sys
+from pathlib import Path
+p = Path(sys.argv[1])
+if not p.exists():
+    sys.exit(1)
+phases = json.loads(p.read_text()).get("phases", {})
+sys.exit(0 if all(phases.get(ph, {}).get("done") for ph in sys.argv[2:]) else 1)
+PY
+}
+
+# Roll-up across all shards, so overall progress is one file rather than 131.
+OVERALL=${OVERALL:-$(dirname "$DATA_ROOT_BASE")/progress.json}
+overall_set() {  # overall_set <done> <skipped> <failed> <total> <clips> <elapsed> <failed-list>
+    python - "$OVERALL" "$@" "$DATA_ROOT_BASE" "$SHARDS" <<'PY'
+import datetime, json, os, sys
+from pathlib import Path
+p = Path(sys.argv[1])
+done, skipped, failed, total, clips, elapsed = (int(x) for x in sys.argv[2:8])
+failed_list, base, spec = sys.argv[8], sys.argv[9], sys.argv[10]
+processed = done + skipped
+d = {
+    "spec": spec, "data_root_base": base,
+    "shards_total": total, "complete": processed, "processed_this_session": done,
+    "skipped_already_done": skipped, "failed": failed,
+    "failed_shards": [int(x) for x in failed_list.split() if x],
+    "npz_this_session": clips,
+    "elapsed_min": round(elapsed / 60, 1),
+    "avg_min_per_shard": round(elapsed / done / 60, 1) if done else None,
+    "eta_hours": round((elapsed / done) * (total - processed - failed) / 3600, 1) if done else None,
+    "updated": datetime.datetime.now().isoformat(timespec="seconds"),
+}
+tmp = p.with_suffix(f".json.tmp{os.getpid()}"); tmp.write_text(json.dumps(d, indent=2)); tmp.replace(p)
+PY
+}
+
+echo "[shards] $N_SHARDS shard(s): ${SHARD_LIST[0]}..${SHARD_LIST[-1]}   phase=$PHASE  parallel=$PARALLEL  cleanup=$CLEANUP"
+echo "[shards] overall progress: $OVERALL   per-shard: ${DATA_ROOT_BASE}<N>/progress.json"
+
+# ---- parallel mode: one shard pipeline per GPU -------------------------------------
+if [[ "$PARALLEL" == "1" ]]; then
+    IFS=',' read -ra GPU_ARR <<< "$GPUS"
+    NG=${#GPU_ARR[@]}
+    echo "[shards] parallel: $NG concurrent pipelines, one GPU each (${GPUS})"
+    T0=$(date +%s)
+    wpids=()
+    for gi in "${!GPU_ARR[@]}"; do
+        (
+            gpu=${GPU_ARR[$gi]}
+            mine=()
+            for ((j = gi; j < N_SHARDS; j += NG)); do mine+=("${SHARD_LIST[$j]}"); done
+            echo "[gpu$gpu] ${#mine[@]} shard(s): ${mine[*]:0:6}$([[ ${#mine[@]} -gt 6 ]] && echo ' ...')"
+            for s in "${mine[@]}"; do
+                root="$(shard_root "$s")"
+                if shard_complete "$s"; then echo "[gpu$gpu] shard $s already complete"; continue; fi
+                t=$(date +%s)
+                if SHARD="$s" RESUME=1 DATA_ROOT="$root" GPUS="$gpu" PARQUET_GPU="$gpu" \
+                       bash "$SHARD_SCRIPT" >> "$root.log" 2>&1; then
+                    echo "[gpu$gpu] shard $s OK in $(( ($(date +%s) - t) / 60 ))m"
+                    [[ "$CLEANUP" == "1" ]] && shard_complete "$s" && rm -rf "$root/download" "$root/raw_videos"
+                else
+                    echo "[gpu$gpu] shard $s FAILED -- see $root.log" >&2
+                fi
+            done
+        ) &
+        wpids+=($!)
+    done
+    for p in "${wpids[@]}"; do wait "$p"; done
+
+    # tally from the per-shard progress files (authoritative, survives restarts)
+    done_n=0; fail_n=0; failed_list=()
+    for s in "${SHARD_LIST[@]}"; do
+        if shard_complete "$s"; then done_n=$((done_n + 1)); else fail_n=$((fail_n + 1)); failed_list+=("$s"); fi
+    done
+    npz_n=$(find "$(dirname "$DATA_ROOT_BASE")" -name '*.npz' -path '*/tracks/*' 2>/dev/null | wc -l || echo 0)
+    overall_set "$done_n" 0 "$fail_n" "$N_SHARDS" "$npz_n" "$(( $(date +%s) - T0 ))" "${failed_list[*]:-}"
+    echo "=============================================================="
+    echo "[shards] done in $(( ($(date +%s) - T0) / 60 ))m: $done_n complete, $fail_n incomplete"
+    echo "[shards] per-shard console logs: ${DATA_ROOT_BASE}<N>.log"
+    (( fail_n > 0 )) && { echo "[shards] incomplete: ${failed_list[*]}"; echo "[shards] re-run to retry"; exit 1; }
+    exit 0
+fi
+T0=$(date +%s)
+n_done=0 n_skip=0 n_fail=0 clips_total=0
+FAILED=()
+
+for s in "${SHARD_LIST[@]}"; do
+    root="$(shard_root "$s")"
+    if shard_complete "$s"; then
+        n_skip=$((n_skip + 1))
+        echo "[shards] shard $s: already complete, skipping"
+        continue
+    fi
+
+    echo "=============================================================="
+    echo "[shards] shard $s  ($((n_done + n_skip + n_fail + 1))/$N_SHARDS)  elapsed $(( ($(date +%s) - T0) / 60 ))m"
+    t=$(date +%s)
+    rc=0
+    SHARD="$s" RESUME=1 DATA_ROOT="$root" bash "$SHARD_SCRIPT" || rc=$?
+    # count npz regardless of outcome: a shard can produce tracks and still fail a later phase
+    c=$(ls "$root"/tracks/*.npz 2>/dev/null | wc -l || echo 0)
+    clips_total=$((clips_total + c))
+    if [[ $rc -eq 0 ]]; then
+        n_done=$((n_done + 1))
+        echo "[shards] shard $s OK in $(( ($(date +%s) - t) / 60 ))m  ($c npz)"
+        if [[ "$CLEANUP" == "1" ]] && shard_complete "$s"; then
+            # only after progress.json confirms completion -- never delete inputs for a
+            # shard that would need re-processing
+            rm -rf "$root/download" "$root/raw_videos"
+            echo "[shards] shard $s: removed download/ and raw_videos/ (tracks kept)"
+        fi
+    else
+        n_fail=$((n_fail + 1)); FAILED+=("$s")
+        echo "[shards] shard $s FAILED -- see $root/logs/" >&2
+        [[ "$STOP_ON_FAIL" == "1" ]] && { echo "[shards] stopping (STOP_ON_FAIL=1)" >&2; break; }
+    fi
+
+    overall_set "$n_done" "$n_skip" "$n_fail" "$N_SHARDS" "$clips_total" \
+                "$(( $(date +%s) - T0 ))" "${FAILED[*]:-}"
+
+    # rolling ETA from shards actually processed this session
+    if (( n_done > 0 )); then
+        avg=$(( ($(date +%s) - T0) / n_done ))
+        left=$(( N_SHARDS - n_done - n_skip - n_fail ))
+        echo "[shards] avg $(( avg / 60 ))m/shard, $left left, ETA $(( avg * left / 3600 ))h"
+    fi
+done
+
+overall_set "$n_done" "$n_skip" "$n_fail" "$N_SHARDS" "$clips_total" \
+            "$(( $(date +%s) - T0 ))" "${FAILED[*]:-}"
+echo "=============================================================="
+echo "[shards] done in $(( ($(date +%s) - T0) / 60 ))m: $n_done processed, $n_skip skipped, $n_fail failed"
+echo "[shards] npz produced this session: $clips_total"
+if (( n_fail > 0 )); then
+    echo "[shards] failed shards: ${FAILED[*]}"
+    echo "[shards] re-run the same command to retry them (completed shards are skipped)"
+    exit 1
+fi

data_pipeline/run_preprocess_track_slurm.sh

@@ -0,0 +1,88 @@
+#!/usr/bin/env bash
+# Multi-node data-parallel i2v_track PREPROCESS over OpenVid-1M.
+# v1_preprocess.py asserts WORLD_SIZE==1 (one GPU per process), so scale-out is
+# N nodes x 4 GPU x 1 proc/GPU = 4N independent single-GPU processes, each over a
+# SHARD of the clips, writing parquet into <combined>/shard_<idx>/. The trainer
+# reads <combined> as ONE dataset (get_parquet_files_and_length os.walks recursively).
+#
+# Usage:
+#   NODES=16 bash data_pipeline/run_preprocess_track_slurm.sh
+set -euo pipefail
+WORK=/mnt/lustre/vlm-s4duan
+DATA_DIR=${DATA_DIR:-$WORK/openvid_1m}
+MODEL=${MODEL:-$WORK/models/trackwan_1.3b_i2v_d64_nobias_init}
+CLIPS_DIR=${CLIPS_DIR:-$DATA_DIR/clips}
+MANIFEST=${MANIFEST:-$DATA_DIR/videos2caption.json}
+COMBINED=${COMBINED:-$DATA_DIR/combined_parquet_dataset}   # trainer points here
+SHARDS_DIR=${SHARDS_DIR:-$DATA_DIR/preprocess_shards}
+NODES=${NODES:-14}; GPUS=${GPUS:-4}
+PARTITION=${PARTITION:-all}
+CONC=${CONC:-$(( NODES * GPUS ))}          # max concurrent shards (= total GPUs)
+# MANY small shards (not one-per-GPU) so a node glitch only loses its in-flight
+# ~CONC/NODES shards, and .done markers make a resubmit skip finished work.
+NUM_SHARDS=${NUM_SHARDS:-1024}             # ~253 clips/shard (fine-grained .done resume)
+# Training-target geometry (must match trainer):
+MAX_H=${MAX_H:-480}; MAX_W=${MAX_W:-832}; NUM_FRAMES=${NUM_FRAMES:-121}
+TRAIN_FPS=${TRAIN_FPS:-24}; NUM_LATENT_T=${NUM_LATENT_T:-31}
+# batch_size MUST be 1: the T5 tokenizer_kwargs pad-free config (configs/models/
+# encoders/t5.py) can't batch variable-length captions with return_tensors="pt".
+BATCH=${BATCH:-1}
+VAE_PREC=${VAE_PREC:-fp32}   # bf16 gives NO speedup here (decode-bound, not VAE-compute-bound)
+NW=${NW:-8}                  # decode-prefetch workers. SAFE now: patched VideoCaptionMergedDataset.__iter__
+                             # to shard by get_worker_info() (else num_workers>0 duplicates). Validated:
+                             # 16 clips -> 16 rows, latents byte-identical to nw=0. Overlaps decode w/ GPU encode.
+TIME=${TIME:-24:00:00}
+CPUS_PER_TASK=$(( 128 / GPUS )); [ "$CPUS_PER_TASK" -lt 1 ] && CPUS_PER_TASK=1
+mkdir -p "$WORK/logs" "$COMBINED"
+
+# 1) Build/refresh shard manifests (idempotent; cheap).
+source "$WORK/FastVideo/.venv/bin/activate"
+python "$WORK/FastVideo/data_pipeline/split_manifest_shards.py" \
+  --manifest "$MANIFEST" --clips-dir "$CLIPS_DIR" \
+  --out-dir "$SHARDS_DIR" --num-shards "$NUM_SHARDS"
+
+echo "array of $NUM_SHARDS shards, up to $CONC concurrent (~$((NODES)) nodes x $GPUS GPU); out=$COMBINED"
+
+# 2) Launch. This cluster gives a GPU job the WHOLE node (all 4 GPU + 144 CPU), so we
+#    use -N nodes x ntasks-per-node=4 (= 4N=56 workers) rather than a job array (which
+#    would get 1 whole node per task = only 14 concurrent). Each worker loops over its
+#    slice shards[PROCID::NW], running a FRESH python per SMALL (~260-clip) shard:
+#    the fresh process bounds the pipeline's ~0.75 GB/clip RSS growth (VAE feature cache)
+#    -> 4 workers x ~214 GB peak = ~856 GB < 979 GB/node. Per-shard .done makes resubmit
+#    skip finished work; --no-requeue avoids a destructive auto-restart on node glitches.
+WORKERS=$(( NODES * GPUS ))
+# CORDON-PROOF: the k8s operator cordons worker pods mid-run (node-pool scale-down);
+# a gang -N$NODES job dies entirely if ANY one node is cordoned. So submit a JOB ARRAY
+# of single-node tasks (--array=0..NODES-1, each -N1): a cordon kills+requeues only ONE
+# array task; the other nodes keep running. Global worker idx = ARRAY_TASK_ID*GPUS+LOCALID
+# over WORKERS=NODES*GPUS; per-shard .done makes every (re)start skip finished work.
+WORKERS=$(( NODES * GPUS ))
+sbatch --array=0-$(( NODES - 1 )) -N1 --gres=gpu:$GPUS --ntasks-per-node=$GPUS --exclusive --mem=0 \
+  -t "$TIME" --requeue -p "$PARTITION" -J preprocess_track \
+  --chdir="$WORK/FastVideo" -o "$WORK/logs/preprocess_track_%A_%a.out" -e "$WORK/logs/preprocess_track_%A_%a.out" \
+  --wrap "srun --chdir=$WORK/FastVideo bash -lc '
+    set -uo pipefail
+    source .venv/bin/activate
+    export HOME=$WORK TRITON_CACHE_DIR=$WORK/.triton XDG_CACHE_HOME=$WORK/.cache \
+      HF_HOME=$WORK/.hf TORCH_HOME=$WORK/.torch TOKENIZERS_PARALLELISM=false
+    export CUDA_VISIBLE_DEVICES=\$(( SLURM_LOCALID % $GPUS ))
+    export WORLD_SIZE=1 RANK=0 LOCAL_RANK=0 MASTER_ADDR=127.0.0.1
+    export MASTER_PORT=\$(( 29500 + SLURM_LOCALID ))
+    W=\$(( SLURM_ARRAY_TASK_ID * $GPUS + SLURM_LOCALID )); NW=$WORKERS
+    echo \"[worker \$W/\$NW arr=\$SLURM_ARRAY_TASK_ID] host=\$(hostname) CVD=\$CUDA_VISIBLE_DEVICES\"
+    for IDX in \$(seq \$W \$NW $(( NUM_SHARDS - 1 ))); do
+      SHARD=\$(printf shard_%05d \$IDX); SDIR=$SHARDS_DIR/\$SHARD; ODIR=$COMBINED/\$SHARD
+      [ -f \"\$ODIR/.done\" ] && continue
+      rm -rf \"\$ODIR\"
+      if python fastvideo/pipelines/preprocess/v1_preprocess.py \
+        --model_path $MODEL --preprocess_task i2v_track \
+        --data_merge_path \$SDIR/merge.txt --output_dir \$ODIR \
+        --max_height $MAX_H --max_width $MAX_W --num_frames $NUM_FRAMES \
+        --train_fps $TRAIN_FPS --num_latent_t $NUM_LATENT_T --vae_precision $VAE_PREC \
+        --preprocess_video_batch_size $BATCH --dataloader_num_workers $NW; then
+        touch \"\$ODIR/.done\"; echo \"[worker \$W] \$SHARD done\"
+      else
+        echo \"[worker \$W] \$SHARD FAILED (rc=\$?), skipping\"
+      fi
+    done
+    echo \"[worker \$W] all shards done\"'"

data_pipeline/run_seg_slurm.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+# Multi-node data-parallel frame-0 segmentation (SAM2.1-b+) over OpenVid-1M.
+# Bare node (torch-only ultralytics needs just the driver + Lustre venv, no container).
+# Each Slurm task = one SAM worker; PROCS_PER_GPU tasks share each GPU.
+#
+# Usage:
+#   DATA_DIR=/mnt/lustre/vlm-s4duan/openvid_1m NODES=6 PROCS_PER_GPU=2 \
+#   bash data_pipeline/run_seg_slurm.sh
+set -euo pipefail
+WORK=/mnt/lustre/vlm-s4duan
+DATA_DIR=${DATA_DIR:-$WORK/openvid_1m}
+NODES=${NODES:-1}; GPUS=${GPUS:-4}; PROCS_PER_GPU=${PROCS_PER_GPU:-2}
+MODEL=${MODEL:-sam2.1_b.pt}; VIDEOS_SUBDIR=${VIDEOS_SUBDIR:-clips}
+TASKS_PER_NODE=$(( GPUS * PROCS_PER_GPU ))
+CPUS_PER_TASK=$(( 128 / TASKS_PER_NODE )); [ "$CPUS_PER_TASK" -lt 1 ] && CPUS_PER_TASK=1
+TIME=${TIME:-12:00:00}
+mkdir -p "$WORK/logs"
+echo "nodes=$NODES gpus/node=$GPUS procs/gpu=$PROCS_PER_GPU -> $((NODES*TASKS_PER_NODE)) workers, model=$MODEL"
+
+sbatch -N "$NODES" --gres=gpu:$GPUS --ntasks-per-node=$TASKS_PER_NODE --exclusive \
+  --cpus-per-task=$CPUS_PER_TASK --mem=0 -t "$TIME" -J seg_sam_dp \
+  --chdir="$WORK/FastVideo" -o "$WORK/logs/seg_sam_dp_%j.out" -e "$WORK/logs/seg_sam_dp_%j.out" \
+  --wrap "srun --chdir=$WORK/FastVideo \
+    bash -lc 'source .venv/bin/activate && \
+      export TORCH_HOME=$WORK/.torch HF_HOME=$WORK/.hf MPLCONFIGDIR=$WORK/.mpl \
+        YOLO_CONFIG_DIR=$WORK/.ultralytics TOKENIZERS_PARALLELISM=false \
+        PYTHONPATH=$WORK/FastVideo:$WORK/FastVideo/data_pipeline && \
+      python data_pipeline/segment_tracks_mp.py \
+        --data-dir $DATA_DIR --videos-subdir $VIDEOS_SUBDIR --model $MODEL --gpus-per-node $GPUS'"

data_pipeline/run_segment_synth_slurm.sh

@@ -0,0 +1,25 @@
+#!/usr/bin/env bash
+# SAM 2.1-b+ segmentation of frame-0 for our synth mp4s. Adds object_ids +
+# track_weights to each tracks .npz. Idempotent (skips npz already labeled).
+set -euo pipefail
+WORK=/mnt/lustre/vlm-s4duan
+: "${DATA_DIR:?}"
+NODES=${NODES:-4}; GPUS=${GPUS:-4}; PROCS_PER_GPU=${PROCS_PER_GPU:-2}
+TASKS_PER_NODE=$(( GPUS * PROCS_PER_GPU ))
+CPUS_PER_TASK=$(( 128 / TASKS_PER_NODE )); [ "$CPUS_PER_TASK" -lt 1 ] && CPUS_PER_TASK=1
+NUM_SHARDS=$(( NODES * TASKS_PER_NODE ))
+
+mkdir -p "$WORK/logs"
+echo "nodes=$NODES gpus/node=$GPUS procs/gpu=$PROCS_PER_GPU -> $NUM_SHARDS workers"
+
+sbatch -N "$NODES" --gres=gpu:$GPUS --ntasks-per-node=$TASKS_PER_NODE --exclusive \
+  --cpus-per-task=$CPUS_PER_TASK --mem=0 -t 4:00:00 -J segment_synth \
+  --chdir="$WORK/FastVideo" -o "$WORK/logs/segment_synth_%j.out" -e "$WORK/logs/segment_synth_%j.out" \
+  --wrap "srun --chdir=$WORK/FastVideo \
+    bash -lc 'source .venv/bin/activate && \
+      export TORCH_HOME=$WORK/.torch HF_HOME=$WORK/.hf HOME=$WORK MPLCONFIGDIR=$WORK/.mpl \
+        YOLO_CONFIG_DIR=$WORK/.ultralytics TOKENIZERS_PARALLELISM=false && \
+      export CUDA_VISIBLE_DEVICES=\$(( SLURM_LOCALID % $GPUS )) && \
+      python data_pipeline/segment_tracks.py \
+        --data-dir $DATA_DIR \
+        --num-shards $NUM_SHARDS --shard \$SLURM_PROCID'"

data_pipeline/run_segment_tracks.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+# Run segment_tracks.py in parallel across 4 GPUs.
+# Usage: bash data_pipeline/run_segment_tracks.sh [extra args]
+# Example: bash data_pipeline/run_segment_tracks.sh --limit 20
+
+DATA_DIR=/home/hal-kevin/data/motion-stream-test
+WORLD_SIZE=4
+LOG_FILE=data_pipeline/segment_tracks.log
+
+> $LOG_FILE  # truncate on each run
+
+echo "[seg] launching $WORLD_SIZE workers... logging to $LOG_FILE"
+
+for RANK in $(seq 0 $((WORLD_SIZE - 1))); do
+    CUDA_VISIBLE_DEVICES=$RANK python -u data_pipeline/segment_tracks.py \
+        --data-dir $DATA_DIR \
+        --videos-subdir roundtrip_videos \
+        --conf 0.75 --iou 0.9 --imgsz 1024 \
+        --vis-override-every 3 --viz \
+        --rank $RANK --world-size $WORLD_SIZE \
+        "$@" \
+        >> $LOG_FILE 2>&1 &
+done
+
+wait
+echo "[seg] all done. log at $LOG_FILE"

data_pipeline/run_synth_gen_slurm.sh

@@ -0,0 +1,62 @@
+#!/usr/bin/env bash
+# Multi-node data-parallel Wan2.2-T2V-A14B synthetic video generation.
+#
+# One VideoGenerator per GPU (num_gpus=1, NO SP/TP -- data-parallel is faster at scale).
+# Node layout: NODES x 4 GPU x 1 proc/GPU = 4*NODES independent single-GPU workers.
+# Worker W in [0, 4*NODES) owns prompts[W :: 4*NODES] (stride slice of a fixed shuffle).
+#
+# CORDON-PROOF: k8s cordons worker pods mid-run; a gang -N$NODES job dies if ANY node is
+# cordoned. So we submit a JOB ARRAY of single-node tasks (--array=0..NODES-1, each -N1):
+# a cordon kills+requeues only ONE array task; the others keep running. Per-video mp4
+# existence makes every (re)start skip finished work -> fully resumable.
+#
+# Usage:
+#   NODES=12 MAX_VIDEOS=100000 bash data_pipeline/run_synth_gen_slurm.sh
+#   NODES=1  MAX_VIDEOS=8  SMOKE=1 bash data_pipeline/run_synth_gen_slurm.sh   # smoke test
+set -euo pipefail
+WORK=/mnt/lustre/vlm-s4duan
+FV=$WORK/FastVideo
+PROMPTS=${PROMPTS:-$FV/examples/dataset/vidprom/prompts/vidprom_filtered_extended.txt}
+OUTDIR=${OUTDIR:-$WORK/data/wan22_synth_720p}
+NODES=${NODES:-12}
+GPUS=${GPUS:-4}
+WORKERS=$(( NODES * GPUS ))
+MAX_VIDEOS=${MAX_VIDEOS:-100000}
+HEIGHT=${HEIGHT:-720}; WIDTH=${WIDTH:-1280}
+NUM_FRAMES=${NUM_FRAMES:-121}; DROP=${DROP:-8}; FPS=${FPS:-16}
+STEPS=${STEPS:-40}; GS=${GS:-4.0}; GS2=${GS2:-3.0}
+SEED_BASE=${SEED_BASE:-1024}; SHUFFLE_SEED=${SHUFFLE_SEED:-1234}
+PARTITION=${PARTITION:-all}
+TIME=${TIME:-24:00:00}
+SMOKE=${SMOKE:-0}
+CPUS_PER_TASK=$(( 128 / GPUS )); [ "$CPUS_PER_TASK" -lt 1 ] && CPUS_PER_TASK=1
+mkdir -p "$WORK/logs" "$OUTDIR"
+
+[ -f "$PROMPTS" ] || { echo "PROMPTS not found: $PROMPTS"; exit 1; }
+echo "NODES=$NODES GPUS=$GPUS -> $WORKERS workers | target=$MAX_VIDEOS videos | out=$OUTDIR"
+echo "gen $((NUM_FRAMES+DROP))f -> keep ${NUM_FRAMES} (drop ${DROP}) @ ${WIDTH}x${HEIGHT} ${FPS}fps, ${STEPS} steps, CFG ${GS}/${GS2}"
+
+sbatch --array=0-$(( NODES - 1 )) -N1 --gres=gpu:$GPUS --ntasks-per-node=$GPUS --exclusive --mem=0 \
+  --cpus-per-task=$CPUS_PER_TASK -t "$TIME" --requeue -p "$PARTITION" -J synth_gen \
+  --chdir="$FV" -o "$WORK/logs/synth_gen_%A_%a.out" -e "$WORK/logs/synth_gen_%A_%a.out" \
+  --wrap "srun --chdir=$FV bash -lc '
+    set -uo pipefail
+    source .venv/bin/activate
+    export HOME=$WORK HF_HOME=$WORK/.hf TORCH_HOME=$WORK/.torch MPLCONFIGDIR=$WORK/.mpl \
+      XDG_CACHE_HOME=$WORK/.cache TOKENIZERS_PARALLELISM=false NCCL_CUMEM_ENABLE=0 \
+      PYTHONPATH=$FV
+    export CUDA_VISIBLE_DEVICES=\$(( SLURM_LOCALID % $GPUS ))
+    export TRITON_CACHE_DIR=/tmp/triton_synth_\${SLURM_LOCALID}
+    export WORLD_SIZE=1 RANK=0 LOCAL_RANK=0 MASTER_ADDR=127.0.0.1 MASTER_PORT=\$(( 29700 + SLURM_LOCALID ))
+    mkdir -p \$TRITON_CACHE_DIR
+    W=\$(( SLURM_ARRAY_TASK_ID * $GPUS + SLURM_LOCALID ))
+    echo \"[launch] worker \$W/$WORKERS host=\$(hostname) CVD=\$CUDA_VISIBLE_DEVICES\"
+    python data_pipeline/gen_synth_worker.py \
+      --prompts $PROMPTS --output-dir $OUTDIR \
+      --worker-id \$W --num-workers $WORKERS --max-videos $MAX_VIDEOS \
+      --height $HEIGHT --width $WIDTH --num-frames $NUM_FRAMES --drop $DROP --fps $FPS \
+      --steps $STEPS --guidance-scale $GS --guidance-scale-2 $GS2 \
+      --seed-base $SEED_BASE --shuffle-seed $SHUFFLE_SEED
+  '"
+echo "submitted. monitor: squeue -u \$USER -n synth_gen ; tail -f $WORK/logs/synth_gen_*.out"
+echo "progress:  python data_pipeline/merge_synth_manifests.py --output-dir $OUTDIR"

data_pipeline/run_synth_pipeline.sh

@@ -0,0 +1,119 @@
+#!/bin/bash
+# Synth toy post-processing: tracks (+SAM) and parquet on the SINGLE synth video dir (not sharded).
+# Mirrors run_openvid_shard.sh's PHASE interface for the 720p/24fps synth set produced by
+# gen_synth_i2v_worker.py, reusing the same extract_tracks.py --segment and v1_preprocess calls.
+#
+#   PHASE=tracks  COMPILE=0 bash data_pipeline/run_synth_pipeline.sh   # CoTracker + SAM -> tracks/*.npz
+#   PHASE=parquet           bash data_pipeline/run_synth_pipeline.sh   # captions + points_path + Stage 5 -> parquet
+#   PHASE=both              bash data_pipeline/run_synth_pipeline.sh
+#
+# Resumable: extract_tracks skips clips whose npz exists (unless FORCE_TRACKS=1); v1_preprocess
+# skips clip ids already written. Run tracks first (across all GPUs), then parquet (single-GPU).
+set -uo pipefail
+cd "$(dirname "$0")/.."
+
+DATA_ROOT=${DATA_ROOT:-/home/hal-kevin/data/motion-stream-synth}
+HEIGHT=${HEIGHT:-720}; WIDTH=${WIDTH:-1280}
+NUM_FRAMES=${NUM_FRAMES:-121}; FPS=${FPS:-24}         # MUST match the generated videos (24fps)
+GRID=${GRID:-50}
+GPUS=${GPUS:-0,1,2,3}
+AMP=${AMP:-1}; COMPILE=${COMPILE:-0}; FORCE_TRACKS=${FORCE_TRACKS:-0}
+# v1_preprocess only uses the VAE / T5 / CLIP from MODEL_PATH (same Wan2.1 VAE across sizes), so
+# any Wan2.1 model dir with those encoders works; override to a lighter one if you have it.
+MODEL_PATH=${MODEL_PATH:-/home/hal-kevin/models/Wan2.1-I2V-14B-720P-Diffusers}
+PARQUET_WORKERS=${PARQUET_WORKERS:-2}
+VID_SUBDIR=videos
+PHASE=${PHASE:-both}
+case "$PHASE" in
+    tracks)  TRACKS=1; PARQUET=0 ;;
+    parquet) TRACKS=0; PARQUET=1 ;;
+    both)    TRACKS=1; PARQUET=1 ;;
+    *) echo "[synth] ERROR: PHASE must be tracks|parquet|both (got '$PHASE')" >&2; exit 1 ;;
+esac
+
+IFS=',' read -ra GPU_ARR <<< "$GPUS"; WORLD_SIZE=${#GPU_ARR[@]}
+LOG_DIR="$DATA_ROOT/logs"; mkdir -p "$LOG_DIR" "$DATA_ROOT/tracks"
+N_VID=$(ls "$DATA_ROOT/$VID_SUBDIR"/*.mp4 2>/dev/null | wc -l || true)
+[ "$N_VID" -gt 0 ] || { echo "[synth] no videos under $DATA_ROOT/$VID_SUBDIR -- run gen_synth_i2v_worker first" >&2; exit 1; }
+echo "[synth] DATA_ROOT=$DATA_ROOT videos=$N_VID PHASE=$PHASE ${WIDTH}x${HEIGHT}@${FPS}fps x${NUM_FRAMES}f GPUS=$GPUS"
+
+# --- tracks (+ SAM object_ids / track_weights, fused) across all GPUs -------------------
+if [[ "$TRACKS" == "1" ]]; then
+    echo "[synth] extracting tracks (+segment) across $WORLD_SIZE GPU(s) ..."
+    SPEED=(); [[ "$AMP" == "1" ]] && SPEED+=(--amp); [[ "$COMPILE" == "1" ]] && SPEED+=(--compile)
+    FORCE_ARGS=(); [[ "$FORCE_TRACKS" == "1" ]] && FORCE_ARGS=(--force)
+    if [[ "$COMPILE" == "1" ]]; then
+        export TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR:-$HOME/.cache/torchinductor}
+        export TRITON_CACHE_DIR=${TRITON_CACHE_DIR:-$HOME/.cache/triton}
+    fi
+    pids=()
+    for i in "${!GPU_ARR[@]}"; do
+        CUDA_VISIBLE_DEVICES=${GPU_ARR[$i]} python -u data_pipeline/extract_tracks.py \
+            --data-dir "$DATA_ROOT" --videos-subdir "$VID_SUBDIR" --out-subdir tracks \
+            --grid-size "$GRID" --device cuda \
+            --detect-entries --sam-conf 0.75 --sam-iou 0.9 --sam-imgsz 1024 \
+            --entry-sample-every 2 --entry-min-area 0.001 --entry-new-area 0.5 \
+            --segment --vis-override-every 3 \
+            "${SPEED[@]}" "${FORCE_ARGS[@]}" \
+            --rank "$i" --world-size "$WORLD_SIZE" \
+            >> "$LOG_DIR/tracks.log" 2>&1 &
+        pids+=($!)
+    done
+    tfail=0; for p in "${pids[@]}"; do wait "$p" || tfail=$((tfail + 1)); done
+    N_NPZ=$(ls "$DATA_ROOT"/tracks/*.npz 2>/dev/null | wc -l || true)
+    echo "[synth] tracks: $N_NPZ/$N_VID npz (failed workers: $tfail) -- log: $LOG_DIR/tracks.log"
+    [[ "$tfail" -gt 0 ]] && echo "[synth] WARNING: some track workers failed; inspect the log and re-run (resumable)"
+fi
+
+# --- parquet (Stage 5): captions from gen + points_path patch + v1_preprocess -----------
+if [[ "$PARQUET" == "1" ]]; then
+    N_NPZ=$(ls "$DATA_ROOT"/tracks/*.npz 2>/dev/null | wc -l || true)
+    if [[ "$N_NPZ" -lt "$N_VID" ]]; then
+        echo "[synth] SKIPPING parquet: tracks incomplete ($N_NPZ/$N_VID) -- run PHASE=tracks first" >&2
+        exit 1
+    fi
+    # 1. compile the gen manifest shards -> videos2caption.json (real captions) + merge.txt
+    python data_pipeline/merge_synth_manifests.py --output-dir "$DATA_ROOT"
+    # 2. patch points_path (tracks/<stem>.npz) into each entry, preserving the captions
+    python - "$DATA_ROOT" <<'PY'
+import json, sys
+from pathlib import Path
+root = Path(sys.argv[1]); j = root / "videos2caption.json"; td = root / "tracks"
+items = json.loads(j.read_text())
+n = 0
+for it in items:
+    # preprocess validation computes num_frames = ceil(fps*duration); the gen manifest omits
+    # 'duration', which makes it 0 and rejects every clip -- derive it from num_frames/fps.
+    if not it.get("duration") and it.get("fps"):
+        it["duration"] = it["num_frames"] / float(it["fps"])
+    npz = td / (Path(it["path"]).stem + ".npz")
+    if npz.exists():
+        it["points_path"] = str(npz.resolve()); n += 1
+j.write_text(json.dumps(items, indent=2))
+print(f"[synth] manifest: {len(items)} entries, {n} with points_path (duration patched)")
+PY
+    # 3. Stage 5. --train_fps MUST equal the generated fps (24) or FrameSamplingStage resamples
+    #    and the latents stop aligning with the tracks. num_latent_t = (num_frames-1)/4 + 1.
+    NLT=$(( (NUM_FRAMES - 1) / 4 + 1 ))
+    PQ_OUT="${PARQUET_ROOT:-$DATA_ROOT/preprocessed_i2v_track}"
+    echo "[synth] Stage 5: v1_preprocess (${HEIGHT}x${WIDTH}, ${NUM_FRAMES}f, num_latent_t=$NLT, train_fps=$FPS) -> $PQ_OUT"
+    CUDA_VISIBLE_DEVICES="${PARQUET_GPU:-${GPU_ARR[0]}}" \
+    torchrun --nproc_per_node=1 -m fastvideo.pipelines.preprocess.v1_preprocess \
+        --model_path "$MODEL_PATH" \
+        --data_merge_path "$DATA_ROOT/merge.txt" \
+        --output_dir "$PQ_OUT" \
+        --preprocess_task i2v_track \
+        --num_frames "$NUM_FRAMES" \
+        --num_latent_t "$NLT" \
+        --train_fps "$FPS" \
+        --max_height "$HEIGHT" \
+        --max_width "$WIDTH" \
+        --preprocess_video_batch_size 1 \
+        --dataloader_num_workers "$PARQUET_WORKERS" \
+        --samples_per_file "${PARQUET_SAMPLES:-64}" \
+        --flush_frequency "${PARQUET_FLUSH:-8}" \
+        2>&1 | tee -a "$LOG_DIR/parquet.log"
+    N_PQ=$(find "$PQ_OUT" -name '*.parquet' 2>/dev/null | wc -l || true)
+    echo "[synth] parquet: $N_PQ file(s) -> $PQ_OUT/combined_parquet_dataset"
+    echo "[synth] point the overfit config data_path at: $PQ_OUT/combined_parquet_dataset"
+fi

data_pipeline/run_tracks_slurm.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+# Multi-node / multi-GPU data-parallel CoTracker on this Slinky GB200 cluster.
+# Each Slurm task = one CoTracker worker (B=1); PROCS_PER_GPU tasks share each GPU.
+# Runs inside the enroot `fvbuild` container (needed for the venv's CUDA runtime).
+#
+# Usage:
+#   VIDEO_LIST=/mnt/lustre/vlm-s4duan/openvid/videos.txt \
+#   OUT_DIR=/mnt/lustre/vlm-s4duan/openvid/tracks \
+#   NODES=4 PROCS_PER_GPU=2 FPS=24 NUM_FRAMES=121 GRID=50 \
+#   bash data_pipeline/run_tracks_slurm.sh
+set -euo pipefail
+WORK=/mnt/lustre/vlm-s4duan
+: "${VIDEO_LIST:?set VIDEO_LIST}"; : "${OUT_DIR:?set OUT_DIR}"
+NODES=${NODES:-1}; GPUS=${GPUS:-4}; PROCS_PER_GPU=${PROCS_PER_GPU:-2}
+FPS=${FPS:-24}; NUM_FRAMES=${NUM_FRAMES:-121}; GRID=${GRID:-50}
+CLIPS_ARG=""; [ -n "${CLIPS_DIR:-}" ] && CLIPS_ARG="--clips-dir $CLIPS_DIR"
+TASKS_PER_NODE=$(( GPUS * PROCS_PER_GPU ))
+CPUS_PER_TASK=$(( 128 / TASKS_PER_NODE )); [ "$CPUS_PER_TASK" -lt 1 ] && CPUS_PER_TASK=1
+
+mkdir -p "$OUT_DIR" "$WORK/logs"
+echo "nodes=$NODES gpus/node=$GPUS procs/gpu=$PROCS_PER_GPU -> $((NODES*TASKS_PER_NODE)) workers"
+
+# Bare-node: torch (self-contained cu128 wheels) + CoTracker need only the driver +
+# the Lustre venv — NO container (avoids a per-node image pull).
+sbatch -N "$NODES" --gres=gpu:$GPUS --ntasks-per-node=$TASKS_PER_NODE --exclusive \
+  --cpus-per-task=$CPUS_PER_TASK --mem=0 -t 24:00:00 -J cotracker_dp \
+  --chdir="$WORK/FastVideo" -o "$WORK/logs/cotracker_dp_%j.out" -e "$WORK/logs/cotracker_dp_%j.out" \
+  --wrap "srun --chdir=$WORK/FastVideo \
+    bash -lc 'source .venv/bin/activate && export TORCH_HOME=$WORK/.torch TOKENIZERS_PARALLELISM=false && \
+      python data_pipeline/extract_tracks_mp.py \
+        --video-list $VIDEO_LIST --out-dir $OUT_DIR $CLIPS_ARG \
+        --gpus-per-node $GPUS --fps $FPS --num-frames $NUM_FRAMES --grid-size $GRID'"

data_pipeline/run_tracks_synth_slurm.sh

@@ -0,0 +1,25 @@
+#!/usr/bin/env bash
+# CoTracker extraction for our synth mp4s. Same pattern as run_tracks_slurm.sh but
+# with the aspect/lowres filters disabled (synth videos are exactly 720x1280) and
+# H/W set to our training resolution 480x832.
+set -euo pipefail
+WORK=/mnt/lustre/vlm-s4duan
+: "${VIDEO_LIST:?}"; : "${OUT_DIR:?}"
+NODES=${NODES:-4}; GPUS=${GPUS:-4}; PROCS_PER_GPU=${PROCS_PER_GPU:-2}
+FPS=${FPS:-24}; NUM_FRAMES=${NUM_FRAMES:-121}; GRID=${GRID:-50}
+HEIGHT=${HEIGHT:-480}; WIDTH=${WIDTH:-832}
+TASKS_PER_NODE=$(( GPUS * PROCS_PER_GPU ))
+CPUS_PER_TASK=$(( 128 / TASKS_PER_NODE )); [ "$CPUS_PER_TASK" -lt 1 ] && CPUS_PER_TASK=1
+
+mkdir -p "$OUT_DIR" "$WORK/logs"
+echo "nodes=$NODES gpus/node=$GPUS procs/gpu=$PROCS_PER_GPU -> $((NODES*TASKS_PER_NODE)) workers at ${HEIGHT}x${WIDTH}"
+
+sbatch -N "$NODES" --gres=gpu:$GPUS --ntasks-per-node=$TASKS_PER_NODE --exclusive \
+  --cpus-per-task=$CPUS_PER_TASK --mem=0 -t 6:00:00 -J cotracker_synth \
+  --chdir="$WORK/FastVideo" -o "$WORK/logs/cotracker_synth_%j.out" -e "$WORK/logs/cotracker_synth_%j.out" \
+  --wrap "srun --chdir=$WORK/FastVideo \
+    bash -lc 'source .venv/bin/activate && export TORCH_HOME=$WORK/.torch TOKENIZERS_PARALLELISM=false && \
+      python data_pipeline/extract_tracks_mp.py \
+        --video-list $VIDEO_LIST --out-dir $OUT_DIR \
+        --gpus-per-node $GPUS --fps $FPS --num-frames $NUM_FRAMES --grid-size $GRID \
+        --height $HEIGHT --width $WIDTH --min-height 0 --aspect-tol 10.0'"

data_pipeline/sampling_viz.py

@@ -0,0 +1,213 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Visualize the training-time track sampler on a preprocessed dataset.
+
+For each clip it shows, straight from the stored ``tracks.npz`` (tracks + visibility +
+``object_ids`` + ``track_weights``):
+  (1) the low-rank informativeness HEATMAP over the 50x50 grid on frame 0,
+  (2) the points the SAMPLER KEEPS (exact port of WanTrack ``_augment_tracks``:
+      >=1 pt per SAM object  +  low-rank-weighted draw (1-uniform_frac)  +  uniform draw),
+  (3) an overlay video of ONLY the kept tracks over the source video, colored by low-rank
+      weight (blue = generic, red = informative), so you can see if good traces are sampled.
+
+CPU only (no GPU / no model needed).
+
+    .venv/bin/python data_pipeline/sampling_viz.py render --data-dir <dataset root> --out <dir> --limit 8
+    .venv/bin/python data_pipeline/sampling_viz.py serve --viz-dir <dir> --share
+"""
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(__file__))
+from track_informativeness import _draw_overlay, _read_frames, _norm  # noqa: E402
+
+
+def sample_kept(tracks: np.ndarray,
+                vis: np.ndarray,
+                oid: np.ndarray | None,
+                tw: np.ndarray | None,
+                k: int,
+                uniform_frac: float = 0.3,
+                diversity: float = 1.0,
+                seed: int = 0) -> np.ndarray:
+    """Exact numpy port of WanTrack._augment_tracks sampling -> kept bool [N]."""
+    rng = np.random.default_rng(seed)
+    N = tracks.shape[1]
+    valid = (vis[0] > 0.5).astype(np.float64)  # only frame-0-visible are valid queries
+    if tw is not None and tw.size == N:
+        w = (0.05 + diversity * tw.astype(np.float64)) * valid
+    else:
+        disp = np.sqrt(((tracks - tracks[0:1])**2).sum(-1))
+        motion = (disp * (vis > 0.5)).max(0)
+        w = (1.0 + diversity * (motion / (motion.mean() + 1e-6))) * valid
+    if w.sum() <= 0:
+        w = valid.copy()
+    if w.sum() <= 0:
+        w = np.ones(N)
+    keep = np.zeros(N, bool)
+    # (a) object coverage: one weighted pick per present segment
+    if oid is not None:
+        for o in np.unique(oid):
+            if int(o) < 0:
+                continue
+            idx = np.where(oid == o)[0]
+            ww = w[idx]
+            ww = np.ones_like(ww) if ww.sum() <= 0 else ww
+            keep[rng.choice(idx, p=ww / ww.sum())] = True
+    # (b) low-rank-weighted draw for (1 - uniform_frac) of the remaining budget
+    rem = max(0, k - int(keep.sum()))
+    n_weighted = rem - int(round(rem * uniform_frac))
+    pool = np.where((~keep) & (w > 0))[0]
+    nw = min(n_weighted, pool.size)
+    if nw > 0:
+        pw = w[pool] / w[pool].sum()
+        keep[rng.choice(pool, size=nw, replace=False, p=pw)] = True
+    # (c) uniform draw to fill up to k (also backfills weighted shortfall)
+    pool2 = np.where((~keep) & (valid > 0))[0]
+    nu = min(max(0, k - int(keep.sum())), pool2.size)
+    if nu > 0:
+        keep[rng.choice(pool2, size=nu, replace=False)] = True
+    return keep
+
+
+def cmd_render(args: argparse.Namespace) -> None:
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    import imageio.v2 as imageio
+    turbo = matplotlib.colormaps["turbo"]
+
+    data = Path(args.data_dir)
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    items = json.loads((data / "videos2caption.json").read_text())[:args.limit]
+    entries = []
+    for i, it in enumerate(items, 1):
+        stem = Path(it["path"]).stem
+        vpath = str(data / "videos" / it["path"])
+        npz = it.get("points_path") or str(data / "tracks" / f"{stem}.npz")
+        frames = _read_frames(vpath)
+        d = np.load(npz)
+        tracks = d["tracks"].astype(np.float32)[:frames.shape[0]]
+        vis = d["visibility"].astype(np.float32)[:frames.shape[0]]
+        oid = d["object_ids"].astype(np.int64) if "object_ids" in d else None
+        tw = d["track_weights"].astype(np.float32) if "track_weights" in d else None
+        N = tracks.shape[1]
+        k = int(args.k)
+        keep = sample_kept(tracks, vis, oid, tw, k, args.uniform_frac, args.diversity, args.seed)
+        twn = _norm(tw) if tw is not None else np.zeros(N)
+        x, y = tracks[0, :, 0], tracks[0, :, 1]
+
+        # ---- figure: (1) lowrank heatmap  (2) kept vs dropped
+        fig, axes = plt.subplots(1, 2, figsize=(17, 5.6))
+        axes[0].imshow(frames[0])
+        sc = axes[0].scatter(x, y, c=twn, cmap="turbo", s=12, vmin=0, vmax=1)
+        axes[0].set_title(f"low-rank informativeness heatmap (all {N} pts)", fontsize=11)
+        axes[0].axis("off")
+        fig.colorbar(sc, ax=axes[0], fraction=0.03)
+        axes[1].imshow(frames[0])
+        axes[1].scatter(x[~keep], y[~keep], c="0.5", s=3, alpha=0.5)  # dropped
+        axes[1].scatter(x[keep],
+                        y[keep],
+                        c=twn[keep],
+                        cmap="turbo",
+                        s=16,
+                        vmin=0,
+                        vmax=1,
+                        edgecolors="white",
+                        linewidths=0.3)  # kept, colored by weight
+        n_obj = len(np.unique(oid[oid >= 0])) if oid is not None else 0
+        n_cov = len(np.unique(oid[keep][oid[keep] >= 0])) if oid is not None else 0
+        axes[1].set_title(f"kept by sampler: {int(keep.sum())}/{N}  (objects covered {n_cov}/{n_obj})", fontsize=11)
+        axes[1].axis("off")
+        fig.suptitle(f"{stem}  |  {(it['cap'][0] if isinstance(it.get('cap'), list) else '')[:80]}", fontsize=10)
+        fig.tight_layout()
+        heat = f"{stem}_sampling.png"
+        fig.savefig(str(out / heat), dpi=80, bbox_inches="tight")
+        plt.close(fig)
+
+        # ---- overlay video: ONLY kept tracks, colored by low-rank weight
+        kidx = np.where(keep)[0]
+        # subsample for legibility if many kept
+        if kidx.size > 700:
+            kidx = kidx[np.linspace(0, kidx.size - 1, 700).round().astype(int)]
+        pcols = (np.array([turbo(v)[:3] for v in twn[kidx]]) * 255).astype(np.uint8)
+        ov = _draw_overlay(frames, tracks[:, kidx].copy(), vis[:, kidx], pcols, 14, 2, 0.5)
+        vid = f"{stem}_kept.mp4"
+        imageio.mimsave(str(out / vid), ov, fps=int(it.get("fps", 24)), macro_block_size=1)
+
+        stats = {
+            "id": stem,
+            "n_points": int(N),
+            "k_sampled": k,
+            "kept": int(keep.sum()),
+            "objects_covered": f"{n_cov}/{n_obj}",
+            "mean_weight_kept": round(float(twn[keep].mean()), 3) if keep.any() else 0.0,
+            "mean_weight_dropped": round(float(twn[~keep].mean()), 3) if (~keep).any() else 0.0,
+            "frac_kept_informative(>0.5)": round(float((twn[keep] > 0.5).mean()), 3) if keep.any() else 0.0,
+        }
+        entries.append({"id": stem, "heat": heat, "video": vid, "stats": stats})
+        print(
+            f"[samp] [{i}/{len(items)}] {stem}: kept {int(keep.sum())}/{N}, cov {n_cov}/{n_obj}, "
+            f"w_kept={stats['mean_weight_kept']} vs w_drop={stats['mean_weight_dropped']}",
+            flush=True)
+
+    (out / "manifest.json").write_text(json.dumps(entries, indent=2))
+    print(f"[samp] rendered {len(entries)} clips -> {out}", flush=True)
+
+
+def cmd_serve(args: argparse.Namespace) -> None:
+    import gradio as gr
+    viz = Path(args.viz_dir).resolve()
+    entries = json.loads((viz / "manifest.json").read_text())
+    gal = [(str(viz / e["heat"]), e["id"]) for e in entries]
+
+    def show(evt: gr.SelectData):
+        e = entries[evt.index]
+        return str(viz / e["heat"]), str(viz / e["video"]), json.dumps(e["stats"], indent=2)
+
+    with gr.Blocks(title="Track sampling viz") as demo:
+        gr.Markdown("### Training-time track sampler: low-rank heatmap + kept points + kept-track overlay\n"
+                    "Left panel: low-rank informativeness heatmap (bright = informative). Right panel: the points "
+                    "the sampler **keeps** (colored by weight, gray = dropped) — should cover every object and "
+                    "favor bright/informative points while keeping some uniform. Click a clip -> the overlay video "
+                    "of ONLY the kept tracks (blue = generic, red = informative). Good sampling = kept traces follow "
+                    "the moving hands/objects. `stats.mean_weight_kept` should exceed `mean_weight_dropped`.")
+        with gr.Row():
+            g = gr.Gallery(value=gal, columns=2, height=620, label="clips (click)")
+            with gr.Column():
+                heat = gr.Image(label="heatmap (left) + kept points (right)")
+                vid = gr.Video(label="kept tracks over video (colored by low-rank weight)")
+                meta = gr.Code(label="stats", language="json")
+        g.select(show, None, [heat, vid, meta])
+    demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share, allowed_paths=[str(viz)])
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub = p.add_subparsers(dest="cmd", required=True)
+    r = sub.add_parser("render")
+    r.add_argument("--data-dir", required=True)
+    r.add_argument("--out", required=True)
+    r.add_argument("--limit", type=int, default=8)
+    r.add_argument("--k", type=int, default=1500, help="tracks to keep (training samples U[1000,2500])")
+    r.add_argument("--uniform-frac", type=float, default=0.3)
+    r.add_argument("--diversity", type=float, default=1.0)
+    r.add_argument("--seed", type=int, default=0)
+    r.set_defaults(func=cmd_render)
+    s = sub.add_parser("serve")
+    s.add_argument("--viz-dir", required=True)
+    s.add_argument("--host", default="0.0.0.0")
+    s.add_argument("--port", type=int, default=7889)
+    s.add_argument("--share", action="store_true")
+    s.set_defaults(func=cmd_serve)
+    a = p.parse_args()
+    a.func(a)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/seg_compare.py

@@ -0,0 +1,377 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Compare segmentation-model *variants* (not just FastSAM configs) for frame-0
+object segmentation + sparse track sampling, on real openvid clips.
+
+``render`` (GPU) runs each model in no-prompt "everything" mode on frame-0 of each
+clip, labels the CoTracker grid points by object (smallest containing mask wins),
+runs the 1-per-object + extras sparse sampler, and writes per (model, clip):
+  - <..>_masks.png    frame-0 with colored SAM masks
+  - <..>_sparse.png   frame-0 with kept sparse tracks (colored by object) vs dropped (gray)
+  - <..>_tracks.mp4   kept sparse tracks over the whole clip
+plus a per-entry .npz (oid/tw/vis0/xy0) so the dashboard can re-sample live.
+
+``serve`` (CPU, login node) is a gradio dashboard: pick a clip, see every model's
+masks / sparse-sampling side by side, tweak num_sampled live, click to inspect.
+
+aarch64-safe: PyAV decode, headless cv2, no decord.
+
+    # render on a GPU node (piggyback job 365)
+    srun --overlap --jobid=365 --ntasks=1 --chdir=/mnt/lustre/vlm-s4duan/FastVideo \
+      bash -lc 'source .venv/bin/activate && export CUDA_VISIBLE_DEVICES=1 \
+        YOLO_CONFIG_DIR=/mnt/lustre/vlm-s4duan/.ultralytics \
+        TORCH_HOME=/mnt/lustre/vlm-s4duan/.torch HF_HOME=/mnt/lustre/vlm-s4duan/.hf \
+        PYTHONPATH=$PWD:$PWD/data_pipeline MPLCONFIGDIR=/mnt/lustre/vlm-s4duan/.mpl && \
+        python -u data_pipeline/seg_compare.py render --limit 8'
+
+    # serve on the login node, then ssh -L 7862:localhost:7862 <login>
+    .venv/bin/python data_pipeline/seg_compare.py serve
+"""
+# NOTE: no ``from __future__ import annotations`` — gradio needs the real
+# gr.SelectData annotation on the click handler.
+import argparse
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(__file__))
+from segment_tracks import object_ids_for_points, lowrank_track_weights  # noqa: E402
+from sparse_sampling_dashboard import sparse_sample  # noqa: E402
+from track_informativeness import _draw_overlay  # noqa: E402
+
+DATA_DEFAULT = "/mnt/lustre/vlm-s4duan/openvid_1m"
+WEIGHTS_DIR = "/mnt/lustre/vlm-s4duan/models/seg"
+OUT_DEFAULT = "/mnt/lustre/vlm-s4duan/seg_compare_out"
+
+# The model zoo. loader in {FastSAM, SAM}. All run no-prompt everything-mode.
+# min_area_frac drops mask specks (< frac of frame); nms_contain drops a mask if
+# > that fraction of it is already covered by a larger kept mask (collapses
+# part-level fragments into their object). max_masks caps to N largest (0=all).
+DEFAULT_MODELS = [
+    {"name": "fastsam-s", "loader": "FastSAM", "weight": "FastSAM-s.pt",
+     "imgsz": 1024, "conf": 0.4, "iou": 0.9, "min_area_frac": 0.0015, "nms_contain": 0.0, "max_masks": 0},
+    {"name": "fastsam-x", "loader": "FastSAM", "weight": "FastSAM-x.pt",
+     "imgsz": 1024, "conf": 0.4, "iou": 0.9, "min_area_frac": 0.0015, "nms_contain": 0.0, "max_masks": 0},
+    {"name": "mobilesam", "loader": "SAM", "weight": "mobile_sam.pt",
+     "imgsz": 1024, "conf": 0.4, "iou": 0.9, "min_area_frac": 0.0015, "nms_contain": 0.0, "max_masks": 0},
+    {"name": "sam-b", "loader": "SAM", "weight": "sam_b.pt",
+     "imgsz": 1024, "conf": 0.4, "iou": 0.9, "min_area_frac": 0.0015, "nms_contain": 0.0, "max_masks": 0},
+    {"name": "sam2-b", "loader": "SAM", "weight": "sam2_b.pt",
+     "imgsz": 1024, "conf": 0.4, "iou": 0.9, "min_area_frac": 0.0015, "nms_contain": 0.0, "max_masks": 0},
+    {"name": "sam2.1-b+", "loader": "SAM", "weight": "sam2.1_b.pt",
+     "imgsz": 1024, "conf": 0.4, "iou": 0.9, "min_area_frac": 0.0015, "nms_contain": 0.0, "max_masks": 0},
+    {"name": "sam2.1-l", "loader": "SAM", "weight": "sam2.1_l.pt",
+     "imgsz": 1024, "conf": 0.4, "iou": 0.9, "min_area_frac": 0.0015, "nms_contain": 0.0, "max_masks": 0},
+]
+
+# sparse-sampler defaults baked into the rendered PNG/mp4 (dashboard can re-sample the PNG live)
+SPARSE_DEFAULT = {"num_sampled": 20, "mode": "weighted", "seed": 0}
+
+
+def read_frames(path: str) -> np.ndarray:
+    """Full clip -> [T,H,W,3] uint8 via PyAV (aarch64-safe)."""
+    import av
+    c = av.open(path)
+    frames = [f.to_ndarray(format="rgb24") for f in c.decode(video=0)]
+    c.close()
+    return np.stack(frames)
+
+
+def colors(n: int) -> np.ndarray:
+    import colorsys
+    return np.array([[int(255 * v) for v in colorsys.hsv_to_rgb((i * 0.61803) % 1.0, 0.65, 1.0)]
+                     for i in range(max(1, n))], np.uint8)
+
+
+def extract_masks(res, H: int, W: int) -> np.ndarray:
+    if not res or res[0].masks is None:
+        return np.zeros((0, H, W), bool)
+    m = res[0].masks.data.cpu().numpy().astype(bool)  # [M,h,w]
+    if m.shape[0] and m.shape[1:] != (H, W):  # nearest resize w/o cv2
+        import torch
+        t = torch.from_numpy(m.astype(np.uint8))[None].float()
+        t = torch.nn.functional.interpolate(t, size=(H, W), mode="nearest")[0]
+        m = t.numpy().astype(bool)
+    return m
+
+
+def filter_masks(masks: np.ndarray, min_area_frac: float, nms_contain: float, max_masks: int) -> np.ndarray:
+    if masks.shape[0] == 0:
+        return masks
+    H, W = masks.shape[1], masks.shape[2]
+    areas = masks.reshape(masks.shape[0], -1).sum(1).astype(np.float64)
+    if min_area_frac > 0:
+        keep = (areas / float(H * W)) >= min_area_frac
+        masks, areas = masks[keep], areas[keep]
+    if masks.shape[0] and nms_contain > 0:  # drop fragments mostly inside a larger kept mask
+        order = np.argsort(-areas)  # large -> small
+        kept_idx, covered = [], np.zeros((H, W), bool)
+        for i in order:
+            m = masks[i]
+            a = areas[i]
+            if a > 0 and (m & covered).sum() / a > nms_contain:
+                continue
+            kept_idx.append(i)
+            covered |= m
+        masks, areas = masks[kept_idx], areas[kept_idx]
+    if max_masks and masks.shape[0] > max_masks:
+        masks = masks[np.argsort(-areas)[:max_masks]]
+    return masks
+
+
+def mask_panel(frame0: np.ndarray, masks: np.ndarray, oid: np.ndarray, xy0: np.ndarray) -> "Image.Image":
+    from PIL import Image, ImageDraw
+    mcols = colors(len(masks) + 1)
+    base = frame0.astype(np.float32)
+    for mi, m in enumerate(masks):
+        base[m] = 0.55 * base[m] + 0.45 * mcols[mi % len(mcols)][None].astype(np.float32)
+    img = Image.fromarray(base.clip(0, 255).astype(np.uint8))
+    draw = ImageDraw.Draw(img)
+    for o in sorted({int(v) for v in np.unique(oid) if int(v) >= 0}):  # 1 dot per object (its centroid pt)
+        idx = np.where(oid == o)[0]
+        cx, cy = float(xy0[idx, 0].mean()), float(xy0[idx, 1].mean())
+        draw.ellipse([cx - 4, cy - 4, cx + 4, cy + 4], fill=(255, 255, 255), outline=(0, 0, 0))
+    return img
+
+
+def sparse_panel(frame0: np.ndarray, xy0: np.ndarray, oid: np.ndarray, keep: np.ndarray, title: str,
+                 out_path: str) -> None:
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    cmap = plt.get_cmap("tab20")
+    fig, ax = plt.subplots(figsize=(9, 5.2))
+    ax.imshow(frame0)
+    x, y = xy0[:, 0], xy0[:, 1]
+    ax.scatter(x[~keep], y[~keep], c="0.5", s=3, alpha=0.30)
+    for o in np.unique(oid[keep]):
+        m = keep & (oid == o)
+        col = "white" if int(o) < 0 else cmap(int(o) % 20)
+        ax.scatter(x[m], y[m], c=[col], s=80, edgecolors="black", linewidths=1.1)
+    ax.set_title(title, fontsize=10)
+    ax.axis("off")
+    fig.savefig(out_path, dpi=90, bbox_inches="tight")
+    plt.close(fig)
+
+
+def cmd_render(args: argparse.Namespace) -> None:
+    import torch
+    import imageio.v2 as imageio
+    from ultralytics import FastSAM, SAM
+
+    data = Path(args.data_dir)
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    clips_dir = data / ("clips" if (data / "clips").exists() else "videos")
+    models_cfg = json.loads(Path(args.models_json).read_text()) if args.models_json else DEFAULT_MODELS
+    if args.models:  # subset by name
+        want = set(args.models.split(","))
+        models_cfg = [m for m in models_cfg if m["name"] in want]
+    loaders = {"FastSAM": FastSAM, "SAM": SAM}
+
+    # load every model once (reused across all clips)
+    print(f"[seg] loading {len(models_cfg)} models ...", flush=True)
+    loaded = {}
+    for mc in models_cfg:
+        wp = Path(WEIGHTS_DIR) / mc["weight"]
+        loaded[mc["name"]] = loaders[mc["loader"]](str(wp) if wp.exists() else mc["weight"])
+
+    items = json.loads((data / "videos2caption.json").read_text())[:args.limit]
+    print(f"[seg] {len(items)} clips x {len(models_cfg)} models", flush=True)
+    entries, clip_ids = [], []
+    for k, it in enumerate(items, 1):
+        stem = Path(it["path"]).stem
+        clip_ids.append(stem)
+        vpath = str(clips_dir / it["path"])
+        npz = it.get("points_path") or str(data / "tracks" / f"{stem}.npz")
+        frames = read_frames(vpath)
+        H, W = frames.shape[1], frames.shape[2]
+        d = np.load(npz)
+        tracks = d["tracks"].astype(np.float32)[:frames.shape[0]]  # [T,N,2] px
+        vis = d["visibility"].astype(np.float32)[:frames.shape[0]]  # [T,N]
+        xy0, vis0 = tracks[0], vis[0]
+        tw = lowrank_track_weights(tracks)  # [N] informativeness in [0,1]
+        # frame-0 png (shared, for dashboard live re-sampling background)
+        from PIL import Image
+        f0png = f"{stem}__frame0.png"
+        Image.fromarray(frames[0]).save(str(out / f0png))
+
+        for mc in models_cfg:
+            t0 = time.time()
+            res = loaded[mc["name"]](frames[0], device=args.device, retina_masks=True,
+                                     imgsz=mc["imgsz"], conf=mc["conf"], iou=mc["iou"], verbose=False)
+            masks = filter_masks(extract_masks(res, H, W), mc["min_area_frac"], mc["nms_contain"], mc["max_masks"])
+            oid = object_ids_for_points(masks, xy0, H, W)
+            n_obj = int(len({int(v) for v in np.unique(oid) if int(v) >= 0}))
+            keep = sparse_sample(oid, tw, vis0, SPARSE_DEFAULT["num_sampled"], SPARSE_DEFAULT["mode"],
+                                 SPARSE_DEFAULT["seed"])
+            dt = time.time() - t0
+            pfx = f"{mc['name']}__{stem}"
+            mask_panel(frames[0], masks, oid, xy0).save(str(out / f"{pfx}_masks.png"))
+            sparse_panel(frames[0], xy0, oid, keep,
+                         f"{stem} | {mc['name']}: {int(keep.sum())} kept ({n_obj} objs +"
+                         f"{SPARSE_DEFAULT['num_sampled']} extra)", str(out / f"{pfx}_sparse.png"))
+            trk_name = ""
+            if not args.no_video:
+                import matplotlib.pyplot as plt
+                cmap = plt.get_cmap("tab20")
+                kidx = np.where(keep)[0]
+                cols = np.zeros((kidx.size, 3), np.uint8)
+                for o in np.unique(oid[kidx]):
+                    mm = oid[kidx] == o
+                    rgba = (1., 1., 1., 1.) if int(o) < 0 else cmap(int(o) % 20)
+                    cols[mm] = (np.array(rgba[:3]) * 255).astype(np.uint8)
+                ov = _draw_overlay(frames, tracks[:, kidx].copy(), vis[:, kidx], cols, 14, 3, 0.5)
+                trk_name = f"{pfx}_tracks.mp4"
+                imageio.mimsave(str(out / trk_name), ov, fps=int(it.get("fps", 24)), macro_block_size=1)
+            # per-entry npz so the dashboard can re-sample live (no GPU)
+            np.savez(str(out / f"{pfx}.npz"), oid=oid.astype(np.int64), tw=tw, vis0=vis0, xy0=xy0)
+            entries.append({
+                "model": mc["name"], "clip": stem,
+                "caption": (it["cap"][0] if isinstance(it.get("cap"), list) else str(it.get("cap", "")))[:140],
+                "n_masks": int(masks.shape[0]), "n_objects": n_obj,
+                "n_labeled": int((oid >= 0).sum()), "n_points": int(oid.shape[0]),
+                "n_kept": int(keep.sum()), "sec": round(dt, 2),
+                "masks_png": f"{pfx}_masks.png", "sparse_png": f"{pfx}_sparse.png",
+                "tracks_mp4": trk_name, "frame0_png": f0png, "npz": f"{pfx}.npz",
+            })
+            print(f"[seg] [{k}/{len(items)}] {stem} [{mc['name']}]: {masks.shape[0]} masks, "
+                  f"{n_obj} objs, {(oid >= 0).sum()}/{oid.shape[0]} pts labeled, {dt:.1f}s", flush=True)
+        del frames
+        torch.cuda.empty_cache()
+
+    manifest = {"models": models_cfg, "clips": clip_ids, "sparse_default": SPARSE_DEFAULT, "entries": entries}
+    (out / "manifest.json").write_text(json.dumps(manifest, indent=2))
+    print(f"[seg] wrote {len(entries)} entries -> {out}/manifest.json", flush=True)
+
+
+def cmd_serve(args: argparse.Namespace) -> None:
+    import gradio as gr
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    viz = Path(args.viz_dir).resolve()
+    manifest = json.loads((viz / "manifest.json").read_text())
+    entries = manifest["entries"]
+    model_names = [m["name"] for m in manifest["models"]]
+    clip_ids = manifest["clips"]
+    by_clip = {c: [e for e in entries if e["clip"] == c] for c in clip_ids}
+    caption_of = {c: (by_clip[c][0]["caption"] if by_clip.get(c) else "") for c in clip_ids}
+    npz_cache = {}
+
+    def _npz(e):
+        p = e["npz"]
+        if p not in npz_cache:
+            npz_cache[p] = {k: v for k, v in np.load(str(viz / p)).items()}
+        return npz_cache[p]
+
+    def masks_gallery(clip):
+        es = sorted(by_clip.get(clip, []), key=lambda e: model_names.index(e["model"]))
+        return [(str(viz / e["masks_png"]), f'{e["model"]}  •  {e["n_masks"]}m / {e["n_objects"]}o') for e in es]
+
+    def resample_png(e, num_sampled, mode, seed):
+        z = _npz(e)
+        oid, tw, vis0, xy0 = z["oid"], z["tw"], z["vis0"], z["xy0"]
+        keep = sparse_sample(oid, tw, vis0, int(num_sampled), mode, int(seed))
+        frame0 = plt.imread(str(viz / e["frame0_png"]))
+        cmap = plt.get_cmap("tab20")
+        fig, ax = plt.subplots(figsize=(9, 5.2))
+        ax.imshow(frame0)
+        x, y = xy0[:, 0], xy0[:, 1]
+        ax.scatter(x[~keep], y[~keep], c="0.5", s=3, alpha=0.30)
+        for o in np.unique(oid[keep]):
+            m = keep & (oid == o)
+            col = "white" if int(o) < 0 else cmap(int(o) % 20)
+            ax.scatter(x[m], y[m], c=[col], s=80, edgecolors="black", linewidths=1.1)
+        n_obj = int(len({int(v) for v in np.unique(oid) if int(v) >= 0}))
+        ax.set_title(f'{e["clip"]} | {e["model"]}: {int(keep.sum())} kept ({n_obj} objs +{num_sampled}, {mode})',
+                     fontsize=10)
+        ax.axis("off")
+        import numpy as _np
+        fig.canvas.draw()
+        buf = _np.asarray(fig.canvas.buffer_rgba())[..., :3].copy()
+        plt.close(fig)
+        return buf
+
+    def sparse_gallery(clip, num_sampled, mode, seed):
+        es = sorted(by_clip.get(clip, []), key=lambda e: model_names.index(e["model"]))
+        return [(resample_png(e, num_sampled, mode, seed), f'{e["model"]}  •  {e["n_kept"]}→') for e in es]
+
+    with gr.Blocks(title="Segmentation model comparison") as demo:
+        gr.Markdown(
+            "## Segmentation model comparison — frame-0 masks + sparse track sampling\n"
+            "Pick a **clip**; every model is shown side by side. **Masks** tab = raw everything-mode "
+            "segmentation (tile label `<#masks>m / <#objects-with-points>o`). **Sparse sampling** tab = "
+            "1 track per object + N extras (tweak live). Fewer, cleaner object masks = better for our "
+            "1-per-object recipe. Click a tile to inspect it big + its track-overlay video.")
+        clip_dd = gr.Dropdown(clip_ids, value=clip_ids[0], label=f"clip (1 of {len(clip_ids)})")
+        clip_cap = gr.Markdown(f"*{caption_of[clip_ids[0]]}*")
+        with gr.Tab("Masks (everything-mode)"):
+            mg = gr.Gallery(value=masks_gallery(clip_ids[0]), columns=3, height=680,
+                            label="SAM masks per model (click a tile)")
+        with gr.Tab("Sparse sampling"):
+            with gr.Row():
+                num_sl = gr.Slider(0, 150, value=manifest["sparse_default"]["num_sampled"], step=1,
+                                   label="num_sampled (extras beyond 1-per-object)")
+                mode_rd = gr.Radio(["weighted", "random"], value=manifest["sparse_default"]["mode"], label="extras mode")
+                seed_sl = gr.Slider(0, 50, value=manifest["sparse_default"]["seed"], step=1, label="seed")
+            sg = gr.Gallery(value=sparse_gallery(clip_ids[0], manifest["sparse_default"]["num_sampled"],
+                                                 manifest["sparse_default"]["mode"],
+                                                 manifest["sparse_default"]["seed"]),
+                            columns=3, height=640, label="sparse-sampled tracks per model (click a tile)")
+        gr.Markdown("### Inspect")
+        with gr.Row():
+            big = gr.Image(label="selected panel", height=460)
+            vid = gr.Video(label="kept-tracks overlay over the clip", height=460)
+        meta = gr.Code(label="metrics", language="json")
+
+        def on_clip(clip, num_sampled, mode, seed):
+            return masks_gallery(clip), sparse_gallery(clip, num_sampled, mode, seed), f"*{caption_of.get(clip, '')}*"
+
+        def on_sparse_ctrl(clip, num_sampled, mode, seed):
+            return sparse_gallery(clip, num_sampled, mode, seed)
+
+        def pick(clip, evt: gr.SelectData):
+            es = sorted(by_clip.get(clip, []), key=lambda e: model_names.index(e["model"]))
+            e = es[evt.index]
+            v = str(viz / e["tracks_mp4"]) if e.get("tracks_mp4") else None
+            return str(viz / e["masks_png"]), v, json.dumps(e, indent=2)
+
+        clip_dd.change(on_clip, [clip_dd, num_sl, mode_rd, seed_sl], [mg, sg, clip_cap])
+        for c in (num_sl, mode_rd, seed_sl):
+            c.change(on_sparse_ctrl, [clip_dd, num_sl, mode_rd, seed_sl], sg)
+        mg.select(pick, clip_dd, [big, vid, meta])
+        sg.select(pick, clip_dd, [big, vid, meta])
+
+    demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share,
+                        allowed_paths=[str(viz)])
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub = p.add_subparsers(dest="cmd", required=True)
+    r = sub.add_parser("render")
+    r.add_argument("--data-dir", default=DATA_DEFAULT)
+    r.add_argument("--out", default=OUT_DEFAULT)
+    r.add_argument("--limit", type=int, default=8)
+    r.add_argument("--device", default="cuda")
+    r.add_argument("--models", default=None, help="comma-separated subset of model names")
+    r.add_argument("--models-json", default=None, help="path to a JSON list of model config dicts")
+    r.add_argument("--no-video", action="store_true", help="skip the track-overlay mp4s (faster)")
+    r.set_defaults(func=cmd_render)
+    s = sub.add_parser("serve")
+    s.add_argument("--viz-dir", default=OUT_DEFAULT)
+    s.add_argument("--host", default="127.0.0.1")
+    s.add_argument("--port", type=int, default=7862)
+    s.add_argument("--share", action="store_true")
+    s.set_defaults(func=cmd_serve)
+    a = p.parse_args()
+    a.func(a)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/seg_smoke.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Smoke-test: run several ultralytics segmentation backends in no-prompt
+'everything' mode on one real openvid frame. Report #masks + latency so we can
+pick which variants are worth a full render. aarch64-safe (PyAV decode, no cv2)."""
+import json
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+
+DATA = Path("/mnt/lustre/vlm-s4duan/openvid_1m")
+WEIGHTS = Path("/mnt/lustre/vlm-s4duan/models/seg")
+WEIGHTS.mkdir(parents=True, exist_ok=True)
+
+# (display name, weight file, loader class)
+MODELS = [
+    ("fastsam-s", "FastSAM-s.pt", "FastSAM"),
+    ("fastsam-x", "FastSAM-x.pt", "FastSAM"),
+    ("mobilesam", "mobile_sam.pt", "SAM"),
+    ("sam-b", "sam_b.pt", "SAM"),
+    ("sam2-t", "sam2_t.pt", "SAM"),
+    ("sam2-b", "sam2_b.pt", "SAM"),
+    ("sam2.1-l", "sam2.1_l.pt", "SAM"),
+]
+
+
+def read_frame0(path: str) -> np.ndarray:
+    import av
+    c = av.open(path)
+    for f in c.decode(video=0):
+        return f.to_ndarray(format="rgb24")
+    raise RuntimeError(f"no frames in {path}")
+
+
+def main() -> None:
+    from ultralytics import FastSAM, SAM  # noqa: F401
+    import torch
+
+    items = json.loads((DATA / "videos2caption.json").read_text())
+    item = items[0]
+    vpath = str(DATA / "clips" / item["path"]) if (DATA / "clips" / item["path"]).exists() \
+        else str(DATA / "videos" / item["path"])
+    frame0 = read_frame0(vpath)
+    H, W = frame0.shape[:2]
+    print(f"frame: {vpath}  {W}x{H}", flush=True)
+
+    loaders = {"FastSAM": FastSAM, "SAM": SAM}
+    rows = []
+    for name, wf, cls in MODELS:
+        wp = WEIGHTS / wf
+        try:
+            t0 = time.time()
+            model = loaders[cls](str(wp) if wp.exists() else wf)
+            # everything-mode: no prompts. imgsz 1024, retina full-res masks.
+            res = model(frame0, device="cuda", retina_masks=True, imgsz=1024,
+                        conf=0.4, iou=0.9, verbose=False)
+            dt = time.time() - t0
+            m = res[0].masks
+            nm = 0 if m is None else int(m.data.shape[0])
+            mh, mw = (0, 0) if m is None else tuple(m.data.shape[1:])
+            print(f"[OK] {name:10s} weights={wf:14s} masks={nm:4d} maskres={mw}x{mh} "
+                  f"time={dt:6.2f}s", flush=True)
+            rows.append((name, nm, dt))
+            # cache downloaded weight into WEIGHTS dir
+            if not wp.exists() and Path(wf).exists():
+                Path(wf).replace(wp)
+            del model
+            torch.cuda.empty_cache()
+        except Exception as e:  # noqa: BLE001
+            print(f"[FAIL] {name:10s} weights={wf:14s} -> {type(e).__name__}: {e}", flush=True)
+    print("\nsummary (name, n_masks, sec):", flush=True)
+    for r in rows:
+        print(f"  {r[0]:10s} {r[1]:4d}  {r[2]:6.2f}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/segment_tracks.py

@@ -0,0 +1,363 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stage 0d: segment frames (FastSAM) and label each CoTracker grid point by object.
+
+Adds an ``object_ids`` array ([N] int, -1 = background/none) to each tracks ``.npz``,
+so the trainer's object-coverage sampling can guarantee >=1 track per object.
+
+Each point is assigned based on its FIRST-VISIBLE frame (from CoTracker visibility):
+FastSAM runs on each unique first-visible frame, and the point is assigned the smallest
+mask containing its position at that frame. This correctly handles objects that enter
+the scene after frame 0.
+
+Run on a GPU node (FastSAM is light). Idempotent (skips npz that already have object_ids)::
+
+    srun --jobid=<job> --overlap --ntasks=1 env CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$PWD \
+      .venv/bin/python data_pipeline/segment_tracks.py --data-dir <dataset root>
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+
+
+def read_all_frames(path: str) -> np.ndarray:
+    """Read all frames from a video, return [T,H,W,3] uint8."""
+    try:
+        from decord import VideoReader, cpu
+        vr = VideoReader(path, ctx=cpu(0))
+        return vr.get_batch(list(range(len(vr)))).asnumpy()
+    except Exception:  # noqa: BLE001
+        import av
+        c = av.open(path)
+        return np.stack([f.to_ndarray(format="rgb24") for f in c.decode(video=0)])
+
+
+def _colors(n: int) -> np.ndarray:
+    import colorsys
+    return np.array(
+        [[int(255 * c) for c in colorsys.hsv_to_rgb((i * 0.61803) % 1.0, 0.65, 1.0)]
+         for i in range(max(1, n))],
+        np.uint8,
+    )
+
+
+def render_viz(frames: np.ndarray, tracks: np.ndarray, vis: np.ndarray,
+               object_ids: np.ndarray, out_path: Path, fps: int = 24) -> None:
+    import imageio.v2 as imageio
+    from fastvideo.train.callbacks.track_validation import _draw_overlay
+
+    objs = sorted(int(o) for o in np.unique(object_ids) if int(o) >= 0)
+    ocols = _colors(len(objs) + 1)
+    N = tracks.shape[1]
+    pcols = np.tile(np.array([[110, 110, 110]], np.uint8), (N, 1))
+    for oi, o in enumerate(objs):
+        pcols[object_ids == o] = ocols[oi % len(ocols)]
+
+    G = int(round(N**0.5))
+    if G * G == N:
+        k = max(1, G // 50)
+        sel = np.arange(N).reshape(G, G)[::k, ::k].reshape(-1)
+    else:
+        st = max(1, N // 1500)
+        sel = np.arange(0, N, st)
+
+    ov = _draw_overlay(frames, tracks[:, sel].copy(), vis[:, sel], pcols[sel], 12, 2, 0.5)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    imageio.mimsave(str(out_path), ov, fps=fps, macro_block_size=1)
+
+
+def render_seg_viz(frame: np.ndarray, masks: np.ndarray, out_path: Path) -> None:
+    import imageio.v2 as imageio
+    img = frame.copy()
+    if masks.shape[0] > 0:
+        colors = _colors(masks.shape[0])
+        for i, mask in enumerate(masks):
+            img[mask] = (img[mask] * 0.5 + colors[i % len(colors)] * 0.5).astype(np.uint8)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    imageio.imwrite(str(out_path), img)
+
+
+def object_ids_for_points(masks: np.ndarray, pts_xy: np.ndarray, H: int, W: int) -> np.ndarray:
+    """masks [M,H,W] bool, pts_xy [N,2] px -> object id per point (-1 none), smallest mask wins."""
+    N = pts_xy.shape[0]
+    oid = np.full(N, -1, np.int64)
+    if masks.size == 0:
+        return oid
+    areas = masks.reshape(masks.shape[0], -1).sum(1)  # [M]
+    order = np.argsort(areas)  # smallest first -> assign, larger won't overwrite
+    xi = np.clip(pts_xy[:, 0].round().astype(int), 0, W - 1)
+    yi = np.clip(pts_xy[:, 1].round().astype(int), 0, H - 1)
+    assigned = np.zeros(N, bool)
+    for m in order:
+        inside = masks[m][yi, xi] & (~assigned)
+        oid[inside] = int(m)
+        assigned |= inside
+    return oid
+
+
+def extract_masks(result, H: int, W: int, min_area_frac: float, max_masks: int) -> np.ndarray:
+    """Pull masks out of a single FastSAM Result, resize if needed, and apply filtering."""
+    masks = np.zeros((0, H, W), bool)
+    if result is not None and result.masks is not None:
+        masks = result.masks.data.cpu().numpy().astype(bool)
+        if masks.shape[0] and masks.shape[1:] != (H, W):
+            import cv2
+            masks = np.stack([
+                cv2.resize(m.astype(np.uint8), (W, H), interpolation=cv2.INTER_NEAREST).astype(bool)
+                for m in masks
+            ])
+    if masks.shape[0] and (min_area_frac > 0 or max_masks):
+        areas = masks.reshape(masks.shape[0], -1).sum(1).astype(np.float64)
+        if min_area_frac > 0:
+            keep = (areas / float(H * W)) >= min_area_frac
+            masks, areas = masks[keep], areas[keep]
+        if max_masks and masks.shape[0] > max_masks:
+            masks = masks[np.argsort(-areas)[:max_masks]]
+    return masks
+
+
+def masks_for_frames(model, frames: np.ndarray, frame_ts: list[int], cache: dict[int, np.ndarray],
+                     args) -> dict[int, np.ndarray]:
+    """Segment the requested frames in batched FastSAM forwards, filling/reusing `cache`."""
+    todo = [t for t in frame_ts if t not in cache]
+    for s in range(0, len(todo), max(1, args.sam_batch)):
+        chunk = todo[s:s + max(1, args.sam_batch)]
+        res = model([frames[t] for t in chunk], device=args.device, retina_masks=True,
+                    imgsz=args.imgsz, conf=args.conf, iou=args.iou, verbose=False)
+        for t, r in zip(chunk, res):
+            cache[t] = extract_masks(r, frames.shape[1], frames.shape[2],
+                                     args.min_area_frac, args.max_masks)
+    return cache
+
+
+def assign_object_ids_multiframe(
+    frame_masks: dict[int, np.ndarray],
+    first_visible: np.ndarray,
+    tracks: np.ndarray,
+    H: int,
+    W: int,
+) -> np.ndarray:
+    """Assign globally-unique object IDs using each point's first-visible frame.
+
+    frame_masks: {frame_idx: masks [M,H,W] bool}
+    first_visible: [N] int, -1 = never visible
+    tracks: [T,N,2] px
+    Returns object_ids [N] int64, -1 = background/never visible.
+    """
+    N = first_visible.shape[0]
+    oid = np.full(N, -1, np.int64)
+    global_offset = 0
+    for frame_t, masks in sorted(frame_masks.items()):
+        point_sel = first_visible == frame_t
+        if point_sel.any() and masks.shape[0] > 0:
+            pts_xy = tracks[frame_t, point_sel]
+            local_oid = object_ids_for_points(masks, pts_xy, H, W)
+            oid[point_sel] = np.where(local_oid >= 0, local_oid + global_offset, -1)
+        global_offset += masks.shape[0]
+    return oid
+
+
+def segment_tracks_arrays(tracks: np.ndarray, vis: np.ndarray, H: int, W: int, get_masks,
+                          vis_override_every: int, verbose: bool = False
+                          ) -> tuple[np.ndarray, int, np.ndarray, np.ndarray, int, list[int]]:
+    """Core of Stage 4: object IDs from first-visible frames, vis override, low-rank weights.
+
+    Shared by segment_tracks.py and extract_tracks.py --segment (fused mode) so both paths
+    produce identical results. ``get_masks(frame_ts)`` must return {frame_t: masks [M,H,W] bool}.
+    tracks: [T,N,2] px. vis: [T,N] bool, updated in place by the override sweep.
+    Returns (object_ids, n_objects, vis, track_weights, n_overrides, unique_frames).
+    """
+    T = tracks.shape[0]
+    # Per-point first-visible frame; -1 for points CoTracker never marks visible
+    ever_visible = vis.any(axis=0)                                          # [N]
+    first_visible = np.where(ever_visible, np.argmax(vis, axis=0), -1)     # [N]
+    unique_frames = sorted(set(first_visible[ever_visible].tolist()))
+
+    if verbose:
+        fv_counts = {f: int((first_visible == f).sum()) for f in unique_frames}
+        print(f"  [seg] first_visible frames: {fv_counts}", flush=True)
+
+    # FastSAM masks for each unique first-visible frame
+    frame_masks = get_masks(unique_frames)
+    if verbose:
+        for frame_t, masks in frame_masks.items():
+            areas = masks.reshape(masks.shape[0], -1).sum(1).tolist() if masks.shape[0] else []
+            print(f"  [seg] frame {frame_t}: {masks.shape[0]} masks, areas={[int(a) for a in areas]}", flush=True)
+
+    oid = assign_object_ids_multiframe(frame_masks, first_visible, tracks, H, W)
+    n_objects = int(np.unique(oid[oid >= 0]).shape[0]) if (oid >= 0).any() else 0
+
+    if verbose:
+        for frame_t in unique_frames:
+            pt_sel = first_visible == frame_t
+            labeled = int((oid[pt_sel] >= 0).sum())
+            print(f"  [seg] frame {frame_t}: {pt_sel.sum()} pts, {labeled} got oid>=0", flush=True)
+
+    # Vis override: run FastSAM every N frames and set vis=True for object points
+    # that fall inside any mask — fixes CoTracker vis=0 on edge-of-frame objects.
+    n_overrides = 0
+    if vis_override_every > 0 and (oid >= 0).any():
+        object_pts = np.where(oid >= 0)[0]
+        override_ts = list(range(0, T, vis_override_every))
+        override_masks = get_masks(override_ts)
+        for frame_t in override_ts:
+            masks = override_masks[frame_t]
+            if masks.shape[0] == 0:
+                continue
+            pts = tracks[frame_t, object_pts]                              # [K,2]
+            xi = np.clip(pts[:, 0].round().astype(int), 0, W - 1)
+            yi = np.clip(pts[:, 1].round().astype(int), 0, H - 1)
+            in_any_mask = masks[:, yi, xi].any(axis=0)                    # [K] bool
+            newly_visible = in_any_mask & ~vis[frame_t, object_pts]
+            vis[frame_t, object_pts] |= in_any_mask
+            n_overrides += int(newly_visible.sum())
+
+        # Fill gaps between True frames caused by the sampling interval.
+        # If vis is True at frame T and True again at T+k (k <= override_every),
+        # the frames in between should also be True — the object didn't disappear.
+        V = vis[:, object_pts]                                             # [T,K]
+        t_idx = np.arange(T)[:, None]
+        prev = np.maximum.accumulate(np.where(V, t_idx, -1), axis=0)
+        nxt = np.minimum.accumulate(np.where(V, t_idx, 2 * T)[::-1], axis=0)[::-1]
+        fill = (~V) & (prev >= 0) & (nxt < T) & ((nxt - prev) <= vis_override_every)
+        vis[:, object_pts] = V | fill
+
+    weights = lowrank_track_weights(tracks)
+    return oid.astype(np.int64), n_objects, vis, weights, n_overrides, unique_frames
+
+
+def lowrank_track_weights(tracks: np.ndarray, rank: int = 3, pct: float = 97.0) -> np.ndarray:
+    """Per-point sampling weight in [0,1] = percentile-normalized low-rank motion residual.
+
+    Stack per-point displacement [N, 2T], subtract the mean trajectory + top-`rank` shared
+    SVD modes (camera / dominant scene motion), and take the residual norm. A point is heavy
+    iff it moves *uniquely* relative to all other points (independent object motion), not just
+    a lot -- so on egocentric/moving-camera clips the head-motion background is down-weighted.
+    """
+    T, N, _ = tracks.shape
+    D = (tracks - tracks[0:1]).transpose(1, 0, 2).reshape(N, 2 * T).astype(np.float64)
+    Dc = D - D.mean(0, keepdims=True)
+    U, S, Vt = np.linalg.svd(Dc, full_matrices=False)
+    r = int(min(rank, S.shape[0]))
+    resid = np.sqrt(((Dc - (U[:, :r] * S[:r]) @ Vt[:r])**2).sum(1))
+    lo, hi = float(resid.min()), float(np.percentile(resid, pct))
+    w = np.zeros_like(resid) if hi <= lo else np.clip((resid - lo) / (hi - lo), 0.0, 1.0)
+    return w.astype(np.float32)
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--data-dir", type=Path, required=True)
+    p.add_argument("--videos-subdir", type=str, default="videos")
+    p.add_argument("--tracks-subdir", type=str, default="tracks")
+    p.add_argument("--manifest", type=str, default="videos2caption.json")
+    p.add_argument("--model", type=str, default="sam2.1_b.pt",
+                   help="ultralytics weight; FastSAM-*.pt use the FastSAM loader, everything else "
+                        "(sam2.1_b.pt, sam2.1_l.pt, sam_b.pt, mobile_sam.pt, ...) use the SAM loader")
+    p.add_argument("--weights-dir", type=str, default="/mnt/lustre/vlm-s4duan/models/seg",
+                   help="dir holding cached weights; falls back to auto-download if missing")
+    p.add_argument("--device", type=str, default="cuda")
+    p.add_argument("--imgsz", type=int, default=1024)
+    p.add_argument("--conf", type=float, default=0.4)
+    p.add_argument("--iou", type=float, default=0.9)
+    p.add_argument("--sam-batch", type=int, default=16,
+                   help="Frames per batched FastSAM forward.")
+    p.add_argument("--min-area-frac", type=float, default=0.0,
+                   help="drop masks smaller than this fraction of the frame")
+    p.add_argument("--max-masks", type=int, default=0, help="keep only the N largest masks (0 = keep all)")
+    p.add_argument("--limit", type=int, default=None)
+    p.add_argument("--index", type=int, nargs="+", default=None, metavar="IDX",
+                   help="Only process videos at these manifest indices (e.g. --index 4 7 12).")
+    p.add_argument("--rank", type=int, default=0, help="GPU rank for sharding (0-indexed).")
+    p.add_argument("--world-size", type=int, default=1, help="Total number of parallel processes.")
+    p.add_argument("--force", action="store_true", help="re-run even if object_ids already present")
+    p.add_argument("--vis-override-every", type=int, default=0,
+                   help="Run FastSAM every N frames and set vis=True for object points inside masks. "
+                        "Fixes CoTracker vis=0 on edge-of-frame objects. 0 = disabled.")
+    p.add_argument("--viz", action="store_true", help="Render a track-overlay mp4 after each video.")
+    p.add_argument("--viz-dir", type=str, default=None,
+                   help="Output directory for viz mp4s (default: <data-dir>/viz).")
+    p.add_argument("--verbose", action="store_true", help="Print per-frame debug info.")
+    args = p.parse_args()
+    if args.viz and args.viz_dir is None:
+        args.viz_dir = str(args.data_dir / "viz")
+
+    # FastSAM-*.pt -> FastSAM loader; everything else (SAM/SAM2/SAM2.1/MobileSAM) -> SAM loader.
+    # Both expose the same no-prompt "everything" call used below.
+    from ultralytics import FastSAM, SAM
+    wp = Path(args.weights_dir) / args.model
+    weight = str(wp) if wp.exists() else args.model
+    model = (FastSAM if Path(args.model).name.lower().startswith("fastsam") else SAM)(weight)
+
+    manifest_path = args.data_dir / args.manifest
+    items = json.loads(manifest_path.read_text()) if manifest_path.exists() else []
+    if args.index is not None:
+        items = [items[i] for i in args.index if i < len(items)]
+    if args.limit:
+        items = items[:args.limit]
+    if args.world_size > 1:
+        items = items[args.rank::args.world_size]
+    n_ok = 0
+    for k, item in enumerate(items, 1):
+        vpath = args.data_dir / args.videos_subdir / item["path"]
+        npz_path = Path(item.get("points_path") or (args.data_dir / args.tracks_subdir / f"{vpath.stem}.npz"))
+        if not npz_path.exists():
+            print(f"[seg] [{k}/{len(items)}] {vpath.name}: no npz, skip", flush=True)
+            continue
+        d = dict(np.load(npz_path))
+        if not args.force and "object_ids" in d and "track_weights" in d:
+            n_ok += 1
+            continue
+
+        tracks = d["tracks"].astype(np.float32)   # [T,N,2] px
+        vis = d["visibility"].astype(bool)         # [T,N]
+        H, W = int(d["height"]), int(d["width"])
+
+        # Decode the video once; frames are reused for segmentation, vis override, and viz.
+        frames = read_all_frames(str(vpath))
+        mask_cache: dict[int, np.ndarray] = {}
+
+        def get_masks(frame_ts: list[int], _frames=frames, _cache=mask_cache) -> dict[int, np.ndarray]:
+            masks_for_frames(model, _frames, frame_ts, _cache, args)
+            return {t: _cache[t] for t in frame_ts}
+
+        oid, n_objects, vis, weights, n_overrides, unique_frames = segment_tracks_arrays(
+            tracks, vis, H, W, get_masks, args.vis_override_every, verbose=args.verbose)
+
+        if args.vis_override_every > 0:
+            d["visibility"] = vis
+        d["object_ids"] = oid
+        d["n_objects"] = np.int64(n_objects)
+        d["track_weights"] = weights
+        tmp = npz_path.with_suffix(".tmp.npz")
+        np.savez(tmp, **d)
+        tmp.replace(npz_path)
+        n_ok += 1
+        if args.viz:
+            stem_dir = Path(args.viz_dir) / vpath.stem
+            stem_dir.mkdir(parents=True, exist_ok=True)
+            render_viz(frames[:tracks.shape[0]], tracks, vis, oid, stem_dir / "tracks.mp4",
+                       fps=int(item.get("fps", 24)))
+            # for label, fidx in [("000", 0), ("mid", T_v // 2), ("last", T_v - 1)]:
+            #     frame = frames[fidx]
+            #     res = model(frame, device=args.device, retina_masks=True,
+            #                 imgsz=args.imgsz, conf=args.conf, iou=args.iou, verbose=False)
+            #     fmasks = extract_masks(res, H, W, args.min_area_frac, args.max_masks)
+            #     render_seg_viz(frame, fmasks, stem_dir / f"seg_frame{label}.jpg")
+            print(f"  viz -> {stem_dir}/", flush=True)
+
+        cov = int((oid >= 0).sum())
+        override_str = f", {n_overrides} vis overrides" if args.vis_override_every > 0 else ""
+        print(
+            f"[seg] [{k}/{len(items)}] {vpath.name}: {n_objects} objs across {len(unique_frames)} frames, "
+            f"{cov}/{oid.shape[0]} grid pts labeled{override_str}, "
+            f"w[mean={d['track_weights'].mean():.3f} >0.5={(d['track_weights'] > 0.5).mean():.2f}]",
+            flush=True)
+    print(f"[seg] done; {n_ok}/{len(items)} npz have object_ids", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/segment_tracks_mp.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Data-parallel frame-0 segmentation (SAM2.1-b+ everything-mode) for OpenVid-1M.
+
+Mirrors extract_tracks_mp.py: each Slurm task segments manifest[shard::num_shards]
+on one pinned GPU and writes object_ids + track_weights into the EXISTING tracks
+.npz (adds keys, keeps tracks/visibility). Idempotent — skips npz that already have
+both keys, so it is resumable.
+
+Sharding + GPU pinning from Slurm env:
+  shard      = SLURM_PROCID     (0..num_shards-1)
+  num_shards = SLURM_NTASKS
+  gpu        = SLURM_LOCALID % gpus_per_node
+
+aarch64-safe: PyAV frame-0 decode, torch mask resize (no cv2/decord).
+
+    srun --ntasks-per-node=<gpus*procs_per_gpu> ... \
+      python data_pipeline/segment_tracks_mp.py --data-dir <root> --videos-subdir clips
+"""
+from __future__ import annotations
+import argparse
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(__file__))
+from segment_tracks import object_ids_for_points, lowrank_track_weights  # noqa: E402
+
+
+def read_frame0(path: str) -> np.ndarray:
+    import av
+    c = av.open(str(path))
+    try:
+        for f in c.decode(video=0):
+            return f.to_ndarray(format="rgb24")
+    finally:
+        c.close()
+    raise RuntimeError(f"no frames in {path}")
+
+
+def extract_masks(res, H: int, W: int) -> np.ndarray:
+    if not res or res[0].masks is None:
+        return np.zeros((0, H, W), bool)
+    import torch
+    m = res[0].masks.data
+    if m.shape[1:] != (H, W):
+        m = torch.nn.functional.interpolate(m[None].float(), size=(H, W), mode="nearest")[0]
+    return m.cpu().numpy().astype(bool)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--data-dir", required=True)
+    ap.add_argument("--videos-subdir", default="clips", help="frame-0 source; MUST match track coords (720p clips)")
+    ap.add_argument("--tracks-subdir", default="tracks")
+    ap.add_argument("--manifest", default="videos2caption.json")
+    ap.add_argument("--model", default="sam2.1_b.pt")
+    ap.add_argument("--weights-dir", default="/mnt/lustre/vlm-s4duan/models/seg")
+    ap.add_argument("--imgsz", type=int, default=1024)
+    ap.add_argument("--conf", type=float, default=0.4)
+    ap.add_argument("--iou", type=float, default=0.9)
+    ap.add_argument("--min-area-frac", type=float, default=0.0)
+    ap.add_argument("--num-shards", type=int, default=int(os.environ.get("SLURM_NTASKS", "1")))
+    ap.add_argument("--shard", type=int, default=int(os.environ.get("SLURM_PROCID", "0")))
+    ap.add_argument("--gpus-per-node", type=int, default=4)
+    ap.add_argument("--limit", type=int, default=None)
+    a = ap.parse_args()
+
+    import torch
+    local = int(os.environ.get("SLURM_LOCALID", str(a.shard)))
+    gpu = local % a.gpus_per_node
+    torch.cuda.set_device(gpu)
+    device = f"cuda:{gpu}"
+
+    from ultralytics import FastSAM, SAM
+    wp = Path(a.weights_dir) / a.model
+    weight = str(wp) if wp.exists() else a.model
+    model = (FastSAM if Path(a.model).name.lower().startswith("fastsam") else SAM)(weight)
+
+    data = Path(a.data_dir)
+    items = json.loads((data / a.manifest).read_text())
+    if a.limit:
+        items = items[:a.limit]
+    mine = items[a.shard::a.num_shards]
+    print(f"[seg {a.shard}/{a.num_shards}] gpu={gpu} localid={local} items={len(mine)}", flush=True)
+
+    t0 = time.time()
+    done = skip = notrk = err = 0
+    for it in mine:
+        stem = Path(it["path"]).stem
+        npz_path = Path(it.get("points_path") or (data / a.tracks_subdir / f"{stem}.npz"))
+        if not npz_path.exists():
+            notrk += 1
+            continue
+        try:
+            d = dict(np.load(npz_path))
+            if "object_ids" in d and "track_weights" in d:
+                skip += 1
+                continue
+            vpath = data / a.videos_subdir / it["path"]
+            frame0 = read_frame0(str(vpath))
+            H, W = frame0.shape[0], frame0.shape[1]
+            res = model(frame0, device=device, retina_masks=True, imgsz=a.imgsz,
+                        conf=a.conf, iou=a.iou, verbose=False)
+            masks = extract_masks(res, H, W)
+            if masks.shape[0] and a.min_area_frac > 0:
+                areas = masks.reshape(masks.shape[0], -1).sum(1).astype(np.float64)
+                masks = masks[(areas / float(H * W)) >= a.min_area_frac]
+            tracks = d["tracks"].astype(np.float32)
+            oid = object_ids_for_points(masks, tracks[0], H, W)
+            d["object_ids"] = oid.astype(np.int64)
+            d["n_objects"] = np.int64(masks.shape[0])
+            d["track_weights"] = lowrank_track_weights(tracks)
+            tmp = npz_path.with_suffix(".seg.tmp.npz")
+            np.savez(tmp, **d)
+            tmp.replace(npz_path)
+            done += 1
+            if done % 200 == 0:
+                r = done / (time.time() - t0)
+                print(f"[seg {a.shard}] {done} done ({r:.2f}/s), skip={skip} notrk={notrk} err={err}", flush=True)
+        except Exception as e:  # noqa: BLE001
+            err += 1
+            print(f"[seg {a.shard}] ERR {stem}: {repr(e)[:120]}", flush=True)
+    dt = time.time() - t0
+    print(f"[seg {a.shard}] DONE done={done} skip={skip} notrk={notrk} err={err} in {dt:.1f}s "
+          f"-> {done / dt if dt > 0 else 0:.3f} clip/s", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/segment_viz.py

@@ -0,0 +1,346 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Visualize / sweep the segmentation pipeline: FastSAM masks + chosen per-object
+points + CoTracker tracks overlaid on the source video.
+
+``render`` (GPU) sweeps a set of FastSAM configs over the clips and writes one
+SAM-panel PNG + one track-overlay mp4 per (config, clip). ``serve`` (no GPU) is a
+gradio gallery with config/clip filters so you can scroll and compare any combo.
+
+    srun ... env CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$PWD .venv/bin/python data_pipeline/segment_viz.py render \
+        --data-dir <dataset root> --out <viz dir> --limit 8
+    .venv/bin/python data_pipeline/segment_viz.py serve --viz-dir <viz dir> --share
+
+Configs default to DEFAULT_SWEEP (below); override with --configs '<json>' or
+--configs-json <file>. Each config: {name, conf, iou, imgsz, min_area_frac, max_masks}.
+``min_area_frac`` drops masks smaller than that fraction of the frame; ``max_masks``
+keeps only the N largest — both fight FastSAM over-segmentation.
+"""
+# NOTE: intentionally no ``from __future__ import annotations`` — gradio needs the
+# real ``gr.SelectData`` annotation object on the click handler to inject the event.
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(__file__))
+
+# name, FastSAM conf/iou/imgsz, then two over-segmentation knobs:
+#   min_area_frac: drop masks whose area < this fraction of the frame (0 = keep all)
+#   max_masks:     keep only the N largest masks (0 = keep all)
+DEFAULT_SWEEP = [
+    {
+        "name": "baseline",
+        "conf": 0.4,
+        "iou": 0.9,
+        "imgsz": 1024,
+        "min_area_frac": 0.0,
+        "max_masks": 0
+    },
+    {
+        "name": "conf0.6",
+        "conf": 0.6,
+        "iou": 0.9,
+        "imgsz": 1024,
+        "min_area_frac": 0.0,
+        "max_masks": 0
+    },
+    {
+        "name": "conf0.75",
+        "conf": 0.75,
+        "iou": 0.9,
+        "imgsz": 1024,
+        "min_area_frac": 0.0,
+        "max_masks": 0
+    },
+    {
+        "name": "iou0.6",
+        "conf": 0.4,
+        "iou": 0.6,
+        "imgsz": 1024,
+        "min_area_frac": 0.0,
+        "max_masks": 0
+    },
+    {
+        "name": "areafloor",
+        "conf": 0.4,
+        "iou": 0.9,
+        "imgsz": 1024,
+        "min_area_frac": 0.006,
+        "max_masks": 0
+    },
+    {
+        "name": "clean",
+        "conf": 0.6,
+        "iou": 0.7,
+        "imgsz": 1024,
+        "min_area_frac": 0.004,
+        "max_masks": 25
+    },
+    {
+        "name": "img1536",
+        "conf": 0.5,
+        "iou": 0.8,
+        "imgsz": 1536,
+        "min_area_frac": 0.003,
+        "max_masks": 0
+    },
+]
+
+
+def _colors(n: int) -> np.ndarray:
+    import colorsys
+    return np.array([[int(255 * c) for c in colorsys.hsv_to_rgb((i * 0.61803) % 1.0, 0.65, 1.0)]
+                     for i in range(max(1, n))], np.uint8)
+
+
+def _read_frames(path: str) -> np.ndarray:
+    try:
+        from decord import VideoReader, cpu
+        vr = VideoReader(path, ctx=cpu(0))
+        return vr.get_batch(list(range(len(vr)))).asnumpy()
+    except Exception:  # noqa: BLE001
+        import av
+        c = av.open(path)
+        return np.stack([f.to_ndarray(format="rgb24") for f in c.decode(video=0)])
+
+
+def _load_configs(args: argparse.Namespace) -> list[dict]:
+    if args.configs:
+        raw = json.loads(args.configs)
+    elif args.configs_json:
+        raw = json.loads(Path(args.configs_json).read_text())
+    else:
+        raw = DEFAULT_SWEEP
+    cfgs = []
+    for i, c in enumerate(raw):
+        cfgs.append({
+            "name": str(c.get("name", f"cfg{i}")),
+            "conf": float(c.get("conf", 0.4)),
+            "iou": float(c.get("iou", 0.9)),
+            "imgsz": int(c.get("imgsz", 1024)),
+            "min_area_frac": float(c.get("min_area_frac", 0.0)),
+            "max_masks": int(c.get("max_masks", 0)),
+        })
+    return cfgs
+
+
+def _extract_masks(res, H: int, W: int) -> np.ndarray:
+    if not res or res[0].masks is None:
+        return np.zeros((0, H, W), bool)
+    masks = res[0].masks.data.cpu().numpy().astype(bool)  # [M,h,w]
+    if masks.shape[0] and masks.shape[1:] != (H, W):
+        import cv2
+        masks = np.stack(
+            [cv2.resize(m.astype(np.uint8), (W, H), interpolation=cv2.INTER_NEAREST).astype(bool) for m in masks])
+    return masks
+
+
+def filter_masks(masks: np.ndarray, min_area_frac: float, max_masks: int) -> np.ndarray:
+    """Drop tiny masks (< min_area_frac of frame) and keep only the N largest."""
+    if masks.shape[0] == 0:
+        return masks
+    H, W = masks.shape[1], masks.shape[2]
+    areas = masks.reshape(masks.shape[0], -1).sum(1).astype(np.float64)
+    if min_area_frac > 0:
+        keep = (areas / float(H * W)) >= min_area_frac
+        masks, areas = masks[keep], areas[keep]
+    if max_masks and masks.shape[0] > max_masks:
+        masks = masks[np.argsort(-areas)[:max_masks]]
+    return masks
+
+
+def cmd_render(args: argparse.Namespace) -> None:
+    from PIL import Image, ImageDraw
+    from ultralytics import FastSAM
+    from fastvideo.train.callbacks.track_validation import _draw_overlay
+    from segment_tracks import assign_object_ids_multiframe, extract_masks
+    import imageio.v2 as imageio
+
+    configs = _load_configs(args)
+    model = FastSAM(args.model)
+    data = Path(args.data_dir)
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    items = json.loads((data / "videos2caption.json").read_text())[:args.limit]
+    print(f"[viz] {len(items)} clips x {len(configs)} configs = {len(items) * len(configs)} panels", flush=True)
+
+    entries: list[dict] = []
+    clip_ids: list[str] = []
+    for k, it in enumerate(items, 1):
+        stem = Path(it["path"]).stem
+        clip_ids.append(stem)
+        vpath = str(data / "videos" / it["path"])
+        npz = it.get("points_path") or str(data / "tracks" / f"{stem}.npz")
+        frames = _read_frames(vpath)
+        H, W = frames.shape[1], frames.shape[2]
+        d = np.load(npz)
+        tracks = d["tracks"].astype(np.float32)[:frames.shape[0]]  # [T,N,2] px
+        vis = d["visibility"].astype(np.float32)[:frames.shape[0]]  # [T,N]
+        disp = np.sqrt(((tracks - tracks[0:1])**2).sum(-1)).max(0)  # [N] max displacement
+
+        ever_visible = vis.astype(bool).any(axis=0)                                      # [N]
+        first_visible = np.where(ever_visible, np.argmax(vis.astype(bool), axis=0), -1)  # [N]
+        unique_frames = sorted(set(first_visible[ever_visible].tolist()))
+
+        for cfg in configs:
+            # Run FastSAM on each unique first-visible frame for this config.
+            frame_masks: dict[int, np.ndarray] = {}
+            for frame_t in unique_frames:
+                res = model(frames[frame_t],
+                            device=args.device,
+                            retina_masks=True,
+                            imgsz=cfg["imgsz"],
+                            conf=cfg["conf"],
+                            iou=cfg["iou"],
+                            verbose=False)
+                frame_masks[frame_t] = extract_masks(res, H, W, cfg["min_area_frac"], cfg["max_masks"])
+
+            oid = assign_object_ids_multiframe(frame_masks, first_visible, tracks, H, W)
+            masks = frame_masks.get(0, np.zeros((0, H, W), bool))  # frame-0 masks for SAM panel
+            objs = sorted(int(o) for o in np.unique(oid) if int(o) >= 0)
+
+            # (1) SAM panel: colored masks + chosen point (highest-motion track) per object
+            mcols = _colors(len(masks) + 1)
+            base = frames[0].astype(np.float32)
+            for mi, m in enumerate(masks):
+                base[m] = 0.55 * base[m] + 0.45 * mcols[mi % len(mcols)][None].astype(np.float32)
+            img = Image.fromarray(base.clip(0, 255).astype(np.uint8))
+            draw = ImageDraw.Draw(img)
+            for o in objs:
+                idx = np.where(oid == o)[0]
+                pick = idx[np.argmax(disp[idx])]
+                x, y = float(tracks[0, pick, 0]), float(tracks[0, pick, 1])
+                draw.ellipse([x - 5, y - 5, x + 5, y + 5], fill=(255, 255, 255), outline=(0, 0, 0))
+            sam_name = f"{cfg['name']}__{stem}_sam.png"
+            img.save(str(out / sam_name))
+
+            # (2) CoTracker tracks over the video, colored by object (background gray)
+            trk_name = ""
+            if not args.no_tracks:
+                ocols = _colors(len(objs) + 1)
+                pcols = np.tile(np.array([[110, 110, 110]], np.uint8), (tracks.shape[1], 1))
+                for oi, o in enumerate(objs):
+                    pcols[oid == o] = ocols[oi % len(ocols)]
+                # grid-aware subsample: a flat row-major stride staggers columns and
+                # looks like half the grid; stride rows AND cols equally instead so the
+                # true (e.g. 50x50) grid stays visible and aligned.
+                N = tracks.shape[1]
+                G = int(round(N**0.5))
+                if G * G == N:
+                    k = max(1, G // 50)  # keep full grid up to 50x50
+                    sel = np.arange(N).reshape(G, G)[::k, ::k].reshape(-1)
+                else:
+                    st = max(1, N // 1500)
+                    sel = np.arange(0, N, st)
+                ov = _draw_overlay(frames, tracks[:, sel].copy(), vis[:, sel], pcols[sel], 12, 2, 0.5)
+                trk_name = f"{cfg['name']}__{stem}_tracks.mp4"
+                imageio.mimsave(str(out / trk_name), ov, fps=int(it.get("fps", 24)), macro_block_size=1)
+
+            entries.append({
+                "config":
+                cfg["name"],
+                "clip":
+                stem,
+                "caption": (it["cap"][0] if isinstance(it.get("cap"), list) else str(it.get("cap", "")))[:120],
+                "n_masks":
+                int(len(masks)),
+                "n_objects":
+                len(objs),
+                "n_labeled":
+                int((oid >= 0).sum()),
+                "n_points":
+                int(oid.shape[0]),
+                "sam":
+                sam_name,
+                "tracks":
+                trk_name,
+            })
+            print(
+                f"[viz] [{k}/{len(items)}] {stem} [{cfg['name']}]: "
+                f"{len(masks)} masks, {len(objs)} objs, {(oid >= 0).sum()}/{oid.shape[0]} pts labeled",
+                flush=True)
+
+    manifest = {"configs": configs, "clips": clip_ids, "entries": entries}
+    (out / "manifest.json").write_text(json.dumps(manifest, indent=2))
+    print(f"[viz] rendered {len(entries)} panels -> {out}", flush=True)
+
+
+def cmd_serve(args: argparse.Namespace) -> None:
+    import gradio as gr
+    viz = Path(args.viz_dir).resolve()
+    manifest = json.loads((viz / "manifest.json").read_text())
+    entries = manifest["entries"]
+    cfg_names = [c["name"] for c in manifest["configs"]]
+    clip_ids = manifest["clips"]
+    cfg_params = {c["name"]: c for c in manifest["configs"]}
+    # closure state: the currently-filtered entry list that the gallery reflects
+    current = {"entries": list(entries)}
+
+    def _label(e: dict) -> str:
+        return f'{e["clip"]} | {e["config"]} | {e["n_masks"]}m/{e["n_objects"]}o'
+
+    def gallery_for(cfg_sel: str, clip_sel: str):
+        es = entries
+        if cfg_sel and cfg_sel != "(all)":
+            es = [e for e in es if e["config"] == cfg_sel]
+        if clip_sel and clip_sel != "(all)":
+            es = [e for e in es if e["clip"] == clip_sel]
+        current["entries"] = es
+        return [(str(viz / e["sam"]), _label(e)) for e in es]
+
+    def show(evt: gr.SelectData):
+        e = current["entries"][evt.index]
+        tv = str(viz / e["tracks"]) if e.get("tracks") else None
+        info = {**e, "config_params": cfg_params.get(e["config"], {})}
+        return str(viz / e["sam"]), tv, json.dumps(info, indent=2)
+
+    with gr.Blocks(title="FastSAM config sweep") as demo:
+        gr.Markdown("### FastSAM config sweep — masks + chosen per-object points + CoTracker tracks\n"
+                    "Filter by **config** and/or **clip**, scroll the gallery, click a tile to inspect. "
+                    "Tile label: `clip | config | <#masks>m/<#objects-with-points>o`. "
+                    "Fewer, cleaner masks = less over-segmentation.")
+        with gr.Row():
+            cfg_dd = gr.Dropdown(["(all)"] + cfg_names, value="(all)", label="config")
+            clip_dd = gr.Dropdown(["(all)"] + clip_ids, value="(all)", label="clip")
+        with gr.Row():
+            g = gr.Gallery(value=gallery_for("(all)", "(all)"), columns=4, height=620, label="panels (click to view)")
+            with gr.Column():
+                samimg = gr.Image(label="frame-0: SAM masks + chosen points")
+                trkvid = gr.Video(label="CoTracker tracks (colored by object)")
+                meta = gr.Code(label="info", language="json")
+        cfg_dd.change(gallery_for, [cfg_dd, clip_dd], g)
+        clip_dd.change(gallery_for, [cfg_dd, clip_dd], g)
+        g.select(show, None, [samimg, trkvid, meta])
+    # allowed_paths: without this gradio blocks serving the PNGs/mp4s that live
+    # outside its app root -> the three side panels error out on click.
+    demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share, allowed_paths=[str(viz)])
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub = p.add_subparsers(dest="cmd", required=True)
+    r = sub.add_parser("render")
+    r.add_argument("--data-dir", required=True)
+    r.add_argument("--out", required=True)
+    r.add_argument("--model", default="FastSAM-s.pt")
+    r.add_argument("--device", default="cuda")
+    r.add_argument("--limit", type=int, default=8)
+    r.add_argument("--configs", default=None, help="inline JSON list of config dicts (overrides sweep)")
+    r.add_argument("--configs-json", default=None, help="path to a JSON list of config dicts")
+    r.add_argument("--no-tracks", action="store_true", help="skip the (slow) track-overlay mp4s")
+    r.set_defaults(func=cmd_render)
+    s = sub.add_parser("serve")
+    s.add_argument("--viz-dir", required=True)
+    s.add_argument("--host", default="0.0.0.0")
+    s.add_argument("--port", type=int, default=7880)
+    s.add_argument("--share", action="store_true")
+    s.set_defaults(func=cmd_serve)
+    a = p.parse_args()
+    a.func(a)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/sparse_sampling_dashboard.py

@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Interactive dashboard for the SPARSE ~1-per-object + few-background sampler.
+
+Per Yongqi's advice, our track budget for the sparse recipe is dramatically smaller
+than MotionStream's 1000-2500: keep 1 track per SAM object, then add ``num_sampled``
+extra points drawn from the rest (either uniformly at random, or weighted by the
+precomputed low-rank informativeness). The dashboard lets us eyeball whether the
+resulting subset actually covers the "action" of each clip before we commit to
+launching training runs on it.
+
+CPU only, gradio share (no GPU needed).
+
+    .venv/bin/python data_pipeline/sparse_sampling_dashboard.py \
+        --data-dir /mnt/weka/home/hao.zhang/shao/data/motion_pipeline/synthetic_toy --share
+"""
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(__file__))
+from track_informativeness import _draw_overlay, _read_frames  # noqa: E402
+
+
+def sparse_sample(oid: np.ndarray, tw: np.ndarray | None, vis0: np.ndarray, num_sampled: int, mode: str,
+                  seed: int) -> np.ndarray:
+    """One point per SAM object (weighted-by-`tw` inside each object) + ``num_sampled`` extras
+    drawn from the rest via ``mode`` in {"random", "weighted"}. Returns kept bool [N]."""
+    rng = np.random.default_rng(int(seed))
+    N = oid.shape[0]
+    valid = vis0 > 0.5
+    keep = np.zeros(N, bool)
+    # (a) 1 per object
+    for o in np.unique(oid):
+        if int(o) < 0:
+            continue
+        idx = np.where((oid == o) & valid)[0]
+        if idx.size == 0:
+            continue
+        if tw is not None and tw[idx].sum() > 0:
+            p = tw[idx] / tw[idx].sum()
+            pick = int(rng.choice(idx, p=p))
+        else:
+            pick = int(rng.choice(idx))
+        keep[pick] = True
+    # (b) num_sampled extras from the remaining valid points
+    pool = np.where((~keep) & valid)[0]
+    k = int(min(max(num_sampled, 0), pool.size))
+    if k > 0:
+        if mode == "weighted" and tw is not None and tw[pool].sum() > 0:
+            p = tw[pool] / tw[pool].sum()
+            picks = rng.choice(pool, size=k, replace=False, p=p)
+        else:
+            picks = rng.choice(pool, size=k, replace=False)
+        keep[picks] = True
+    return keep
+
+
+def _load(data_dir: Path) -> list[dict]:
+    man = json.loads((data_dir / "videos2caption.json").read_text())
+    labels_dir = data_dir / "sam_labels"
+    items = []
+    for it in man:
+        stem = Path(it["path"]).stem
+        vpath = str(data_dir / "videos" / it["path"])
+        npzp = it.get("points_path") or str(data_dir / "tracks" / f"{stem}.npz")
+        d = np.load(npzp)
+        lp = labels_dir / f"{stem}.npy"
+        labels = np.load(lp) if lp.exists() else None  # [H,W] int16, -1 = background
+        items.append({
+            "stem": stem,
+            "caption": it["cap"][0] if isinstance(it.get("cap"), list) else "",
+            "fps": int(it.get("fps", 24)),
+            "vpath": vpath,
+            "tracks": d["tracks"].astype(np.float32),
+            "vis": d["visibility"].astype(np.float32),
+            "oid": d["object_ids"].astype(np.int64) if "object_ids" in d else None,
+            "tw": d["track_weights"].astype(np.float32) if "track_weights" in d else None,
+            "labels": labels,
+        })
+    return items
+
+
+def build_ui(items: list[dict], out_dir: Path):
+    import gradio as gr
+    import imageio.v2 as imageio
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    def _render_frame(clip_idx: int, num_sampled: int, mode: str, seed: int, mask_alpha: float = 0.45):
+        it = items[int(clip_idx)]
+        frames = _read_frames(it["vpath"])
+        tracks = it["tracks"][:frames.shape[0]]
+        vis = it["vis"][:frames.shape[0]]
+        oid = it["oid"]
+        tw = it["tw"]
+        labels = it.get("labels")
+        vis0 = vis[0]
+        if oid is None:
+            return None, "No object_ids in npz — run segment_tracks first."
+        keep = sparse_sample(oid, tw, vis0, int(num_sampled), mode, int(seed))
+        n_kept = int(keep.sum())
+        n_obj = int(len(np.unique(oid[oid >= 0])))
+        # Blend SAM segmentation on top of frame 0: color each segment by the same tab20
+        # index we use for the tracked points (mask id == object id).
+        base = frames[0].astype(np.float32)
+        if labels is not None:
+            cmap = plt.get_cmap("tab20")
+            H, W = base.shape[:2]
+            overlay = np.zeros_like(base)
+            uniq = np.unique(labels)
+            for u in uniq:
+                if int(u) < 0:
+                    continue
+                col = np.array(cmap(int(u) % 20)[:3]) * 255
+                m = labels == u
+                overlay[m] = col
+            mask = (labels >= 0)[..., None]
+            base = np.where(mask, (1 - mask_alpha) * base + mask_alpha * overlay, base)
+        base = base.clip(0, 255).astype(np.uint8)
+        fig, ax = plt.subplots(1, 1, figsize=(9, 5.5))
+        ax.imshow(base)
+        x, y = tracks[0, :, 0], tracks[0, :, 1]
+        ax.scatter(x[~keep], y[~keep], c="0.5", s=3, alpha=0.35)  # dropped
+        # kept: color per unique object id (SAME palette as the mask overlay above)
+        cmap = plt.get_cmap("tab20")
+        for o in np.unique(oid[keep]):
+            mask_pts = keep & (oid == o)
+            col = "white" if int(o) < 0 else cmap(int(o) % 20)
+            ax.scatter(x[mask_pts], y[mask_pts], c=[col], s=90, edgecolors="black", linewidths=1.2)
+        ax.set_title(
+            f"{it['stem']}: {n_kept} tracks kept  ({n_obj} SAM objects, +{num_sampled} background, "
+            f"{mode}, seed {seed})",
+            fontsize=10)
+        ax.axis("off")
+        img_path = str(out_dir / f"preview_{it['stem']}.png")
+        fig.savefig(img_path, dpi=90, bbox_inches="tight")
+        plt.close(fig)
+        info = json.dumps(
+            {
+                "caption": it["caption"],
+                "n_kept": n_kept,
+                "n_objects_present": n_obj,
+                "kept_from_objects": int(len(np.unique(oid[keep][oid[keep] >= 0]))),
+                "kept_from_background": int((keep & (oid < 0)).sum()),
+            },
+            indent=2)
+        return img_path, info
+
+    def _render_video(clip_idx: int, num_sampled: int, mode: str, seed: int):
+        it = items[int(clip_idx)]
+        frames = _read_frames(it["vpath"])
+        tracks = it["tracks"][:frames.shape[0]]
+        vis = it["vis"][:frames.shape[0]]
+        oid = it["oid"]
+        tw = it["tw"]
+        keep = sparse_sample(oid, tw, vis[0], int(num_sampled), mode, int(seed))
+        kidx = np.where(keep)[0]
+        if kidx.size == 0:
+            return None
+        cmap = plt.get_cmap("tab20")
+        cols = np.zeros((kidx.size, 3), dtype=np.uint8)
+        for o in np.unique(oid[kidx]):
+            m = (oid[kidx] == o)
+            rgba = (1.0, 1.0, 1.0, 1.0) if int(o) < 0 else cmap(int(o) % 20)
+            cols[m] = (np.array(rgba[:3]) * 255).astype(np.uint8)
+        ov = _draw_overlay(frames, tracks[:, kidx].copy(), vis[:, kidx], cols, 14, 3, 0.5)
+        out = str(out_dir / f"preview_{it['stem']}_overlay.mp4")
+        imageio.mimsave(out, ov, fps=int(it.get("fps", 24)), macro_block_size=1)
+        return out
+
+    with gr.Blocks(title="Sparse track sampler dashboard") as demo:
+        gr.Markdown("### Sparse-sampling recipe: 1 track per SAM object + `num_sampled` extras\n"
+                    "Frame 0 shows all 2500 grid tracks in gray, and the kept ones colored by object. "
+                    "Set `num_sampled` = 0 for pure 1-per-object; raise it to add background context. "
+                    "Toggle `weighted` to bias the extras toward high-lowrank-informativeness points. "
+                    "The overlay video shows only the kept tracks over the whole clip.")
+        with gr.Row():
+            with gr.Column(scale=1):
+                clip_dd = gr.Dropdown([f"{i:02d} — {items[i]['stem']}" for i in range(len(items))],
+                                      value=f"00 — {items[0]['stem']}",
+                                      label="Clip",
+                                      type="index")
+                num_slider = gr.Slider(0, 200, value=20, step=1, label="num_sampled (extras)")
+                mode_radio = gr.Radio(["random", "weighted"], value="weighted", label="extras sampling mode")
+                seed_slider = gr.Slider(0, 100, value=0, step=1, label="seed")
+                info_box = gr.Code(label="info", language="json")
+                vid_btn = gr.Button("Render overlay video", variant="primary")
+            with gr.Column(scale=2):
+                frame_img = gr.Image(label="frame 0 with kept tracks", height=460)
+                overlay_vid = gr.Video(label="kept-tracks overlay video", height=460)
+
+        for src in [clip_dd, num_slider, mode_radio, seed_slider]:
+            src.change(_render_frame, [clip_dd, num_slider, mode_radio, seed_slider], [frame_img, info_box])
+        vid_btn.click(_render_video, [clip_dd, num_slider, mode_radio, seed_slider], overlay_vid)
+        demo.load(_render_frame, [clip_dd, num_slider, mode_radio, seed_slider], [frame_img, info_box])
+    return demo
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--data-dir", required=True)
+    p.add_argument("--out-dir",
+                   default="/mnt/weka/home/hao.zhang/shao/data/motion_pipeline/sparse_sampling_dashboard_out")
+    p.add_argument("--host", default="0.0.0.0")
+    p.add_argument("--port", type=int, default=7894)
+    p.add_argument("--share", action="store_true")
+    args = p.parse_args()
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    items = _load(Path(args.data_dir))
+    demo = build_ui(items, out_dir)
+    demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share, allowed_paths=[str(out_dir)])
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/split_manifest_shards.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Split a videos2caption.json manifest into NUM_SHARDS shards for data-parallel preprocess.
+
+The i2v_track preprocess reads a merge.txt (``<clips_dir>,<json>``) whose JSON is a
+*list* of dicts (json.load, NOT jsonlines) — see
+fastvideo/dataset/preprocessing_datasets.py:_load_raw_data. So each shard is itself a
+JSON array + its own merge.txt, and each shard is fed to one single-GPU v1_preprocess.py
+process.
+
+Also injects a ``duration`` field when missing: FrameSamplingStage.should_keep drops any
+row with ``duration is None`` (preprocessing_datasets.py:183), and OpenVid's
+videos2caption.json carries only num_frames/fps. duration = num_frames / fps.
+"""
+from __future__ import annotations
+import argparse
+import json
+from pathlib import Path
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--manifest", required=True, help="videos2caption.json (JSON array)")
+    ap.add_argument("--clips-dir", required=True, help="dir with the .mp4 clips (merge.txt col 1)")
+    ap.add_argument("--out-dir", required=True, help="output shards dir (shard_XXXXX/ created inside)")
+    ap.add_argument("--num-shards", type=int, required=True)
+    a = ap.parse_args()
+
+    items = json.load(open(a.manifest))
+    filled = 0
+    for it in items:
+        if it.get("duration") is None:
+            nf, fps = it.get("num_frames"), it.get("fps")
+            if nf and fps:
+                it["duration"] = float(nf) / float(fps)
+                filled += 1
+
+    out = Path(a.out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    # Contiguous shards keep each shard's clips locally clustered on disk.
+    n = a.num_shards
+    total = len(items)
+    per = (total + n - 1) // n
+    written = 0
+    for k in range(n):
+        chunk = items[k * per:(k + 1) * per]
+        sd = out / f"shard_{k:05d}"
+        sd.mkdir(parents=True, exist_ok=True)
+        mani = sd / "videos2caption.json"
+        mani.write_text(json.dumps(chunk, indent=2))
+        (sd / "merge.txt").write_text(f"{a.clips_dir},{mani}")
+        written += len(chunk)
+    print(f"[split] {total} rows ({filled} got injected duration) -> {n} shards "
+          f"(~{per}/shard) under {out}; wrote {written} rows total")
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/synthetic_tracks.py

@@ -0,0 +1,289 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Author synthetic point-track motion controls for WanTrack.
+
+A "track" is just a per-point (x, y) trajectory over time, so we can *author*
+counterfactual motion directly (no CoTracker needed -- that is only for
+*extracting* tracks from real video). These authored tracks are the test-time
+control signal: feed them to the model and check the generated motion follows.
+
+All coordinates are NORMALIZED to [0, 1] (x/width, y/height) -- the same space
+the preprocessing stores and the model trains on. Outputs:
+    tracks      float32 [T, N, 2]   normalized (x, y) per point per frame
+    visibility  float32 [T, N]      1 = active/visible, 0 = inactive/occluded
+
+Modes mirror MotionStream's control surface:
+  - ``pan``    : global translation (camera/scene pan)
+  - ``zoom``   : scale toward/away from a center (dolly in/out)
+  - ``rotate`` : rotate the grid about a center
+  - ``drag``   : move only points inside a region (the sparse "drag" handle),
+                 with a smooth spatial falloff; everything else stays put
+  - ``swirl``  : rotation whose angle decays with radius
+  - ``static`` : no motion (tests "does text+first-frame alone reproduce motion?")
+
+``select_*`` helpers turn a dense field into a SPARSE control by zeroing the
+visibility of unselected points -- this is how MotionStream lets you drag just a
+few points. (Requires a model trained with point-subsampling aug to work well.)
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+
+
+# ----------------------------------------------------------------------
+# Grid + small helpers
+# ----------------------------------------------------------------------
+def make_grid(grid_size: int = 50) -> np.ndarray:
+    """Frame-0 normalized grid positions, returned as [N, 2] (x, y) in [0, 1].
+
+    Row-major (matches CoTracker's grid and our 50x50 preprocessing), so
+    ``reshape(grid_size, grid_size, 2)`` is [row(y), col(x)].
+    """
+    ys, xs = np.meshgrid(
+        np.linspace(0.0, 1.0, grid_size, dtype=np.float32),
+        np.linspace(0.0, 1.0, grid_size, dtype=np.float32),
+        indexing="ij",
+    )
+    return np.stack([xs.reshape(-1), ys.reshape(-1)], axis=-1)  # [N, 2]
+
+
+def _ramp(num_frames: int, ease: bool = True) -> np.ndarray:
+    """Time ramp in [0, 1] over ``num_frames`` (optionally smooth ease-in-out)."""
+    t = np.linspace(0.0, 1.0, num_frames, dtype=np.float32)
+    if ease:
+        t = t * t * (3.0 - 2.0 * t)  # smoothstep
+    return t
+
+
+def _visibility_from_bounds(tracks: np.ndarray, margin: float = 0.02) -> np.ndarray:
+    """1 where a point is inside [(-m), (1+m)] in both axes, else 0."""
+    lo, hi = -margin, 1.0 + margin
+    inb = (tracks[..., 0] >= lo) & (tracks[..., 0] <= hi) & (tracks[..., 1] >= lo) & (tracks[..., 1] <= hi)
+    return inb.astype(np.float32)
+
+
+# ----------------------------------------------------------------------
+# Motion fields  ->  (tracks [T,N,2], visibility [T,N])
+# ----------------------------------------------------------------------
+def static(grid: np.ndarray, num_frames: int) -> tuple[np.ndarray, np.ndarray]:
+    tracks = np.repeat(grid[None], num_frames, axis=0)
+    return tracks, np.ones((num_frames, grid.shape[0]), np.float32)
+
+
+def pan(grid: np.ndarray, num_frames: int, dx: float, dy: float,
+        ease: bool = True) -> tuple[np.ndarray, np.ndarray]:
+    """Translate all points by (dx, dy) (normalized units) over the clip."""
+    r = _ramp(num_frames, ease)[:, None, None]
+    tracks = grid[None] + r * np.array([dx, dy], np.float32)[None, None]
+    return tracks.astype(np.float32), _visibility_from_bounds(tracks)
+
+
+def zoom(grid: np.ndarray, num_frames: int, scale: float,
+         center: tuple[float, float] = (0.5, 0.5), ease: bool = True
+         ) -> tuple[np.ndarray, np.ndarray]:
+    """Scale toward (scale<1) or away from (scale>1) ``center`` (dolly)."""
+    c = np.array(center, np.float32)[None, None]
+    r = _ramp(num_frames, ease)[:, None, None]
+    s = 1.0 + (scale - 1.0) * r
+    tracks = c + (grid[None] - c) * s
+    return tracks.astype(np.float32), _visibility_from_bounds(tracks)
+
+
+def rotate(grid: np.ndarray, num_frames: int, degrees: float,
+           center: tuple[float, float] = (0.5, 0.5), ease: bool = True
+           ) -> tuple[np.ndarray, np.ndarray]:
+    c = np.array(center, np.float32)[None]
+    r = _ramp(num_frames, ease)
+    out = np.empty((num_frames, grid.shape[0], 2), np.float32)
+    rel = grid - c
+    for t in range(num_frames):
+        a = np.deg2rad(degrees) * r[t]
+        ca, sa = np.cos(a), np.sin(a)
+        rot = np.array([[ca, -sa], [sa, ca]], np.float32)
+        out[t] = rel @ rot.T + c
+    return out, _visibility_from_bounds(out)
+
+
+def drag(grid: np.ndarray, num_frames: int, *, center: tuple[float, float],
+         dx: float, dy: float, radius: float = 0.15, falloff: str = "smooth",
+         ease: bool = True) -> tuple[np.ndarray, np.ndarray]:
+    """Move only points within ``radius`` of ``center`` by (dx, dy).
+
+    ``falloff='smooth'`` weights displacement by a cosine window of the distance
+    (handle-like drag); ``'hard'`` moves all in-radius points equally. Points
+    outside the radius stay put. Visibility stays 1 for all points (dense field
+    with a localized motion) -- use ``select_radius`` to make it *sparse*.
+    """
+    c = np.array(center, np.float32)[None]
+    d = np.linalg.norm(grid - c, axis=-1)  # [N]
+    if falloff == "smooth":
+        w = np.clip(1.0 - d / max(radius, 1e-6), 0.0, 1.0)
+        w = 0.5 - 0.5 * np.cos(np.pi * w)  # smooth 0->1
+    else:
+        w = (d <= radius).astype(np.float32)
+    r = _ramp(num_frames, ease)[:, None, None]
+    disp = (w[:, None] * np.array([dx, dy], np.float32)[None])[None]  # [1,N,2]
+    tracks = grid[None] + r * disp
+    return tracks.astype(np.float32), _visibility_from_bounds(tracks)
+
+
+def swirl(grid: np.ndarray, num_frames: int, degrees: float,
+          center: tuple[float, float] = (0.5, 0.5), radius: float = 0.5,
+          ease: bool = True) -> tuple[np.ndarray, np.ndarray]:
+    """Rotation whose angle decays linearly to 0 at ``radius`` from center."""
+    c = np.array(center, np.float32)[None]
+    rel = grid - c
+    dist = np.linalg.norm(rel, axis=-1)
+    decay = np.clip(1.0 - dist / max(radius, 1e-6), 0.0, 1.0)
+    rmp = _ramp(num_frames, ease)
+    out = np.empty((num_frames, grid.shape[0], 2), np.float32)
+    for t in range(num_frames):
+        a = np.deg2rad(degrees) * rmp[t] * decay
+        ca, sa = np.cos(a), np.sin(a)
+        x = rel[:, 0] * ca - rel[:, 1] * sa
+        y = rel[:, 0] * sa + rel[:, 1] * ca
+        out[t] = np.stack([x, y], -1) + c
+    return out, _visibility_from_bounds(out)
+
+
+# ----------------------------------------------------------------------
+# Sparsify (sparse "drag a few points" control)
+# ----------------------------------------------------------------------
+def select_radius(visibility: np.ndarray, grid: np.ndarray, *,
+                  center: tuple[float, float], radius: float) -> np.ndarray:
+    """Zero visibility outside ``radius`` of ``center`` (keep only a local handle)."""
+    c = np.array(center, np.float32)[None]
+    keep = (np.linalg.norm(grid - c, axis=-1) <= radius).astype(np.float32)  # [N]
+    return visibility * keep[None]
+
+
+def select_random(visibility: np.ndarray, k: int, seed: int = 0) -> np.ndarray:
+    """Keep only ``k`` random points active (rest visibility 0)."""
+    rng = np.random.default_rng(seed)
+    n = visibility.shape[1]
+    keep = np.zeros(n, np.float32)
+    keep[rng.choice(n, size=min(k, n), replace=False)] = 1.0
+    return visibility * keep[None]
+
+
+def select_stride(visibility: np.ndarray, grid_size: int, stride: int) -> np.ndarray:
+    """Keep a coarse sub-grid (every ``stride``-th point in both dims) active."""
+    n = visibility.shape[1]
+    mask2d = np.zeros((grid_size, grid_size), np.float32)
+    mask2d[::stride, ::stride] = 1.0
+    return visibility * mask2d.reshape(1, n)
+
+
+# ----------------------------------------------------------------------
+# Motion transfer: reuse tracks extracted from a real/other video
+# ----------------------------------------------------------------------
+def from_npz(npz_path: str, num_frames: int | None = None
+             ) -> tuple[np.ndarray, np.ndarray]:
+    """Load tracks from an ``extract_tracks.py`` npz, normalized to [0,1]."""
+    data = np.load(npz_path)
+    tr = data["tracks"].astype(np.float32)            # [T, N, 2] pixels
+    vis = data["visibility"].astype(np.float32)       # [T, N]
+    w = float(data["width"]) if "width" in data else 1.0
+    h = float(data["height"]) if "height" in data else 1.0
+    tr = tr.copy()
+    tr[..., 0] /= max(w, 1e-6)
+    tr[..., 1] /= max(h, 1e-6)
+    if num_frames is not None:
+        tr, vis = tr[:num_frames], vis[:num_frames]
+    return tr, vis
+
+
+# ----------------------------------------------------------------------
+# Convenience: named presets
+# ----------------------------------------------------------------------
+def preset(name: str, num_frames: int, grid_size: int = 50, *,
+           strength: float = 0.25) -> tuple[np.ndarray, np.ndarray]:
+    """Author a named motion preset. ``strength`` scales translations/zoom."""
+    g = make_grid(grid_size)
+    name = name.lower()
+    if name == "static":
+        return static(g, num_frames)
+    if name in ("pan_right", "pan_left", "pan_up", "pan_down"):
+        dx = {"pan_right": strength, "pan_left": -strength}.get(name, 0.0)
+        dy = {"pan_down": strength, "pan_up": -strength}.get(name, 0.0)
+        return pan(g, num_frames, dx, dy)
+    if name == "zoom_in":
+        return zoom(g, num_frames, 1.0 + strength)
+    if name == "zoom_out":
+        return zoom(g, num_frames, 1.0 - strength)
+    if name in ("rotate_cw", "rotate_ccw"):
+        return rotate(g, num_frames, 30.0 * (1 if name == "rotate_cw" else -1))
+    if name == "swirl":
+        return swirl(g, num_frames, 60.0)
+    if name == "drag_center_right":
+        return drag(g, num_frames, center=(0.5, 0.5), dx=strength, dy=0.0, radius=0.2)
+    raise ValueError(f"unknown preset {name!r}")
+
+
+PRESETS = ["static", "pan_right", "pan_left", "pan_up", "pan_down",
+           "zoom_in", "zoom_out", "rotate_cw", "rotate_ccw", "swirl",
+           "drag_center_right"]
+
+
+def to_pixel(tracks: np.ndarray, height: int, width: int) -> np.ndarray:
+    """Normalized [0,1] tracks -> pixel coords for overlay/EPE."""
+    out = tracks.copy()
+    out[..., 0] *= width
+    out[..., 1] *= height
+    return out
+
+
+# ----------------------------------------------------------------------
+# CLI: overlay a preset on a first frame (sanity check, no model/GPU)
+# ----------------------------------------------------------------------
+def _overlay_preview(frame0: np.ndarray, tracks_px: np.ndarray, vis: np.ndarray,
+                     stride: int = 3) -> np.ndarray:
+    import colorsys
+
+    from PIL import Image, ImageDraw
+    h, w, _ = frame0.shape
+    img = Image.fromarray(frame0).convert("RGB")
+    draw = ImageDraw.Draw(img)
+    n = tracks_px.shape[1]
+    gs = int(round(n ** 0.5))
+    keep = np.zeros(n, bool)
+    keep.reshape(gs, gs)[::stride, ::stride] = True
+    for pi in range(n):
+        if not keep[pi]:
+            continue
+        col = tuple(int(c * 255) for c in colorsys.hsv_to_rgb((pi % gs) / gs, 1.0, 1.0))
+        pts = [(float(tracks_px[t, pi, 0]), float(tracks_px[t, pi, 1]))
+               for t in range(tracks_px.shape[0]) if vis[t, pi] >= 0.5]
+        if len(pts) >= 2:
+            draw.line(pts, fill=col, width=1)
+        if pts:
+            x, y = pts[-1]
+            draw.ellipse([x - 2, y - 2, x + 2, y + 2], fill=col)
+    return np.asarray(img)
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description="Preview a synthetic track preset over a first frame.")
+    p.add_argument("--frame", type=Path, required=True, help="First-frame image (png/jpg).")
+    p.add_argument("--preset", type=str, default="pan_right", choices=PRESETS)
+    p.add_argument("--num-frames", type=int, default=121)
+    p.add_argument("--grid-size", type=int, default=50)
+    p.add_argument("--strength", type=float, default=0.25)
+    p.add_argument("--out", type=Path, required=True)
+    args = p.parse_args()
+
+    import imageio.v2 as imageio
+    frame0 = np.asarray(imageio.imread(args.frame))[..., :3]
+    h, w, _ = frame0.shape
+    tracks, vis = preset(args.preset, args.num_frames, args.grid_size, strength=args.strength)
+    overlay = _overlay_preview(frame0, to_pixel(tracks, h, w), vis)
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    imageio.imwrite(str(args.out), overlay)
+    print(f"[synthetic-tracks] preset={args.preset} frames={args.num_frames} "
+          f"points={tracks.shape[1]} -> {args.out}")
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/test_controllability.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Controllability test for an overfitted TrackWan model.
+
+Disentangles "did it overfit on the video content" from "did it learn to follow
+the tracks" by holding the content conditioning (first frame + text) fixed and
+varying ONLY the control tracks:
+
+  - gt        : the clip's own tracks            (reconstruction; should match GT)
+  - none      : no tracks at all                 (if motion still appears -> content overfit)
+  - pan_*/zoom/drag : authored counterfactuals   (if motion follows -> real control)
+  - swap      : another clip's tracks            (motion transfer)
+
+For each, it generates a video, overlays the control tracks, saves an mp4, and
+computes CoTracker EPE (input tracks vs tracks re-extracted from the generation).
+Low EPE on counterfactuals = genuine track control.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+import imageio.v2 as imageio
+import numpy as np
+import torch
+
+sys.path.insert(0, os.path.dirname(__file__))
+import synthetic_tracks as st  # noqa: E402
+import trackwan_infer as twi  # noqa: E402
+
+
+def _overlay(frames_thwc: np.ndarray, tracks_norm: np.ndarray, vis: np.ndarray, stride: int = 3) -> list[np.ndarray]:
+    from fastvideo.train.callbacks.track_validation import (_draw_overlay, _grid_colors, _subsample)
+    T, H, W, _ = frames_thwc.shape
+    tt = min(T, tracks_norm.shape[0])
+    fr, tr, vs = frames_thwc[:tt], tracks_norm[:tt], vis[:tt]
+    grid = int(round(tr.shape[1]**0.5))
+    tr, vs = _subsample(tr, vs, grid, stride)
+    colors = _grid_colors(grid, stride)
+    if colors.shape[0] != tr.shape[1]:
+        colors = _grid_colors(int(round(tr.shape[1]**0.5)) or 1, 1)[:tr.shape[1]]
+    trpx = tr.copy()
+    trpx[..., 0] *= W
+    trpx[..., 1] *= H
+    return _draw_overlay(fr, trpx, vs, colors, 12, 2, 0.5)
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--export", required=True, help="dcp_to_diffusers export dir")
+    p.add_argument("--yaml", default="examples/train/scenario/worldmodel/finetune_wantrack_i2v.yaml")
+    p.add_argument("--data",
+                   default="/mnt/weka/home/hao.zhang/shao/data/motion_pipeline/"
+                   "wan22_a14b_720p_24fps/preprocessed_i2v_track_funinp/combined_parquet_dataset")
+    p.add_argument("--out", required=True)
+    p.add_argument("--clips", type=int, nargs="+", default=[0, 1])
+    p.add_argument("--steps", type=int, default=30)
+    p.add_argument("--seed", type=int, default=1000)
+    p.add_argument("--fps", type=int, default=24)
+    args = p.parse_args()
+
+    os.makedirs(args.out, exist_ok=True)
+    model, tc = twi.load_trackwan(args.export, args.yaml)
+    text_len = int(tc.pipeline_config.text_encoder_configs[0].arch_config.text_len)
+    samples = twi.load_conditioning_from_parquet(args.data, args.clips, text_len)
+
+    from fastvideo.eval.metrics.motion.cotracker_epe.metric import (compute_epe, load_cotracker)
+    ct = load_cotracker(model.device)
+
+    # latent T -> pixel T
+    num_lat_t = samples[0]["first_frame_latent"].shape[2]
+    ratio = int(tc.pipeline_config.vae_config.arch_config.temporal_compression_ratio)
+    Tpx = (num_lat_t - 1) * ratio + 1
+
+    results = []
+    for ci, s in zip(args.clips, samples, strict=False):
+        gt_tracks = s["track_points"][0].numpy()[:Tpx]  # [T,N,2] normalized
+        gt_vis = s["track_visibility"][0].numpy()[:Tpx]
+        swap = samples[(args.clips.index(ci) + 1) % len(samples)]
+        swap_tracks = swap["track_points"][0].numpy()[:Tpx]
+        swap_vis = swap["track_visibility"][0].numpy()[:Tpx]
+
+        g = st.make_grid(50)
+        pan_t, pan_v = st.pan(g, Tpx, 0.25, 0.0)
+        zoom_t, zoom_v = st.zoom(g, Tpx, 1.25)
+        drag_t, drag_v = st.drag(g, Tpx, center=(0.5, 0.5), dx=0.3, dy=0.0, radius=0.2)
+        drag_v_sparse = st.select_radius(drag_v, g, center=(0.5, 0.5), radius=0.2)
+
+        controls = {
+            "gt": (gt_tracks, gt_vis),
+            "none": (None, None),
+            "pan_right": (pan_t, pan_v),
+            "zoom_in": (zoom_t, zoom_v),
+            "drag_dense": (drag_t, drag_v),
+            "drag_sparse": (drag_t, drag_v_sparse),
+            "swap": (swap_tracks, swap_vis),
+        }
+
+        # GT reference (decode the real clip)
+        ref = twi.decode_reference(model, s["vae_latent"])
+        imageio.mimsave(os.path.join(args.out, f"clip{ci}_reference.mp4"),
+                        _overlay(ref, gt_tracks, gt_vis),
+                        fps=args.fps,
+                        macro_block_size=1)
+
+        for name, (tr, vs) in controls.items():
+            tp = torch.from_numpy(tr)[None].float() if tr is not None else None
+            tv = torch.from_numpy(vs)[None].float() if vs is not None else None
+            lat = twi.generate(model,
+                               first_frame_latent=s["first_frame_latent"],
+                               text_embedding=s["text_embedding"],
+                               text_attention_mask=s["text_attention_mask"],
+                               track_points=tp,
+                               track_visibility=tv,
+                               clip_feature=s["clip_feature"],
+                               num_steps=args.steps,
+                               seed=args.seed)
+            frames = twi.decode_to_pixels(model, lat)
+            ov_tr = tr if tr is not None else np.zeros((Tpx, 1, 2), np.float32)
+            ov_vs = vs if vs is not None else np.zeros((Tpx, 1), np.float32)
+            imageio.mimsave(os.path.join(args.out, f"clip{ci}_{name}.mp4"),
+                            _overlay(frames, ov_tr, ov_vs),
+                            fps=args.fps,
+                            macro_block_size=1)
+
+            epe = None
+            if tr is not None:
+                H, W = frames.shape[1], frames.shape[2]
+                tr_px = tr.copy()
+                tr_px[..., 0] *= W
+                tr_px[..., 1] *= H
+                epe = compute_epe(frames, tr_px, vs, ct, model.device)["epe"]
+            results.append((ci, name, epe))
+            print(f"[clip {ci}] {name:12s} EPE={epe if epe is None else round(epe,2)}", flush=True)
+
+    print("\n=== EPE summary (lower = follows control better) ===")
+    for ci, name, epe in results:
+        print(f"  clip {ci}  {name:12s}  {('n/a' if epe is None else f'{epe:7.2f} px')}")
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/test_motion_cfg.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Quick motion-CFG test: does guidance>1 improve action adherence on an existing ckpt?"""
+import argparse
+import os
+import sys
+import numpy as np
+import torch
+
+sys.path.insert(0, os.path.dirname(__file__))
+import synthetic_tracks as st
+import trackwan_infer as twi
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--export", required=True)
+    p.add_argument(
+        "--data",
+        default=
+        "/mnt/weka/home/hao.zhang/shao/data/motion_pipeline/droid_track_200/preprocessed_i2v_track/combined_parquet_dataset"
+    )
+    p.add_argument("--yaml", default="examples/train/scenario/worldmodel/finetune_wantrack_droid_overfit.yaml")
+    p.add_argument("--clips", type=int, nargs="+", default=[0, 1])
+    p.add_argument("--scales", type=float, nargs="+", default=[1.0, 2.0, 3.0])
+    p.add_argument("--steps", type=int, default=30)
+    a = p.parse_args()
+    model, tc = twi.load_trackwan(a.export, a.yaml)
+    text_len = int(tc.pipeline_config.text_encoder_configs[0].arch_config.text_len)
+    samples = twi.load_conditioning_from_parquet(a.data, a.clips, text_len)
+    from fastvideo.eval.metrics.motion.cotracker_epe.metric import compute_epe, load_cotracker
+    ct = load_cotracker(model.device)
+    nlt = samples[0]["first_frame_latent"].shape[2]
+    ratio = int(tc.pipeline_config.vae_config.arch_config.temporal_compression_ratio)
+    Tpx = (nlt - 1) * ratio + 1
+    # grid-snapped dense drag control (moves a radius-0.2 region by +0.3 x)
+    g = st.make_grid(50)
+    dt, dv = st.drag(g, Tpx, center=(0.5, 0.5), dx=0.3, dy=0.0, radius=0.2)
+    disp = np.sqrt(((dt - dt[0:1])**2).sum(-1))
+    moving = disp.max(0) > 8 / 832.0  # normalized thresh ~ px
+    mv = np.where(moving)[0]
+    print(f"{'clip':>4} {'scale':>6} {'EPE_all':>8} {'EPE_moving':>11}")
+    for ci, s in zip(a.clips, samples, strict=False):
+        tp = torch.from_numpy(dt)[None].float()
+        tvv = torch.from_numpy(dv)[None].float()
+        for sc in a.scales:
+            lat = twi.generate(model,
+                               first_frame_latent=s["first_frame_latent"],
+                               text_embedding=s["text_embedding"],
+                               text_attention_mask=s["text_attention_mask"],
+                               track_points=tp,
+                               track_visibility=tvv,
+                               clip_feature=s["clip_feature"],
+                               num_steps=a.steps,
+                               seed=1000,
+                               guidance_scale=sc)
+            fr = twi.decode_to_pixels(model, lat)
+            H, W = fr.shape[1], fr.shape[2]
+            dpx = dt.copy()
+            dpx[..., 0] *= W
+            dpx[..., 1] *= H
+            epe_all = compute_epe(fr, dpx, dv, ct, model.device)["epe"]
+            epe_mv = compute_epe(fr, dpx[:, mv], dv[:, mv], ct, model.device)["epe"]
+            print(f"{ci:>4} {sc:>6.1f} {epe_all:>8.2f} {epe_mv:>11.2f}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/track_informativeness.py

@@ -0,0 +1,477 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Diagnostic bake-off: which point tracks are *informative* vs *generic*?
+
+A point that moves a lot is not necessarily important. If the camera pans, every
+background point moves, but that motion is shared by the crowd, so any one background
+point is redundant. We compare several ways to score / select useful points, from
+simple statistics to clustering and learned embeddings.
+
+Per-point scalar scores (heatmaps, bright = informative):
+  abs_disp          max_t ||p_t - p_0||                       naive raw motion
+  affine_residual   residual after a robust per-frame global   removes camera pan/zoom/rot
+                    affine x0 -> p_t
+  lowrank_residual  residual after subtracting mean trajectory  unique vs ALL points, but
+                    + top-k shared SVD modes (ONE subspace)     assumes a single subspace
+  motion_outlier    1 - HDBSCAN membership prob on trajectory   density outlier in motion space
+                    embeddings                                  (#1 motion clustering)
+  subspace_resid    per-cluster low-rank residual after         union-of-subspaces; the
+                    Spectral motion segmentation                per-object lowrank (#2 SSC-lite)
+  dino_uniqueness   1 - membership prob clustering DINOv2       semantic ⊕ motion (#3)
+                    appearance ⊕ trajectory embeddings
+
+Grouping / subset panels:
+  motion_clusters   HDBSCAN labels on trajectory embeddings     motion groups (#1)
+  subspace_clusters Spectral motion segmentation labels         object subspaces (#2)
+  dpp / fps / greedy  diverse K-subset selection                what sampling wants (#4)
+
+``render`` (CPU; DINOv2 on CPU) writes per clip a big multi-panel PNG + an overlay mp4
+colored by low-rank informativeness + stats. ``serve`` is a gradio gallery (--share).
+
+    .venv/bin/python data_pipeline/track_informativeness.py render --data-dir <root> --out <dir> --limit 8
+    .venv/bin/python data_pipeline/track_informativeness.py serve --viz-dir <dir> --share
+"""
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(__file__))
+
+K_SELECT = 64  # size of the illustrative "useful subset" for the selection methods
+
+
+# ---------------------------------------------------------------------------- overlay / io
+def _draw_overlay(frames: Any,
+                  tracks: Any,
+                  vis: Any,
+                  colors: Any,
+                  tail: int = 12,
+                  radius: int = 2,
+                  vis_thresh: float = 0.5) -> list:
+    """Self-contained track overlay (dots + short tails); no fastvideo/triton dependency."""
+    from PIL import Image, ImageDraw
+    T, N, _ = tracks.shape
+    out = []
+    for t in range(T):
+        img = Image.fromarray(np.ascontiguousarray(frames[t]))
+        dr = ImageDraw.Draw(img)
+        t0 = max(0, t - tail)
+        for i in range(N):
+            col = (int(colors[i][0]), int(colors[i][1]), int(colors[i][2]))
+            if t - t0 >= 1:
+                dr.line([(float(p[0]), float(p[1])) for p in tracks[t0:t + 1, i]], fill=col, width=1)
+            if vis[t, i] > vis_thresh:
+                x, y = float(tracks[t, i, 0]), float(tracks[t, i, 1])
+                dr.ellipse([x - radius, y - radius, x + radius, y + radius], fill=col)
+        out.append(np.asarray(img))
+    return out
+
+
+def _read_frames(path: str) -> np.ndarray:
+    try:
+        from decord import VideoReader, cpu
+        vr = VideoReader(path, ctx=cpu(0))
+        return vr.get_batch(list(range(len(vr)))).asnumpy()
+    except Exception:  # noqa: BLE001
+        import av
+        c = av.open(path)
+        return np.stack([f.to_ndarray(format="rgb24") for f in c.decode(video=0)])
+
+
+def _norm(x: np.ndarray, pct: float = 97.0) -> np.ndarray:
+    lo = float(x.min())
+    hi = float(np.percentile(x, pct))
+    return np.zeros_like(x) if hi <= lo else np.clip((x - lo) / (hi - lo), 0.0, 1.0)
+
+
+# ---------------------------------------------------------------------------- statistical scores
+def abs_disp(tracks: np.ndarray) -> np.ndarray:
+    return np.sqrt(((tracks - tracks[0:1])**2).sum(-1)).max(0)
+
+
+def affine_residual(tracks: np.ndarray, vis: np.ndarray, iters: int = 3) -> np.ndarray:
+    T, N, _ = tracks.shape
+    X0 = np.concatenate([tracks[0], np.ones((N, 1), np.float32)], 1)  # [N,3]
+    resid = np.zeros((T, N), np.float32)
+    for t in range(1, T):
+        pt = tracks[t]
+        w = (vis[t] > 0.5).astype(np.float64)
+        if w.sum() < 6:
+            w = np.ones(N)
+        for _ in range(iters):
+            W = X0.T * w
+            A = np.linalg.lstsq((W @ X0), (W @ pt), rcond=None)[0]
+            r = np.sqrt(((X0 @ A - pt)**2).sum(-1)) + 1e-6
+            delta = 1.5 * np.median(r[w > 0]) + 1e-3
+            w = np.where(r <= delta, 1.0, delta / r) * (vis[t] > 0.5)
+        resid[t] = r
+    return np.median(resid[1:], 0)
+
+
+def global_affine_motion(tracks: np.ndarray, vis: np.ndarray) -> float:
+    """Camera-fixedness proxy: median per-frame magnitude of the fitted GLOBAL affine
+    displacement (in px), normalized by frame diagonal. ~0 => fixed camera; large => the
+    whole scene translates/zooms (moving/egocentric camera). Robust (Huber) fit."""
+    T, N, _ = tracks.shape
+    X0 = np.concatenate([tracks[0], np.ones((N, 1), np.float32)], 1)
+    mags = []
+    for t in range(1, T):
+        pt = tracks[t]
+        w = (vis[t] > 0.5).astype(np.float64)
+        if w.sum() < 6:
+            w = np.ones(N)
+        for _ in range(2):
+            Wm = X0.T * w
+            A = np.linalg.lstsq((Wm @ X0), (Wm @ pt), rcond=None)[0]
+            r = np.sqrt(((X0 @ A - pt)**2).sum(-1)) + 1e-6
+            delta = 1.5 * np.median(r[w > 0]) + 1e-3
+            w = np.where(r <= delta, 1.0, delta / r) * (vis[t] > 0.5)
+        mags.append(float(np.median(np.sqrt(((X0 @ A - tracks[0])**2).sum(-1)))))
+    diag = float(np.hypot(tracks[..., 0].max(), tracks[..., 1].max()) + 1e-6)
+    return round(float(np.median(mags)) / diag, 4)
+
+
+def displacement_matrix(tracks: np.ndarray) -> np.ndarray:
+    """[N, 2T] per-point displacement relative to frame 0."""
+    T, N, _ = tracks.shape
+    return (tracks - tracks[0:1]).transpose(1, 0, 2).reshape(N, 2 * T)
+
+
+def lowrank_residual(tracks: np.ndarray, rank: int = 3) -> np.ndarray:
+    D = displacement_matrix(tracks)
+    Dc = D - D.mean(0, keepdims=True)
+    U, S, Vt = np.linalg.svd(Dc, full_matrices=False)
+    r = min(rank, S.shape[0])
+    return np.sqrt(((Dc - (U[:, :r] * S[:r]) @ Vt[:r])**2).sum(-1))
+
+
+# ---------------------------------------------------------------------------- embeddings + clustering
+def traj_embedding(tracks: np.ndarray, n_pca: int = 16) -> np.ndarray:
+    from sklearn.decomposition import PCA
+    D = displacement_matrix(tracks)
+    D = (D - D.mean(0)) / (D.std(0) + 1e-6)
+    d = int(min(n_pca, min(D.shape) - 1))
+    return PCA(n_components=d, random_state=0).fit_transform(D).astype(np.float32)
+
+
+def motion_cluster(emb: np.ndarray):
+    """#1: HDBSCAN on trajectory embeddings -> labels, outlierness score (1 - membership prob)."""
+    from sklearn.cluster import HDBSCAN
+    mcs = max(10, emb.shape[0] // 100)
+    cl = HDBSCAN(min_cluster_size=mcs, min_samples=5).fit(emb)
+    return cl.labels_, (1.0 - cl.probabilities_).astype(np.float32)
+
+
+def subspace_cluster(tracks: np.ndarray, emb: np.ndarray, k: int):
+    """#2: Spectral motion segmentation into k groups; per-cluster low-rank residual."""
+    from sklearn.cluster import SpectralClustering
+    k = int(np.clip(k, 2, 8))
+    lab = SpectralClustering(n_clusters=k,
+                             affinity="nearest_neighbors",
+                             n_neighbors=15,
+                             assign_labels="cluster_qr",
+                             random_state=0).fit_predict(emb)
+    D = displacement_matrix(tracks)
+    resid = np.zeros(lab.shape[0], np.float32)
+    for c in np.unique(lab):
+        idx = np.where(lab == c)[0]
+        Z = D[idx] - D[idx].mean(0, keepdims=True)
+        r = min(4, min(Z.shape) - 1)
+        if r > 0:
+            U, S, Vt = np.linalg.svd(Z, full_matrices=False)
+            resid[idx] = np.sqrt(((Z - (U[:, :r] * S[:r]) @ Vt[:r])**2).sum(-1))
+    return lab, resid
+
+
+def dino_point_features(frame0: np.ndarray, xy: np.ndarray, model, device: str) -> np.ndarray:
+    """#3: sample DINOv2 patch tokens at each frame-0 point (bilinear). Returns [N,C] L2-normed."""
+    import torch
+    import torch.nn.functional as F
+    H, W = frame0.shape[0], frame0.shape[1]
+    gw, gh = max(1, round(W / 14)), max(1, round(H / 14))
+    rw, rh = gw * 14, gh * 14
+    img = torch.from_numpy(frame0).float().permute(2, 0, 1)[None] / 255.0
+    img = F.interpolate(img, size=(rh, rw), mode="bilinear", align_corners=False)
+    mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
+    std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
+    img = ((img - mean) / std).to(device)
+    with torch.no_grad():
+        tok = model.forward_features(img)["x_norm_patchtokens"][0]  # [gh*gw, C]
+    C = tok.shape[-1]
+    grid = tok.reshape(gh, gw, C).permute(2, 0, 1)[None]  # [1,C,gh,gw]
+    # normalized sample coords in [-1,1] from pixel xy
+    gx = (xy[:, 0] / W) * 2 - 1
+    gy = (xy[:, 1] / H) * 2 - 1
+    samp = torch.from_numpy(np.stack([gx, gy], -1)).float().view(1, -1, 1, 2).to(device)
+    feat = F.grid_sample(grid, samp, mode="bilinear", align_corners=False)[0, :, :, 0].T  # [N,C]
+    feat = feat.cpu().numpy()
+    return (feat / (np.linalg.norm(feat, axis=1, keepdims=True) + 1e-6)).astype(np.float32)
+
+
+def dino_motion_cluster(dino_feat: np.ndarray, traj_emb: np.ndarray):
+    """Cluster fused (semantic ⊕ motion) embedding -> labels, uniqueness (1 - membership prob)."""
+    from sklearn.cluster import HDBSCAN
+    from sklearn.decomposition import PCA
+    dz = PCA(n_components=int(min(16, min(dino_feat.shape) - 1)), random_state=0).fit_transform(dino_feat)
+    dz = (dz - dz.mean(0)) / (dz.std(0) + 1e-6)
+    tz = (traj_emb - traj_emb.mean(0)) / (traj_emb.std(0) + 1e-6)
+    fused = np.concatenate([dz, tz], 1)
+    cl = HDBSCAN(min_cluster_size=max(10, fused.shape[0] // 100), min_samples=5).fit(fused)
+    return cl.labels_, (1.0 - cl.probabilities_).astype(np.float32)
+
+
+# ---------------------------------------------------------------------------- subset selection (#4)
+def select_fps(emb: np.ndarray, K: int, seed: int) -> np.ndarray:
+    sel = [int(seed)]
+    d = np.linalg.norm(emb - emb[seed], axis=1)
+    for _ in range(K - 1):
+        i = int(np.argmax(d))
+        sel.append(i)
+        d = np.minimum(d, np.linalg.norm(emb - emb[i], axis=1))
+    return np.array(sorted(set(sel)))
+
+
+def select_dpp(quality: np.ndarray, emb: np.ndarray, K: int) -> np.ndarray:
+    """Greedy MAP for a DPP with L = diag(q) S diag(q), S a Gaussian similarity kernel."""
+    from scipy.spatial.distance import cdist
+    q = _norm(quality) + 0.05
+    d2 = cdist(emb, emb, "sqeuclidean")
+    S = np.exp(-d2 / (np.median(d2) + 1e-6))
+    L = (q[:, None] * q[None, :]) * S
+    N = L.shape[0]
+    cis = np.zeros((K, N))
+    di2s = np.diag(L).copy()
+    j = int(np.argmax(di2s))
+    sel = [j]
+    for it in range(1, K):
+        k = it - 1
+        ei = (L[j] - cis[:k, j] @ cis[:k]) / np.sqrt(max(di2s[j], 1e-12))
+        cis[k] = ei
+        di2s = di2s - ei**2
+        di2s[sel] = -np.inf
+        j = int(np.argmax(di2s))
+        if di2s[j] <= 1e-10:
+            break
+        sel.append(j)
+    return np.array(sorted(sel))
+
+
+def select_greedy_recon(xy0: np.ndarray, D: np.ndarray, K: int) -> np.ndarray:
+    """Add the point whose displacement is least predictable by RBF-interpolating the
+    already-selected points over frame-0 positions -> targets motion discontinuities."""
+    from scipy.spatial.distance import cdist
+    sel = [int(np.argmax(np.linalg.norm(D - D.mean(0), axis=1)))]
+    for _ in range(K - 1):
+        S = np.array(sel)
+        d2 = cdist(xy0, xy0[S], "sqeuclidean")
+        W = np.exp(-d2 / (np.median(d2) + 1e-6))
+        W /= (W.sum(1, keepdims=True) + 1e-9)
+        err = np.linalg.norm(D - W @ D[S], axis=1)
+        err[S] = -1
+        sel.append(int(np.argmax(err)))
+    return np.array(sorted(sel))
+
+
+# ---------------------------------------------------------------------------- render
+def _spearman(a, b) -> float:
+    from scipy.stats import spearmanr
+    return round(float(spearmanr(a, b).correlation), 3)
+
+
+def cmd_render(args: argparse.Namespace) -> None:
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    import imageio.v2 as imageio
+
+    dino_model, dino_dev = None, "cpu"
+    if not args.no_dino:
+        try:
+            import torch
+            dino_dev = "cuda" if torch.cuda.is_available() else "cpu"
+            dino_model = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14").to(dino_dev).eval()
+            print(f"[info] DINOv2 loaded on {dino_dev}", flush=True)
+        except Exception as e:  # noqa: BLE001
+            print(f"[info] DINOv2 unavailable ({type(e).__name__}: {e}); skipping semantic panels", flush=True)
+
+    data = Path(args.data_dir)
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    items = json.loads((data / "videos2caption.json").read_text())[:args.limit]
+    turbo = matplotlib.colormaps["turbo"]
+    entries = []
+    for k, it in enumerate(items, 1):
+        stem = Path(it["path"]).stem
+        vpath = str(data / "videos" / it["path"])
+        npz = it.get("points_path") or str(data / "tracks" / f"{stem}.npz")
+        frames = _read_frames(vpath)
+        d = np.load(npz)
+        tracks = d["tracks"].astype(np.float32)[:frames.shape[0]]
+        vis = d["visibility"].astype(np.float32)[:frames.shape[0]]
+        N = tracks.shape[1]
+        x, y = tracks[0, :, 0], tracks[0, :, 1]
+        D = displacement_matrix(tracks)
+        emb = traj_embedding(tracks)
+
+        # scalar scores
+        scores = {
+            "abs_disp": abs_disp(tracks),
+            "affine_residual": affine_residual(tracks, vis),
+            "lowrank_residual": lowrank_residual(tracks),
+        }
+        m_lab, m_out = motion_cluster(emb)
+        scores["motion_outlier"] = m_out
+        k_est = len(np.unique(m_lab[m_lab >= 0])) or 4
+        s_lab, s_res = subspace_cluster(tracks, emb, k_est)
+        scores["subspace_resid"] = s_res
+        d_lab = None
+        if dino_model is not None:
+            try:
+                dfeat = dino_point_features(frames[0], tracks[0], dino_model, dino_dev)
+                d_lab, d_uniq = dino_motion_cluster(dfeat, emb)
+                scores["dino_uniqueness"] = d_uniq
+            except Exception as e:  # noqa: BLE001
+                print(f"[info]   dino features failed for {stem}: {e}", flush=True)
+
+        # subset selections (#4)
+        seed = int(np.argmax(scores["lowrank_residual"]))
+        subsets = {
+            "dpp": select_dpp(scores["lowrank_residual"], emb, K_SELECT),
+            "fps": select_fps(emb, K_SELECT, seed),
+            "greedy_recon": select_greedy_recon(tracks[0], D, K_SELECT),
+        }
+
+        # ---- figure: row1 scalar scores, row2 cluster maps + subsets
+        label_panels = [("motion_clusters", m_lab), ("subspace_clusters", s_lab)]
+        if d_lab is not None:
+            label_panels.append(("dino_clusters", d_lab))
+        panels = ([("score", n, s)
+                   for n, s in scores.items()] + [("labels", n, lab)
+                                                  for n, lab in label_panels] + [("subset", n, idx)
+                                                                                 for n, idx in subsets.items()])
+        ncol = 5
+        nrow = int(np.ceil(len(panels) / ncol))
+        fig, axes = plt.subplots(nrow, ncol, figsize=(4.4 * ncol, 3.1 * nrow))
+        axes = np.array(axes).reshape(-1)
+        for ax in axes:
+            ax.axis("off")
+        for ax, (kind, name, val) in zip(axes, panels, strict=False):
+            ax.imshow(frames[0])
+            if kind == "score":
+                ax.scatter(x, y, c=_norm(val), cmap="turbo", s=10, vmin=0, vmax=1)
+                ax.set_title(f"{name}", fontsize=10)
+            elif kind == "labels":
+                noise = val < 0
+                ax.scatter(x[noise], y[noise], c="0.5", s=6)
+                ax.scatter(x[~noise], y[~noise], c=val[~noise], cmap="tab20", s=10)
+                ax.set_title(f"{name} ({len(np.unique(val[val >= 0]))} groups)", fontsize=10)
+            else:  # subset
+                ax.scatter(x, y, c="0.4", s=3)
+                ax.scatter(x[val], y[val], c="red", s=26, edgecolors="white", linewidths=0.4)
+                ax.set_title(f"{name} (K={len(val)})", fontsize=10)
+        fig.suptitle(f"{stem}  |  {(it['cap'][0] if isinstance(it.get('cap'), list) else '')[:90]}", fontsize=11)
+        fig.tight_layout()
+        heat = f"{stem}_scores.png"
+        fig.savefig(str(out / heat), dpi=80, bbox_inches="tight")
+        plt.close(fig)
+
+        # ---- overlay colored by lowrank informativeness
+        lr = _norm(scores["lowrank_residual"])
+        pcols = (np.array([turbo(v)[:3] for v in lr]) * 255).astype(np.uint8)
+        G = int(round(N**0.5))
+        sel = (np.arange(N).reshape(G, G)[::max(1, G // 25), ::max(1, G // 25)].reshape(-1) if G *
+               G == N else np.arange(0, N, max(1, N // 625)))
+        ov = _draw_overlay(frames, tracks[:, sel].copy(), vis[:, sel], pcols[sel], 14, 2, 0.5)
+        vid = f"{stem}_lowrank.mp4"
+        imageio.mimsave(str(out / vid), ov, fps=int(it.get("fps", 24)), macro_block_size=1)
+
+        # ---- stats
+        names = list(scores.keys())
+        rho = {f"{a}|{b}": _spearman(scores[a], scores[b]) for i, a in enumerate(names) for b in names[i + 1:]}
+        lr_raw = scores["lowrank_residual"]
+        lo_info = lr_raw <= np.percentile(lr_raw, 50)
+
+        def _cov_generic(idx: Any, m_lab: Any = m_lab, lo_info: Any = lo_info) -> dict:
+            grp = m_lab[idx]
+            ngrp = len(np.unique(grp[grp >= 0]))
+            tot = max(len(np.unique(m_lab[m_lab >= 0])), 1)
+            return {"coverage_of_motion_groups": f"{ngrp}/{tot}", "frac_generic": round(float(lo_info[idx].mean()), 3)}
+
+        stats = {
+            "id": stem,
+            "n_points": int(N),
+            "camera_motion": global_affine_motion(tracks, vis),  # ~0 = fixed cam, high = moving
+            "n_motion_groups": int(len(np.unique(m_lab[m_lab >= 0]))),
+            "noise_frac": round(float((m_lab < 0).mean()), 3),
+            "spearman": rho,
+            "subset_quality": {
+                name: _cov_generic(idx)
+                for name, idx in subsets.items()
+            },
+        }
+        entries.append({"id": stem, "heat": heat, "video": vid, "stats": stats})
+        print(
+            f"[info] [{k}/{len(items)}] {stem}: cam_motion={stats['camera_motion']}, "
+            f"{stats['n_motion_groups']} motion groups, noise={stats['noise_frac']}, "
+            f"rho(abs|lowrank)={rho.get('abs_disp|lowrank_residual')}",
+            flush=True)
+
+    (out / "manifest.json").write_text(json.dumps(entries, indent=2))
+    print(f"[info] rendered {len(entries)} clips -> {out}", flush=True)
+
+
+def cmd_serve(args: argparse.Namespace) -> None:
+    import gradio as gr
+    viz = Path(args.viz_dir).resolve()
+    entries = json.loads((viz / "manifest.json").read_text())
+    gal = [(str(viz / e["heat"]), e["id"]) for e in entries]
+
+    def show(evt: gr.SelectData):
+        e = entries[evt.index]
+        return str(viz / e["heat"]), str(viz / e["video"]), json.dumps(e["stats"], indent=2)
+
+    with gr.Blocks(title="Track informativeness bake-off") as demo:
+        gr.Markdown(
+            "### Which point tracks are useful? Statistics vs clustering vs embeddings vs subset selection\n"
+            "**Row 1 (scores, bright = informative):** `abs_disp` raw motion · `affine_residual` after "
+            "removing camera affine · `lowrank_residual` unique vs all (one subspace) · `motion_outlier` "
+            "HDBSCAN density outlier · `subspace_resid` per-object lowrank · `dino_uniqueness` semantic⊕motion. "
+            "**Row 2:** `motion_clusters` / `subspace_clusters` / `dino_clusters` (colored groups, gray = noise) "
+            "and the diverse K-subsets `dpp` / `fps` / `greedy_recon` (red = chosen). "
+            "Click a clip -> panels + overlay colored by low-rank informativeness. `stats.subset_quality` shows "
+            "how well each subset covers the motion groups and how many chosen points are generic.")
+        with gr.Row():
+            g = gr.Gallery(value=gal, columns=2, height=640, label="clips (click)")
+            with gr.Column():
+                heat = gr.Image(label="method panels")
+                vid = gr.Video(label="overlay colored by low-rank informativeness")
+                meta = gr.Code(label="stats", language="json")
+        g.select(show, None, [heat, vid, meta])
+    demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share, allowed_paths=[str(viz)])
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub = p.add_subparsers(dest="cmd", required=True)
+    r = sub.add_parser("render")
+    r.add_argument("--data-dir", required=True)
+    r.add_argument("--out", required=True)
+    r.add_argument("--limit", type=int, default=8)
+    r.add_argument("--no-dino", action="store_true", help="skip the DINOv2 semantic panels")
+    r.set_defaults(func=cmd_render)
+    s = sub.add_parser("serve")
+    s.add_argument("--viz-dir", required=True)
+    s.add_argument("--host", default="0.0.0.0")
+    s.add_argument("--port", type=int, default=7883)
+    s.add_argument("--share", action="store_true")
+    s.set_defaults(func=cmd_serve)
+    a = p.parse_args()
+    a.func(a)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/track_loop.sh

@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+# Overlap tracking with the slow self-healing download: idempotent multi-node passes
+# over the growing videos/ dir until download done (98 parts) AND ~all tracked.
+# MUST run inside srun on a compute node (login-node `sleep` gets killed by the harness).
+# Waits for any in-flight tracking pass first, so it's safe to relaunch (no double-submit).
+set +e
+WORK=/mnt/lustre/vlm-s4duan
+cd "$WORK/FastVideo"
+DC(){ ls "$WORK"/openvid_1m/_extracted/*.done 2>/dev/null | wc -l; }
+RUNNING(){ squeue -u vlm-s4duan -h -n cotracker_dp -o %i 2>/dev/null; }
+NODES=${NODES:-8}
+while true; do
+  while [ -n "$(RUNNING)" ]; do sleep 90; done            # let in-flight pass finish
+  ls "$WORK"/openvid_1m/videos/*.mp4 > "$WORK"/openvid_1m/videos_all.txt 2>/dev/null
+  nv=$(wc -l < "$WORK"/openvid_1m/videos_all.txt 2>/dev/null); nv=${nv:-0}
+  nt=$(ls "$WORK"/openvid_1m/tracks/*.npz 2>/dev/null | wc -l)
+  echo "[$(date +%H:%M)] parts=$(DC)/98 videos=$nv tracks=$nt"
+  if [ "$(DC)" -ge 98 ] && [ "$nt" -ge "$((nv - nv/40 - 200))" ]; then
+    echo "TRACK_LOOP_DONE videos=$nv tracks=$nt"; break; fi
+  if [ "$nv" -le "$((nt + 300))" ]; then echo "  little new; wait 300s"; sleep 300; continue; fi
+  jid=$(CLIPS_DIR="$WORK"/openvid_1m/clips VIDEO_LIST="$WORK"/openvid_1m/videos_all.txt \
+        OUT_DIR="$WORK"/openvid_1m/tracks NODES=$NODES PROCS_PER_GPU=2 FPS=24 NUM_FRAMES=121 GRID=50 \
+        bash data_pipeline/run_tracks_slurm.sh 2>&1 | grep -oE "[0-9]+$" | tail -1)
+  echo "  submitted track job $jid over $nv videos"
+  [ -z "$jid" ] && sleep 180
+done

data_pipeline/track_sensitivity.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Single-checkpoint *swap sensitivity* (one process per checkpoint).
+
+Linearized, generation-free version of the controllability EPE test: at a fixed
+noise + timestep, measure how much the predicted velocity ``v`` moves when we swap
+ONLY one conditioner and hold the rest fixed:
+
+  S_track = ||v(swap tracks)      - v(base)|| / ||v(base)||
+  S_ff    = ||v(swap first-frame) - v(base)|| / ||v(base)||
+  S_text  = ||v(zero text)        - v(base)|| / ||v(base)||
+
+Averaged over a few timesteps and the given clips. This captures the *backbone's*
+responsiveness to each input (unlike weight-norm or input-projection magnitude,
+which only see the input side). Expectation for the control50 run: S_track high at
+step 500, decaying by 6000 (control prior forgotten), while S_ff stays high.
+
+One checkpoint per process (loads via the framework loader, so weights/DTensor are
+handled correctly). Launch once per --step across GPUs.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from typing import Any
+
+import torch
+
+sys.path.insert(0, os.path.dirname(__file__))
+import trackwan_infer as twi  # noqa: E402
+
+
+def _v(model: Any, latents: torch.Tensor, ff: torch.Tensor, tp: torch.Tensor, tv: torch.Tensor, txt: torch.Tensor,
+       mask: torch.Tensor, img: torch.Tensor, ts: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    from fastvideo.forward_context import set_forward_context
+    cond20 = model._build_i2v_cond_concat(ff)
+    model_in = torch.cat([latents.to(dtype), cond20], dim=1)
+    with torch.no_grad(), torch.autocast(model.device.type, dtype=dtype), \
+            set_forward_context(current_timestep=ts, attn_metadata=None):
+        v = model.transformer(hidden_states=model_in,
+                              encoder_hidden_states=txt,
+                              encoder_attention_mask=mask,
+                              timestep=ts,
+                              encoder_hidden_states_image=img,
+                              track_points=tp,
+                              track_visibility=tv,
+                              return_dict=False)
+    return v.float()
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--export", required=True, help="single export dir (export_<step>)")
+    p.add_argument("--step", type=int, default=-1, help="label only")
+    p.add_argument("--yaml", default="examples/train/scenario/worldmodel/finetune_wantrack_i2v.yaml")
+    p.add_argument("--data",
+                   default="/mnt/weka/home/hao.zhang/shao/data/motion_pipeline/"
+                   "wan22_a14b_720p_24fps/preprocessed_i2v_track_funinp/combined_parquet_dataset")
+    p.add_argument("--clips", type=int, nargs="+", default=[0, 1])
+    p.add_argument("--timestep-idx", type=int, nargs="+", default=[2, 15, 27])
+    p.add_argument("--seed", type=int, default=1000)
+    args = p.parse_args()
+
+    dtype = torch.bfloat16
+    model, tc = twi.load_trackwan(args.export, args.yaml)
+    text_len = int(tc.pipeline_config.text_encoder_configs[0].arch_config.text_len)
+    samples = twi.load_conditioning_from_parquet(args.data, args.clips, text_len)
+    device = model.device
+
+    from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+        FlowMatchEulerDiscreteScheduler, )
+    sched = FlowMatchEulerDiscreteScheduler(shift=float(model.timestep_shift))
+    sched.set_timesteps(30, device=device)
+    tlist = [sched.timesteps[i].reshape(1).to(device, dtype) for i in args.timestep_idx]
+
+    pairs = []
+    for k, ci in enumerate(args.clips):
+        s = samples[k]
+        sw = samples[(k + 1) % len(samples)]
+        _, _, T, H, W = s["first_frame_latent"].shape
+        g = torch.Generator(device="cpu").manual_seed(args.seed + ci)
+        latents = torch.randn((1, 16, T, H, W), generator=g, dtype=torch.float32).to(device)
+
+        def d(x: torch.Tensor) -> torch.Tensor:
+            return x.to(device, dtype)
+
+        pairs.append(
+            dict(lat=latents,
+                 ff=d(s["first_frame_latent"]),
+                 ff_sw=d(sw["first_frame_latent"]),
+                 tp=d(s["track_points"]),
+                 tv=d(s["track_visibility"]),
+                 tp_sw=d(sw["track_points"]),
+                 tv_sw=d(sw["track_visibility"]),
+                 txt=d(s["text_embedding"]),
+                 mask=d(s["text_attention_mask"]),
+                 img=d(s["clip_feature"])))
+
+    st = sf = sx = 0.0
+    n = 0
+    for pr in pairs:
+        for ts in tlist:
+            vb = _v(model, pr["lat"], pr["ff"], pr["tp"], pr["tv"], pr["txt"], pr["mask"], pr["img"], ts, dtype)
+            nb = vb.norm() + 1e-6
+            vt = _v(model, pr["lat"], pr["ff"], pr["tp_sw"], pr["tv_sw"], pr["txt"], pr["mask"], pr["img"], ts, dtype)
+            vf = _v(model, pr["lat"], pr["ff_sw"], pr["tp"], pr["tv"], pr["txt"], pr["mask"], pr["img"], ts, dtype)
+            vx = _v(model, pr["lat"], pr["ff"], pr["tp"], pr["tv"], torch.zeros_like(pr["txt"]), pr["mask"], pr["img"],
+                    ts, dtype)
+            st += float((vt - vb).norm() / nb)
+            sf += float((vf - vb).norm() / nb)
+            sx += float((vx - vb).norm() / nb)
+            n += 1
+    print(f"RESULT step={args.step} S_track={st/n:.4f} S_ff={sf/n:.4f} S_text={sx/n:.4f}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/trackwan_infer.py

@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Standalone TrackWan inference: generate a video from (first frame + text + tracks).
+
+Loads a checkpoint exported with ``dcp_to_diffusers`` into the *same*
+``WanTrackModel`` wrapper used for training (so train/inference parity is exact),
+then runs a flow-matching denoise loop feeding the point-track control. Used by
+the controllability tests and the Gradio app.
+
+Backend only -- no track authoring here (see ``synthetic_tracks.py``) and no
+metrics (see ``fastvideo/eval/metrics/motion/cotracker_epe``).
+"""
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import numpy as np
+import torch
+
+
+def _ensure_dist_env() -> None:
+    for k, v in [("RANK", "0"), ("LOCAL_RANK", "0"), ("WORLD_SIZE", "1"), ("MASTER_ADDR", "127.0.0.1"),
+                 ("MASTER_PORT", "29600")]:
+        os.environ.setdefault(k, v)
+
+
+def load_trackwan(export_dir: str, yaml_path: str) -> tuple[Any, Any]:
+    """Build a ``WanTrackModel`` with the exported (trained) weights on 1 GPU."""
+    _ensure_dist_env()
+    from fastvideo.distributed import (
+        maybe_init_distributed_environment_and_model_parallel, )
+    from fastvideo.train.utils.config import load_run_config
+    from fastvideo.train.models.wantrack.wantrack import WanTrackModel
+
+    maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1)
+
+    cfg = load_run_config(yaml_path)
+    tc = cfg.training
+    tc.distributed.tp_size = 1
+    tc.distributed.sp_size = 1
+    tc.distributed.num_gpus = 1
+    tc.distributed.hsdp_replicate_dim = 1
+    tc.distributed.hsdp_shard_dim = 1
+    tc.model_path = export_dir  # vae + components loaded from the export
+
+    model = WanTrackModel(
+        init_from=export_dir,
+        training_config=tc,
+        trainable=False,
+        flow_shift=float(tc.pipeline_config.flow_shift),
+        enable_gradient_checkpointing_type=None,
+    )
+    model.init_preprocessors(tc)
+    model.transformer.eval()
+    return model, tc
+
+
+def load_conditioning_from_parquet(data_path: str, indices: list[int], text_len: int) -> list[dict[str, Any]]:
+    """Pull (first_frame_latent, text, vae_latent, GT tracks, caption) for clips."""
+    import glob
+
+    import pyarrow.parquet as pq
+
+    from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v_track
+    from fastvideo.dataset.utils import collate_rows_from_parquet_schema
+
+    files = sorted(glob.glob(os.path.join(data_path, "**", "*.parquet"), recursive=True))
+    rows: list[dict[str, Any]] = []
+    for f in files:
+        rows.extend(pq.read_table(f).to_pylist())
+    sel = [rows[i] for i in indices]
+    batch = collate_rows_from_parquet_schema(sel,
+                                             pyarrow_schema_i2v_track,
+                                             text_padding_length=int(text_len),
+                                             cfg_rate=0.0)
+    infos = batch.get("info_list") or [{} for _ in sel]
+    out = []
+    for i in range(len(sel)):
+        out.append({
+            "text_embedding": batch["text_embedding"][i:i + 1].clone(),
+            "text_attention_mask": batch["text_attention_mask"][i:i + 1].clone(),
+            "vae_latent": batch["vae_latent"][i:i + 1].clone(),
+            "first_frame_latent": batch["first_frame_latent"][i:i + 1].clone(),
+            "clip_feature": batch["clip_feature"][i:i + 1].clone(),  # [1,SeqLen,Dim] CLIP frame-0
+            "track_points": batch["track_points"][i:i + 1].clone(),  # [1,T,N,2] normalized
+            "track_visibility": batch["track_visibility"][i:i + 1].clone(),  # [1,T,N]
+            "caption": str(infos[i].get("caption", "")),
+        })
+    return out
+
+
+@torch.no_grad()
+def generate(model: Any,
+             *,
+             first_frame_latent: torch.Tensor,
+             text_embedding: torch.Tensor,
+             text_attention_mask: torch.Tensor,
+             track_points: torch.Tensor | None,
+             track_visibility: torch.Tensor | None,
+             clip_feature: torch.Tensor | None = None,
+             num_steps: int = 30,
+             seed: int = 0,
+             guidance_scale: float = 1.0) -> torch.Tensor:
+    """Denoise from noise -> normalized latents [1,16,T,H,W]. Tracks may be None.
+
+    ``guidance_scale`` > 1 = MOTION classifier-free guidance:
+    v = v(no-tracks) + s*(v(tracks) - v(no-tracks)). The no-tracks branch is the
+    base first-frame+text I2V prediction (track channels zeroed), so it works on any
+    checkpoint (no training-time motion dropout needed)."""
+    from fastvideo.forward_context import set_forward_context
+    from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+        FlowMatchEulerDiscreteScheduler, )
+
+    device = model.device
+    dtype = torch.bfloat16
+    ff = first_frame_latent.to(device, dtype)  # [1,16,T,H,W] (already normalized)
+    cond20 = model._build_i2v_cond_concat(ff)  # [1,20,T,H,W]
+    txt = text_embedding.to(device, dtype)
+    mask = text_attention_mask.to(device, dtype)
+    tp = track_points.to(device, dtype) if track_points is not None else None
+    tv = track_visibility.to(device, dtype) if track_visibility is not None else None
+    img = clip_feature.to(device, dtype) if clip_feature is not None else None
+
+    _, _, T, H, W = ff.shape
+    gen = torch.Generator(device="cpu").manual_seed(int(seed))
+    latents = torch.randn((1, 16, T, H, W), generator=gen, dtype=torch.float32).to(device)
+
+    sched = FlowMatchEulerDiscreteScheduler(shift=float(model.timestep_shift))
+    sched.set_timesteps(int(num_steps), device=device)
+    use_cfg = guidance_scale != 1.0 and tp is not None
+    for t in sched.timesteps:
+        model_in = torch.cat([latents.to(dtype), cond20], dim=1)  # [1,36,T,H,W]
+        ts = t.reshape(1).to(device, dtype)
+
+        def _fwd(tpp: torch.Tensor | None,
+                 tvv: torch.Tensor | None,
+                 mi: torch.Tensor = model_in,
+                 tsv: torch.Tensor = ts) -> torch.Tensor:
+            with torch.autocast(device.type, dtype=dtype), set_forward_context(current_timestep=tsv,
+                                                                               attn_metadata=None):
+                return model.transformer(hidden_states=mi,
+                                         encoder_hidden_states=txt,
+                                         encoder_attention_mask=mask,
+                                         timestep=tsv,
+                                         encoder_hidden_states_image=img,
+                                         track_points=tpp,
+                                         track_visibility=tvv,
+                                         return_dict=False)
+
+        if use_cfg:
+            v_cond = _fwd(tp, tv).float()
+            v_uncond = _fwd(None, None).float()  # base I2V (track channels -> 0)
+            v = v_uncond + guidance_scale * (v_cond - v_uncond)
+        else:
+            v = _fwd(tp, tv).float()
+        latents = sched.step(v, t, latents.float(), return_dict=False)[0]
+    return latents
+
+
+@torch.no_grad()
+def decode_to_pixels(model: Any, latents: torch.Tensor) -> np.ndarray:
+    """Normalized latents [1,16,T,H,W] -> uint8 frames [T,H,W,3]."""
+    px = model.decode_latents(latents.permute(0, 2, 1, 3, 4))[0]  # [3,T,H,W] in [0,1]
+    video = (px.clamp(0, 1).float().cpu().numpy() * 255.0).astype(np.uint8)
+    return np.transpose(video, (1, 2, 3, 0))  # [T,H,W,3]
+
+
+@torch.no_grad()
+def decode_reference(model: Any, vae_latent: torch.Tensor) -> np.ndarray:
+    """Decode the raw GT vae_latent (needs normalize first) -> uint8 [T,H,W,3]."""
+    from fastvideo.training.training_utils import normalize_dit_input
+    raw = vae_latent.to(model.device, torch.bfloat16)
+    norm = normalize_dit_input("wan", raw, model.vae)
+    px = model.decode_latents(norm.permute(0, 2, 1, 3, 4))[0]
+    video = (px.clamp(0, 1).float().cpu().numpy() * 255.0).astype(np.uint8)
+    return np.transpose(video, (1, 2, 3, 0))

data_pipeline/two_track_val.py

@@ -0,0 +1,325 @@
+#!/usr/bin/env python3
+"""2-track validation ablation.
+
+For each val sample, keep exactly 2 tracks (one per top-motion object) and zero
+the rest. This is the extreme sparse-inference case: what the user actually gets
+when they draw 2 traces at inference. Compares checkpoints side-by-side.
+
+Uses paper CFG (wt=3.0, wm=1.5) unless overridden. Runs `_sample`-style denoise
+matching track_validation.py exactly.
+
+Usage:
+  python data_pipeline/two_track_val.py \
+    --model-dir /mnt/lustre/vlm-s4duan/exports/merged_bias_ckpt4800 \
+    --yaml examples/train/scenario/worldmodel/finetune_wantrack_synth_stage2_paperLR.yaml \
+    --val-parquet /mnt/lustre/vlm-s4duan/val_examples_mixed/combined_parquet_dataset \
+    --wandb-run-name twotrack_mb4800
+"""
+from __future__ import annotations
+
+import argparse
+import glob
+import os
+import sys
+import time
+from pathlib import Path
+
+import imageio.v2 as imageio
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__)))
+import trackwan_infer as twi  # noqa: E402
+from fastvideo.forward_context import set_forward_context  # noqa: E402
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (  # noqa: E402
+    FlowMatchEulerDiscreteScheduler,
+)
+
+
+def load_val_samples(parquet_dir: str, text_len: int) -> list:
+    from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v_track
+    from fastvideo.dataset.utils import collate_rows_from_parquet_schema
+    fs = sorted(glob.glob(os.path.join(parquet_dir, "**", "*.parquet"), recursive=True))
+    if not fs:
+        raise FileNotFoundError(parquet_dir)
+    rows = []
+    for f in fs:
+        rows.extend(pq.read_table(f).to_pylist())
+    batch = collate_rows_from_parquet_schema(rows, pyarrow_schema_i2v_track,
+                                             text_padding_length=int(text_len), cfg_rate=0.0)
+    infos = batch.get("info_list") or [{} for _ in rows]
+    return [{
+        "id": str(infos[i].get("id", f"clip{i:03d}")) if i < len(infos) else f"clip{i:03d}",
+        "file_name": rows[i].get("file_name", ""),
+        "text_embedding": batch["text_embedding"][i:i + 1].clone(),
+        "text_attention_mask": batch["text_attention_mask"][i:i + 1].clone(),
+        "first_frame_latent": batch["first_frame_latent"][i:i + 1].clone(),
+        "clip_feature": batch["clip_feature"][i:i + 1].clone(),
+        "vae_latent": batch["vae_latent"][i:i + 1].clone(),
+        "track_points": batch["track_points"][i:i + 1].clone(),  # [1,T,N,2] normalised
+        "track_visibility": batch["track_visibility"][i:i + 1].clone(),
+        "object_ids": batch["object_ids"][i:i + 1].clone() if "object_ids" in batch else None,
+        "caption": str(infos[i].get("caption", "") if i < len(infos) else ""),
+    } for i in range(len(rows))]
+
+
+def pick_two_tracks(sample: dict, seed: int, num_bg: int = 0,
+                    num_fg: int = 2) -> tuple[torch.Tensor, torch.Tensor, list]:
+    """Return (tp_2, tv_2, info) — ``num_fg`` foreground tracks + ``num_bg`` bg anchor tracks.
+
+    Foreground picks require the track to STAY IN FRAME [0,1] for all 121 frames — no
+    point in showing traces that fly off the canvas.
+    """
+    tp = sample["track_points"].float()  # [1,T,N,2]
+    tv = sample["track_visibility"].float()  # [1,T,N]
+    oid = sample["object_ids"]  # [1,N]
+    if oid is None:
+        raise ValueError("no object_ids in val sample")
+    oid = oid[0].to(torch.int64)  # [N]
+
+    tp0 = tp[0]  # [T, N, 2]
+    N = tp0.shape[1]
+    motion_per_track = torch.linalg.norm(tp0[-1] - tp0[0], dim=-1)  # [N]
+
+    # Which tracks stay fully within [0,1] on both x and y across ALL frames?
+    in_frame = ((tp0 >= 0.0) & (tp0 <= 1.0)).all(dim=-1).all(dim=0)  # [N]
+
+    # Rank object_ids by max in-frame within-object track motion.
+    fg_oids = torch.unique(oid[oid >= 0])
+    obj_best = []  # (oid, best_track_idx, best_motion)
+    for o in fg_oids.tolist():
+        idxs = ((oid == o) & in_frame & (tv[0, 0] > 0.5)).nonzero(as_tuple=True)[0]
+        if idxs.numel() == 0:
+            continue
+        m = motion_per_track[idxs]
+        best = idxs[int(torch.argmax(m).item())]
+        obj_best.append((o, int(best.item()), float(m.max().item())))
+    obj_best.sort(key=lambda x: -x[2])
+    picked_tracks = [idx for _, idx, _ in obj_best[:num_fg]]
+    info = [{"oid": int(o), "motion_max": mm, "track_idx": pk}
+            for (o, pk, mm) in obj_best[:num_fg]]
+    rng = np.random.RandomState(seed)
+    vis0 = tv[0, 0]  # [N] visibility at frame 0
+
+    # Add num_bg background anchor tracks — near-static (bottom of motion), visible at frame 0.
+    bg_picked = []
+    if num_bg > 0:
+        bg_pool = ((oid == -1) & (vis0 > 0.5)).nonzero(as_tuple=True)[0]
+        if bg_pool.numel() > 0:
+            # Sort bg tracks by motion (ascending), take lowest-motion ones and pick uniformly
+            bg_motion = motion_per_track[bg_pool]
+            n_pool = int(min(bg_pool.numel(), max(num_bg * 20, num_bg)))
+            top_idxs = torch.argsort(bg_motion)[:n_pool]
+            candidate = bg_pool[top_idxs]
+            if candidate.numel() >= num_bg:
+                sel = rng.choice(candidate.numel(), num_bg, replace=False)
+                bg_picked = [int(candidate[k].item()) for k in sel]
+            else:
+                bg_picked = [int(k.item()) for k in candidate]
+            for k in bg_picked:
+                info.append({"oid": -1, "motion_max": float(motion_per_track[k].item()), "track_idx": k})
+
+    # Build zeroed visibility: only picked tracks visible.
+    tv_new = torch.zeros_like(tv)
+    for k in picked_tracks + bg_picked:
+        tv_new[0, :, k] = tv[0, :, k]
+
+    return tp, tv_new, info
+
+
+@torch.no_grad()
+def sample_denoise(model, sample: dict, tp: torch.Tensor, tv: torch.Tensor,
+                   *, w_text: float, w_motion: float, num_steps: int, seed: int) -> np.ndarray:
+    """Reproduce track_validation._sample denoise EXACTLY."""
+    device = model.device
+    dtype = torch.bfloat16
+
+    ff = sample["first_frame_latent"].to(device, dtype)
+    txt = sample["text_embedding"].to(device, dtype)
+    mask = sample["text_attention_mask"].to(device, dtype)
+    clip = sample["clip_feature"].to(device, dtype)
+    tp_d = tp.to(device, dtype)
+    tv_d = tv.to(device, dtype)
+    cond20 = model._build_i2v_cond_concat(ff)
+
+    _, _, T, H, W = ff.shape
+    gen = torch.Generator(device="cpu").manual_seed(int(seed))
+    latents = torch.randn((1, 16, T, H, W), generator=gen, dtype=torch.float32).to(device)
+    sched = FlowMatchEulerDiscreteScheduler(shift=float(model.timestep_shift))
+    sched.set_timesteps(int(num_steps), device=device)
+
+    cfg_on = (w_text != 1.0) or (w_motion != 1.0)
+    txt_null = torch.zeros_like(txt) if cfg_on else None
+
+    def _fwd(text_e, tp_e, tv_e, mi, ts_):
+        with torch.autocast(device.type, dtype=dtype), \
+             set_forward_context(current_timestep=ts_, attn_metadata=None):
+            return model.transformer(hidden_states=mi, encoder_hidden_states=text_e,
+                                     encoder_attention_mask=mask, timestep=ts_,
+                                     encoder_hidden_states_image=clip,
+                                     track_points=tp_e, track_visibility=tv_e,
+                                     return_dict=False)
+
+    for tt in sched.timesteps:
+        mi = torch.cat([latents.to(dtype), cond20], dim=1)
+        ts_ = tt.reshape(1).to(device, dtype)
+        v_full = _fwd(txt, tp_d, tv_d, mi, ts_)
+        if not cfg_on:
+            v = v_full
+        elif w_text != 1.0 and w_motion == 1.0:
+            v_no_text = _fwd(txt_null, tp_d, tv_d, mi, ts_)
+            v = v_no_text + w_text * (v_full - v_no_text)
+        elif w_text == 1.0 and w_motion != 1.0:
+            v_no_motion = _fwd(txt, None, None, mi, ts_)
+            v = v_no_motion + w_motion * (v_full - v_no_motion)
+        else:
+            v_no_text = _fwd(txt_null, tp_d, tv_d, mi, ts_)
+            v_no_motion = _fwd(txt, None, None, mi, ts_)
+            alpha = w_text / max(w_text + w_motion, 1e-6)
+            v_base = alpha * v_no_text + (1.0 - alpha) * v_no_motion
+            v = v_base + w_text * (v_full - v_no_text) + w_motion * (v_full - v_no_motion)
+        latents = sched.step(v.float(), tt, latents.float(), return_dict=False)[0]
+
+    return twi.decode_to_pixels(model, latents)
+
+
+def draw_track_overlay(video: np.ndarray, tp: torch.Tensor, tv: torch.Tensor,
+                       info: list | None = None) -> np.ndarray:
+    """Draw kept tracks on the video. Foreground = red/blue, background anchors = gray."""
+    from PIL import Image, ImageDraw
+    T, H, W, _ = video.shape
+    tp_np = tp[0].cpu().numpy()  # [T, N, 2]
+    tv_np = tv[0].cpu().numpy()  # [T, N]
+    N = tp_np.shape[1]
+    kept = [k for k in range(N) if tv_np[:, k].max() > 0.5]
+    fg_colors = [(255, 60, 60), (60, 200, 255)]
+    bg_color = (180, 180, 180)
+    # figure out which are fg vs bg by looking at info if given
+    is_bg_by_idx = {}
+    if info is not None:
+        for entry in info:
+            is_bg_by_idx[entry["track_idx"]] = (entry["oid"] == -1)
+    out = []
+    fg_seen = 0
+    for t in range(T):
+        img = Image.fromarray(np.ascontiguousarray(video[t]))
+        dr = ImageDraw.Draw(img)
+        fg_seen = 0
+        for k in kept:
+            if tv_np[t, k] <= 0.5:
+                continue
+            x = float(tp_np[t, k, 0] * W)
+            y = float(tp_np[t, k, 1] * H)
+            if is_bg_by_idx.get(k, False):
+                col = bg_color
+                r = 3
+            else:
+                col = fg_colors[fg_seen % len(fg_colors)]
+                fg_seen += 1
+                r = 5
+            dr.ellipse([x - r, y - r, x + r, y + r], fill=col, outline=(255, 255, 255))
+        out.append(np.asarray(img))
+    return np.stack(out)
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--model-dir", required=True)
+    p.add_argument("--yaml", required=True)
+    p.add_argument("--val-parquet", required=True)
+    p.add_argument("--wandb-project", default="wantrack-bidir")
+    p.add_argument("--wandb-run-name", required=True)
+    p.add_argument("--out-dir", default="/mnt/lustre/vlm-s4duan/two_track_val")
+    p.add_argument("--w-text", type=float, default=3.0)
+    p.add_argument("--w-motion", type=float, default=1.5)
+    p.add_argument("--steps", type=int, default=30)
+    p.add_argument("--seed", type=int, default=1000)
+    p.add_argument("--fps", type=int, default=24)
+    p.add_argument("--num-fg-tracks", type=int, default=2,
+                   help="Number of foreground (moving object) tracks to keep.")
+    p.add_argument("--num-bg-tracks", type=int, default=0,
+                   help="Extra static bg anchor tracks in addition to the fg ones.")
+    p.add_argument("--swap-text-from", type=int, default=None,
+                   help="For every sample, replace text_embedding+mask with THIS sample idx's text.")
+    p.add_argument("--swap-clip-from", type=int, default=None,
+                   help="For every sample, replace clip_feature with THIS sample idx's clip.")
+    p.add_argument("--swap-ff-from", type=int, default=None,
+                   help="For every sample, replace first_frame_latent with THIS sample idx's ff.")
+    args = p.parse_args()
+
+    out_dir = Path(args.out_dir) / args.wandb_run_name
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    import wandb
+
+    print(f"[2trk] loading model {args.model_dir}", flush=True)
+    model, tc = twi.load_trackwan(args.model_dir, args.yaml)
+    text_len = int(getattr(tc.data, "text_padding_length", 256))
+    print(f"[2trk] text_len={text_len}", flush=True)
+
+    print(f"[2trk] loading val samples", flush=True)
+    samples = load_val_samples(args.val_parquet, text_len)
+    print(f"[2trk] {len(samples)} val samples", flush=True)
+
+    run = wandb.init(project=args.wandb_project, name=args.wandb_run_name,
+                     config={"model_dir": args.model_dir, "w_text": args.w_text,
+                             "w_motion": args.w_motion, "steps": args.steps})
+
+    for i, s in enumerate(samples):
+        sid = s["id"] or f"clip{i:03d}"
+        print(f"\n[2trk] === sample {i}: {sid} ({s['file_name'][:40]}) ===", flush=True)
+        # Optional cross-sample field swaps (isolate which condition drives quality)
+        if args.swap_text_from is not None:
+            src = samples[args.swap_text_from]
+            s = {**s, "text_embedding": src["text_embedding"].clone(),
+                       "text_attention_mask": src["text_attention_mask"].clone()}
+            print(f"[2trk]   SWAPPED text from sample {args.swap_text_from}", flush=True)
+        if args.swap_clip_from is not None:
+            src = samples[args.swap_clip_from]
+            s = {**s, "clip_feature": src["clip_feature"].clone()}
+            print(f"[2trk]   SWAPPED clip from sample {args.swap_clip_from}", flush=True)
+        if args.swap_ff_from is not None:
+            src = samples[args.swap_ff_from]
+            s = {**s, "first_frame_latent": src["first_frame_latent"].clone()}
+            print(f"[2trk]   SWAPPED first_frame_latent from sample {args.swap_ff_from}", flush=True)
+        tp2, tv2, info = pick_two_tracks(s, seed=args.seed + i,
+                                         num_fg=args.num_fg_tracks, num_bg=args.num_bg_tracks)
+        print(f"[2trk]   picked: {info}", flush=True)
+
+        t0 = time.time()
+        frames = sample_denoise(model, s, tp2, tv2, w_text=args.w_text, w_motion=args.w_motion,
+                                num_steps=args.steps, seed=args.seed + i)
+        dt = time.time() - t0
+        print(f"[2trk]   denoise {dt:.1f}s -> frames {frames.shape}", flush=True)
+
+        overlay = draw_track_overlay(frames, tp2, tv2, info)
+        fn_raw = out_dir / f"{sid}_gen.mp4"
+        fn_ov = out_dir / f"{sid}_gen_overlay.mp4"
+        imageio.mimsave(str(fn_raw), frames, fps=args.fps, macro_block_size=1)
+        imageio.mimsave(str(fn_ov), overlay, fps=args.fps, macro_block_size=1)
+
+        # GT for reference
+        try:
+            ref = twi.decode_reference(model, s["vae_latent"].to(model.device, torch.bfloat16))
+            ref_ov = draw_track_overlay(ref, tp2, tv2, info)
+            fn_ref = out_dir / f"{sid}_gt_overlay.mp4"
+            imageio.mimsave(str(fn_ref), ref_ov, fps=args.fps, macro_block_size=1)
+            wandb.log({
+                f"sample_{i:02d}_gen": wandb.Video(str(fn_ov), fps=args.fps, format="mp4"),
+                f"sample_{i:02d}_gt": wandb.Video(str(fn_ref), fps=args.fps, format="mp4"),
+                "sample": i, "caption": s["caption"][:200], "picked": str(info),
+            })
+        except Exception as e:
+            print(f"[2trk]   gt fail: {e}", flush=True)
+            wandb.log({
+                f"sample_{i:02d}_gen": wandb.Video(str(fn_ov), fps=args.fps, format="mp4"),
+                "sample": i, "caption": s["caption"][:200], "picked": str(info),
+            })
+
+    wandb.finish()
+    print("[2trk] DONE", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data_pipeline/upload_parquets.sh

@@ -0,0 +1,79 @@
+#!/bin/bash
+# Publish the processed (parquet) dataset to a HuggingFace dataset repo as a directory tree,
+# mirroring noctuashap/openvid-wantrack-processed's layout (raw parquet files, not tarred).
+# Uses `hf upload-large-folder` -- resumable and built for multi-TB uploads: re-running skips
+# files already on the Hub, so a killed upload just continues.
+#
+# Usage:
+#   REPO=FastVideo/openvid-wantrack-processed-v2 bash data_pipeline/upload_parquets.sh
+#   ... DRY_RUN=1 ...      # verify completeness + print the command, upload nothing
+set -uo pipefail
+
+REPO=${REPO:?set REPO=<owner>/<name>}
+PARQUET_ROOT=${PARQUET_ROOT:-/home/hal-shared/motionstream/data/openvid-wantrack-parquets}
+SRC_ROOT=${SRC_ROOT:-/home/hal-shared/motionstream/data/openvid-wantrack}
+PRIVATE=${PRIVATE:-1}
+NUM_WORKERS=${NUM_WORKERS:-8}
+README=${README:-data_pipeline/notes/processed_dataset_README.md}
+DRY_RUN=${DRY_RUN:-0}
+
+cd "$(dirname "$0")/.."
+HF=$(command -v hf || command -v huggingface-cli) || { echo "[pub] ERROR: hf CLI not found" >&2; exit 1; }
+
+# --- safety: verify the shards PRESENT under PARQUET_ROOT are duplicate-free before publishing.
+# Validates whatever is present (no hardcoded 260, no source-video cross-check), so partial /
+# derivative sets like the bf16 copy work; still catches real corruption via duplicate ids.
+# Set VERIFY=0 to skip verification entirely.
+if [[ "${VERIFY:-1}" == "1" ]]; then
+echo "[pub] verifying present shards are duplicate-free before upload ..."
+python - "$PARQUET_ROOT" <<'PY'
+import glob, os, sys
+import pyarrow.parquet as pq
+pbase = sys.argv[1]
+shards = sorted(d for d in os.listdir(pbase) if d.startswith("shard"))
+bad=[]; total=0; nsh=0
+for s in shards:
+    fs=glob.glob(f"{pbase}/{s}/**/*.parquet", recursive=True)
+    if not fs:
+        continue
+    ids=[i for f in fs for i in pq.read_table(f, columns=["id"]).column("id").to_pylist()]
+    if len(ids)>0 and len(ids)==len(set(ids)):
+        total+=len(ids); nsh+=1
+    else:
+        bad.append((s, len(ids), len(set(ids))))
+if bad:
+    print(f"[pub] REFUSING: {len(bad)} shard(s) empty/duplicate-id: {bad[:8]}")
+    sys.exit(1)
+print(f"[pub] OK: {nsh} shard(s) present, {total:,} clips (no duplicate ids)")
+PY
+[[ $? -eq 0 ]] || { echo "[pub] aborted -- fix the shards above, then re-run" >&2; exit 1; }
+else
+  echo "[pub] VERIFY=0 -> skipping shard verification"
+fi
+
+n_files=$(find "$PARQUET_ROOT" -name '*.parquet' | wc -l)
+size=$(du -sh "$PARQUET_ROOT" 2>/dev/null | cut -f1)
+echo "[pub] repo=$REPO  files=$n_files  size=$size  private=$PRIVATE"
+
+if [[ "$DRY_RUN" == "1" ]]; then
+    echo "[pub] DRY RUN -- would run:"
+    echo "  $HF upload-large-folder $REPO $PARQUET_ROOT --repo-type dataset --include '*.parquet' --num-workers $NUM_WORKERS"
+    echo "  (+ README.md upload)"
+    exit 0
+fi
+
+# --- create repo + upload README once -------------------------------------------------
+vis=(); [[ "$PRIVATE" == "1" ]] && vis=(--private)
+"$HF" repo create "$REPO" --repo-type dataset "${vis[@]}" 2>/dev/null \
+    && echo "[pub] created $REPO" || echo "[pub] repo exists (ok)"
+[[ -f "$README" ]] && "$HF" upload "$REPO" "$README" README.md --repo-type dataset >/dev/null 2>&1 \
+    && echo "[pub] README uploaded"
+
+# --- upload the parquet tree (resumable) ----------------------------------------------
+# --include '*.parquet' skips any stray files; the shard*/combined_parquet_dataset/... tree
+# is preserved in the repo. Re-run this exact command to resume after any interruption.
+echo "[pub] uploading parquet tree (resumable; re-run to continue if interrupted) ..."
+"$HF" upload-large-folder "$REPO" "$PARQUET_ROOT" \
+    --repo-type dataset --include '*.parquet' --num-workers "$NUM_WORKERS"
+
+echo "[pub] done -> https://huggingface.co/datasets/$REPO"

data_pipeline/upload_tracks.sh

@@ -0,0 +1,115 @@
+#!/bin/bash
+# Package each shard's tracks/ into a tar and upload to a HuggingFace dataset repo,
+# mirroring noctuashap/openvid-wantrack-tracks layout: one tars-NNNNN.tar per shard,
+# each holding ~1000 .npz (flat, no directory prefix).
+#
+# Only shards whose progress.json marks tracks done are packaged. Resumable: a shard
+# already present in the repo (checked via the HF API) is skipped.
+#
+# Usage:
+#   REPO=<user-or-org>/<name> SHARDS=0-170 bash data_pipeline/upload_tracks.sh
+#   REPO=FastVideo/openvid-wantrack-tracks-v2 SHARDS=0-170 PRIVATE=1 bash data_pipeline/upload_tracks.sh
+#   ... DRY_RUN=1 ...    # build tars + report, do NOT create repo or upload
+set -uo pipefail
+
+REPO=${REPO:?set REPO=<owner>/<name>}
+SHARDS=${SHARDS:-0-259}
+DATA_ROOT_BASE=${DATA_ROOT_BASE:-/home/hal-shared/motionstream/data/openvid-wantrack/shard}
+STAGING=${STAGING:-/home/hal-shared/motionstream/data/openvid-wantrack/_upload_tars}
+PRIVATE=${PRIVATE:-1}          # create the repo private by default; you flip it public in the UI
+KEEP_TARS=${KEEP_TARS:-0}      # 1 = keep local tar after upload (default: delete to save disk)
+DRY_RUN=${DRY_RUN:-0}
+PREFIX=${PREFIX:-tracks}       # tar basename: ${PREFIX}-00042.tar
+README=${README:-data_pipeline/notes/tracks_dataset_README.md}  # uploaded as README.md if present
+
+cd "$(dirname "$0")/.."
+shard_root() { printf "%s%03d" "$DATA_ROOT_BASE" "$1"; }
+mkdir -p "$STAGING"
+
+# expand "0-9,20,30-35"
+expand() {
+    local tok lo hi out=(); IFS=',' read -ra toks <<< "$1"
+    for tok in "${toks[@]}"; do
+        if [[ "$tok" =~ ^([0-9]+)-([0-9]+)$ ]]; then
+            for ((i=${BASH_REMATCH[1]}; i<=${BASH_REMATCH[2]}; i++)); do out+=("$i"); done
+        elif [[ "$tok" =~ ^[0-9]+$ ]]; then out+=("$tok")
+        else echo "[upload] ERROR: bad SHARDS token '$tok'" >&2; exit 1; fi
+    done
+    printf '%s\n' "${out[@]}"
+}
+mapfile -t LIST < <(expand "$SHARDS")
+
+tracks_done() {  # tracks_done <n>
+    python - "$(shard_root "$1")/progress.json" <<'PY'
+import json,sys
+from pathlib import Path
+p=Path(sys.argv[1])
+sys.exit(0 if p.exists() and json.loads(p.read_text()).get("phases",{}).get("tracks",{}).get("done") else 1)
+PY
+}
+
+# --- ensure repo exists (unless dry run) -------------------------------------------
+if [[ "$DRY_RUN" != "1" ]]; then
+    vis=(); [[ "$PRIVATE" == "1" ]] && vis=(--private)
+    hf repo create "$REPO" --repo-type dataset "${vis[@]}" 2>/dev/null \
+        && echo "[upload] created dataset repo $REPO" \
+        || echo "[upload] repo $REPO already exists (ok)"
+    # names of files already in the repo (via the hub API), to skip re-upload on resume
+    mapfile -t REMOTE < <(python - "$REPO" <<'PY'
+import sys
+from huggingface_hub import HfApi
+try:
+    print("\n".join(HfApi().list_repo_files(sys.argv[1], repo_type="dataset")))
+except Exception:
+    pass
+PY
+)
+    remote_has() { printf '%s\n' "${REMOTE[@]:-}" | grep -qx "$1"; }
+
+    # upload README once (if present and not already there)
+    if [[ -f "$README" ]] && ! remote_has "README.md"; then
+        hf upload "$REPO" "$README" "README.md" --repo-type dataset >/dev/null 2>&1 \
+            && echo "[upload] README.md uploaded" || echo "[upload] WARN: README upload failed" >&2
+    fi
+else
+    remote_has() { return 1; }
+fi
+
+n_up=0 n_skip=0 n_todo=0
+T0=$(date +%s)
+for s in "${LIST[@]}"; do
+    root="$(shard_root "$s")"
+    tar_name=$(printf "%s-%05d.tar" "$PREFIX" "$s")
+    if ! tracks_done "$s"; then continue; fi
+    n_todo=$((n_todo+1))
+    if remote_has "$tar_name"; then n_skip=$((n_skip+1)); echo "[upload] $tar_name already in repo, skip"; continue; fi
+
+    ntracks=$(ls "$root"/tracks/*.npz 2>/dev/null | wc -l || echo 0)
+    [[ "$ntracks" -gt 0 ]] || { echo "[upload] WARN shard $s: no npz despite progress=done, skip" >&2; continue; }
+
+    tar_path="$STAGING/$tar_name"
+    # Clip names start with '---', so a glob/ls would feed tar filenames it reads as options.
+    # find -print0 | tar --null -T - is dash-safe; paths come out './name' (matches the
+    # reference repo's layout). Exclude any leftover *.tmp.npz from an interrupted write.
+    echo "[upload] packing shard $s: $ntracks npz -> $tar_name"
+    ( cd "$root/tracks" && find . -maxdepth 1 -type f -name '*.npz' ! -name '*.tmp.npz' -print0 ) \
+        | tar -cf "$tar_path" --null -C "$root/tracks" -T -
+    sz=$(du -h "$tar_path" 2>/dev/null | cut -f1)
+
+    if [[ "$DRY_RUN" == "1" ]]; then
+        echo "[upload] DRY_RUN: built $tar_name ($sz), not uploading"
+        [[ "$KEEP_TARS" == "1" ]] || rm -f "$tar_path"
+        continue
+    fi
+
+    if hf upload "$REPO" "$tar_path" "$tar_name" --repo-type dataset >/dev/null 2>&1; then
+        n_up=$((n_up+1)); echo "[upload] shard $s -> $tar_name ($sz) uploaded"
+        [[ "$KEEP_TARS" == "1" ]] || rm -f "$tar_path"
+    else
+        echo "[upload] ERROR uploading $tar_name (kept at $tar_path)" >&2
+    fi
+done
+
+echo "[upload] done in $(( ($(date +%s)-T0)/60 ))m: $n_up uploaded, $n_skip already present, $n_todo eligible"
+[[ "$DRY_RUN" == "1" ]] && echo "[upload] (dry run: repo not created, nothing uploaded)"
+echo "[upload] repo: https://huggingface.co/datasets/$REPO"

data_pipeline/val_compare.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Re-run validation as a clean GT-vs-generated side-by-side (no track overlay).
+
+For each clip: decode the GT clip, generate with its GT tracks, and write a
+left=GT / right=generated side-by-side video + a first-frame comparison PNG, so
+the first-frame match and overall reconstruction are easy to eyeball.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+import imageio.v2 as imageio
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(__file__))
+import trackwan_infer as twi  # noqa: E402
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--export", required=True)
+    p.add_argument("--yaml", default="examples/train/scenario/worldmodel/finetune_wantrack_i2v.yaml")
+    p.add_argument("--data", default="/mnt/weka/home/hao.zhang/shao/data/motion_pipeline/"
+                   "wan22_a14b_720p_24fps/preprocessed_i2v_track/combined_parquet_dataset")
+    p.add_argument("--out", default="/mnt/weka/home/hao.zhang/shao/data/motion_pipeline/val_compare_4k")
+    p.add_argument("--fig", default="research_log/figures")
+    p.add_argument("--clips", type=int, nargs="+", default=[0, 1])
+    p.add_argument("--steps", type=int, default=30)
+    p.add_argument("--seed", type=int, default=1000)
+    args = p.parse_args()
+    os.makedirs(args.out, exist_ok=True)
+    os.makedirs(args.fig, exist_ok=True)
+
+    model, tc = twi.load_trackwan(args.export, args.yaml)
+    text_len = int(tc.pipeline_config.text_encoder_configs[0].arch_config.text_len)
+    ratio = int(tc.pipeline_config.vae_config.arch_config.temporal_compression_ratio)
+    samples = twi.load_conditioning_from_parquet(args.data, args.clips, text_len)
+
+    for ci, s in zip(args.clips, samples):
+        Tpx = (s["first_frame_latent"].shape[2] - 1) * ratio + 1
+        gt = twi.decode_reference(model, s["vae_latent"])                       # [T,H,W,3]
+        lat = twi.generate(model, first_frame_latent=s["first_frame_latent"],
+                           text_embedding=s["text_embedding"], text_attention_mask=s["text_attention_mask"],
+                           track_points=s["track_points"][:, :Tpx], track_visibility=s["track_visibility"][:, :Tpx],
+                           num_steps=args.steps, seed=args.seed)
+        gen = twi.decode_to_pixels(model, lat)
+        T = min(len(gt), len(gen))
+        side = np.concatenate([gt[:T], gen[:T]], axis=2)                        # [T,H,2W,3] (GT|gen)
+        imageio.mimsave(os.path.join(args.out, f"val_clip{ci}_GTleft_GENright.mp4"),
+                        list(side), fps=24, macro_block_size=1)
+        ff = np.concatenate([gt[0], gen[0]], axis=1)
+        imageio.imwrite(os.path.join(args.fig, f"val_firstframe_clip{ci}_GTleft_GENright.png"), ff)
+        mse0 = float(((gt[0].astype(np.float32) - gen[0].astype(np.float32)) ** 2).mean())
+        print(f"[clip {ci}] first-frame MSE(GT,gen)={mse0:8.1f}  -> {args.out}/val_clip{ci}_GTleft_GENright.mp4", flush=True)
+
+
+if __name__ == "__main__":
+    main()