271 lines
8.8 KiB
Python
271 lines
8.8 KiB
Python
|
"""Load Qwen2.5-Omni-3B and extract causal audio-video features.
|
||
|
|
|
||
|
|
This research example does not generate text or speech: only the Thinker
|
||
|
|
(multimodal understanding model) is loaded.
|
||
|
|
|
||
|
|
Example:
|
||
|
|
python qwen_omni_features.py --video clip.mp4 \
|
||
|
|
--use-audio-in-video --output clip_features.pt
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any, Literal, Mapping
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from qwen_omni_utils import process_mm_info
|
||
|
|
from transformers import (
|
||
|
|
Qwen2_5OmniProcessor,
|
||
|
|
Qwen2_5OmniThinkerForConditionalGeneration,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
MODEL_ID = "Qwen/Qwen2.5-Omni-3B"
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class OmniFeatures:
|
||
|
|
"""Features returned for one or more audio-video examples."""
|
||
|
|
|
||
|
|
token_features: torch.Tensor # [B, T, D]
|
||
|
|
attention_mask: torch.Tensor # [B, T]
|
||
|
|
last_token: torch.Tensor # [B, D]
|
||
|
|
tail_mean: torch.Tensor # [B, D]
|
||
|
|
hidden_states: tuple[torch.Tensor, ...] | None = None
|
||
|
|
|
||
|
|
|
||
|
|
def choose_dtype(dtype: Literal["auto", "bfloat16", "float16", "float32"]):
|
||
|
|
if dtype == "auto":
|
||
|
|
return "auto"
|
||
|
|
return getattr(torch, dtype)
|
||
|
|
|
||
|
|
|
||
|
|
def load_qwen_omni(
|
||
|
|
model_id: str = MODEL_ID,
|
||
|
|
*,
|
||
|
|
dtype: Literal["auto", "bfloat16", "float16", "float32"] = "auto",
|
||
|
|
device_map: str | Mapping[str, str | int] = "auto",
|
||
|
|
use_flash_attention_2: bool = False,
|
||
|
|
) -> tuple[Qwen2_5OmniThinkerForConditionalGeneration, Qwen2_5OmniProcessor]:
|
||
|
|
"""Load the multimodal Thinker and processor without the Talker."""
|
||
|
|
|
||
|
|
model_kwargs: dict[str, Any] = {
|
||
|
|
"torch_dtype": choose_dtype(dtype),
|
||
|
|
"device_map": device_map,
|
||
|
|
}
|
||
|
|
if use_flash_attention_2:
|
||
|
|
model_kwargs["attn_implementation"] = "flash_attention_2"
|
||
|
|
|
||
|
|
model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
|
||
|
|
model_id, **model_kwargs
|
||
|
|
)
|
||
|
|
processor = Qwen2_5OmniProcessor.from_pretrained(model_id)
|
||
|
|
model.eval()
|
||
|
|
return model, processor
|
||
|
|
|
||
|
|
|
||
|
|
def build_conversation(
|
||
|
|
*,
|
||
|
|
video_path: str | Path | None = None,
|
||
|
|
audio_path: str | Path | None = None,
|
||
|
|
prompt: str = (
|
||
|
|
"Observe the conversation up to the current time and represent the "
|
||
|
|
"multimodal evidence useful for predicting the next speaking turn."
|
||
|
|
),
|
||
|
|
) -> list[dict[str, Any]]:
|
||
|
|
"""Build one Qwen chat containing video, optional audio, and instruction."""
|
||
|
|
|
||
|
|
if video_path is None and audio_path is None:
|
||
|
|
raise ValueError("At least one of video_path and audio_path is required")
|
||
|
|
|
||
|
|
content: list[dict[str, str]] = []
|
||
|
|
if video_path is not None:
|
||
|
|
content.append({"type": "video", "video": str(Path(video_path).resolve())})
|
||
|
|
if audio_path is not None:
|
||
|
|
content.append({"type": "audio", "audio": str(Path(audio_path).resolve())})
|
||
|
|
content.append({"type": "text", "text": prompt})
|
||
|
|
|
||
|
|
return [
|
||
|
|
{
|
||
|
|
"role": "system",
|
||
|
|
"content": [{
|
||
|
|
"type": "text",
|
||
|
|
"text": (
|
||
|
|
"You are a causal multimodal conversation observer. "
|
||
|
|
"Use only information available up to the current time."
|
||
|
|
),
|
||
|
|
}],
|
||
|
|
},
|
||
|
|
{"role": "user", "content": content},
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def prepare_inputs(
|
||
|
|
processor: Qwen2_5OmniProcessor,
|
||
|
|
conversations: list[list[dict[str, Any]]],
|
||
|
|
*,
|
||
|
|
use_audio_in_video: bool,
|
||
|
|
model_device: torch.device,
|
||
|
|
model_dtype: torch.dtype,
|
||
|
|
):
|
||
|
|
"""Convert a batch of conversations into model tensors."""
|
||
|
|
|
||
|
|
texts = processor.apply_chat_template(
|
||
|
|
conversations, add_generation_prompt=True, tokenize=False
|
||
|
|
)
|
||
|
|
audios, images, videos = process_mm_info(
|
||
|
|
conversations, use_audio_in_video=use_audio_in_video
|
||
|
|
)
|
||
|
|
inputs = processor(
|
||
|
|
text=texts,
|
||
|
|
audio=audios,
|
||
|
|
images=images,
|
||
|
|
videos=videos,
|
||
|
|
padding=True,
|
||
|
|
return_tensors="pt",
|
||
|
|
use_audio_in_video=use_audio_in_video,
|
||
|
|
)
|
||
|
|
# BatchFeature.to(dtype) casts floating tensors only; IDs/masks stay integer.
|
||
|
|
return inputs.to(model_device).to(model_dtype)
|
||
|
|
|
||
|
|
|
||
|
|
def _last_valid_token(states: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
|
||
|
|
positions = torch.arange(states.shape[1], device=states.device).unsqueeze(0)
|
||
|
|
indices = positions.masked_fill(~mask.bool(), -1).max(dim=1).values.clamp_min(0)
|
||
|
|
batch_indices = torch.arange(states.shape[0], device=states.device)
|
||
|
|
return states[batch_indices, indices]
|
||
|
|
|
||
|
|
|
||
|
|
def _tail_mean(
|
||
|
|
states: torch.Tensor, mask: torch.Tensor, tail_tokens: int
|
||
|
|
) -> torch.Tensor:
|
||
|
|
if tail_tokens <= 0:
|
||
|
|
raise ValueError("tail_tokens must be positive")
|
||
|
|
# Ranking valid tokens from the right supports both left and right padding.
|
||
|
|
reverse_ranks = mask.long().flip(1).cumsum(dim=1).flip(1)
|
||
|
|
tail_mask = mask.bool() & (reverse_ranks <= tail_tokens)
|
||
|
|
weights = tail_mask.unsqueeze(-1).to(states.dtype)
|
||
|
|
return (states * weights).sum(dim=1) / weights.sum(dim=1).clamp_min(1)
|
||
|
|
|
||
|
|
|
||
|
|
@torch.inference_mode()
|
||
|
|
def extract_features(
|
||
|
|
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
||
|
|
processor: Qwen2_5OmniProcessor,
|
||
|
|
conversations: list[list[dict[str, Any]]],
|
||
|
|
*,
|
||
|
|
use_audio_in_video: bool = False,
|
||
|
|
layer: int = -1,
|
||
|
|
tail_tokens: int = 8,
|
||
|
|
return_all_layers: bool = False,
|
||
|
|
output_device: str | torch.device = "cpu",
|
||
|
|
) -> OmniFeatures:
|
||
|
|
"""Extract fused Thinker states without generation.
|
||
|
|
|
||
|
|
``last_token`` is the recommended one-vector-per-window baseline. The
|
||
|
|
Thinker is causal, so it summarizes all earlier prompt and AV tokens.
|
||
|
|
"""
|
||
|
|
|
||
|
|
inputs = prepare_inputs(
|
||
|
|
processor,
|
||
|
|
conversations,
|
||
|
|
use_audio_in_video=use_audio_in_video,
|
||
|
|
model_device=model.device,
|
||
|
|
model_dtype=model.dtype,
|
||
|
|
)
|
||
|
|
# A direct forward ends at the observation boundary; generate() would append
|
||
|
|
# answer tokens and change which representation is being extracted.
|
||
|
|
outputs = model(
|
||
|
|
**inputs,
|
||
|
|
use_audio_in_video=use_audio_in_video,
|
||
|
|
output_hidden_states=True,
|
||
|
|
use_cache=False,
|
||
|
|
return_dict=True,
|
||
|
|
)
|
||
|
|
if outputs.hidden_states is None:
|
||
|
|
raise RuntimeError("The model did not return hidden states")
|
||
|
|
|
||
|
|
states = outputs.hidden_states[layer]
|
||
|
|
mask = inputs["attention_mask"]
|
||
|
|
return OmniFeatures(
|
||
|
|
token_features=states.to(output_device),
|
||
|
|
attention_mask=mask.to(output_device),
|
||
|
|
last_token=_last_valid_token(states, mask).to(output_device),
|
||
|
|
tail_mean=_tail_mean(states, mask, tail_tokens).to(output_device),
|
||
|
|
hidden_states=(
|
||
|
|
tuple(x.to(output_device) for x in outputs.hidden_states)
|
||
|
|
if return_all_layers else None
|
||
|
|
),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def parse_args() -> argparse.Namespace:
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument("--video", type=Path)
|
||
|
|
parser.add_argument("--audio", type=Path)
|
||
|
|
parser.add_argument("--output", type=Path, required=True)
|
||
|
|
parser.add_argument("--model-id", default=MODEL_ID)
|
||
|
|
parser.add_argument("--prompt", default=None)
|
||
|
|
parser.add_argument("--use-audio-in-video", action="store_true")
|
||
|
|
parser.add_argument("--tail-tokens", type=int, default=8)
|
||
|
|
parser.add_argument("--layer", type=int, default=-1)
|
||
|
|
parser.add_argument(
|
||
|
|
"--dtype",
|
||
|
|
choices=("auto", "bfloat16", "float16", "float32"),
|
||
|
|
default="auto",
|
||
|
|
)
|
||
|
|
parser.add_argument("--flash-attention-2", action="store_true")
|
||
|
|
parser.add_argument("--device", default="cuda:0")
|
||
|
|
return parser.parse_args()
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
args = parse_args()
|
||
|
|
if args.video is None and args.audio is None:
|
||
|
|
raise SystemExit("Pass --video, --audio, or both")
|
||
|
|
if args.use_audio_in_video and args.video is None:
|
||
|
|
raise SystemExit("--use-audio-in-video requires --video")
|
||
|
|
|
||
|
|
model, processor = load_qwen_omni(
|
||
|
|
args.model_id,
|
||
|
|
dtype=args.dtype,
|
||
|
|
device_map={"": args.device},
|
||
|
|
use_flash_attention_2=args.flash_attention_2,
|
||
|
|
)
|
||
|
|
conversation_kwargs: dict[str, Any] = {
|
||
|
|
"video_path": args.video,
|
||
|
|
"audio_path": args.audio,
|
||
|
|
}
|
||
|
|
if args.prompt is not None:
|
||
|
|
conversation_kwargs["prompt"] = args.prompt
|
||
|
|
conversation = build_conversation(**conversation_kwargs)
|
||
|
|
features = extract_features(
|
||
|
|
model,
|
||
|
|
processor,
|
||
|
|
[conversation],
|
||
|
|
use_audio_in_video=args.use_audio_in_video,
|
||
|
|
layer=args.layer,
|
||
|
|
tail_tokens=args.tail_tokens,
|
||
|
|
)
|
||
|
|
|
||
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
torch.save({
|
||
|
|
"model_id": args.model_id,
|
||
|
|
"layer": args.layer,
|
||
|
|
"token_features": features.token_features,
|
||
|
|
"attention_mask": features.attention_mask,
|
||
|
|
"last_token": features.last_token,
|
||
|
|
"tail_mean": features.tail_mean,
|
||
|
|
}, args.output)
|
||
|
|
print(f"Saved features to {args.output}")
|
||
|
|
print(f"token_features: {tuple(features.token_features.shape)}")
|
||
|
|
print(f"last_token: {tuple(features.last_token.shape)}")
|
||
|
|
print(f"tail_mean: {tuple(features.tail_mean.shape)}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|