#!/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. """ 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("--decoder", choices=("decord", "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, 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, ) ] 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") 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) 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, overwrite=args.overwrite, ) cache_seconds = time.perf_counter() - cache_started 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, "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"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" ) 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()