From 71125f80ad2f1ee473a7d46c81930e8c61681bb6 Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Wed, 23 Sep 2026 20:20:06 +0800 Subject: [PATCH] [New Model] Ming-Image (inclusionAI/Ming-Image-0.1-Design) Support (#8021) Signed-off-by: Yuanheng --- recipes/inclusionAI/Ming-Image.md | 141 ++++++ .../cache/test_teacache_extractors.py | 57 +++ .../z_image/test_zimage_tp_constraints.py | 41 +- .../ming_image/test_config_and_pipeline.py | 286 ++++++++++++ tests/model_tests/diffusion/test_alignment.py | 2 + vllm_omni/config/pipeline_registry.py | 2 + vllm_omni/deploy/ming_image.yaml | 46 ++ .../diffusion/cache/teacache/extractors.py | 1 + vllm_omni/diffusion/data.py | 58 ++- vllm_omni/diffusion/forward_context.py | 8 + .../ming_flash_omni/condition_encoder.py | 16 +- .../diffusion/models/ming_image/__init__.py | 4 + .../models/ming_image/component_config.py | 37 ++ .../diffusion/models/ming_image/condition.py | 88 ++++ .../diffusion/models/ming_image/pipeline.py | 425 ++++++++++++++++++ .../models/ming_image/transformer.py | 52 +++ .../models/z_image/pipeline_z_image.py | 17 +- .../models/z_image/z_image_transformer.py | 102 ++++- vllm_omni/diffusion/registry.py | 10 + .../modeling_bailing_moe_v2.py | 19 +- .../models/ming_image/__init__.py | 2 + .../models/ming_image/checkpoint.py | 45 ++ .../model_executor/models/ming_image/model.py | 223 +++++++++ .../models/ming_image/pipeline.py | 46 ++ vllm_omni/model_executor/models/registry.py | 5 + .../stage_input_processors/ming_image.py | 110 +++++ vllm_omni/model_extras/ming_image.py | 29 ++ vllm_omni/model_extras/registry.py | 19 + .../configs/ming_flash_omni.py | 25 +- .../transformers_utils/processors/__init__.py | 3 + .../transformers_utils/processors/ming.py | 84 +++- 31 files changed, 1952 insertions(+), 51 deletions(-) create mode 100644 recipes/inclusionAI/Ming-Image.md create mode 100644 tests/model_executor/models/ming_image/test_config_and_pipeline.py create mode 100644 vllm_omni/deploy/ming_image.yaml create mode 100644 vllm_omni/diffusion/models/ming_image/__init__.py create mode 100644 vllm_omni/diffusion/models/ming_image/component_config.py create mode 100644 vllm_omni/diffusion/models/ming_image/condition.py create mode 100644 vllm_omni/diffusion/models/ming_image/pipeline.py create mode 100644 vllm_omni/diffusion/models/ming_image/transformer.py create mode 100644 vllm_omni/model_executor/models/ming_image/__init__.py create mode 100644 vllm_omni/model_executor/models/ming_image/checkpoint.py create mode 100644 vllm_omni/model_executor/models/ming_image/model.py create mode 100644 vllm_omni/model_executor/models/ming_image/pipeline.py create mode 100644 vllm_omni/model_executor/stage_input_processors/ming_image.py create mode 100644 vllm_omni/model_extras/ming_image.py diff --git a/recipes/inclusionAI/Ming-Image.md b/recipes/inclusionAI/Ming-Image.md new file mode 100644 index 000000000..060c47570 --- /dev/null +++ b/recipes/inclusionAI/Ming-Image.md @@ -0,0 +1,141 @@ +# Ming-Image 0.1 Design + +> Text-to-image, image editing, and layer decomposition. + +## Summary + +- Vendor: inclusionAI +- Models: `inclusionAI/Ming-Image-0.1-Design` and `inclusionAI/Ming-Image-0.1-Design-Layer` +- Runtime: vLLM-Omni two-stage online serving +- API: OpenAI-compatible chat completions + +## Architecture + +Stage 0 uses a Qwen2.5-VL vision tower plus a 20-layer BailingMoeV2 language model. +It appends 256 learned image-query tokens and exports the final query states plus direct VLM states from layers 5, 12, and 20. +The query-token checkpoint is in the root sibling `mlp/` directory, so remote stage-0 materialization downloads both `mllm/` and `mlp/`. + +Stage 1 projects those conditions through the checkpoint connector and runs a 30-layer Z-Image DiT with a Qwen-Image RGBA VAE. + +For Design-Layer, the first returned image is the reconstructed composite and the remaining images are the requested layers. + +## CUDA + +### Environment + +- OS: Linux +- Python: 3.10+ +- CUDA: 13.0 +- vLLM version: 0.29.0 +- vLLM-Omni version or commit: 59db6a428 +- 2x H100 80GB (1xH100 to be validated) + +### Commands + +```bash +MODEL=inclusionAI/Ming-Image-0.1-Design +vllm serve "$MODEL" --omni --deploy-config vllm_omni/deploy/ming_image.yaml --port 8091 +``` + +For layer decomposition, set `MODEL` to `inclusionAI/Ming-Image-0.1-Design-Layer`. + +## Text-to-image + +Note that a prompt refiner is expected to describe the prompts with details; we will refine with more example inputs soon. + +```bash +curl -s http://127.0.0.1:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "inclusionAI/Ming-Image-0.1-Design", + "messages": [{"role": "user", "content": "A clean editorial botanical poster"}], + "modalities": ["image"], + "extra_body": { + "height": 1024, "width": 1024, + "steps": 12, "cfg": 1.0, "seed": 42 + } + }' \ + | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \ + | base64 -d > ming_design_smoke.png +``` + +## Image editing + +Pass one input image and an editing instruction: + +```bash +MODEL=inclusionAI/Ming-Image-0.1-Design +INPUT_IMAGE=/path/to/input.png + +jq -n \ + --arg model "$MODEL" \ + --rawfile image <(base64 -w0 "$INPUT_IMAGE") \ + '{ + model: $model, + messages: [{role: "user", content: [ + {type: "image_url", image_url: {url: ("data:image/png;base64," + $image)}}, + {type: "text", text: "Change the background to blue"} + ]}], + modalities: ["image"], + extra_body: { + height: 1024, width: 1024, + steps: 12, cfg: 1.0, seed: 42 + } + }' | +curl -sS http://127.0.0.1:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + --data-binary @- \ + | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \ + | base64 -d > ming_edit.png +``` + +## Layer decomposition + +Use the Design-Layer checkpoint and set `INPUT_IMAGE` to a local flattened design image. +Note that the prompt should better depict each layer to be decomposed, we will refine with more example inputs soon. + +```bash +MODEL=inclusionAI/Ming-Image-0.1-Design-Layer +INPUT_IMAGE=/path/to/input.png +PROMPT="Decompose this design into editable visual layers" + +jq -n \ + --arg model "$MODEL" \ + --rawfile image <(base64 -w0 "$INPUT_IMAGE") \ + --arg prompt "$PROMPT" \ + '{ + model: $model, + messages: [{role: "user", content: [ + {type: "image_url", image_url: {url: ("data:image/png;base64," + $image)}}, + {type: "text", text: $prompt} + ]}], + modalities: ["image"], + extra_body: { + num_layers: 6, + height: 1024, width: 1024, + steps: 12, cfg: 2.0, seed: 42 + } + }' | +curl -sS http://127.0.0.1:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + --data-binary @- > response.json + +jq -r '.choices[0].message.content[].image_url.url | split(",")[1]' response.json | + nl -v 0 | + while read -r index data; do + printf "%s" "$data" | base64 -d > "ming_layer_${index}.png" + done +``` + +## Notes + +- The default deployment keeps Stage 0 eager and enables dynamic regional + compilation with CUDA Graph Trees for the repeated Stage 1 DiT blocks. + The first request for each new image shape or layer count pays compilation + and graph-capture cost; warm up every production shape before measuring or + serving latency-sensitive traffic. +- Only one reference image and one request at a time are currently supported. +- Design-Layer requires a reference image except during warmup. +- A non-empty `negative_prompt` is rejected; Ming-Image uses zero negative conditioning. +- Height and width must be divisible by 16. +- Returned images retain the checkpoint's four RGBA channels. diff --git a/tests/diffusion/cache/test_teacache_extractors.py b/tests/diffusion/cache/test_teacache_extractors.py index 55b92b820..7dd613f5f 100644 --- a/tests/diffusion/cache/test_teacache_extractors.py +++ b/tests/diffusion/cache/test_teacache_extractors.py @@ -33,6 +33,7 @@ from vllm_omni.diffusion.cache.teacache.extractors import ( extract_flux2_klein_context, extract_flux_context, extract_minimax_h3_context, + extract_zimage_context, ) from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformer2DModel from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import ( @@ -53,6 +54,62 @@ def setup_tp_group(): yield +@pytest.mark.cpu +def test_zimage_extractor_accepts_extended_patchify_output(): + class _Block(nn.Module): + def __init__(self): + super().__init__() + self.attention_norm1 = nn.Identity() + + def adaLN_modulation(self, value): + return value.new_zeros((value.shape[0], 16)) + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([_Block()]) + self.noise_refiner = nn.ModuleList() + self.context_refiner = nn.ModuleList() + self.all_x_embedder = nn.ModuleDict({"2-1": nn.Linear(2, 4)}) + self.cap_embedder = nn.Linear(3, 4) + self.x_pad_token = nn.Parameter(torch.zeros(1, 4)) + self.cap_pad_token = nn.Parameter(torch.zeros(1, 4)) + self.t_scale = 1.0 + + @staticmethod + def t_embedder(timestep): + return timestep.new_ones((timestep.shape[0], 4)) + + @staticmethod + def rope_embedder(position_ids): + rope = torch.zeros((position_ids.shape[0], 2)) + return rope, rope + + @staticmethod + def patchify_and_embed(*_args): + positions = [torch.zeros((32, 3), dtype=torch.int32)] + masks = [torch.zeros(32, dtype=torch.bool)] + return ( + [torch.zeros((32, 2))], + [torch.zeros((32, 3))], + [(1, 1, 1)], + positions, + positions, + masks, + masks, + [], + ) + + context = extract_zimage_context( + _Model(), + x=[torch.zeros((1, 1, 1, 1))], + t=torch.zeros(1), + cap_feats=[torch.zeros((1, 3))], + ) + + assert context.modulated_input.shape == (1, 64, 4) + + class BaseExtractorTest(ABC): """Base class for testing TeaCache extractors. diff --git a/tests/diffusion/models/z_image/test_zimage_tp_constraints.py b/tests/diffusion/models/z_image/test_zimage_tp_constraints.py index b87d51c36..5e529d4fd 100644 --- a/tests/diffusion/models/z_image/test_zimage_tp_constraints.py +++ b/tests/diffusion/models/z_image/test_zimage_tp_constraints.py @@ -1,10 +1,47 @@ -import pytest +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -from vllm_omni.diffusion.models.z_image.z_image_transformer import validate_zimage_tp_constraints +import pytest +import torch + +from vllm_omni.diffusion.models.z_image.z_image_transformer import ( + UnifiedPrepare, + validate_zimage_tp_constraints, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_unified_prepare_applies_padding_masks_before_sharding(): + x = torch.zeros((2, 4, 2)) + cap = torch.zeros((2, 3, 2)) + x_rope = torch.zeros((2, 4, 1)) + cap_rope = torch.zeros((2, 3, 1)) + x_mask = torch.tensor([[1, 0, 1, 0], [1, 1, 0, 0]], dtype=torch.bool) + cap_mask = torch.tensor([[1, 1, 0], [1, 0, 1]], dtype=torch.bool) + + *_, unified_mask = UnifiedPrepare()( + x, + x_rope, + x_rope, + cap, + cap_rope, + cap_rope, + [4, 2], + [3, 3], + x_mask, + cap_mask, + ) + + assert torch.equal( + unified_mask, + torch.tensor( + [[1, 0, 1, 0, 1, 1, 0], [1, 1, 1, 0, 1, 0, 0]], + dtype=torch.bool, + ), + ) + + def test_validate_zimage_tp_constraints_tp2_ok(): ffn_hidden_dim, final_out_dims, supported_tp = validate_zimage_tp_constraints( dim=3840, diff --git a/tests/model_executor/models/ming_image/test_config_and_pipeline.py b/tests/model_executor/models/ming_image/test_config_and_pipeline.py new file mode 100644 index 000000000..b1350e7e8 --- /dev/null +++ b/tests/model_executor/models/ming_image/test_config_and_pipeline.py @@ -0,0 +1,286 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.models.ming_image.pipeline import ( + MingImageDiffusionPipeline, + _validate_variant_config, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_executor.models.ming_flash_omni.ming_flash_omni_thinker import ( + MingFlashOmniThinkerForConditionalGeneration, +) +from vllm_omni.model_executor.models.ming_image import checkpoint +from vllm_omni.model_executor.models.ming_image.model import MingImageForConditionalGeneration +from vllm_omni.model_executor.models.ming_image.pipeline import MING_IMAGE_PIPELINE +from vllm_omni.model_executor.stage_input_processors.ming_image import thinker2image +from vllm_omni.model_extras import ( + get_extra_body_params, + should_init_extra_args_for_non_diffusion_stages, +) +from vllm_omni.transformers_utils.configs.ming_flash_omni import BailingMM2Config + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_config_selects_qwen25_vision_tower(): + config = BailingMM2Config(vision_config={"model_type": "qwen2_5_vit"}, llm_config={}) + + assert type(config.vision_config).__name__ == "Qwen2_5_VLVisionConfig" + + +@pytest.mark.parametrize( + ("transformer_config", "layered"), + ( + ( + SimpleNamespace( + _class_name="DiffusionTransformer", + alignment_padding_mode="zero_masked", + multi_frame_output=False, + ), + False, + ), + ( + SimpleNamespace( + _class_name="DiffusionTransformer", + alignment_padding_mode="learned", + multi_frame_output=True, + ), + True, + ), + ), +) +def test_variant_config_without_model_index(transformer_config, layered): + assert _validate_variant_config(None, transformer_config) is layered + + +def test_variant_config_rejects_mixed_semantics(): + with pytest.raises(ValueError, match="must use either"): + _validate_variant_config( + {"_class_name": "MingImageDiffusionPipeline"}, + SimpleNamespace( + _class_name="DiffusionTransformer", + alignment_padding_mode="learned", + multi_frame_output=False, + ), + ) + + +def test_layer_latent_frames_flatten_once_in_frame_major_order(): + latents = torch.arange(2 * 4 * 3 * 2 * 2).reshape(2, 4, 3, 2, 2) + + pipeline = MingImageDiffusionPipeline.__new__(MingImageDiffusionPipeline) + flat = pipeline._flatten_latent_frames(latents) + + assert flat.shape == (6, 4, 2, 2) + assert torch.equal(flat[0], latents[0, :, 0]) + assert torch.equal(flat[1], latents[1, :, 0]) + assert torch.equal(flat[2], latents[0, :, 1]) + + +def test_layer_pipeline_allows_missing_reference_only_for_dummy_run(): + pipeline = MingImageDiffusionPipeline.__new__(MingImageDiffusionPipeline) + pipeline.is_layer_decomposition = True + pipeline._configure_output_frames( + reference=None, + num_layers=1, + is_dummy_run=True, + ) + assert pipeline._num_frames_per_prompt == 2 + + with pytest.raises(ValueError, match="requires a reference image"): + pipeline._configure_output_frames( + reference=None, + num_layers=1, + is_dummy_run=False, + ) + + +def test_pipeline_rejects_multiple_outputs_per_prompt(): + pipeline = MingImageDiffusionPipeline.__new__(MingImageDiffusionPipeline) + request = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt="test", + sampling_params=OmniDiffusionSamplingParams(num_outputs_per_prompt=2), + request_id="test-request", + ) + ] + ) + + with pytest.raises(ValueError, match="num_outputs_per_prompt=1 only, got 2"): + pipeline.forward(request) + + +def _source_output(prefix_len: int = 4): + prompt_ids = [11] * prefix_len + [157158] + [157157] * 256 + [157159] + length = len(prompt_ids) + multimodal = { + "final_hidden_states": torch.arange(length * 2048).view(length, 2048), + } + for layer in (5, 12, 20): + multimodal[f"hidden_states_{layer}"] = torch.full((length, 2048), layer) + return SimpleNamespace( + prompt_token_ids=prompt_ids, + outputs=[SimpleNamespace(multimodal_output=multimodal)], + ) + + +def test_bridge_extracts_query_and_direct_conditions(): + result = thinker2image( + [_source_output()], + prompt={"multi_modal_data": {"img2img": "reference"}}, + sampling_params=SimpleNamespace(extra_args={"num_layers": 3}), + ) + + extra = result[0]["extra"] + assert extra["query_hidden_states"].shape == (256, 2048) + assert extra["direct_hidden_states"].shape == (4, 6144) + assert extra["reference_image"] == "reference" + assert extra["num_layers"] == 3 + assert set(extra["direct_hidden_states"][:, :2048].unique().tolist()) == {5} + + +def test_bridge_rejects_text_negative_condition(): + with pytest.raises(ValueError, match="does not accept negative_prompt"): + thinker2image( + [_source_output()], + prompt={"negative_prompt": "bad"}, + sampling_params=SimpleNamespace(extra_args={}), + ) + + +def test_two_stage_topology_and_request_metadata(): + assert MING_IMAGE_PIPELINE.model_type == "ming_image" + assert [stage.model_stage for stage in MING_IMAGE_PIPELINE.stages] == ["mllm", "dit"] + assert MING_IMAGE_PIPELINE.diffusers_class_name == "MingImageDiffusionPipeline" + assert MING_IMAGE_PIPELINE.diffusers_class_aliases == ("MingImageLayeredDiffusionPipeline",) + assert MING_IMAGE_PIPELINE.stages[0].model_subdir == "mllm" + assert MING_IMAGE_PIPELINE.stages[0].model_path_resolver.endswith(".resolve_ming_image_model_root") + assert MING_IMAGE_PIPELINE.stages[1].model_arch == "MingImageDiffusionPipeline" + assert MING_IMAGE_PIPELINE.hf_architectures == ("MingImageForConditionalGeneration",) + for class_name in MING_IMAGE_PIPELINE.diffusers_class_aliases + (MING_IMAGE_PIPELINE.diffusers_class_name,): + assert "num_layers" in get_extra_body_params(class_name) + assert should_init_extra_args_for_non_diffusion_stages(class_name) + + +def test_stage1_compile_uses_configured_dynamic_regional_cuda_graph(monkeypatch): + pipeline = MingImageDiffusionPipeline.__new__(MingImageDiffusionPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + diffusion_compile_granularity="regional", + diffusion_compile_dynamic=True, + ) + pipeline.transformer = torch.nn.Identity() + captured = {} + + def _regionally_compile(model, **kwargs): + captured.update(kwargs) + return model + + monkeypatch.setattr( + "vllm_omni.diffusion.models.ming_image.pipeline.regionally_compile", + _regionally_compile, + ) + + pipeline.setup_compile() + + assert captured == { + "mode": "reduce-overhead", + "fullgraph": True, + "dynamic": True, + } + assert pipeline._uses_cudagraph_trees + + +def test_explicit_pipeline_loads_component_config_without_model_index(monkeypatch): + transformer_config = { + "_class_name": "DiffusionTransformer", + "alignment_padding_mode": "zero_masked", + "multi_frame_output": False, + } + + def _get_hf_file_to_dict(filename, model, revision=None): + del model, revision + return transformer_config if filename == "transformer/config.json" else None + + monkeypatch.setattr( + "vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + "vllm.transformers_utils.config.get_hf_file_to_dict", + _get_hf_file_to_dict, + ) + + config = OmniDiffusionConfig( + model="org/model", + model_class_name="MingImageDiffusionPipeline", + ) + config.enrich_config() + + assert config.model_class_name == "MingImageDiffusionPipeline" + assert config.tf_model_config.alignment_padding_mode == "zero_masked" + assert config.tf_model_config.multi_frame_output is False + + +def test_checkpoint_resolver_downloads_mllm_and_sibling_mlp(monkeypatch): + captured = {} + + def _download(**kwargs): + captured.update(kwargs) + return "/cache/model" + + monkeypatch.setattr(checkpoint, "download_weights_from_hf_specific", _download) + + assert checkpoint.resolve_ming_image_model_root("org/model", "rev", None) == "/cache/model" + assert captured["allow_patterns"] == ["mllm/**", "mlp/**"] + assert captured["revision"] == "rev" + assert captured["require_all"] is True + + +def test_compute_logits_accepts_vllm_v1_signature(): + class _LanguageModel(torch.nn.Module): + def compute_logits(self, hidden_states, sampling_metadata): + assert sampling_metadata is None + return hidden_states + 1 + + model = MingImageForConditionalGeneration.__new__(MingImageForConditionalGeneration) + torch.nn.Module.__init__(model) + model.language_model = _LanguageModel() + + hidden_states = torch.zeros(2, 3) + assert torch.equal(model.compute_logits(hidden_states), hidden_states + 1) + + +def test_embedding_only_run_uses_empty_modality_masks(): + model = MingImageForConditionalGeneration.__new__(MingImageForConditionalGeneration) + torch.nn.Module.__init__(model) + image_mask, audio_mask = model._compute_modality_masks(None, torch.empty(7, 2048)) + + assert image_mask.shape == (7,) + assert not image_mask.any() + assert torch.equal(audio_mask, image_mask) + + +def test_mllm_reports_root_mlp_query_tokens_as_loaded(monkeypatch): + model = MingImageForConditionalGeneration.__new__(MingImageForConditionalGeneration) + torch.nn.Module.__init__(model) + model.query_tokens_dict = torch.nn.ParameterDict({"16x16": torch.nn.Parameter(torch.empty(1), requires_grad=False)}) + monkeypatch.setattr( + MingFlashOmniThinkerForConditionalGeneration, + "load_weights", + lambda self, weights: {"language_model.weight"}, + ) + + assert model.load_weights([]) == { + "language_model.weight", + "query_tokens_dict.16x16", + } diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index f7f1855cc..6c39d702c 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -53,6 +53,8 @@ EXCLUDED_MODELS = [ "BooguImageTurboPipeline", "LancePipeline", "MingImagePipeline", + "MingImageDiffusionPipeline", + "MingImageLayeredDiffusionPipeline", "InternVLAA1Pipeline", "StableDiffusion3Pipeline", "HunyuanImage3ForCausalMM", diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 4876bd0e1..997a6991f 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -86,6 +86,7 @@ from vllm_omni.model_executor.models.ming_flash_omni.pipeline import ( MING_FLASH_OMNI_THINKER_ONLY_PIPELINE, MING_FLASH_OMNI_TTS_PIPELINE, ) +from vllm_omni.model_executor.models.ming_image.pipeline import MING_IMAGE_PIPELINE from vllm_omni.model_executor.models.ming_tts.pipeline import ( MING_TTS_MOE_PIPELINE, MING_TTS_PIPELINE, @@ -187,6 +188,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "ming_flash_omni_tts": MING_FLASH_OMNI_TTS_PIPELINE, "ming_flash_omni_thinker_only": MING_FLASH_OMNI_THINKER_ONLY_PIPELINE, "ming_flash_omni_image": MING_FLASH_OMNI_IMAGE_PIPELINE, + "ming_image": MING_IMAGE_PIPELINE, "moss_tts_nano": MOSS_TTS_NANO_PIPELINE, "minimax_h3_disaggregated": MINIMAX_H3_PIPELINE, "omnivoice": OMNIVOICE_PIPELINE, diff --git a/vllm_omni/deploy/ming_image.yaml b/vllm_omni/deploy/ming_image.yaml new file mode 100644 index 000000000..15709bd51 --- /dev/null +++ b/vllm_omni/deploy/ming_image.yaml @@ -0,0 +1,46 @@ +# Ming-Image Design and Design-Layer two-stage deployment. +pipeline: ming_image +async_chunk: false + +trust_remote_code: false +distributed_executor_backend: mp +dtype: bfloat16 +enable_prefix_caching: false + +stages: + - stage_id: 0 + enforce_eager: true + tensor_parallel_size: 1 + devices: "0" + gpu_memory_utilization: 0.55 + max_num_seqs: 1 + output_connectors: + to_stage_1: shared_memory_connector + default_sampling_params: + max_tokens: 1 + temperature: 0 + seed: 42 + + - stage_id: 1 + enforce_eager: false + diffusion_compile_granularity: regional + diffusion_compile_dynamic: true + tensor_parallel_size: 1 + devices: "1" + max_num_seqs: 1 + input_connectors: + from_stage_0: shared_memory_connector + default_sampling_params: + seed: 42 + +connectors: + shared_memory_connector: + name: SharedMemoryConnector + +edges: + - from: 0 + to: 1 + window_size: -1 + max_inflight: 1 + +platforms: {} diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index a7539b34a..d1b04151d 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -444,6 +444,7 @@ def extract_zimage_context( cap_pos_ids, x_inner_pad_mask, cap_inner_pad_mask, + _cap_feats_2, ) = module.patchify_and_embed(x, cap_feats, patch_size, f_patch_size) # Process image patches through embedder and noise refiner diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index eb025a617..3180dae67 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1487,6 +1487,29 @@ class OmniDiffusionConfig: self.max_multimodal_image_inputs = metadata.max_multimodal_image_inputs self.supports_mixed_reference_inputs = metadata.supports_mixed_reference_inputs + def _load_component_transformer_config(self) -> bool: + """Load a registered pipeline's component-level DiT configuration.""" + from vllm.transformers_utils.config import get_hf_file_to_dict + + from vllm_omni.model_extras import get_transformer_config_subfolder + + transformer_subfolder = get_transformer_config_subfolder( + self.model_class_name, + model=self.model, + revision=self.revision, + ) + tf_config_dict = get_hf_file_to_dict( + f"{transformer_subfolder}/config.json", + self.model, + revision=self.revision, + ) + if tf_config_dict is None: + tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model, revision=self.revision) + if tf_config_dict is None: + return False + self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict)) + return True + @staticmethod def _looks_like_lance_subfolder(model: str | None) -> bool: """Return True when ``--model`` points at a Lance per-component subfolder. @@ -1553,23 +1576,7 @@ class OmniDiffusionConfig: exc, ) else: - from vllm_omni.model_extras import get_transformer_config_subfolder - - transformer_subfolder = get_transformer_config_subfolder( - self.model_class_name, - model=self.model, - revision=self.revision, - ) - tf_config_dict = get_hf_file_to_dict( - f"{transformer_subfolder}/config.json", - self.model, - revision=self.revision, - ) - if tf_config_dict is None: - tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model, revision=self.revision) - if tf_config_dict is not None: - self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict)) - else: + if not self._load_component_transformer_config(): self.set_tf_model_config(TransformerConfig()) else: raise FileNotFoundError("Diffusers pipeline index not found") @@ -1605,6 +1612,23 @@ class OmniDiffusionConfig: self.set_tf_model_config(TransformerConfig()) self.update_multimodal_support() return + # An explicit topology or CLI override can select a native, + # registered pipeline whose repository contains component + # configs but no root HF config or Diffusers index. Trust + # that explicit class selection only when its component + # transformer config is present. + from vllm_omni.diffusion.registry import DiffusionModelRegistry + + if ( + self.model_class_name in DiffusionModelRegistry.get_supported_archs() + and self._load_component_transformer_config() + ): + logger.info( + "Using explicitly selected diffusion pipeline %r with component-level configuration.", + self.model_class_name, + ) + self.update_multimodal_support() + return raise ValueError(f"Could not find config.json or a Diffusers pipeline index for {self.model}") self.set_tf_model_config(TransformerConfig.from_dict(cfg)) diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index 496f384d7..009c851ca 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/vllm_omni/diffusion/forward_context.py @@ -48,6 +48,8 @@ class ForwardContext: total_denoise_steps: int | None = None # Per-request reference latent for img2img DiT models (e.g. Ming) ref_latent: torch.Tensor | None = None + # Per-request projected direct-VLM condition (e.g., Ming-Image). For now for bsz 1. + direct_condition: torch.Tensor | None = None # whether to split the text embed in sequence parallel, if True, the text embed will be split in sequence parallel # Sequence Parallel padding support @@ -377,3 +379,9 @@ def set_forward_context_ref_latent(ref_latent: torch.Tensor | None) -> None: """ if _forward_context is not None: _forward_context.ref_latent = ref_latent + + +def set_forward_context_direct_condition(direct_condition: torch.Tensor | None) -> None: + """Set the projected direct-VLM condition on the active context.""" + if _forward_context is not None: + _forward_context.direct_condition = direct_condition diff --git a/vllm_omni/diffusion/models/ming_flash_omni/condition_encoder.py b/vllm_omni/diffusion/models/ming_flash_omni/condition_encoder.py index 6ebca4f5c..6c762a346 100644 --- a/vllm_omni/diffusion/models/ming_flash_omni/condition_encoder.py +++ b/vllm_omni/diffusion/models/ming_flash_omni/condition_encoder.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The vLLM-Omni team. # # Adapted from Ming repository (inclusionAI/Ming) — the ``get_condition_embeds @@ -65,12 +66,16 @@ class MingConditionEncoder(nn.Module): thinker_hidden_size: int = 4096, device: torch.device | str | None = None, dtype: torch.dtype | None = None, + normalize_output: bool = True, + strict_loading: bool = False, ) -> None: super().__init__() self.config = image_gen_config self.thinker_hidden_size = thinker_hidden_size self._target_device = torch.device(device) if device is not None else None self._target_dtype = dtype + self.normalize_output = normalize_output + self.strict_loading = strict_loading # Populated lazily by ``load_from_checkpoint`` to keep this module # cheap to construct (useful for dummy-init paths and unit tests). @@ -195,6 +200,8 @@ class MingConditionEncoder(nn.Module): condition path to be meaningful. """ if not mlp_path.exists(): + if self.strict_loading: + raise FileNotFoundError(f"Missing Ming MLP directory: {mlp_path}") logger.warning( "[MingConditionEncoder] mlp/ subfolder missing at %s — proj/norm " "will stay randomly initialized. EXPECT BAD IMAGES until this " @@ -206,6 +213,8 @@ class MingConditionEncoder(nn.Module): try: from safetensors.torch import load_file # type: ignore except ImportError: + if self.strict_loading: + raise logger.exception("[MingConditionEncoder] safetensors not installed") return @@ -213,6 +222,8 @@ class MingConditionEncoder(nn.Module): if not candidates: candidates = sorted(mlp_path.glob("*.bin")) if not candidates: + if self.strict_loading: + raise FileNotFoundError(f"No Ming MLP weights under {mlp_path}") logger.warning("[MingConditionEncoder] no weight files under %s", mlp_path) return @@ -258,6 +269,8 @@ class MingConditionEncoder(nn.Module): ok_ow = _copy(self.proj_out.weight, "proj_out.weight") ok_ob = _copy(self.proj_out.bias, "proj_out.bias") if not (ok_w and ok_b and ok_ow and ok_ob): + if self.strict_loading: + raise ValueError(f"Ming condition projection weights are incomplete or incompatible under {mlp_path}.") logger.error( "[MingConditionEncoder] proj_in/proj_out NOT fully loaded; diffusion conditioning will be garbage." ) @@ -342,7 +355,8 @@ class MingConditionEncoder(nn.Module): # L2 normalize + ×1000 rescale. Ming's ``get_condition_embeds_for_image_gen`` # applies F.normalize; ``ZImageModel_withMLP`` then rescales by 1000 # when ``text_encoder_norm=True``. We fold both into one place. - cap_feats = torch.nn.functional.normalize(cap_feats, dim=-1) + if self.normalize_output: + cap_feats = torch.nn.functional.normalize(cap_feats, dim=-1) if self.config.text_encoder_norm: cap_feats = cap_feats * 1000.0 diff --git a/vllm_omni/diffusion/models/ming_image/__init__.py b/vllm_omni/diffusion/models/ming_image/__init__.py new file mode 100644 index 000000000..2423ec1fb --- /dev/null +++ b/vllm_omni/diffusion/models/ming_image/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Ming-Image specific diffusion components.""" diff --git a/vllm_omni/diffusion/models/ming_image/component_config.py b/vllm_omni/diffusion/models/ming_image/component_config.py new file mode 100644 index 000000000..25b02daf6 --- /dev/null +++ b/vllm_omni/diffusion/models/ming_image/component_config.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Validation for Ming-Image component configuration.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + + +def load_mlp_config(model_path: str | Path) -> dict[str, Any]: + path = Path(model_path) / "mlp" / "config.json" + if not path.is_file(): + raise FileNotFoundError(f"Ming-Image checkpoint is missing MLP config: {path}") + data = json.loads(path.read_text()) + required = { + "diffusion_c_input_dim", + "diffusion_inner_dim", + "img_gen_scales", + "selected_hidden_states_layers", + "use_identity_mlp", + "use_learnable_token_condition", + "use_vlm_directvlm_condition", + } + missing = sorted(required - data.keys()) + if missing: + raise ValueError(f"Ming-Image MLP config is missing keys: {missing}") + if data["img_gen_scales"] != [16]: + raise ValueError(f"Only img_gen_scales=[16] is supported, got {data['img_gen_scales']!r}.") + if data["selected_hidden_states_layers"] != [5, 12, 20]: + raise ValueError("Ming-Image requires selected hidden layers [5, 12, 20].") + if not data["use_identity_mlp"]: + raise NotImplementedError("Ming-Image non-identity diffusion MLP is not supported.") + if not data["use_learnable_token_condition"] or not data["use_vlm_directvlm_condition"]: + raise ValueError("Ming-Image requires both query-token and direct-VLM conditioning.") + return data diff --git a/vllm_omni/diffusion/models/ming_image/condition.py b/vllm_omni/diffusion/models/ming_image/condition.py new file mode 100644 index 000000000..bf9dd54e8 --- /dev/null +++ b/vllm_omni/diffusion/models/ming_image/condition.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Ming-Image query and direct-VLM condition projectors.""" + +from __future__ import annotations + +from pathlib import Path + +import torch +import torch.nn as nn +from safetensors.torch import load_file + +from vllm_omni.diffusion.models.ming_flash_omni.condition_encoder import MingConditionEncoder +from vllm_omni.transformers_utils.configs.ming_flash_omni import MingImageGenConfig + +from .component_config import load_mlp_config + + +class MingImageConditioning(nn.Module): + def __init__( + self, + model_path: str | Path, + *, + device: torch.device | str, + dtype: torch.dtype, + ) -> None: + super().__init__() + model_path = Path(model_path) + mlp_config = load_mlp_config(model_path) + config = MingImageGenConfig( + diffusion_c_input_dim=int(mlp_config["diffusion_c_input_dim"]), + thinker_hidden_size=2048, + img_gen_scales=list(mlp_config["img_gen_scales"]), + text_encoder_norm=bool(mlp_config.get("text_encoder_norm", False)), + ) + self.query_encoder = MingConditionEncoder( + config, + thinker_hidden_size=2048, + device=device, + dtype=dtype, + normalize_output=bool(mlp_config.get("connector_norm", False)), + strict_loading=True, + ) + self.query_encoder.load_from_checkpoint(model_path) + + selected_layers = mlp_config["selected_hidden_states_layers"] + direct_input_dim = 2048 * len(selected_layers) + direct_output_dim = int(mlp_config["diffusion_inner_dim"]) + self.direct_projector = nn.Sequential( + nn.RMSNorm(direct_input_dim, eps=1e-5), + nn.Linear(direct_input_dim, direct_output_dim, bias=True), + ).to(device=device, dtype=dtype) + self._load_direct_weights(model_path / "mlp" / "model.safetensors") + + def _load_direct_weights(self, path: Path) -> None: + if not path.is_file(): + raise FileNotFoundError(f"Missing Ming-Image direct projector weights: {path}") + state = load_file(str(path)) + mapping = { + "0.weight": "proj_directvlm.0.weight", + "1.weight": "proj_directvlm.1.weight", + "1.bias": "proj_directvlm.1.bias", + } + own_state = self.direct_projector.state_dict() + projected: dict[str, torch.Tensor] = {} + for target, source in mapping.items(): + if source not in state: + raise ValueError(f"Missing Ming-Image MLP weight {source!r}.") + if own_state[target].shape != state[source].shape: + raise ValueError( + f"Ming-Image MLP shape mismatch for {source}: " + f"expected {tuple(own_state[target].shape)}, got {tuple(state[source].shape)}." + ) + projected[target] = state[source] + self.direct_projector.load_state_dict(projected, strict=True) + + def forward( + self, + query_hidden_states: torch.Tensor, + direct_hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return ( + self.query_encoder(query_hidden_states), + self.direct_projector(direct_hidden_states), + ) + + +__all__ = ["MingImageConditioning"] diff --git a/vllm_omni/diffusion/models/ming_image/pipeline.py b/vllm_omni/diffusion/models/ming_image/pipeline.py new file mode 100644 index 000000000..872de839d --- /dev/null +++ b/vllm_omni/diffusion/models/ming_image/pipeline.py @@ -0,0 +1,425 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Diffusion pipeline for Ming-Image Design and Design-Layer checkpoints.""" + +from __future__ import annotations + +import logging +import os +from collections.abc import Iterable +from typing import Any, ClassVar + +import torch +import torch.nn as nn +from diffusers.image_processor import VaeImageProcessor +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from PIL import Image +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.compile import regionally_compile +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import ( + DistributedAutoencoderKLQwenImage, +) +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.forward_context import ( + set_forward_context_direct_condition, + set_forward_context_ref_latent, +) +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch +from vllm_omni.diffusion.models.ming_image.condition import MingImageConditioning +from vllm_omni.diffusion.models.ming_image.transformer import MingImageTransformer2DModel +from vllm_omni.diffusion.models.z_image.pipeline_z_image import ZImagePipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index +from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific + +logger = logging.getLogger(__name__) + +_DESIGN_PIPELINE = "MingImageDiffusionPipeline" +_LAYERED_PIPELINE = "MingImageLayeredDiffusionPipeline" +_VENDOR_TRANSFORMER_CLASS = "DiffusionTransformer" +_DIFFUSION_REQUIRED_PATTERNS = [ + "model_index.json", + "scheduler/**", + "transformer/**", + "vae/**", + "mlp/**", + "connector/**", +] + + +def _validate_variant_config( + model_index: dict[str, Any] | None, + transformer_config: Any, +) -> bool: + """Validate checkpoint-owned variant metadata and return layered mode.""" + declared_transformer = getattr(transformer_config, "_class_name", None) + if declared_transformer != _VENDOR_TRANSFORMER_CLASS: + raise ValueError( + "Ming-Image transformer/config.json must preserve the vendor " + f"_class_name={_VENDOR_TRANSFORMER_CLASS!r}, got {declared_transformer!r}. " + "vLLM-Omni loads those weights through MingImageTransformer2DModel." + ) + + alignment_padding_mode = getattr(transformer_config, "alignment_padding_mode", None) + multi_frame_output = getattr(transformer_config, "multi_frame_output", None) + variant = (alignment_padding_mode, multi_frame_output) + if variant == ("zero_masked", False): + is_layer_decomposition = False + elif variant == ("learned", True): + is_layer_decomposition = True + else: + raise ValueError( + "Ming-Image transformer/config.json must use either " + "alignment_padding_mode='zero_masked' with multi_frame_output=False " + "or alignment_padding_mode='learned' with multi_frame_output=True; " + f"got alignment_padding_mode={alignment_padding_mode!r}, " + f"multi_frame_output={multi_frame_output!r}." + ) + + declared_pipeline = (model_index or {}).get("_class_name") + if declared_pipeline is not None: + if declared_pipeline not in {_DESIGN_PIPELINE, _LAYERED_PIPELINE}: + raise ValueError( + "Ming-Image model_index.json must declare " + f"{_DESIGN_PIPELINE!r} or {_LAYERED_PIPELINE!r}, got {declared_pipeline!r}." + ) + expected_pipeline = _LAYERED_PIPELINE if is_layer_decomposition else _DESIGN_PIPELINE + if declared_pipeline != expected_pipeline: + raise ValueError( + f"Ming-Image model_index.json declares {declared_pipeline!r}, but " + f"transformer/config.json describes {expected_pipeline!r}." + ) + return is_layer_decomposition + + +class MingImageDiffusionPipeline(ZImagePipeline): + """Ming-Image component adapter around the canonical Z-Image loop.""" + + supports_request_batch = False + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["conditioning"] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None: + del prefix + nn.Module.__init__(self) + + model_path = od_config.model + if not os.path.exists(model_path): + model_path = download_weights_from_hf_specific( + model_name_or_path=model_path, + cache_dir=None, + allow_patterns=_DIFFUSION_REQUIRED_PATTERNS, + revision=od_config.revision, + require_all=True, + ) + local_files_only = os.path.isdir(model_path) + dtype = od_config.dtype + + self.od_config = od_config + self._execution_device = get_local_device() + self.device = self._execution_device + model_index = get_diffusion_model_index(model_path, revision=od_config.revision) or {} + transformer_config = od_config.tf_model_config + self.is_layer_decomposition = _validate_variant_config( + model_index, + transformer_config, + ) + + self.default_num_inference_steps = 12 + self.default_guidance_scale = 2.0 if self.is_layer_decomposition else 1.0 + self._num_frames_per_prompt = 1 + self._pending_prompt_embeds: list[torch.Tensor] | None = None + self._pending_negative_prompt_embeds: list[torch.Tensor] | None = None + self._uses_cudagraph_trees = False + + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=model_path, + subfolder="transformer", + revision=od_config.revision, + prefix="transformer.", + fall_back_to_pt=True, + ) + ] + subfolders = ["scheduler", "transformer", "vae"] + + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model_path, + subfolder="scheduler", + local_files_only=local_files_only, + ) + + self.vae = from_pretrained_with_prefetch( + DistributedAutoencoderKLQwenImage.from_pretrained, + model_path, + subfolder="vae", + prefetch_list=subfolders, + local_files_only=local_files_only, + torch_dtype=dtype, + ).to(self.device) + self.vae.eval() + vae_channels = int(getattr(self.vae.config, "input_channels", 3)) + if vae_channels != 4: + raise ValueError(f"Ming-Image requires an RGBA VAE, got input_channels={vae_channels}.") + + transformer_kwargs = get_transformer_config_kwargs( + od_config.tf_model_config, + MingImageTransformer2DModel, + ) + self.transformer = MingImageTransformer2DModel( + quant_config=od_config.quantization_config, + **transformer_kwargs, + ) + self.conditioning = MingImageConditioning( + model_path, + device=self.device, + dtype=dtype, + ) + self.text_encoder = None + self.tokenizer = None + + self.vae_scale_factor = 2 ** len(self.vae.config.temperal_downsample) + self.image_processor = VaeImageProcessor( + vae_scale_factor=self.vae_scale_factor * 2, + do_convert_rgb=False, + ) + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler + ) + + def setup_compile(self) -> None: + # Keep request preparation, scheduling, and VAE work eager, + # while capturing repeated DiT blocks for CUDAGraph Trees replay. + if self.od_config.diffusion_compile_granularity != "regional": + logger.warning( + "Ming-Image CUDA Graph uses regional DiT compilation; diffusion_compile_granularity=%r is ignored.", + self.od_config.diffusion_compile_granularity, + ) + self.transformer = regionally_compile( + self.transformer, + mode="reduce-overhead", + fullgraph=True, + dynamic=self.od_config.diffusion_compile_dynamic, + ) + self._uses_cudagraph_trees = True + + def encode_prompt(self, *args, **kwargs): + del args, kwargs + if self._pending_prompt_embeds is None or self._pending_negative_prompt_embeds is None: + raise RuntimeError("Ming-Image conditioning is only available during forward.") + return self._pending_prompt_embeds, self._pending_negative_prompt_embeds + + def prepare_latents(self, batch_size, *args, **kwargs): + frames = self._num_frames_per_prompt + flat = super().prepare_latents(batch_size * frames, *args, **kwargs) + return torch.stack(flat.chunk(frames, dim=0), dim=2) + + def _encode_reference( + self, + reference: Any | None, + height: int, + width: int, + ) -> torch.Tensor | None: + if reference is None: + return None + if isinstance(reference, list): + if len(reference) != 1: + raise ValueError("Ming-Image currently accepts exactly one reference image.") + reference = reference[0] + if isinstance(reference, Image.Image): + reference = reference.convert("RGBA") + if not isinstance(reference, torch.Tensor): + reference = self.image_processor.preprocess(reference, height=height, width=width) + if reference.ndim != 4 or reference.shape[1] != 4: + raise ValueError("Ming-Image reference must be a batched RGBA tensor or an RGBA-compatible image.") + reference = reference.to(device=self.device, dtype=self.vae.dtype).unsqueeze(2) + latent = self.vae.encode(reference).latent_dist.mode() + return (latent - self.vae.config.shift_factor) * self.vae.config.scaling_factor + + @staticmethod + def _get_prompt_extra(req: DiffusionRequestBatch) -> dict[str, Any]: + if not req.prompts: + return {} + prompt = req.prompts[0] + if isinstance(prompt, dict): + return dict(prompt.get("extra") or {}) + if hasattr(prompt, "_asdict"): + return dict(prompt._asdict().get("extra") or {}) + return {} + + def _configure_output_frames( + self, + *, + reference: Any | None, + num_layers: int, + is_dummy_run: bool, + ) -> None: + if num_layers < 1: + raise ValueError("num_layers must be at least 1.") + if self.is_layer_decomposition: + if reference is None and not is_dummy_run: + raise ValueError("Ming-Image Design-Layer requires a reference image.") + self._num_frames_per_prompt = num_layers + 1 + return + if num_layers != 1: + raise ValueError("Ming-Image Design supports exactly one output frame.") + self._num_frames_per_prompt = 1 + + @torch.inference_mode() + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + sampling = req.sampling_params + if sampling.num_outputs_per_prompt != 1: + # TODO(yuanheng-zhao): enable after supporting batching + raise ValueError( + f"Ming-Image currently supports num_outputs_per_prompt=1 only, got {sampling.num_outputs_per_prompt}." + ) + + extra = self._get_prompt_extra(req) + extra_args = sampling.extra_args or {} + + is_dummy_run = req.is_dummy_run() + query_hidden = extra.get("query_hidden_states") + direct_hidden = extra.get("direct_hidden_states") + if query_hidden is None or direct_hidden is None: + if not is_dummy_run: + raise ValueError("Ming-Image requests require query and direct conditions.") + logger.warning("Ming-Image conditions are absent during warmup; using zero tensors.") + query_hidden = torch.zeros((256, 2048), device=self.device, dtype=self.od_config.dtype) + direct_hidden = torch.zeros((1, 6144), device=self.device, dtype=self.od_config.dtype) + if not isinstance(query_hidden, torch.Tensor) or not isinstance(direct_hidden, torch.Tensor): + raise TypeError("Ming-Image query and direct conditions must be tensors.") + if query_hidden.ndim == 2: + query_hidden = query_hidden.unsqueeze(0) + if direct_hidden.ndim == 2: + direct_hidden = direct_hidden.unsqueeze(0) + query_hidden = query_hidden.to(device=self.device, dtype=self.od_config.dtype) + direct_hidden = direct_hidden.to(device=self.device, dtype=self.od_config.dtype) + cap_feats, direct_condition = self.conditioning(query_hidden, direct_hidden) + + reference = extra.get("reference_image") + num_layers = int(extra_args.get("num_layers", extra.get("num_layers", 1))) + self._configure_output_frames( + reference=reference, + num_layers=num_layers, + is_dummy_run=is_dummy_run, + ) + + height = int(extra_args.get("height") or sampling.height or 1024) + width = int(extra_args.get("width") or sampling.width or 1024) + steps = int(extra_args.get("steps") or sampling.num_inference_steps or self.default_num_inference_steps) + cfg = float( + extra_args["cfg"] + if extra_args.get("cfg") is not None + else sampling.guidance_scale + if sampling.guidance_scale is not None + else self.default_guidance_scale + ) + seed = extra_args.get("seed", sampling.seed) + generator = ( + torch.Generator(device=self.device).manual_seed(int(seed)) if seed is not None else sampling.generator + ) + + ref_latent = self._encode_reference(reference, height, width) + positive = [item for item in cap_feats] + negative = [torch.zeros_like(item) for item in positive] + self._pending_prompt_embeds = positive + self._pending_negative_prompt_embeds = negative + + apply_cfg = cfg > 0 + context_direct = ( + torch.cat([direct_condition, torch.zeros_like(direct_condition)], dim=0) if apply_cfg else direct_condition + ) + context_ref = ref_latent + if apply_cfg and context_ref is not None: + context_ref = context_ref.repeat(2, 1, 1, 1, 1) + + inner_sampling = OmniDiffusionSamplingParams( + height=height, + width=width, + num_inference_steps=steps, + guidance_scale=cfg, + generator=generator, + output_type="latent", + ) + inner_req = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt={"prompt": ""}, + sampling_params=inner_sampling, + request_id=req.request_id or "ming-image", + ) + ] + ) + + set_forward_context_ref_latent(context_ref) + set_forward_context_direct_condition(context_direct) + try: + latent_output = super().forward(inner_req) + if not isinstance(latent_output.output, torch.Tensor): + raise TypeError("Ming-Image denoising must return latent tensors.") + image = self._decode_latent_frames(latent_output.output) + return DiffusionOutput( + output=image, + stage_durations=latent_output.stage_durations, + ) + finally: + set_forward_context_ref_latent(None) + set_forward_context_direct_condition(None) + self._pending_prompt_embeds = None + self._pending_negative_prompt_embeds = None + self._num_frames_per_prompt = 1 + + @staticmethod + def _flatten_latent_frames(latents: torch.Tensor) -> torch.Tensor: + if latents.ndim == 4: + return latents + if latents.ndim != 5: + raise ValueError(f"Expected 4D or 5D Ming-Image latents, got {tuple(latents.shape)}") + batch, channels, frames, height, width = latents.shape + return latents.permute(2, 0, 1, 3, 4).reshape( + frames * batch, + channels, + height, + width, + ) + + def _decode_latent_frames(self, latents: torch.Tensor) -> torch.Tensor: + latents = self._flatten_latent_frames(latents).to(self.vae.dtype).unsqueeze(2) + latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor + image = self.vae.decode(latents, return_dict=False)[0] + if image.ndim == 5: + if image.shape[2] != 1: + raise ValueError( + f"Ming-Image VAE returned multiple decoded frames for a single latent frame: {tuple(image.shape)}" + ) + image = image.squeeze(2) + return image + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loaded = AutoWeightsLoader(self).load_weights(weights) + loaded |= {f"vae.{name}" for name, _ in self.vae.named_parameters()} + loaded |= {f"conditioning.{name}" for name, _ in self.conditioning.named_parameters()} + return loaded + + +def get_ming_image_post_process_func(od_config: OmniDiffusionConfig): + del od_config + image_processor = VaeImageProcessor(vae_scale_factor=16, do_convert_rgb=False) + + def post_process(images: torch.Tensor): + if images.ndim == 5: + images = images.permute(0, 2, 1, 3, 4).flatten(0, 1) + return image_processor.postprocess(images.float()) + + return post_process + + +__all__ = ["MingImageDiffusionPipeline", "get_ming_image_post_process_func"] diff --git a/vllm_omni/diffusion/models/ming_image/transformer.py b/vllm_omni/diffusion/models/ming_image/transformer.py new file mode 100644 index 000000000..8dc2e183a --- /dev/null +++ b/vllm_omni/diffusion/models/ming_image/transformer.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Ming-Image adapter for the shared Z-Image transformer.""" + +from __future__ import annotations + +import torch + +from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available +from vllm_omni.diffusion.models.z_image.z_image_transformer import ZImageTransformer2DModel + + +class MingImageTransformer2DModel(ZImageTransformer2DModel): + """Runtime adapter for vendor DiffusionTransformer checkpoints. + + This adapter reuses the shared Z-Image DiT and injects Ming-specific request conditions. + """ + + def forward( + self, + x: list[torch.Tensor], + t, + cap_feats: list[torch.Tensor], + patch_size=2, + f_patch_size=1, + ): + generated_frames = [item.shape[1] for item in x] + ref_x = None + cap_feats_2 = None + if is_forward_context_available(): + context = get_forward_context() + if context.ref_latent is not None: + ref_x = [item.to(device=x[0].device, dtype=x[0].dtype) for item in context.ref_latent.unbind(dim=0)] + if context.direct_condition is not None: + cap_feats_2 = [ + item.to(device=x[0].device, dtype=x[0].dtype) for item in context.direct_condition.unbind(dim=0) + ] + + output, metadata = super().forward( + x, + t, + cap_feats, + patch_size=patch_size, + f_patch_size=f_patch_size, + ref_x=ref_x, + cap_feats_2=cap_feats_2, + ) + output = [item[:, :frames] for item, frames in zip(output, generated_frames)] + return output, metadata + + +__all__ = ["MingImageTransformer2DModel"] diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py index b77d2265c..b40fd57bf 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved. # @@ -584,7 +584,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen # 5. Prepare timesteps if image is None: - image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2) + # for both [B, C, H, W] and multi-layer/frame [B, C, F, H, W] + image_seq_len = (latents.shape[-2] // 2) * (latents.shape[-1] // 2) mu = calculate_shift( image_seq_len, self.scheduler.config.get("base_image_seq_len", 256), @@ -643,7 +644,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen latents_typed = latents.to(self.od_config.dtype) if apply_cfg: - latent_model_input = latents_typed.repeat(2, 1, 1, 1) + repeat_dims = (2,) + (1,) * (latents_typed.ndim - 1) + latent_model_input = latents_typed.repeat(*repeat_dims) prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds timestep_model_input = timestep.repeat(2) else: @@ -651,9 +653,13 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen prompt_embeds_model_input = prompt_embeds timestep_model_input = timestep - latent_model_input = latent_model_input.unsqueeze(2) + if latent_model_input.ndim == 4: + latent_model_input = latent_model_input.unsqueeze(2) latent_model_input_list = list(latent_model_input.unbind(dim=0)) + if getattr(self, "_uses_cudagraph_trees", False): + # check Ming-Image + torch.compiler.cudagraph_mark_step_begin() model_out_list = self.transformer( latent_model_input_list, timestep_model_input, @@ -690,7 +696,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen else: noise_pred = torch.stack([t.float() for t in model_out_list], dim=0) - noise_pred = noise_pred.squeeze(2) + if latents.ndim == 4: + noise_pred = noise_pred.squeeze(2) noise_pred = -noise_pred # compute the previous noisy sample x_t -> x_t-1 diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index f25c01fca..900e2e2cd 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # _sp_plan definition adapted from HuggingFace diffusers library (_cp_plan) # Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved. @@ -56,6 +56,8 @@ from vllm_omni.model_executor.layers.timestep_embedding import timestep_embeddin ADALN_EMBED_DIM = 256 SEQ_MULTI_OF = 32 +LEARNED_PADDING = "learned" +ZERO_MASKED_PADDING = "zero_masked" logger = init_logger(__name__) @@ -81,6 +83,8 @@ class UnifiedPrepare(nn.Module): cap_sin: torch.Tensor, x_item_seqlens: list[int], cap_item_seqlens: list[int], + x_attn_mask: torch.Tensor, + cap_attn_mask: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Combine x and cap tensors into unified sequences. @@ -110,8 +114,9 @@ class UnifiedPrepare(nn.Module): unified_cos = pad_sequence(unified_cos, batch_first=True, padding_value=0.0) unified_sin = pad_sequence(unified_sin, batch_first=True, padding_value=0.0) unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(unified_item_seqlens): - unified_attn_mask[i, :seq_len] = 1 + for i, (x_len, cap_len) in enumerate(zip(x_item_seqlens, cap_item_seqlens)): + unified_attn_mask[i, :x_len] = x_attn_mask[i, :x_len] + unified_attn_mask[i, x_len : x_len + cap_len] = cap_attn_mask[i, :cap_len] return unified, unified_cos, unified_sin, unified_attn_mask @@ -671,6 +676,8 @@ class ZImageTransformer2DModel(CachedTransformer): axes_dims=[32, 48, 48], axes_lens=[1024, 512, 512], quant_config: "QuantizationConfig | None" = None, + alignment_padding_mode: str = LEARNED_PADDING, + multi_frame_output: bool = True, ) -> None: super().__init__() # NOTE: `DiffusersPipelineLoader.load_model()` initializes this module @@ -685,6 +692,12 @@ class ZImageTransformer2DModel(CachedTransformer): self.dim = dim self.n_heads = n_heads + if alignment_padding_mode not in {LEARNED_PADDING, ZERO_MASKED_PADDING}: + raise ValueError(f"Unsupported Z-Image alignment padding mode: {alignment_padding_mode!r}") + if type(multi_frame_output) is not bool: + raise ValueError("multi_frame_output must be boolean.") + self.alignment_padding_mode = alignment_padding_mode + self.multi_frame_output = multi_frame_output self.rope_theta = rope_theta self.t_scale = t_scale self.gradient_checkpointing = False @@ -780,8 +793,12 @@ class ZImageTransformer2DModel(CachedTransformer): ), ) - self.x_pad_token = nn.Parameter(torch.empty((1, dim))) - self.cap_pad_token = nn.Parameter(torch.empty((1, dim))) + if alignment_padding_mode == LEARNED_PADDING: + self.x_pad_token = nn.Parameter(torch.empty((1, dim))) + self.cap_pad_token = nn.Parameter(torch.empty((1, dim))) + else: + self.register_parameter("x_pad_token", None) + self.register_parameter("cap_pad_token", None) self.layers = nn.ModuleList( [ @@ -823,6 +840,8 @@ class ZImageTransformer2DModel(CachedTransformer): .permute(6, 0, 3, 1, 4, 2, 5) .reshape(self.out_channels, F, H, W) ) + if not self.multi_frame_output: + x[i] = x[i][:, :1] return x @staticmethod @@ -840,6 +859,8 @@ class ZImageTransformer2DModel(CachedTransformer): all_cap_feats: list[torch.Tensor], patch_size: int, f_patch_size: int, + all_image_ref: list[torch.Tensor | None] | None = None, + all_cap_feats_2: list[torch.Tensor | None] | None = None, ): pH = pW = patch_size pF = f_patch_size @@ -852,10 +873,16 @@ class ZImageTransformer2DModel(CachedTransformer): all_cap_pos_ids = [] all_cap_pad_mask = [] all_cap_feats_out = [] + all_cap_feats_2_out = [] - for i, (image, cap_feat) in enumerate(zip(all_image, all_cap_feats)): + if all_image_ref is None: + all_image_ref = [None] * len(all_image) + if all_cap_feats_2 is None: + all_cap_feats_2 = [None] * len(all_image) + + for image, cap_feat, cap_feat_2, image_ref in zip(all_image, all_cap_feats, all_cap_feats_2, all_image_ref): ### Process Caption - cap_ori_len = len(cap_feat) + cap_ori_len = len(cap_feat) + (len(cap_feat_2) if cap_feat_2 is not None else 0) cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF # padded position ids cap_padded_pos_ids = self.create_coordinate_grid( @@ -875,13 +902,17 @@ class ZImageTransformer2DModel(CachedTransformer): ) ) # padded feature - cap_padded_feat = torch.cat( - [cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)], - dim=0, - ) - all_cap_feats_out.append(cap_padded_feat) + if cap_feat_2 is None: + all_cap_feats_out.append(torch.cat([cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)], dim=0)) + else: + all_cap_feats_out.append(cap_feat) + all_cap_feats_2_out.append( + torch.cat([cap_feat_2, torch.zeros_like(cap_feat_2[-1:]).repeat(cap_padding_len, 1)]) + ) ### Process Image + if image_ref is not None: + image = torch.cat([image, image_ref], dim=1) C, F, H, W = image.size() all_image_size.append((F, H, W)) F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW @@ -931,6 +962,7 @@ class ZImageTransformer2DModel(CachedTransformer): all_cap_pos_ids, all_image_pad_mask, all_cap_pad_mask, + all_cap_feats_2_out, ) def forward( @@ -940,6 +972,8 @@ class ZImageTransformer2DModel(CachedTransformer): cap_feats: list[torch.Tensor], patch_size=2, f_patch_size=1, + ref_x: list[torch.Tensor | None] | None = None, + cap_feats_2: list[torch.Tensor] | None = None, ): assert patch_size in self.all_patch_size assert f_patch_size in self.all_f_patch_size @@ -957,7 +991,8 @@ class ZImageTransformer2DModel(CachedTransformer): cap_pos_ids, x_inner_pad_mask, cap_inner_pad_mask, - ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size) + cap_feats_2, + ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size, ref_x, cap_feats_2) # x embed & refine x_item_seqlens = [len(_) for _ in x] @@ -971,9 +1006,14 @@ class ZImageTransformer2DModel(CachedTransformer): adaln_input = t.type_as(x) # Use torch.where instead of x[mask]= to avoid aten::index_put_/nonzero and cudaStreamSynchronize x_pad_mask = torch.cat(x_inner_pad_mask) + x_padding = ( + self.x_pad_token.expand(x.shape[0], -1) + if self.alignment_padding_mode == LEARNED_PADDING + else torch.zeros_like(x) + ) x = torch.where( x_pad_mask.unsqueeze(1).expand_as(x), - self.x_pad_token.expand(x.shape[0], -1), + x_padding, x, ) x = list(x.split(x_item_seqlens, dim=0)) @@ -987,22 +1027,36 @@ class ZImageTransformer2DModel(CachedTransformer): x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device) for i, seq_len in enumerate(x_item_seqlens): x_attn_mask[i, :seq_len] = 1 + if self.alignment_padding_mode == ZERO_MASKED_PADDING: + x_attn_mask[i, :seq_len].masked_fill_(x_inner_pad_mask[i], False) for layer in self.noise_refiner: x = layer(x, x_attn_mask, x_cos, x_sin, adaln_input) # cap embed & refine cap_item_seqlens = [len(_) for _ in cap_feats] + cap_feats = torch.cat(cap_feats, dim=0) + cap_feats = self.cap_embedder(cap_feats) + if cap_feats_2: + cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0)) + if len(cap_feats) != len(cap_feats_2): + raise ValueError("Primary and direct caption conditions must have equal batch size.") + cap_feats = [torch.cat([primary, direct], dim=0) for primary, direct in zip(cap_feats, cap_feats_2)] + cap_item_seqlens = [len(item) for item in cap_feats] + cap_feats = torch.cat(cap_feats, dim=0) assert all(_ % SEQ_MULTI_OF == 0 for _ in cap_item_seqlens) cap_max_item_seqlen = max(cap_item_seqlens) - cap_feats = torch.cat(cap_feats, dim=0) - cap_feats = self.cap_embedder(cap_feats) # Use torch.where instead of cap_feats[mask]= to avoid aten::index_put_/nonzero and cudaStreamSynchronize cap_pad_mask = torch.cat(cap_inner_pad_mask) + cap_padding = ( + self.cap_pad_token.expand(cap_feats.shape[0], -1) + if self.alignment_padding_mode == LEARNED_PADDING + else torch.zeros_like(cap_feats) + ) cap_feats = torch.where( cap_pad_mask.unsqueeze(1).expand_as(cap_feats), - self.cap_pad_token.expand(cap_feats.shape[0], -1), + cap_padding, cap_feats, ) cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0)) @@ -1016,6 +1070,8 @@ class ZImageTransformer2DModel(CachedTransformer): cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device) for i, seq_len in enumerate(cap_item_seqlens): cap_attn_mask[i, :seq_len] = 1 + if self.alignment_padding_mode == ZERO_MASKED_PADDING: + cap_attn_mask[i, :seq_len].masked_fill_(cap_inner_pad_mask[i], False) for layer in self.context_refiner: cap_feats = layer(cap_feats, cap_attn_mask, cap_cos, cap_sin) @@ -1023,9 +1079,17 @@ class ZImageTransformer2DModel(CachedTransformer): # Prepare unified tensors via UnifiedPrepare module # This enables _cp_plan to shard outputs via split_output=True unified, unified_cos, unified_sin, unified_attn_mask = self.unified_prepare( - x, x_cos, x_sin, cap_feats, cap_cos, cap_sin, x_item_seqlens, cap_item_seqlens + x, + x_cos, + x_sin, + cap_feats, + cap_cos, + cap_sin, + x_item_seqlens, + cap_item_seqlens, + x_attn_mask, + cap_attn_mask, ) - # Main transformer blocks for layer in self.layers: unified = layer(unified, unified_attn_mask, unified_cos, unified_sin, adaln_input) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index a9cb0e411..bbbd6e134 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -204,6 +204,11 @@ _DIFFUSION_MODELS = { "pipeline_ming_imagegen", "MingImagePipeline", ), + "MingImageDiffusionPipeline": ( + "ming_image", + "pipeline", + "MingImageDiffusionPipeline", + ), "SanaWmPipeline": ( "sana_wm", "pipeline_sana_wm", @@ -390,6 +395,7 @@ _DIFFUSION_MODELS = { "Krea2Pipeline", ), } +_DIFFUSION_MODELS["MingImageLayeredDiffusionPipeline"] = _DIFFUSION_MODELS["MingImageDiffusionPipeline"] DiffusionModelRegistry = _ModelRegistry( @@ -663,6 +669,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "BagelPipeline": "get_bagel_post_process_func", "LancePipeline": "get_lance_post_process_func", "MingImagePipeline": "get_ming_image_post_process_func", + "MingImageDiffusionPipeline": "get_ming_image_post_process_func", "InternVLAA1Pipeline": "get_internvla_a1_post_process_func", "Pi0Pipeline": "get_pi0_post_process_func", "Pi05Pipeline": "get_pi05_post_process_func", @@ -696,6 +703,9 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "Krea2Pipeline": "get_krea2_post_process_func", "HunyuanImage3ForCausalMM": "get_hunyuan_image3_post_process_func", } +_DIFFUSION_POST_PROCESS_FUNCS["MingImageLayeredDiffusionPipeline"] = _DIFFUSION_POST_PROCESS_FUNCS[ + "MingImageDiffusionPipeline" +] _DIFFUSION_IR_OP_PRIORITY_FUNCS = { # arch: ir_op_priority_func diff --git a/vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py b/vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py index c685bb9e7..ff29730b8 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The vLLM-Omni team. # Copyright 2023 Antgroup and The HuggingFace Inc. team. All rights reserved. # Adapted from Ming @@ -728,7 +729,8 @@ class BailingMoeV2Model(nn.Module): inputs_embeds: torch.Tensor | None = None, image_mask: torch.Tensor | None = None, audio_mask: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: + capture_layers: tuple[int, ...] | None = None, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, dict[int, torch.Tensor]]: if get_pp_group().is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds @@ -740,7 +742,14 @@ class BailingMoeV2Model(nn.Module): hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] - for layer in self.layers[self.start_layer : self.end_layer]: + capture_set = set(capture_layers or ()) + captured: dict[int, torch.Tensor] = {} + for layer_idx, layer in enumerate( + self.layers[self.start_layer : self.end_layer], + start=self.start_layer, + ): + if layer_idx in capture_set: + captured[layer_idx] = hidden_states if residual is None else hidden_states + residual hidden_states, residual = layer( positions, hidden_states, @@ -753,6 +762,10 @@ class BailingMoeV2Model(nn.Module): return IntermediateTensors({"hidden_states": hidden_states, "residual": residual}) hidden_states, _ = self.norm(hidden_states, residual) + if self.end_layer in capture_set: + captured[self.end_layer] = hidden_states + if capture_layers: + return hidden_states, captured return hidden_states @@ -811,6 +824,7 @@ class BailingMoeV2ForCausalLM(nn.Module, CustomProcessMixin): inputs_embeds: torch.Tensor | None = None, image_mask: torch.Tensor | None = None, audio_mask: torch.Tensor | None = None, + capture_layers: tuple[int, ...] | None = None, ): hidden_states = self.model( input_ids=input_ids, @@ -819,6 +833,7 @@ class BailingMoeV2ForCausalLM(nn.Module, CustomProcessMixin): inputs_embeds=inputs_embeds, image_mask=image_mask, audio_mask=audio_mask, + capture_layers=capture_layers, ) return hidden_states diff --git a/vllm_omni/model_executor/models/ming_image/__init__.py b/vllm_omni/model_executor/models/ming_image/__init__.py new file mode 100644 index 000000000..4244ce4b4 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_image/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/vllm_omni/model_executor/models/ming_image/checkpoint.py b/vllm_omni/model_executor/models/ming_image/checkpoint.py new file mode 100644 index 000000000..461766f6e --- /dev/null +++ b/vllm_omni/model_executor/models/ming_image/checkpoint.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Checkpoint resolution for the Ming-Image MLLM stage.""" + +from pathlib import Path + +from vllm_omni.model_executor.model_loader.weight_utils import ( + download_weights_from_hf_specific, +) + +_MLLM_REQUIRED_PATTERNS = [ + "mllm/**", + "mlp/**", +] + + +def resolve_ming_image_model_root( + model: str, + revision: str | None, + task_type: str | None, +) -> str: + """Materialize the MLLM and its sibling query-token checkpoint. + + The autoregressive weights and tokenizer live under mllm/, + while the learned image-query tokens consumed by that stage live under mlp/. + Stage subdirectory resolution runs after this hook, so return the repository + root and let model_subdir/tokenizer_subdir select mllm/. + """ + del task_type + path = Path(model) + if path.is_dir(): + if path.name == "mllm" and (path / "config.json").is_file(): + return str(path.parent) + return str(path) + + return download_weights_from_hf_specific( + model_name_or_path=model, + cache_dir=None, + allow_patterns=_MLLM_REQUIRED_PATTERNS, + revision=revision, + require_all=True, + ) + + +__all__ = ["resolve_ming_image_model_root"] diff --git a/vllm_omni/model_executor/models/ming_image/model.py b/vllm_omni/model_executor/models/ming_image/model.py new file mode 100644 index 000000000..962414c98 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_image/model.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""vLLM-native MLLM stage for inclusionAI Ming-Image.""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as F +from vllm.config import VllmConfig +from vllm.inputs import MultiModalDataDict +from vllm.model_executor.models.qwen2_5_vl import Qwen2_5_VisionTransformer +from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.parse import MultiModalDataParser + +from vllm_omni.model_executor.models.ming_flash_omni.ming_flash_omni_thinker import ( + MingFlashOmniThinkerDummyInputsBuilder, + MingFlashOmniThinkerForConditionalGeneration, + MingFlashOmniThinkerMultiModalProcessor, + MingFlashOmniThinkerProcessingInfo, +) +from vllm_omni.model_executor.models.ming_flash_omni.modeling_bailing_moe_v2 import ( + BailingMoeV2ForCausalLM, +) +from vllm_omni.model_executor.models.ming_flash_omni.projectors import VisionProjector +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.transformers_utils.configs.ming_flash_omni import BailingMM2Config +from vllm_omni.transformers_utils.processors.ming import MingImageProcessor + + +class _MingImageDataParser(MultiModalDataParser): + def parse_mm_data(self, mm_data): + normalized = dict(mm_data) + if "img2img" in normalized: + references = normalized.pop("img2img") + references = references if isinstance(references, list) else [references] + existing = normalized.get("image", []) + existing = existing if isinstance(existing, list) else [existing] + normalized["image"] = existing + references + return super().parse_mm_data(normalized) + + +class MingImageProcessingInfo(MingFlashOmniThinkerProcessingInfo): + def get_hf_config(self) -> BailingMM2Config: + return self.ctx.get_hf_config(BailingMM2Config) + + def get_hf_processor(self, **kwargs: object): + return self.ctx.get_hf_processor(MingImageProcessor, **kwargs) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"image": None, "img2img": 1} + + def get_mm_max_tokens_per_item(self, seq_len, mm_counts): + counts = dict(mm_counts or {}) + counts["image"] = counts.get("image", 0) + counts.pop("img2img", 0) + return super(MingFlashOmniThinkerProcessingInfo, self).get_mm_max_tokens_per_item( + seq_len=seq_len, + mm_counts=counts, + ) + + def get_data_parser(self): + return _MingImageDataParser() + + +class MingImageDummyInputsBuilder(MingFlashOmniThinkerDummyInputsBuilder): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + count = mm_counts.get("image", 0) + mm_counts.get("img2img", 0) + return self.info.get_hf_processor().image_token * count + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options=None, + ) -> MultiModalDataDict: + count = mm_counts.get("image", 0) + mm_counts.get("img2img", 0) + return { + "image": self._get_dummy_images( + width=448, + height=448, + num_images=count, + ) + } + + +class MingImageMultiModalProcessor(MingFlashOmniThinkerMultiModalProcessor): + pass + + +@MULTIMODAL_REGISTRY.register_processor( + MingImageMultiModalProcessor, + info=MingImageProcessingInfo, + dummy_inputs=MingImageDummyInputsBuilder, +) +class MingImageForConditionalGeneration(MingFlashOmniThinkerForConditionalGeneration): + """Ming-Image MLLM stage with Qwen2.5-VL and direct-state capture.""" + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.": "language_model.", + "linear_proj.": "linear_proj.proj.", + }, + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + torch.nn.Module.__init__(self) + thinker_config = vllm_config.model_config.hf_config + if not isinstance(thinker_config, BailingMM2Config): + raise TypeError(f"Expected BailingMM2Config, got {type(thinker_config).__name__}") + if thinker_config.llm_config is None or thinker_config.vision_config is None: + raise ValueError("Ming-Image requires llm_config and vision_config.") + + llm_config = thinker_config.llm_config + self.config = llm_config + self.thinker_config = thinker_config + self.have_multimodal_outputs = True + + with self._mark_language_model(vllm_config): + self.language_model = BailingMoeV2ForCausalLM( + vllm_config=vllm_config.with_hf_config(llm_config), + prefix=maybe_prefix(prefix, "llm"), + ) + + with self._mark_tower_model(vllm_config, "image"): + self.vision = Qwen2_5_VisionTransformer( + thinker_config.vision_config, + norm_eps=llm_config.rms_norm_eps, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "vision"), + ) + self.linear_proj = VisionProjector( + vision_dim=thinker_config.vision_config.out_hidden_size, + llm_dim=llm_config.hidden_size, + mlp_depth=getattr(thinker_config, "mlp_depth", 2), + ) + + self.audio = None + self.linear_proj_audio = None + self.query_tokens_dict = torch.nn.ParameterDict() + component_path = Path(vllm_config.model_config.model) + model_root = component_path.parent if component_path.name == "mllm" else component_path + self._load_image_gen_query_tokens(str(model_root)) + if not self.query_tokens_dict: + raise ValueError(f"Ming-Image query tokens were not found under {model_root / 'mlp'}") + + self.capture_layers = (5, 12, llm_config.num_hidden_layers) + self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors + + def extract_image_feature(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: + with torch.amp.autocast(pixel_values.device.type, dtype=torch.bfloat16): + image_embeds = self.vision(pixel_values, grid_thw=grid_thw.tolist()) + return F.normalize(self.linear_proj(image_embeds), dim=-1) + + def compute_logits( + self, + hidden_states: torch.Tensor, + sampling_metadata=None, + ) -> torch.Tensor | None: + # vLLM v1 calls this with hidden states only; the reused Bailing model + # still accepts the legacy metadata argument (which it does not use). + return self.language_model.compute_logits(hidden_states, sampling_metadata) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loaded_weights = super().load_weights(weights) + # Query tokens live in the root-level mlp checkpoint and are loaded by + # ``_load_image_gen_query_tokens`` during construction, outside vLLM's + # main checkpoint iterator. Include them in the returned set so strict + # post-load validation accounts for that checkpoint-owned component. + loaded_weights.update(f"query_tokens_dict.{scale}" for scale in self.query_tokens_dict) + return loaded_weights + + def _compute_modality_masks( + self, + input_ids: torch.Tensor | None, + inputs_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if input_ids is None: + if inputs_embeds is None: + raise ValueError("input_ids and inputs_embeds cannot both be None") + empty_mask = torch.zeros( + inputs_embeds.shape[:-1], + device=inputs_embeds.device, + dtype=torch.bool, + ) + return empty_mask, empty_mask + return input_ids == self.config.image_patch_token, torch.zeros_like(input_ids, dtype=torch.bool) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors=None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> OmniOutput: + image_mask, audio_mask = self._compute_modality_masks(input_ids, inputs_embeds) + result = self.language_model.forward( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + image_mask=image_mask, + audio_mask=audio_mask, + capture_layers=self.capture_layers, + ) + if not isinstance(result, tuple): + return OmniOutput(text_hidden_states=result, multimodal_outputs={"final_hidden_states": result}) + + hidden_states, captured = result + multimodal_outputs: dict[str, Any] = {"final_hidden_states": hidden_states} + for layer_idx, layer_hidden in captured.items(): + multimodal_outputs[f"hidden_states_{layer_idx}"] = layer_hidden + return OmniOutput( + text_hidden_states=hidden_states, + multimodal_outputs=multimodal_outputs, + ) + + +__all__ = ["MingImageForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/ming_image/pipeline.py b/vllm_omni/model_executor/models/ming_image/pipeline.py new file mode 100644 index 000000000..c4b8157ed --- /dev/null +++ b/vllm_omni/model_executor/models/ming_image/pipeline.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Two-stage Ming-Image topology.""" + +from vllm_omni.config.stage_config import PipelineConfig, StageExecutionType, StagePipelineConfig + +_PROC = "vllm_omni.model_executor.stage_input_processors.ming_image" +_CHECKPOINT = "vllm_omni.model_executor.models.ming_image.checkpoint" + +MING_IMAGE_PIPELINE = PipelineConfig( + model_type="ming_image", + default_deploy_config_name="ming_image.yaml", + model_arch="MingImageForConditionalGeneration", + hf_architectures=("MingImageForConditionalGeneration",), + diffusers_class_name="MingImageDiffusionPipeline", + diffusers_class_aliases=("MingImageLayeredDiffusionPipeline",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="mllm", + model_arch="MingImageForConditionalGeneration", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=False, + owns_tokenizer=True, + requires_multimodal_data=True, + engine_output_type="latent", + model_subdir="mllm", + tokenizer_subdir="mllm", + sampling_constraints={"detokenize": False}, + model_path_resolver=f"{_CHECKPOINT}.resolve_ming_image_model_root", + ), + StagePipelineConfig( + stage_id=1, + model_stage="dit", + model_arch="MingImageDiffusionPipeline", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(0,), + final_output=True, + final_output_type="image", + custom_process_input_func=f"{_PROC}.thinker2image", + ), + ), +) + +__all__ = ["MING_IMAGE_PIPELINE"] diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 71b2792e3..a0a812343 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -425,6 +425,11 @@ _OMNI_MODELS = { "ming_flash_omni_thinker", "MingFlashOmniThinkerForConditionalGeneration", ), + "MingImageForConditionalGeneration": ( + "ming_image", + "model", + "MingImageForConditionalGeneration", + ), "MingFlashOmniTalkerForConditionalGeneration": ( "ming_flash_omni", "ming_flash_omni_talker", diff --git a/vllm_omni/model_executor/stage_input_processors/ming_image.py b/vllm_omni/model_executor/stage_input_processors/ming_image.py new file mode 100644 index 000000000..59e8e4f94 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/ming_image.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Strict MLLM-to-diffusion bridge for Ming-Image.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from vllm_omni.model_executor.stage_input_processors.ming_flash_omni import _ensure_list + +_IMAGE_PATCH_TOKEN_ID = 157157 +_IMAGE_END_TOKEN_ID = 157159 +_NUM_QUERY_TOKENS = 256 +_CAPTURE_LAYERS = (5, 12, 20) + + +def _output_payload(source_output: Any) -> tuple[list[int], dict[str, torch.Tensor]]: + if not getattr(source_output, "outputs", None): + raise ValueError("Ming-Image thinker produced no outputs.") + output = source_output.outputs[0] + multimodal = getattr(output, "multimodal_output", None) or {} + return _ensure_list(source_output.prompt_token_ids), multimodal + + +def _validate_query_suffix(prompt_ids: list[int]) -> tuple[int, int]: + expected = _NUM_QUERY_TOKENS + 2 + if len(prompt_ids) < expected: + raise ValueError("Ming-Image prompt is shorter than its image-query suffix.") + start = len(prompt_ids) - _NUM_QUERY_TOKENS - 1 + end = len(prompt_ids) - 1 + if prompt_ids[end] != _IMAGE_END_TOKEN_ID or any(token != _IMAGE_PATCH_TOKEN_ID for token in prompt_ids[start:end]): + raise ValueError("Ming-Image query suffix must end with 256 image-patch tokens and an image-end token.") + return start, end + + +def _as_sequence(hidden: torch.Tensor, prompt_length: int, name: str) -> torch.Tensor: + if hidden.ndim == 3: + if hidden.shape[0] != 1: + raise ValueError(f"{name} must have batch size 1, got {tuple(hidden.shape)}") + hidden = hidden[0] + if hidden.ndim != 2 or hidden.shape[0] != prompt_length: + raise ValueError(f"{name} shape {tuple(hidden.shape)} does not match prompt length {prompt_length}.") + return hidden + + +def _extract_conditions(source_output: Any) -> tuple[torch.Tensor, torch.Tensor]: + prompt_ids, multimodal = _output_payload(source_output) + query_start, query_end = _validate_query_suffix(prompt_ids) + final_hidden = multimodal.get("final_hidden_states") + if not isinstance(final_hidden, torch.Tensor): + raise ValueError("Ming-Image thinker output is missing final_hidden_states.") + final_hidden = _as_sequence(final_hidden, len(prompt_ids), "final_hidden_states") + query_hidden = final_hidden[query_start:query_end].detach().contiguous() + + selected: list[torch.Tensor] = [] + for layer_idx in _CAPTURE_LAYERS: + value = multimodal.get(f"hidden_states_{layer_idx}") + if not isinstance(value, torch.Tensor): + raise ValueError(f"Ming-Image thinker output is missing hidden_states_{layer_idx}.") + selected.append(_as_sequence(value, len(prompt_ids), f"hidden_states_{layer_idx}")) + + # Vendor labels mark the full generated query block, including delimiters. + # Direct VLM conditioning is restricted to the original prompt/reference + # prefix and therefore excludes all 258 appended tokens. + direct_end = len(prompt_ids) - (_NUM_QUERY_TOKENS + 2) + direct_hidden = torch.cat([value[:direct_end] for value in selected], dim=-1) + return query_hidden, direct_hidden.detach().contiguous() + + +def thinker2image( + source_outputs: list[Any], + prompt: Any | None = None, + requires_multimodal_data: bool = False, + sampling_params: Any | None = None, +) -> list[dict[str, Any]]: + del requires_multimodal_data + if len(source_outputs) != 1: + raise ValueError(f"Ming-Image expects one thinker output, got {len(source_outputs)}.") + + extra_args = getattr(sampling_params, "extra_args", None) or {} + negative_prompt = extra_args.get("negative_prompt") + if isinstance(prompt, dict): + negative_prompt = prompt.get("negative_prompt", negative_prompt) + if isinstance(negative_prompt, str) and negative_prompt.strip(): + raise ValueError("Ming-Image uses zero negative conditioning and does not accept negative_prompt.") + + query_hidden, direct_hidden = _extract_conditions(source_outputs[0]) + extra: dict[str, Any] = { + "query_hidden_states": query_hidden, + "direct_hidden_states": direct_hidden, + } + + if isinstance(prompt, dict): + mm_data = prompt.get("multi_modal_data") or {} + reference = mm_data.get("img2img", mm_data.get("image")) + if isinstance(reference, list): + reference = reference[0] if reference else None + if reference is not None: + extra["reference_image"] = reference + + num_layers = int(extra_args.get("num_layers", 1)) + if num_layers < 1: + raise ValueError("num_layers must be at least 1.") + extra["num_layers"] = num_layers + return [{"prompt": "", "extra": extra}] + + +__all__ = ["thinker2image"] diff --git a/vllm_omni/model_extras/ming_image.py b/vllm_omni/model_extras/ming_image.py new file mode 100644 index 000000000..51ccc974e --- /dev/null +++ b/vllm_omni/model_extras/ming_image.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from vllm_omni.model_extras.ming_flash_omni import ( + build_image_to_image_prompt, + build_text_to_image_prompt, +) + +MING_IMAGE_EXTRA_BODY_PARAMS = frozenset( + { + "height", + "width", + "steps", + "cfg", + "seed", + "num_layers", + "negative_prompt", + } +) +MING_IMAGE_EXTRA_OUTPUT_PARAMS: frozenset[str] = frozenset() +MING_IMAGE_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True + +__all__ = [ + "MING_IMAGE_EXTRA_BODY_PARAMS", + "MING_IMAGE_EXTRA_OUTPUT_PARAMS", + "MING_IMAGE_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES", + "build_image_to_image_prompt", + "build_text_to_image_prompt", +] diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 9c4a6eed6..2ba1ee01e 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -75,6 +75,17 @@ from vllm_omni.model_extras.ming_flash_omni import ( from vllm_omni.model_extras.ming_flash_omni import ( build_text_to_image_prompt as build_ming_flash_omni_text_to_image_prompt, ) +from vllm_omni.model_extras.ming_image import ( + MING_IMAGE_EXTRA_BODY_PARAMS, + MING_IMAGE_EXTRA_OUTPUT_PARAMS, + MING_IMAGE_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, +) +from vllm_omni.model_extras.ming_image import ( + build_image_to_image_prompt as build_ming_image_image_to_image_prompt, +) +from vllm_omni.model_extras.ming_image import ( + build_text_to_image_prompt as build_ming_image_text_to_image_prompt, +) from vllm_omni.model_extras.sana_video import SANA_VIDEO_EXTRA_BODY_PARAMS from vllm_omni.model_extras.sensenova_u1 import ( SENSENOVA_U1_EXTRA_BODY_PARAMS, @@ -299,7 +310,15 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { "text_to_image_prompt_builder": build_ming_flash_omni_text_to_image_prompt, "image_to_image_prompt_builder": build_ming_flash_omni_image_to_image_prompt, }, + "MingImageDiffusionPipeline": { + "extra_body_params": MING_IMAGE_EXTRA_BODY_PARAMS, + "extra_output_params": MING_IMAGE_EXTRA_OUTPUT_PARAMS, + "init_extra_args_for_non_diffusion_stages": MING_IMAGE_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, + "text_to_image_prompt_builder": build_ming_image_text_to_image_prompt, + "image_to_image_prompt_builder": build_ming_image_image_to_image_prompt, + }, } +_EXTRA_SPECS["MingImageLayeredDiffusionPipeline"] = _EXTRA_SPECS["MingImageDiffusionPipeline"] for model_class_name in ("LTX2Pipeline", "LTX2TwoStagePipeline"): _EXTRA_SPECS[model_class_name]["transformer_config_subfolder_resolver"] = ltx_transformer_config_subfolder diff --git a/vllm_omni/transformers_utils/configs/ming_flash_omni.py b/vllm_omni/transformers_utils/configs/ming_flash_omni.py index 6799a411e..e78d66d78 100644 --- a/vllm_omni/transformers_utils/configs/ming_flash_omni.py +++ b/vllm_omni/transformers_utils/configs/ming_flash_omni.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The vLLM-Omni team. # Copyright 2024 ANT Group and the HuggingFace Inc. team. All rights reserved. # @@ -19,7 +20,13 @@ import os from typing import Any, ClassVar -from transformers import AutoConfig, AutoTokenizer, PretrainedConfig, PreTrainedTokenizerFast +from transformers import ( + AutoConfig, + AutoTokenizer, + PretrainedConfig, + PreTrainedTokenizerFast, + Qwen2_5_VLVisionConfig, +) from transformers.utils import logging logger = logging.get_logger(__name__) @@ -251,9 +258,19 @@ class BailingMM2Config(PretrainedConfig): **kwargs, ): self.audio_config = WhisperEncoderConfig(**audio_config) if isinstance(audio_config, dict) else audio_config - self.vision_config = ( - Qwen3VLMoeVisionConfig(**vision_config) if isinstance(vision_config, dict) else vision_config - ) + if isinstance(vision_config, dict): + vision_config = dict(vision_config) + vision_model_type = vision_config.get("model_type") + if vision_model_type in {"qwen2_5_vit", "qwen2_5_vl"}: + # 1) Ming-Image uses the Qwen2.5-VL vision tower while + # 2) Ming-flash-omni-2.0 uses Qwen3-VL. Keep the composition config shared + # and select the concrete tower from ckpt metadata. + vision_config.pop("model_type", None) + self.vision_config = Qwen2_5_VLVisionConfig(**vision_config) + else: + self.vision_config = Qwen3VLMoeVisionConfig(**vision_config) + else: + self.vision_config = vision_config self.llm_config = BailingMoeV2Config(**llm_config) if isinstance(llm_config, dict) else llm_config self.mlp_depth = mlp_depth super().__init__(**kwargs) diff --git a/vllm_omni/transformers_utils/processors/__init__.py b/vllm_omni/transformers_utils/processors/__init__.py index 52ca65753..ec9ae133b 100644 --- a/vllm_omni/transformers_utils/processors/__init__.py +++ b/vllm_omni/transformers_utils/processors/__init__.py @@ -1,12 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The vLLM-Omni team. from vllm_omni.transformers_utils.processors.ming import ( MingFlashOmniProcessor, + MingImageProcessor, MingWhisperFeatureExtractor, ) __all__ = [ "MingFlashOmniProcessor", + "MingImageProcessor", "MingWhisperFeatureExtractor", ] diff --git a/vllm_omni/transformers_utils/processors/ming.py b/vllm_omni/transformers_utils/processors/ming.py index aa59e5324..5b6d852de 100644 --- a/vllm_omni/transformers_utils/processors/ming.py +++ b/vllm_omni/transformers_utils/processors/ming.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The vLLM-Omni team. # Copyright 2024 ANT Group and the HuggingFace Inc. team. # @@ -18,7 +19,7 @@ from typing import Any import numpy as np import torch -from transformers import AutoFeatureExtractor, AutoProcessor +from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, Qwen2VLImageProcessor from transformers.feature_extraction_utils import BatchFeature, FeatureExtractionMixin from transformers.processing_utils import ProcessorMixin from transformers.tokenization_utils_base import PreTokenizedInput, TextInput @@ -481,6 +482,87 @@ class MingFlashOmniProcessor(ProcessorMixin): return list(dict.fromkeys(names)) +class MingImageProcessor(ProcessorMixin): + """Image-only Bailing processor used by Ming-Image ckpts. + + Equivalent to Qwen2-VL preprocessing. Construct the upstream component explicitly + in case model repo auto_map points at files that are not bundled in the ckpt. + """ + + attributes = ["image_processor", "tokenizer"] + image_processor_class = "Qwen2VLImageProcessor" + tokenizer_class = "AutoTokenizer" + + def __init__(self, image_processor=None, tokenizer=None, merge_size: int = 2, **kwargs): + if image_processor is None or tokenizer is None: + raise ValueError("MingImageProcessor requires image_processor and tokenizer.") + self.spatial_merge_size = int(getattr(image_processor, "merge_size", merge_size)) + self.image_token = PLACEHOLDER_IMAGE_TOKEN_IN_TEXT + super().__init__(image_processor=image_processor, tokenizer=tokenizer, **kwargs) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): + image_processor = Qwen2VLImageProcessor.from_pretrained( + pretrained_model_name_or_path, + *args, + **kwargs, + ) + tokenizer = AutoTokenizer.from_pretrained( + pretrained_model_name_or_path, + *args, + **kwargs, + ) + return cls(image_processor=image_processor, tokenizer=tokenizer) + + def __call__(self, text, images=None, **kwargs) -> BatchFeature: + if isinstance(text, str): + text = [text] + data: dict[str, Any] = {} + if images is not None: + image_outputs = self.image_processor( + images=images, + return_tensors="pt", + **kwargs.get("images_kwargs", {}), + ) + data.update(image_outputs) + text = self._expand_image_tokens(text, image_outputs["image_grid_thw"]) + data.update( + self.tokenizer( + text, + return_tensors="pt", + **kwargs.get("text_kwargs", {}), + ) + ) + return BatchFeature(data=data) + + def _expand_image_tokens(self, text: list[str], image_grid_thw: torch.Tensor) -> list[str]: + num_patches = image_grid_thw.prod(dim=1) // (self.spatial_merge_size**2) + image_index = 0 + expanded: list[str] = [] + for sample in text: + for _ in range(sample.count(self.image_token)): + replacement = ( + DEFAULT_IM_START_TOKEN + + DEFAULT_IMAGE_PATCH_TOKEN * int(num_patches[image_index].item()) + + DEFAULT_IM_END_TOKEN + + "\n" + ) + sample = sample.replace(self.image_token, replacement, 1) + image_index += 1 + expanded.append(sample) + return expanded + + def batch_decode(self, *args, **kwargs): + return self.tokenizer.batch_decode(*args, **kwargs) + + def decode(self, *args, **kwargs): + return self.tokenizer.decode(*args, **kwargs) + + @property + def model_input_names(self): + return list(dict.fromkeys(self.tokenizer.model_input_names + self.image_processor.model_input_names)) + + # transformers >= 5.x requires the config *class* (not a string) as the first # arg to AutoFeatureExtractor/AutoProcessor.register — it does `key.__module__` # internally. Passing "MingWhisperFeatureExtractor" raises: