[New Model] Ming-Image (inclusionAI/Ming-Image-0.1-Design) Support (#8021)

Signed-off-by: Yuanheng <jonathan.zhaoyh@gmail.com>

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.
 

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,