583 lines
25 KiB
Python
583 lines
25 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""
|
|||
|
|
SNAC tokenization for FineVideo-VLA activities.
|
|||
|
|
|
|||
|
|
Reads final_dataset_adaptive JSONL files, extracts audio from .mp4 videos,
|
|||
|
|
and tokenizes each activity segment with SNAC_24kHz in "listen" format.
|
|||
|
|
|
|||
|
|
Listen format — 3 tokens per SNAC base frame (base rate = 12.5 Hz → 37.5 tokens/sec):
|
|||
|
|
token_1 = codes[0][i] + 128266 → <snac_128266> .. <snac_132361>
|
|||
|
|
token_2 = codes[1][2*i] + 128266 + 4096 → <snac_132362> .. <snac_136457>
|
|||
|
|
token_3 = codes[1][2*i+1] + 128266 +16384 → <snac_144650> .. <snac_148745>
|
|||
|
|
|
|||
|
|
Same Orpheus offset scheme as MixtureVitae-Omni → tokens are directly compatible.
|
|||
|
|
Total unique SNAC token strings: 3 × 4096 = 12,288.
|
|||
|
|
|
|||
|
|
Output: {OUTPUT_DIR}/{video_id}_snac.jsonl
|
|||
|
|
One line per activity:
|
|||
|
|
{
|
|||
|
|
"video_id":"...", "activity_id":"...", "start_sec":1.0, "end_sec":8.9,
|
|||
|
|
"has_agent": true,
|
|||
|
|
"snac_by_chunk": {
|
|||
|
|
"0": ["<snac_130055>", "<snac_133001>", "<snac_145000>", ...], // ~9-10 tokens
|
|||
|
|
"1": [...],
|
|||
|
|
...
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
snac_by_chunk keys are chunk_idx (integer as string), aligned to the same 8-frame
|
|||
|
|
grid as cosmos/avclm/agent. Phase7 reads chunk_idx → snac tokens directly.
|
|||
|
|
|
|||
|
|
Two modes:
|
|||
|
|
--build-tasks Scan all final_dataset_adaptive files, write snac_task_list.json.
|
|||
|
|
Run once on login node (or task 0) before the array job.
|
|||
|
|
(default) Load task list, process this SLURM task's slice of videos.
|
|||
|
|
|
|||
|
|
SLURM usage:
|
|||
|
|
SLURM_ARRAY_TASK_ID = task index (0-based)
|
|||
|
|
SLURM_ARRAY_TASK_COUNT = total number of tasks in the array
|
|||
|
|
|
|||
|
|
Local test:
|
|||
|
|
python pipeline_pose/snac_finevideo.py --build-tasks
|
|||
|
|
python pipeline_pose/snac_finevideo.py # task_id=0, num_tasks=1
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import argparse
|
|||
|
|
import glob
|
|||
|
|
import json
|
|||
|
|
import logging
|
|||
|
|
import math
|
|||
|
|
import multiprocessing
|
|||
|
|
import os
|
|||
|
|
import subprocess
|
|||
|
|
import sys
|
|||
|
|
import time
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
# ── Paths ─────────────────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
VIDEO_DIR = "/e/data1/datasets/playground/mmlaion/shared/nguyen38/videos_staging"
|
|||
|
|
# Still scanned for activity time_range_sec + has_agent -- those don't depend
|
|||
|
|
# on window size. chunk_timing (per-window breakdown) is IGNORED (see
|
|||
|
|
# _scan_one_rank_file docstring) because this file is window=8-based (the
|
|||
|
|
# pre-pivot merge) and Phase 6 hasn't rerun at window=24 yet; n_chunks is
|
|||
|
|
# instead recomputed independently from CHUNK_SIZE below, same pattern as
|
|||
|
|
# data_prep/omnivideo_100k/snac_omnivideo.py.
|
|||
|
|
INPUT_GLOB = ("/e/data1/datasets/playground/mmlaion/shared/nguyen38/FineVideo-VLA/"
|
|||
|
|
"final_dataset_adaptive/final_vla_adaptive_rank_*.jsonl")
|
|||
|
|
OUTPUT_DIR = ("/e/data1/datasets/playground/mmlaion/shared/nguyen38/FineVideo-VLA/snac_tokens_w24")
|
|||
|
|
TASK_CACHE = ("/e/data1/datasets/playground/mmlaion/shared/nguyen38/FineVideo-VLA/"
|
|||
|
|
"snac_task_list_w24.json")
|
|||
|
|
HF_CACHE = "/e/project1/reformo/nguyen38/jupiter_cache/huggingface"
|
|||
|
|
SNAC_MODEL = "hubertsiuzdak/snac_24khz"
|
|||
|
|
SAMPLE_RATE = 24000
|
|||
|
|
TARGET_FPS = 30
|
|||
|
|
CHUNK_SIZE = 24 # 2026-07-23 window=24 pivot -- must match step_a_tokenize_video.py's CHUNK_SIZE
|
|||
|
|
|
|||
|
|
# ── SNAC listen-format offsets (matches MixtureVitae-Omni) ───────────────────
|
|||
|
|
|
|||
|
|
OFFSET_L0 = 128266 # codes[0] base
|
|||
|
|
OFFSET_L1A = 128266 + 4096 # codes[1] even frames → 132362
|
|||
|
|
OFFSET_L1B = 128266 + 4 * 4096 # codes[1] odd frames → 144650
|
|||
|
|
|
|||
|
|
# 2026-07-23: speak-format offsets for codes[2] (fine, 50Hz, 4 sub-positions
|
|||
|
|
# per base frame) -- corrected to match the REAL scheme in Huu/Chien's
|
|||
|
|
# production snac_gpu.py on Leonardo (pipeline_video/snac_gpu.py), the
|
|||
|
|
# Orpheus-standard SNAC packing layout. Sub-codes 0/1 sit between L1A and
|
|||
|
|
# L1B (in what was wrongly assumed to be an unused gap); sub-codes 2/3 sit
|
|||
|
|
# after L1B. See data_prep/laion_emotional_roleplay/tokenize_snac.py's
|
|||
|
|
# OFFSET_L2 docstring for the full correction history. Must match exactly.
|
|||
|
|
OFFSET_L2 = [136458, 140554, 148746, 152842]
|
|||
|
|
|
|||
|
|
# ── Logging ───────────────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
logging.basicConfig(
|
|||
|
|
level=logging.INFO,
|
|||
|
|
format="%(asctime)s %(levelname)s %(message)s",
|
|||
|
|
datefmt="%H:%M:%S",
|
|||
|
|
stream=sys.stdout,
|
|||
|
|
)
|
|||
|
|
log = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|||
|
|
# Audio extraction
|
|||
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def _find_ffmpeg() -> str:
|
|||
|
|
"""Return path to ffmpeg binary, trying imageio_ffmpeg as fallback."""
|
|||
|
|
try:
|
|||
|
|
subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
|
|||
|
|
return "ffmpeg"
|
|||
|
|
except (FileNotFoundError, subprocess.CalledProcessError):
|
|||
|
|
pass
|
|||
|
|
try:
|
|||
|
|
import imageio_ffmpeg
|
|||
|
|
return imageio_ffmpeg.get_ffmpeg_exe()
|
|||
|
|
except ImportError:
|
|||
|
|
pass
|
|||
|
|
raise RuntimeError("ffmpeg not found. Install ffmpeg or imageio_ffmpeg.")
|
|||
|
|
|
|||
|
|
_FFMPEG = None
|
|||
|
|
|
|||
|
|
def get_ffmpeg() -> str:
|
|||
|
|
global _FFMPEG
|
|||
|
|
if _FFMPEG is None:
|
|||
|
|
_FFMPEG = _find_ffmpeg()
|
|||
|
|
return _FFMPEG
|
|||
|
|
|
|||
|
|
|
|||
|
|
def extract_full_audio(video_path: str) -> np.ndarray | None:
|
|||
|
|
"""
|
|||
|
|
Extract full mono 24 kHz PCM audio from a video file.
|
|||
|
|
Pipes raw float32 PCM from ffmpeg directly to a numpy array — no temp files.
|
|||
|
|
Returns float32 array or None on any failure.
|
|||
|
|
"""
|
|||
|
|
cmd = [
|
|||
|
|
get_ffmpeg(), "-y",
|
|||
|
|
"-i", video_path,
|
|||
|
|
"-vn", # strip video
|
|||
|
|
"-ac", "1", # mono
|
|||
|
|
"-ar", str(SAMPLE_RATE), # 24 kHz
|
|||
|
|
"-f", "f32le", # raw float32 PCM
|
|||
|
|
"-",
|
|||
|
|
]
|
|||
|
|
try:
|
|||
|
|
result = subprocess.run(cmd, capture_output=True, timeout=300)
|
|||
|
|
if result.returncode != 0 or not result.stdout:
|
|||
|
|
return None
|
|||
|
|
audio = np.frombuffer(result.stdout, dtype=np.float32).copy()
|
|||
|
|
return audio if len(audio) > 0 else None
|
|||
|
|
except Exception as e:
|
|||
|
|
log.debug(f"ffmpeg failed for {video_path}: {e}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def slice_audio(
|
|||
|
|
audio: np.ndarray,
|
|||
|
|
start_sec: float,
|
|||
|
|
end_sec: float,
|
|||
|
|
sr: int = SAMPLE_RATE,
|
|||
|
|
) -> np.ndarray:
|
|||
|
|
"""Slice a float32 audio array to [start_sec, end_sec]."""
|
|||
|
|
s = max(0, int(start_sec * sr))
|
|||
|
|
e = min(len(audio), int(end_sec * sr))
|
|||
|
|
return audio[s:e]
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|||
|
|
# SNAC tokenization
|
|||
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def encode_listen(audio: np.ndarray, model, device: str) -> list[str]:
|
|||
|
|
"""
|
|||
|
|
Encode a float32 audio array with SNAC_24kHz, return listen-format tokens.
|
|||
|
|
|
|||
|
|
Listen format (3 tokens per base frame):
|
|||
|
|
<snac_{codes[0][i] + 128266}>
|
|||
|
|
<snac_{codes[1][2i] + 132362}>
|
|||
|
|
<snac_{codes[1][2i+1] + 144650}>
|
|||
|
|
|
|||
|
|
SNAC_24kHz hierarchy:
|
|||
|
|
codes[0] — base codebook, 12.5 Hz
|
|||
|
|
codes[1] — mid codebook, 25.0 Hz (2× codes[0])
|
|||
|
|
codes[2] — fine codebook, 50.0 Hz (4× codes[0], not used in listen)
|
|||
|
|
|
|||
|
|
Listen format ignores codes[2] (fine detail) to keep token count low
|
|||
|
|
(~37.5 tokens/sec vs 87.5 for full speak format). Matches MV-Omni.
|
|||
|
|
"""
|
|||
|
|
tensor = torch.from_numpy(audio).unsqueeze(0).unsqueeze(0).to(device) # (1,1,T)
|
|||
|
|
with torch.inference_mode():
|
|||
|
|
codes = model.encode(tensor) # list: [codes[0], codes[1], codes[2]]
|
|||
|
|
|
|||
|
|
c0 = codes[0] # (1, N0)
|
|||
|
|
c1 = codes[1] # (1, N1), N1 == 2*N0
|
|||
|
|
|
|||
|
|
n0 = c0.shape[1]
|
|||
|
|
tokens: list[str] = []
|
|||
|
|
for i in range(n0):
|
|||
|
|
i1a = 2 * i
|
|||
|
|
i1b = 2 * i + 1
|
|||
|
|
if i1b >= c1.shape[1]:
|
|||
|
|
break # boundary guard: shouldn't happen for valid audio
|
|||
|
|
tokens.append(f"<snac_{c0[0, i].item() + OFFSET_L0}>")
|
|||
|
|
tokens.append(f"<snac_{c1[0, i1a].item() + OFFSET_L1A}>")
|
|||
|
|
tokens.append(f"<snac_{c1[0, i1b].item() + OFFSET_L1B}>")
|
|||
|
|
return tokens
|
|||
|
|
|
|||
|
|
|
|||
|
|
def encode_speak(audio: np.ndarray, model, device: str) -> list[str]:
|
|||
|
|
"""
|
|||
|
|
Encode a float32 audio array with SNAC_24kHz, return full speak-format
|
|||
|
|
tokens (7 tokens per base frame, +133% vs encode_listen()) -- 2026-07-22
|
|||
|
|
(REPORT.md #37), decided after a real audio A/B
|
|||
|
|
(tools/snac_l2_experiment.py) showed audibly better reconstruction.
|
|||
|
|
|
|||
|
|
Speak format (2026-07-23, Leo-matched order -- L2 interleaved between
|
|||
|
|
L1a and L1b, not appended after):
|
|||
|
|
<snac_{codes[0][i] + OFFSET_L0}>
|
|||
|
|
<snac_{codes[1][2i] + OFFSET_L1A}>
|
|||
|
|
<snac_{codes[2][4i] + OFFSET_L2[0]}>
|
|||
|
|
<snac_{codes[2][4i+1] + OFFSET_L2[1]}>
|
|||
|
|
<snac_{codes[1][2i+1] + OFFSET_L1B}>
|
|||
|
|
<snac_{codes[2][4i+2] + OFFSET_L2[2]}>
|
|||
|
|
<snac_{codes[2][4i+3] + OFFSET_L2[3]}>
|
|||
|
|
"""
|
|||
|
|
tensor = torch.from_numpy(audio).unsqueeze(0).unsqueeze(0).to(device)
|
|||
|
|
with torch.inference_mode():
|
|||
|
|
codes = model.encode(tensor)
|
|||
|
|
|
|||
|
|
c0, c1, c2 = codes[0], codes[1], codes[2]
|
|||
|
|
n0 = c0.shape[1]
|
|||
|
|
tokens: list[str] = []
|
|||
|
|
for i in range(n0):
|
|||
|
|
i1a, i1b = 2 * i, 2 * i + 1
|
|||
|
|
i2 = [4 * i + k for k in range(4)]
|
|||
|
|
if i1b >= c1.shape[1] or i2[-1] >= c2.shape[1]:
|
|||
|
|
break
|
|||
|
|
tokens.append(f"<snac_{c0[0, i].item() + OFFSET_L0}>")
|
|||
|
|
tokens.append(f"<snac_{c1[0, i1a].item() + OFFSET_L1A}>")
|
|||
|
|
tokens.append(f"<snac_{c2[0, i2[0]].item() + OFFSET_L2[0]}>")
|
|||
|
|
tokens.append(f"<snac_{c2[0, i2[1]].item() + OFFSET_L2[1]}>")
|
|||
|
|
tokens.append(f"<snac_{c1[0, i1b].item() + OFFSET_L1B}>")
|
|||
|
|
tokens.append(f"<snac_{c2[0, i2[2]].item() + OFFSET_L2[2]}>")
|
|||
|
|
tokens.append(f"<snac_{c2[0, i2[3]].item() + OFFSET_L2[3]}>")
|
|||
|
|
return tokens
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|||
|
|
# Task list building (pre-processing step)
|
|||
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def _scan_one_rank_file(fpath: str) -> dict:
|
|||
|
|
"""
|
|||
|
|
Scan one final_dataset_adaptive rank file for activity time boundaries.
|
|||
|
|
Returns {video_id: [activity_dict, ...]} for ALL activities with a valid
|
|||
|
|
time_range_sec.
|
|||
|
|
|
|||
|
|
Each activity dict: activity_id, start_sec, end_sec, has_agent.
|
|||
|
|
|
|||
|
|
2026-07-23: no longer reads chunk_timing (this file is the pre-window=24-
|
|||
|
|
pivot merge; its chunk_timing reflects the OLD 8-frame grid). n_chunks is
|
|||
|
|
instead recomputed in process_video() from (end_sec-start_sec) and the
|
|||
|
|
current CHUNK_SIZE=24 -- same independent-recompute pattern
|
|||
|
|
data_prep/omnivideo_100k/snac_omnivideo.py uses, so this script no longer
|
|||
|
|
needs to wait on Phase 6 merge to re-run at window=24 before it can align
|
|||
|
|
correctly.
|
|||
|
|
|
|||
|
|
We tokenize ALL activities (not just agent) because:
|
|||
|
|
- Non-agent activities have seed2+cosmos → seed2+cosmos+snac trains modality transitions
|
|||
|
|
- Agent-only activities = only 14% of total; skipping the rest wastes 86% of this GPU run
|
|||
|
|
"""
|
|||
|
|
tasks: dict = {}
|
|||
|
|
try:
|
|||
|
|
with open(fpath, "r", errors="replace") as f:
|
|||
|
|
for line in f:
|
|||
|
|
line = line.strip()
|
|||
|
|
if not line:
|
|||
|
|
continue
|
|||
|
|
try:
|
|||
|
|
rec = json.loads(line)
|
|||
|
|
except json.JSONDecodeError:
|
|||
|
|
continue
|
|||
|
|
vid = rec.get("video_id", "")
|
|||
|
|
if not vid:
|
|||
|
|
continue
|
|||
|
|
for scene in rec.get("scenes", []):
|
|||
|
|
for act in scene.get("activities", []):
|
|||
|
|
tr = act.get("time_range_sec")
|
|||
|
|
if not tr or len(tr) < 2:
|
|||
|
|
continue
|
|||
|
|
has_agent = "<agent>" in act.get("video_tokens", "")
|
|||
|
|
tasks.setdefault(vid, []).append({
|
|||
|
|
"activity_id": act.get("activity_id", ""),
|
|||
|
|
"start_sec": float(tr[0]),
|
|||
|
|
"end_sec": float(tr[1]),
|
|||
|
|
"has_agent": has_agent,
|
|||
|
|
})
|
|||
|
|
except Exception as e:
|
|||
|
|
log.warning(f"Error scanning {fpath}: {e}")
|
|||
|
|
return tasks
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_task_list(input_glob: str, cache_path: str, workers: int = 8) -> dict:
|
|||
|
|
"""
|
|||
|
|
Scan all final_dataset_adaptive rank files in parallel to build a task list.
|
|||
|
|
Saves result to cache_path as JSON.
|
|||
|
|
Returns {video_id: [activity_dicts]}.
|
|||
|
|
|
|||
|
|
This is I/O-heavy (~657 GB total) — use multiprocessing to parallelize.
|
|||
|
|
Estimated wall time: 5–15 min with 8 workers on shared filesystem.
|
|||
|
|
"""
|
|||
|
|
rank_files = sorted(glob.glob(input_glob))
|
|||
|
|
if not rank_files:
|
|||
|
|
raise FileNotFoundError(f"No files matched: {input_glob}")
|
|||
|
|
log.info(f"Scanning {len(rank_files)} rank files with {workers} workers...")
|
|||
|
|
t0 = time.time()
|
|||
|
|
|
|||
|
|
with multiprocessing.Pool(workers) as pool:
|
|||
|
|
partial_results = pool.map(_scan_one_rank_file, rank_files)
|
|||
|
|
|
|||
|
|
# merge
|
|||
|
|
all_tasks: dict = {}
|
|||
|
|
for partial in partial_results:
|
|||
|
|
for vid, acts in partial.items():
|
|||
|
|
all_tasks.setdefault(vid, []).extend(acts)
|
|||
|
|
|
|||
|
|
# deduplicate activities by activity_id (in case of overlap across ranks)
|
|||
|
|
for vid in all_tasks:
|
|||
|
|
seen = set()
|
|||
|
|
deduped = []
|
|||
|
|
for act in all_tasks[vid]:
|
|||
|
|
key = act["activity_id"]
|
|||
|
|
if key not in seen:
|
|||
|
|
seen.add(key)
|
|||
|
|
deduped.append(act)
|
|||
|
|
all_tasks[vid] = deduped
|
|||
|
|
|
|||
|
|
log.info(
|
|||
|
|
f"Task list built: {len(all_tasks)} videos, "
|
|||
|
|
f"{sum(len(v) for v in all_tasks.values())} activities "
|
|||
|
|
f"({time.time()-t0:.0f}s)"
|
|||
|
|
)
|
|||
|
|
os.makedirs(os.path.dirname(cache_path), exist_ok=True)
|
|||
|
|
with open(cache_path, "w") as f:
|
|||
|
|
json.dump(all_tasks, f)
|
|||
|
|
log.info(f"Saved task list → {cache_path}")
|
|||
|
|
return all_tasks
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|||
|
|
# Chunk alignment
|
|||
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def split_snac_by_chunks(tokens: list[str], n_chunks: int) -> dict[int, list[str]]:
|
|||
|
|
"""
|
|||
|
|
Split a flat SNAC listen token list evenly across n_chunks video chunks.
|
|||
|
|
|
|||
|
|
Why: SNAC rate (12.5 Hz base × 3 tokens = 37.5 tok/s) does not divide evenly
|
|||
|
|
by the video chunk rate (30fps / 8 = 3.75 Hz). Encoding the full activity once
|
|||
|
|
preserves audio context; then we split by chunk count snapping to 3-token
|
|||
|
|
boundaries (one SNAC base frame = 3 listen tokens).
|
|||
|
|
|
|||
|
|
Per 8-frame chunk at 30fps (0.267s): ~3.33 SNAC base frames → 9–10 listen tokens.
|
|||
|
|
"""
|
|||
|
|
n_tokens = len(tokens)
|
|||
|
|
n_base = n_tokens // 3 # truncate to complete base frames
|
|||
|
|
tokens = tokens[:n_base * 3]
|
|||
|
|
|
|||
|
|
result: dict[int, list[str]] = {}
|
|||
|
|
for k in range(n_chunks):
|
|||
|
|
start_base = round(k * n_base / n_chunks)
|
|||
|
|
end_base = round((k + 1) * n_base / n_chunks)
|
|||
|
|
result[k] = tokens[start_base * 3 : end_base * 3]
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|||
|
|
# Per-video processing
|
|||
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def process_video(
|
|||
|
|
video_id: str,
|
|||
|
|
activities: list[dict],
|
|||
|
|
model,
|
|||
|
|
device: str,
|
|||
|
|
video_dir: str,
|
|||
|
|
output_dir: str,
|
|||
|
|
skip_existing: bool,
|
|||
|
|
encode_fn=encode_listen,
|
|||
|
|
) -> dict:
|
|||
|
|
"""
|
|||
|
|
Tokenize all activities for one video.
|
|||
|
|
|
|||
|
|
Steps:
|
|||
|
|
1. Check skip — if output file exists and skip_existing, return immediately.
|
|||
|
|
2. Extract full audio from .mp4 once (1 ffmpeg call per video).
|
|||
|
|
3. For each activity: slice audio by time_range_sec, run SNAC encode once,
|
|||
|
|
then split the flat token list across chunks (preserving audio context).
|
|||
|
|
4. Write all results to {output_dir}/{video_id}_snac.jsonl.
|
|||
|
|
|
|||
|
|
Output per activity: snac_by_chunk {chunk_idx → [tokens]}
|
|||
|
|
Phase6 merge uses this directly to inject SNAC tokens per 24-frame chunk,
|
|||
|
|
aligned with the cosmos/avclm/agent tokens that fire at the same chunk
|
|||
|
|
boundaries (n_chunks recomputed from CHUNK_SIZE, not read from a file).
|
|||
|
|
|
|||
|
|
Returns stats: {ok, skipped_vid, failed_audio, failed_snac, tokens}
|
|||
|
|
"""
|
|||
|
|
out_path = os.path.join(output_dir, f"{video_id}_snac.jsonl")
|
|||
|
|
if skip_existing and os.path.exists(out_path):
|
|||
|
|
return {"ok": 0, "skipped_vid": len(activities), "failed_audio": 0,
|
|||
|
|
"failed_snac": 0, "tokens": 0}
|
|||
|
|
|
|||
|
|
video_path = os.path.join(video_dir, f"{video_id}.mp4")
|
|||
|
|
if not os.path.exists(video_path):
|
|||
|
|
return {"ok": 0, "skipped_vid": 0, "failed_audio": len(activities),
|
|||
|
|
"failed_snac": 0, "tokens": 0}
|
|||
|
|
|
|||
|
|
# Extract full audio once
|
|||
|
|
full_audio = extract_full_audio(video_path)
|
|||
|
|
if full_audio is None:
|
|||
|
|
log.warning(f"No audio: {video_path}")
|
|||
|
|
return {"ok": 0, "skipped_vid": 0, "failed_audio": len(activities),
|
|||
|
|
"failed_snac": 0, "tokens": 0}
|
|||
|
|
|
|||
|
|
stats = {"ok": 0, "skipped_vid": 0, "failed_audio": 0, "failed_snac": 0, "tokens": 0}
|
|||
|
|
rows = []
|
|||
|
|
|
|||
|
|
for act in activities:
|
|||
|
|
segment = slice_audio(full_audio, act["start_sec"], act["end_sec"])
|
|||
|
|
if len(segment) < int(SAMPLE_RATE * 0.1): # skip segments < 100 ms
|
|||
|
|
stats["failed_audio"] += 1
|
|||
|
|
continue
|
|||
|
|
try:
|
|||
|
|
flat_tokens = encode_fn(segment, model, device)
|
|||
|
|
except Exception as e:
|
|||
|
|
log.warning(f"SNAC failed {video_id}/{act['activity_id']}: {e}")
|
|||
|
|
stats["failed_snac"] += 1
|
|||
|
|
continue
|
|||
|
|
if not flat_tokens:
|
|||
|
|
stats["failed_snac"] += 1
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# Split flat token list into per-chunk dicts, aligned to the same
|
|||
|
|
# 24-frame grid as cosmos/avclm/agent. n_chunks recomputed
|
|||
|
|
# independently from CHUNK_SIZE (not from a merged file's
|
|||
|
|
# chunk_timing -- see _scan_one_rank_file docstring).
|
|||
|
|
total_frames = max(1, round((act["end_sec"] - act["start_sec"]) * TARGET_FPS))
|
|||
|
|
n_chunks = math.ceil(total_frames / CHUNK_SIZE)
|
|||
|
|
by_chunk = split_snac_by_chunks(flat_tokens, n_chunks)
|
|||
|
|
snac_by_chunk = {str(k): v for k, v in by_chunk.items()}
|
|||
|
|
|
|||
|
|
rows.append({
|
|||
|
|
"video_id": video_id,
|
|||
|
|
"activity_id": act["activity_id"],
|
|||
|
|
"start_sec": round(act["start_sec"], 4),
|
|||
|
|
"end_sec": round(act["end_sec"], 4),
|
|||
|
|
"has_agent": act.get("has_agent", False),
|
|||
|
|
"snac_by_chunk": snac_by_chunk,
|
|||
|
|
})
|
|||
|
|
stats["ok"] += 1
|
|||
|
|
stats["tokens"] += len(flat_tokens)
|
|||
|
|
|
|||
|
|
if rows:
|
|||
|
|
with open(out_path, "w") as f:
|
|||
|
|
for row in rows:
|
|||
|
|
f.write(json.dumps(row) + "\n")
|
|||
|
|
|
|||
|
|
return stats
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|||
|
|
# Main
|
|||
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def parse_args():
|
|||
|
|
p = argparse.ArgumentParser(description="SNAC tokenization for FineVideo-VLA")
|
|||
|
|
p.add_argument("--build-tasks", action="store_true",
|
|||
|
|
help="Scan final_dataset_adaptive and write snac_task_list.json, then exit.")
|
|||
|
|
p.add_argument("--input-glob", default=INPUT_GLOB)
|
|||
|
|
p.add_argument("--output-dir", default=OUTPUT_DIR)
|
|||
|
|
p.add_argument("--video-dir", default=VIDEO_DIR)
|
|||
|
|
p.add_argument("--task-cache", default=TASK_CACHE)
|
|||
|
|
p.add_argument("--hf-cache", default=HF_CACHE)
|
|||
|
|
p.add_argument("--scan-workers", type=int, default=8,
|
|||
|
|
help="CPU workers for --build-tasks scan (default 8)")
|
|||
|
|
p.add_argument("--no-skip", action="store_true",
|
|||
|
|
help="Re-process videos even if output file exists")
|
|||
|
|
p.add_argument("--format", choices=["listen", "speak"], default="listen",
|
|||
|
|
help="listen = L0+L1 only (current production, 3 tok/base-frame); "
|
|||
|
|
"speak = full L0+L1+L2 (2026-07-22, 7 tok/base-frame, +133%% tokens). "
|
|||
|
|
"speak requires L2 tokens added to the tokenizer vocab first.")
|
|||
|
|
return p.parse_args()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
args = parse_args()
|
|||
|
|
skip_existing = not args.no_skip
|
|||
|
|
|
|||
|
|
# ── Set HF cache ─────────────────────────────────────────────────────────
|
|||
|
|
os.environ.setdefault("HF_HOME", args.hf_cache)
|
|||
|
|
os.makedirs(args.hf_cache, exist_ok=True)
|
|||
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
# ── Mode: build task list ─────────────────────────────────────────────────
|
|||
|
|
if args.build_tasks:
|
|||
|
|
build_task_list(args.input_glob, args.task_cache, workers=args.scan_workers)
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# ── Mode: tokenize ────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
# SLURM array vars
|
|||
|
|
task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", "0"))
|
|||
|
|
num_tasks = int(os.environ.get("SLURM_ARRAY_TASK_COUNT", "1"))
|
|||
|
|
|
|||
|
|
# Load task list (must exist — run --build-tasks first)
|
|||
|
|
if not os.path.exists(args.task_cache):
|
|||
|
|
log.error(
|
|||
|
|
f"Task list not found: {args.task_cache}\n"
|
|||
|
|
f"Run first: python pipeline_pose/snac_finevideo.py --build-tasks"
|
|||
|
|
)
|
|||
|
|
sys.exit(1)
|
|||
|
|
with open(args.task_cache) as f:
|
|||
|
|
all_tasks = json.load(f)
|
|||
|
|
|
|||
|
|
all_vids = sorted(all_tasks.keys())
|
|||
|
|
my_vids = all_vids[task_id::num_tasks]
|
|||
|
|
log.info(
|
|||
|
|
f"Task {task_id}/{num_tasks}: {len(my_vids)}/{len(all_vids)} videos "
|
|||
|
|
f"skip_existing={skip_existing}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# ── Load SNAC model ───────────────────────────────────────────────────────
|
|||
|
|
from snac import SNAC # imported here to avoid load cost during --build-tasks
|
|||
|
|
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
|||
|
|
log.info(f"Loading SNAC model ({SNAC_MODEL}) on {device}...")
|
|||
|
|
t_load = time.time()
|
|||
|
|
os.environ["HF_HOME"] = args.hf_cache # ensure hub/ subdir is found
|
|||
|
|
model = SNAC.from_pretrained(SNAC_MODEL,
|
|||
|
|
local_files_only=True).eval().to(device)
|
|||
|
|
log.info(f"SNAC loaded ({time.time()-t_load:.1f}s)")
|
|||
|
|
|
|||
|
|
# ── Process videos ────────────────────────────────────────────────────────
|
|||
|
|
cumul = {"ok": 0, "skipped_vid": 0, "failed_audio": 0, "failed_snac": 0, "tokens": 0}
|
|||
|
|
t_start = time.time()
|
|||
|
|
|
|||
|
|
encode_fn = encode_speak if args.format == "speak" else encode_listen
|
|||
|
|
|
|||
|
|
for idx, vid in enumerate(my_vids, 1):
|
|||
|
|
s = process_video(
|
|||
|
|
vid, all_tasks[vid], model, device,
|
|||
|
|
args.video_dir, args.output_dir, skip_existing,
|
|||
|
|
encode_fn=encode_fn,
|
|||
|
|
)
|
|||
|
|
for k in cumul:
|
|||
|
|
cumul[k] += s[k]
|
|||
|
|
|
|||
|
|
if idx % 100 == 0 or idx == len(my_vids):
|
|||
|
|
elapsed = time.time() - t_start
|
|||
|
|
rate = idx / elapsed
|
|||
|
|
eta = (len(my_vids) - idx) / rate if rate > 0 else 0
|
|||
|
|
log.info(
|
|||
|
|
f"[{idx:5d}/{len(my_vids)}] vid={vid} "
|
|||
|
|
f"ok={s['ok']} skip={s['skipped_vid']} "
|
|||
|
|
f"fail_audio={s['failed_audio']} fail_snac={s['failed_snac']} "
|
|||
|
|
f"rate={rate:.1f}vid/s ETA={eta/60:.0f}m "
|
|||
|
|
f"total_tokens={cumul['tokens']:,}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
elapsed = time.time() - t_start
|
|||
|
|
log.info(
|
|||
|
|
f"DONE task {task_id}: "
|
|||
|
|
f"ok={cumul['ok']} skipped={cumul['skipped_vid']} "
|
|||
|
|
f"fail_audio={cumul['failed_audio']} fail_snac={cumul['failed_snac']} "
|
|||
|
|
f"tokens={cumul['tokens']:,} wall={elapsed:.0f}s"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|