193 lines
8.5 KiB
Python
193 lines
8.5 KiB
Python
"""Strict no-box full-frame Qwen AV sidecar contract and loader."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
from local_turn_features import FramePersonFeatures, reference_alignment
|
|
|
|
|
|
FULLFRAME_AV_FORMAT = "qwen-fullframe-av-tokens-v1"
|
|
|
|
|
|
def load_fullframe_av_features(
|
|
path: Path,
|
|
reference: dict[str, Any],
|
|
*,
|
|
sample_id: str,
|
|
dataset: str,
|
|
expected_contract: dict[str, Any] | None = None,
|
|
history_max_gap_seconds: float | None = None,
|
|
) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]:
|
|
"""Load aligned Qwen audio and ragged full-frame spatial token sets."""
|
|
|
|
times, indices = reference_alignment(reference)
|
|
payload = torch.load(path, map_location="cpu", weights_only=False)
|
|
if not isinstance(payload, dict) or payload.get("format") != FULLFRAME_AV_FORMAT:
|
|
raise ValueError(f"Invalid full-frame feature format: {path}")
|
|
if payload.get("sample_id") != sample_id or payload.get("dataset") != dataset:
|
|
raise ValueError(f"Full-frame sample/dataset mismatch: {path}")
|
|
if dataset != "easycom":
|
|
raise ValueError("Full-frame F2/F3 currently support EasyCom only")
|
|
if payload.get("partial", False):
|
|
raise ValueError(f"Partial smoke-test features cannot be used for training: {path}")
|
|
local_indices = torch.as_tensor(payload.get("label_indices", []))
|
|
local_times = torch.as_tensor(payload.get("times", []), dtype=torch.float64)
|
|
if local_indices.dtype != torch.long or not torch.equal(indices, local_indices):
|
|
raise ValueError(f"Full-frame label_indices mismatch: {path}")
|
|
if local_times.shape != times.shape or not torch.allclose(
|
|
local_times, times, atol=1e-7, rtol=0
|
|
):
|
|
raise ValueError(f"Full-frame times mismatch: {path}")
|
|
config = payload.get("config", {})
|
|
if (
|
|
config.get("causal_policy") != "strict-half-open-fullframe-grid"
|
|
or config.get("uses_speech_labels") is not False
|
|
or config.get("uses_boxes") is not False
|
|
or config.get("visual_feature_kind")
|
|
!= "qwen-vision-post-merger-spatial-token"
|
|
or config.get("visual_source")
|
|
!= "easycom-full-frame-after-session-banner-crop"
|
|
or config.get("audio_source") != "reused-qwen-local-mixed-audio"
|
|
):
|
|
raise ValueError(f"Full-frame AV features lack the causal/no-box contract: {path}")
|
|
contract = {"format": FULLFRAME_AV_FORMAT, "modalities": "av", "config": config}
|
|
if expected_contract is not None and contract != expected_contract:
|
|
raise ValueError(f"Full-frame contract differs across sessions/checkpoint: {path}")
|
|
|
|
audio = payload.get("audio")
|
|
audio_valid = payload.get("audio_valid")
|
|
audio_dim = config.get("audio_dim")
|
|
if not isinstance(audio_dim, int) or audio_dim <= 0:
|
|
raise ValueError(f"Missing audio_dim: {path}")
|
|
if (
|
|
not isinstance(audio, torch.Tensor)
|
|
or audio.shape != (len(times), audio_dim)
|
|
or not audio.is_floating_point()
|
|
or not torch.isfinite(audio).all()
|
|
):
|
|
raise ValueError(f"Invalid full-frame sidecar audio: {path}")
|
|
if (
|
|
not isinstance(audio_valid, torch.Tensor)
|
|
or audio_valid.dtype != torch.bool
|
|
or audio_valid.shape != (len(times),)
|
|
):
|
|
raise ValueError(f"Invalid full-frame audio_valid mask: {path}")
|
|
if torch.count_nonzero(audio[~audio_valid]):
|
|
raise ValueError(f"Invalid audio rows must be stored as zeros: {path}")
|
|
|
|
visual_dim = config.get("visual_dim")
|
|
visual_fps = config.get("visual_fps")
|
|
tokens = payload.get("visual_tokens")
|
|
if not isinstance(visual_dim, int) or visual_dim <= 0:
|
|
raise ValueError(f"Missing visual_dim: {path}")
|
|
if not isinstance(visual_fps, (int, float)) or not math.isfinite(visual_fps) or visual_fps <= 0:
|
|
raise ValueError(f"Invalid visual_fps: {path}")
|
|
if not isinstance(tokens, dict):
|
|
raise ValueError(f"Missing full-frame visual tokens: {path}")
|
|
features = tokens.get("features")
|
|
frame_times = tokens.get("frame_times")
|
|
frame_offsets = tokens.get("frame_offsets")
|
|
grid_shapes = tokens.get("grid_shapes")
|
|
source_clip_indices = tokens.get("frame_source_clip_indices")
|
|
if (
|
|
not isinstance(features, torch.Tensor)
|
|
or features.ndim != 2
|
|
or features.shape[1] != visual_dim
|
|
or not features.is_floating_point()
|
|
or not torch.isfinite(features).all()
|
|
):
|
|
raise ValueError(f"Invalid full-frame token features: {path}")
|
|
if (
|
|
not isinstance(frame_times, torch.Tensor)
|
|
or frame_times.ndim != 1
|
|
or not frame_times.is_floating_point()
|
|
):
|
|
raise ValueError(f"Invalid full-frame frame_times: {path}")
|
|
frame_times = frame_times.to(dtype=torch.float64)
|
|
if (
|
|
not len(frame_times)
|
|
or not torch.isfinite(frame_times).all()
|
|
or bool((frame_times < 0).any())
|
|
or bool((frame_times[1:] <= frame_times[:-1]).any())
|
|
):
|
|
raise ValueError(f"Full-frame times must be finite/increasing: {path}")
|
|
expected_frame_count = max(
|
|
0,
|
|
math.ceil(float(times[-1]) * float(visual_fps) - 1e-8),
|
|
)
|
|
expected_frame_times = (
|
|
torch.arange(expected_frame_count, dtype=torch.float64)
|
|
/ float(visual_fps)
|
|
)
|
|
if frame_times.shape != expected_frame_times.shape or not torch.allclose(
|
|
frame_times, expected_frame_times, atol=1e-8, rtol=0
|
|
):
|
|
raise ValueError(f"Full-frame times are not the fixed causal grid: {path}")
|
|
if (
|
|
not isinstance(frame_offsets, torch.Tensor)
|
|
or frame_offsets.dtype != torch.long
|
|
or frame_offsets.shape != (len(frame_times) + 1,)
|
|
or int(frame_offsets[0]) != 0
|
|
or int(frame_offsets[-1]) != len(features)
|
|
or bool((frame_offsets[1:] <= frame_offsets[:-1]).any())
|
|
):
|
|
raise ValueError(f"Invalid full-frame frame_offsets: {path}")
|
|
if (
|
|
not isinstance(grid_shapes, torch.Tensor)
|
|
or grid_shapes.dtype != torch.long
|
|
or grid_shapes.shape != (len(frame_times), 2)
|
|
or bool((grid_shapes <= 0).any())
|
|
):
|
|
raise ValueError(f"Invalid full-frame grid_shapes: {path}")
|
|
if not torch.equal(frame_offsets[1:] - frame_offsets[:-1], grid_shapes.prod(-1)):
|
|
raise ValueError(f"Token counts do not match full-frame grids: {path}")
|
|
if (
|
|
not isinstance(source_clip_indices, torch.Tensor)
|
|
or source_clip_indices.dtype != torch.long
|
|
or source_clip_indices.shape != (len(frame_times),)
|
|
or bool((source_clip_indices < 0).any())
|
|
or bool((source_clip_indices[1:] < source_clip_indices[:-1]).any())
|
|
):
|
|
raise ValueError(f"Invalid full-frame source-clip indices: {path}")
|
|
if float(frame_times[-1]) >= float(times[-1]) - 1e-8:
|
|
raise ValueError(f"Full-frame timeline extends past the reference cutoff: {path}")
|
|
|
|
row_to_frame = torch.searchsorted(frame_times, times, right=False) - 1
|
|
chosen = torch.full(times.shape, float("nan"), dtype=torch.float64)
|
|
valid_rows = row_to_frame >= 0
|
|
chosen[valid_rows] = frame_times[row_to_frame[valid_rows]]
|
|
if bool((chosen[valid_rows] >= times[valid_rows]).any()):
|
|
raise ValueError(f"Full-frame mapping selected boundary/future evidence: {path}")
|
|
|
|
previous = None
|
|
if history_max_gap_seconds is not None:
|
|
if not math.isfinite(history_max_gap_seconds) or history_max_gap_seconds <= 0:
|
|
raise ValueError("history_max_gap_seconds must be finite and positive")
|
|
previous = torch.full((len(features),), -1, dtype=torch.long)
|
|
for frame in range(1, len(frame_times)):
|
|
gap = float(frame_times[frame] - frame_times[frame - 1])
|
|
same_clip = bool(source_clip_indices[frame] == source_clip_indices[frame - 1])
|
|
same_grid = torch.equal(grid_shapes[frame], grid_shapes[frame - 1])
|
|
if not same_clip or not same_grid or gap > history_max_gap_seconds + 1e-8:
|
|
continue
|
|
current_start, current_end = int(frame_offsets[frame]), int(frame_offsets[frame + 1])
|
|
old_start, old_end = int(frame_offsets[frame - 1]), int(frame_offsets[frame])
|
|
if current_end - current_start != old_end - old_start:
|
|
raise AssertionError("Equal grids have unequal token counts")
|
|
previous[current_start:current_end] = torch.arange(old_start, old_end)
|
|
|
|
frames = FramePersonFeatures(
|
|
features,
|
|
frame_offsets,
|
|
row_to_frame.to(torch.long),
|
|
chosen,
|
|
previous,
|
|
)
|
|
packed_audio = torch.cat((audio.float(), audio_valid[:, None].float()), dim=-1)
|
|
return packed_audio, frames, contract
|