icassp2027-experiment/cache_omni_av_pieces.py

#!/usr/bin/env python3
"""Cache exact Qwen2.5-Omni AV pieces for Thinker LoRA training.

For every source-aligned 2-second block this script caches ten causal phases:
0.2, 0.4, ..., 2.0 seconds.  The 2.0-second phase is reused as completed
history; phases 0.2--1.8 are used only for the current partial block.  No
future samples enter a phase.

The optional verification compares a full-sequence reconstruction from the
cache with the existing open-stream KV replay.  Run this 24-second pilot before
caching a full EasyCom session.  ``--av-batch-size`` batches equal-duration
phases from different source blocks, increasing frozen AV-tower occupancy while
preserving the exact per-block cache schema and causal timestamps.

The audio optimization flags keep the official Qwen/Whisper feature extractor
but may run its STFT on CUDA and bound the otherwise model-maximum zero padding.
Verification deliberately recomputes its oracle with the original processor
defaults, so an optimization cannot silently redefine the extracted feature.
"""

from __future__ import annotations

import argparse
import math
import platform
import time
from pathlib import Path
from typing import Any, Sequence

import numpy as np
import torch
import torch.nn.functional as F
from tqdm.auto import tqdm

import streaming_turn_model as core
from easycom_head import NUM_HORIZONS
from omni_cached_av import (
    SCHEMA_VERSION,
    CachedBlockStore,
    anchor_intervals,
    cache_block_filename,
    cache_size_bytes,
    forward_last_state,
    make_anchor,
    pieces_to_cache,
    reconstruct_open_stream_sequence,
    save_json,
)
from streaming_turn_model_directly import CausalSourceAV, _easycom_directory_specs


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-id", required=True)
    parser.add_argument("--easycom-session-dir", type=Path, required=True)
    parser.add_argument("--target", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument(
        "--end",
        type=float,
        default=24.0,
        help="Dense session endpoint in seconds; use 24 for the first pilot",
    )
    parser.add_argument("--step-seconds", type=float, default=0.2)
    parser.add_argument("--native-block-seconds", type=float, default=2.0)
    parser.add_argument("--memory-seconds", type=float, default=8.0)
    parser.add_argument("--video-fps", type=float, default=2.0)
    parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28)
    parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28)
    parser.add_argument("--frame-cache-size", type=int, default=64)
    parser.add_argument(
        "--av-batch-size",
        type=int,
        default=1,
        help=(
            "Microbatch source blocks that have the same 0.2s phase through "
            "the processor and frozen AV towers; 1 preserves serial behavior"
        ),
    )
    parser.add_argument(
        "--audio-feature-device",
        default=None,
        help=(
            "Device used by the native Whisper log-Mel extractor, for example "
            "model, cuda:0, or cpu; omitted preserves processor defaults"
        ),
    )
    parser.add_argument(
        "--audio-right-padding-samples",
        type=int,
        default=None,
        help=(
            "Replace model-maximum audio padding with actual samples plus this "
            "many zeros, rounded up to the Whisper hop; omitted preserves the "
            "original maximum-length padding"
        ),
    )
    parser.add_argument(
        "--decoder",
        choices=("decord", "torchcodec", "pynvcodec"),
        default="decord",
    )
    parser.add_argument("--decoder-gpu-id", type=int, default=0)
    parser.add_argument("--decoder-threads", type=int, default=1)
    parser.add_argument("--max-open-readers", type=int, default=1)
    parser.add_argument(
        "--dtype",
        choices=("auto", "bfloat16", "float16", "float32"),
        default="bfloat16",
    )
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--flash-attention-2", action="store_true")
    parser.add_argument(
        "--verify-every",
        type=int,
        default=10,
        help="Compare every Nth 5-Hz state with exact KV replay; 0 disables",
    )
    parser.add_argument(
        "--minimum-cosine",
        type=float,
        default=0.999,
        help="Fail if cached full-forward/reference KV cosine is below this",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite this cache's known block/static/manifest files",
    )
    return parser.parse_args(argv)


def _integer_ratio(numerator: float, denominator: float, name: str) -> int:
    ratio = numerator / denominator
    rounded = round(ratio)
    if not math.isclose(ratio, rounded, rel_tol=0.0, abs_tol=1e-6):
        raise ValueError(f"{name} must be integral, got {ratio}")
    return int(rounded)


def _load_targets(path: Path) -> np.ndarray:
    targets = np.load(path, mmap_mode="r")
    if targets.ndim != 2 or targets.shape[1] != 3:
        raise ValueError(f"Expected target [T,3], got {targets.shape}: {path}")
    return targets


def _safe_output_path(path: Path, overwrite: bool) -> None:
    if path.exists() and not overwrite:
        raise FileExistsError(f"Refusing to overwrite {path}; pass --overwrite")


def _model_device(model: torch.nn.Module) -> torch.device:
    try:
        return next(model.parameters()).device
    except StopIteration as exc:
        raise RuntimeError("Model has no parameters") from exc


def _one_video_item(value: Any) -> Any:
    if isinstance(value, (list, tuple)):
        if len(value) != 1:
            raise ValueError(f"Expected one prepared video, got {len(value)}")
        return value[0]
    return value


@torch.inference_mode()
def _prepare_pieces_batch(
    *,
    model: torch.nn.Module,
    processor: Any,
    media: Any,
    config: core.StreamConfig,
    intervals: list[tuple[float, float]],
) -> list[core.OpenMediaPieces]:
    """Encode equal-duration AV intervals as one independent-sample batch."""

    if not intervals:
        raise ValueError("Cannot prepare an empty AV batch")
    durations = [end - start for start, end in intervals]
    if any(
        not math.isclose(value, durations[0], rel_tol=0.0, abs_tol=1e-6)
        for value in durations[1:]
    ):
        raise ValueError(f"AV microbatch intervals have mixed durations: {durations}")

    texts: list[str] = []
    audios: list[Any] = []
    videos: list[Any] = []
    for interval in intervals:
        conversation, manual_audios = core.build_chunk_conversation(
            video_path=None,
            audio_path=None,
            intervals=[interval],
            use_audio_in_video=True,
            config=config,
            fixed_grid_av=media,
        )
        if manual_audios is None or len(manual_audios) != 1:
            raise RuntimeError("Direct-media interval did not provide one audio slice")
        text = core._single_template_text(
            processor.apply_chat_template(
                conversation,
                add_generation_prompt=False,
                tokenize=False,
            )
        )
        _, images, prepared_videos = core.process_mm_info(
            conversation,
            use_audio_in_video=False,
        )
        if images is not None and len(images):
            raise RuntimeError("AV cache batch unexpectedly produced image inputs")
        if prepared_videos is None:
            raise RuntimeError("AV cache batch did not produce video inputs")
        texts.append(text)
        audios.append(manual_audios[0])
        videos.append(_one_video_item(prepared_videos))

    batch = processor(
        text=texts,
        audio=audios,
        videos=videos,
        padding=True,
        use_audio_in_video=True,
        fps=config.video_fps,
    )
    batch = core._processor_batch_convert_to_tensors(batch)
    batch = core._processor_batch_to_device(batch, model.device)
    batch = core._processor_batch_to_dtype(batch, model.dtype)
    inputs = dict(batch)
    input_ids = inputs["input_ids"]
    if input_ids.shape[0] != len(intervals):
        raise RuntimeError(
            f"Processor returned batch {input_ids.shape[0]} for {len(intervals)} intervals"
        )

    embeddings = core._multimodal_embeddings(model, inputs)
    positions = core._local_position_ids(
        model,
        inputs,
        use_audio_in_video=True,
    )
    attention_mask = inputs.get("attention_mask")
    audio_id = core._config_token_id(model.config, "audio")
    video_id = core._config_token_id(model.config, "video")
    positions_per_chunk = core._tmrope_positions_per_chunk(model)
    pieces_batch: list[core.OpenMediaPieces] = []
    for row in range(len(intervals)):
        payload_mask = (input_ids[row] == audio_id) | (input_ids[row] == video_id)
        indices = payload_mask.nonzero(as_tuple=False).flatten()
        if indices.numel() == 0:
            raise ValueError(f"Prepared AV batch row {row} has no payload tokens")
        first = int(indices[0].item())
        last = int(indices[-1].item())
        if not bool(payload_mask[first : last + 1].all().item()):
            raise RuntimeError(f"AV payload is not contiguous in batch row {row}")
        if attention_mask is None:
            valid_end = input_ids.shape[1]
        else:
            valid = attention_mask[row].nonzero(as_tuple=False).flatten()
            if valid.numel() == 0:
                raise RuntimeError(f"Prepared AV batch row {row} is fully padded")
            valid_end = int(valid[-1].item()) + 1
        if valid_end <= last + 1:
            raise RuntimeError(f"Prepared AV batch row {row} has no closure tokens")
        pieces = core.OpenMediaPieces(
            prefix_embeddings=embeddings[row : row + 1, :first, :],
            prefix_position_ids=positions[:, row : row + 1, :first],
            payload_embeddings=embeddings[row : row + 1, first : last + 1, :],
            payload_position_ids=positions[:, row : row + 1, first : last + 1],
            closure_ids=input_ids[row : row + 1, last + 1 : valid_end],
            closure_position_ids=positions[
                :, row : row + 1, last + 1 : valid_end
            ],
        )
        temporal = pieces.payload_position_ids[0]
        temporal_span = int(temporal.max().item() - temporal.min().item())
        if temporal_span >= positions_per_chunk:
            raise RuntimeError(
                "One batched AV phase spans multiple native TMRoPE chunks: "
                f"interval={intervals[row]}, span={temporal_span}, "
                f"chunk={positions_per_chunk}"
            )
        pieces_batch.append(pieces)
    return pieces_batch


@torch.inference_mode()
def _cache_blocks(
    *,
    extractor: core.OpenStreamReplayExtractor,
    model: torch.nn.Module,
    processor: Any,
    media: Any,
    config: core.StreamConfig,
    output_dir: Path,
    num_steps: int,
    steps_per_block: int,
    step_seconds: float,
    block_seconds: float,
    av_batch_size: int,
    audio_feature_device: str | torch.device | None,
    audio_right_padding_samples: int | None,
    overwrite: bool,
) -> list[dict[str, Any]]:
    num_blocks = math.ceil(num_steps / steps_per_block)
    records: list[dict[str, Any]] = []
    input_embeddings = model.get_input_embeddings()
    progress = tqdm(
        total=num_blocks,
        desc=f"cache AV blocks (batch={av_batch_size})",
        unit="block",
    )
    for tile_start in range(0, num_blocks, av_batch_size):
        block_indices = list(
            range(tile_start, min(tile_start + av_batch_size, num_blocks))
        )
        variants_by_block: dict[int, dict[str, dict[str, torch.Tensor]]] = {
            index: {} for index in block_indices
        }
        lengths_by_block: dict[int, dict[str, int]] = {
            index: {} for index in block_indices
        }
        phase_counts = {
            index: min(
                steps_per_block,
                num_steps - index * steps_per_block,
            )
            for index in block_indices
        }
        for block_index in block_indices:
            _safe_output_path(
                output_dir / cache_block_filename(block_index), overwrite
            )

        for tail_steps in range(1, steps_per_block + 1):
            active_indices = [
                index
                for index in block_indices
                if tail_steps <= phase_counts[index]
            ]
            if not active_indices:
                continue
            intervals = [
                (
                    index * block_seconds,
                    index * block_seconds + tail_steps * step_seconds,
                )
                for index in active_indices
            ]
            if len(active_indices) == 1 or av_batch_size == 1:
                pieces_batch = [
                    extractor._prepare_pieces(
                        video_path=None,
                        audio_path=None,
                        interval=intervals[0],
                        use_audio_in_video=True,
                        audio_feature_device=audio_feature_device,
                        audio_right_padding_samples=audio_right_padding_samples,
                    )
                ]
            else:
                pieces_batch = _prepare_pieces_batch(
                    model=model,
                    processor=processor,
                    media=media,
                    config=config,
                    intervals=intervals,
                )
            if len(pieces_batch) != len(active_indices):
                raise RuntimeError("Batched AV piece count changed unexpectedly")
            for block_index, pieces in zip(active_indices, pieces_batch):
                closure_embeddings = input_embeddings(pieces.closure_ids)
                cached = pieces_to_cache(
                    pieces,
                    closure_embeddings=closure_embeddings,
                    floating_dtype=torch.bfloat16,
                )
                variants_by_block[block_index][str(tail_steps)] = cached
                lengths_by_block[block_index][str(tail_steps)] = int(
                    cached["prefix_embeddings"].shape[1]
                    + cached["payload_embeddings"].shape[1]
                    + cached["closure_embeddings"].shape[1]
                )

        for block_index in block_indices:
            block_start = block_index * block_seconds
            path = output_dir / cache_block_filename(block_index)
            torch.save(
                {
                    "schema_version": SCHEMA_VERSION,
                    "block_index": block_index,
                    "start_seconds": block_start,
                    "end_seconds": min(
                        block_start + block_seconds, num_steps * step_seconds
                    ),
                    "variants": variants_by_block[block_index],
                },
                path,
            )
            records.append(
                {
                    "block_index": block_index,
                    "path": path.name,
                    "phase_count": phase_counts[block_index],
                    "token_lengths_without_query": lengths_by_block[block_index],
                    "size_bytes": path.stat().st_size,
                }
            )
        progress.update(len(block_indices))
    progress.close()
    return records


@torch.inference_mode()
def _verify_cache(
    *,
    model: torch.nn.Module,
    extractor: core.OpenStreamReplayExtractor,
    cache_dir: Path,
    static: dict[str, Any],
    anchors: list[dict[str, Any]],
    verify_every: int,
    minimum_cosine: float,
    block_seconds: float,
    step_seconds: float,
) -> list[dict[str, float]]:
    if verify_every <= 0:
        return []
    selected = [
        anchor
        for anchor in anchors
        if (int(anchor["step_index"]) + 1) % verify_every == 0
    ]
    if anchors and anchors[-1] not in selected:
        selected.append(anchors[-1])
    store = CachedBlockStore(cache_dir, max_blocks=max(2, int(static["memory_blocks"])))
    device = _model_device(model)
    checks: list[dict[str, float]] = []
    for anchor in tqdm(selected, desc="verify cached replay", unit="state"):
        embeddings, positions = reconstruct_open_stream_sequence(
            store=store, static=static, anchor=anchor
        )
        embeddings = embeddings.unsqueeze(0).to(device=device, dtype=model.dtype)
        positions = positions.unsqueeze(1).to(device=device)
        attention_mask = torch.ones(
            1, embeddings.shape[1], dtype=torch.long, device=device
        )
        cached_state = forward_last_state(
            model.model,
            inputs_embeds=embeddings,
            position_ids=positions,
            attention_mask=attention_mask,
            readout_indices=torch.tensor(
                [embeddings.shape[1] - 1], dtype=torch.long, device=device
            ),
        )
        exact_state = extractor.state_from_scratch(
            video_path=None,
            audio_path=None,
            intervals=anchor_intervals(
                anchor,
                block_seconds=block_seconds,
                step_seconds=step_seconds,
            ),
            use_audio_in_video=True,
        )
        cached_float = cached_state.float()
        exact_float = exact_state.float()
        check = {
            "time_seconds": float(anchor["time_seconds"]),
            "cosine": float(
                F.cosine_similarity(cached_float, exact_float, dim=-1).item()
            ),
            "mean_abs": float((cached_float - exact_float).abs().mean().item()),
            "max_abs": float((cached_float - exact_float).abs().max().item()),
        }
        checks.append(check)
        print(
            f"verify t={check['time_seconds']:.1f}s "
            f"cosine={check['cosine']:.8f} "
            f"mean_abs={check['mean_abs']:.6f} "
            f"max_abs={check['max_abs']:.6f}",
            flush=True,
        )
        if check["cosine"] < minimum_cosine:
            raise RuntimeError(
                "Cached full-sequence reconstruction disagrees with open-stream "
                f"KV replay at t={check['time_seconds']:.1f}s: "
                f"cosine={check['cosine']:.8f} < {minimum_cosine:.8f}"
            )
    return checks


def main(argv: Sequence[str] | None = None) -> None:
    args = parse_args(argv)
    core._configure_qwen_logging()
    if args.end <= 0:
        raise ValueError("--end must be positive")
    if args.av_batch_size <= 0:
        raise ValueError("--av-batch-size must be positive")
    if (
        args.audio_right_padding_samples is not None
        and args.audio_right_padding_samples < 0
    ):
        raise ValueError("--audio-right-padding-samples must be non-negative")
    if args.av_batch_size != 1 and (
        args.audio_feature_device is not None
        or args.audio_right_padding_samples is not None
    ):
        raise ValueError(
            "Optimized audio preprocessing currently requires --av-batch-size 1; "
            "the HF multimodal batch path was slower and changed cached states"
        )
    steps_per_block = _integer_ratio(
        args.native_block_seconds, args.step_seconds, "native block / step"
    )
    memory_blocks = _integer_ratio(
        args.memory_seconds, args.native_block_seconds, "memory / native block"
    )
    num_steps = _integer_ratio(args.end, args.step_seconds, "end / step")
    targets = _load_targets(args.target)
    if num_steps > len(targets):
        raise ValueError(
            f"Requested {num_steps} cached steps but target has only {len(targets)}"
        )

    output_dir = args.output_dir.expanduser().resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    manifest_path = output_dir / "manifest.json"
    static_path = output_dir / "static.pt"
    _safe_output_path(manifest_path, args.overwrite)
    _safe_output_path(static_path, args.overwrite)

    config = core.StreamConfig(
        step_seconds=args.step_seconds,
        native_block_seconds=args.native_block_seconds,
        memory_seconds=args.memory_seconds,
        video_fps=args.video_fps,
        video_min_pixels=args.video_min_pixels,
        video_max_pixels=args.video_max_pixels,
    )
    config.validate()
    specs = _easycom_directory_specs(args.easycom_session_dir)
    media = CausalSourceAV(
        specs,
        output_fps=args.video_fps,
        required_duration=args.end,
        frame_cache_size=args.frame_cache_size,
        decoder_backend=args.decoder,
        decoder_gpu_id=args.decoder_gpu_id,
        decoder_threads=args.decoder_threads,
        max_open_readers=args.max_open_readers,
    )
    print(
        f"Direct media ready: duration={args.end:.1f}s, "
        f"grid={media.expected_grid_frames} frames, audio={len(media.audio)} samples"
    )

    model, processor = core.load_qwen_omni(
        args.model_id,
        dtype=args.dtype,
        device_map=args.device,
        use_flash_attention_2=args.flash_attention_2,
    )
    model.requires_grad_(False)
    audio_feature_device: str | torch.device | None = args.audio_feature_device
    if audio_feature_device == "model":
        audio_feature_device = _model_device(model)
    audio_padding_description = "model maximum"
    if args.audio_right_padding_samples is not None:
        hop_length = int(getattr(processor.feature_extractor, "hop_length", 1))
        rounded_padding = (
            math.ceil(args.audio_right_padding_samples / hop_length) * hop_length
        )
        audio_padding_description = (
            f"actual + >= {args.audio_right_padding_samples} samples "
            f"(hop-aligned, typically +{rounded_padding})"
        )
    print(
        "Audio preprocessing: "
        f"feature_device={audio_feature_device or 'processor default'}, "
        f"padding={audio_padding_description}"
    )
    extractor = core.OpenStreamReplayExtractor(
        model, processor, config, fixed_grid_av=media
    )
    with torch.inference_mode():
        query_ids = extractor.query_ids
        query_embeddings = model.get_input_embeddings()(query_ids)
    positions_per_chunk = core._tmrope_positions_per_chunk(model)
    static = {
        "schema_version": SCHEMA_VERSION,
        "steps_per_block": steps_per_block,
        "memory_blocks": memory_blocks,
        "positions_per_chunk": positions_per_chunk,
        "query_ids": query_ids.detach().cpu().long(),
        "query_embeddings": query_embeddings.detach().cpu().to(torch.bfloat16),
    }
    torch.save(static, static_path)

    cache_started = time.perf_counter()
    block_records = _cache_blocks(
        extractor=extractor,
        model=model,
        processor=processor,
        media=media,
        config=config,
        output_dir=output_dir,
        num_steps=num_steps,
        steps_per_block=steps_per_block,
        step_seconds=args.step_seconds,
        block_seconds=args.native_block_seconds,
        av_batch_size=args.av_batch_size,
        audio_feature_device=audio_feature_device,
        audio_right_padding_samples=args.audio_right_padding_samples,
        overwrite=args.overwrite,
    )
    cache_seconds = time.perf_counter() - cache_started
    print(
        f"AV cache generation: {cache_seconds:.3f}s, "
        f"{len(block_records) / cache_seconds:.3f} block/s, "
        f"{args.end / cache_seconds:.3f}x media realtime"
    )
    anchors = [
        make_anchor(
            step_index=step_index,
            steps_per_block=steps_per_block,
            memory_blocks=memory_blocks,
            step_seconds=args.step_seconds,
            num_target_steps=len(targets),
            num_horizons=NUM_HORIZONS,
        )
        for step_index in range(num_steps)
    ]
    checks = _verify_cache(
        model=model,
        extractor=extractor,
        cache_dir=output_dir,
        static=static,
        anchors=anchors,
        verify_every=args.verify_every,
        minimum_cosine=args.minimum_cosine,
        block_seconds=args.native_block_seconds,
        step_seconds=args.step_seconds,
    )

    manifest = {
        "schema_version": SCHEMA_VERSION,
        "cache_kind": "qwen2.5-omni-frozen-av-open-stream-phases",
        "session": args.easycom_session_dir.resolve().name,
        "model_id": args.model_id,
        "target_path": str(args.target.expanduser().resolve()),
        "target_steps": int(len(targets)),
        "cached_steps": num_steps,
        "cached_end_seconds": args.end,
        "hidden_size": int(query_embeddings.shape[-1]),
        "floating_dtype": "bfloat16",
        "step_seconds": args.step_seconds,
        "native_block_seconds": args.native_block_seconds,
        "memory_seconds": args.memory_seconds,
        "steps_per_block": steps_per_block,
        "memory_blocks": memory_blocks,
        "positions_per_chunk": positions_per_chunk,
        "video_fps": args.video_fps,
        "video_min_pixels": args.video_min_pixels,
        "video_max_pixels": args.video_max_pixels,
        "use_audio_in_video": True,
        "interval_convention": "half_open_[start,end)",
        "phase_policy": "cache_0.2_to_2.0s;_2.0s_is_completed_history",
        "av_batch_size": args.av_batch_size,
        "audio_feature_device": (
            None if audio_feature_device is None else str(audio_feature_device)
        ),
        "audio_right_padding_samples": args.audio_right_padding_samples,
        "verification_reference_audio_preprocessing": (
            "processor-default-device-and-model-maximum-padding"
        ),
        "cache_generation_seconds": cache_seconds,
        "cache_media_realtime_factor": args.end / cache_seconds,
        "cache_blocks_per_second": len(block_records) / cache_seconds,
        "label_policy": "target[i+1:i+11]",
        "static_path": static_path.name,
        "blocks": block_records,
        "anchors": anchors,
        "verification": checks,
        "media": media.metadata(),
        "software": {
            "python": platform.python_version(),
            "torch": torch.__version__,
        },
    }
    save_json(manifest_path, manifest)
    size = cache_size_bytes(
        [static_path, *[output_dir / record["path"] for record in block_records]]
    )
    print(
        f"Saved {len(block_records)} blocks / {len(anchors)} anchors to {output_dir}"
    )
    print(f"Cache tensor size: {size / (1024 ** 3):.3f} GiB")
    if checks:
        print(
            "Verification: "
            f"lowest cosine={min(item['cosine'] for item in checks):.8f}, "
            f"worst max_abs={max(item['max_abs'] for item in checks):.6f}"
        )


if __name__ == "__main__":
    main()