Initial commit: Omni-only turn-taking experiment scaffold.
Add EasyCom prep, Qwen2.5-Omni feature extraction, streaming turn model prototype, and docs for the ICASSP 2027 ego-speak study. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Xin Wang
@@ -0,0 +1,270 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user