From efda9369a8082c96ebf55497880a695d84dbefc2 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 9 Sep 2026 22:25:05 +0800 Subject: [PATCH] refactor: reduce copied vLLM 0.29 compatibility code Reuse upstream weight mapping and multimodal processing, consolidate Omni processor bridges, and trim unused Hunyuan internals. Route no-media kwargs through the active renderer and preserve GLM generation prompt formatting. Validation: 453 CPU tests passed, 7 skipped; processor imports, Ruff, and diff checks passed. GPU and end-to-end CI not rerun. Signed-off-by: tzhouam --- .../test_async_omni_engine_stage_init.py | 63 +-- tests/inputs/test_mm_processor.py | 44 ++ tests/inputs/test_preprocess.py | 112 +++- .../models/glm_image/test_glm_image_ar.py | 38 ++ .../test_long_audio_processing.py | 28 + .../qwen2_5_omni/test_per_video_audio_mask.py | 7 + .../qwen3_omni/test_qwen3_omni_moe_lora.py | 4 +- .../models/test_weight_loader.py | 17 +- .../hunyuan_image3/pipeline_hunyuan_image3.py | 10 +- vllm_omni/engine/stage_init_utils.py | 7 +- vllm_omni/entrypoints/openai/serving_chat.py | 2 +- vllm_omni/inputs/mm_processor.py | 39 ++ vllm_omni/inputs/preprocess.py | 404 ++------------ .../models/audex/audex_omni_processing.py | 21 +- .../model_executor/models/bagel/bagel.py | 11 +- .../models/cosyvoice3/cosyvoice3.py | 24 +- .../models/covo_audio/covo_audio.py | 20 +- .../models/gepard/gepard_talker.py | 9 +- .../models/glm_image/glm_image_ar.py | 502 +----------------- .../model_executor/models/glm_tts/glm_tts.py | 24 +- .../higgs_audio_v3/higgs_audio_v3_talker.py | 8 +- .../hunyuan_image3/_hunyuan_v1_vendored.py | 209 +------- .../models/hunyuan_image3/hunyuan_image3.py | 32 +- .../models/mimo_audio/mimo_audio.py | 20 +- .../ming_flash_omni/ming_flash_omni_talker.py | 13 +- .../minicpmo_4_5/minicpmo_4_5_omni_llm.py | 31 +- .../models/minimax_music3/talker.py | 8 +- .../models/omnivoice/omnivoice.py | 24 +- .../qwen2_5_omni/qwen2_5_omni_talker.py | 11 +- .../qwen2_5_omni/qwen2_5_omni_thinker.py | 11 +- .../models/qwen2_5_omni/qwen2_old.py | 12 +- .../models/qwen3_omni/qwen3_omni_code2wav.py | 10 +- .../qwen3_omni/qwen3_omni_moe_talker.py | 13 +- .../qwen3_omni/qwen3_omni_moe_thinker.py | 93 +--- .../models/qwen3_tts/qwen3_tts_code2wav.py | 9 +- .../models/qwen3_tts/qwen3_tts_talker.py | 10 +- .../model_executor/models/weight_loader.py | 54 -- 37 files changed, 448 insertions(+), 1506 deletions(-) create mode 100644 tests/inputs/test_mm_processor.py create mode 100644 vllm_omni/inputs/mm_processor.py delete mode 100644 vllm_omni/model_executor/models/weight_loader.py diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index e6f7f8a8b..79317805f 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -1593,73 +1593,42 @@ def test_model_path_resolver_is_generic_and_model_owned(tmp_path): assert "model_path_resolver" not in engine_args -def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch): +def test_build_stage0_input_processor_wraps_active_renderer(monkeypatch): import vllm_omni.engine.stage_init_utils as init_mod + from vllm_omni.inputs.preprocess import OmniRenderer + + original = object() class DummyInputProcessor: def __init__(self, vllm_config, renderer=None): - self.vllm_config = vllm_config - self.renderer = renderer or object() - self.input_preprocessor = None - - class DummyOmniInputPreprocessor: - def __init__(self, vllm_config, renderer=None): - self.vllm_config = vllm_config - self.renderer = renderer + self.renderer = renderer or original monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor) - monkeypatch.setattr(init_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor) - - input_processor = build_stage0_input_processor( + processor = build_stage0_input_processor( types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {})) ) - - assert isinstance(input_processor.input_preprocessor, DummyOmniInputPreprocessor) - assert input_processor.input_preprocessor.renderer is input_processor.renderer + assert isinstance(processor.renderer, OmniRenderer) + assert processor.renderer._renderer is original + assert not hasattr(processor, "input_preprocessor") -def test_build_stage0_input_processor_does_not_resolve_tokenizer_when_skipped( - monkeypatch, -): +def test_build_stage0_input_processor_does_not_resolve_tokenizer_when_skipped(monkeypatch): import vllm_omni.engine.stage_init_utils as init_mod - token_only_renderer = object() - seen = {} + original = object() class DummyInputProcessor: def __init__(self, vllm_config, renderer=None): - if renderer is None: - raise AssertionError("skip_tokenizer_init must supply a renderer") - seen["renderer"] = renderer + assert renderer is original self.renderer = renderer - self.input_preprocessor = None - - class DummyOmniInputPreprocessor: - def __init__(self, vllm_config, renderer=None): - seen["preprocessor_renderer"] = renderer monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor) - monkeypatch.setattr( - init_mod, - "_build_token_only_renderer", - lambda _config: token_only_renderer, - ) - monkeypatch.setattr( - init_mod, - "OmniInputPreprocessor", - DummyOmniInputPreprocessor, - ) - + monkeypatch.setattr(init_mod, "_build_token_only_renderer", lambda _: original) config = types.SimpleNamespace( - model_config=types.SimpleNamespace( - skip_tokenizer_init=True, - try_get_generation_config=lambda: {}, - ) + model_config=types.SimpleNamespace(skip_tokenizer_init=True, try_get_generation_config=lambda: {}) ) - build_stage0_input_processor(config) - - assert seen["renderer"] is token_only_renderer - assert seen["preprocessor_renderer"] is token_only_renderer + processor = build_stage0_input_processor(config) + assert processor.renderer._renderer is original def test_inject_kv_stage_info_infers_sender_tp_topology(): diff --git a/tests/inputs/test_mm_processor.py b/tests/inputs/test_mm_processor.py new file mode 100644 index 000000000..813d7f436 --- /dev/null +++ b/tests/inputs/test_mm_processor.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from transformers import BatchFeature + +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("counts", [{}, {"audio": 1, "image": 0}]) +@pytest.mark.parametrize("original_text", [None, "", "voice cloning text"]) +def test_custom_processor_bridge_preserves_prompt_kwargs_and_passthrough(counts, original_text): + class Processor(OmniMultiModalProcessor): + def _get_mm_fields_config(self, *args): + return {} + + def _get_prompt_updates(self, *args): + return [] + + processor = object.__new__(Processor) + processor.dummy_inputs = SimpleNamespace(get_dummy_text=Mock(return_value="dummy")) + processor._get_hf_mm_data = Mock(return_value=({"audios": ["wave"]}, {"embedding": "preserve"})) + processor._call_hf_processor = Mock(return_value=BatchFeature({"features": "processed"})) + items = Mock() + items.get_all_counts.return_value = counts + kwargs = {"sampling_rate": 16000} + if original_text is not None: + kwargs[processor._OMNI_PROMPT_TEXT_KEY] = original_text + before = dict(kwargs) + result = processor._apply_hf_processor_main(items, kwargs) + items.select.assert_called_once_with({k for k, v in counts.items() if v > 0}) + processor._call_hf_processor.assert_called_once_with( + "dummy" if original_text is None else original_text, + {"audios": ["wave"]}, + {"sampling_rate": 16000}, + {}, + ) + assert dict(result) == {"features": "processed", "embedding": "preserve"} + assert kwargs == before diff --git a/tests/inputs/test_preprocess.py b/tests/inputs/test_preprocess.py index db8817cf2..1b5e057fb 100644 --- a/tests/inputs/test_preprocess.py +++ b/tests/inputs/test_preprocess.py @@ -1,35 +1,99 @@ # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for OmniInputPreprocessor._process_text routing.""" +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exercise Omni routing through the active upstream rendering pipeline.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock import pytest -from pytest_mock import MockerFixture +import torch +from vllm.inputs import tokens_input +from vllm.pooling_params import PoolingParams +from vllm.renderers import BaseRenderer +from vllm.renderers.params import TokenizeParams +from vllm.v1.engine.input_processor import InputProcessor + +from vllm_omni.inputs.preprocess import OmniRenderer pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -class TestProcessTextMmProcessorKwargsRouting: - """Presence-based routing: an explicit empty ``mm_processor_kwargs`` dict - on the prompt still routes through ``_process_multimodal``. Required for - AR-based image generation (e.g. GLM-Image t2i) where the HF processor - supplies the image-generation scaffold from its own defaults and no extra - kwargs are needed from the caller. - """ +class _Renderer(BaseRenderer): + def __init__(self): + self.model_config = SimpleNamespace(is_encoder_decoder=False, enable_prompt_embeds=True) + self.tokenizer = None + self.default_cmpl_tok_params = TokenizeParams(max_total_tokens=128) + self._tokenize_prompt = lambda prompt, params: {**prompt, "prompt_token_ids": [1, 2, 3]} + self._tokenize_prompt_async = AsyncMock(side_effect=self._tokenize_prompt) + self._process_multimodal = Mock(return_value=tokens_input([1, 2, 3, 99])) + self._process_multimodal_async = AsyncMock(side_effect=self._process_multimodal) - @pytest.fixture - def preprocessor(self, mocker: MockerFixture): - from vllm_omni.inputs.preprocess import OmniInputPreprocessor + def render_messages(self, messages, params): + return messages, {"prompt": "hello"} - instance = object.__new__(OmniInputPreprocessor) - instance._process_multimodal = mocker.MagicMock(return_value={}) - instance._tokenize_prompt = mocker.MagicMock(return_value=[1, 2, 3]) - return instance - def test_empty_mm_processor_kwargs_routes_to_multimodal(self, preprocessor): - preprocessor._process_text({"prompt": "hello", "mm_processor_kwargs": {}}) - assert preprocessor._process_multimodal.called - assert not preprocessor._tokenize_prompt.called +@pytest.fixture +def renderer(): + return OmniRenderer(_Renderer()) - def test_missing_mm_processor_kwargs_routes_to_tokenize(self, preprocessor): - preprocessor._process_text({"prompt": "hello"}) - assert preprocessor._tokenize_prompt.called - assert not preprocessor._process_multimodal.called + +@pytest.mark.parametrize("kwargs", [{}, {"target_h": 512, "target_w": 768}]) +@pytest.mark.parametrize("tokenized", [False, True]) +def test_process_inputs_routes_no_media_processor_kwargs(renderer, kwargs, tokenized): + # Keep real process_inputs/render_cmpl; stub unrelated config validation. + processor = object.__new__(InputProcessor) + processor.renderer = renderer + processor.model_config = renderer.model_config + processor.vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace( + data_parallel_size=1, + data_parallel_size_local=1, + local_engines_only=False, + ) + ) + processor._validate_params = Mock() + processor._validate_lora = Mock() + processor._validate_model_inputs = Mock() + prompt = {"prompt_token_ids": [1, 2, 3]} if tokenized else {"prompt": "hello"} + prompt.update(mm_processor_kwargs=kwargs, cache_salt="salt") + request = processor.process_inputs("image-request", prompt, PoolingParams(), ("embed",)) + assert request.prompt_token_ids == [1, 2, 3, 99] + assert request.cache_salt == "salt" + renderer._renderer._process_multimodal.assert_called_once_with( + [1, 2, 3], + {}, + mm_processor_kwargs=kwargs, + mm_uuids=None, + skip_mm_cache=False, + ) + + +@pytest.mark.parametrize("async_mode", [False, True]) +@pytest.mark.parametrize("kind", ["text", "tokens", "media", "kwargs", "embeds"]) +def test_renderer_preserves_routes_extras_and_cache_policy(renderer, async_mode, kind): + prompt = {"prompt": "hello", "additional_information": {"speaker": 1}, "model_intermediate_buffer": {"x": 2}} + if kind == "tokens": + prompt["prompt_token_ids"] = [1, 2, 3] + elif kind == "media": + prompt["multi_modal_data"] = {"image": "image"} + elif kind == "kwargs": + prompt["mm_processor_kwargs"] = {} + elif kind == "embeds": + prompt["prompt_embeds"] = torch.ones(3, 4) + if async_mode: + (result,) = asyncio.run(renderer.render_cmpl_async([prompt], skip_mm_cache=True)) + else: + (result,) = renderer.render_cmpl([prompt], skip_mm_cache=True) + assert result["additional_information"] == {"speaker": 1} + assert result["model_intermediate_buffer"] == {"x": 2} + assert result["prompt"] == "hello" + if kind in ("media", "kwargs"): + assert result["prompt_token_ids"] == [1, 2, 3, 99] + assert renderer._renderer._process_multimodal.call_args.kwargs["skip_mm_cache"] is True + else: + renderer._renderer._process_multimodal.assert_not_called() + if kind == "embeds": + torch.testing.assert_close(result["prompt_embeds"], prompt["prompt_embeds"]) + else: + assert result["prompt_token_ids"] == [1, 2, 3] diff --git a/tests/model_executor/models/glm_image/test_glm_image_ar.py b/tests/model_executor/models/glm_image/test_glm_image_ar.py index 5ccc9fc62..936980ef8 100644 --- a/tests/model_executor/models/glm_image/test_glm_image_ar.py +++ b/tests/model_executor/models/glm_image/test_glm_image_ar.py @@ -91,6 +91,44 @@ GlmImageRotaryEmbedding = _ar_mod.GlmImageRotaryEmbedding pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.mark.parametrize("num_images", [0, 1, 2]) +@pytest.mark.parametrize("preformatted", [False, True]) +def test_generation_prompt_ids_preserve_hf_target_scaffold(num_images, preformatted): + from functools import partial + from types import SimpleNamespace + + from transformers.models.glm_image.processing_glm_image import GlmImageProcessor + from vllm.multimodal.parse import MultiModalDataItems + from vllm.multimodal.processing import BaseMultiModalProcessor, ProcessorInputs + + processor = object.__new__(GlmImageMultiModalProcessor) + symbols = SimpleNamespace(grid_bos_token="", grid_eos_token="", bos_token="") + hf_processor = SimpleNamespace( + _build_prompt_with_target_shape=partial(GlmImageProcessor._build_prompt_with_target_shape, symbols) + ) + tokenizer = SimpleNamespace(encode=lambda text, **kwargs: list(text.encode())) + processor.info = SimpleNamespace( + get_tokenizer=lambda: tokenizer, + get_hf_processor=lambda: hf_processor, + get_hf_config=lambda: _make_hf_config(), + ) + items = MagicMock(spec=MultiModalDataItems) + items.get_all_counts.return_value = {"image": num_images} + suffix = "16 24" + ("13 19" if not num_images else "") + "" + text_ids = list(b"a red bird") + expected = [167855] * num_images + text_ids + list(suffix.encode()) + inputs = ProcessorInputs( + prompt=expected if preformatted else text_ids, + mm_data_items=items, + hf_processor_mm_kwargs={"target_h": 512, "target_w": 768}, + ) + with patch.object(BaseMultiModalProcessor, "apply", side_effect=lambda inputs, timing: inputs) as upstream: + result = processor.apply(inputs, MagicMock()) + assert result.prompt == expected + assert inputs.prompt == (expected if preformatted else text_ids) + upstream.assert_called_once() + + def test_glm_image_ar_omits_hidden_pooler_payload(): assert GlmImageForConditionalGeneration.omni_pooler_payload_include_hidden is False diff --git a/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py b/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py index a76f752c5..eaa2ff2b3 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py +++ b/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py @@ -43,6 +43,34 @@ _FEAT_DIM = 128 _CHUNK_FRAMES = 3000 +@pytest.mark.parametrize("version", [(2, 5), (4, 5)]) +def test_prompt_batching_delegates_upstream_without_tts_kwargs(version): + from unittest.mock import MagicMock + + info = SimpleNamespace( + get_model_version=lambda: version, + get_hf_processor=MagicMock(return_value=object()), + ctx=SimpleNamespace(call_hf_processor=MagicMock()), + ) + info.ctx.call_hf_processor.side_effect = ( + [{"features": [1]}, {"features": [2]}] if version == (2, 5) else [{"features": [1, 2]}] + ) + kwargs = {"use_tts": True, "sampling_rate": 16000} + result = MiniCPMO45OmniLLMMultiModalProcessor._call_hf_processor_on_prompts( + SimpleNamespace(info=info), ["one", "two"], {"audios": [10, 20]}, kwargs, out_keys={"features"} + ) + assert result == {"features": [1, 2]} + assert kwargs == {"use_tts": True, "sampling_rate": 16000} + for call in info.ctx.call_hf_processor.call_args_list: + assert call.args[2] == {"sampling_rate": 16000} + payloads = [call.args[1] for call in info.ctx.call_hf_processor.call_args_list] + assert payloads == ( + [{"text": "one", "audios": 10}, {"text": "two", "audios": 20}] + if version == (2, 5) + else [{"text": ["one", "two"], "audios": [10, 20]}] + ) + + class TestMiniCPMOFieldConfig: def test_short_audios_stay_batched(self) -> None: # two <=30s audios: one feature entry + one 1-element lens per audio diff --git a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py index 2683a8d9d..54efc99e4 100644 --- a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py +++ b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py @@ -620,10 +620,17 @@ def test_derive_audio_from_video_placeholders_only_pairs_true_videos(): def test_qwen3_processor_inherits_vllm_omni_per_video_helpers(): + from vllm.model_executor.models.qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerMultiModalProcessor as UpstreamQwen3Processor, + ) + assert issubclass( Qwen3OmniMoeThinkerMultiModalProcessor, Qwen2_5OmniThinkerMultiModalProcessor, ) + assert Qwen3OmniMoeThinkerMultiModalProcessor._get_hf_mm_data is UpstreamQwen3Processor._get_hf_mm_data + mro = Qwen3OmniMoeThinkerMultiModalProcessor.__mro__ + assert mro.index(Qwen2_5OmniThinkerMultiModalProcessor) < mro.index(UpstreamQwen3Processor) def test_qwen3_prompt_updates_prefer_hf_second_per_grid_ts_from_outputs(): diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py index 75d81eed5..4ff6d5096 100644 --- a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py @@ -113,7 +113,9 @@ def _save_peft_adapter(tmp_path: Path) -> Path: return adapter_dir -def test_peft_expert_lora_loads_into_thinker(tmp_path: Path): +def test_peft_expert_lora_loads_into_thinker(tmp_path: Path, monkeypatch): + # This CPU harness tests weight mapping, not CUDA pinned-memory allocation. + monkeypatch.setattr("vllm.lora.lora_model.PIN_MEMORY", False) adapter_dir = _save_peft_adapter(tmp_path) peft_helper = PEFTHelper.from_local_dir(str(adapter_dir), max_position_embeddings=128) lora_model = LoRAModel.from_local_checkpoint( diff --git a/tests/model_executor/models/test_weight_loader.py b/tests/model_executor/models/test_weight_loader.py index 458eae25f..34d1cd28a 100644 --- a/tests/model_executor/models/test_weight_loader.py +++ b/tests/model_executor/models/test_weight_loader.py @@ -4,9 +4,7 @@ import pytest import torch from torch import nn -from vllm.model_executor.models.utils import WeightsMapper - -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -20,21 +18,18 @@ class _TinyModel(nn.Module): self.rotary_embed.register_parameter("inv_freq", nn.Parameter(torch.zeros(1))) -def test_auto_weights_loader_preserves_removed_skip_arguments() -> None: +def test_upstream_mapper_preserves_stage_exclusions_after_renaming() -> None: model = _TinyModel() model.skipped.weight.data.zero_() - loaded = AutoWeightsLoader( - model, - skip_prefixes=["skipped."], - skip_substrs=["rotary_embed.inv_freq"], - ).load_weights( + loaded = AutoWeightsLoader(model).load_weights( [ ("checkpoint.kept.weight", torch.ones(1, 1)), - ("skipped.weight", torch.full((1, 1), 2.0)), + ("checkpoint.skipped.weight", torch.full((1, 1), 2.0)), ("rotary_embed.inv_freq", torch.ones(1)), ], - mapper=WeightsMapper(orig_to_new_prefix={"checkpoint.": ""}), + mapper=WeightsMapper(orig_to_new_prefix={"checkpoint.": ""}) + | WeightsMapper(orig_to_new_prefix={"skipped.": None}, orig_to_new_substr={"rotary_embed.inv_freq": None}), ) assert loaded == {"kept.weight"} diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index a2f9b0699..fcfd48b4f 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -17,7 +17,7 @@ from transformers.generation.configuration_utils import GenerationConfig from transformers.generation.utils import ALL_CACHE_NAMES, GenerationMixin from transformers.utils.generic import ModelOutput from vllm.config.vllm import get_current_vllm_config -from vllm.model_executor.models.utils import WeightsMapper +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm.transformers_utils.config import get_config from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig @@ -33,7 +33,6 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import Siglip2VisionTransformer -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader from . import request_layout as request_layout_utils from .hunyuan_image3_tokenizer import TokenizerEncodeOutput, TokenizerWrapper @@ -455,11 +454,10 @@ class HunyuanImage3Pipeline( # Note: guidance_emb and timestep_r_emb are no longer skipped # to support HunyuanImage-3.0-Distil and MeanFlow distilled models - loader = AutoWeightsLoader( - self, - skip_prefixes=skip_prefixes, + loader = AutoWeightsLoader(self) + return loader.load_weights( + weights, mapper=WeightsMapper(orig_to_new_prefix={name: None for name in (skip_prefixes or ())}) ) - return loader.load_weights(weights) def prepare_seed(self, seed=None, batch_size=1): # random seed diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 765549bcb..9bdcd647f 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -54,7 +54,7 @@ from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices from vllm_omni.entrypoints.utils import filter_dataclass_kwargs from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams -from vllm_omni.inputs.preprocess import OmniInputPreprocessor +from vllm_omni.inputs.preprocess import OmniRenderer from vllm_omni.outputs.output_processor import MultimodalOutputProcessor from vllm_omni.platforms import current_omni_platform from vllm_omni.quantization.inc_config import OmniINCConfig @@ -1563,10 +1563,7 @@ def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor: ) else: input_processor = InputProcessor(vllm_config=stage_vllm_config) - input_processor.input_preprocessor = OmniInputPreprocessor( - vllm_config=stage_vllm_config, - renderer=input_processor.renderer, - ) + input_processor.renderer = OmniRenderer(input_processor.renderer) return input_processor diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 6b50b3887..f5cb24efd 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -702,7 +702,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if negative_prompt is not None: tprompt["negative_prompt"] = negative_prompt # Always attach mm_processor_kwargs (possibly empty) so - # OmniInputPreprocessor._process_text routes through the + # OmniRenderer routes through the # multimodal processor path. Without it, the preprocessor # falls back to plain _tokenize_prompt and AR-based image-gen # models like GLM-Image never see their image-generation diff --git a/vllm_omni/inputs/mm_processor.py b/vllm_omni/inputs/mm_processor.py new file mode 100644 index 000000000..bb02fe961 --- /dev/null +++ b/vllm_omni/inputs/mm_processor.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Multimodal data plumbing for Omni's custom, non-HF processors.""" + +from collections.abc import Mapping +from typing import TypeVar + +from transformers import BatchFeature +from vllm.multimodal.parse import MultiModalDataItems +from vllm.multimodal.processing import BaseMultiModalProcessor, BaseProcessingInfo + +_I = TypeVar("_I", bound=BaseProcessingInfo) + + +class OmniMultiModalProcessor(BaseMultiModalProcessor[_I]): + """Share data/passthrough handling, leaving model processing in its owner. + + These processors implement ``_call_hf_processor`` locally, sometimes + without any HF processor object. Prompt tokens are prepared separately + in ``apply``; the upstream base still owns caching and prompt updates. + """ + + _OMNI_PROMPT_TEXT_KEY = "_vllm_omni_original_prompt_text" + + def _apply_hf_processor_main( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> BatchFeature: + kwargs = dict(hf_processor_mm_kwargs) + prompt = kwargs.pop(self._OMNI_PROMPT_TEXT_KEY, None) + if prompt is None: + prompt = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) + valid_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) + mm_data, passthrough = self._get_hf_mm_data(valid_items) + result = self._call_hf_processor(str(prompt), dict(mm_data), kwargs, {}) + result.update(passthrough) + return result diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 6d6fe1d8a..28fe29236 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -3,379 +3,61 @@ from typing import Any -from typing_extensions import assert_never -from vllm.inputs import ( - EmbedsInput, - MultiModalDataDict, - MultiModalInput, - MultiModalUUIDDict, - SingletonInput, - TextPrompt, - TokensInput, - TokensPrompt, - tokens_input, -) -from vllm.logger import init_logger -from vllm.renderers import BaseRenderer, renderer_from_config -from vllm.renderers.inputs import SingletonDictPrompt - -from vllm_omni.inputs.data import ( - OmniEmbedsPrompt, - OmniTextPrompt, - OmniTokenInputs, - OmniTokensPrompt, - token_inputs_omni, -) - -logger = init_logger(__name__) +from vllm.renderers import BaseRenderer -class _InputPreprocessor: - """Compatibility shim for the upstream-removed ``InputPreprocessor``. +class OmniRenderer: + """Delegate tokenization/chat rendering while extending engine-input routing. - Upstream vLLM moved raw-prompt preprocessing into - :class:`vllm.renderers.BaseRenderer` and removed - ``vllm.inputs.preprocess.InputPreprocessor``. vLLM-Omni still exposes an - omni-specific preprocessor object for stage-0 input processing, so keep the - same thin renderer-delegating surface locally and adapt the few calls whose - renderer signature changed upstream. + Keep the original renderer's tokenizer, executors, and multimodal caches. + The upstream pipeline methods run on this adapter so both sync and async + entrypoints reach the Omni singleton processing below. """ - def __init__( - self, - vllm_config: Any, - renderer: BaseRenderer | None = None, - ) -> None: - self.model_config = vllm_config.model_config - self.renderer = renderer or renderer_from_config(vllm_config) + render_cmpl = BaseRenderer.render_cmpl + render_cmpl_async = BaseRenderer.render_cmpl_async + render_chat = BaseRenderer.render_chat + render_chat_async = BaseRenderer.render_chat_async + process_for_engine = BaseRenderer.process_for_engine + process_for_engine_async = BaseRenderer.process_for_engine_async + _process_enc_dec = BaseRenderer._process_enc_dec + _process_enc_dec_async = BaseRenderer._process_enc_dec_async - @property - def tokenizer(self) -> Any | None: - return self.renderer.tokenizer + def __init__(self, renderer: BaseRenderer) -> None: + self._renderer = renderer - def _tokenize_prompt( - self, - prompt: str, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> list[int]: - """Apply the model's tokenizer to a text prompt.""" - tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(**(tokenization_kwargs or {})) - - tok_prompt = self.renderer._tokenize_singleton_prompt( - TextPrompt(prompt=prompt), - tok_params, - ) - - return tok_prompt["prompt_token_ids"] - - def _process_multimodal( - self, - prompt: str | list[int], - mm_data: MultiModalDataDict, - mm_processor_kwargs: dict[str, Any] | None = None, - tokenization_kwargs: dict[str, Any] | None = None, - *, - mm_uuids: MultiModalUUIDDict | None = None, - ) -> MultiModalInput: - """Apply the model's multi-modal processor to a multi-modal prompt. - - The upstream renderer no longer accepts raw ``str`` prompts or - ``tokenization_kwargs`` in ``_process_multimodal``; tokenize text - prompts here first to preserve the legacy ``InputPreprocessor`` API. - """ - if isinstance(prompt, str): - prompt = self._tokenize_prompt(prompt, tokenization_kwargs) - - return self.renderer._process_multimodal( - prompt, - mm_data, - mm_uuids=mm_uuids, - mm_processor_kwargs=mm_processor_kwargs, - ) - - def _process_embeds( - self, - parsed_content: OmniEmbedsPrompt, - ) -> EmbedsInput: - return self.renderer._process_embeds(parsed_content) - - def _truncate_inputs( - self, - inputs: list[int], - tokenization_kwargs: dict[str, Any] | None = None, - ) -> list[int]: - tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(**(tokenization_kwargs or {})) - - tok_prompt = self.renderer._tokenize_singleton_prompt( - TokensPrompt(prompt_token_ids=inputs), - tok_params, - ) - - return tok_prompt["prompt_token_ids"] - - def _process_tokens( - self, - parsed_content: OmniTokensPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> TokensInput | MultiModalInput: - prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs) - - inputs: TokensInput | MultiModalInput - if multi_modal_data := parsed_content.get("multi_modal_data"): - inputs = self._process_multimodal( - prompt_token_ids, - multi_modal_data, - parsed_content.get("mm_processor_kwargs"), - tokenization_kwargs=tokenization_kwargs, - mm_uuids=parsed_content.get("multi_modal_uuids"), - ) - else: - inputs = tokens_input(prompt_token_ids) - - if prompt_text := parsed_content.get("prompt"): - inputs["prompt"] = prompt_text - if cache_salt := parsed_content.get("cache_salt"): - inputs["cache_salt"] = cache_salt + def __getattr__(self, name: str) -> Any: + return getattr(self._renderer, name) + @staticmethod + def _with_omni_extras(inputs, prompt): + for key in ("prompt", "cache_salt", "additional_information", "model_intermediate_buffer"): + if key in prompt: + inputs[key] = prompt[key] return inputs - def _process_text( - self, - parsed_content: OmniTextPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - *, - mm_uuids: MultiModalUUIDDict | None = None, - ) -> TokensInput | MultiModalInput: - prompt_text = parsed_content["prompt"] - - inputs: TokensInput | MultiModalInput - if multi_modal_data := parsed_content.get("multi_modal_data"): - inputs = self._process_multimodal( - prompt_text, - multi_modal_data, - parsed_content.get("mm_processor_kwargs") or {}, - tokenization_kwargs=tokenization_kwargs, - mm_uuids=mm_uuids, - ) - else: - prompt_token_ids = self._tokenize_prompt( - prompt_text, - tokenization_kwargs=tokenization_kwargs, - ) - inputs = tokens_input(prompt_token_ids) - - inputs["prompt"] = prompt_text - - if cache_salt := parsed_content.get("cache_salt"): - inputs["cache_salt"] = cache_salt - - return inputs - - def _prompt_to_llm_inputs( - self, - prompt: SingletonDictPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - *, - mm_uuids: MultiModalUUIDDict | None = None, - ) -> SingletonInput: - """Extract the singleton inputs from a prompt.""" - if "prompt_embeds" in prompt: - return self._process_embeds(prompt) # type: ignore[arg-type] - - if "prompt_token_ids" in prompt: - return self._process_tokens( - prompt, # type: ignore[arg-type] - tokenization_kwargs=tokenization_kwargs, - ) - - if "prompt" in prompt: - return self._process_text( - prompt, # type: ignore[arg-type] - tokenization_kwargs=tokenization_kwargs, - mm_uuids=mm_uuids, - ) - - assert_never(prompt) # type: ignore[arg-type] - - -class OmniInputPreprocessor(_InputPreprocessor): - """Input preprocessor for omni models. - - Extends the base InputPreprocessor to handle omni-specific input - types including prompt embeddings and additional information payloads. - Supports processing tokens, embeddings, text, and multimodal inputs. - """ - - def _process_text( - self, - parsed_content: OmniTextPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - *, - mm_uuids: Any | None = None, - ) -> OmniTokenInputs | MultiModalInput: - """Process text prompts with support for mm_processor_kwargs. - - Extends base class to support mm_processor_kwargs without multi_modal_data. - This is needed for models like GLM-Image where text-to-image generation - requires processor kwargs (target_h, target_w) to format the prompt. - """ - prompt_text = parsed_content["prompt"] - mm_processor_kwargs = parsed_content.get("mm_processor_kwargs") or {} - # When the deprecated raw-prompt path is used, process_inputs does - # not pass mm_uuids to preprocess(). Fall back to reading it from - # the prompt dict so the Renderer's _validate_mm_uuids can see it. - effective_mm_uuids = mm_uuids or parsed_content.get("multi_modal_uuids") - - inputs: OmniTokenInputs | MultiModalInput - if multi_modal_data := parsed_content.get("multi_modal_data"): - inputs = self._process_multimodal( - prompt_text, - multi_modal_data, - mm_processor_kwargs, - tokenization_kwargs=tokenization_kwargs, - mm_uuids=effective_mm_uuids, - ) - prompt_embeds = parsed_content.get("prompt_embeds") - if prompt_embeds is not None: - inputs["prompt_embeds"] = prompt_embeds - additional_information = parsed_content.get("additional_information") - if additional_information is not None: - inputs["additional_information"] = additional_information - model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") - if model_intermediate_buffer is not None: - inputs["model_intermediate_buffer"] = model_intermediate_buffer - elif "mm_processor_kwargs" in parsed_content: - # Presence — not truthiness. An explicitly-set empty dict still - # signals "route through the multimodal processor" (needed for - # AR-based image-gen where the HF processor supplies its own - # defaults and scaffold). - inputs = self._process_multimodal( - prompt_text, + def _process_singleton(self, prompt, *, skip_mm_cache: bool = False): + if "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data"): + inputs = self._renderer._process_multimodal( + prompt["prompt_token_ids"], {}, - mm_processor_kwargs, - tokenization_kwargs=tokenization_kwargs, - mm_uuids=effective_mm_uuids, + mm_processor_kwargs=prompt["mm_processor_kwargs"], + mm_uuids=prompt.get("multi_modal_uuids"), + skip_mm_cache=skip_mm_cache, ) else: - prompt_token_ids = self._tokenize_prompt( - prompt_text, - tokenization_kwargs=tokenization_kwargs, + inputs = self._renderer._process_singleton(prompt, skip_mm_cache=skip_mm_cache) + return self._with_omni_extras(inputs, prompt) + + async def _process_singleton_async(self, prompt, *, skip_mm_cache: bool = False): + if "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data"): + inputs = await self._renderer._process_multimodal_async( + prompt["prompt_token_ids"], + {}, + mm_processor_kwargs=prompt["mm_processor_kwargs"], + mm_uuids=prompt.get("multi_modal_uuids"), + skip_mm_cache=skip_mm_cache, ) - inputs = token_inputs_omni( - prompt_token_ids, - prompt_embeds=parsed_content.get("prompt_embeds"), - additional_information=parsed_content.get("additional_information"), - model_intermediate_buffer=parsed_content.get("model_intermediate_buffer"), - ) - prompt_embeds = parsed_content.get("prompt_embeds") - if prompt_embeds is not None: - inputs["prompt_embeds"] = prompt_embeds - additional_information = parsed_content.get("additional_information") - if additional_information is not None: - inputs["additional_information"] = additional_information - model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") - if model_intermediate_buffer is not None: - inputs["model_intermediate_buffer"] = model_intermediate_buffer - if cache_salt := parsed_content.get("cache_salt"): - inputs["cache_salt"] = cache_salt - - return inputs - - def _process_tokens( - self, - parsed_content: OmniTokensPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> OmniTokenInputs | MultiModalInput: - prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs) - prompt_embeds = parsed_content.get("prompt_embeds") - additional_information = parsed_content.get("additional_information") - model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") - - multi_modal_data = parsed_content.get("multi_modal_data") - - inputs: OmniTokenInputs | MultiModalInput - if multi_modal_data: - inputs = self._process_multimodal( - prompt_token_ids, - multi_modal_data, - parsed_content.get("mm_processor_kwargs"), - tokenization_kwargs=tokenization_kwargs, - mm_uuids=parsed_content.get("multi_modal_uuids"), - ) - else: - inputs = token_inputs_omni( - prompt_token_ids=prompt_token_ids, - prompt_embeds=prompt_embeds, - additional_information=additional_information, - model_intermediate_buffer=model_intermediate_buffer, - ) - if prompt_embeds is not None: - inputs["prompt_embeds"] = prompt_embeds - if additional_information is not None: - inputs["additional_information"] = additional_information - if model_intermediate_buffer is not None: - inputs["model_intermediate_buffer"] = model_intermediate_buffer - if prompt_text := parsed_content.get("prompt"): - inputs["prompt"] = prompt_text - if cache_salt := parsed_content.get("cache_salt"): - inputs["cache_salt"] = cache_salt - - return inputs - - def _process_embeds( - self, - parsed_content: OmniEmbedsPrompt, - ) -> EmbedsInput: - """Process embeddings prompt with omni-specific extensions. - - Extends base _process_embeds to handle additional_information payload - for direct transfer between pipeline stages. - """ - # Call parent implementation for base embeds processing - inputs = super()._process_embeds(parsed_content) - - # Add omni-specific additional_information if present - additional_information = parsed_content.get("additional_information") - if additional_information is not None: - inputs["additional_information"] = additional_information # type: ignore[typeddict-unknown-key] - model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") - if model_intermediate_buffer is not None: - inputs["model_intermediate_buffer"] = model_intermediate_buffer # type: ignore[typeddict-unknown-key] - - return inputs - - def _prompt_to_llm_inputs( - self, - prompt: SingletonDictPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - *, - mm_uuids: Any | None = None, - ) -> SingletonInput: - """ - Extract the singleton inputs from a prompt. - - Arguments: - - * prompt: single encoder or decoder input prompt - - Returns: - - * [`SingletonInput`][vllm.inputs.engine.SingletonInput] instance - """ - if "prompt_embeds" in prompt: - return self._process_embeds(prompt) # type: ignore[arg-type] - - if "prompt_token_ids" in prompt: - return self._process_tokens( - prompt, # type: ignore[arg-type] - ) - - if "prompt" in prompt: - return self._process_text( - prompt, # type: ignore[arg-type] - tokenization_kwargs=tokenization_kwargs, - mm_uuids=mm_uuids, - ) - - assert_never(prompt) # type: ignore[arg-type] + inputs = await self._renderer._process_singleton_async(prompt, skip_mm_cache=skip_mm_cache) + return self._with_omni_extras(inputs, prompt) diff --git a/vllm_omni/model_executor/models/audex/audex_omni_processing.py b/vllm_omni/model_executor/models/audex/audex_omni_processing.py index 52d2bec56..7b80f53f8 100644 --- a/vllm_omni/model_executor/models/audex/audex_omni_processing.py +++ b/vllm_omni/model_executor/models/audex/audex_omni_processing.py @@ -33,13 +33,14 @@ from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, PromptReplacement, PromptUpdate, PromptUpdateDetails, ) +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor + SOUND_TOKEN = "" SOUND_START_TOKEN = "" SOUND_END_TOKEN = "" @@ -176,7 +177,7 @@ class AudexDummyInputsBuilder(BaseDummyInputsBuilder[AudexProcessingInfo]): return {"audio": self._get_dummy_audios(length=audio_len, num_audios=num_audios, overrides=overrides)} -class AudexMultiModalProcessor(BaseMultiModalProcessor[AudexProcessingInfo]): +class AudexMultiModalProcessor(OmniMultiModalProcessor[AudexProcessingInfo]): def _call_hf_processor( self, prompt: str, @@ -222,22 +223,6 @@ class AudexMultiModalProcessor(BaseMultiModalProcessor[AudexProcessingInfo]): tensor_type="pt", ) - def _apply_hf_processor_main( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - processed_data = self._call_hf_processor( - self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()), - mm_data, - hf_processor_mm_kwargs, - {}, - ) - processed_data.update(passthrough_data) - return processed_data - def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index e8ad443e4..3c65f83aa 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -21,6 +21,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.bagel import BagelForConditionalGeneration from vllm.model_executor.models.interfaces import MultiModalEmbeddings from vllm.model_executor.models.qwen2 import Qwen2DecoderLayer, Qwen2MLP +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, @@ -53,7 +54,6 @@ from vllm_omni.diffusion.models.bagel.bagel_transformer import ( TimestepEmbedder, ) from vllm_omni.diffusion.models.bagel.pipeline_bagel import default_ae_params -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader class OmniBagelProcessor(BagelProcessor): @@ -1164,11 +1164,12 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration): filtered_weights.append((mapped_name, tensor)) loader = AutoWeightsLoader( - self, - skip_prefixes=["vit_pos_embed.pos_embed"], - ignore_unexpected_prefixes=["vae.", "latent_pos_embed.", "time_embedder.", "vae2llm."], + self, ignore_unexpected_prefixes=["vae.", "latent_pos_embed.", "time_embedder.", "vae2llm."] + ) + loaded = loader.load_weights( + filtered_weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"vit_pos_embed.pos_embed": None}), ) - loaded = loader.load_weights(filtered_weights, mapper=self.hf_to_vllm_mapper) loaded |= self._load_moe_gen_weights(moe_gen_weights) diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 8ec88b070..1c9ca47a6 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -25,7 +25,6 @@ from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptIndexTargets, @@ -39,6 +38,7 @@ from vllm.v1.sample.ops.topk_topp_sampler import random_sample from vllm.v1.sample.sampler import Sampler from vllm_omni.data_entry_keys import EmbeddingsStruct, OmniPayloadStruct, to_dict, to_struct +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer from vllm_omni.model_executor.models.cosyvoice3.utils import ( concat_text_with_prompt_ids, @@ -109,9 +109,7 @@ class CosyVoice3MultiModalProcessingInfo(BaseProcessingInfo): ) -class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModalProcessingInfo]): - _OMNI_PROMPT_TEXT_KEY = "_vllm_omni_original_prompt_text" - +class CosyVoice3MultiModalProcessor(OmniMultiModalProcessor[CosyVoice3MultiModalProcessingInfo]): def apply(self, inputs: ProcessorInputs, timing_ctx): tokenizer = self.info.get_tokenizer() prompt_text = tokenizer.decode(inputs.prompt, skip_special_tokens=False) @@ -405,24 +403,6 @@ class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModal return ft - def _apply_hf_processor_main( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - mm_kwargs = dict(hf_processor_mm_kwargs) - prompt_text = str(mm_kwargs.pop(self._OMNI_PROMPT_TEXT_KEY, "")) - valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - processed_data = self._call_hf_processor( - prompt_text, - mm_data, - mm_kwargs, - {}, - ) - processed_data.update(passthrough_data) - return processed_data - def _get_mm_fields_config( self, hf_inputs: "BatchFeature", diff --git a/vllm_omni/model_executor/models/covo_audio/covo_audio.py b/vllm_omni/model_executor/models/covo_audio/covo_audio.py index 4560b2a36..e907ed1b1 100644 --- a/vllm_omni/model_executor/models/covo_audio/covo_audio.py +++ b/vllm_omni/model_executor/models/covo_audio/covo_audio.py @@ -20,7 +20,6 @@ from vllm.multimodal.inputs import MultiModalFieldConfig from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptReplacement, @@ -29,6 +28,7 @@ from vllm.multimodal.processing import ( from vllm.sequence import IntermediateTensors from vllm.v1.sample.metadata import SamplingMetadata +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights @@ -109,7 +109,7 @@ class CovoAudioDummyInputsBuilder(BaseDummyInputsBuilder[CovoAudioProcessingInfo ) -class CovoAudioMultiModalProcessor(BaseMultiModalProcessor[CovoAudioProcessingInfo]): +class CovoAudioMultiModalProcessor(OmniMultiModalProcessor[CovoAudioProcessingInfo]): def _hf_processor_applies_updates(self, prompt_text, mm_items, hf_processor_mm_kwargs, tokenization_kwargs) -> bool: return False @@ -155,22 +155,6 @@ class CovoAudioMultiModalProcessor(BaseMultiModalProcessor[CovoAudioProcessingIn tensor_type=None, ) - def _apply_hf_processor_main( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - processed_data = self._call_hf_processor( - self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()), - mm_data, - hf_processor_mm_kwargs, - {}, - ) - processed_data.update(passthrough_data) - return processed_data - def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm_omni/model_executor/models/gepard/gepard_talker.py b/vllm_omni/model_executor/models/gepard/gepard_talker.py index 4b4c89cf9..37c1e66a9 100644 --- a/vllm_omni/model_executor/models/gepard/gepard_talker.py +++ b/vllm_omni/model_executor/models/gepard/gepard_talker.py @@ -22,13 +22,12 @@ import torch.nn as nn from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.qwen3_5 import Qwen3_5ForCausalLM -from vllm.model_executor.models.utils import maybe_prefix +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper, maybe_prefix from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig from vllm_omni.model_executor.models.gepard.nanocodec import NanoCodec from vllm_omni.model_executor.models.output_templates import OmniOutput -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -200,8 +199,10 @@ class GepardTalkerForConditionalGeneration(nn.Module): # NanoCodec loads lazily on first decode, so load_format=dummy and # NeMo-less environments never pay for it. - loader = AutoWeightsLoader(self, skip_prefixes=["mtp.", "ref_compressor."]) - return loader.load_weights(iter(rest)) + loader = AutoWeightsLoader(self) + return loader.load_weights( + iter(rest), mapper=WeightsMapper(orig_to_new_prefix={"mtp.": None, "ref_compressor.": None}) + ) def _get_or_create_state(self, request_id: str) -> _GepardState: st = self._active_states.get(request_id) diff --git a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py b/vllm_omni/model_executor/models/glm_image/glm_image_ar.py index 05ef9830a..fb8ce83b3 100644 --- a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py +++ b/vllm_omni/model_executor/models/glm_image/glm_image_ar.py @@ -82,7 +82,6 @@ from vllm.multimodal.inputs import ( MultiModalKwargsItems, ) from vllm.multimodal.parse import ( - ImageProcessorItems, MultiModalDataItems, MultiModalDataParser, ) @@ -360,282 +359,36 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo return super()._cached_apply_hf_processor(inputs, timing_ctx) def apply(self, inputs, timing_ctx): - """Ensure token prompts contain one source-image placeholder each.""" + """Add the generation scaffold before upstream expands source images.""" num_images = inputs.mm_data_items.get_all_counts().get("image", 0) + prompt_ids = list(inputs.prompt) if num_images: image_token_id = getattr( self.info.get_hf_config(), "image_token_id", 167855, ) - prompt_ids = list(inputs.prompt) missing = num_images - prompt_ids.count(image_token_id) if missing > 0: - inputs = replace( - inputs, - prompt=[image_token_id] * missing + prompt_ids, - ) - return super().apply(inputs, timing_ctx) + prompt_ids = [image_token_id] * missing + prompt_ids - def _call_hf_processor( - self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], - ) -> BatchFeature: - """ - Call the HuggingFace processor. - - For text-to-image mode (no images), we need to: - 1. Build the prompt with target grid dimensions - 2. Build the image_grid_thw tensor for M-RoPE position encoding - - For image-to-image mode: - 1. Process source images through the image processor - 2. Build prompt with image placeholders expanded - 3. Build image_grid_thw including source and target grids - """ processor = self.info.get_hf_processor() - - # Get target dimensions from mm_kwargs or use defaults - target_h = mm_kwargs.get("target_h", 1024) if mm_kwargs else 1024 - target_w = mm_kwargs.get("target_w", 1024) if mm_kwargs else 1024 - - logger.debug( - f"_call_hf_processor: target dimensions for generation: {target_h}x{target_w}, mm_kwargs={mm_kwargs}" - ) - - if not mm_data or not mm_data.get("images"): - # Text-to-image mode - if processor is not None: - # Build messages format expected by processor - messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}] - - # Use apply_chat_template which handles target dimensions - hf_inputs = processor.apply_chat_template( - messages, - tokenize=True, - target_h=target_h, - target_w=target_w, - return_dict=True, - return_tensors="pt", - ) - - return hf_inputs - else: - # Fallback: just tokenize (this won't work properly for generation) - tokenizer = self.info.get_tokenizer() - prompt_ids = tokenizer.encode(prompt) - return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt") - - # Image-to-image mode - # NOTE: Use "images" (plural) - this is what vLLM's ImageProcessorItems.get_processor_data() returns - images = mm_data.get("images") - if not isinstance(images, list): - images = [images] - - logger.debug( - f"_call_hf_processor i2i: num_images={len(images)}, image_types={[type(img).__name__ for img in images]}" - ) - - if processor is not None: - # Build messages with image objects directly in content - # This is how GlmImageProcessor expects images - embedded in the content dict - # NOT as a separate images= parameter - # - # IMPORTANT: Remove <|image|> placeholders from prompt since apply_chat_template - # will automatically insert them for each image in content. Having both leads to - # index out of bounds when processing image_grid_thw. - clean_prompt = prompt.replace("<|image|>", "") - content = [] - for img in images: - content.append({"type": "image", "image": img}) - content.append({"type": "text", "text": clean_prompt}) - messages = [{"role": "user", "content": content}] - - logger.debug(f"_call_hf_processor: calling apply_chat_template with {len(images)} images in content") - - # Use apply_chat_template - processor will process images when they're in content - hf_inputs = processor.apply_chat_template( - messages, - tokenize=True, - target_h=target_h, - target_w=target_w, - return_dict=True, - return_tensors="pt", - ) - - logger.debug(f"_call_hf_processor: apply_chat_template returned keys: {list(hf_inputs.keys())}") - - # IMPORTANT (i2i): vLLM multimodal encoder must see source-only grids - # (matching pixel_values and number of images), but M-RoPE needs full - # grids (source + target) to compute correct decode positions. - image_grid_thw = hf_inputs.get("image_grid_thw") - if image_grid_thw is not None: - # Preserve full grids for M-RoPE. - hf_inputs["mrope_image_grid_thw"] = image_grid_thw - - # Expose source-only grids for MM. - # In most i2i requests, we process one prompt at a time here, - # so `len(images)` is the number of source images. - num_source_images = len(images) - if image_grid_thw.shape[0] != num_source_images: - source_grids = image_grid_thw[:num_source_images] - hf_inputs["image_grid_thw"] = source_grids - logger.debug( - "_call_hf_processor: adjusted image_grid_thw for MM from %s to %s (num_source_images=%d)", - tuple(image_grid_thw.shape), - tuple(source_grids.shape), - num_source_images, - ) - - # Debug: Analyze input_ids for image tokens - input_ids = hf_inputs.get("input_ids") - if input_ids is not None: - if hasattr(input_ids, "tolist"): - ids_list = input_ids.tolist() - if isinstance(ids_list[0], list): - ids_list = ids_list[0] # Unbatch - else: - ids_list = list(input_ids) - - # Get image token ID from config - hf_config = self.info.get_hf_config() - image_token_id = getattr(hf_config, "image_token_id", 167855) - - # Count image tokens - image_token_count = ids_list.count(image_token_id) - logger.debug( - f"_call_hf_processor: input_ids length={len(ids_list)}, " - f"image_token_id={image_token_id}, " - f"image_token_count={image_token_count}" - ) - - # Log first/last few tokens to understand structure - logger.debug(f"_call_hf_processor: first 20 tokens: {ids_list[:20]}") - logger.debug(f"_call_hf_processor: last 20 tokens: {ids_list[-20:]}") - - # Find positions of image tokens - image_positions = [i for i, t in enumerate(ids_list) if t == image_token_id] - if image_positions: - logger.debug(f"_call_hf_processor: image token positions (first 10): {image_positions[:10]}") - - return hf_inputs - else: - # Fallback without processor - this is not ideal but prevents crashes - logger.warning("GlmImageProcessor not available, using fallback for i2i") - tokenizer = self.info.get_tokenizer() - hf_config = self.info.get_hf_config() - - # Get image token - image_token_id = getattr(hf_config, "image_token_id", 167855) - try: - image_token = tokenizer.convert_ids_to_tokens(image_token_id) - except Exception: - image_token = "<|image|>" - - # Build prompt with image placeholders - image_placeholders = image_token * len(images) - full_prompt = f"{image_placeholders}{prompt}" - prompt_ids = tokenizer.encode(full_prompt) - return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt") - - def _apply_hf_processor_mm_only( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> BatchFeature: - """ - Apply the HF processor on the multi-modal data only. - - GLM-Image requires special handling because apply_chat_template always - adds a target <|image|> placeholder in addition to source image placeholders. - This causes an IndexError when the HF processor tries to find grid info - for the target placeholder (which doesn't exist for source-only processing). - - Solution: Call the image processor directly to get pixel_values and - image_grid_thw, bypassing apply_chat_template's target handling. - """ - mm_counts = mm_items.get_all_counts() - num_images = mm_counts.get("image", 0) - - if num_images == 0: - # No images - call parent implementation - return super()._apply_hf_processor_mm_only( - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - - # For i2i mode, we need to process images directly with the image processor - # to avoid the apply_chat_template target placeholder issue - processor = self.info.get_hf_processor() - image_processor = processor.image_processor - - # Get images from mm_items - images = mm_items.get_items("image", ImageProcessorItems) - image_list = [images.get(i) for i in range(images.get_count())] - - logger.debug(f"_apply_hf_processor_mm_only: processing {len(image_list)} images directly") - - # Process images directly with image processor - image_inputs = image_processor( - images=image_list, - return_tensors="pt", - ) - - # Get grid info for source images only (no target) - pixel_values = image_inputs.get("pixel_values") - image_grid_thw = image_inputs.get("image_grid_thw") - if image_grid_thw is not None and image_grid_thw.shape[0] != num_images: - # Be defensive: some processors may include extra target grids. - image_grid_thw = image_grid_thw[:num_images] - image_inputs["image_grid_thw"] = image_grid_thw - - logger.debug( - f"_apply_hf_processor_mm_only: pixel_values shape=\ - {pixel_values.shape if pixel_values is not None else None}, " - f"image_grid_thw shape={image_grid_thw.shape if image_grid_thw is not None else None}" - ) - - # Build input_ids with image token placeholders - # The _get_prompt_updates returns PromptReplacement(target=[image_token_id], ...) - # which needs to find image tokens in input_ids to replace them. - # We need to include one image_token_id per image so the replacement can work. tokenizer = self.info.get_tokenizer() - image_token_id = tokenizer.convert_tokens_to_ids("<|image|>") - - # Build input_ids: [image_token] * num_images + tokenized text - # This way _apply_prompt_updates can find the image tokens and replace them - dummy_text = self.dummy_inputs.get_dummy_text(mm_counts) - text_ids = tokenizer.encode(dummy_text, add_special_tokens=False) - input_ids = [image_token_id] * num_images + text_ids - - logger.debug( - f"_apply_hf_processor_mm_only: built input_ids with {num_images} image tokens + {len(text_ids)} text tokens" + # Reuse HF's shape formatter without decoding/re-tokenizing the user's + # tokens. Text-to-image needs both target and preview grids; i2i needs + # only the target. Grid metadata alone does not supply these AR tokens. + target_grid = self._build_generation_grids(inputs.hf_processor_mm_kwargs)[0] + suffix, *_ = processor._build_prompt_with_target_shape( + "", + height=int(target_grid[1]) * 32, + width=int(target_grid[2]) * 32, + is_text_to_image=num_images == 0, ) - - return BatchFeature( - dict( - input_ids=[input_ids], - pixel_values=pixel_values, - image_grid_thw=image_grid_thw, - ), - tensor_type="pt", - ) - - def _apply_hf_processor_text_only( - self, prompt_text: str, hf_processor_mm_kwargs: Mapping[str, object], tokenization_kwargs: Mapping[str, object] - ) -> list[int]: - prompt_ids, _, _ = super()._apply_hf_processor_text_mm( - prompt_text=prompt_text, - mm_items=MultiModalDataItems({}), - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - return prompt_ids + suffix_ids = tokenizer.encode(suffix, add_special_tokens=False) + if prompt_ids[-len(suffix_ids) :] != suffix_ids: + prompt_ids.extend(suffix_ids) + inputs = replace(inputs, prompt=prompt_ids) + return super().apply(inputs, timing_ctx) def _build_generation_grids(self, hf_processor_mm_kwargs: Mapping[str, object]) -> torch.Tensor: """Build generation grids for M-RoPE decode positions. @@ -678,174 +431,6 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo dtype=torch.long, ) - def _apply_hf_processor_main( - self, - prompt: str | list[int], - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - *, - enable_hf_prompt_update: bool, - ) -> tuple[list[int], BatchFeature, bool]: - """ - Override to handle GLM-Image i2i mode correctly. - - Problem: When vLLM processes cached mm items (enable_hf_prompt_update=False), - the base implementation: - 1. Gets prompt_ids from _apply_hf_processor_text_only (no image tokens) - 2. Gets mm_data from _apply_hf_processor_mm_only - 3. Returns is_update_applied=False - - This causes _apply_prompt_updates to fail because prompt_ids has no image tokens. - - Solution: For i2i mode, we build prompt_ids that include image placeholders, - and return is_update_applied=False so _apply_prompt_updates can expand them. - """ - mm_counts = mm_items.get_all_counts() - num_images = mm_counts.get("image", 0) - - logger.debug(f"_apply_hf_processor_main: mm_counts={mm_counts}, num_images={num_images}") - - if num_images == 0 and isinstance(prompt, str): - # t2i mode or normal flow - use parent implementation - prompt_ids = self._apply_hf_processor_text_only( - prompt_text=prompt, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - mm_processed_data = self._apply_hf_processor_mm_only( - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - - # t2i has no source images, so mm features cannot provide image_grid_thw. - # Provide explicit generation grids for M-RoPE to avoid fallback token parsing - # (which can degrade high-resolution spatial positions, e.g. 1920x1920). - try: - mrope_grid_thw = self._build_generation_grids(hf_processor_mm_kwargs) - mm_processed_data["mrope_image_grid_thw"] = mrope_grid_thw - logger.debug( - "_apply_hf_processor_main t2i: mrope_image_grid_thw=%s", - mrope_grid_thw.tolist(), - ) - except Exception as e: - logger.warning("_apply_hf_processor_main t2i: failed to set mrope_image_grid_thw: %s", e) - - return prompt_ids, mm_processed_data, False - - # i2i mode: use unified HF processor path only. - # This avoids drift between duplicated manual/HF i2i implementations. - logger.debug( - "_apply_hf_processor_main: i2i mode (enable_hf_prompt_update=%s), num_images=%s", - enable_hf_prompt_update, - num_images, - ) - - if not isinstance(prompt, str): - # Online OpenAI chat preprocessing can arrive here with tokenized - # prompts (list[int]) before serving_chat replaces engine prompt - # with the clean text prompt. Do not fail the whole request. - logger.warning( - "_apply_hf_processor_main i2i: got tokenized prompt type=%s; " - "using compatibility path for preprocessing", - type(prompt).__name__, - ) - - prompt_ids = list(prompt) - mm_processed_data = self._apply_hf_processor_mm_only( - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - - # Preserve full grids for M-RoPE decode (source + target), while - # keeping image_grid_thw source-only for MM batching. - try: - source_grid_thw = mm_processed_data.get("image_grid_thw") - if source_grid_thw is not None and isinstance(source_grid_thw, torch.Tensor): - target_h = ( - hf_processor_mm_kwargs.get("target_h") - if isinstance(hf_processor_mm_kwargs.get("target_h"), int) - else None - ) - target_w = ( - hf_processor_mm_kwargs.get("target_w") - if isinstance(hf_processor_mm_kwargs.get("target_w"), int) - else None - ) - if target_h is None or target_w is None: - target_h = ( - hf_processor_mm_kwargs.get("height") - if isinstance(hf_processor_mm_kwargs.get("height"), int) - else 1024 - ) - target_w = ( - hf_processor_mm_kwargs.get("width") - if isinstance(hf_processor_mm_kwargs.get("width"), int) - else 1024 - ) - - factor = 32 - token_h = max(1, target_h // factor) - token_w = max(1, target_w // factor) - target_grid = torch.tensor([[1, token_h, token_w]], dtype=source_grid_thw.dtype) - mm_processed_data["mrope_image_grid_thw"] = torch.cat([source_grid_thw, target_grid], dim=0) - except Exception: - pass - - # Prompt updates will expand image placeholders in this compatibility path. - return prompt_ids, mm_processed_data, False - - images = mm_items.get_items("image", ImageProcessorItems) - image_list = [images.get(i) for i in range(images.get_count())] - if not image_list: - raise ValueError("GLM-Image i2i requires at least one source image in mm_items") - - hf_inputs = self._call_hf_processor( - prompt=prompt, - mm_data={"images": image_list}, - mm_kwargs=hf_processor_mm_kwargs, - tok_kwargs=tokenization_kwargs, - ) - - input_ids = hf_inputs.get("input_ids") - if input_ids is None: - raise ValueError("HF i2i processor returned no input_ids") - - if isinstance(input_ids, torch.Tensor): - prompt_ids = input_ids[0].tolist() if input_ids.dim() > 1 else input_ids.tolist() - else: - prompt_ids = ( - input_ids[0] - if isinstance(input_ids, list) and input_ids and isinstance(input_ids[0], list) - else list(input_ids) - ) - - mm_processed_data = BatchFeature(dict(), tensor_type="pt") - for key in ("pixel_values", "image_grid_thw", "mrope_image_grid_thw"): - value = hf_inputs.get(key) - if value is not None: - mm_processed_data[key] = value - - image_grid_thw = mm_processed_data.get("image_grid_thw") - mrope_grid_thw = mm_processed_data.get("mrope_image_grid_thw") - hf_config = self.info.get_hf_config() - image_token_id = getattr(hf_config, "image_token_id", 167855) - image_token_count = prompt_ids.count(image_token_id) - logger.debug( - "_apply_hf_processor_main i2i(HF): num_images=%s, prompt_len=%s, image_token_count=%s, " - "source_grid_shape=%s, mrope_grid_shape=%s", - num_images, - len(prompt_ids), - image_token_count, - tuple(image_grid_thw.shape) if image_grid_thw is not None else None, - tuple(mrope_grid_thw.shape) if mrope_grid_thw is not None else None, - ) - - # HF processor already expanded image placeholders in input_ids. - return prompt_ids, mm_processed_data, True - def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, @@ -891,7 +476,7 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo Get the multimodal field configuration. For GLM-Image i2i mode: - - image_grid_thw has been sliced in _call_hf_processor to only include source images + - image_grid_thw contains only source grids from _apply_hf_processor_main - pixel_values has shape [total_patches, C, H, W] - only for source images For t2i mode: @@ -899,7 +484,7 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo """ result = {} - # Get image_grid_thw if present (already sliced in _call_hf_processor) + # Get the source-only image_grid_thw if present. image_grid_thw = hf_inputs.get("image_grid_thw") if "pixel_values" in hf_inputs and image_grid_thw is not None: @@ -916,7 +501,7 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo result["pixel_values"] = MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes) - # Register image_grid_thw - it's been sliced in _call_hf_processor + # Register image_grid_thw - it's been sliced in _apply_hf_processor_main # to only include source image grids, so batching will work correctly result["image_grid_thw"] = MultiModalFieldConfig.batched("image") @@ -931,37 +516,6 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo return result - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - """ - Return whether the HF processor applies prompt updates. - - For GLM-Image i2i mode, the HF processor's apply_chat_template already - expands <|image|> to N tokens (e.g., 4096 for 64x64 grid). - - By returning True, we tell vLLM that HF processor DID apply prompt updates, - so vLLM will use _find_mm_placeholders to locate the expanded tokens - instead of trying to apply replacements. - - For t2i mode (no images), there are no image placeholders to expand. - """ - # Check if we have images (i2i mode) - num_images = mm_items.get_all_counts().get("image", 0) - if num_images > 0: - logger.debug( - f"_hf_processor_applies_updates: returning True for i2i mode " - f"(num_images={num_images}) - HF processor already expanded tokens" - ) - return True - - # For t2i mode (no images), use default behavior - return True - def _get_prompt_updates( self, mm_items: MultiModalDataItems, @@ -971,17 +525,8 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo """ Get prompt updates for image tokens. - For GLM-Image image-to-image mode, the HF processor's apply_chat_template - already expands each <|image|> placeholder to the correct number of - image tokens (grid_h * grid_w tokens per source image). - - The HF processor does: - 1. Replace each <|image|> with num_image_tokens copies of <|placeholder|> - 2. Replace all <|placeholder|> back to <|image|> - - So the tokenized input already has the expanded tokens. We use - target=[image_token_id] to match each occurrence of the image token, - similar to how Qwen2VL handles this pattern. + Expand each source-image placeholder into grid_h * grid_w image tokens + through vLLM's token-based prompt updates. We use image_grid_thw from out_mm_kwargs to get the actual processed grid size, following the Qwen2VL pattern. This is critical because the HF processor @@ -992,8 +537,7 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo """ hf_config = self.info.get_hf_config() - # Get image token ID - this is the token that appears multiple times - # in the tokenized input after HF processor expansion + # Match the source-image placeholder in the tokenized input. image_token_id = getattr(hf_config, "image_token_id", 167855) # Debug: log mm_items info diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts.py b/vllm_omni/model_executor/models/glm_tts/glm_tts.py index e57489b82..1a2da6581 100644 --- a/vllm_omni/model_executor/models/glm_tts/glm_tts.py +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts.py @@ -39,7 +39,6 @@ from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptIndexTargets, @@ -51,6 +50,7 @@ from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor from vllm_omni.model_executor.models.common.nucleus_ras_sampling import ras_sample_one as _ras_sample_one from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform @@ -364,7 +364,7 @@ class GLMTTSMultiModalProcessingInfo(BaseProcessingInfo): ) -class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessingInfo]): +class GLMTTSMultiModalProcessor(OmniMultiModalProcessor[GLMTTSMultiModalProcessingInfo]): """GLM-TTS voice-clone processor. Unlike CosyVoice3, GLM-TTS prompt speech tokens are normal Llama vocab IDs @@ -373,8 +373,6 @@ class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessi handoff. """ - _OMNI_PROMPT_TEXT_KEY = "_vllm_omni_original_prompt_text" - def apply(self, inputs: ProcessorInputs, timing_ctx): source_tokenizer = self.info.get_tokenizer() prompt_text = source_tokenizer.decode( @@ -630,24 +628,6 @@ class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessi } ) - def _apply_hf_processor_main( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - mm_kwargs = dict(hf_processor_mm_kwargs) - prompt_text = str(mm_kwargs.pop(self._OMNI_PROMPT_TEXT_KEY, "")) - valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - processed_data = self._call_hf_processor( - prompt_text, - mm_data, - mm_kwargs, - {}, - ) - processed_data.update(passthrough_data) - return processed_data - def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py index d4d903733..738e469f1 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py @@ -1669,8 +1669,10 @@ class _BackboneWrapper(nn.Module): self.config = config def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader + from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper skip = ["lm_head."] if getattr(self.config, "tie_word_embeddings", False) else None - loader = AutoWeightsLoader(self, skip_prefixes=skip) - return loader.load_weights(weights) + loader = AutoWeightsLoader(self) + return loader.load_weights( + weights, mapper=WeightsMapper(orig_to_new_prefix={name: None for name in (skip or ())}) + ) diff --git a/vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py b/vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py index 93e41955f..cd1546505 100644 --- a/vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py @@ -8,13 +8,12 @@ # subclasses the removed HunYuanModel / HunYuanSparseMoeBlock / HunYuanMLP, # so those internals are preserved here to keep the omni diffusion model # working against current upstream. +# Weight loading, expert mapping, and QKV splitting are owned by Omni's +# HunyuanModel subclass; their unused upstream implementations are omitted. """Inference-only HunYuan V1 internals vendored for vllm-omni HunyuanImage3.""" -import typing -from collections.abc import Callable, Iterable from itertools import islice -import regex as re import torch from torch import nn from transformers import PretrainedConfig @@ -29,7 +28,6 @@ from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( FusedMoEFactory, - fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -44,14 +42,9 @@ from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) from vllm.model_executor.models.interfaces import EagleModelMixin from vllm.model_executor.models.utils import ( PPMissingLayer, - is_pp_missing_parameter, make_layers, ) from vllm.sequence import IntermediateTensors @@ -630,201 +623,3 @@ class HunYuanModel(nn.Module, EagleModelMixin): if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states return hidden_states - - def _split_qkv_weight(self, qkv: torch.Tensor): - num_attention_heads = self.config.num_attention_heads - num_kv_heads = getattr(self.config, "num_key_value_heads", self.config.num_attention_heads) - num_key_value_groups = num_attention_heads // num_kv_heads - hidden_size = self.config.hidden_size - - if hasattr(self.config, "head_dim"): - attention_head_dim = self.config.head_dim - elif hasattr(self.config, "attention_head_dim"): - attention_head_dim = self.config.attention_head_dim - else: - attention_head_dim = self.config.hidden_size // num_attention_heads - - qkv = qkv.reshape(num_kv_heads, num_key_value_groups + 2, attention_head_dim, hidden_size) - q, k, v = torch.split(qkv, (num_key_value_groups, 1, 1), dim=1) - q = q.reshape(-1, hidden_size) - k = k.reshape(-1, hidden_size) - v = v.reshape(-1, hidden_size) - return torch.concat((q, k, v)) - - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: - if _is_moe(self.config): - # Params for weights, fp8 weight scales, fp8 activation scales - # (param_name, weight_name, expert_id, shard_id) - return fused_moe_make_expert_params_mapping( - self, - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.num_experts, - num_redundant_experts=self.num_redundant_experts, - ) - else: - return [] - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): - cla_factor = _get_cla_factor(self.config) - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".qkv_proj", ".q_proj", "q"), - (".qkv_proj", ".k_proj", "k"), - (".qkv_proj", ".v_proj", "v"), - (".gate_up_proj", ".gate_proj", 0), - (".gate_up_proj", ".up_proj", 1), - ] - - num_attention_heads = self.config.num_attention_heads - num_kv_heads = getattr(self.config, "num_key_value_heads", self.config.num_attention_heads) - split_params_mapping = [ - (".gate_up_proj", ".gate_and_up_proj", 2, [(1, 1), (0, 1)], None), - ( - ".qkv_proj", - ".qkv_proj", - num_attention_heads + num_kv_heads * 2, - [("q", num_attention_heads), ("k", num_kv_heads), ("v", num_kv_heads)], - self._split_qkv_weight, - ), - ] - - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - expert_params_mapping = self.get_expert_mapping() - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - if "gate_proj_bias" in name: - name = name.replace("gate_proj_bias", "gate_proj.bias") - if "up_proj_bias" in name: - name = name.replace("up_proj_bias", "up_proj.bias") - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: - # Models trained using ColossalAI may include these tensors in - # the checkpoint. Skip them. - continue - # With tie_word_embeddings, we can skip lm_head.weight - # The weight might appear unnecessarily in the files if the model is - # processed with quantization, LoRA, fine-tuning, etc. - if self.config.tie_word_embeddings and "lm_head.weight" in name: - continue - - is_found = False - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - if "mlp.experts" in name: - continue - # cross layer only have q_proj, skip qkv pack - if weight_name == ".q_proj": - match = re.search(r"layers\.\d+", name) - if match: - layer_id = int(match.group(0).split(".")[-1]) - if cla_factor > 1 and layer_id % cla_factor != 0: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - loaded_params.add(name) - is_found = True - break - if is_found: - continue - - for ( - param_name, - weight_name, - den, - split_param, - func, - ) in split_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - - if is_pp_missing_parameter(name, self): - continue - - assert loaded_weight.shape[0] % den == 0 - units = loaded_weight.shape[0] // den - - param = params_dict[name] - weight_loader = param.weight_loader - offset = 0 - for shard_id, num in split_param: - new_offset = offset + num * units - if func: - weight_loader(param, func(loaded_weight)[offset:new_offset], shard_id) - else: - weight_loader(param, loaded_weight[offset:new_offset], shard_id) - offset = new_offset - - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - is_expert_weight = False - for mapping in expert_params_mapping: - param_name, weight_name, expert_id, shard_id = mapping - if weight_name not in name: - continue - # this is an expert weight and should not be - # attempted to load as other weights later - is_expert_weight = True - - # Do not modify `name` since the loop may continue here - # Instead, create a new variable - name_mapped = name.replace(weight_name, param_name) - if is_pp_missing_parameter(name_mapped, self): - continue - param = params_dict[name_mapped] - # We should ask the weight loader to return success or not - # here since otherwise we may skip experts with other - # available replicas. - weight_loader = typing.cast(Callable[..., bool], param.weight_loader) - success = weight_loader( - param, - loaded_weight, - name_mapped, - shard_id=shard_id, - expert_id=expert_id, - return_success=True, - ) - if success: - name = name_mapped - break - else: - if is_expert_weight: - # We've checked that this is an expert weight - # However it's not mapped locally to this rank - # So we simply skip it - continue - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - - if is_pp_missing_parameter(name, self): - continue - - if "mlp.gate.wg." in name: - name = name.replace("wg.", "") - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index 82907c3ad..f330de97b 100755 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -54,7 +54,9 @@ from vllm.model_executor.models.interfaces import ( _require_is_multimodal, ) from vllm.model_executor.models.utils import ( + AutoWeightsLoader, PPMissingLayer, + WeightsMapper, _merge_multimodal_embeddings, is_pp_missing_parameter, maybe_prefix, @@ -70,7 +72,6 @@ from vllm.multimodal.parse import ( ) from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, PromptReplacement, PromptUpdate, @@ -83,6 +84,7 @@ from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor from vllm_omni.model_executor.models.hunyuan_image3._hunyuan_v1_vendored import ( HunYuanMLP, HunYuanModel, @@ -92,7 +94,6 @@ from vllm_omni.model_executor.models.hunyuan_image3._hunyuan_v1_vendored import ) from vllm_omni.model_executor.models.hunyuan_image3.autoencoder_kl_3d import AutoencoderKLConv3D from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import LightProjector, Siglip2VisionTransformer -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader logger = init_logger(__name__) @@ -1027,7 +1028,7 @@ class HunyuanImage3DummyInputsBuilder(BaseDummyInputsBuilder[HunyuanImage3Proces } -class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3ProcessingInfo]): +class HunyuanImage3MultiModalProcessor(OmniMultiModalProcessor[HunyuanImage3ProcessingInfo]): """Multimodal processor for HunyuanImage3 model.""" def _call_hf_processor( @@ -1046,22 +1047,6 @@ class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3Proc batch_feature["vae_generator_seed"] = torch.full((len(images),), int(vae_generator_seed), dtype=torch.long) return batch_feature - def _apply_hf_processor_main( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - processed_data = self._call_hf_processor( - self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()), - mm_data, - hf_processor_mm_kwargs, - {}, - ) - processed_data.update(passthrough_data) - return processed_data - def _hf_processor_applies_updates( self, prompt_text: str, @@ -2196,12 +2181,11 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo "timestep_r_emb", ] skip_prefixes.extend(unexpected_keywords) - loader = AutoWeightsLoader( - self, - skip_prefixes=skip_prefixes, - ) + loader = AutoWeightsLoader(self) - loaded_params = loader.load_weights(weights) + loaded_params = loader.load_weights( + weights, mapper=WeightsMapper(orig_to_new_prefix={name: None for name in (skip_prefixes or ())}) + ) return loaded_params def get_language_model(self) -> torch.nn.Module: diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py index 6e4592dec..4dbc9a455 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py @@ -30,7 +30,6 @@ from vllm.multimodal.inputs import ( from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptReplacement, @@ -45,6 +44,7 @@ from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import ( NO_INTERLEAVE_NEXT_TOKEN_ID, @@ -351,7 +351,7 @@ class MiMoAudioDataParser(MultiModalDataParser): return AudioProcessorItems(new_audios) -class MiMoAudioLLMMultiModalProcessor(BaseMultiModalProcessor[MiMoAudioLLMProcessingInfo]): +class MiMoAudioLLMMultiModalProcessor(OmniMultiModalProcessor[MiMoAudioLLMProcessingInfo]): def _call_hf_processor( self, prompt: str, @@ -398,22 +398,6 @@ class MiMoAudioLLMMultiModalProcessor(BaseMultiModalProcessor[MiMoAudioLLMProces tensor_type=None, ) - def _apply_hf_processor_main( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - processed_data = self._call_hf_processor( - self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()), - dict(mm_data), - hf_processor_mm_kwargs, - {}, - ) - processed_data.update(passthrough_data) - return processed_data - def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py index afc426ddc..d34cce465 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py @@ -22,6 +22,7 @@ from transformers import AutoTokenizer, Qwen2Config, Qwen2Model from transformers.utils.hub import cached_file from vllm.config import VllmConfig from vllm.logger import init_logger +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin @@ -30,7 +31,6 @@ from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE, AudioVAEConfig from vllm_omni.model_executor.models.ming_tts.constants import SPEAKER_EMBEDDING_DIM from vllm_omni.model_executor.models.output_templates import OmniOutput -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader from vllm_omni.transformers_utils.configs.ming_flash_omni import MingFlashOmniTalkerConfig from .prompt_utils import DEFAULT_PROMPT as MING_DEFAULT_PROMPT @@ -510,12 +510,13 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin) if self._standalone: weights = self._iter_talker_safetensors() - loader = AutoWeightsLoader( - self, - skip_prefixes=["audio_vae."], # loaded separately - skip_substrs=["rotary_embed.inv_freq"], # non-persistent buffer + loader = AutoWeightsLoader(self) + loaded = loader.load_weights( + weights, + mapper=WeightsMapper( + orig_to_new_prefix={"audio_vae.": None}, orig_to_new_substr={"rotary_embed.inv_freq": None} + ), ) - loaded = loader.load_weights(weights) logger.info("Loaded %d talker weights from checkpoint", len(loaded)) if self.audio_vae is not None and self._vae_weight_source is not None: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 88eddbf19..315ae159d 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -21,7 +21,6 @@ import inspect import math import os import warnings -from collections import defaultdict from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass, replace from functools import partial @@ -3596,29 +3595,15 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn *, out_keys: set[str], ) -> dict[str, NestedTensors]: - mm_kwargs = {k: v for k, v in mm_kwargs.items() if k != "use_tts"} - # This processor supports zipping prompt and mm_data together - if self.info.get_model_version() in {(2, 6), (4, 0), (4, 5)}: - inputs = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**mm_kwargs), - dict(text=prompts, **mm_data), - mm_kwargs, - ) - else: - inputs = defaultdict[str, list[torch.Tensor]](list) + from vllm.model_executor.models.minicpmv import MiniCPMVMultiModalProcessor - for i, prompt in enumerate(prompts): - inputs_one = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**mm_kwargs), - dict(text=prompt, **{key: value[i] for key, value in mm_data.items()}), - mm_kwargs, - ) - - for k, v in inputs_one.items(): - assert len(v) == 1, (k, len(v)) - inputs[k].append(v[0]) - - return {k: inputs[k] for k in out_keys} + return MiniCPMVMultiModalProcessor._call_hf_processor_on_prompts( + self, + prompts, + mm_data, + {key: value for key, value in mm_kwargs.items() if key != "use_tts"}, + out_keys=out_keys, + ) def _get_prompt_updates( self, diff --git a/vllm_omni/model_executor/models/minimax_music3/talker.py b/vllm_omni/model_executor/models/minimax_music3/talker.py index 5deda49c7..9baa7315c 100644 --- a/vllm_omni/model_executor/models/minimax_music3/talker.py +++ b/vllm_omni/model_executor/models/minimax_music3/talker.py @@ -37,10 +37,10 @@ from vllm.logger import init_logger from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.qwen3 import Qwen3Model +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm.v1.outputs import SamplerOutput from vllm_omni.model_executor.models.output_templates import OmniOutput -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader from .chunking import chunk_windows from .constants import ( @@ -980,8 +980,10 @@ class MiniMaxMusic3TalkerForConditionalGeneration(nn.Module): through vLLM. The audio embedding table and the depth decoder live in a separate component of the repo and are loaded explicitly. """ - loader = AutoWeightsLoader(self, skip_prefixes=["audio_embeddings.", "rvq_decoder."]) - loaded = loader.load_weights(weights) + loader = AutoWeightsLoader(self) + loaded = loader.load_weights( + weights, mapper=WeightsMapper(orig_to_new_prefix={"audio_embeddings.": None, "rvq_decoder.": None}) + ) loaded |= load_depth_decoder_weights( self.vllm_config, audio_embeddings=self.audio_embeddings, diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py index 2f95fed86..42a8e6ea8 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice.py @@ -29,7 +29,6 @@ from vllm.multimodal.inputs import ( from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptIndexTargets, @@ -38,6 +37,7 @@ from vllm.multimodal.processing import ( ) from vllm.sequence import IntermediateTensors +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig @@ -64,15 +64,13 @@ class OmniVoiceMultiModalProcessingInfo(BaseProcessingInfo): ) -class OmniVoiceMultiModalProcessor(BaseMultiModalProcessor[OmniVoiceMultiModalProcessingInfo]): +class OmniVoiceMultiModalProcessor(OmniMultiModalProcessor[OmniVoiceMultiModalProcessingInfo]): """Processes text + optional reference audio for OmniVoice. For voice cloning: text + reference audio → tokenized reference For auto voice: text only """ - _OMNI_PROMPT_TEXT_KEY = "_vllm_omni_original_prompt_text" - def apply(self, inputs: ProcessorInputs, timing_ctx): tokenizer = self.info.get_tokenizer() prompt_text = tokenizer.decode(inputs.prompt, skip_special_tokens=False) @@ -205,24 +203,6 @@ class OmniVoiceMultiModalProcessor(BaseMultiModalProcessor[OmniVoiceMultiModalPr ) return ft - def _apply_hf_processor_main( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - mm_kwargs = dict(hf_processor_mm_kwargs) - prompt_text = str(mm_kwargs.pop(self._OMNI_PROMPT_TEXT_KEY, "")) - valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - processed_data = self._call_hf_processor( - prompt_text, - mm_data, - mm_kwargs, - {}, - ) - processed_data.update(passthrough_data) - return processed_data - def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py index 0088fcbcb..b3e916f25 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py @@ -11,6 +11,7 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.interfaces import SupportsPP from vllm.model_executor.models.utils import ( + AutoWeightsLoader, WeightsMapper, init_vllm_registered_model, maybe_prefix, @@ -20,8 +21,6 @@ from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader - class Qwen2_5OmniTalkerForConditionalGeneration( nn.Module, @@ -143,11 +142,11 @@ class Qwen2_5OmniTalkerForConditionalGeneration( return self.language_model.sample(logits, sampling_metadata) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=["thinker.", "token2wav."], + loader = AutoWeightsLoader(self) + loaded = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"thinker.": None, "token2wav.": None}), ) - loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) # Log load summary try: total_bytes = 0 diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index f606f968e..6bbefceab 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py @@ -63,6 +63,7 @@ from vllm.model_executor.models.qwen2_audio import ( _get_feat_extract_output_lengths, ) from vllm.model_executor.models.utils import ( + AutoWeightsLoader, WeightsMapper, init_vllm_registered_model, maybe_prefix, @@ -87,7 +88,6 @@ from vllm.multimodal.processing.processor import ( from vllm.sequence import IntermediateTensors from vllm.utils.collection_utils import is_list_of -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader from vllm_omni.quantization.component_config import ( resolve_encoder_quant_config, ) @@ -1398,11 +1398,12 @@ class Qwen2_5OmniThinkerForConditionalGeneration( if self.visual is None: skip_prefixes.extend(["visual."]) - loader = AutoWeightsLoader( - self, - skip_prefixes=skip_prefixes, + loader = AutoWeightsLoader(self) + loaded_weights = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) + | WeightsMapper(orig_to_new_prefix={name: None for name in (skip_prefixes or ())}), ) - loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) return loaded_weights diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py index 24730f6ae..e9e92c904 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py @@ -21,7 +21,9 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead, from vllm.model_executor.model_loader.weight_utils import default_weight_loader, maybe_remap_kv_scale_name from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP from vllm.model_executor.models.utils import ( + AutoWeightsLoader, PPMissingLayer, + WeightsMapper, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, @@ -33,8 +35,6 @@ from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader - logger = init_logger(__name__) @@ -445,8 +445,8 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): return next_tokens def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), + loader = AutoWeightsLoader(self) + return loader.load_weights( + weights, + mapper=WeightsMapper(orig_to_new_prefix={"lm_head.": None} if self.config.tie_word_embeddings else {}), ) - return loader.load_weights(weights) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py index 0921b0902..67352eecb 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -23,6 +23,7 @@ from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( from vllm.config import VllmConfig # type: ignore from vllm.logger import init_logger # type: ignore from vllm.model_executor.models.utils import ( # type: ignore + AutoWeightsLoader, WeightsMapper, ) @@ -30,7 +31,6 @@ from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta from vllm_omni.model_executor.models.qwen3_omni.quantization import ( Qwen3OmniNestedSupportsQuant, ) -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader logger = init_logger(__name__) @@ -336,11 +336,11 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights from HuggingFace checkpoint.""" - loader = AutoWeightsLoader( - self, - skip_prefixes=["thinker.", "talker."], # Already loaded above + loader = AutoWeightsLoader(self) + loaded = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"thinker.": None, "talker.": None}), ) - loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) # Log load summary try: diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index c7c0db00f..ea3e2fcfe 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -15,6 +15,7 @@ from vllm.model_executor.models.interfaces import ( SupportsPP, ) from vllm.model_executor.models.utils import ( + AutoWeightsLoader, WeightsMapper, maybe_prefix, ) @@ -29,7 +30,6 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mt from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( Qwen3MoeLLMForCausalLM, ) -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader from vllm_omni.quantization.component_config import ComponentQuantizationConfig logger = init_logger(__name__) @@ -302,13 +302,12 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( to vLLM's internal structure. Code predictor weights are routed to its custom loader for vocab extension support. """ - loader = AutoWeightsLoader( - self, - skip_prefixes=["thinker.", "code2wav."], - # "code_predictor."], - ) + loader = AutoWeightsLoader(self) # Don't apply mapper again since we already did it - loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + loaded = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"thinker.": None, "code2wav.": None}), + ) # Log load summary try: diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index 8e0e7deb3..8329bbc91 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -32,7 +32,6 @@ import torch import torch.nn as nn import torch.nn.functional as F from transformers import PretrainedConfig -from transformers.feature_extraction_utils import BatchFeature from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( Qwen3OmniMoeConfig, Qwen3OmniMoeThinkerConfig, @@ -93,10 +92,14 @@ from vllm.model_executor.models.qwen3_omni_moe_thinker import ( from vllm.model_executor.models.qwen3_omni_moe_thinker import ( Qwen3OmniMoeAudioEncoderLayer as _Qwen3OmniMoeAudioEncoderLayer, ) +from vllm.model_executor.models.qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerMultiModalProcessor as _UpstreamQwen3Processor, +) from vllm.model_executor.models.qwen3_omni_moe_thinker import ( _get_feat_extract_output_lengths, ) from vllm.model_executor.models.utils import ( + AutoWeightsLoader, WeightsMapper, _merge_multimodal_embeddings, maybe_prefix, @@ -125,7 +128,6 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( from vllm_omni.model_executor.models.qwen3_omni.quantization import ( Qwen3OmniNestedSupportsQuant, ) -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader from vllm_omni.quantization.component_config import ( PRE_QUANTIZED_METHODS, ComponentQuantizationConfig, @@ -698,83 +700,10 @@ Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder class Qwen3OmniMoeThinkerMultiModalProcessor( Qwen2_5OmniThinkerMultiModalProcessor, + _UpstreamQwen3Processor, ): - def _get_hf_mm_data( - self, - mm_items: MultiModalDataItems, - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - processor_data, passthrough_data = super()._get_hf_mm_data(mm_items) - - audios = processor_data.get("audios") - if audios: - feature_extractor = self.info.get_feature_extractor() - hop_length = feature_extractor.hop_length - - def pad_to_hop_length(x: np.ndarray) -> np.ndarray: - length = x.shape[-1] - if length % hop_length != 0: - pad_length = hop_length - (length % hop_length) - x = np.pad( - x, - (0, pad_length), - mode="constant", - constant_values=0, - ) - return x - - # To make sure the cache works with padding=True, pre-pad audio - # to a multiple of the feature extractor's hop length. - processor_data = dict(processor_data) - processor_data["audios"] = [ - pad_to_hop_length(audio) if isinstance(audio, np.ndarray) else (pad_to_hop_length(audio[0]), audio[1]) - for audio in audios - ] - - return processor_data, passthrough_data - - def _apply_hf_processor_main( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> BatchFeature: - if mm_items.get_all_counts().get("audio", 0): - # TODO(Isotr0py): Remove this patch after upstream fix PR - # released and Transformers version update: - # https://github.com/huggingface/transformers/pull/41473 - hf_processor_mm_kwargs = dict(hf_processor_mm_kwargs) - hf_processor_mm_kwargs["audio_kwargs"] = dict(hf_processor_mm_kwargs.get("audio_kwargs") or {}) - hf_processor_mm_kwargs["text_kwargs"] = dict(hf_processor_mm_kwargs.get("text_kwargs") or {}) - - processed_data = super()._apply_hf_processor_main( - mm_items, - hf_processor_mm_kwargs, - ) - - if "audio_feature_lengths" in processed_data and "feature_attention_mask" in processed_data: - valid_mm_items = mm_items.select({k for k, c in mm_items.get_all_counts().items() if c > 0}) - processor_data, _ = self._get_hf_mm_data(valid_mm_items) - audios = processor_data.get("audios", []) - if audios: - feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs) - hop_length = feature_extractor.hop_length - - audio_num_frames = [] - for audio in audios: - audio_length = len(audio[0]) if isinstance(audio, tuple) else len(audio) - num_frame = ( - audio_length // hop_length if audio_length % hop_length == 0 else audio_length // hop_length - 1 - ) - if hf_processor_mm_kwargs.get("truncation", False): - num_frame = min( - num_frame, - feature_extractor.n_samples // hop_length, - ) - audio_num_frames.append(num_frame) - - processed_data["feature_attention_mask"] = [torch.ones(num_frame) for num_frame in audio_num_frames] - processed_data["audio_feature_lengths"] = torch.tensor(audio_num_frames) - - return processed_data + # Preserve Omni's per-video overrides while inheriting upstream Qwen3's + # audio processing through the shared, cooperative Qwen2.5 base. def _maybe_apply_prompt_updates( self, @@ -1518,11 +1447,11 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=["talker.", "code2wav."], + loader = AutoWeightsLoader(self) + loaded_weights = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"talker.": None, "code2wav.": None}), ) - loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) return loaded_weights diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index 5864948f5..f1692f73d 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py @@ -14,9 +14,9 @@ from vllm.config import VllmConfig from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.model_loader import DefaultModelLoader +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm_omni.model_executor.models.output_templates import OmniOutput -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import ( @@ -555,10 +555,9 @@ class Qwen3TTSCode2Wav(nn.Module): subfolder="speech_tokenizer", ) subfolder_weights = model_loader._get_weights_iterator(source) - loaded = AutoWeightsLoader( - self, - skip_prefixes=["encoder."], - ).load_weights(subfolder_weights) + loaded = AutoWeightsLoader(self).load_weights( + subfolder_weights, mapper=WeightsMapper(orig_to_new_prefix={"encoder.": None}) + ) device = self.vllm_config.device_config.device self.decoder.to(device=device, dtype=torch.float32) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index 4024e3536..dfb24516a 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -20,13 +20,12 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.model_executor.models.qwen3 import Qwen3Model -from vllm.model_executor.models.utils import PPMissingLayer, WeightsMapper, maybe_prefix +from vllm.model_executor.models.utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper, maybe_prefix from vllm.multimodal.audio import AudioResampler from vllm.sequence import IntermediateTensors from vllm_omni.data_entry_keys import OmniPayload from vllm_omni.model_executor.models.output_templates import OmniOutput -from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader from vllm_omni.utils.speaker_cache import ( get_speaker_cache, iter_custom_voice_profiles, @@ -1104,10 +1103,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): subfolder="speech_tokenizer", ) subfolder_weights = model_loader._get_weights_iterator(source) - enc_loaded = AutoWeightsLoader( - self, - skip_prefixes=["decoder."], - ).load_weights(subfolder_weights) + enc_loaded = AutoWeightsLoader(self).load_weights( + subfolder_weights, mapper=WeightsMapper(orig_to_new_prefix={"decoder.": None}) + ) loaded |= enc_loaded # AutoWeightsLoader only loads parameters; the encoder's VQ diff --git a/vllm_omni/model_executor/models/weight_loader.py b/vllm_omni/model_executor/models/weight_loader.py deleted file mode 100644 index a8e898a86..000000000 --- a/vllm_omni/model_executor/models/weight_loader.py +++ /dev/null @@ -1,54 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -"""Compatibility helpers for vLLM model weight loading.""" - -from collections.abc import Iterable - -import torch -from torch import nn -from vllm.model_executor.models.utils import ( - AutoWeightsLoader as _VllmAutoWeightsLoader, -) -from vllm.model_executor.models.utils import ( - WeightsMapper, -) - - -class AutoWeightsLoader(_VllmAutoWeightsLoader): - """Preserve vLLM's pre-0.29 weight-filtering arguments. - - vLLM removed ``skip_prefixes`` and ``skip_substrs`` from - :class:`AutoWeightsLoader` in favor of mapper-based filtering. Omni models - still use those arguments to exclude checkpoint components owned by a - different stage. Convert the filters to a ``WeightsMapper`` so the old - behavior is retained without bypassing the current loader implementation. - """ - - def __init__( - self, - module: nn.Module, - *, - skip_prefixes: list[str] | None = None, - skip_substrs: list[str] | None = None, - ignore_unexpected_prefixes: list[str] | None = None, - ignore_unexpected_suffixes: list[str] | None = None, - ) -> None: - super().__init__( - module, - ignore_unexpected_prefixes=ignore_unexpected_prefixes, - ignore_unexpected_suffixes=ignore_unexpected_suffixes, - ) - self._skip_mapper = WeightsMapper( - orig_to_new_prefix={prefix: None for prefix in skip_prefixes or ()}, - orig_to_new_substr={substring: None for substring in skip_substrs or ()}, - ) - - def load_weights( - self, - weights: Iterable[tuple[str, torch.Tensor]], - *, - mapper: WeightsMapper | None = None, - ) -> set[str]: - mapper = (mapper or WeightsMapper()) | self._skip_mapper - return super().load_weights(weights, mapper=mapper)