[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.