212 lines
11 KiB
Python
212 lines
11 KiB
Python
"""Shared AV preparation, independent strict-window Thinker replay banks.
|
|
|
|
Only encoder-native block embeddings are shared. High-layer KV is never
|
|
cropped from a longer context to stand in for a shorter-context extraction.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass, field, replace
|
|
from pathlib import Path
|
|
from typing import Any, Sequence
|
|
|
|
import torch
|
|
from tqdm.auto import tqdm
|
|
|
|
import streaming_turn_model as core
|
|
|
|
DEFAULT_MEMORY_SECONDS = (2.0, 4.0, 6.0, 8.0, 10.0, 12.0)
|
|
|
|
|
|
def validate_memory_seconds(values: Sequence[float], native_seconds: float) -> tuple[float, ...]:
|
|
"""Reject ambiguous/non-native budgets before loading media or a model."""
|
|
if not math.isfinite(native_seconds) or native_seconds <= 0:
|
|
raise ValueError("native block seconds must be finite and positive")
|
|
caps = tuple(float(value) for value in values)
|
|
if not caps or len(set(caps)) != len(caps):
|
|
raise ValueError("memory seconds list must be nonempty with no duplicates")
|
|
for cap in caps:
|
|
if not math.isfinite(cap) or cap < native_seconds:
|
|
raise ValueError("every memory cap must be finite and at least one native block")
|
|
blocks = core._require_integer_ratio(cap, native_seconds, "memory / native block")
|
|
if cap != blocks * native_seconds:
|
|
raise ValueError("memory caps must be exact native-block multiples, not rounded budgets")
|
|
return tuple(sorted(caps))
|
|
|
|
|
|
def context_directory(cap: float) -> str:
|
|
number = str(int(cap)) if float(cap).is_integer() else repr(float(cap))
|
|
return f"context_{number}s"
|
|
|
|
|
|
def output_paths(output: Path, caps: Sequence[float]) -> dict[float, Path]:
|
|
return {cap: output.parent / context_directory(cap) / output.name for cap in caps}
|
|
|
|
|
|
@dataclass
|
|
class _Bank:
|
|
helper: core.OpenStreamReplayExtractor
|
|
history_start: float
|
|
committed_cache: Any = None
|
|
payload_base: int | None = None
|
|
states: list[torch.Tensor] = field(default_factory=list)
|
|
checks: list[dict[str, Any]] = field(default_factory=list)
|
|
history_starts: list[float] = field(default_factory=list)
|
|
rebuilds: int = 0
|
|
replayed_blocks: int = 0
|
|
|
|
|
|
class MultiWindowOpenStreamReplayExtractor(core.OpenStreamReplayExtractor):
|
|
"""Same tail-aware semantics as separate single-window extractions.
|
|
|
|
Each bank has an independent KV history. All banks consume the same prepared
|
|
current tail and a bounded shared ring of completed AV block embeddings.
|
|
Temporary tails, chat closure and readout queries are not committed until a
|
|
native boundary; only the completed media payload is committed there.
|
|
"""
|
|
|
|
def __init__(self, model, processor, config, *, memory_seconds_list, fixed_grid_av=None):
|
|
self.memory_seconds_list = validate_memory_seconds(
|
|
memory_seconds_list, config.native_block_seconds
|
|
)
|
|
super().__init__(
|
|
model, processor, replace(config, memory_seconds=max(self.memory_seconds_list)),
|
|
fixed_grid_av=fixed_grid_av,
|
|
)
|
|
|
|
@torch.inference_mode()
|
|
def extract(
|
|
self, *, video_path: str | Path | None, audio_path: str | Path | None,
|
|
start: float, end: float, use_audio_in_video: bool,
|
|
verify_every: int = 0, verify_continuous_every: int = 0,
|
|
show_progress: bool = True,
|
|
) -> dict[float, core.ExtractionResult]:
|
|
if verify_every < 0 or verify_continuous_every < 0:
|
|
raise ValueError("verification intervals must be non-negative")
|
|
if not math.isclose(start, 0.0, abs_tol=1e-9) or not math.isclose(
|
|
self.config.step_seconds, 0.2, abs_tol=1e-9
|
|
):
|
|
raise ValueError("multiwindow dense banks require start=0 and step_seconds=0.2")
|
|
if not math.isfinite(end) or end <= start:
|
|
raise ValueError("end must be finite and greater than start")
|
|
core._require_integer_ratio(end - start, self.config.step_seconds, "duration / step")
|
|
times = core._time_grid(start, end, self.config.step_seconds)
|
|
native = self.config.native_block_seconds
|
|
positions_per_chunk = core._tmrope_positions_per_chunk(self.model)
|
|
banks = {
|
|
cap: _Bank(
|
|
helper=core.OpenStreamReplayExtractor(
|
|
self.model, self.processor, replace(self.config, memory_seconds=cap),
|
|
fixed_grid_av=self.fixed_grid_av,
|
|
),
|
|
history_start=start,
|
|
)
|
|
for cap in self.memory_seconds_list
|
|
}
|
|
completed = core.CompletedBlockPieceCache(
|
|
origin=start, block_seconds=native,
|
|
max_blocks=round(max(self.memory_seconds_list) / native) - 1,
|
|
)
|
|
block_start = start
|
|
prepared_seconds = 0.0
|
|
with tqdm(total=len(times), desc="multiwindow-replay", unit="cutoff",
|
|
disable=not show_progress, dynamic_ncols=True) as progress:
|
|
for step, cutoff in enumerate(times, start=1):
|
|
for cap, bank in banks.items():
|
|
if bank.committed_cache is not None and block_start - bank.history_start >= cap - 1e-6:
|
|
bank.history_start = block_start - (cap - native)
|
|
intervals = core._block_intervals(bank.history_start, block_start, native)
|
|
bank.rebuilds += 1
|
|
bank.replayed_blocks += len(intervals)
|
|
if intervals:
|
|
bank.committed_cache, _, bank.payload_base, _, _ = bank.helper._cache_from_intervals(
|
|
video_path=video_path, audio_path=audio_path, intervals=intervals,
|
|
use_audio_in_video=use_audio_in_video,
|
|
prepared_blocks=completed.get_many(intervals),
|
|
)
|
|
else:
|
|
bank.committed_cache = None
|
|
bank.payload_base = None
|
|
|
|
# Exactly one production processor/tower invocation per cutoff.
|
|
pieces = self._prepare_pieces(
|
|
video_path=video_path, audio_path=audio_path,
|
|
interval=(block_start, cutoff), use_audio_in_video=use_audio_in_video,
|
|
)
|
|
prepared_seconds += cutoff - block_start
|
|
boundary = math.isclose(cutoff, block_start + native, abs_tol=1e-6)
|
|
for cap, bank in banks.items():
|
|
if bank.payload_base is None:
|
|
bank.payload_base = int(pieces.payload_position_ids.min().item())
|
|
block_index = round((block_start - bank.history_start) / native)
|
|
media_cache, _, closure_positions = core._append_open_pieces(
|
|
self.model, pieces, past_key_values=core.clone_cache(bank.committed_cache),
|
|
rope_offset=bank.payload_base + block_index * positions_per_chunk,
|
|
include_prefix=bank.committed_cache is None,
|
|
)
|
|
state = core._read_open_state(
|
|
self.model, media_cache, closure_ids=pieces.closure_ids,
|
|
closure_position_ids=closure_positions, query_ids=bank.helper.query_ids,
|
|
last_n_layers=self.config.readout_last_n_layers,
|
|
layer_fusion=self.config.readout_layer_fusion,
|
|
)
|
|
bank.states.append(state.cpu())
|
|
bank.history_starts.append(bank.history_start)
|
|
intervals = core._block_intervals(bank.history_start, cutoff, native)
|
|
if verify_every and step % verify_every == 0:
|
|
# Deliberately independent processor/tower calls, not the shared ring.
|
|
reference = bank.helper.state_from_scratch(
|
|
video_path=video_path, audio_path=audio_path, intervals=intervals,
|
|
use_audio_in_video=use_audio_in_video,
|
|
)
|
|
bank.checks.append(core._compare_states(cutoff, state, reference, kind="cache_replay"))
|
|
if verify_continuous_every and step % verify_continuous_every == 0:
|
|
reference = bank.helper.continuous_reference.state_for_intervals(
|
|
video_path=video_path, audio_path=audio_path,
|
|
intervals=[(bank.history_start, cutoff)], use_audio_in_video=use_audio_in_video,
|
|
)
|
|
bank.checks.append(core._compare_states(
|
|
cutoff, state, reference, kind="continuous_window_semantics"
|
|
))
|
|
if boundary:
|
|
bank.committed_cache = media_cache
|
|
if boundary:
|
|
completed.put(core.CachedOpenMediaBlock(block_start, cutoff, pieces))
|
|
block_start = cutoff
|
|
progress.set_postfix_str(f"media={cutoff:.1f}s banks={len(banks)}", refresh=False)
|
|
progress.update(1)
|
|
|
|
time_tensor = torch.tensor(times, dtype=torch.float64)
|
|
boundary_mask = core._periodic_boundary_mask(times, origin=start, period=native)
|
|
results = {}
|
|
for cap, bank in banks.items():
|
|
history = torch.tensor(bank.history_starts, dtype=torch.float64)
|
|
span = time_tensor - history
|
|
if (span <= 0).any() or (span > cap + 1e-6).any():
|
|
raise RuntimeError(f"invalid context span in {cap:g}s bank")
|
|
results[cap] = core.ExtractionResult(
|
|
times=time_tensor.clone(), states=torch.cat(bank.states, dim=0),
|
|
verification=bank.checks,
|
|
diagnostics={
|
|
"history_start_times": history,
|
|
"actual_context_seconds": span,
|
|
"state_source_times": time_tensor.clone(),
|
|
"state_age_seconds": torch.zeros_like(time_tensor),
|
|
"initialization_mask": torch.zeros(len(times), dtype=torch.bool),
|
|
"state_update_mask": torch.ones(len(times), dtype=torch.bool),
|
|
"payload_commit_mask": boundary_mask.clone(),
|
|
"native_boundary_mask": boundary_mask.clone(),
|
|
"production_processor_calls_shared": len(times),
|
|
"production_prepared_media_seconds_shared": prepared_seconds,
|
|
"production_thinker_tail_calls": len(times),
|
|
"rolling_rebuilds": bank.rebuilds,
|
|
"replayed_complete_blocks": bank.replayed_blocks,
|
|
"shared_completed_block_capacity": round(max(self.memory_seconds_list) / native) - 1,
|
|
"memory_seconds_list": list(self.memory_seconds_list),
|
|
"sharing": "AV_embeddings_only_independent_Thinker_KV",
|
|
"verification_excluded_from_production_counters": True,
|
|
},
|
|
)
|
|
return results
|