"""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()