From a224c577b08109083d68e21b2a1c56319ff24ab5 Mon Sep 17 00:00:00 2001
From: Aryan Kumar <aryan5v@users.noreply.github.com>
Date: Thu, 29 Oct 2026 22:13:55 -0700
Subject: [PATCH] [feat]: Pika H3 Flex Ref2VA coverage rows from the OmniRef v4
 sources

Re-encode references and Qwen3-VL conditioning with Flex's tier sizing (1024 short edge
below 768p) and add audio-reference, mixed image+video, 6-9 image and two-video rows,
reusing the v4 target latents.
---
 .../train/models/test_pika_h3_flex_qad.py     |  66 +++
 .../build_flex_ref2va_rows.py                 | 445 ++++++++++++++++++
 2 files changed, 511 insertions(+)
 create mode 100644 scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py

diff --git a/fastvideo/tests/train/models/test_pika_h3_flex_qad.py b/fastvideo/tests/train/models/test_pika_h3_flex_qad.py
index 1de5e0843..aa6bf4231 100644
--- a/fastvideo/tests/train/models/test_pika_h3_flex_qad.py
+++ b/fastvideo/tests/train/models/test_pika_h3_flex_qad.py
@@ -255,3 +255,69 @@ def test_calibration_choice_all_tokens_or_balanced():
     assert report["b0.ff.fc_in"]["best"] == "0.5x_max"
     table, _ = calibrate.choose_table(mse, candidates, "all")
     assert table == {"b0.ff.fc_in": 10.0}
+
+
+def _builder_module():
+    path = Path(__file__).resolve().parents[4] / "scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py"
+    spec = importlib.util.spec_from_file_location("build_flex_ref2va_rows", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+_PROMPT = "subject_definitions: <Picture 1> is the exact first frame of the target video.\n\nsummary: A dog runs."
+
+
+def _v4_entry(case: str, refs: list[dict], clip: str = "asset-" + "a" * 32 + "_f000001_f000124_x") -> dict:
+    return {
+        "case": case,
+        "id": f"{case}:{clip}" if case in ("first_frame", "first_last_frame", "storyboard") else clip,
+        "resolution": "480p",
+        "prompt": _PROMPT,
+        "references": refs,
+        "target": {"height": 480, "width": 832, "num_frames": 124},
+        "target_clip_id": clip,
+    }
+
+
+def test_builder_adds_subject_definitions_in_order():
+    builder = _builder_module()
+    after = builder._subject(_PROMPT, "<Audio 1> is the soundtrack of the target video.")
+    assert after.startswith("subject_definitions: <Picture 1> is the exact first frame of the target video. <Audio 1>")
+    assert after.endswith("\n\nsummary: A dog runs.")
+    before = builder._subject(_PROMPT, "<Picture 1> is the first frame.", before=True)
+    assert before.startswith("subject_definitions: <Picture 1> is the first frame. <Picture 1> is the exact")
+
+
+def test_builder_variants_carry_flex_reference_mixes(tmp_path):
+    import random as stdlib_random
+    builder = _builder_module()
+    image = [{"kind": "image", "path": "/x/keyframes/480p/c/first.png", "crop_left": 11, "crop_right": 11}]
+    entry = _v4_entry("first_frame", image)
+    audio = builder.variant("audio_reference", entry, tmp_path, stdlib_random.Random(0), [])
+    assert [r["kind"] for r in audio["references"]] == ["image", "audio"]
+    assert audio["id"].startswith("audio_reference:") and audio["case"] == "audio_reference"
+    assert audio["references"][1]["path"].endswith("target_clips/480p/" + entry["target_clip_id"] + ".mp4")
+    retier = builder.variant("retier", entry, tmp_path, stdlib_random.Random(0), [])
+    assert retier["id"] == entry["id"] and retier["case"] == "first_frame" and retier["prompt"] == _PROMPT
+    board = _v4_entry("storyboard", image * 3)
+    many = builder.variant("many_images", board, tmp_path, stdlib_random.Random(0), [])
+    assert 6 <= len(many["references"]) <= 9
+    assert many["prompt"].count("is a frame of the target video") == len(many["references"]) - 3
+
+
+def test_builder_continuation_variants(tmp_path):
+    import random as stdlib_random
+    builder = _builder_module()
+    clip = "asset-" + "b" * 32 + "_omniref_continue_1"
+    video = [{"kind": "video", "path": "/x/a/preceding.mp4", "use_embedded_audio": True}]
+    entry = _v4_entry("continue_scene", video, clip=clip)
+    entry["target_clip_id"] = clip + "__target"
+    mixed = builder.variant("mixed_image_video", entry, tmp_path, stdlib_random.Random(0), [])
+    assert [r["kind"] for r in mixed["references"]] == ["frame", "video"]
+    assert mixed["references"][0]["path"].endswith(f"continuation_clips/480p/asset-{'b' * 32}/{clip}/target.mp4")
+    sibling = _v4_entry("continue_shot", [{"kind": "video", "path": "/x/b/preceding.mp4"}], clip=clip + "2")
+    multi = builder.variant("multi_video", entry, tmp_path, stdlib_random.Random(0), [entry, sibling])
+    assert [r["path"] for r in multi["references"]] == ["/x/a/preceding.mp4", "/x/b/preceding.mp4"]
+    assert multi["references"][1]["audio"] is False
+    assert builder.variant("multi_video", entry, tmp_path, stdlib_random.Random(0), [entry]) is None
diff --git a/scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py b/scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py
new file mode 100644
index 000000000..0c0bb04d1
--- /dev/null
+++ b/scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py
@@ -0,0 +1,445 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pika H3 Flex Ref2VA rows from the OmniRef v4 sources: Flex reference sizing plus the inputs v4 lacks.
+
+The OmniRef v4 encoded rows (the FastH3 OmniRef QAD data) size every image reference at a 2048 short edge.
+Pika H3 Flex uses 1024 below its 768p tier and 2048 at 768p, and it accepts input mixes v4 never encoded. This
+builds packed Ref2VA rows (``pyarrow_schema_ref2va``) for the Flex QAD from the v4 raw manifests and media:
+
+``retier``             the v4 row's own references and prompt, re-encoded with Flex's tier sizing; same case name
+                       and id as the v4 row, so a manifest listed before v4 replaces it (``load_manifest_groups``
+                       keeps the first entry per id and resolution).
+``audio_reference``    first-frame rows plus the target's soundtrack as a standalone audio reference.
+``mixed_image_video``  continuation rows plus the target's first frame as an image ahead of the preceding video.
+``many_images``        storyboard rows extended with frames of the target to 6-9 image references.
+``multi_video``        continuation rows plus a second, silent preceding clip of the same asset.
+
+Only references and the Qwen3-VL conditioning are encoded (with the released Ref2VA encode path of
+``preprocess_minimax_h3_ref2va``; the VAE, audio VAE and text encoder stay resident). Target video and audio latents
+are copied from the v4 row of the same target: QAD regresses on the teacher, so they only carry the geometry.
+
+    python scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py plan --v4-root <v4 run dir> \\
+        --plan-config <QAD yaml>... --extra-per-group 480p=80 768p=20 --output-dir <out>
+    CUDA_VISIBLE_DEVICES=<g> python .../build_flex_ref2va_rows.py encode --model-path <merged Flex> \\
+        --output-dir <out> --shard <g> --num-shards <n>
+    python .../build_flex_ref2va_rows.py finalize --output-dir <out>
+"""
+from __future__ import annotations
+
+import argparse
+import functools
+import hashlib
+import json
+import random
+import re
+from collections import defaultdict
+from pathlib import Path
+from typing import Any
+
+V4_CASES = ("first_frame", "first_last_frame", "storyboard", "continue_scene", "continue_shot")
+IMAGE_CASES = ("first_frame", "first_last_frame", "storyboard")
+CONTINUE_CASES = ("continue_scene", "continue_shot")
+NEW_CASES = ("audio_reference", "mixed_image_video", "many_images", "multi_video")
+MAX_IMAGES = 9
+MAX_VIDEO_SECONDS = 15.0
+FPS = 24.0
+_ASSET = re.compile(r"(asset-[0-9a-f]{32})")
+
+
+# ---------------------------------------------------------------------- planning (CPU)
+def load_v4(v4_root: Path) -> tuple[dict[tuple[str, str], dict[str, Any]], dict[tuple[str, str], dict[str, Any]]]:
+    """``(raw, encoded)`` v4 entries keyed by (id, resolution)."""
+    raw: dict[tuple[str, str], dict[str, Any]] = {}
+    for path in sorted((v4_root / "manifests").glob("shard-*.jsonl")):
+        for line in path.read_text().splitlines():
+            entry = json.loads(line)
+            raw[(entry["id"], entry["resolution"])] = entry
+    encoded: dict[tuple[str, str], dict[str, Any]] = {}
+    for line in (v4_root / "encoded" / "full" / "MANIFEST.jsonl").read_text().splitlines():
+        entry = json.loads(line)
+        encoded[(entry["id"], entry["resolution"])] = entry
+    return raw, encoded
+
+
+def qad_plan_ids(config: str, overrides: list[str]) -> set[tuple[str, str]]:
+    """(id, resolution) of every training and held-out row the QAD config would read from the v4 manifest."""
+    from fastvideo.train.models.minimax_h3.omniref_data import (OMNIREF_CASES, limit_frames, load_eval_rows,
+                                                                load_manifest_groups, select_rows, top_up_eval_rows)
+    from fastvideo.train.utils.config import load_run_config
+
+    cfg = load_run_config(config, overrides=overrides).raw["models"]["student"]["data"]
+    cases = tuple(cfg.get("cases", OMNIREF_CASES))
+    resolutions = tuple(cfg.get("eval_resolutions", ("480p", )))
+    eval_per_group = cfg.get("eval_per_group", 2)
+    groups = load_manifest_groups(list(cfg["manifests"]), cases, int(cfg.get("max_frames", 243)))
+    eval_rows = load_eval_rows(cfg["eval_manifest"], cases, resolutions, eval_per_group,
+                               int(cfg.get("eval_max_frames", 0)))
+    if isinstance(eval_per_group, dict):
+        eval_groups = limit_frames({k: v for k, v in groups.items() if k[1] in resolutions},
+                                   int(cfg.get("eval_max_frames", 0)))
+        eval_rows = top_up_eval_rows(eval_rows, eval_groups, int(cfg.get("eval_seed", 20261008)), dict(eval_per_group))
+    plan = select_rows(groups, int(cfg.get("seed", 20261007)), dict(cfg.get("per_group", {"480p": 400})),
+                       frozenset(row.source for row in eval_rows))
+    ids = set()
+    for row in [*plan, *eval_rows]:
+        entry_id = f"{row.case}:{row.source}" if row.case in IMAGE_CASES else row.source
+        ids.add((entry_id, row.resolution))
+    return ids
+
+
+def _keyframe_source(path: str) -> str:
+    """The 768p copy of a keyframe (same frame, more pixels) when it exists; Flex upsamples 480p sources."""
+    hi = path.replace("/keyframes/480p/", "/keyframes/768p/")
+    return hi if hi != path and Path(hi).is_file() else path
+
+
+def target_path(entry: dict[str, Any], v4_assets: Path) -> str:
+    clip = entry["target_clip_id"]
+    if clip.endswith("__target"):
+        clip = clip[:-len("__target")]
+        asset = _ASSET.search(clip).group(1)  # type: ignore[union-attr]
+        return str(v4_assets / "continuation_clips" / entry["resolution"] / asset / clip / "target.mp4")
+    return str(v4_assets / "target_clips" / entry["resolution"] / f"{clip}.mp4")
+
+
+def _subject(prompt: str, added: str, *, before: bool = False) -> str:
+    """Add one definition to the prompt's ``subject_definitions:`` paragraph."""
+    head, sep, rest = prompt.partition("subject_definitions:")
+    if not sep:
+        raise ValueError("prompt has no subject_definitions section")
+    body, gap, tail = rest.partition("\n\n")
+    body = f" {added}{body}" if before else f"{body.rstrip()} {added}"
+    return f"{head}{sep}{body}{gap}{tail}"
+
+
+def _ref(kind: str, **fields: Any) -> dict[str, Any]:
+    return {"kind": kind, **fields}
+
+
+def v4_references(entry: dict[str, Any], v4_assets: Path) -> list[dict[str, Any]]:
+    """The v4 entry's references as builder specs (image crops and video-frame indices kept)."""
+    refs = []
+    for ref in entry["references"]:
+        if ref["kind"] == "image":
+            refs.append(_ref("image", path=_keyframe_source(ref["path"]), crop=[ref.get("crop_left", 0),
+                                                                              ref.get("crop_right", 0)]))
+        elif ref["kind"] == "video_frame":
+            refs.append(_ref("frame", path=ref["path"], frame=int(ref["frame_index"]),
+                             crop=[ref.get("crop_left", 0), ref.get("crop_right", 0)]))
+        elif ref["kind"] == "video":
+            refs.append(_ref("video", path=ref["path"], audio=bool(ref.get("use_embedded_audio", False))))
+        else:
+            raise ValueError(f"{entry['id']}: unsupported v4 reference kind {ref['kind']!r}")
+    return refs
+
+
+def variant(case: str, entry: dict[str, Any], v4_assets: Path, rng: random.Random,
+            siblings: list[dict[str, Any]]) -> dict[str, Any] | None:
+    """One output row spec for ``case`` built from v4 ``entry`` (None when the entry cannot carry it)."""
+    refs = v4_references(entry, v4_assets)
+    prompt = entry["prompt"]
+    target = target_path(entry, v4_assets)
+    frames = int(entry["target"]["num_frames"])
+    if case == "retier":
+        pass
+    elif case == "audio_reference":
+        refs = [*refs, _ref("audio", path=target)]
+        prompt = _subject(prompt, "<Audio 1> is the soundtrack of the target video.")
+    elif case == "mixed_image_video":
+        refs = [_ref("frame", path=target, frame=0, crop=[0, 0]), *refs]
+        prompt = _subject(prompt, "<Picture 1> is the exact first frame of the target video.", before=True)
+    elif case == "many_images":
+        have = len(refs)
+        want = rng.randint(max(6, have + 1), MAX_IMAGES)
+        picks = sorted(rng.sample(range(1, frames - 1), want - have))
+        refs = [*refs, *(_ref("frame", path=target, frame=f, crop=[0, 0]) for f in picks)]
+        for index, frame in enumerate(picks, start=have + 1):
+            prompt = _subject(prompt, f"<Picture {index}> is a frame of the target video at {frame / FPS:.2f} s.")
+    elif case == "multi_video":
+        own = refs[0]["path"]
+        others = [s for s in siblings if s["references"][0]["path"] != own]
+        if not others:
+            return None
+        other = rng.choice(others)
+        refs = [*refs, _ref("video", path=other["references"][0]["path"], audio=False)]
+        prompt = _subject(prompt, "<Video 2> is an earlier shot from the same scene, without sound.")
+    else:
+        raise ValueError(f"unknown case {case!r}")
+    source_id = entry["id"].split(":", 1)[-1]
+    out_case = entry["case"] if case == "retier" else case
+    return {
+        "id": entry["id"] if case == "retier" else f"{case}:{source_id}",
+        "case": out_case,
+        "source_case": entry["case"],
+        "v4_id": entry["id"],
+        "resolution": entry["resolution"],
+        "prompt": prompt,
+        "references": refs,
+        "target": entry["target"],
+    }
+
+
+def plan_rows(v4_root: Path, v4_assets: Path, retier_ids: set[tuple[str, str]], extra: dict[str, int],
+              seed: int, exclude_sources: set[str]) -> list[dict[str, Any]]:
+    raw, encoded = load_v4(v4_root)
+    specs = [variant("retier", raw[key], v4_assets, random.Random(0), []) for key in sorted(retier_ids)
+             if key in raw and key[1] == "480p" and raw[key]["case"] in IMAGE_CASES]
+    by_asset: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
+    for (entry_id, res), entry in raw.items():
+        if entry["case"] in CONTINUE_CASES:
+            by_asset[(_ASSET.search(entry_id).group(1), res)].append(entry)  # type: ignore[union-attr]
+    sources = {"audio_reference": ("first_frame", ), "mixed_image_video": CONTINUE_CASES,
+               "many_images": ("storyboard", ), "multi_video": CONTINUE_CASES}
+    for case in NEW_CASES:
+        for res, count in sorted(extra.items()):
+            pool = sorted((e for (i, r), e in raw.items()
+                           if r == res and e["case"] in sources[case] and (i, r) in encoded
+                           and i.split(":", 1)[-1] not in exclude_sources), key=lambda e: e["id"])
+            rng = random.Random(f"{seed}:{case}:{res}")
+            rng.shuffle(pool)
+            made = 0
+            for entry in pool:
+                if made >= count:
+                    break
+                asset = _ASSET.search(entry["id"]).group(1)  # type: ignore[union-attr]
+                spec = variant(case, entry, v4_assets, rng, by_asset.get((asset, res), []))
+                if spec is not None:
+                    specs.append(spec)
+                    made += 1
+    for spec in specs:
+        spec["v4_parquet"] = encoded[(spec["v4_id"], spec["resolution"])]["parquet"]
+    return specs
+
+
+# ---------------------------------------------------------------------- encoding (GPU)
+def _crop(image: Any, crop: list[int]) -> Any:
+    left, right = int(crop[0]), int(crop[1])
+    return image if not (left or right) else image.crop((left, 0, image.width - right, image.height))
+
+
+@functools.lru_cache(maxsize=1)
+def _decoded(path: str) -> Any:
+    from fastvideo.pipelines.basic.minimax_h3.reference import decode_reference_video
+    return decode_reference_video(path)
+
+
+def _frame(path: str, index: int) -> Any:
+    from PIL import Image
+
+    frames = _decoded(path)[0]
+    return Image.fromarray(frames[min(index, len(frames) - 1)])
+
+
+def build_references(spec: dict[str, Any]) -> list[Any]:
+    """``MiniMaxH3Reference`` list in declared order; a silent video drops its container soundtrack."""
+    from PIL import Image
+
+    from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3Reference
+
+    refs = []
+    for ref in spec["references"]:
+        if ref["kind"] == "image":
+            refs.append(MiniMaxH3Reference(source=_crop(Image.open(ref["path"]).convert("RGB"), ref["crop"])))
+        elif ref["kind"] == "frame":
+            refs.append(MiniMaxH3Reference(source=_crop(_frame(ref["path"], ref["frame"]), ref["crop"])))
+        elif ref["kind"] == "video" and ref["audio"]:
+            refs.append(MiniMaxH3Reference(source=ref["path"], media_type="video"))
+        elif ref["kind"] == "video":
+            frames, fps, _ = _decoded(ref["path"])
+            refs.append(MiniMaxH3Reference(source=frames, media_type="video", fps=fps))
+        elif ref["kind"] == "audio":
+            _, _, soundtrack = _decoded(ref["path"])
+            if soundtrack is None:
+                raise ValueError(f"{ref['path']} has no soundtrack")
+            refs.append(MiniMaxH3Reference(source=soundtrack[0], media_type="audio", sample_rate=soundtrack[1]))
+        else:
+            raise ValueError(f"unknown reference kind {ref['kind']!r}")
+    return refs
+
+
+def check_limits(references: list[Any]) -> None:
+    from fastvideo.pipelines.basic.minimax_h3.reference import (MINIMAX_H3_MAX_REFERENCE_AUDIOS,
+                                                                MINIMAX_H3_MAX_REFERENCE_IMAGES,
+                                                                MINIMAX_H3_MAX_REFERENCE_VIDEOS,
+                                                                MINIMAX_H3_MAX_REFERENCES)
+    counts = defaultdict(int)
+    for ref in references:
+        counts[ref.media_type] += 1
+    if (counts["image"] > MINIMAX_H3_MAX_REFERENCE_IMAGES or counts["video"] > MINIMAX_H3_MAX_REFERENCE_VIDEOS
+            or counts["audio"] > MINIMAX_H3_MAX_REFERENCE_AUDIOS or len(references) > MINIMAX_H3_MAX_REFERENCES):
+        raise ValueError(f"reference counts {dict(counts)} exceed Flex's limits")
+    video_seconds = sum(ref.frames.shape[0] / FPS for ref in references if ref.media_type == "video")
+    if video_seconds > MAX_VIDEO_SECONDS + 1e-6:
+        raise ValueError(f"combined reference video is {video_seconds:.2f} s (> {MAX_VIDEO_SECONDS} s)")
+
+
+_COMPONENTS: dict[str, Any] = {}
+
+
+@functools.cache
+def _encoders(model_path: str) -> tuple[Any, Any, Any]:
+    """``(encoding module, model_index, fastvideo_args)``; every component loads once and stays resident.
+
+    The released encode functions load their component per call; routing ``_load_component`` through a cache keeps
+    the VAE, audio VAE and Qwen3-VL resident across rows without changing their numerics.
+    """
+    from fastvideo.pipelines.preprocess import preprocess_minimax_h3_ref2va as encoding
+    from fastvideo.utils import verify_model_config_and_directory
+
+    encoding._init_single_process_distributed()
+    fastvideo_args, _ = encoding._build_fastvideo_args(Path(model_path))
+    loader = encoding._load_component
+
+    def cached(name: str, path: Path, index: dict[str, Any], args: Any) -> Any:
+        if name not in _COMPONENTS:
+            _COMPONENTS[name] = loader(name, path, index, args)
+        return _COMPONENTS[name]
+
+    encoding._load_component = cached
+    return encoding, verify_model_config_and_directory(model_path), fastvideo_args
+
+
+def encode_spec(spec: dict[str, Any], model_path: str) -> dict[str, Any]:
+    """The v4 row of the same target with this spec's references, prompt and Flex-sized conditioning."""
+    import pyarrow.parquet as pq
+
+    from fastvideo.pipelines.basic.minimax_h3.reference import prepare_reference
+    from fastvideo.pipelines.basic.pika_h3_flex.buckets import PIKA_TIERS, tier_for_canvas
+
+    encoding, model_index, fastvideo_args = _encoders(model_path)
+    target = spec["target"]
+    tier = PIKA_TIERS[tier_for_canvas(int(target["height"]), int(target["width"]))]
+    prepared = [
+        prepare_reference(ref,
+                          int(target["num_frames"]),
+                          encoding.AUDIO_SAMPLE_RATE,
+                          reference_short_edge=tier.reference_short_edge,
+                          reference_canvas_short_edge=tier.short_edge,
+                          reference_canvas_max_pixels=tier.max_pixels) for ref in build_references(spec)
+    ]
+    check_limits(prepared)
+    path = Path(model_path)
+    visual = encoding.encode_ref_visual_anchor(prepared, path, model_index, fastvideo_args, (1, 2, 2), augment=False)
+    audio = encoding.encode_ref_audio_anchor(prepared, path, model_index, fastvideo_args)
+    text, tags = encoding.encode_ref2va_conditioning(spec["prompt"], prepared, path, model_index, fastvideo_args)
+    metadata = encoding._canonical_prepared_references(prepared)
+    record = pq.ParquetFile(spec["v4_parquet"]).read().to_pylist()[0]
+    for name, tensor in (("text_embedding", text), ("text_token_tags", tags), ("reference_video_latent_rows", visual),
+                         ("reference_audio_latent_rows", audio)):
+        # Every packed tensor is stored as float32 (the readers decode tags as float32 and round).
+        array = tensor.detach().float().cpu().contiguous()
+        record[f"{name}_bytes"] = array.numpy().tobytes()
+        record[f"{name}_shape"] = list(array.shape)
+        record[f"{name}_dtype"] = "float32"
+    for field, key in (("reference_media_types", "media_type"), ("reference_has_audio", "has_audio"),
+                       ("reference_num_latent_frames", "num_latent_frames"),
+                       ("reference_latent_heights", "latent_height"), ("reference_latent_widths", "latent_width"),
+                       ("reference_num_audio_latents", "num_audio_latents")):
+        record[field] = [ref[key] for ref in metadata]
+    record["id"] = spec["id"]
+    record["caption"] = spec["prompt"]
+    return record
+
+
+def run_encode(args: argparse.Namespace) -> None:
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+
+    from fastvideo.dataset.dataloader.schema import pyarrow_schema_ref2va
+
+    out = Path(args.output_dir)
+    specs = json.loads((out / "plan.json").read_text())[args.shard::args.num_shards]
+    manifest = out / f"MANIFEST-shard{args.shard:02d}.jsonl"
+    done = {json.loads(line)["id"] + "|" + json.loads(line)["resolution"]
+            for line in manifest.read_text().splitlines()} if manifest.exists() else set()
+    for index, spec in enumerate(specs):
+        if f"{spec['id']}|{spec['resolution']}" in done:
+            continue
+        try:
+            record = encode_spec(spec, args.model_path)
+        except Exception as error:  # a bad source clip must not stop the shard; it is reported and skipped
+            print(f"SKIP {spec['id']} {spec['resolution']}: {type(error).__name__}: {error}", flush=True)
+            with (out / f"skipped-shard{args.shard:02d}.jsonl").open("a") as handle:
+                handle.write(json.dumps({"id": spec["id"], "resolution": spec["resolution"], "error": str(error)}) +
+                             "\n")
+            continue
+        bucket = re.search(r"bucket=[^/]+", spec["v4_parquet"]).group(0)  # type: ignore[union-attr]
+        digest = hashlib.sha256(f"{spec['id']}|{spec['resolution']}".encode()).hexdigest()[:24]
+        path = out / "parquet" / spec["case"] / bucket / f"{digest}.parquet"
+        path.parent.mkdir(parents=True, exist_ok=True)
+        pq.write_table(pa.Table.from_pylist([record], schema=pyarrow_schema_ref2va), path)
+        pq.ParquetFile(path).read()  # round trip
+        entry = {"case": spec["case"], "id": spec["id"], "parquet": str(path), "resolution": spec["resolution"],
+                 "source_case": spec["source_case"]}
+        with manifest.open("a") as handle:
+            handle.write(json.dumps(entry) + "\n")
+        print(f"encoded {index + 1}/{len(specs)} {spec['case']} {spec['resolution']} {spec['id']}", flush=True)
+
+
+def run_plan(args: argparse.Namespace, overrides: list[str]) -> None:
+    from fastvideo.train.models.minimax_h3.omniref_data import load_eval_rows
+
+    retier: set[tuple[str, str]] = set()
+    for config in args.plan_config:
+        retier |= qad_plan_ids(config, overrides)
+        for frames in args.plan_max_frames:
+            retier |= qad_plan_ids(config, [*overrides, "--models.student.data.max_frames", str(frames)])
+    extra = {item.split("=")[0]: int(item.split("=")[1]) for item in args.extra_per_group}
+    exclude = {row.source for row in load_eval_rows(args.eval_manifest, V4_CASES, ("480p", "768p"), 99, 0)}
+    specs = plan_rows(Path(args.v4_root), Path(args.v4_assets), retier, extra, args.seed, exclude)
+    out = Path(args.output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    (out / "plan.json").write_text(json.dumps(specs))
+    counts: dict[str, int] = defaultdict(int)
+    for spec in specs:
+        counts[f"{spec['case'] if spec['case'] in NEW_CASES else 'retier/' + spec['case']}/{spec['resolution']}"] += 1
+    print(f"{len(specs)} rows: {dict(sorted(counts.items()))}")
+
+
+def run_finalize(args: argparse.Namespace) -> None:
+    """Merge the shard manifests; with ``--eval-manifest``, also write the held-out manifest with retiered rows."""
+    out = Path(args.output_dir)
+    lines = []
+    for path in sorted(out.glob("MANIFEST-shard*.jsonl")):
+        lines += [line for line in path.read_text().splitlines() if line.strip()]
+    (out / "MANIFEST.jsonl").write_text("\n".join(lines) + "\n")
+    skipped = sum(len(p.read_text().splitlines()) for p in out.glob("skipped-shard*.jsonl"))
+    print(f"MANIFEST.jsonl: {len(lines)} rows ({skipped} skipped)")
+    if args.eval_manifest:
+        encoded = {(e["id"], e["resolution"]): e["parquet"] for e in map(json.loads, lines)}
+        items = json.loads(Path(args.eval_manifest).read_text())
+        replaced = 0
+        for item in items:
+            key = (f"{item['case']}:{item['source']}" if item["case"] in IMAGE_CASES else item["source"],
+                   item["resolution"])
+            if key in encoded:
+                item["parquet"] = encoded[key]
+                replaced += 1
+        (out / "eval-manifest.json").write_text(json.dumps(items, indent=1))
+        print(f"eval-manifest.json: {replaced} of {len(items)} held-out rows retiered")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("stage", choices=("plan", "encode", "finalize"))
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--v4-root", help="OmniRef v4 run dir (manifests/, encoded/full/MANIFEST.jsonl)")
+    parser.add_argument("--v4-assets", help="OmniRef v4 media root (target_clips/, continuation_clips/, keyframes/)")
+    parser.add_argument("--plan-config", nargs="*", default=[], help="QAD configs whose 480p image rows to retier")
+    parser.add_argument("--plan-max-frames", nargs="*", type=int, default=[], help="extra max_frames plans to cover")
+    parser.add_argument("--eval-manifest", help="held-out manifest; its sources are excluded from new cases")
+    parser.add_argument("--extra-per-group", nargs="*", default=["480p=80", "768p=20"])
+    parser.add_argument("--seed", type=int, default=20261030)
+    parser.add_argument("--model-path", help="merged Flex Ref2VA model (encode)")
+    parser.add_argument("--shard", type=int, default=0)
+    parser.add_argument("--num-shards", type=int, default=1)
+    args, overrides = parser.parse_known_args()
+    if args.stage == "plan":
+        run_plan(args, overrides)
+    elif args.stage == "encode":
+        run_encode(args)
+    else:
+        run_finalize(args)
+
+
+if __name__ == "__main__":
+    main()