709 lines
26 KiB
Python
709 lines
26 KiB
Python
#!/usr/bin/env python3
|
|
"""Cache exact Qwen2.5-Omni AV pieces for Thinker LoRA training.
|
|
|
|
For every source-aligned 2-second block this script caches ten causal phases:
|
|
0.2, 0.4, ..., 2.0 seconds. The 2.0-second phase is reused as completed
|
|
history; phases 0.2--1.8 are used only for the current partial block. No
|
|
future samples enter a phase.
|
|
|
|
The optional verification compares a full-sequence reconstruction from the
|
|
cache with the existing open-stream KV replay. Run this 24-second pilot before
|
|
caching a full EasyCom session. ``--av-batch-size`` batches equal-duration
|
|
phases from different source blocks, increasing frozen AV-tower occupancy while
|
|
preserving the exact per-block cache schema and causal timestamps.
|
|
|
|
The audio optimization flags keep the official Qwen/Whisper feature extractor
|
|
but may run its STFT on CUDA and bound the otherwise model-maximum zero padding.
|
|
Verification deliberately recomputes its oracle with the original processor
|
|
defaults, so an optimization cannot silently redefine the extracted feature.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import math
|
|
import platform
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Sequence
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from tqdm.auto import tqdm
|
|
|
|
import streaming_turn_model as core
|
|
from easycom_head import NUM_HORIZONS
|
|
from omni_cached_av import (
|
|
SCHEMA_VERSION,
|
|
CachedBlockStore,
|
|
anchor_intervals,
|
|
cache_block_filename,
|
|
cache_size_bytes,
|
|
forward_last_state,
|
|
make_anchor,
|
|
pieces_to_cache,
|
|
reconstruct_open_stream_sequence,
|
|
save_json,
|
|
)
|
|
from streaming_turn_model_directly import CausalSourceAV, _easycom_directory_specs
|
|
|
|
|
|
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--model-id", required=True)
|
|
parser.add_argument("--easycom-session-dir", type=Path, required=True)
|
|
parser.add_argument("--target", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument(
|
|
"--end",
|
|
type=float,
|
|
default=24.0,
|
|
help="Dense session endpoint in seconds; use 24 for the first pilot",
|
|
)
|
|
parser.add_argument("--step-seconds", type=float, default=0.2)
|
|
parser.add_argument("--native-block-seconds", type=float, default=2.0)
|
|
parser.add_argument("--memory-seconds", type=float, default=8.0)
|
|
parser.add_argument("--video-fps", type=float, default=2.0)
|
|
parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28)
|
|
parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28)
|
|
parser.add_argument("--frame-cache-size", type=int, default=64)
|
|
parser.add_argument(
|
|
"--av-batch-size",
|
|
type=int,
|
|
default=1,
|
|
help=(
|
|
"Microbatch source blocks that have the same 0.2s phase through "
|
|
"the processor and frozen AV towers; 1 preserves serial behavior"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--audio-feature-device",
|
|
default=None,
|
|
help=(
|
|
"Device used by the native Whisper log-Mel extractor, for example "
|
|
"model, cuda:0, or cpu; omitted preserves processor defaults"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--audio-right-padding-samples",
|
|
type=int,
|
|
default=None,
|
|
help=(
|
|
"Replace model-maximum audio padding with actual samples plus this "
|
|
"many zeros, rounded up to the Whisper hop; omitted preserves the "
|
|
"original maximum-length padding"
|
|
),
|
|
)
|
|
parser.add_argument("--decoder", choices=("decord", "pynvcodec"), default="decord")
|
|
parser.add_argument("--decoder-gpu-id", type=int, default=0)
|
|
parser.add_argument("--decoder-threads", type=int, default=1)
|
|
parser.add_argument("--max-open-readers", type=int, default=1)
|
|
parser.add_argument(
|
|
"--dtype",
|
|
choices=("auto", "bfloat16", "float16", "float32"),
|
|
default="bfloat16",
|
|
)
|
|
parser.add_argument("--device", default="cuda:0")
|
|
parser.add_argument("--flash-attention-2", action="store_true")
|
|
parser.add_argument(
|
|
"--verify-every",
|
|
type=int,
|
|
default=10,
|
|
help="Compare every Nth 5-Hz state with exact KV replay; 0 disables",
|
|
)
|
|
parser.add_argument(
|
|
"--minimum-cosine",
|
|
type=float,
|
|
default=0.999,
|
|
help="Fail if cached full-forward/reference KV cosine is below this",
|
|
)
|
|
parser.add_argument(
|
|
"--overwrite",
|
|
action="store_true",
|
|
help="Overwrite this cache's known block/static/manifest files",
|
|
)
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def _integer_ratio(numerator: float, denominator: float, name: str) -> int:
|
|
ratio = numerator / denominator
|
|
rounded = round(ratio)
|
|
if not math.isclose(ratio, rounded, rel_tol=0.0, abs_tol=1e-6):
|
|
raise ValueError(f"{name} must be integral, got {ratio}")
|
|
return int(rounded)
|
|
|
|
|
|
def _load_targets(path: Path) -> np.ndarray:
|
|
targets = np.load(path, mmap_mode="r")
|
|
if targets.ndim != 2 or targets.shape[1] != 3:
|
|
raise ValueError(f"Expected target [T,3], got {targets.shape}: {path}")
|
|
return targets
|
|
|
|
|
|
def _safe_output_path(path: Path, overwrite: bool) -> None:
|
|
if path.exists() and not overwrite:
|
|
raise FileExistsError(f"Refusing to overwrite {path}; pass --overwrite")
|
|
|
|
|
|
def _model_device(model: torch.nn.Module) -> torch.device:
|
|
try:
|
|
return next(model.parameters()).device
|
|
except StopIteration as exc:
|
|
raise RuntimeError("Model has no parameters") from exc
|
|
|
|
|
|
def _one_video_item(value: Any) -> Any:
|
|
if isinstance(value, (list, tuple)):
|
|
if len(value) != 1:
|
|
raise ValueError(f"Expected one prepared video, got {len(value)}")
|
|
return value[0]
|
|
return value
|
|
|
|
|
|
@torch.inference_mode()
|
|
def _prepare_pieces_batch(
|
|
*,
|
|
model: torch.nn.Module,
|
|
processor: Any,
|
|
media: Any,
|
|
config: core.StreamConfig,
|
|
intervals: list[tuple[float, float]],
|
|
) -> list[core.OpenMediaPieces]:
|
|
"""Encode equal-duration AV intervals as one independent-sample batch."""
|
|
|
|
if not intervals:
|
|
raise ValueError("Cannot prepare an empty AV batch")
|
|
durations = [end - start for start, end in intervals]
|
|
if any(
|
|
not math.isclose(value, durations[0], rel_tol=0.0, abs_tol=1e-6)
|
|
for value in durations[1:]
|
|
):
|
|
raise ValueError(f"AV microbatch intervals have mixed durations: {durations}")
|
|
|
|
texts: list[str] = []
|
|
audios: list[Any] = []
|
|
videos: list[Any] = []
|
|
for interval in intervals:
|
|
conversation, manual_audios = core.build_chunk_conversation(
|
|
video_path=None,
|
|
audio_path=None,
|
|
intervals=[interval],
|
|
use_audio_in_video=True,
|
|
config=config,
|
|
fixed_grid_av=media,
|
|
)
|
|
if manual_audios is None or len(manual_audios) != 1:
|
|
raise RuntimeError("Direct-media interval did not provide one audio slice")
|
|
text = core._single_template_text(
|
|
processor.apply_chat_template(
|
|
conversation,
|
|
add_generation_prompt=False,
|
|
tokenize=False,
|
|
)
|
|
)
|
|
_, images, prepared_videos = core.process_mm_info(
|
|
conversation,
|
|
use_audio_in_video=False,
|
|
)
|
|
if images is not None and len(images):
|
|
raise RuntimeError("AV cache batch unexpectedly produced image inputs")
|
|
if prepared_videos is None:
|
|
raise RuntimeError("AV cache batch did not produce video inputs")
|
|
texts.append(text)
|
|
audios.append(manual_audios[0])
|
|
videos.append(_one_video_item(prepared_videos))
|
|
|
|
batch = processor(
|
|
text=texts,
|
|
audio=audios,
|
|
videos=videos,
|
|
padding=True,
|
|
use_audio_in_video=True,
|
|
fps=config.video_fps,
|
|
)
|
|
batch = core._processor_batch_convert_to_tensors(batch)
|
|
batch = core._processor_batch_to_device(batch, model.device)
|
|
batch = core._processor_batch_to_dtype(batch, model.dtype)
|
|
inputs = dict(batch)
|
|
input_ids = inputs["input_ids"]
|
|
if input_ids.shape[0] != len(intervals):
|
|
raise RuntimeError(
|
|
f"Processor returned batch {input_ids.shape[0]} for {len(intervals)} intervals"
|
|
)
|
|
|
|
embeddings = core._multimodal_embeddings(model, inputs)
|
|
positions = core._local_position_ids(
|
|
model,
|
|
inputs,
|
|
use_audio_in_video=True,
|
|
)
|
|
attention_mask = inputs.get("attention_mask")
|
|
audio_id = core._config_token_id(model.config, "audio")
|
|
video_id = core._config_token_id(model.config, "video")
|
|
positions_per_chunk = core._tmrope_positions_per_chunk(model)
|
|
pieces_batch: list[core.OpenMediaPieces] = []
|
|
for row in range(len(intervals)):
|
|
payload_mask = (input_ids[row] == audio_id) | (input_ids[row] == video_id)
|
|
indices = payload_mask.nonzero(as_tuple=False).flatten()
|
|
if indices.numel() == 0:
|
|
raise ValueError(f"Prepared AV batch row {row} has no payload tokens")
|
|
first = int(indices[0].item())
|
|
last = int(indices[-1].item())
|
|
if not bool(payload_mask[first : last + 1].all().item()):
|
|
raise RuntimeError(f"AV payload is not contiguous in batch row {row}")
|
|
if attention_mask is None:
|
|
valid_end = input_ids.shape[1]
|
|
else:
|
|
valid = attention_mask[row].nonzero(as_tuple=False).flatten()
|
|
if valid.numel() == 0:
|
|
raise RuntimeError(f"Prepared AV batch row {row} is fully padded")
|
|
valid_end = int(valid[-1].item()) + 1
|
|
if valid_end <= last + 1:
|
|
raise RuntimeError(f"Prepared AV batch row {row} has no closure tokens")
|
|
pieces = core.OpenMediaPieces(
|
|
prefix_embeddings=embeddings[row : row + 1, :first, :],
|
|
prefix_position_ids=positions[:, row : row + 1, :first],
|
|
payload_embeddings=embeddings[row : row + 1, first : last + 1, :],
|
|
payload_position_ids=positions[:, row : row + 1, first : last + 1],
|
|
closure_ids=input_ids[row : row + 1, last + 1 : valid_end],
|
|
closure_position_ids=positions[
|
|
:, row : row + 1, last + 1 : valid_end
|
|
],
|
|
)
|
|
temporal = pieces.payload_position_ids[0]
|
|
temporal_span = int(temporal.max().item() - temporal.min().item())
|
|
if temporal_span >= positions_per_chunk:
|
|
raise RuntimeError(
|
|
"One batched AV phase spans multiple native TMRoPE chunks: "
|
|
f"interval={intervals[row]}, span={temporal_span}, "
|
|
f"chunk={positions_per_chunk}"
|
|
)
|
|
pieces_batch.append(pieces)
|
|
return pieces_batch
|
|
|
|
|
|
@torch.inference_mode()
|
|
def _cache_blocks(
|
|
*,
|
|
extractor: core.OpenStreamReplayExtractor,
|
|
model: torch.nn.Module,
|
|
processor: Any,
|
|
media: Any,
|
|
config: core.StreamConfig,
|
|
output_dir: Path,
|
|
num_steps: int,
|
|
steps_per_block: int,
|
|
step_seconds: float,
|
|
block_seconds: float,
|
|
av_batch_size: int,
|
|
audio_feature_device: str | torch.device | None,
|
|
audio_right_padding_samples: int | None,
|
|
overwrite: bool,
|
|
) -> list[dict[str, Any]]:
|
|
num_blocks = math.ceil(num_steps / steps_per_block)
|
|
records: list[dict[str, Any]] = []
|
|
input_embeddings = model.get_input_embeddings()
|
|
progress = tqdm(
|
|
total=num_blocks,
|
|
desc=f"cache AV blocks (batch={av_batch_size})",
|
|
unit="block",
|
|
)
|
|
for tile_start in range(0, num_blocks, av_batch_size):
|
|
block_indices = list(
|
|
range(tile_start, min(tile_start + av_batch_size, num_blocks))
|
|
)
|
|
variants_by_block: dict[int, dict[str, dict[str, torch.Tensor]]] = {
|
|
index: {} for index in block_indices
|
|
}
|
|
lengths_by_block: dict[int, dict[str, int]] = {
|
|
index: {} for index in block_indices
|
|
}
|
|
phase_counts = {
|
|
index: min(
|
|
steps_per_block,
|
|
num_steps - index * steps_per_block,
|
|
)
|
|
for index in block_indices
|
|
}
|
|
for block_index in block_indices:
|
|
_safe_output_path(
|
|
output_dir / cache_block_filename(block_index), overwrite
|
|
)
|
|
|
|
for tail_steps in range(1, steps_per_block + 1):
|
|
active_indices = [
|
|
index
|
|
for index in block_indices
|
|
if tail_steps <= phase_counts[index]
|
|
]
|
|
if not active_indices:
|
|
continue
|
|
intervals = [
|
|
(
|
|
index * block_seconds,
|
|
index * block_seconds + tail_steps * step_seconds,
|
|
)
|
|
for index in active_indices
|
|
]
|
|
if len(active_indices) == 1 or av_batch_size == 1:
|
|
pieces_batch = [
|
|
extractor._prepare_pieces(
|
|
video_path=None,
|
|
audio_path=None,
|
|
interval=intervals[0],
|
|
use_audio_in_video=True,
|
|
audio_feature_device=audio_feature_device,
|
|
audio_right_padding_samples=audio_right_padding_samples,
|
|
)
|
|
]
|
|
else:
|
|
pieces_batch = _prepare_pieces_batch(
|
|
model=model,
|
|
processor=processor,
|
|
media=media,
|
|
config=config,
|
|
intervals=intervals,
|
|
)
|
|
if len(pieces_batch) != len(active_indices):
|
|
raise RuntimeError("Batched AV piece count changed unexpectedly")
|
|
for block_index, pieces in zip(active_indices, pieces_batch):
|
|
closure_embeddings = input_embeddings(pieces.closure_ids)
|
|
cached = pieces_to_cache(
|
|
pieces,
|
|
closure_embeddings=closure_embeddings,
|
|
floating_dtype=torch.bfloat16,
|
|
)
|
|
variants_by_block[block_index][str(tail_steps)] = cached
|
|
lengths_by_block[block_index][str(tail_steps)] = int(
|
|
cached["prefix_embeddings"].shape[1]
|
|
+ cached["payload_embeddings"].shape[1]
|
|
+ cached["closure_embeddings"].shape[1]
|
|
)
|
|
|
|
for block_index in block_indices:
|
|
block_start = block_index * block_seconds
|
|
path = output_dir / cache_block_filename(block_index)
|
|
torch.save(
|
|
{
|
|
"schema_version": SCHEMA_VERSION,
|
|
"block_index": block_index,
|
|
"start_seconds": block_start,
|
|
"end_seconds": min(
|
|
block_start + block_seconds, num_steps * step_seconds
|
|
),
|
|
"variants": variants_by_block[block_index],
|
|
},
|
|
path,
|
|
)
|
|
records.append(
|
|
{
|
|
"block_index": block_index,
|
|
"path": path.name,
|
|
"phase_count": phase_counts[block_index],
|
|
"token_lengths_without_query": lengths_by_block[block_index],
|
|
"size_bytes": path.stat().st_size,
|
|
}
|
|
)
|
|
progress.update(len(block_indices))
|
|
progress.close()
|
|
return records
|
|
|
|
|
|
@torch.inference_mode()
|
|
def _verify_cache(
|
|
*,
|
|
model: torch.nn.Module,
|
|
extractor: core.OpenStreamReplayExtractor,
|
|
cache_dir: Path,
|
|
static: dict[str, Any],
|
|
anchors: list[dict[str, Any]],
|
|
verify_every: int,
|
|
minimum_cosine: float,
|
|
block_seconds: float,
|
|
step_seconds: float,
|
|
) -> list[dict[str, float]]:
|
|
if verify_every <= 0:
|
|
return []
|
|
selected = [
|
|
anchor
|
|
for anchor in anchors
|
|
if (int(anchor["step_index"]) + 1) % verify_every == 0
|
|
]
|
|
if anchors and anchors[-1] not in selected:
|
|
selected.append(anchors[-1])
|
|
store = CachedBlockStore(cache_dir, max_blocks=max(2, int(static["memory_blocks"])))
|
|
device = _model_device(model)
|
|
checks: list[dict[str, float]] = []
|
|
for anchor in tqdm(selected, desc="verify cached replay", unit="state"):
|
|
embeddings, positions = reconstruct_open_stream_sequence(
|
|
store=store, static=static, anchor=anchor
|
|
)
|
|
embeddings = embeddings.unsqueeze(0).to(device=device, dtype=model.dtype)
|
|
positions = positions.unsqueeze(1).to(device=device)
|
|
attention_mask = torch.ones(
|
|
1, embeddings.shape[1], dtype=torch.long, device=device
|
|
)
|
|
cached_state = forward_last_state(
|
|
model.model,
|
|
inputs_embeds=embeddings,
|
|
position_ids=positions,
|
|
attention_mask=attention_mask,
|
|
readout_indices=torch.tensor(
|
|
[embeddings.shape[1] - 1], dtype=torch.long, device=device
|
|
),
|
|
)
|
|
exact_state = extractor.state_from_scratch(
|
|
video_path=None,
|
|
audio_path=None,
|
|
intervals=anchor_intervals(
|
|
anchor,
|
|
block_seconds=block_seconds,
|
|
step_seconds=step_seconds,
|
|
),
|
|
use_audio_in_video=True,
|
|
)
|
|
cached_float = cached_state.float()
|
|
exact_float = exact_state.float()
|
|
check = {
|
|
"time_seconds": float(anchor["time_seconds"]),
|
|
"cosine": float(
|
|
F.cosine_similarity(cached_float, exact_float, dim=-1).item()
|
|
),
|
|
"mean_abs": float((cached_float - exact_float).abs().mean().item()),
|
|
"max_abs": float((cached_float - exact_float).abs().max().item()),
|
|
}
|
|
checks.append(check)
|
|
print(
|
|
f"verify t={check['time_seconds']:.1f}s "
|
|
f"cosine={check['cosine']:.8f} "
|
|
f"mean_abs={check['mean_abs']:.6f} "
|
|
f"max_abs={check['max_abs']:.6f}",
|
|
flush=True,
|
|
)
|
|
if check["cosine"] < minimum_cosine:
|
|
raise RuntimeError(
|
|
"Cached full-sequence reconstruction disagrees with open-stream "
|
|
f"KV replay at t={check['time_seconds']:.1f}s: "
|
|
f"cosine={check['cosine']:.8f} < {minimum_cosine:.8f}"
|
|
)
|
|
return checks
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> None:
|
|
args = parse_args(argv)
|
|
core._configure_qwen_logging()
|
|
if args.end <= 0:
|
|
raise ValueError("--end must be positive")
|
|
if args.av_batch_size <= 0:
|
|
raise ValueError("--av-batch-size must be positive")
|
|
if (
|
|
args.audio_right_padding_samples is not None
|
|
and args.audio_right_padding_samples < 0
|
|
):
|
|
raise ValueError("--audio-right-padding-samples must be non-negative")
|
|
if args.av_batch_size != 1 and (
|
|
args.audio_feature_device is not None
|
|
or args.audio_right_padding_samples is not None
|
|
):
|
|
raise ValueError(
|
|
"Optimized audio preprocessing currently requires --av-batch-size 1; "
|
|
"the HF multimodal batch path was slower and changed cached states"
|
|
)
|
|
steps_per_block = _integer_ratio(
|
|
args.native_block_seconds, args.step_seconds, "native block / step"
|
|
)
|
|
memory_blocks = _integer_ratio(
|
|
args.memory_seconds, args.native_block_seconds, "memory / native block"
|
|
)
|
|
num_steps = _integer_ratio(args.end, args.step_seconds, "end / step")
|
|
targets = _load_targets(args.target)
|
|
if num_steps > len(targets):
|
|
raise ValueError(
|
|
f"Requested {num_steps} cached steps but target has only {len(targets)}"
|
|
)
|
|
|
|
output_dir = args.output_dir.expanduser().resolve()
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
manifest_path = output_dir / "manifest.json"
|
|
static_path = output_dir / "static.pt"
|
|
_safe_output_path(manifest_path, args.overwrite)
|
|
_safe_output_path(static_path, args.overwrite)
|
|
|
|
config = core.StreamConfig(
|
|
step_seconds=args.step_seconds,
|
|
native_block_seconds=args.native_block_seconds,
|
|
memory_seconds=args.memory_seconds,
|
|
video_fps=args.video_fps,
|
|
video_min_pixels=args.video_min_pixels,
|
|
video_max_pixels=args.video_max_pixels,
|
|
)
|
|
config.validate()
|
|
specs = _easycom_directory_specs(args.easycom_session_dir)
|
|
media = CausalSourceAV(
|
|
specs,
|
|
output_fps=args.video_fps,
|
|
required_duration=args.end,
|
|
frame_cache_size=args.frame_cache_size,
|
|
decoder_backend=args.decoder,
|
|
decoder_gpu_id=args.decoder_gpu_id,
|
|
decoder_threads=args.decoder_threads,
|
|
max_open_readers=args.max_open_readers,
|
|
)
|
|
print(
|
|
f"Direct media ready: duration={args.end:.1f}s, "
|
|
f"grid={media.expected_grid_frames} frames, audio={len(media.audio)} samples"
|
|
)
|
|
|
|
model, processor = core.load_qwen_omni(
|
|
args.model_id,
|
|
dtype=args.dtype,
|
|
device_map=args.device,
|
|
use_flash_attention_2=args.flash_attention_2,
|
|
)
|
|
model.requires_grad_(False)
|
|
audio_feature_device: str | torch.device | None = args.audio_feature_device
|
|
if audio_feature_device == "model":
|
|
audio_feature_device = _model_device(model)
|
|
audio_padding_description = "model maximum"
|
|
if args.audio_right_padding_samples is not None:
|
|
hop_length = int(getattr(processor.feature_extractor, "hop_length", 1))
|
|
rounded_padding = (
|
|
math.ceil(args.audio_right_padding_samples / hop_length) * hop_length
|
|
)
|
|
audio_padding_description = (
|
|
f"actual + >= {args.audio_right_padding_samples} samples "
|
|
f"(hop-aligned, typically +{rounded_padding})"
|
|
)
|
|
print(
|
|
"Audio preprocessing: "
|
|
f"feature_device={audio_feature_device or 'processor default'}, "
|
|
f"padding={audio_padding_description}"
|
|
)
|
|
extractor = core.OpenStreamReplayExtractor(
|
|
model, processor, config, fixed_grid_av=media
|
|
)
|
|
with torch.inference_mode():
|
|
query_ids = extractor.query_ids
|
|
query_embeddings = model.get_input_embeddings()(query_ids)
|
|
positions_per_chunk = core._tmrope_positions_per_chunk(model)
|
|
static = {
|
|
"schema_version": SCHEMA_VERSION,
|
|
"steps_per_block": steps_per_block,
|
|
"memory_blocks": memory_blocks,
|
|
"positions_per_chunk": positions_per_chunk,
|
|
"query_ids": query_ids.detach().cpu().long(),
|
|
"query_embeddings": query_embeddings.detach().cpu().to(torch.bfloat16),
|
|
}
|
|
torch.save(static, static_path)
|
|
|
|
cache_started = time.perf_counter()
|
|
block_records = _cache_blocks(
|
|
extractor=extractor,
|
|
model=model,
|
|
processor=processor,
|
|
media=media,
|
|
config=config,
|
|
output_dir=output_dir,
|
|
num_steps=num_steps,
|
|
steps_per_block=steps_per_block,
|
|
step_seconds=args.step_seconds,
|
|
block_seconds=args.native_block_seconds,
|
|
av_batch_size=args.av_batch_size,
|
|
audio_feature_device=audio_feature_device,
|
|
audio_right_padding_samples=args.audio_right_padding_samples,
|
|
overwrite=args.overwrite,
|
|
)
|
|
cache_seconds = time.perf_counter() - cache_started
|
|
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"
|
|
)
|
|
anchors = [
|
|
make_anchor(
|
|
step_index=step_index,
|
|
steps_per_block=steps_per_block,
|
|
memory_blocks=memory_blocks,
|
|
step_seconds=args.step_seconds,
|
|
num_target_steps=len(targets),
|
|
num_horizons=NUM_HORIZONS,
|
|
)
|
|
for step_index in range(num_steps)
|
|
]
|
|
checks = _verify_cache(
|
|
model=model,
|
|
extractor=extractor,
|
|
cache_dir=output_dir,
|
|
static=static,
|
|
anchors=anchors,
|
|
verify_every=args.verify_every,
|
|
minimum_cosine=args.minimum_cosine,
|
|
block_seconds=args.native_block_seconds,
|
|
step_seconds=args.step_seconds,
|
|
)
|
|
|
|
manifest = {
|
|
"schema_version": SCHEMA_VERSION,
|
|
"cache_kind": "qwen2.5-omni-frozen-av-open-stream-phases",
|
|
"session": args.easycom_session_dir.resolve().name,
|
|
"model_id": args.model_id,
|
|
"target_path": str(args.target.expanduser().resolve()),
|
|
"target_steps": int(len(targets)),
|
|
"cached_steps": num_steps,
|
|
"cached_end_seconds": args.end,
|
|
"hidden_size": int(query_embeddings.shape[-1]),
|
|
"floating_dtype": "bfloat16",
|
|
"step_seconds": args.step_seconds,
|
|
"native_block_seconds": args.native_block_seconds,
|
|
"memory_seconds": args.memory_seconds,
|
|
"steps_per_block": steps_per_block,
|
|
"memory_blocks": memory_blocks,
|
|
"positions_per_chunk": positions_per_chunk,
|
|
"video_fps": args.video_fps,
|
|
"video_min_pixels": args.video_min_pixels,
|
|
"video_max_pixels": args.video_max_pixels,
|
|
"use_audio_in_video": True,
|
|
"interval_convention": "half_open_[start,end)",
|
|
"phase_policy": "cache_0.2_to_2.0s;_2.0s_is_completed_history",
|
|
"av_batch_size": args.av_batch_size,
|
|
"audio_feature_device": (
|
|
None if audio_feature_device is None else str(audio_feature_device)
|
|
),
|
|
"audio_right_padding_samples": args.audio_right_padding_samples,
|
|
"verification_reference_audio_preprocessing": (
|
|
"processor-default-device-and-model-maximum-padding"
|
|
),
|
|
"cache_generation_seconds": cache_seconds,
|
|
"cache_media_realtime_factor": args.end / cache_seconds,
|
|
"cache_blocks_per_second": len(block_records) / cache_seconds,
|
|
"label_policy": "target[i+1:i+11]",
|
|
"static_path": static_path.name,
|
|
"blocks": block_records,
|
|
"anchors": anchors,
|
|
"verification": checks,
|
|
"media": media.metadata(),
|
|
"software": {
|
|
"python": platform.python_version(),
|
|
"torch": torch.__version__,
|
|
},
|
|
}
|
|
save_json(manifest_path, manifest)
|
|
size = cache_size_bytes(
|
|
[static_path, *[output_dir / record["path"] for record in block_records]]
|
|
)
|
|
print(
|
|
f"Saved {len(block_records)} blocks / {len(anchors)} anchors to {output_dir}"
|
|
)
|
|
print(f"Cache tensor size: {size / (1024 ** 3):.3f} GiB")
|
|
if checks:
|
|
print(
|
|
"Verification: "
|
|
f"lowest cosine={min(item['cosine'] for item in checks):.8f}, "
|
|
f"worst max_abs={max(item['max_abs'] for item in checks):.6f}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|