403 lines
16 KiB
Python
403 lines
16 KiB
Python
|
"""Extract causal full-frame Qwen vision tokens plus matched Qwen audio.
|
||
|
|
|
||
|
|
This EasyCom-only extractor never reads face/head boxes or speaker labels. It
|
||
|
|
reuses the audited recent mixed-audio vectors from a local-Qwen sidecar and
|
||
|
|
encodes complete camera frames on a fixed causal grid. The output supports F2
|
||
|
|
latest-frame spatial attention and F3 same-grid-position temporal modeling.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import os
|
||
|
|
import tempfile
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
from tqdm.auto import tqdm
|
||
|
|
|
||
|
|
from easycom_head import normalize_sample_name, resolve_feature_path
|
||
|
|
from extract_turn_local_features import (
|
||
|
|
model_config_digest,
|
||
|
|
validate_local_transformers_version,
|
||
|
|
validate_source_timeline,
|
||
|
|
)
|
||
|
|
from fullframe_qwen_av import FULLFRAME_AV_FORMAT
|
||
|
|
from local_turn_features import causal_frame_ticks, load_local_features, reference_alignment
|
||
|
|
|
||
|
|
|
||
|
|
class FrozenQwenFullFrameEncoder:
|
||
|
|
"""Keep only the frozen Qwen vision tower and merger on the GPU."""
|
||
|
|
|
||
|
|
def __init__(self, args: argparse.Namespace) -> None:
|
||
|
|
import transformers
|
||
|
|
from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniThinkerForConditionalGeneration
|
||
|
|
|
||
|
|
self.transformers_version = validate_local_transformers_version(transformers.__version__)
|
||
|
|
self.device = torch.device(args.device)
|
||
|
|
if self.device.type == "cuda" and not torch.cuda.is_available():
|
||
|
|
raise RuntimeError(f"CUDA requested but unavailable: {args.device}")
|
||
|
|
dtype_key = "dtype" if self.transformers_version == "4.57.6" else "torch_dtype"
|
||
|
|
kwargs = {
|
||
|
|
dtype_key: getattr(torch, args.dtype),
|
||
|
|
"device_map": {"": "cpu"},
|
||
|
|
"output_loading_info": True,
|
||
|
|
}
|
||
|
|
if args.flash_attention_2:
|
||
|
|
kwargs["attn_implementation"] = "flash_attention_2"
|
||
|
|
model, info = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
|
||
|
|
args.model_id, **kwargs
|
||
|
|
)
|
||
|
|
bad = [key for key in info.get("missing_keys", []) if key.startswith("visual.")]
|
||
|
|
if bad or info.get("mismatched_keys") or info.get("error_msgs"):
|
||
|
|
raise RuntimeError(f"Incomplete Qwen vision weights: {info}")
|
||
|
|
model.requires_grad_(False).eval()
|
||
|
|
self.visual_dim = int(model.config.vision_config.out_hidden_size)
|
||
|
|
self.model_config_sha256 = model_config_digest(model.config.to_dict())
|
||
|
|
model.model = None
|
||
|
|
model.lm_head = None
|
||
|
|
model.audio_tower = None
|
||
|
|
model.visual.to(self.device)
|
||
|
|
self.model = model
|
||
|
|
self.processor = Qwen2_5OmniProcessor.from_pretrained(
|
||
|
|
args.model_id, use_fast=False
|
||
|
|
)
|
||
|
|
self.max_pixels = args.max_pixels
|
||
|
|
self.image_batch_size = args.image_batch_size
|
||
|
|
|
||
|
|
@torch.inference_mode()
|
||
|
|
def image_tokens(
|
||
|
|
self, images: list[Any]
|
||
|
|
) -> tuple[list[torch.Tensor], list[tuple[int, int]]]:
|
||
|
|
pieces: list[torch.Tensor] = []
|
||
|
|
grids: list[tuple[int, int]] = []
|
||
|
|
for start in range(0, len(images), self.image_batch_size):
|
||
|
|
batch = self.processor.image_processor(
|
||
|
|
images=images[start : start + self.image_batch_size],
|
||
|
|
min_pixels=4 * 28 * 28,
|
||
|
|
max_pixels=self.max_pixels,
|
||
|
|
return_tensors="pt",
|
||
|
|
)
|
||
|
|
grid = batch["image_grid_thw"].to(self.device)
|
||
|
|
pixels = batch["pixel_values"].to(
|
||
|
|
self.device, dtype=self.model.visual.dtype
|
||
|
|
)
|
||
|
|
tokens = self.model.get_image_features(pixels, grid)
|
||
|
|
merge = int(self.model.config.vision_config.spatial_merge_size)
|
||
|
|
counts = (grid.prod(-1) // (merge * merge)).tolist()
|
||
|
|
if not isinstance(tokens, torch.Tensor) or tokens.shape != (
|
||
|
|
sum(counts), self.visual_dim
|
||
|
|
):
|
||
|
|
raise RuntimeError(f"Unexpected Qwen image output: {getattr(tokens, 'shape', None)}")
|
||
|
|
pieces.extend(part.float().cpu() for part in tokens.split(counts))
|
||
|
|
for row in grid.tolist():
|
||
|
|
temporal, height, width = (int(value) for value in row)
|
||
|
|
if temporal != 1 or height % merge or width % merge:
|
||
|
|
raise RuntimeError(f"Unexpected image grid before merger: {row}")
|
||
|
|
grids.append((height // merge, width // merge))
|
||
|
|
if len(pieces) != len(images) or len(grids) != len(images):
|
||
|
|
raise AssertionError("Qwen did not return one token set per image")
|
||
|
|
return pieces, grids
|
||
|
|
|
||
|
|
|
||
|
|
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument("--data-root", type=Path, required=True)
|
||
|
|
parser.add_argument("--feature-root", type=Path, required=True)
|
||
|
|
parser.add_argument("--target-root", type=Path, required=True)
|
||
|
|
parser.add_argument("--audio-feature-root", type=Path, required=True)
|
||
|
|
parser.add_argument("--output-root", type=Path, required=True)
|
||
|
|
parser.add_argument("--samples", nargs="+", required=True)
|
||
|
|
parser.add_argument("--model-id", default="Qwen/Qwen2.5-Omni-3B")
|
||
|
|
parser.add_argument("--encoder-id")
|
||
|
|
parser.add_argument("--device", default="cuda:0")
|
||
|
|
parser.add_argument(
|
||
|
|
"--dtype", choices=("float32", "float16", "bfloat16"), default="bfloat16"
|
||
|
|
)
|
||
|
|
parser.add_argument("--flash-attention-2", action="store_true")
|
||
|
|
parser.add_argument("--video-fps", type=float, default=2.0)
|
||
|
|
parser.add_argument("--max-pixels", type=int, default=224 * 224)
|
||
|
|
parser.add_argument("--image-batch-size", type=int, default=8)
|
||
|
|
parser.add_argument(
|
||
|
|
"--decoder-backend",
|
||
|
|
choices=("decord", "torchcodec", "pynvcodec"),
|
||
|
|
default="decord",
|
||
|
|
)
|
||
|
|
parser.add_argument("--decoder-threads", type=int, default=1)
|
||
|
|
parser.add_argument("--decoder-gpu-id", type=int, default=0)
|
||
|
|
parser.add_argument("--max-steps", type=int)
|
||
|
|
parser.add_argument("--dry-run", action="store_true")
|
||
|
|
parser.add_argument("--overwrite", action="store_true")
|
||
|
|
parser.add_argument("--no-progress", action="store_true")
|
||
|
|
args = parser.parse_args(argv)
|
||
|
|
if not np.isfinite(args.video_fps) or args.video_fps <= 0:
|
||
|
|
parser.error("--video-fps must be finite and positive")
|
||
|
|
if args.max_pixels < 4 * 28 * 28 or args.image_batch_size <= 0:
|
||
|
|
parser.error("Require max-pixels >= 3136 and positive image-batch-size")
|
||
|
|
if args.max_steps is not None and args.max_steps <= 0:
|
||
|
|
parser.error("--max-steps must be positive")
|
||
|
|
protected = {
|
||
|
|
args.feature_root.resolve(),
|
||
|
|
args.audio_feature_root.resolve(),
|
||
|
|
args.target_root.resolve(),
|
||
|
|
}
|
||
|
|
if args.output_root.resolve() in protected:
|
||
|
|
parser.error("--output-root must not overwrite an input root")
|
||
|
|
return args
|
||
|
|
|
||
|
|
|
||
|
|
def _build_media(args: argparse.Namespace, sample: str, end: float):
|
||
|
|
from streaming_turn_model_directly import CausalSourceAV, _easycom_directory_specs
|
||
|
|
|
||
|
|
specs = _easycom_directory_specs(
|
||
|
|
args.data_root / "Main" / "Video_Compressed" / sample
|
||
|
|
)
|
||
|
|
return CausalSourceAV(
|
||
|
|
specs,
|
||
|
|
output_fps=args.video_fps,
|
||
|
|
required_duration=end,
|
||
|
|
decoder_backend=args.decoder_backend,
|
||
|
|
decoder_threads=args.decoder_threads,
|
||
|
|
decoder_gpu_id=args.decoder_gpu_id,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _requested_config(
|
||
|
|
args: argparse.Namespace,
|
||
|
|
encoder: FrozenQwenFullFrameEncoder,
|
||
|
|
audio_contract: dict[str, Any],
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
source = audio_contract["config"]
|
||
|
|
if source.get("causal_policy") != "half-open-trailing-window":
|
||
|
|
raise ValueError("Source Qwen audio is not a half-open causal window")
|
||
|
|
if source.get("uses_speech_labels") is not False:
|
||
|
|
raise ValueError("Source Qwen audio must not use speech labels")
|
||
|
|
audio_dim = source.get("audio_dim")
|
||
|
|
if not isinstance(audio_dim, int) or audio_dim <= 0:
|
||
|
|
raise ValueError("Source Qwen audio contract lacks audio_dim")
|
||
|
|
encoder_id = args.encoder_id or args.model_id
|
||
|
|
if source.get("encoder_id") != encoder_id:
|
||
|
|
raise ValueError(
|
||
|
|
f"Vision/audio encoder identities differ: {encoder_id!r} vs "
|
||
|
|
f"{source.get('encoder_id')!r}"
|
||
|
|
)
|
||
|
|
if source.get("model_config_sha256") != encoder.model_config_sha256:
|
||
|
|
raise ValueError("Vision/audio Qwen model configs differ")
|
||
|
|
if source.get("transformers") != encoder.transformers_version:
|
||
|
|
raise ValueError("Vision/audio Transformers versions differ")
|
||
|
|
return {
|
||
|
|
"encoder_id": encoder_id,
|
||
|
|
"model_config_sha256": encoder.model_config_sha256,
|
||
|
|
"transformers": encoder.transformers_version,
|
||
|
|
"dtype": args.dtype,
|
||
|
|
"attention": "flash_attention_2" if args.flash_attention_2 else "default",
|
||
|
|
"visual_dim": encoder.visual_dim,
|
||
|
|
"audio_dim": audio_dim,
|
||
|
|
"visual_fps": float(args.video_fps),
|
||
|
|
"max_pixels": int(args.max_pixels),
|
||
|
|
"visual_feature_kind": "qwen-vision-post-merger-spatial-token",
|
||
|
|
"visual_source": "easycom-full-frame-after-session-banner-crop",
|
||
|
|
"visual_pooling": "none-ragged-spatial-tokens",
|
||
|
|
"token_grid_identity": "merged-grid-row-column-within-source-clip",
|
||
|
|
"audio_source": "reused-qwen-local-mixed-audio",
|
||
|
|
"audio_window_seconds": source.get("audio_window_seconds"),
|
||
|
|
"source_audio_contract": audio_contract,
|
||
|
|
"uses_boxes": False,
|
||
|
|
"uses_speech_labels": False,
|
||
|
|
"causal_policy": "strict-half-open-fullframe-grid",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _extract_visual_tokens(
|
||
|
|
media,
|
||
|
|
encoder: FrozenQwenFullFrameEncoder,
|
||
|
|
*,
|
||
|
|
end: float,
|
||
|
|
fps: float,
|
||
|
|
batch_size: int,
|
||
|
|
no_progress: bool,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
ticks = list(causal_frame_ticks(end, fps))
|
||
|
|
frame_times = torch.tensor(ticks, dtype=torch.float64) / fps
|
||
|
|
feature_parts: list[torch.Tensor] = []
|
||
|
|
grid_shapes: list[tuple[int, int]] = []
|
||
|
|
source_clip_indices: list[int] = []
|
||
|
|
offsets = [0]
|
||
|
|
progress = tqdm(
|
||
|
|
range(0, len(ticks), batch_size),
|
||
|
|
desc="full-frame Qwen vision",
|
||
|
|
unit="batch",
|
||
|
|
disable=no_progress,
|
||
|
|
)
|
||
|
|
for start in progress:
|
||
|
|
batch_ticks = ticks[start : start + batch_size]
|
||
|
|
images, sources = [], []
|
||
|
|
for tick in batch_ticks:
|
||
|
|
time = tick / fps
|
||
|
|
source = media._source_frame_for_time(time)
|
||
|
|
images.append(media.frame_at_or_before(time))
|
||
|
|
sources.append(source)
|
||
|
|
token_sets, grids = encoder.image_tokens(images)
|
||
|
|
for tick, tokens, grid, source in zip(batch_ticks, token_sets, grids, sources):
|
||
|
|
if tokens.shape != (grid[0] * grid[1], encoder.visual_dim):
|
||
|
|
raise RuntimeError(
|
||
|
|
f"Qwen token/grid mismatch at tick {tick}: "
|
||
|
|
f"{tuple(tokens.shape)} vs {grid}"
|
||
|
|
)
|
||
|
|
if not torch.isfinite(tokens).all():
|
||
|
|
raise RuntimeError(f"Non-finite full-frame tokens at tick {tick}")
|
||
|
|
stored = tokens.to(torch.bfloat16)
|
||
|
|
feature_parts.append(stored)
|
||
|
|
grid_shapes.append(grid)
|
||
|
|
source_clip_indices.append(int(source[0]))
|
||
|
|
offsets.append(offsets[-1] + len(stored))
|
||
|
|
features = (
|
||
|
|
torch.cat(feature_parts, dim=0)
|
||
|
|
if feature_parts
|
||
|
|
else torch.empty(0, encoder.visual_dim, dtype=torch.bfloat16)
|
||
|
|
)
|
||
|
|
return {
|
||
|
|
"features": features,
|
||
|
|
"frame_times": frame_times,
|
||
|
|
"frame_offsets": torch.tensor(offsets, dtype=torch.long),
|
||
|
|
"grid_shapes": torch.tensor(grid_shapes, dtype=torch.long),
|
||
|
|
"frame_source_clip_indices": torch.tensor(
|
||
|
|
source_clip_indices, dtype=torch.long
|
||
|
|
),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def main(argv: list[str] | None = None) -> None:
|
||
|
|
args = parse_args(argv)
|
||
|
|
samples = [normalize_sample_name(value, "easycom") for value in args.samples]
|
||
|
|
if len(samples) != len(set(samples)):
|
||
|
|
raise ValueError("Duplicate sessions")
|
||
|
|
references, audio_paths = {}, {}
|
||
|
|
for sample in samples:
|
||
|
|
references[sample] = resolve_feature_path(args.feature_root, sample, dataset="easycom")
|
||
|
|
audio_paths[sample] = resolve_feature_path(
|
||
|
|
args.audio_feature_root, sample, dataset="easycom"
|
||
|
|
)
|
||
|
|
for path in (
|
||
|
|
references[sample],
|
||
|
|
audio_paths[sample],
|
||
|
|
args.target_root / f"{sample}.npy",
|
||
|
|
):
|
||
|
|
if not path.is_file():
|
||
|
|
raise FileNotFoundError(path)
|
||
|
|
destination = args.output_root / f"{sample}.pt"
|
||
|
|
if destination.exists() and not args.overwrite:
|
||
|
|
raise FileExistsError(
|
||
|
|
f"{destination} exists; select a new root or pass --overwrite"
|
||
|
|
)
|
||
|
|
|
||
|
|
audio_contract: dict[str, Any] | None = None
|
||
|
|
session_steps = {}
|
||
|
|
for sample in samples:
|
||
|
|
reference = torch.load(references[sample], map_location="cpu", weights_only=False)
|
||
|
|
times, _ = reference_alignment(reference)
|
||
|
|
target = np.load(
|
||
|
|
args.target_root / f"{sample}.npy", allow_pickle=False, mmap_mode="r"
|
||
|
|
)
|
||
|
|
if target.shape != (len(times), 3):
|
||
|
|
raise ValueError(f"Reference/target mismatch for {sample}")
|
||
|
|
packed_audio, audio_contract = load_local_features(
|
||
|
|
audio_paths[sample],
|
||
|
|
reference,
|
||
|
|
sample_id=sample,
|
||
|
|
dataset="easycom",
|
||
|
|
modalities="audio",
|
||
|
|
expected_contract=audio_contract,
|
||
|
|
)
|
||
|
|
session_steps[sample] = len(times)
|
||
|
|
steps = min(len(times), args.max_steps or len(times))
|
||
|
|
print(
|
||
|
|
f"{sample}: {steps} states, {steps / 5:g}s; "
|
||
|
|
f"matched audio={packed_audio.shape[1] - 1}D",
|
||
|
|
flush=True,
|
||
|
|
)
|
||
|
|
del reference, target, packed_audio
|
||
|
|
if args.dry_run:
|
||
|
|
print("Dry run only: source video decode and Qwen vision were not checked.")
|
||
|
|
return
|
||
|
|
|
||
|
|
args.output_root.mkdir(parents=True, exist_ok=True)
|
||
|
|
encoder = FrozenQwenFullFrameEncoder(args)
|
||
|
|
assert audio_contract is not None
|
||
|
|
config = _requested_config(args, encoder, audio_contract)
|
||
|
|
for position, sample in enumerate(samples, 1):
|
||
|
|
print(f"[{position}/{len(samples)}] {sample}", flush=True)
|
||
|
|
reference = torch.load(references[sample], map_location="cpu", weights_only=False)
|
||
|
|
full_times, full_indices = reference_alignment(reference)
|
||
|
|
packed_audio, current_contract = load_local_features(
|
||
|
|
audio_paths[sample],
|
||
|
|
reference,
|
||
|
|
sample_id=sample,
|
||
|
|
dataset="easycom",
|
||
|
|
modalities="audio",
|
||
|
|
expected_contract=audio_contract,
|
||
|
|
)
|
||
|
|
if current_contract != audio_contract:
|
||
|
|
raise AssertionError("Audio contract changed after preflight")
|
||
|
|
full_steps = session_steps[sample]
|
||
|
|
steps = min(full_steps, args.max_steps or full_steps)
|
||
|
|
times, indices, packed_audio = (
|
||
|
|
full_times[:steps],
|
||
|
|
full_indices[:steps],
|
||
|
|
packed_audio[:steps],
|
||
|
|
)
|
||
|
|
end = steps / 5
|
||
|
|
media = _build_media(args, sample, end)
|
||
|
|
validate_source_timeline(reference, media)
|
||
|
|
visual_tokens = _extract_visual_tokens(
|
||
|
|
media,
|
||
|
|
encoder,
|
||
|
|
end=end,
|
||
|
|
fps=args.video_fps,
|
||
|
|
batch_size=args.image_batch_size,
|
||
|
|
no_progress=args.no_progress,
|
||
|
|
)
|
||
|
|
audio = packed_audio[:, :-1].to(torch.bfloat16)
|
||
|
|
audio_valid = packed_audio[:, -1].bool()
|
||
|
|
if torch.count_nonzero(audio[~audio_valid]):
|
||
|
|
raise ValueError("Invalid inherited audio rows are not zero")
|
||
|
|
output = {
|
||
|
|
"format": FULLFRAME_AV_FORMAT,
|
||
|
|
"dataset": "easycom",
|
||
|
|
"sample_id": sample,
|
||
|
|
"times": times,
|
||
|
|
"label_indices": indices,
|
||
|
|
"partial": steps != full_steps,
|
||
|
|
"audio": audio,
|
||
|
|
"audio_valid": audio_valid,
|
||
|
|
"visual_tokens": visual_tokens,
|
||
|
|
"config": config,
|
||
|
|
"reference_features": str(references[sample].resolve()),
|
||
|
|
"source_audio_features": str(audio_paths[sample].resolve()),
|
||
|
|
"direct_media": media.metadata(),
|
||
|
|
}
|
||
|
|
destination = args.output_root / f"{sample}.pt"
|
||
|
|
with tempfile.NamedTemporaryFile(
|
||
|
|
dir=args.output_root,
|
||
|
|
prefix=f".{sample}.",
|
||
|
|
suffix=".tmp",
|
||
|
|
delete=False,
|
||
|
|
) as handle:
|
||
|
|
temporary = Path(handle.name)
|
||
|
|
try:
|
||
|
|
torch.save(output, temporary)
|
||
|
|
os.replace(temporary, destination)
|
||
|
|
finally:
|
||
|
|
temporary.unlink(missing_ok=True)
|
||
|
|
print(
|
||
|
|
f"Saved {destination}: frames={len(visual_tokens['frame_times'])}, "
|
||
|
|
f"tokens={len(visual_tokens['features'])}, partial={output['partial']}",
|
||
|
|
flush=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|