#!/usr/bin/env python3 """Extract strictly causal MiniCPM-o 4.5 states from original AV media. MiniCPM-o is used only as a frozen input feature extractor. This script never calls ``streaming_generate`` and therefore never inserts model-generated listen/speak/text/audio tokens into the recorded conversation. Two representations are implemented: ``native-1hz-hold`` Commit one native one-second AV chunk at each integer-second boundary and forward-hold the newest completed state on the 5 Hz label grid. Rows before the first completed chunk use a frozen all-zero initialization. ``cached-partial-5hz`` Keep completed one-second chunks in the official streaming cache. At every non-boundary 200-ms cutoff, snapshot the complete MiniCPM streaming state, append only the current causal partial AV tail with ``is_last_chunk=True``, read its state, and restore the snapshot. At a one-second boundary, append the complete chunk to the main cache. The branch is admitted as ``mode=causal-prefix`` only after mandatory cache replay, future-audio intervention, and native-boundary checks pass. The saved payload follows the existing global-feature contract consumed by ``easycom_head.py``: dense zero-based ``label_indices`` and ``states`` aligned to the 5 Hz target rows. """ from __future__ import annotations import argparse import importlib.metadata import math from contextlib import AbstractContextManager from dataclasses import asdict, dataclass from pathlib import Path from typing import Any, Sequence import numpy as np import torch from tqdm import tqdm import streaming_turn_model_directly as direct_media MODEL_ID = "openbmb/MiniCPM-o-4_5" TESTED_TRANSFORMERS_VERSION = "4.51.0" FEATURE_FORMAT = "minicpmo-causal-global-v1" EPSILON = 1e-7 @dataclass(frozen=True) class MiniCPMStreamConfig: """Frozen extraction and verification settings.""" step_seconds: float = 0.2 native_chunk_seconds: float = 1.0 system_prompt: str = "Streaming Omni Conversation." max_slice_nums: int = 1 boundary_max_abs: float = 1e-3 boundary_min_cosine: float = 0.9999 replay_max_abs: float = 1e-3 replay_min_cosine: float = 0.9999 def validate(self) -> None: if self.step_seconds <= 0: raise ValueError("step_seconds must be positive") if self.native_chunk_seconds <= 0: raise ValueError("native_chunk_seconds must be positive") ratio = self.native_chunk_seconds / self.step_seconds if not math.isclose(ratio, round(ratio), abs_tol=1e-9): raise ValueError("native chunk / step must be integral") if self.max_slice_nums <= 0: raise ValueError("max_slice_nums must be positive") for name, value in ( ("boundary_max_abs", self.boundary_max_abs), ("replay_max_abs", self.replay_max_abs), ): if value < 0: raise ValueError(f"{name} must be non-negative") for name, value in ( ("boundary_min_cosine", self.boundary_min_cosine), ("replay_min_cosine", self.replay_min_cosine), ): if not -1.0 <= value <= 1.0: raise ValueError(f"{name} must be in [-1,1]") @dataclass class ExtractionResult: times: torch.Tensor states: torch.Tensor state_source_times: torch.Tensor native_update_mask: torch.Tensor initialization_mask: torch.Tensor verification: list[dict[str, Any]] partial_tail_admissible: bool def _time_grid(start: float, end: float, step: float) -> list[float]: """Return right-edge cutoffs ``start + k*step`` through ``end``.""" count_f = (end - start) / step count = int(round(count_f)) if not math.isclose(count_f, count, rel_tol=0.0, abs_tol=1e-6): raise ValueError( f"duration must be an integer number of {step:g}s steps: " f"start={start:g}, end={end:g}" ) return [start + (index + 1) * step for index in range(count)] def _is_native_boundary( current_time: float, *, origin: float, chunk_seconds: float, ) -> bool: units = (current_time - origin) / chunk_seconds return math.isclose(units, round(units), rel_tol=0.0, abs_tol=1e-6) def _comparison( cached: torch.Tensor, reference: torch.Tensor, *, kind: str, current_time: float, max_abs_tolerance: float, min_cosine: float, ) -> dict[str, Any]: if cached.shape != reference.shape: raise RuntimeError( f"{kind} shape mismatch at {current_time:g}s: " f"{tuple(cached.shape)} vs {tuple(reference.shape)}" ) cached_f = cached.float().reshape(1, -1) reference_f = reference.float().reshape(1, -1) difference = (cached_f - reference_f).abs() max_abs = float(difference.max().item()) if difference.numel() else 0.0 mean_abs = float(difference.mean().item()) if difference.numel() else 0.0 cached_norm = float(cached_f.norm().item()) reference_norm = float(reference_f.norm().item()) if torch.equal(cached_f, reference_f): cosine = 1.0 elif cached_norm == 0.0 and reference_norm == 0.0: cosine = 1.0 elif cached_norm == 0.0 or reference_norm == 0.0: cosine = 0.0 else: cosine = float( torch.nn.functional.cosine_similarity( cached_f, reference_f, dim=-1 ).item() ) passed = max_abs <= max_abs_tolerance and cosine >= min_cosine return { "kind": kind, "time": float(current_time), "max_abs": max_abs, "mean_abs": mean_abs, "cosine": cosine, "max_abs_tolerance": float(max_abs_tolerance), "min_cosine": float(min_cosine), "passed": bool(passed), } class _FinalHiddenCapture(AbstractContextManager["_FinalHiddenCapture"]): """Capture the final normalized Qwen3 state used by MiniCPM-o's LM head.""" def __init__(self, model: Any) -> None: try: self.module = model.llm.model.norm except AttributeError as exc: raise RuntimeError( "MiniCPM-o readout contract changed: expected model.llm.model.norm" ) from exc self.value: torch.Tensor | None = None self._handle: Any | None = None def __enter__(self) -> "_FinalHiddenCapture": def hook(_module: Any, _inputs: Any, output: Any) -> None: tensor = output[0] if isinstance(output, (tuple, list)) else output if not isinstance(tensor, torch.Tensor) or tensor.ndim != 3: raise RuntimeError( "Unexpected MiniCPM-o final-norm output: " f"{type(tensor)!r} shape={getattr(tensor, 'shape', None)}" ) self.value = tensor[:, -1, :].detach().float().cpu() self._handle = self.module.register_forward_hook(hook) return self def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> None: if self._handle is not None: self._handle.remove() self._handle = None def validate_transformers_version( version: str, *, allow_untested: bool = False, ) -> str: if version != TESTED_TRANSFORMERS_VERSION and not allow_untested: raise RuntimeError( "MiniCPM-o 4.5 officially pins transformers==4.51.0; " f"found {version}. Pass --allow-untested-transformers only for a " "separate compatibility audit." ) return version def _cache_length(cache: Any | None) -> int: if cache is None: return 0 if hasattr(cache, "get_seq_length"): return int(cache.get_seq_length()) if hasattr(cache, "key_cache") and cache.key_cache: return int(cache.key_cache[0].shape[-2]) if isinstance(cache, (tuple, list)) and cache: return int(cache[0][0].shape[-2]) raise TypeError(f"Unsupported cache type: {type(cache)!r}") def _cache_checksum(cache: Any | None) -> float | None: """Match the first-key checksum stored by MiniCPM's snapshot object.""" if cache is None: return None if hasattr(cache, "self_attention_cache"): cache = cache.self_attention_cache if hasattr(cache, "key_cache") and cache.key_cache: return float(cache.key_cache[0].sum().item()) if isinstance(cache, (tuple, list)) and cache: return float(cache[0][0].sum().item()) raise TypeError(f"Unsupported cache type: {type(cache)!r}") class MiniCPMInputOnlyEngine: """Version-gated adapter around MiniCPM-o's simplex streaming prefill.""" def __init__( self, model: Any, processor: Any, config: MiniCPMStreamConfig, *, session_id: str, ) -> None: self.model = model self.processor = processor self.config = config self.session_id = session_id self.last_restore_audit: dict[str, Any] | None = None try: self.hidden_size = int(model.llm.config.hidden_size) except (AttributeError, TypeError, ValueError) as exc: raise RuntimeError("Cannot resolve MiniCPM-o LLM hidden size") from exc for name in ( "streaming_prefill", "save_speculative_snapshot", "restore_speculative_snapshot", "reset_session", ): if not hasattr(model, name): raise RuntimeError( f"MiniCPM-o revision lacks required input-only API: {name}" ) def reset(self) -> None: self.model.reset_session(reset_token2wav_cache=False) self.model.streaming_prefill( session_id=self.session_id, msgs=[ { "role": "system", "content": [self.config.system_prompt], } ], omni_mode=True, max_slice_nums=self.config.max_slice_nums, use_tts_template=False, enable_thinking=False, is_last_chunk=False, processor=self.processor, ) def _prefill( self, *, frame: Any, audio: np.ndarray, is_last_chunk: bool, ) -> torch.Tensor: if not isinstance(audio, np.ndarray) or audio.ndim != 1 or not len(audio): raise ValueError("MiniCPM-o audio must be a non-empty mono ndarray") audio = np.asarray(audio, dtype=np.float32) with _FinalHiddenCapture(self.model) as capture: self.model.streaming_prefill( session_id=self.session_id, msgs=[{"role": "user", "content": [frame, audio]}], omni_mode=True, max_slice_nums=self.config.max_slice_nums, use_tts_template=False, enable_thinking=False, is_last_chunk=is_last_chunk, processor=self.processor, ) if capture.value is None: raise RuntimeError("MiniCPM-o prefill produced no captured LLM state") if capture.value.shape != (1, self.hidden_size): raise RuntimeError( "Unexpected MiniCPM-o state shape: " f"{tuple(capture.value.shape)}, expected (1,{self.hidden_size})" ) if not torch.isfinite(capture.value).all(): raise RuntimeError("MiniCPM-o state contains NaN/Inf") return capture.value def commit_state(self, *, frame: Any, audio: np.ndarray) -> torch.Tensor: return self._prefill(frame=frame, audio=audio, is_last_chunk=False) def temporary_state(self, *, frame: Any, audio: np.ndarray) -> torch.Tensor: self.last_restore_audit = None before_llm = _cache_length(getattr(self.model, "llm_past_key_values", None)) before_audio = _cache_length( getattr(self.model, "audio_past_key_values", None) ) before_chunk = int(getattr(self.model, "audio_chunk_idx", 0)) before_history = tuple( id(item) for item in getattr(self.model, "_omni_chunk_history", ()) ) before_position = int( getattr(self.model, "streaming_position_offset", 0) ) before_window_enabled = bool( getattr(self.model, "streaming_window_enabled", False) ) snapshot = self.model.save_speculative_snapshot() try: # A speculative prefill must never invoke MiniCPM's destructive # round-pruning path: the upstream snapshot truncates appended KV # tokens but cannot recreate a round that pruning removed. Native # committed chunks still retain the model's normal window policy. if hasattr(self.model, "streaming_window_enabled"): self.model.streaming_window_enabled = False try: state = self._prefill( frame=frame, audio=audio, is_last_chunk=True, ) finally: if hasattr(self.model, "streaming_window_enabled"): self.model.streaming_window_enabled = before_window_enabled finally: restored = self.model.restore_speculative_snapshot(snapshot) if not restored: raise RuntimeError("MiniCPM-o failed to restore a partial-tail snapshot") after_llm = _cache_length(getattr(self.model, "llm_past_key_values", None)) after_audio = _cache_length( getattr(self.model, "audio_past_key_values", None) ) after_chunk = int(getattr(self.model, "audio_chunk_idx", 0)) after_history = tuple( id(item) for item in getattr(self.model, "_omni_chunk_history", ()) ) after_position = int( getattr(self.model, "streaming_position_offset", 0) ) after_window_enabled = bool( getattr(self.model, "streaming_window_enabled", False) ) after_llm_checksum = _cache_checksum( getattr(self.model, "llm_past_key_values", None) ) after_audio_checksum = _cache_checksum( getattr(self.model, "audio_past_key_values", None) ) expected_llm_checksum = getattr(snapshot, "llm_cache_checksum", None) expected_audio_checksum = getattr(snapshot, "audio_cache_checksum", None) structure_matches = ( after_llm, after_audio, after_chunk, after_history, after_position, after_window_enabled, ) == ( before_llm, before_audio, before_chunk, before_history, before_position, before_window_enabled, ) checksum_matches = ( after_llm_checksum == expected_llm_checksum and after_audio_checksum == expected_audio_checksum ) self.last_restore_audit = { "kind": "snapshot_restore", "llm_cache_length": after_llm, "audio_cache_length": after_audio, "audio_chunk_index": after_chunk, "streaming_position_offset": after_position, "window_pruning_suspended": bool(before_window_enabled), "window_policy_restored": bool( after_window_enabled == before_window_enabled ), "llm_checksum_matches": bool( after_llm_checksum == expected_llm_checksum ), "audio_checksum_matches": bool( after_audio_checksum == expected_audio_checksum ), "history_matches": bool(after_history == before_history), "passed": bool(structure_matches and checksum_matches), } if not self.last_restore_audit["passed"]: raise RuntimeError( "Partial-tail restore changed committed MiniCPM-o state: " f"{self.last_restore_audit}" ) return state class MiniCPMFeatureExtractor: """Schedule native commits and temporary tails on a causal AV reader.""" def __init__( self, reader: Any, engine: Any, config: MiniCPMStreamConfig, ) -> None: config.validate() self.reader = reader self.engine = engine self.config = config def _frame_audio(self, start: float, end: float) -> tuple[Any, np.ndarray]: # Audio uses [start,end). Excluding a source frame exactly at ``end`` # gives the same half-open cutoff contract as the Qwen extractor. # CausalSourceAV admits an EPSILON tolerance, so nextafter alone is not # enough to exclude a frame whose PTS is exactly the right edge. frame_time = max(start, end - max(4 * direct_media.EPSILON, 1e-6)) frame = self.reader.frame_at_or_before(frame_time) audio = self.reader.audio_slice(start, end) return frame, audio def _commit(self, start: float, end: float) -> torch.Tensor: frame, audio = self._frame_audio(start, end) return self.engine.commit_state(frame=frame, audio=audio) def _temporary(self, start: float, end: float) -> torch.Tensor: frame, audio = self._frame_audio(start, end) return self.engine.temporary_state(frame=frame, audio=audio) def _replay_to( self, *, origin: float, current_time: float, ) -> tuple[torch.Tensor, float]: """Cleanly replay completed chunks and the optional current tail.""" self.engine.reset() cursor = origin state: torch.Tensor | None = None while cursor + self.config.native_chunk_seconds <= current_time + EPSILON: next_cursor = cursor + self.config.native_chunk_seconds state = self._commit(cursor, next_cursor) cursor = next_cursor if cursor < current_time - EPSILON: state = self._temporary(cursor, current_time) if state is None: state = torch.zeros(1, self.engine.hidden_size, dtype=torch.float32) return state, cursor def extract_native_hold( self, *, start: float, end: float, show_progress: bool, ) -> ExtractionResult: times = _time_grid(start, end, self.config.step_seconds) self.engine.reset() latest = torch.zeros(1, self.engine.hidden_size, dtype=torch.float32) source_time = start chunk_start = start states: list[torch.Tensor] = [] source_times: list[float] = [] updates: list[bool] = [] initialization: list[bool] = [] with tqdm( total=len(times), desc="MiniCPM native 1Hz", unit="state", disable=not show_progress, dynamic_ncols=True, ) as progress: for current_time in times: update = _is_native_boundary( current_time, origin=start, chunk_seconds=self.config.native_chunk_seconds, ) if update: latest = self._commit(chunk_start, current_time) chunk_start = current_time source_time = current_time states.append(latest.cpu()) source_times.append(source_time) updates.append(update) initialization.append(source_time <= start + EPSILON) progress.update(1) return ExtractionResult( times=torch.tensor(times, dtype=torch.float64), states=torch.cat(states, dim=0), state_source_times=torch.tensor(source_times, dtype=torch.float64), native_update_mask=torch.tensor(updates, dtype=torch.bool), initialization_mask=torch.tensor(initialization, dtype=torch.bool), verification=[], partial_tail_admissible=True, ) def extract_cached_partial( self, *, start: float, end: float, verify_every: int, future_intervention_every: int, allow_unverified: bool, show_progress: bool, ) -> ExtractionResult: if verify_every < 0 or future_intervention_every < 0: raise ValueError("verification intervals must be non-negative") times = _time_grid(start, end, self.config.step_seconds) self.engine.reset() chunk_start = start states: list[torch.Tensor] = [] updates: list[bool] = [] checks: list[dict[str, Any]] = [] with tqdm( total=len(times), desc="MiniCPM cached partial 5Hz", unit="state", disable=not show_progress, dynamic_ncols=True, ) as progress: for step_index, current_time in enumerate(times, start=1): boundary = _is_native_boundary( current_time, origin=start, chunk_seconds=self.config.native_chunk_seconds, ) if boundary: # Compare a temporary closed view of exactly the same full # chunk with the native non-final commit before admitting # this 5 Hz representation. temporary = self._temporary(chunk_start, current_time) restore_audit = getattr( self.engine, "last_restore_audit", None ) if restore_audit is not None: checks.append( {**restore_audit, "time": float(current_time)} ) state = self._commit(chunk_start, current_time) comparison = _comparison( state, temporary, kind="native_boundary_equivalence", current_time=current_time, max_abs_tolerance=self.config.boundary_max_abs, min_cosine=self.config.boundary_min_cosine, ) checks.append(comparison) if not comparison["passed"] and not allow_unverified: raise RuntimeError( "MiniCPM partial-tail/native boundary check failed " f"at {current_time:g}s: {comparison}" ) chunk_start = current_time else: state = self._temporary(chunk_start, current_time) restore_audit = getattr( self.engine, "last_restore_audit", None ) if restore_audit is not None: checks.append( {**restore_audit, "time": float(current_time)} ) # Always audit the first row. Positive intervals add sampled # checks without turning a full session into O(T^2) replay. replay_due = step_index == 1 or ( verify_every > 0 and step_index % verify_every == 0 ) if replay_due: reference, rebuilt_until = self._replay_to( origin=start, current_time=current_time, ) comparison = _comparison( state, reference, kind="clean_stream_replay", current_time=current_time, max_abs_tolerance=self.config.replay_max_abs, min_cosine=self.config.replay_min_cosine, ) comparison["rebuilt_committed_until"] = float(rebuilt_until) checks.append(comparison) if not comparison["passed"] and not allow_unverified: raise RuntimeError( "MiniCPM clean replay check failed at " f"{current_time:g}s: {comparison}" ) future_due = ( current_time < end - EPSILON and ( step_index == 1 or ( future_intervention_every > 0 and step_index % future_intervention_every == 0 ) ) ) if future_due: first = int( math.floor( current_time * direct_media.SAMPLE_RATE + EPSILON ) ) last = min( len(self.reader.audio), first + direct_media.SAMPLE_RATE, ) original = self.reader.audio[first:last].copy() if last > first: self.reader.audio[first:last] = np.linspace( -0.25, 0.25, last - first, dtype=np.float32, ) try: intervened, _ = self._replay_to( origin=start, current_time=current_time, ) finally: self.reader.audio[first:last] = original comparison = _comparison( state, intervened, kind="future_audio_intervention", current_time=current_time, max_abs_tolerance=0.0, min_cosine=1.0, ) checks.append(comparison) if not comparison["passed"] and not allow_unverified: raise RuntimeError( "Future audio changed a causal MiniCPM state at " f"{current_time:g}s: {comparison}" ) states.append(state.cpu()) updates.append(boundary) progress.update(1) kinds = {check["kind"] for check in checks if check.get("passed")} mandatory = { "native_boundary_equivalence", "clean_stream_replay", "future_audio_intervention", "snapshot_restore", } admissible = mandatory.issubset(kinds) and all( bool(check.get("passed")) for check in checks ) if not admissible and not allow_unverified: missing = sorted(mandatory - kinds) raise RuntimeError( "MiniCPM partial-tail branch lacks mandatory verification; " f"missing={missing}. Run a >=1.2s smoke test or pass " "--allow-unverified-partial-tail to save a deliberately " "non-loadable audit artifact." ) return ExtractionResult( times=torch.tensor(times, dtype=torch.float64), states=torch.cat(states, dim=0), state_source_times=torch.tensor(times, dtype=torch.float64), native_update_mask=torch.tensor(updates, dtype=torch.bool), initialization_mask=torch.zeros(len(times), dtype=torch.bool), verification=checks, partial_tail_admissible=admissible, ) def load_minicpmo( model_id: str, *, revision: str | None, dtype: str, device: str, attention: str, allow_untested_transformers: bool, ) -> tuple[Any, Any, dict[str, Any]]: """Load the frozen AV backbone without the TTS module.""" installed = importlib.metadata.version("transformers") validate_transformers_version( installed, allow_untested=allow_untested_transformers, ) from transformers import AutoModel, AutoProcessor dtype_value: Any = { "auto": "auto", "bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32, }[dtype] common: dict[str, Any] = { "trust_remote_code": True, } if revision is not None: common["revision"] = revision processor = AutoProcessor.from_pretrained(model_id, **common) model = AutoModel.from_pretrained( model_id, torch_dtype=dtype_value, attn_implementation=attention, init_vision=True, init_audio=True, init_tts=False, device_map={"": device}, **common, ) model.eval() model.requires_grad_(False) resolved_revision = getattr(model.config, "_commit_hash", None) metadata = { "transformers_version": installed, "requested_revision": revision, "resolved_revision": resolved_revision, "attention_implementation": attention, "dtype": dtype, "device": device, "init_tts": False, } return model, processor, metadata def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) source = parser.add_mutually_exclusive_group(required=True) source.add_argument("--video", type=Path) source.add_argument("--easycom-session-dir", type=Path) source.add_argument("--easycom-manifest", type=Path) parser.add_argument("--session") parser.add_argument("--audio", type=Path) parser.add_argument("--crop") parser.add_argument("--target", type=Path) parser.add_argument("--output", type=Path) parser.add_argument( "--mode", choices=("native-1hz-hold", "cached-partial-5hz"), required=True, ) parser.add_argument("--model-id", default=MODEL_ID) parser.add_argument( "--revision", help="Pinned Hugging Face revision; strongly recommended for paper runs", ) parser.add_argument("--start", type=float, default=0.0) parser.add_argument("--end", type=float) parser.add_argument("--step-seconds", type=float, default=0.2) parser.add_argument("--native-chunk-seconds", type=float, default=1.0) parser.add_argument( "--system-prompt", default="Streaming Omni Conversation.", ) parser.add_argument("--max-slice-nums", type=int, default=1) parser.add_argument("--boundary-max-abs", type=float, default=1e-3) parser.add_argument("--boundary-min-cosine", type=float, default=0.9999) parser.add_argument("--replay-max-abs", type=float, default=1e-3) parser.add_argument("--replay-min-cosine", type=float, default=0.9999) parser.add_argument( "--verify-every", type=int, default=0, help="Additional clean replays every N rows; the first row is always checked", ) parser.add_argument( "--future-intervention-every", type=int, default=0, help=( "Additional future-audio checks every N rows; the first eligible " "row is always checked" ), ) parser.add_argument("--allow-unverified-partial-tail", action="store_true") parser.add_argument("--frame-cache-size", type=int, default=8) 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=2) parser.add_argument( "--dtype", choices=("auto", "bfloat16", "float16", "float32"), default="bfloat16", ) parser.add_argument( "--attention", choices=("sdpa", "flash_attention_2"), default="sdpa", ) parser.add_argument("--device", default="cuda:0") parser.add_argument("--allow-untested-transformers", action="store_true") parser.add_argument("--media-only", action="store_true") parser.add_argument("--no-progress", action="store_true") return parser.parse_args(argv) def _validate_args(args: argparse.Namespace) -> None: if args.output is None and not args.media_only: raise SystemExit("--output is required unless --media-only is used") if args.audio is not None and args.video is None: raise SystemExit("--audio is supported only with --video") if args.session is not None and args.easycom_manifest is None: raise SystemExit("--session is used only with --easycom-manifest") if args.easycom_manifest is not None and args.session is None: raise SystemExit("--easycom-manifest requires --session Session_N") if args.frame_cache_size < 0: raise SystemExit("--frame-cache-size must be non-negative") if args.verify_every < 0 or args.future_intervention_every < 0: raise SystemExit("verification intervals must be non-negative") if not args.media_only and args.revision is None: print( "WARNING: --revision is not pinned; this run is not paper-reproducible.", flush=True, ) def _build_reader( args: argparse.Namespace, ) -> tuple[Any, Path | None, int | None, float, str, Any]: manifest_record = None manifest_target = None if args.video is not None: specs = [ direct_media.SourceClipSpec( path=args.video, crop=direct_media._parse_crop(args.crop), audio_path=args.audio, ) ] source_kind = "single_original_video" elif args.easycom_session_dir is not None: if args.crop is not None: raise SystemExit( "EasyCom direct extraction uses the fixed rightmost-1920 crop" ) specs = direct_media._easycom_directory_specs(args.easycom_session_dir) source_kind = "easycom_dense_original_clips" else: if args.crop is not None: raise SystemExit("EasyCom manifest supplies crop metadata") specs, manifest_record, manifest_target = direct_media._manifest_session_specs( args.easycom_manifest.expanduser().resolve(), args.session ) source_kind = "easycom_manifest_original_clips" target_path, target_steps, end = direct_media._resolve_duration_and_target( args, manifest_target ) reader = direct_media.CausalSourceAV( specs, output_fps=1.0 / args.step_seconds, required_duration=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, ) is_easycom = ( args.easycom_session_dir is not None or args.easycom_manifest is not None ) if is_easycom: for clip in reader.clips: if clip.height != direct_media.EASYCOM_HEIGHT: raise RuntimeError( f"Expected EasyCom height {direct_media.EASYCOM_HEIGHT}, " f"got {clip.height}: {clip.path}" ) if not math.isclose(clip.source_fps, 20.0, abs_tol=1e-6): raise RuntimeError( f"Expected EasyCom source FPS 20, got {clip.source_fps:g}: " f"{clip.path}" ) if target_steps is not None: source_frames = sum(clip.num_frames for clip in reader.clips) media_steps = source_frames // 4 if media_steps != target_steps: raise RuntimeError( f"EasyCom source gives {media_steps} 5-Hz rows, target has " f"{target_steps}" ) direct_media._print_media_summary(reader, args.start, end) return reader, target_path, target_steps, end, source_kind, manifest_record def main( argv: Sequence[str] | None = None, *, model: Any | None = None, processor: Any | None = None, ) -> None: args = parse_args(argv) _validate_args(args) if (model is None) != (processor is None): raise ValueError("model and processor must be supplied together") config = MiniCPMStreamConfig( step_seconds=args.step_seconds, native_chunk_seconds=args.native_chunk_seconds, system_prompt=args.system_prompt, max_slice_nums=args.max_slice_nums, boundary_max_abs=args.boundary_max_abs, boundary_min_cosine=args.boundary_min_cosine, replay_max_abs=args.replay_max_abs, replay_min_cosine=args.replay_min_cosine, ) config.validate() if not math.isclose(args.native_chunk_seconds, 1.0, abs_tol=1e-9): raise SystemExit( "MiniCPM-o 4.5 native streaming is fixed at 1.0s; " "--native-chunk-seconds is exposed only to reject protocol drift" ) ( reader, target_path, target_steps, end, source_kind, manifest_record, ) = _build_reader(args) if args.media_only: probe_end = min(end, args.start + args.step_seconds) reader.frame_at_or_before(float(np.nextafter(probe_end, -np.inf))) reader.audio_slice(args.start, probe_end) print(f"Decoded MiniCPM media probe successfully with {args.decoder}") return model_metadata: dict[str, Any] if model is None: model, processor, model_metadata = load_minicpmo( args.model_id, revision=args.revision, dtype=args.dtype, device=args.device, attention=args.attention, allow_untested_transformers=args.allow_untested_transformers, ) else: model.eval() model.requires_grad_(False) model_metadata = { "transformers_version": importlib.metadata.version("transformers"), "requested_revision": args.revision, "resolved_revision": getattr(model.config, "_commit_hash", None), "preloaded": True, } session_name = args.session or ( args.easycom_session_dir.name if args.easycom_session_dir is not None else args.video.stem ) engine = MiniCPMInputOnlyEngine( model, processor, config, session_id=f"feature-extract-{session_name}", ) extractor = MiniCPMFeatureExtractor(reader, engine, config) if args.mode == "native-1hz-hold": result = extractor.extract_native_hold( start=args.start, end=end, show_progress=not args.no_progress, ) else: result = extractor.extract_cached_partial( start=args.start, end=end, verify_every=args.verify_every, future_intervention_every=args.future_intervention_every, allow_unverified=args.allow_unverified_partial_tail, show_progress=not args.no_progress, ) if target_steps is not None and len(result.times) != target_steps: raise RuntimeError( f"Extractor returned {len(result.times)} states, expected {target_steps}" ) loadable_mode = ( "causal-prefix" if result.partial_tail_admissible else "minicpmo-unverified-partial-tail" ) payload = { "format": FEATURE_FORMAT, "model_id": args.model_id, "mode": loadable_mode, "extractor_mode": args.mode, "stream_representation": ( "native_1s_input_only_state_forward_held_to_5hz" if args.mode == "native-1hz-hold" else "committed_1s_cache_plus_temporary_causal_partial_tail_at_5hz" ), "config": asdict(config), "model_metadata": model_metadata, "strict_causal": True, "generation_called": False, "tts_loaded": False, "partial_tail_admissible": result.partial_tail_admissible, "initialization_policy": ( "zeros_until_first_completed_native_chunk" if args.mode == "native-1hz-hold" else None ), "start": args.start, "end": end, "num_steps": len(result.times), "label_indices": torch.arange(len(result.times), dtype=torch.long), "valid_10_horizon_steps": max(0, len(result.times) - 10), "times": result.times, "state_source_times": result.state_source_times, "native_update_mask": result.native_update_mask, "initialization_mask": result.initialization_mask, "states": result.states, "verification": result.verification, "target_path": str(target_path) if target_path is not None else None, "source_kind": source_kind, "source_manifest_record": manifest_record, "direct_media": reader.metadata(), "video_sampling_grid": "latest_source_pts_strictly_before_each_cutoff", "audio_interval": "half_open_[start,end)", } args.output.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, args.output) print(f"Saved {len(result.times)} MiniCPM states to {args.output}") print(f"states: {tuple(result.states.shape)}") print(f"mode: {loadable_mode}; extractor_mode: {args.mode}") if result.verification: kinds = sorted({item["kind"] for item in result.verification}) for kind in kinds: items = [item for item in result.verification if item["kind"] == kind] worst_abs = max(float(item["max_abs"]) for item in items) worst_cosine = min(float(item["cosine"]) for item in items) print( f" {kind}: n={len(items)}, max_abs={worst_abs:.6g}, " f"min_cosine={worst_cosine:.9f}" ) if __name__ == "__main__": main()