223 lines
7.0 KiB
Python
223 lines
7.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
# Copyright 2026 The vLLM team.
|
|
# Copyright 2026 NVIDIA CORPORATION and the HuggingFace Inc. team. All rights
|
|
# reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from transformers import PretrainedConfig
|
|
|
|
from tests.models.registry import HF_EXAMPLE_MODELS
|
|
|
|
|
|
class MockMusicFlamingoConfig(PretrainedConfig):
|
|
model_type = "musicflamingo"
|
|
|
|
def __init__(self, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.audio_config = PretrainedConfig()
|
|
self.text_config = PretrainedConfig()
|
|
|
|
|
|
class MockMusicFlamingoProcessor:
|
|
def __init__(self):
|
|
self.audio_token = "<sound>"
|
|
self.audio_token_id = 12345
|
|
self.audio_bos_token = "<|sound_bos|>"
|
|
self.audio_bos_token_id = 12346
|
|
self.audio_eos_token = "<|sound_eos|>"
|
|
self.audio_eos_token_id = 12347
|
|
self.max_audio_len = 1200
|
|
self.feature_extractor = MockFeatureExtractor()
|
|
|
|
|
|
class MockFeatureExtractor:
|
|
def __init__(self):
|
|
self.sampling_rate = 16000
|
|
self.chunk_length = 30
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_ctx():
|
|
config = MockMusicFlamingoConfig()
|
|
|
|
ctx = MagicMock()
|
|
ctx.get_hf_config.return_value = config
|
|
ctx.get_hf_processor.return_value = MockMusicFlamingoProcessor()
|
|
ctx.model_config.hf_config = config
|
|
return ctx
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def check_transformers_version():
|
|
model_info = HF_EXAMPLE_MODELS.get_hf_info("MusicFlamingoForConditionalGeneration")
|
|
model_info.check_transformers_version(on_fail="skip")
|
|
|
|
|
|
def test_musicflamingo_chunk_counting_uses_rote_timestamps(mock_ctx, monkeypatch):
|
|
from vllm.model_executor.models.musicflamingo import (
|
|
MusicFlamingoDummyInputsBuilder,
|
|
MusicFlamingoMultiModalProcessor,
|
|
MusicFlamingoProcessingInfo,
|
|
)
|
|
|
|
info = MusicFlamingoProcessingInfo(mock_ctx)
|
|
processor = MusicFlamingoMultiModalProcessor(
|
|
info, MusicFlamingoDummyInputsBuilder(info)
|
|
)
|
|
|
|
sr = 16000
|
|
audio_1 = np.zeros(30 * sr)
|
|
audio_2 = np.zeros(45 * sr)
|
|
|
|
mm_data = {"audio": [audio_1, audio_2]}
|
|
prompt = "<|user|>Listen.<|end|>"
|
|
|
|
from vllm.multimodal.processing import BaseMultiModalProcessor
|
|
|
|
def mock_base_call(self, prompt, mm_data, mm_kwargs, tok_kwargs):
|
|
del self, prompt, mm_data, mm_kwargs, tok_kwargs
|
|
return {
|
|
"input_ids": [1, 2, 3],
|
|
"input_features": torch.randn(3, 80, 3000),
|
|
"rote_timestamps": torch.randn(3, 750),
|
|
}
|
|
|
|
monkeypatch.setattr(BaseMultiModalProcessor, "_call_hf_processor", mock_base_call)
|
|
|
|
processed = processor._call_hf_processor(prompt, mm_data, {}, {})
|
|
|
|
chunk_counts = processed["chunk_counts"]
|
|
|
|
assert chunk_counts.tolist() == [1, 2]
|
|
assert "rote_timestamps" in processed
|
|
|
|
|
|
def test_musicflamingo_dummy_text_uses_plain_audio_tokens(mock_ctx):
|
|
from vllm.model_executor.models.musicflamingo import (
|
|
MusicFlamingoDummyInputsBuilder,
|
|
MusicFlamingoProcessingInfo,
|
|
)
|
|
|
|
info = MusicFlamingoProcessingInfo(mock_ctx)
|
|
builder = MusicFlamingoDummyInputsBuilder(info)
|
|
|
|
assert builder.get_dummy_text({"audio": 2}) == "<sound><sound>"
|
|
|
|
|
|
def test_musicflamingo_audio_feature_pipeline_matches_hf_small_config():
|
|
from transformers.models.musicflamingo import (
|
|
modeling_musicflamingo as hf_musicflamingo_modeling,
|
|
)
|
|
from transformers.models.musicflamingo.configuration_musicflamingo import (
|
|
MusicFlamingoConfig,
|
|
)
|
|
|
|
from vllm.model_executor.models.audioflamingo3 import (
|
|
_build_audio_encoder_attention_mask,
|
|
_flatten_valid_audio_embeddings,
|
|
)
|
|
from vllm.model_executor.models.musicflamingo import (
|
|
MusicFlamingoEncoder,
|
|
MusicFlamingoMultiModalProjector,
|
|
MusicFlamingoRotaryEmbedding,
|
|
apply_rotary_time_emb,
|
|
)
|
|
|
|
text_config = {
|
|
"model_type": "qwen2",
|
|
"intermediate_size": 64,
|
|
"initializer_range": 0.02,
|
|
"hidden_size": 32,
|
|
"max_position_embeddings": 1024,
|
|
"num_hidden_layers": 2,
|
|
"num_attention_heads": 4,
|
|
"num_key_value_heads": 2,
|
|
"vocab_size": 128,
|
|
"pad_token_id": 1,
|
|
"use_mrope": False,
|
|
}
|
|
audio_config = {
|
|
"hidden_size": 16,
|
|
"num_attention_heads": 4,
|
|
"intermediate_size": 32,
|
|
"num_hidden_layers": 2,
|
|
"num_mel_bins": 80,
|
|
"max_source_positions": 1500,
|
|
"dropout": 0.0,
|
|
"attention_dropout": 0.0,
|
|
"activation_dropout": 0.0,
|
|
"encoder_layerdrop": 0.0,
|
|
}
|
|
|
|
torch.manual_seed(0)
|
|
config = MusicFlamingoConfig(
|
|
text_config=text_config,
|
|
audio_config=audio_config,
|
|
audio_token_id=0,
|
|
head_dim=8,
|
|
rope_parameters={"rope_type": "default", "rope_theta": 2048},
|
|
)
|
|
hf_model = hf_musicflamingo_modeling.MusicFlamingoForConditionalGeneration(
|
|
config
|
|
).eval()
|
|
|
|
vllm_encoder = MusicFlamingoEncoder(config.audio_config).eval()
|
|
vllm_encoder.load_state_dict(hf_model.audio_tower.state_dict())
|
|
|
|
vllm_projector = MusicFlamingoMultiModalProjector(config).eval()
|
|
vllm_projector.load_state_dict(hf_model.multi_modal_projector.state_dict())
|
|
|
|
vllm_rope = MusicFlamingoRotaryEmbedding(config).eval()
|
|
vllm_rope.load_state_dict(hf_model.pos_emb.state_dict(), strict=False)
|
|
|
|
input_features = torch.randn(3, 80, 3000)
|
|
feature_attention_mask = torch.zeros(3, 3000, dtype=torch.bool)
|
|
feature_attention_mask[0, :3000] = True
|
|
feature_attention_mask[1, :2500] = True
|
|
feature_attention_mask[2, :1500] = True
|
|
rote_timestamps = (
|
|
torch.arange(750, dtype=torch.float32).unsqueeze(0).repeat(3, 1) * 0.04
|
|
)
|
|
|
|
hf_output = hf_model.get_audio_features(
|
|
input_features,
|
|
feature_attention_mask,
|
|
rote_timestamps=rote_timestamps,
|
|
return_dict=True,
|
|
).pooler_output
|
|
vllm_attention_mask = _build_audio_encoder_attention_mask(
|
|
feature_attention_mask,
|
|
dtype=vllm_encoder.conv1.weight.dtype,
|
|
device=vllm_encoder.conv1.weight.device,
|
|
)
|
|
vllm_hidden_states = vllm_encoder(
|
|
input_features,
|
|
attention_mask=vllm_attention_mask,
|
|
)
|
|
cos, sin = vllm_rope(rote_timestamps, seq_len=vllm_hidden_states.shape[-2])
|
|
vllm_hidden_states = apply_rotary_time_emb(vllm_hidden_states, cos, sin)
|
|
vllm_output, _ = _flatten_valid_audio_embeddings(
|
|
vllm_projector(vllm_hidden_states),
|
|
feature_attention_mask,
|
|
)
|
|
|
|
torch.testing.assert_close(vllm_output, hf_output)
|