diff --git a/cache_omni_av_pieces.py b/cache_omni_av_pieces.py index 1343b1e..77a4e5e 100644 --- a/cache_omni_av_pieces.py +++ b/cache_omni_av_pieces.py @@ -8,7 +8,9 @@ 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. +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 @@ -16,6 +18,7 @@ from __future__ import annotations import argparse import math import platform +import time from pathlib import Path from typing import Any, Sequence @@ -60,6 +63,15 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: 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) @@ -118,70 +130,256 @@ def _model_device(model: torch.nn.Module) -> torch.device: 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() - for block_index in tqdm(range(num_blocks), desc="cache AV blocks", unit="block"): - first_step = block_index * steps_per_block - phase_count = min(steps_per_block, num_steps - first_step) - block_start = block_index * block_seconds - variants: dict[str, dict[str, torch.Tensor]] = {} - token_lengths: dict[str, int] = {} - for tail_steps in range(1, phase_count + 1): - interval_end = block_start + tail_steps * step_seconds - pieces = extractor._prepare_pieces( - video_path=None, - audio_path=None, - interval=(block_start, interval_end), - use_audio_in_video=True, - ) - closure_embeddings = input_embeddings(pieces.closure_ids) - cached = pieces_to_cache( - pieces, - closure_embeddings=closure_embeddings, - floating_dtype=torch.bfloat16, - ) - variants[str(tail_steps)] = cached - token_lengths[str(tail_steps)] = int( - cached["prefix_embeddings"].shape[1] - + cached["payload_embeddings"].shape[1] - + cached["closure_embeddings"].shape[1] - ) - path = output_dir / cache_block_filename(block_index) - _safe_output_path(path, overwrite) - 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, - }, - path, - ) - records.append( - { - "block_index": block_index, - "path": path.name, - "phase_count": phase_count, - "token_lengths_without_query": token_lengths, - "size_bytes": path.stat().st_size, - } + 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 @@ -270,6 +468,8 @@ def main(argv: Sequence[str] | None = None) -> None: 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" ) @@ -339,16 +539,22 @@ def main(argv: Sequence[str] | None = None) -> None: } 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, @@ -395,6 +601,10 @@ def main(argv: Sequence[str] | None = None) -> None: "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, @@ -413,6 +623,11 @@ def main(argv: Sequence[str] | None = None) -> None: 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(