[Feat] Add support and testcase for HunyuanImage-3.0-Instruct-Distil (#4048)

Signed-off-by: zengchuang <zengchuang3@huawei.com>
Signed-off-by: zengchuang-hw <zengchuang1005@gmail.com>

tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py

@@ -24,6 +24,7 @@ from vllm_omni.diffusion.models.hunyuan_image3.request_layout import (
     build_hunyuan_diffusion_kv_requests,
     extract_hunyuan_prompt_inputs,
     hunyuan_num_image_tokens,
+    hunyuan_num_special_tokens,
     normalize_hunyuan_cot_text,
     prepare_hunyuan_layout,
 )
@@ -114,7 +115,7 @@ class _FakeImageProcessor:
         self.image_sizes: list[tuple[int, int]] = []
         self.vision_encoder_processor = SimpleNamespace(patch_size=1)
 
-    def build_image_info(self, image_size):
+    def build_image_info(self, image_size, **kwargs):
         self.image_sizes.append(image_size)
         return ImageInfo(
             image_type="gen_image",
@@ -125,6 +126,7 @@ class _FakeImageProcessor:
             image_token_length=self.image_token_length,
             base_size=1024,
             ratio_index=0,
+            **kwargs,
         )
 
 
@@ -166,6 +168,29 @@ def _request(
     )
 
 
+def test_distilled_layout_uses_embedded_cfg_and_meanflow_tokens() -> None:
+    tokenizer, image_processor = _components([12])
+    request = _request(guidance_scale=2.5)
+
+    prepared_layout = prepare_hunyuan_layout(
+        request,
+        tokenizer_wrapper=tokenizer,
+        image_processor=image_processor,
+        generation_config=SimpleNamespace(sequence_template="instruct", drop_think=False),
+        image_base_size=1024,
+        cfg_distilled=True,
+        use_meanflow=True,
+    )
+
+    image_info = prepared_layout.generated_image_info
+    assert image_info.add_guidance_token
+    assert image_info.add_timestep_r_token
+    assert tokenizer.calls[0]["cfg_factor"] == 1
+    assert hunyuan_num_special_tokens(image_info) == 3
+    assert hunyuan_num_image_tokens(image_info) == 19
+    assert len(build_hunyuan_diffusion_kv_requests(request, prepared_layout)) == 1
+
+
 def test_builds_kv_request_lengths_without_model_execution() -> None:
     tokenizer, image_processor = _components([12])
     request = _request(guidance_scale=1.0)
@@ -385,9 +410,13 @@ def test_prepared_model_inputs_match_local_tokenization(
     torch.testing.assert_close(prepared_mask, local_mask)
     assert prepared_inputs["full_attn_spans"] == local_inputs["full_attn_spans"]
 
-    num_image_tokens = hunyuan_num_image_tokens(prepared_layout.generated_image_info)
-    local_inputs.update(attention_mask=local_mask, num_image_tokens=num_image_tokens)
-    prepared_inputs.update(attention_mask=prepared_mask, num_image_tokens=num_image_tokens)
+    image_info = prepared_layout.generated_image_info
+    generation_token_counts = {
+        "num_image_tokens": hunyuan_num_image_tokens(image_info),
+        "num_special_tokens": hunyuan_num_special_tokens(image_info),
+    }
+    local_inputs.update(attention_mask=local_mask, **generation_token_counts)
+    prepared_inputs.update(attention_mask=prepared_mask, **generation_token_counts)
     local_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), local_inputs)
     prepared_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), prepared_inputs)