Add no-box full-frame Qwen AV baselines
This commit is contained in:
Xin Wang
402
extract_fullframe_qwen_av_features.py
Normal file
402
extract_fullframe_qwen_av_features.py
Normal file
@@ -0,0 +1,402 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user