203 lines
9.1 KiB
Python
203 lines
9.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
General-purpose Cosmos detokenizer -- turns `<cosmos_N>` tokens (from any of
|
|
our flattened datasets: FineVideo-VLA, OmniVideo-100K, ...) back into an
|
|
actual video.
|
|
|
|
Verified 20/07/2026 (see PROGRESS_VI.md entry) against a real 200-token
|
|
chunk from omnivideo_100k_final: for the DV8x16x16 checkpoint, encode() of
|
|
an 8-frame/160x160 chunk produces indices of shape (1, 2, 10, 10) == 200
|
|
flat tokens -- this is CHUNK_TOKENS/CHUNK_GRID below, hardcoded for this
|
|
checkpoint (would need updating if a different Cosmos-Tokenizer-* variant
|
|
is ever used).
|
|
|
|
This is a lossy, generative-but-deterministic *reconstruction* (the neural
|
|
decoder's output, not the original pixels) -- expect visible blur/artifacts,
|
|
this is the accepted tradeoff for a much lower token count than storing raw
|
|
pixels. Contrast with decode_avclm.py in this same directory, which is
|
|
byte-exact (BPE over raw H.264 bytes, not a neural codec).
|
|
|
|
Token values in `<cosmos_N>` are the RAW encoder codebook indices with no
|
|
vocab offset added (verified: flatten scripts across the project just wrap
|
|
the same digit string in `<cosmos_{n}>`, never add an offset) -- so N can be
|
|
fed straight into decode() after reshaping.
|
|
|
|
Usage:
|
|
# From a raw list of ints:
|
|
python tools/decode/decode_cosmos.py --tokens 18697,55801,44451,... --output out.mp4
|
|
|
|
# From a flattened JSONL record (extracts the Nth 200-token chunk):
|
|
python tools/decode/decode_cosmos.py --input-jsonl path.jsonl --record-id VIDEO_ID \
|
|
--chunk-index 0 --output out.mp4
|
|
"""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
|
|
# 2026-07-23: PROTOTYPE_DIR only exists on the internal cluster (not part of
|
|
# the public github.com/TieuDaoChanNhan/finevideo-vla repo -- 0 files under
|
|
# prototype/ are git-tracked). External users have no way to reach a
|
|
# hardcoded local path, so the checkpoint is now fetched from NVIDIA's public
|
|
# HF repo on first use instead (cached under the standard HF_HOME, same as
|
|
# any other from_pretrained() call) -- falls back to the local cluster copy
|
|
# first if it happens to be present, to avoid a redundant download here.
|
|
PROTOTYPE_DIR = "/e/project1/reformo/nguyen38/prototype"
|
|
_LOCAL_CHECKPOINT_DEC = os.path.join(
|
|
PROTOTYPE_DIR, "pretrained_ckpts/Cosmos-Tokenizer-DV8x16x16/decoder.jit"
|
|
)
|
|
COSMOS_HF_REPO = "nvidia/Cosmos-Tokenizer-DV8x16x16"
|
|
CHUNK_GRID = (2, 10, 10) # (T', H', W') per 8-frame/160x160 input chunk, this checkpoint
|
|
CHUNK_TOKENS = CHUNK_GRID[0] * CHUNK_GRID[1] * CHUNK_GRID[2] # 200
|
|
|
|
|
|
def _resolve_checkpoint_dec() -> str:
|
|
if os.path.exists(_LOCAL_CHECKPOINT_DEC):
|
|
return _LOCAL_CHECKPOINT_DEC
|
|
from huggingface_hub import hf_hub_download
|
|
print(f"Local checkpoint not found -- downloading decoder.jit from {COSMOS_HF_REPO} "
|
|
f"(~350MB, cached for future runs)...")
|
|
return hf_hub_download(repo_id=COSMOS_HF_REPO, filename="decoder.jit")
|
|
|
|
_COSMOS_ATOMIC_RE = re.compile(r"<cosmos_(\d+)>")
|
|
_COSMOS_RAW_BLOCK_RE = re.compile(r"<cosmos>(.*?)</cosmos>", re.DOTALL)
|
|
|
|
|
|
def extract_chunk_tokens(text: str, chunk_index: int) -> list:
|
|
"""Pull out the Nth CHUNK_TOKENS-sized slice of cosmos ids from a record.
|
|
Supports both formats found in this project:
|
|
- flattened/atomic: `<cosmos_N> <cosmos_N> ...` (post-flatten output,
|
|
e.g. omnivideo_100k_final, FineVideo's megatron_dataset_*)
|
|
- raw pre-flatten block: `<cosmos>N N N...</cosmos>` (Step A's own
|
|
output before any flatten script runs, e.g. FineVideo's
|
|
training_ready_rank_*.jsonl activity.video_tokens, or OmniVideo's
|
|
omnivideo_100k_video_flat) -- chunk_index selects which <cosmos>
|
|
block (each block is already exactly one chunk, verified 200 tokens
|
|
both for FineVideo and OmniVideo real data).
|
|
|
|
In the flattened/atomic format, chunks with <50% cosmos dropout keep-rate
|
|
have gaps (some chunks entirely missing cosmos), so chunk_index there
|
|
means "Nth cosmos chunk present in the stream", not "Nth temporal chunk
|
|
of the video" -- the two only coincide if dropout happened to keep every
|
|
chunk up to that point.
|
|
"""
|
|
raw_blocks = _COSMOS_RAW_BLOCK_RE.findall(text)
|
|
if raw_blocks:
|
|
if chunk_index >= len(raw_blocks):
|
|
raise ValueError(f"Requested chunk {chunk_index} but only {len(raw_blocks)} <cosmos> blocks in this record")
|
|
chunk = [int(x) for x in raw_blocks[chunk_index].split() if x.isdigit()]
|
|
if len(chunk) != CHUNK_TOKENS:
|
|
raise ValueError(f"<cosmos> block {chunk_index} has {len(chunk)} tokens, expected {CHUNK_TOKENS}")
|
|
return chunk
|
|
|
|
all_ids = [int(x) for x in _COSMOS_ATOMIC_RE.findall(text)]
|
|
start = chunk_index * CHUNK_TOKENS
|
|
end = start + CHUNK_TOKENS
|
|
chunk = all_ids[start:end]
|
|
if len(chunk) != CHUNK_TOKENS:
|
|
raise ValueError(
|
|
f"Requested chunk {chunk_index} needs tokens [{start}:{end}) but only "
|
|
f"{len(all_ids)} cosmos tokens total in this record."
|
|
)
|
|
return chunk
|
|
|
|
|
|
def _record_text(rec: dict) -> str:
|
|
"""Flat records (OmniVideo-100K, Megatron-flattened FineVideo): {"text": ...}.
|
|
Raw FineVideo Step A records (training_ready_rank_*.jsonl) instead nest
|
|
per-activity `video_tokens` under scenes[].activities[] -- concatenate
|
|
all of them for the record so chunk_index can walk the whole video."""
|
|
if "text" in rec:
|
|
return rec["text"]
|
|
if "scenes" in rec:
|
|
return "".join(
|
|
act.get("video_tokens", "")
|
|
for scene in rec["scenes"]
|
|
for act in scene.get("activities", [])
|
|
)
|
|
raise KeyError("Record has neither 'text' nor 'scenes' -- unrecognized schema")
|
|
|
|
|
|
def load_tokens_from_jsonl(path: str, record_id: str, chunk_index: int) -> list:
|
|
with open(path, encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
rec = json.loads(line)
|
|
rid = rec.get("video_id", rec.get("id"))
|
|
if rid == record_id:
|
|
return extract_chunk_tokens(_record_text(rec), chunk_index)
|
|
raise KeyError(f"record_id={record_id!r} not found in {path}")
|
|
|
|
|
|
def decode_cosmos_chunk(token_ids: list, output_path: str, fps: int = 6) -> None:
|
|
"""token_ids: exactly CHUNK_TOKENS (200) raw cosmos codebook indices."""
|
|
if len(token_ids) != CHUNK_TOKENS:
|
|
raise ValueError(f"Expected exactly {CHUNK_TOKENS} tokens, got {len(token_ids)}")
|
|
output_path = os.path.abspath(output_path) # resolve before any cwd assumptions below
|
|
|
|
# 2026-07-23: cosmos_tokenizer itself is also vendored (tools/decode/vendor/),
|
|
# not pip-installable -- see vendor/cosmos_tokenizer/NOTICE.md. No os.chdir()
|
|
# needed anymore (that was only for PROTOTYPE_DIR's relative lookups).
|
|
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "vendor"))
|
|
import imageio_ffmpeg
|
|
import torch
|
|
import torchvision.transforms as T
|
|
from cosmos_tokenizer.video_lib import CausalVideoTokenizer
|
|
|
|
checkpoint_dec = _resolve_checkpoint_dec()
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
dec = CausalVideoTokenizer(checkpoint_dec=checkpoint_dec).to(device)
|
|
|
|
indices = torch.tensor(token_ids, dtype=torch.int64, device=device).view(1, *CHUNK_GRID)
|
|
with torch.no_grad():
|
|
out = dec.decode(indices) # (1, 3, T, H, W), range ~[-1, 1]
|
|
|
|
out = ((out.float() + 1.0) / 2.0).clamp(0, 1).squeeze(0) # (3, T, H, W)
|
|
n_frames = out.shape[1]
|
|
|
|
frame_dir = f"/tmp/cosmos_decode_{os.getpid()}"
|
|
os.makedirs(frame_dir, exist_ok=True)
|
|
to_pil = T.ToPILImage()
|
|
for i in range(n_frames):
|
|
to_pil(out[:, i, :, :].cpu()).save(f"{frame_dir}/frame_{i:02d}.png")
|
|
|
|
ffmpeg_bin = imageio_ffmpeg.get_ffmpeg_exe()
|
|
subprocess.run(
|
|
[ffmpeg_bin, "-y", "-framerate", str(fps), "-i", f"{frame_dir}/frame_%02d.png",
|
|
"-vf", "scale=320:320:flags=neighbor", "-pix_fmt", "yuv420p", output_path],
|
|
check=True, capture_output=True,
|
|
)
|
|
for i in range(n_frames):
|
|
os.remove(f"{frame_dir}/frame_{i:02d}.png")
|
|
os.rmdir(frame_dir)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
ap.add_argument("--tokens", help="Comma-separated list of exactly 200 raw cosmos ids")
|
|
ap.add_argument("--input-jsonl", help="Flattened JSONL file to pull tokens from")
|
|
ap.add_argument("--record-id", help="video_id/id field to select within --input-jsonl")
|
|
ap.add_argument("--chunk-index", type=int, default=0,
|
|
help="Which 200-token chunk (in order of appearance) to decode, 0-indexed")
|
|
ap.add_argument("--fps", type=int, default=6, help="Output mp4 framerate (decoded frames are few, slow fps for visibility)")
|
|
ap.add_argument("--output", required=True)
|
|
args = ap.parse_args()
|
|
|
|
if args.tokens:
|
|
token_ids = [int(x) for x in args.tokens.split(",")]
|
|
elif args.input_jsonl and args.record_id:
|
|
token_ids = load_tokens_from_jsonl(args.input_jsonl, args.record_id, args.chunk_index)
|
|
else:
|
|
ap.error("Provide either --tokens or (--input-jsonl and --record-id)")
|
|
|
|
decode_cosmos_chunk(token_ids, args.output, fps=args.fps)
|
|
print(f"Saved: {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|