[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/assets/hunyuan_image3/hunyuan_image3_instruct_cot_ref.txt

@@ -0,0 +1 @@
+首先，我分析所有输入图像：图像1是一个圆形的logo，设计现代且抽象。它由不同色调的蓝色（深蓝、中蓝、浅蓝）和白色构成，这些色块以流畅的曲线相互交织，形成一个动态的、类似旋涡或波浪的图案。整个logo是扁平化的矢量图形，背景为纯黑色。图像2展示了四个并排摆放的卡通动物造型冰箱贴，<recaption>这幅图像以产品摄影的精致风格，呈现了一枚根据`image_1`标志定制的圆形珐琅冰箱贴。最终图像使用`image_2`的分辨率。冰箱贴居中放置在`image_2`的灰色织物背景上，其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝和白色构成的动态旋涡图案。整个冰箱贴被赋予了`image_2`中冰箱贴特有的高级质感：表面覆盖着一层光滑如镜的珐琅釉面，反射出柔和而清晰的高光；图案的每一个色块边缘都由纤细的抛光金属边框精确勾勒，增强了立体感。柔和的顶光在冰箱贴的弧形边缘上形成平滑的过渡，并在其下方投下淡淡的、轮廓模糊的阴影，使其与织物背景无缝融合，营造出一种真实、静谧的视觉效果。<relation_1>最终图像完整保留了`image_1`中标志的全部设计元素。这包括其完美的圆形轮廓，以及内部由深蓝、中蓝、浅蓝和白色组成的精确旋涡状图案布局、形状和色彩关系。</relation_1><relation_2>最终图像的分辨率、背景和材质均来自`image_2`。背景中灰色织物的纹理和质感被完整保留。冰箱贴的材质被完美重构，精确复刻了`image_2`中冰箱贴所展示的光滑珐琅质感、抛光金属边框的视觉效果，以及整体柔和、均匀的布光环境和由此产生的自然阴影。</relation_2></recaption><answer><boi><img_size_1024><img_ratio_36><timestep>[<img>]{3600}<eoi></answer>

tests/assets/hunyuan_image3/hunyuan_image_distill_cot_ref.txt

@@ -0,0 +1 @@
+首先，我分析所有输入图像：图像1是一个圆形logo，背景为黑色，主体由深蓝、中蓝、浅蓝和白色组成的抽象曲线构成，整体呈现出一种动态的、类似旋涡的视觉效果。图像2展示了四个并排摆放在灰色织物表面上的冰箱贴。这些冰箱贴的造型分别是长颈鹿、小鸟、蜜蜂和海豚，其关键特征是材质：它们具有光滑、高反光的珐琅质感，边缘有明显的银色金属包边，呈现出立体感。用户的指令是“基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴”。这个指令要求将一个二维的平面设计（logo）转化为一个具有特定三维物理属性（冰箱贴材质）的实体物品。核心在于保留logo的设计精髓，同时精确复刻参考图像中冰箱贴的材质和工艺感。为了构建最终的答案图像，我会将图一的圆形logo作为新冰箱贴的主体图案。这个圆形冰箱贴将拥有与图二冰箱贴一致的厚度和立体感。其表面会覆盖一层高光泽的透明釉质，完美再现珐琅的质感。logo内部的蓝色和白色区域将保持其原有的色彩，但会因为表面的光泽而出现明亮的高光反射。整个圆形的边缘会包裹一圈纤细而有光泽的银色金属边框，这不仅增加了产品的精致感，也与图二中的金属包边工艺相呼应。最后，将这个新制作的圆形冰箱贴放置在图二所示的灰色织物背景上，确保其光照、阴影和反射效果与周围环境和其他冰箱贴完全一致，使其看起来就像是和它们一起被拍摄的，从而创造出一个逼真且风格统一的视觉效果。</think><recaption>这幅图像以产品摄影的精准风格，展示了一枚根据`image_1`的logo设计定制的圆形冰箱贴。最终图像使用`image_2`的分辨率。该冰箱贴居中放置于`image_2`的灰色织物背景之上，其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝及白色构成的动态旋涡状抽象图案。整个冰箱贴被赋予了与`image_2`中冰箱贴完全一致的高级珐琅质感，表面覆盖着一层晶莹剔透的高光釉面，在柔和的顶光照射下，呈现出清晰而温润的高光反射。一圈纤薄而精致的银色金属边框紧密包裹着圆形主体，其金属光泽与釉面的质感形成精妙对比。冰箱贴在织物背景上投下了一层极其自然的柔和阴影，证明了其与环境光影的无缝融合，整体视觉效果真实可信，宛如一件精心制作的实体工艺品。<relation_1>最终图像完整保留了`image_1`中logo的全部核心设计元素。这包括其完美的圆形轮廓，以及内部由深蓝、中蓝、浅蓝和白色四种颜色构成的抽象旋涡图案的精确形态、色彩和布局。</relation_1><relation_2>最终图像的分辨率、背景和材质风格完全来自`image_2`。这包括其作为背景的灰色织物表面的独特纹理，以及应用于新冰箱贴上的、与`image_2`中冰箱贴完全相同的珐琅材质（高光釉面质感）和银色金属边框的物理属性。画面的整体光照环境，包括柔和的光源方向和质感，也与`image_2`保持一致。</relation_2></recaption><answer><boi><img_size_1024><img_ratio_36><timestep><guidance><timestep_r>[<img>]{3600}<eoi></answer>

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)
 

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_distil_inputs.py

@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for HunyuanImage3 prepare_inputs_for_generation to ensure
+distilled model parameters (guidance, timesteps_r) are properly passed."""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
+    HunyuanImage3Text2ImagePipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class MockHunyuanImage3Pipeline:
+    """Minimal mock that replicates prepare_inputs_for_generation method."""
+
+    def __init__(self):
+        # Mock config for distilled models
+        self.config = type("Config", (), {})()
+        self.hf_config = type("HFConfig", (), {"cfg_distilled": True, "use_meanflow": True})()
+
+    # Bind the real method from the pipeline class
+    from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
+        HunyuanImage3Pipeline as _Real,
+    )
+
+    prepare_inputs_for_generation = _Real.prepare_inputs_for_generation
+    _update_model_kwargs_for_generation = _Real._update_model_kwargs_for_generation
+
+
+class TestPrepareInputsForGenerationDistilled:
+    """Tests that distilled model parameters are passed through
+    prepare_inputs_for_generation."""
+
+    def setup_method(self):
+        self.pipeline = MockHunyuanImage3Pipeline()
+
+    def _create_mock_kwargs(self, include_distilled_params: bool = True) -> dict:
+        """Create mock kwargs for prepare_inputs_for_generation."""
+        kwargs = {
+            "position_ids": torch.tensor([[0, 1, 2]]),
+            "custom_pos_emb": (torch.tensor([[0.0]]), torch.tensor([[1.0]])),
+            "mode": "gen_image",
+            "images": torch.randn(1, 4, 64, 64),
+            "image_mask": torch.ones(1, 64, 64),
+            "timestep": torch.tensor([0.5]),
+            "gen_timestep_scatter_index": torch.tensor([[5]]),
+            "cond_vae_images": None,
+            "cond_timestep": None,
+            "cond_vae_image_mask": None,
+            "cond_vit_images": None,
+            "cond_vit_image_mask": None,
+            "vit_kwargs": None,
+            "cond_timestep_scatter_index": None,
+            "query_lens": [10],
+            "seq_lens": [100],
+            "num_image_tokens": 4096,
+            "num_special_tokens": 3,
+            "ar_kv_reuse_len": 0,
+            "full_attn_spans": None,
+            "use_cache": True,
+        }
+        if include_distilled_params:
+            # CFG distilled model parameters
+            kwargs["guidance"] = torch.tensor([2500.0])
+            kwargs["guidance_scatter_index"] = torch.tensor([[4]])
+            # MeanFlow distilled model parameters
+            kwargs["timesteps_r"] = torch.tensor([0.3])
+            kwargs["timesteps_r_scatter_index"] = torch.tensor([[3]])
+        else:
+            kwargs["guidance"] = None
+            kwargs["guidance_scatter_index"] = None
+            kwargs["timesteps_r"] = None
+            kwargs["timesteps_r_scatter_index"] = None
+        return kwargs
+
+    def test_guidance_passed_to_model_inputs(self):
+        """Test that guidance is passed through prepare_inputs_for_generation."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "guidance" in model_inputs, "guidance should be in model_inputs"
+        assert model_inputs["guidance"] is not None, "guidance should not be None"
+        assert torch.equal(model_inputs["guidance"], kwargs["guidance"])
+
+    def test_guidance_scatter_index_passed_to_model_inputs(self):
+        """Test that guidance_scatter_index is passed through."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "guidance_scatter_index" in model_inputs
+        assert model_inputs["guidance_scatter_index"] is not None
+        assert torch.equal(model_inputs["guidance_scatter_index"], kwargs["guidance_scatter_index"])
+
+    def test_timesteps_r_passed_to_model_inputs(self):
+        """Test that timesteps_r is passed through for MeanFlow models."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "timesteps_r" in model_inputs, "timesteps_r should be in model_inputs"
+        assert model_inputs["timesteps_r"] is not None
+        assert torch.equal(model_inputs["timesteps_r"], kwargs["timesteps_r"])
+
+    def test_timesteps_r_scatter_index_passed_to_model_inputs(self):
+        """Test that timesteps_r_scatter_index is passed through."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "timesteps_r_scatter_index" in model_inputs
+        assert model_inputs["timesteps_r_scatter_index"] is not None
+        assert torch.equal(model_inputs["timesteps_r_scatter_index"], kwargs["timesteps_r_scatter_index"])
+
+    def test_num_special_tokens_passed_to_model_inputs(self):
+        """Test that num_special_tokens is passed through."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "num_special_tokens" in model_inputs
+        assert model_inputs["num_special_tokens"] == kwargs["num_special_tokens"]
+
+    def test_decode_update_keeps_all_special_token_positions(self):
+        """Test decode step uses timestep, guidance, and timestep_r positions."""
+        model_kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        model_kwargs.update(
+            {
+                "tokenizer_output": object(),
+                "attention_mask": torch.ones(1, 1, 7, 7, dtype=torch.bool),
+                "image_mask": torch.tensor([[False, False, False, False, True, True, True]]),
+                "gen_timestep_scatter_index": torch.tensor([[1]]),
+                "guidance_scatter_index": torch.tensor([[2]]),
+                "timesteps_r_scatter_index": torch.tensor([[3]]),
+            }
+        )
+
+        updated = self.pipeline._update_model_kwargs_for_generation({}, model_kwargs)
+
+        assert torch.equal(updated["position_ids"], torch.tensor([[1, 2, 3, 4, 5, 6]]))
+        assert updated["num_special_tokens"] == 3
+        assert updated["guidance_scatter_index"] is model_kwargs["guidance_scatter_index"]
+        assert updated["timesteps_r_scatter_index"] is model_kwargs["timesteps_r_scatter_index"]
+        assert updated["attention_mask"].shape == (1, 1, 6, 7)
+
+    def test_distilled_params_none_when_not_provided(self):
+        """Test that distilled params are None when not provided."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=False)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        # All distilled params should be present but None
+        assert "guidance" in model_inputs
+        assert model_inputs["guidance"] is None
+        assert "guidance_scatter_index" in model_inputs
+        assert model_inputs["guidance_scatter_index"] is None
+        assert "timesteps_r" in model_inputs
+        assert model_inputs["timesteps_r"] is None
+        assert "timesteps_r_scatter_index" in model_inputs
+        assert model_inputs["timesteps_r_scatter_index"] is None
+
+    def test_all_distilled_keys_present_in_model_inputs(self):
+        """Test that all 4 distilled keys are always present (even if None)."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=False)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        required_distilled_keys = [
+            "guidance",
+            "guidance_scatter_index",
+            "timesteps_r",
+            "timesteps_r_scatter_index",
+        ]
+        for key in required_distilled_keys:
+            assert key in model_inputs, f"{key} must be in model_inputs dict"
+
+
+def test_ar_kv_reuse_rebases_distilled_scatter_indices():
+    pipeline = object.__new__(HunyuanImage3Text2ImagePipeline)
+    model_kwargs = {
+        "query_lens": [8],
+        "attention_mask": torch.ones(1, 1, 8, 12, dtype=torch.bool),
+        "position_ids": torch.arange(8).unsqueeze(0),
+        "image_mask": torch.ones(1, 8, dtype=torch.bool),
+        "gen_timestep_scatter_index": torch.tensor([[5]]),
+        "guidance_scatter_index": torch.tensor([[6]]),
+        "timesteps_r_scatter_index": torch.tensor([[7]]),
+    }
+
+    truncated = pipeline._truncate_reused_prefix(
+        input_ids=torch.arange(8).unsqueeze(0),
+        model_kwargs=model_kwargs,
+        positive_reuse_len=3,
+    )
+
+    assert truncated.tolist() == [[3, 4, 5, 6, 7]]
+    assert model_kwargs["gen_timestep_scatter_index"].tolist() == [[2]]
+    assert model_kwargs["guidance_scatter_index"].tolist() == [[3]]
+    assert model_kwargs["timesteps_r_scatter_index"].tolist() == [[4]]
+
+
+def test_cfg_parallel_splits_meanflow_scatter_indices():
+    model_kwargs = {
+        "timesteps_r_scatter_index": torch.tensor([[1], [2], [3], [4]]),
+    }
+
+    HunyuanImage3Text2ImagePipeline._split_model_kwargs_for_cfg_parallel(
+        model_kwargs,
+        batch_size=2,
+        cfg_rank=1,
+    )
+
+    assert model_kwargs["timesteps_r_scatter_index"].tolist() == [[3], [4]]

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py

@@ -38,6 +38,7 @@ def _pipeline():
         cache_backend=None,
         diffusion_kv_cache_skip_step_indices=None,
     )
+    pipeline.hf_config = SimpleNamespace(cfg_distilled=False, use_meanflow=False)
     pipeline._pipeline = SimpleNamespace()
     return pipeline
 
@@ -451,3 +452,44 @@ def test_denoise_step_uses_input_batch_group_order_and_splits_back(monkeypatch):
     assert states[1].extra[_STEP_MODEL_KWARGS]["attention_mask"].shape == (2, 1, 2, 6)
     assert states[0].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(2, 4)], [(2, 4)]]
     assert states[1].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(4, 6)], [(4, 6)]]
+
+
+def test_distilled_step_supplies_guidance_and_meanflow_timestep(monkeypatch):
+    pipeline = _pipeline()
+    pipeline.hf_config = SimpleNamespace(cfg_distilled=True, use_meanflow=True)
+    monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda self: torch.device("cpu")))
+    state = _state("distilled", 1)
+    state.latents = torch.zeros(1, 1)
+    state.extra[_STEP_GUIDANCE_SCALE] = 2.5
+    state.extra[_STEP_MODEL_KWARGS].update(
+        {
+            "attention_mask": torch.ones(1, 1, 2, 4, dtype=torch.bool),
+            "full_attn_spans": [[(2, 4)]],
+        }
+    )
+    state.extra[_STEP_PROMPT_KV] = [
+        {
+            "key": torch.zeros(1, 2, 1, 1),
+            "value": torch.zeros(1, 2, 1, 1),
+            "lens": torch.tensor([2]),
+        }
+    ]
+    state.scheduler = SimpleNamespace(get_timestep_r=lambda _timestep: torch.tensor(0.25))
+    captured = {}
+
+    pipeline._restore_prompt_kv_cache = lambda *_args: None
+
+    def fake_prepare_inputs(input_ids, images, timestep, **model_kwargs):
+        del input_ids, images, timestep
+        captured.update(model_kwargs)
+        return {"model_kwargs": model_kwargs}
+
+    pipeline.prepare_inputs_for_generation = fake_prepare_inputs
+    pipeline.forward_call = lambda **_kwargs: {"diffusion_prediction": torch.tensor([[1.0]])}
+    pipeline._update_model_kwargs_for_generation = lambda _output, model_kwargs: model_kwargs
+
+    output = pipeline.denoise_step(InputBatch.make_batch([state]))
+
+    torch.testing.assert_close(output, torch.tensor([[1.0]]))
+    torch.testing.assert_close(captured["guidance"], torch.tensor([2500.0], dtype=torch.bfloat16))
+    torch.testing.assert_close(captured["timesteps_r"], torch.tensor([0.25]))

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -94,15 +94,28 @@ def _default_ar_dit_devices() -> tuple[str, str]:
 
 
 AR_DEVICES, DIT_DEVICES = _default_ar_dit_devices()
-MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
-HUNYUAN_IMAGE_REF_PATH = get_asset_path("hunyuan/hunyuan_image_ref.png")
-NUM_INFERENCE_STEPS = 50
+
+# Model configuration based on HUNYUAN_MODEL_PATH
+DISTIL_MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct-Distil"
+INSTRUCT_MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
+MODEL_PATH = os.environ.get("HUNYUAN_MODEL_PATH", DISTIL_MODEL_NAME)
+
+# Determine model type and set parameters accordingly
+IS_DISTIL = "distil" in MODEL_PATH.lower()
+MODEL_NAME = DISTIL_MODEL_NAME if IS_DISTIL else INSTRUCT_MODEL_NAME
+NUM_INFERENCE_STEPS = 8 if IS_DISTIL else 50
 GUIDANCE_SCALE = 2.5
 
+# Asset paths based on model type
+ASSET_SUBDIR = "hunyuan_image3"
+REF_IMAGE_NAME = "hunyuan_image_distill_ref.png" if IS_DISTIL else "hunyuan_image_instruct_ref.png"
+COT_REF_NAME = "hunyuan_image_distill_cot_ref.txt" if IS_DISTIL else "hunyuan_image3_instruct_cot_ref.txt"
+HUNYUAN_IMAGE_REF_PATH = get_asset_path(f"{ASSET_SUBDIR}/{REF_IMAGE_NAME}")
+HUNYUAN_COT_REF_PATH = get_asset_path(f"{ASSET_SUBDIR}/{COT_REF_NAME}")
+
 # ============================================================================
 # Constants
 # ============================================================================
-MODEL_PATH = os.environ.get("HUNYUAN_MODEL_PATH", MODEL_NAME)
 # Test input
 PROMPT = "基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴"
 TEST_IMAGE_URLS = [
@@ -113,23 +126,52 @@ SEED = 42
 AR_TP_SIZE = len(AR_DEVICES.split(","))
 DIT_TP_SIZE = len(DIT_DEVICES.split(","))
 
-# Precision thresholds
+# Precision thresholds organized by model variant and test type
 THRESHOLDS = {
-    # AR text comparison
-    "text_prefix_match": 10,  # First 10 characters must match exactly
-    "cot_semantic_sim": 0.9,  # Full CoT semantic similarity
-    # Image comparison
-    "clip_score": 90,  # CLIP image semantic similarity
-    "ssim": 0.26,  # Structural similarity
-    "psnr": 12.5,  # Peak signal-to-noise ratio (dB)
+    "instruct": {
+        # AR text comparison
+        "text_prefix_match": 10,  # First 10 characters must match exactly
+        "cot_semantic_sim": 0.9,  # Full CoT semantic similarity
+        # Image comparison
+        "clip_score": 90,  # CLIP image semantic similarity
+        "ssim": 0.26,  # Structural similarity
+        "psnr": 12.5,  # Peak signal-to-noise ratio (dB)
+    },
+    # The 8-step distilled model has more structural variance across GPU
+    # runtimes; CLIP and PSNR retain the same semantic and pixel safeguards.
+    "distil": {
+        "text_prefix_match": 10,
+        "cot_semantic_sim": 0.9,
+        "clip_score": 90,
+        "ssim": 0.23,
+        "psnr": 12.5,
+    },
+    # Thresholds for quantized DiT accuracy tests (bf16 vs quant comparison)
+    "quant": {
+        "psnr": 10.0,
+        "ssim": 0.20,
+        "clip_score": 20.0,  # Minimum CLIP score for quantized output
+        "clip_score_drop": float(
+            os.environ.get("HUNYUAN_IMAGE3_QUANT_CLIP_SCORE_DROP_THRESHOLD", "5.0")
+        ),  # Max allowed drop from bf16
+    },
 }
 
+
+# Helper to get thresholds for current model variant
+def _get_thresholds() -> dict:
+    """Return appropriate thresholds based on test context."""
+    return THRESHOLDS["distil" if IS_DISTIL else "instruct"]
+
+
+def _get_quant_thresholds() -> dict:
+    """Return thresholds for quantization accuracy tests."""
+    return THRESHOLDS["quant"]
+
+
+# Quant test configuration
 QUANT_PROMPT = "A brown and white dog is running on the grass."
 QUANT_HEIGHT, QUANT_WIDTH = 1024, 1024
-QUANT_PSNR_THRESHOLD = 10.0
-QUANT_SSIM_THRESHOLD = 0.20
-QUANT_CLIP_SCORE_THRESHOLD = 20.0
-QUANT_CLIP_SCORE_DROP_THRESHOLD = float(os.environ.get("HUNYUAN_IMAGE3_QUANT_CLIP_SCORE_DROP_THRESHOLD", "5.0"))
 QUANT_RUN_ENV = "HUNYUAN_IMAGE3_RUN_QUANT_ACCURACY"
 QUANT_BF16_ENV = "HUNYUAN_IMAGE3_BF16_MODEL"
 QUANT_FP8_ENV = "HUNYUAN_IMAGE3_FP8_MODEL"
@@ -167,7 +209,6 @@ _DEPLOY_CONFIG = {
                 "top_p": 1,
                 "top_k": -1,
                 "max_tokens": 8192,
-                "stop_token_ids": [128025],
                 "detokenize": True,
                 "skip_special_tokens": False,
                 "include_stop_str_in_output": True,
@@ -246,19 +287,6 @@ def _quant_accuracy_cases() -> list[pytest.ParameterSet]:
     return params
 
 
-# fmt: off
-COT_REF = ("首先，我分析所有输入图像：图像1是一个圆形的logo，设计现代且抽象。它由不同色调的蓝色（深蓝、中蓝、浅蓝）和白色构成，这些色块以流畅的曲线相互交织，形成一个动态的、类似旋涡或波浪的图案。整个logo是扁平化的矢量图形，背景为纯黑色。图像2展示了四个并排摆放的卡通动物造型冰箱贴，"
-           "它们被放置在灰色的织物背景上。这些冰箱贴的关键特征是其材质：它们具有光滑、高光的珐琅或烤漆质感，边缘有明显的金属包边，整体呈现出一种立体的、有厚度的实体感。用户的指令是“基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴”。这个指令要求将一个二维的平面设计（logo）"
-           "转化为一个具有特定物理属性（材质和立体感）的三维物体。核心任务是保留logo的视觉识别性，同时赋予其冰箱贴的实体质感。为了构建答案图像，我会将图一的圆形logo作为基础形状。然后，我会将图二中冰箱贴的材质特性应用到这个logo上。具体来说，logo中的每一个色块（深蓝、中蓝、浅蓝、白色）"
-           "都会被渲染成具有高光泽度的珐琅质感，表面会反射出柔和的环境光，形成自然的高光。logo中不同颜色区域之间的分界线，将被处理成纤细的、带有金属光泽的凸起边缘，这既能清晰地勾勒出图案，也符合珐琅工艺品的典型特征。整个冰箱贴会呈现出轻微的厚度和圆润的边缘，使其看起来像一个真实的、可触摸的物体。"
-           "最后，将这个制作完成的冰箱贴放置在图二所示的灰色织物背景上，并为其添加一个微妙的、柔和的阴影，以增强其立体感和与背景的融合度，最终呈现出一个精致、逼真的产品展示图。</think><recaption>这幅图像以产品摄影的精致风格，呈现了一枚根据`image_1`标志定制的圆形珐琅冰箱贴。最终图像使用`image_2`的分辨率。"
-           "冰箱贴居中放置在`image_2`的灰色织物背景上，其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝和白色构成的动态旋涡图案。整个冰箱贴被赋予了`image_2`中冰箱贴特有的高级质感：表面覆盖着一层光滑如镜的珐琅釉面，反射出柔和而清晰的高光；图案的每一个色块边缘都由纤细的抛光金属边框精确勾勒，增强了立体感。"
-           "柔和的顶光在冰箱贴的弧形边缘上形成平滑的过渡，并在其下方投下淡淡的、轮廓模糊的阴影，使其与织物背景无缝融合，营造出一种真实、静谧的视觉效果。<relation_1>最终图像完整保留了`image_1`中标志的全部设计元素。这包括其完美的圆形轮廓，以及内部由深蓝、中蓝、浅蓝和白色组成的精确旋涡状图案布局、形状和色彩关系。"
-           "</relation_1><relation_2>最终图像的分辨率、背景和材质均来自`image_2`。背景中灰色织物的纹理和质感被完整保留。冰箱贴的材质被完美重构，精确复刻了`image_2`中冰箱贴所展示的光滑珐琅质感、抛光金属边框的视觉效果，以及整体柔和、均匀的布光环境和由此产生的自然阴影。</relation_2></recaption><answer><boi>"
-           "<img_size_1024><img_ratio_36><timestep>[<img>]{3600}<eoi></answer>")
-# fmt: on
-
-
 def _make_config(enable_kv_reuse: bool, path: Path) -> None:
     config = copy.deepcopy(_DEPLOY_CONFIG)
     config["stages"][0]["omni_kv_config"] = {"need_send_cache": enable_kv_reuse}
@@ -432,34 +460,36 @@ def test_image_to_image_alignment_online(accuracy_artifact_root: Path) -> None:
     online_cot = online_cot.lstrip("\n")
     scorer = SemanticSimilarityScorer()
     clip_scorer = CLIPScorer()
-    cot_results = scorer.text_similarity(online_cot, COT_REF)
+    cot_ref = HUNYUAN_COT_REF_PATH.read_text(encoding="utf-8")
+    cot_results = scorer.text_similarity(online_cot, cot_ref)
     image_ref = Image.open(HUNYUAN_IMAGE_REF_PATH).convert("RGB")
     image_clip_score = clip_scorer.image_image_score(online_image, image_ref)
     ssim_value, psnr_value = compute_image_ssim_psnr(prediction=online_image, reference=image_ref, compare_mode="RGB")
 
+    thresholds = _get_thresholds()
     table = [
-        ["COT similarity to reference", f"{cot_results['cot_semantic_sim']:.4f}", 0.9644],
-        ["COT prefix match", f"{cot_results['text_prefix_match_count']:.4f}", 29],
-        ["Image-Image similarity", f"{image_clip_score:.4f}", 94.5538],
-        ["SSIM", f"{ssim_value:.4f}", 0.242],
-        ["PSNR (dB)", f"{psnr_value:.2f}", 14.1],
+        ["COT similarity to reference", f"{cot_results['cot_semantic_sim']:.4f}", thresholds["cot_semantic_sim"]],
+        ["COT prefix match", f"{cot_results['text_prefix_match_count']:.4f}", thresholds["text_prefix_match"]],
+        ["Image-Image similarity", f"{image_clip_score:.4f}", thresholds["clip_score"]],
+        ["SSIM", f"{ssim_value:.4f}", thresholds["ssim"]],
+        ["PSNR (dB)", f"{psnr_value:.2f}", thresholds["psnr"]],
     ]
-    print("[ONLINE] " + tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid"))
+    print("[ONLINE] " + tabulate(table, headers=["Metric", "Value", "Threshold"], tablefmt="grid"))
 
     _assert_cot_structural_markers(online_cot, "ONLINE")
 
-    assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], (
-        f"[ONLINE] COT semantic similarity {cot_results['cot_semantic_sim']:.4f} below threshold {THRESHOLDS['cot_semantic_sim']}"
+    assert cot_results["cot_semantic_sim"] >= thresholds["cot_semantic_sim"], (
+        f"[ONLINE] COT semantic similarity {cot_results['cot_semantic_sim']:.4f} below threshold {thresholds['cot_semantic_sim']}"
     )
-    assert cot_results["text_prefix_match_count"] >= THRESHOLDS["text_prefix_match"], (
-        f"[ONLINE] COT prefix match {cot_results['text_prefix_match_count']} below threshold {THRESHOLDS['text_prefix_match']}"
+    assert cot_results["text_prefix_match_count"] >= thresholds["text_prefix_match"], (
+        f"[ONLINE] COT prefix match {cot_results['text_prefix_match_count']} below threshold {thresholds['text_prefix_match']}"
     )
-    assert image_clip_score >= THRESHOLDS["clip_score"], (
-        f"[ONLINE] Image-Image similarity {image_clip_score:.4f} below threshold {THRESHOLDS['clip_score']}"
+    assert image_clip_score >= thresholds["clip_score"], (
+        f"[ONLINE] Image-Image similarity {image_clip_score:.4f} below threshold {thresholds['clip_score']}"
     )
-    assert ssim_value >= THRESHOLDS["ssim"], f"[ONLINE] SSIM {ssim_value:.4f} below threshold {THRESHOLDS['ssim']}"
-    assert psnr_value >= THRESHOLDS["psnr"], (
-        f"[ONLINE] PSNR {psnr_value:.2f} dB below threshold {THRESHOLDS['psnr']} dB"
+    assert ssim_value >= thresholds["ssim"], f"[ONLINE] SSIM {ssim_value:.4f} below threshold {thresholds['ssim']}"
+    assert psnr_value >= thresholds["psnr"], (
+        f"[ONLINE] PSNR {psnr_value:.2f} dB below threshold {thresholds['psnr']} dB"
     )
 
 
@@ -546,34 +576,36 @@ def test_image_to_image_alignment(accuracy_artifact_root: Path) -> None:
 
     scorer = SemanticSimilarityScorer()
     clip_scorer = CLIPScorer()
-    cot_results = scorer.text_similarity(omni_cot, COT_REF)
+    cot_ref = HUNYUAN_COT_REF_PATH.read_text(encoding="utf-8")
+    cot_results = scorer.text_similarity(omni_cot, cot_ref)
     image_ref = Image.open(HUNYUAN_IMAGE_REF_PATH).convert("RGB")
     image_clip_score = clip_scorer.image_image_score(omni_image, image_ref)
     ssim_value, psnr_value = compute_image_ssim_psnr(prediction=omni_image, reference=image_ref, compare_mode="RGB")
 
+    thresholds = _get_thresholds()
     table = [
-        ["COT similarity to reference", f"{cot_results['cot_semantic_sim']:.4f}", 0.9644],
-        ["COT prefix match", f"{cot_results['text_prefix_match_count']:.4f}", 29],
-        ["Image-Image similarity", f"{image_clip_score:.4f}", 94.5538],
-        ["SSIM", f"{ssim_value:.4f}", 0.242],
-        ["PSNR (dB)", f"{psnr_value:.2f}", 14.1],
+        ["COT similarity to reference", f"{cot_results['cot_semantic_sim']:.4f}", thresholds["cot_semantic_sim"]],
+        ["COT prefix match", f"{cot_results['text_prefix_match_count']:.4f}", thresholds["text_prefix_match"]],
+        ["Image-Image similarity", f"{image_clip_score:.4f}", thresholds["clip_score"]],
+        ["SSIM", f"{ssim_value:.4f}", thresholds["ssim"]],
+        ["PSNR (dB)", f"{psnr_value:.2f}", thresholds["psnr"]],
     ]
 
-    print(tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid"))
+    print(tabulate(table, headers=["Metric", "Value", "Threshold"], tablefmt="grid"))
 
     _assert_cot_structural_markers(omni_cot, "OFFLINE")
 
-    assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], (
-        f"COT semantic similarity {cot_results['cot_semantic_sim']:.4f} is below threshold {THRESHOLDS['cot_semantic_sim']}"
+    assert cot_results["cot_semantic_sim"] >= thresholds["cot_semantic_sim"], (
+        f"COT semantic similarity {cot_results['cot_semantic_sim']:.4f} is below threshold {thresholds['cot_semantic_sim']}"
     )
-    assert cot_results["text_prefix_match_count"] >= THRESHOLDS["text_prefix_match"], (
-        f"COT prefix match count {cot_results['text_prefix_match_count']} is below threshold {THRESHOLDS['text_prefix_match']}"
+    assert cot_results["text_prefix_match_count"] >= thresholds["text_prefix_match"], (
+        f"COT prefix match count {cot_results['text_prefix_match_count']} is below threshold {thresholds['text_prefix_match']}"
     )
-    assert image_clip_score >= THRESHOLDS["clip_score"], (
-        f"Image-Image similarity{image_clip_score:.4f} is below threshold {THRESHOLDS['clip_score']}"
+    assert image_clip_score >= thresholds["clip_score"], (
+        f"Image-Image similarity{image_clip_score:.4f} is below threshold {thresholds['clip_score']}"
     )
-    assert ssim_value >= THRESHOLDS["ssim"], f"SSIM {ssim_value:.4f} is below threshold {THRESHOLDS['ssim']}"
-    assert psnr_value >= THRESHOLDS["psnr"], f"PSNR {psnr_value:.2f} dB is below threshold {THRESHOLDS['psnr']} dB"
+    assert ssim_value >= thresholds["ssim"], f"SSIM {ssim_value:.4f} is below threshold {thresholds['ssim']}"
+    assert psnr_value >= thresholds["psnr"], f"PSNR {psnr_value:.2f} dB is below threshold {thresholds['psnr']} dB"
 
 
 @pytest.mark.parametrize("case", _quant_accuracy_cases())
@@ -607,12 +639,13 @@ def test_quantized_dit_matches_bf16_accuracy(
         prediction=quant_image,
         reference=bf16_image,
     )
+    quant_thresholds = _get_quant_thresholds()
     assert_similarity(
         model_name=f"{MODEL_NAME} {case.name} vs bf16",
         vllm_image=quant_image,
         diffusers_image=bf16_image,
-        ssim_threshold=QUANT_SSIM_THRESHOLD,
-        psnr_threshold=QUANT_PSNR_THRESHOLD,
+        ssim_threshold=quant_thresholds["ssim"],
+        psnr_threshold=quant_thresholds["psnr"],
         width=QUANT_WIDTH,
         height=QUANT_HEIGHT,
     )
@@ -647,15 +680,15 @@ def test_quantized_dit_matches_bf16_accuracy(
     print(f"  bf16 model:       {bf16_model}")
     print(f"  quant model:      {quant_model}")
     print(f"  BF16 CLIP score:  {bf16_clip_score:.4f}")
-    print(f"  quant CLIP score: {quant_clip_score:.4f} threshold>={QUANT_CLIP_SCORE_THRESHOLD:.4f}")
-    print(f"  CLIP score drop:  {clip_score_drop:.4f} threshold<={QUANT_CLIP_SCORE_DROP_THRESHOLD:.4f}")
+    print(f"  quant CLIP score: {quant_clip_score:.4f} threshold>={quant_thresholds['clip_score']:.4f}")
+    print(f"  CLIP score drop:  {clip_score_drop:.4f} threshold<={quant_thresholds['clip_score_drop']:.4f}")
     print(f"  metrics:          {metrics_path}")
 
-    assert quant_clip_score >= QUANT_CLIP_SCORE_THRESHOLD, (
+    assert quant_clip_score >= quant_thresholds["clip_score"], (
         f"{case.name} CLIP score below threshold: got {quant_clip_score:.4f}, "
-        f"expected >= {QUANT_CLIP_SCORE_THRESHOLD:.4f}"
+        f"expected >= {quant_thresholds['clip_score']:.4f}"
     )
-    assert clip_score_drop <= QUANT_CLIP_SCORE_DROP_THRESHOLD, (
+    assert clip_score_drop <= quant_thresholds["clip_score_drop"], (
         f"{case.name} CLIP score drop too large: got {clip_score_drop:.4f}, "
-        f"expected <= {QUANT_CLIP_SCORE_DROP_THRESHOLD:.4f}"
+        f"expected <= {quant_thresholds['clip_score_drop']:.4f}"
     )

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py

@@ -22,6 +22,7 @@ class TokenizerEncodeOutput(BaseOutput):
     tokens: torch.Tensor | None = None
     timestep_scatter_index: torch.Tensor | None = None
     guidance_scatter_index: torch.Tensor | None = None
+    timesteps_r_scatter_index: torch.Tensor | None = None  # For MeanFlow distilled models
     text_slices: list[slice] | None = None
     gen_image_slices: list[slice] | None = None
     joint_image_slices: list[slice] | None = None
@@ -196,6 +197,7 @@ class TokenizerWrapper:
         extra_token_pos,
         add_timestep_token=False,
         add_image_shape_token=False,
+        add_timestep_r_token=False,
         base_size=None,
         ratio_idx=None,
         image_type=None,
@@ -221,6 +223,10 @@ class TokenizerWrapper:
             token_seq.extend([self.special_token_map["<guidance>"]])
             extra_token_pos["guidance"].append(token_count)
             token_count += 1
+        if add_timestep_r_token:
+            token_seq.extend([self.special_token_map["<timestep_r>"]])
+            extra_token_pos["gen_timestep_r"].append(token_count)
+            token_count += 1
         return token_count
 
     def encode_sequence(
@@ -230,6 +236,7 @@ class TokenizerWrapper:
         total_length=None,
         add_timestep_token=False,
         add_guidance_token=False,
+        add_timestep_r_token=False,
         last_key_only_prefix=False,
         add_eos=True,
         use_front_boi_token=True,
@@ -329,6 +336,7 @@ class TokenizerWrapper:
                     2
                     + (1 if source.get("timestep", add_timestep_token) else 0)
                     + (1 if source.get("guidance", add_guidance_token) else 0)
+                    + (1 if source.get("timestep_r", add_timestep_r_token) else 0)
                     + (2 if source.get("image_shape", add_image_shape_token) else 0)
                 )
                 if drop_last is True and token_count + extra_count + source["length"] > total_length:
@@ -344,6 +352,7 @@ class TokenizerWrapper:
                     extra_token_pos=extra_token_pos,
                     add_timestep_token=source.get("timestep", add_timestep_token),
                     add_guidance_token=source.get("guidance", add_guidance_token),
+                    add_timestep_r_token=source.get("timestep_r", add_timestep_r_token),
                     add_image_shape_token=source.get("image_shape", add_image_shape_token),
                     base_size=source.get("base_size"),
                     ratio_idx=source.get("ratio_idx"),
@@ -739,6 +748,7 @@ class TokenizerWrapper:
                         length=section["token_length"],
                         timestep=section.get("add_timestep_token", False),
                         guidance=section.get("add_guidance_token", False),
+                        timestep_r=section.get("add_timestep_r_token", False),
                         front_boi=section.get("use_front_boi_token", False),
                         image_shape=section.get("add_image_shape_token", False),
                         base_size=section.get("base_size"),
@@ -777,6 +787,11 @@ class TokenizerWrapper:
         guidance_scatter_index = (
             torch.tensor(extra_token_pos["guidance"], dtype=torch.long) if "guidance" in extra_token_pos else None
         )
+        timesteps_r_scatter_index = (
+            torch.tensor(extra_token_pos["gen_timestep_r"], dtype=torch.long)
+            if "gen_timestep_r" in extra_token_pos
+            else None
+        )
         cond_timestep_scatter_index = (
             torch.tensor(extra_token_pos["cond_timestep"], dtype=torch.long)
             if "cond_timestep" in extra_token_pos
@@ -847,6 +862,7 @@ class TokenizerWrapper:
             tokens=full_seq_token_tensor,
             timestep_scatter_index=timestep_scatter_index,
             guidance_scatter_index=guidance_scatter_index,
+            timesteps_r_scatter_index=timesteps_r_scatter_index,
             text_slices=text_slices,
             gen_image_slices=gen_image_slices,
             joint_image_slices=joint_image_slices,

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -630,6 +630,7 @@ class ImageInfo:
 
         self.add_timestep_token = kwargs.get("add_timestep_token", True)
         self.add_guidance_token = kwargs.get("add_guidance_token", False)
+        self.add_timestep_r_token = kwargs.get("add_timestep_r_token", False)
         self.use_front_boi_token = kwargs.get("use_front_boi_token", True)
         self.add_image_shape_token = kwargs.get("add_image_shape_token", True)
 
@@ -667,6 +668,7 @@ class ImageInfo:
                 token_length=self.image_token_length,
                 add_timestep_token=self.add_timestep_token,
                 add_guidance_token=self.add_guidance_token,
+                add_timestep_r_token=self.add_timestep_r_token,
                 use_front_boi_token=self.use_front_boi_token,
                 add_image_shape_token=self.add_image_shape_token,
                 base_size=self.base_size,
@@ -1442,6 +1444,9 @@ class HunyuanImage3Config(PretrainedConfig):
         vit=None,
         vit_processor=None,
         vit_aligner=None,
+        # Distillation flags
+        cfg_distilled: bool = False,
+        use_meanflow: bool = False,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -1521,6 +1526,10 @@ class HunyuanImage3Config(PretrainedConfig):
         self.patch_embed_hidden_dim = patch_embed_hidden_dim
         self.image_base_size = image_base_size
 
+        # Distillation flags
+        self.cfg_distilled = cfg_distilled
+        self.use_meanflow = use_meanflow
+
         # token id
         self.eod_token_id = eod_token_id
         self.im_start_id = im_start_id
@@ -1555,7 +1564,7 @@ class HunyuanImage3ImageProcessor:
         )
         self.vision_encoder_processor = Siglip2ImageProcessorFast.from_dict(config.vit_processor)
 
-    def build_image_info(self, image_size):
+    def build_image_info(self, image_size, add_guidance_token=False, add_timestep_r_token=False):
         # parse image size (HxW, H:W, or <img_ratio_i>)
         if isinstance(image_size, str):
             if image_size.startswith("<img_ratio_"):
@@ -1592,6 +1601,8 @@ class HunyuanImage3ImageProcessor:
             token_height=token_height,
             base_size=base_size,
             ratio_index=ratio_idx,
+            add_guidance_token=add_guidance_token,
+            add_timestep_r_token=add_timestep_r_token,
         )
         return image_info
 
@@ -2244,6 +2255,8 @@ class HunyuanImage3Model(nn.Module):
             "timestep_emb",
             "time_embed",
             "time_embed_2",
+            # These embeddings belong to the outer pipeline model rather than
+            # this transformer core and are loaded by the outer weight loader.
             "guidance_emb",
             "timestep_r_emb",
         ]
@@ -2802,6 +2815,8 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
             "position_ids",
             "image_mask",
             "gen_timestep_scatter_index",
+            "guidance_scatter_index",
+            "timesteps_r_scatter_index",
             "cond_vae_image_mask",
             "cond_vit_image_mask",
             "cond_timestep_scatter_index",
@@ -2971,9 +2986,12 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
         model_kwargs["attention_mask"] = model_kwargs["attention_mask"][:, :, positive_reuse_len:, :]
         model_kwargs["position_ids"] = model_kwargs["position_ids"][:, positive_reuse_len:]
 
-        # Shift image_mask and gen_timestep_scatter_index to match truncated sequence
+        # Shift generated-image placeholder indices to match the truncated sequence.
         model_kwargs["image_mask"] = model_kwargs["image_mask"][:, positive_reuse_len:]
         model_kwargs["gen_timestep_scatter_index"] = model_kwargs["gen_timestep_scatter_index"] - positive_reuse_len
+        for key in ("guidance_scatter_index", "timesteps_r_scatter_index"):
+            if model_kwargs.get(key) is not None:
+                model_kwargs[key] = model_kwargs[key] - positive_reuse_len
         model_kwargs["ar_kv_reuse_offset"] = positive_reuse_len
 
         # cond-image have computed in ar, we may skip it by index in the future.
@@ -3016,7 +3034,8 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
         self.model.inject_ar_kv_into_layers(ar_kv_data, positive_reuse_len)
 
         # 3. negative cfg prefill
-        if self.do_classifier_free_guidance:
+        # For CFG distilled models, skip negative CFG prefill (cfg_factor=1, no negative prompt)
+        if self.do_classifier_free_guidance and not self.model.config.cfg_distilled:
             self._maybe_run_negative_cfg_prefill(
                 input_ids=input_ids,
                 model_kwargs=model_kwargs,
@@ -3124,7 +3143,12 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
         self._guidance_rescale = guidance_rescale
 
         # Detect CFG parallel configuration (only 2-branch layout is supported)
-        cfg_parallel_ready = self.do_classifier_free_guidance and get_classifier_free_guidance_world_size() == 2
+        # For CFG distilled models, skip CFG parallel (cfg_factor=1, no negative branch)
+        cfg_parallel_ready = (
+            self.do_classifier_free_guidance
+            and not self.model.config.cfg_distilled
+            and get_classifier_free_guidance_world_size() == 2
+        )
 
         # Define call parameters
         device = self._execution_device
@@ -3177,7 +3201,11 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
             attention_mask = attention_mask[s]
             self._split_model_kwargs_for_cfg_parallel(model_kwargs, batch_size, cfg_rank)
         else:
-            cfg_factor = 1 + self.do_classifier_free_guidance
+            # For CFG distilled models, cfg_factor is always 1 (CFG embedded in model)
+            if self.model.config.cfg_distilled:
+                cfg_factor = 1
+            else:
+                cfg_factor = 1 + self.do_classifier_free_guidance
             cfg_rank = None
 
         b, _, q_len1, seq_len = attention_mask.shape
@@ -3219,6 +3247,13 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                     # Sequential CFG: double the batch
                     latent_model_input = torch.cat([latents] * cfg_factor)
 
+                if self.model.config.use_meanflow:
+                    r = self.scheduler.get_timestep_r(t)
+                    r_expand = r.repeat(latent_model_input.shape[0])
+                else:
+                    r_expand = None
+                model_kwargs["timesteps_r"] = r_expand
+
                 t_expand = t.repeat(latent_model_input.shape[0])
 
                 # ---- TeaCache: decide whether to compute or reuse ----
@@ -3240,6 +3275,13 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                     tc_prev_mod = cur_mod.detach()
 
                 if should_compute:
+                    # Handle guidance for CFG distilled models
+                    if self.model.config.cfg_distilled:
+                        model_kwargs["guidance"] = torch.tensor(
+                            [1000.0 * self._guidance_scale],
+                            device=self.device,
+                            dtype=torch.bfloat16,
+                        ).repeat(latent_model_input.shape[0])
                     model_inputs = self.model.prepare_inputs_for_generation(
                         input_ids,
                         images=latent_model_input,
@@ -3259,7 +3301,12 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                     pred = tc_prev_pred
 
                 # Perform guidance
-                if cfg_parallel_ready:
+                # For CFG distilled models, guidance is already embedded in the model,
+                # so we skip the explicit CFG computation
+                if self.model.config.cfg_distilled:
+                    # CFG distilled: guidance is handled internally via guidance_emb
+                    pass
+                elif cfg_parallel_ready:
                     # CFG parallel: all_gather → all ranks combine locally (no broadcast needed)
                     gathered = cfg_group.all_gather(pred, separate_tensors=True)
                     pred = self.cfg_operator(gathered[0], gathered[1], self.guidance_scale, step=i)

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -75,6 +75,21 @@ _STEP_PROMPT_KV = "hunyuan_prompt_kv"
 _HUNYUAN_DEFAULT_OUTPUT_TYPE = "pil"
 
 
+def _config_flag(config: Any, name: str) -> bool:
+    """Read an optional Hunyuan config flag with legacy-safe defaults."""
+
+    return bool(getattr(config, name, False))
+
+
+class HunyuanImage3FlowMatchEulerDiscreteScheduler(FlowMatchEulerDiscreteScheduler):
+    def get_timestep_r(self, timestep: float | torch.FloatTensor):
+        if self.step_index is None:
+            self._init_step_index(timestep)
+        if self.step_index + 1 < len(self.timesteps):
+            return self.timesteps[self.step_index + 1]
+        return self.timesteps.new_zeros(())
+
+
 def default(val, d):
     return val if val is not None else d
 
@@ -264,6 +279,8 @@ def get_hunyuan_image_3_pre_process_func(od_config: OmniDiffusionConfig):
                 image_processor=image_processor,
                 generation_config=generation_config,
                 image_base_size=hf_config.image_base_size,
+                cfg_distilled=_config_flag(hf_config, "cfg_distilled"),
+                use_meanflow=_config_flag(hf_config, "use_meanflow"),
             )
             request.prepared_layout = prepared_layout
             request.diffusion_kv_requests = request_layout_utils.build_hunyuan_diffusion_kv_requests(
@@ -398,6 +415,11 @@ class HunyuanImage3Pipeline(
             out_norm=True,
         )
         self.time_embed_2 = TimestepEmbedder(hidden_size=self.hf_config.hidden_size)
+        # Distillation-specific embedding modules
+        if _config_flag(self.hf_config, "cfg_distilled"):
+            self.guidance_emb = TimestepEmbedder(hidden_size=self.hf_config.hidden_size)
+        if _config_flag(self.hf_config, "use_meanflow"):
+            self.timestep_r_emb = TimestepEmbedder(hidden_size=self.hf_config.hidden_size)
         self.lm_head = nn.Linear(self.hf_config.hidden_size, self.hf_config.vocab_size, bias=False)
         self.vllm_config = get_current_vllm_config()
         self.post_init()
@@ -419,6 +441,8 @@ class HunyuanImage3Pipeline(
             "model.ln_f",
             "time_embed",
             "time_embed_2",
+            "guidance_emb",
+            "timestep_r_emb",
             "final_layer.model",
         ]
         device = get_local_device()
@@ -428,11 +452,8 @@ class HunyuanImage3Pipeline(
             if mod:
                 mod.to(device)
 
-        unexpected_keywords = [
-            "guidance_emb",
-            "timestep_r_emb",
-        ]
-        skip_prefixes.extend(unexpected_keywords)
+        # Note: guidance_emb and timestep_r_emb are no longer skipped
+        # to support HunyuanImage-3.0-Distil and MeanFlow distilled models
         loader = AutoWeightsLoader(
             self,
             skip_prefixes=skip_prefixes,
@@ -452,7 +473,7 @@ class HunyuanImage3Pipeline(
     def pipeline(self):
         if self._pipeline is None:
             # shift hard code
-            self.scheduler = FlowMatchEulerDiscreteScheduler(
+            self.scheduler = HunyuanImage3FlowMatchEulerDiscreteScheduler(
                 num_train_timesteps=1000,
                 shift=self.generation_config.flow_shift,
                 use_dynamic_shifting=False,
@@ -902,6 +923,70 @@ class HunyuanImage3Pipeline(
 
         return x
 
+    def instantiate_guidance_tokens(
+        self,
+        x: torch.Tensor,
+        guidance: BatchRaggedTensor | None = None,
+        guidance_scatter_index: torch.Tensor | None = None,
+    ):
+        """Instantiate guidance embedding tokens for CFG distilled models.
+
+        Args:
+            x: Input sequence tensor (batch_size, seq_len, n_embd)
+            guidance: Guidance scale values (as tensor, typically guidance_scale * 1000)
+            guidance_scatter_index: Index positions to scatter guidance embeddings
+        """
+        if guidance is None or guidance_scatter_index is None:
+            return x
+        if not hasattr(self, "guidance_emb"):
+            return x
+
+        batch_size, seq_len, n_embd = x.shape
+        guidance_src = self.guidance_emb(guidance.reshape(-1))  # (bsz * n, n_embd)
+        x.scatter_(
+            dim=1,
+            index=guidance_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd),
+            src=guidance_src.reshape(batch_size, -1, n_embd),
+        )
+        return x
+
+    def instantiate_timestep_r_tokens(
+        self,
+        x: torch.Tensor,
+        timesteps_r: BatchRaggedTensor | None = None,
+        timesteps_r_scatter_index: torch.Tensor | None = None,
+    ):
+        """Instantiate timestep_r embedding tokens for MeanFlow distilled models.
+
+        Args:
+            x: Input sequence tensor (batch_size, seq_len, n_embd)
+            timesteps_r: Timestep r values for MeanFlow
+            timesteps_r_scatter_index: Index positions to scatter embeddings
+        """
+        if timesteps_r is None or timesteps_r_scatter_index is None:
+            return x
+        if not hasattr(self, "timestep_r_emb"):
+            return x
+
+        batch_size, seq_len, n_embd = x.shape
+
+        if isinstance(timesteps_r, list):
+            for i, timestep_r in enumerate(timesteps_r):
+                timestep_r_src = self.timestep_r_emb(timestep_r)  # (n, n_embd)
+                x[i : i + 1].scatter_(
+                    dim=1,
+                    index=timesteps_r_scatter_index[i].unsqueeze(0).unsqueeze(-1).repeat(1, 1, n_embd),
+                    src=timestep_r_src.reshape(1, -1, n_embd),
+                )
+        else:
+            timesteps_r_src = self.timestep_r_emb(timesteps_r.reshape(-1))  # (bsz * n, n_embd)
+            x.scatter_(
+                dim=1,
+                index=timesteps_r_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd),
+                src=timesteps_r_src.reshape(batch_size, -1, n_embd),
+            )
+        return x
+
     def instantiate_vit_image_tokens(
         self,
         x: torch.Tensor,
@@ -936,12 +1021,12 @@ class HunyuanImage3Pipeline(
 
         return x
 
-    def ragged_final_layer(self, x, image_mask, timestep, token_h, token_w, first_step):
+    def ragged_final_layer(self, x, image_mask, timestep, token_h, token_w, first_step, num_special_tokens: int = None):
         bsz, seq_len, n_embd = x.shape
         if first_step:
             image_output = x.masked_select(image_mask.unsqueeze(-1).bool()).reshape(bsz, -1, n_embd)
         else:
-            image_output = x[:, 1:, :]
+            image_output = x[:, num_special_tokens:, :]
         timestep_emb = self.time_embed_2(timestep)
         pred = self.final_layer(image_output, timestep_emb, token_h, token_w)
         return pred
@@ -1130,7 +1215,12 @@ class HunyuanImage3Pipeline(
                     batch_gen_image_info = [prepared_layout.generated_image_info]
                 else:
                     batch_gen_image_info = [
-                        self.image_processor.build_image_info(image_size) for _ in range(batch_size)
+                        self.image_processor.build_image_info(
+                            image_size,
+                            add_guidance_token=_config_flag(getattr(self, "hf_config", None), "cfg_distilled"),
+                            add_timestep_r_token=_config_flag(getattr(self, "hf_config", None), "use_meanflow"),
+                        )
+                        for _ in range(batch_size)
                     ]
 
             if batch_cond_image_info is not None:
@@ -1146,7 +1236,11 @@ class HunyuanImage3Pipeline(
             generator = [torch.Generator(self.device).manual_seed(seed) for seed in seeds]
 
         # 3. apply chat template
+        # For CFG distilled models, cfg_factor is always 1 (CFG is embedded in the model)
+        # For non-distilled models, cfg_factor depends on guidance_scale
         cfg_factor = {"gen_text": 1, "gen_image": 1 + int(guidance_scale > 1.0)}
+        if _config_flag(getattr(self, "hf_config", None), "cfg_distilled"):
+            cfg_factor["gen_image"] = 1
         bot_task = kwargs.pop("bot_task", "auto")
         if prepared_layout is None:
             # Pull template options from generation_config only when the
@@ -1268,6 +1362,14 @@ class HunyuanImage3Pipeline(
             if vit_kwargs is not None
             else None,
             cond_timestep_scatter_index=to_device(output.cond_timestep_scatter_index, device),
+            # for CFG distilled models
+            guidance_scatter_index=to_device(output.guidance_scatter_index, device)
+            if hasattr(output, "guidance_scatter_index") and output.guidance_scatter_index is not None
+            else None,
+            # for MeanFlow distilled models
+            timesteps_r_scatter_index=to_device(output.timesteps_r_scatter_index, device)
+            if hasattr(output, "timesteps_r_scatter_index") and output.timesteps_r_scatter_index is not None
+            else None,
             # for inner usage
             tokenizer_output=output,
             batch_gen_image_info=batch_gen_image_info,
@@ -1350,9 +1452,16 @@ class HunyuanImage3Pipeline(
                 "cond_vit_image_mask": kwargs.get("cond_vit_image_mask"),
                 "vit_kwargs": kwargs.get("vit_kwargs"),
                 "cond_timestep_scatter_index": kwargs.get("cond_timestep_scatter_index"),
+                # For CFG distilled models
+                "guidance": kwargs.get("guidance"),
+                "guidance_scatter_index": kwargs.get("guidance_scatter_index"),
+                # For MeanFlow distilled models
+                "timesteps_r": kwargs.get("timesteps_r"),
+                "timesteps_r_scatter_index": kwargs.get("timesteps_r_scatter_index"),
                 "query_lens": kwargs.get("query_lens"),
                 "seq_lens": kwargs.get("seq_lens"),
                 "num_image_tokens": kwargs.get("num_image_tokens"),
+                "num_special_tokens": kwargs.get("num_special_tokens"),
                 "ar_kv_reuse_len": kwargs.get("ar_kv_reuse_len", 0),
                 "full_attn_spans": kwargs.get("full_attn_spans"),
             }
@@ -1372,6 +1481,11 @@ class HunyuanImage3Pipeline(
             "mode": mode,
             "custom_pos_emb": model_kwargs["custom_pos_emb"],
             "num_image_tokens": model_kwargs["num_image_tokens"],
+            "num_special_tokens": model_kwargs["num_special_tokens"],
+            "guidance": model_kwargs.get("guidance"),
+            "guidance_scatter_index": model_kwargs.get("guidance_scatter_index"),
+            "timesteps_r": model_kwargs.get("timesteps_r"),
+            "timesteps_r_scatter_index": model_kwargs.get("timesteps_r_scatter_index"),
         }
         if "full_attn_spans" in model_kwargs:
             updated_model_kwargs["full_attn_spans"] = model_kwargs["full_attn_spans"]
@@ -1400,15 +1514,24 @@ class HunyuanImage3Pipeline(
                 bsz, seq_len = image_mask.shape
                 offset = model_kwargs.get("ar_kv_reuse_offset", 0)  # should be an absolute position.
                 index = torch.arange(offset, offset + seq_len, device=image_mask.device).unsqueeze(0).repeat(bsz, 1)
-                position_ids = index.masked_select(image_mask.bool()).reshape(bsz, -1)
-                timestep_position_ids = index[
-                    torch.arange(bsz), model_kwargs["gen_timestep_scatter_index"][:, -1]
-                ].unsqueeze(-1)
-                updated_model_kwargs["position_ids"] = torch.cat([timestep_position_ids, position_ids], dim=1)
+                image_position_ids = index.masked_select(image_mask.bool()).reshape(bsz, -1)
+                special_position_ids = [
+                    index[torch.arange(bsz), model_kwargs["gen_timestep_scatter_index"][:, -1]].unsqueeze(-1)
+                ]
+                if model_kwargs.get("guidance_scatter_index") is not None:
+                    special_position_ids.append(
+                        index[torch.arange(bsz), model_kwargs["guidance_scatter_index"][:, -1]].unsqueeze(-1)
+                    )
+                if model_kwargs.get("timesteps_r_scatter_index") is not None:
+                    special_position_ids.append(
+                        index[torch.arange(bsz), model_kwargs["timesteps_r_scatter_index"][:, -1]].unsqueeze(-1)
+                    )
+                special_position_ids = torch.cat(special_position_ids, dim=1)
+                updated_model_kwargs["position_ids"] = torch.cat([special_position_ids, image_position_ids], dim=1)
 
                 # attention mask
                 mask_list = []
-                current_starts = timestep_position_ids.reshape(-1)
+                current_starts = special_position_ids[:, 0]
                 max_current_start = int(current_starts.max().item())
                 for attention_mask_i, position_ids_i, current_start_i in zip(
                     model_kwargs["attention_mask"], updated_model_kwargs["position_ids"], current_starts
@@ -1461,6 +1584,7 @@ class HunyuanImage3Pipeline(
 
             image_info: ImageInfo = batch_gen_image_info[0]
             kwargs["num_image_tokens"] = request_layout_utils.hunyuan_num_image_tokens(image_info)
+            kwargs["num_special_tokens"] = request_layout_utils.hunyuan_num_special_tokens(image_info)
             # 50 and 5.0 hard code
             results = self.pipeline(
                 batch_size=len(batch_gen_image_info),
@@ -1514,9 +1638,16 @@ class HunyuanImage3Pipeline(
         query_lens: list[int] | None = None,
         seq_lens: list[int] | None = None,
         num_image_tokens: int | None = None,
+        num_special_tokens: int | None = None,
         uncond_cfg_prefill: bool = False,
         ar_kv_reuse_len: int = 0,
         full_attn_spans: list[list[tuple[int, int]]] | None = None,
+        # for CFG distilled models
+        guidance: torch.Tensor | None = None,
+        guidance_scatter_index: torch.Tensor | None = None,
+        # for MeanFlow distilled models
+        timesteps_r: torch.Tensor | None = None,
+        timesteps_r_scatter_index: torch.Tensor | None = None,
     ) -> tuple | CausalMMOutputWithPast:
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         # Sanity Check of Inputs
@@ -1578,11 +1709,23 @@ class HunyuanImage3Pipeline(
                     inputs_embeds, images, timestep, image_mask
                 )
                 inputs_embeds = self.instantiate_timestep_tokens(inputs_embeds, timestep, gen_timestep_scatter_index)
+                # Instantiate guidance and timestep_r tokens for distilled models
+                inputs_embeds = self.instantiate_guidance_tokens(inputs_embeds, guidance, guidance_scatter_index)
+                inputs_embeds = self.instantiate_timestep_r_tokens(
+                    inputs_embeds, timesteps_r, timesteps_r_scatter_index
+                )
             else:
                 t_emb = self.time_embed(timestep)
                 image_emb, token_h, token_w = self.patch_embed(images, t_emb)
                 timestep_emb = self.timestep_emb(timestep).reshape(bsz, -1, n_embd)
-                inputs_embeds = torch.cat([timestep_emb, image_emb], dim=1)
+                extra_embeds = []
+                if hasattr(self, "guidance_emb") and guidance is not None:
+                    guidance_emb = self.guidance_emb(guidance.reshape(-1)).reshape(bsz, -1, n_embd)
+                    extra_embeds.append(guidance_emb)
+                if hasattr(self, "timestep_r_emb") and timesteps_r is not None:
+                    timesteps_r_emb = self.timestep_r_emb(timesteps_r.reshape(-1)).reshape(bsz, -1, n_embd)
+                    extra_embeds.append(timesteps_r_emb)
+                inputs_embeds = torch.cat([timestep_emb, *extra_embeds, image_emb], dim=1)
 
         # Instantiate placeholder tokens: <timestep>, <img> for cond images
         # Should only run once with kv-cache enabled.
@@ -1641,7 +1784,7 @@ class HunyuanImage3Pipeline(
             )
             hidden_states = hidden_states.reshape(bsz, seq_len, n_embd)
             diffusion_prediction = self.ragged_final_layer(
-                hidden_states, image_mask, timestep, token_h, token_w, first_step
+                hidden_states, image_mask, timestep, token_h, token_w, first_step, num_special_tokens
             )
 
         if not return_dict:
@@ -1771,6 +1914,7 @@ class HunyuanImage3Pipeline(
         target_height = int(image_info.image_height)
         target_width = int(image_info.image_width)
         model_kwargs["num_image_tokens"] = request_layout_utils.hunyuan_num_image_tokens(image_info)
+        model_kwargs["num_special_tokens"] = request_layout_utils.hunyuan_num_special_tokens(image_info)
 
         timesteps, _ = retrieve_timesteps(
             self.scheduler,
@@ -1819,13 +1963,14 @@ class HunyuanImage3Pipeline(
         state.timesteps = timesteps
         state.step_index = 0
         state.scheduler = req_scheduler
-        state.do_true_cfg = guidance_scale > 1.0
+        cfg_factor = request_layout_utils.hunyuan_cfg_factor(image_info, guidance_scale)
+        state.do_true_cfg = cfg_factor > 1
         state.extra = {
             _STEP_MODEL_KWARGS: model_kwargs,
             _STEP_INPUT_IDS: input_ids,
             _STEP_GENERATOR: model_kwargs["generator"],
             _STEP_GUIDANCE_SCALE: guidance_scale,
-            _STEP_CFG_FACTOR: 1 + int(guidance_scale > 1.0),
+            _STEP_CFG_FACTOR: cfg_factor,
             _STEP_OUTPUT_SIZE: (target_height, target_width),
             _STEP_COT_TEXT_LIST: cot_text_list,
             _STEP_AR_KV: self._snapshot_injected_ar_kv(),
@@ -1990,6 +2135,21 @@ class HunyuanImage3Pipeline(
             row_branches,
             first_step,
         )
+        if _config_flag(self.hf_config, "cfg_distilled"):
+            model_kwargs["guidance"] = torch.tensor(
+                [1000.0 * state.extra[_STEP_GUIDANCE_SCALE] for state in states] * cfg_factor,
+                device=latents.device,
+                dtype=torch.bfloat16,
+            )
+        if _config_flag(self.hf_config, "use_meanflow"):
+            model_kwargs["timesteps_r"] = torch.cat(
+                [
+                    state.scheduler.get_timestep_r(state.current_timestep).reshape(1).to(device=latents.device)
+                    for _ in range(cfg_factor)
+                    for state in states
+                ],
+                dim=0,
+            )
         if first_step:
             self._restore_injected_ar_kv(states, row_state_indexes, row_branches)
         else:

vllm_omni/diffusion/models/hunyuan_image3/request_layout.py

@@ -50,6 +50,7 @@ def _image_info_to_payload(image_info: ImageInfo) -> dict[str, Any]:
         "ratio_index": _to_python_scalar(image_info.ratio_index),
         "add_timestep_token": image_info.add_timestep_token,
         "add_guidance_token": image_info.add_guidance_token,
+        "add_timestep_r_token": image_info.add_timestep_r_token,
         "use_front_boi_token": image_info.use_front_boi_token,
         "add_image_shape_token": image_info.add_image_shape_token,
     }
@@ -68,6 +69,7 @@ def _image_info_from_payload(payload: dict[str, Any]) -> ImageInfo:
         ratio_index=payload.get("ratio_index"),
         add_timestep_token=payload.get("add_timestep_token", True),
         add_guidance_token=payload.get("add_guidance_token", False),
+        add_timestep_r_token=payload.get("add_timestep_r_token", False),
         use_front_boi_token=payload.get("use_front_boi_token", True),
         add_image_shape_token=payload.get("add_image_shape_token", True),
     )
@@ -197,7 +199,21 @@ def resolve_hunyuan_guidance_scale(sampling: Any, default_scale: float = 5.0) ->
 def hunyuan_num_image_tokens(image_info: ImageInfo) -> int:
     """Return the generated-image span overwritten on every denoise step."""
 
-    return int(image_info.image_token_length) + int(image_info.add_timestep_token) + int(image_info.add_guidance_token)
+    return int(image_info.image_token_length) + hunyuan_num_special_tokens(image_info)
+
+
+def hunyuan_num_special_tokens(image_info: ImageInfo) -> int:
+    """Return the generated-image prefix tokens emitted before latent tokens."""
+
+    return (
+        int(image_info.add_timestep_token) + int(image_info.add_guidance_token) + int(image_info.add_timestep_r_token)
+    )
+
+
+def hunyuan_cfg_factor(image_info: ImageInfo, guidance_scale: float) -> int:
+    """Return the execution branch count for standard or embedded CFG."""
+
+    return 1 if image_info.add_guidance_token else 1 + int(guidance_scale > 1.0)
 
 
 def build_hunyuan_batch_rope_image_info(
@@ -242,6 +258,8 @@ def prepare_hunyuan_layout(
     image_processor: HunyuanImage3ImageProcessor,
     generation_config: GenerationConfig,
     image_base_size: int,
+    cfg_distilled: bool = False,
+    use_meanflow: bool = False,
 ) -> HunyuanPreparedLayout:
     """Build the CPU token/image layout reused by Scheduler and Worker."""
 
@@ -261,7 +279,12 @@ def prepare_hunyuan_layout(
     height = sampling.height or 1024
     width = sampling.width or 1024
     guidance_scale = resolve_hunyuan_guidance_scale(sampling)
-    generated_image_info = image_processor.build_image_info((height, width))
+    image_info_kwargs: dict[str, bool] = {}
+    if cfg_distilled:
+        image_info_kwargs["add_guidance_token"] = True
+    if use_meanflow:
+        image_info_kwargs["add_timestep_r_token"] = True
+    generated_image_info = image_processor.build_image_info((height, width), **image_info_kwargs)
     result = tokenizer_wrapper.apply_chat_template(
         batch_prompt=prompt,
         mode="gen_image",
@@ -273,7 +296,7 @@ def prepare_hunyuan_layout(
         bot_task=tokenizer_bot_task,
         image_base_size=image_base_size,
         sequence_template=getattr(generation_config, "sequence_template", "pretrain"),
-        cfg_factor=1 + int(guidance_scale > 1.0),
+        cfg_factor=hunyuan_cfg_factor(generated_image_info, guidance_scale),
         drop_think=getattr(generation_config, "drop_think", False),
     )
     tokenizer_output = result["output"]
@@ -337,7 +360,10 @@ def build_hunyuan_diffusion_kv_requests(
     """Build one persistent Scheduler KV request per Hunyuan execution row."""
 
     tokenizer_output = prepared_layout.tokenizer_output
-    cfg_factor = 1 + int(resolve_hunyuan_guidance_scale(request.sampling_params) > 1.0)
+    cfg_factor = hunyuan_cfg_factor(
+        prepared_layout.generated_image_info,
+        resolve_hunyuan_guidance_scale(request.sampling_params),
+    )
     if prepared_layout.num_branches != cfg_factor:
         raise ValueError(
             "Hunyuan tokenizer sequence count does not match CFG execution: "