1090 lines
41 KiB
Python
1090 lines
41 KiB
Python
|
#!/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()
|