diff --git a/tests/assets/hunyuan_image3/hunyuan_image3_instruct_cot_ref.txt b/tests/assets/hunyuan_image3/hunyuan_image3_instruct_cot_ref.txt
new file mode 100644
index 000000000..fb4005138
--- /dev/null
+++ b/tests/assets/hunyuan_image3/hunyuan_image3_instruct_cot_ref.txt
@@ -0,0 +1 @@
+首先，我分析所有输入图像：图像1是一个圆形的logo，设计现代且抽象。它由不同色调的蓝色（深蓝、中蓝、浅蓝）和白色构成，这些色块以流畅的曲线相互交织，形成一个动态的、类似旋涡或波浪的图案。整个logo是扁平化的矢量图形，背景为纯黑色。图像2展示了四个并排摆放的卡通动物造型冰箱贴，这幅图像以产品摄影的精致风格，呈现了一枚根据`image_1`标志定制的圆形珐琅冰箱贴。最终图像使用`image_2`的分辨率。冰箱贴居中放置在`image_2`的灰色织物背景上，其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝和白色构成的动态旋涡图案。整个冰箱贴被赋予了`image_2`中冰箱贴特有的高级质感：表面覆盖着一层光滑如镜的珐琅釉面，反射出柔和而清晰的高光；图案的每一个色块边缘都由纤细的抛光金属边框精确勾勒，增强了立体感。柔和的顶光在冰箱贴的弧形边缘上形成平滑的过渡，并在其下方投下淡淡的、轮廓模糊的阴影，使其与织物背景无缝融合，营造出一种真实、静谧的视觉效果。最终图像完整保留了`image_1`中标志的全部设计元素。这包括其完美的圆形轮廓，以及内部由深蓝、中蓝、浅蓝和白色组成的精确旋涡状图案布局、形状和色彩关系。最终图像的分辨率、背景和材质均来自`image_2`。背景中灰色织物的纹理和质感被完整保留。冰箱贴的材质被完美重构，精确复刻了`image_2`中冰箱贴所展示的光滑珐琅质感、抛光金属边框的视觉效果，以及整体柔和、均匀的布光环境和由此产生的自然阴影。[![]()
]{3600}
diff --git a/tests/assets/hunyuan_image3/hunyuan_image_distill_cot_ref.txt b/tests/assets/hunyuan_image3/hunyuan_image_distill_cot_ref.txt
new file mode 100644
index 000000000..3cc04fa5c
--- /dev/null
+++ b/tests/assets/hunyuan_image3/hunyuan_image_distill_cot_ref.txt
@@ -0,0 +1 @@
+首先，我分析所有输入图像：图像1是一个圆形logo，背景为黑色，主体由深蓝、中蓝、浅蓝和白色组成的抽象曲线构成，整体呈现出一种动态的、类似旋涡的视觉效果。图像2展示了四个并排摆放在灰色织物表面上的冰箱贴。这些冰箱贴的造型分别是长颈鹿、小鸟、蜜蜂和海豚，其关键特征是材质：它们具有光滑、高反光的珐琅质感，边缘有明显的银色金属包边，呈现出立体感。用户的指令是“基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴”。这个指令要求将一个二维的平面设计（logo）转化为一个具有特定三维物理属性（冰箱贴材质）的实体物品。核心在于保留logo的设计精髓，同时精确复刻参考图像中冰箱贴的材质和工艺感。为了构建最终的答案图像，我会将图一的圆形logo作为新冰箱贴的主体图案。这个圆形冰箱贴将拥有与图二冰箱贴一致的厚度和立体感。其表面会覆盖一层高光泽的透明釉质，完美再现珐琅的质感。logo内部的蓝色和白色区域将保持其原有的色彩，但会因为表面的光泽而出现明亮的高光反射。整个圆形的边缘会包裹一圈纤细而有光泽的银色金属边框，这不仅增加了产品的精致感，也与图二中的金属包边工艺相呼应。最后，将这个新制作的圆形冰箱贴放置在图二所示的灰色织物背景上，确保其光照、阴影和反射效果与周围环境和其他冰箱贴完全一致，使其看起来就像是和它们一起被拍摄的，从而创造出一个逼真且风格统一的视觉效果。这幅图像以产品摄影的精准风格，展示了一枚根据`image_1`的logo设计定制的圆形冰箱贴。最终图像使用`image_2`的分辨率。该冰箱贴居中放置于`image_2`的灰色织物背景之上，其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝及白色构成的动态旋涡状抽象图案。整个冰箱贴被赋予了与`image_2`中冰箱贴完全一致的高级珐琅质感，表面覆盖着一层晶莹剔透的高光釉面，在柔和的顶光照射下，呈现出清晰而温润的高光反射。一圈纤薄而精致的银色金属边框紧密包裹着圆形主体，其金属光泽与釉面的质感形成精妙对比。冰箱贴在织物背景上投下了一层极其自然的柔和阴影，证明了其与环境光影的无缝融合，整体视觉效果真实可信，宛如一件精心制作的实体工艺品。最终图像完整保留了`image_1`中logo的全部核心设计元素。这包括其完美的圆形轮廓，以及内部由深蓝、中蓝、浅蓝和白色四种颜色构成的抽象旋涡图案的精确形态、色彩和布局。最终图像的分辨率、背景和材质风格完全来自`image_2`。这包括其作为背景的灰色织物表面的独特纹理，以及应用于新冰箱贴上的、与`image_2`中冰箱贴完全相同的珐琅材质（高光釉面质感）和银色金属边框的物理属性。画面的整体光照环境，包括柔和的光源方向和质感，也与`image_2`保持一致。[![]()
]{3600}
diff --git a/tests/assets/hunyuan_image3/hunyuan_image_distill_ref.png b/tests/assets/hunyuan_image3/hunyuan_image_distill_ref.png
new file mode 100644
index 000000000..a352826c0
Binary files /dev/null and b/tests/assets/hunyuan_image3/hunyuan_image_distill_ref.png differ
diff --git a/tests/assets/hunyuan_image3/hunyuan_image_instruct_ref.png b/tests/assets/hunyuan_image3/hunyuan_image_instruct_ref.png
new file mode 100644
index 000000000..f3787ef0c
Binary files /dev/null and b/tests/assets/hunyuan_image3/hunyuan_image_instruct_ref.png differ
diff --git a/tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py b/tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py
index 3876762fa..2a17cd691 100644
--- a/tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py
+++ b/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)
diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_distil_inputs.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_distil_inputs.py
new file mode 100644
index 000000000..7df82b767
--- /dev/null
+++ b/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]]
diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py
index b62fe950e..b840da6c0 100644
--- a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py
+++ b/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]))
diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py
index eed0ef6e8..3e37f22b6 100644
--- a/tests/e2e/accuracy/test_hunyuan_image3.py
+++ b/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中不同颜色区域之间的分界线，将被处理成纤细的、带有金属光泽的凸起边缘，这既能清晰地勾勒出图案，也符合珐琅工艺品的典型特征。整个冰箱贴会呈现出轻微的厚度和圆润的边缘，使其看起来像一个真实的、可触摸的物体。"
- "最后，将这个制作完成的冰箱贴放置在图二所示的灰色织物背景上，并为其添加一个微妙的、柔和的阴影，以增强其立体感和与背景的融合度，最终呈现出一个精致、逼真的产品展示图。这幅图像以产品摄影的精致风格，呈现了一枚根据`image_1`标志定制的圆形珐琅冰箱贴。最终图像使用`image_2`的分辨率。"
- "冰箱贴居中放置在`image_2`的灰色织物背景上，其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝和白色构成的动态旋涡图案。整个冰箱贴被赋予了`image_2`中冰箱贴特有的高级质感：表面覆盖着一层光滑如镜的珐琅釉面，反射出柔和而清晰的高光；图案的每一个色块边缘都由纤细的抛光金属边框精确勾勒，增强了立体感。"
- "柔和的顶光在冰箱贴的弧形边缘上形成平滑的过渡，并在其下方投下淡淡的、轮廓模糊的阴影，使其与织物背景无缝融合，营造出一种真实、静谧的视觉效果。最终图像完整保留了`image_1`中标志的全部设计元素。这包括其完美的圆形轮廓，以及内部由深蓝、中蓝、浅蓝和白色组成的精确旋涡状图案布局、形状和色彩关系。"
- "最终图像的分辨率、背景和材质均来自`image_2`。背景中灰色织物的纹理和质感被完整保留。冰箱贴的材质被完美重构，精确复刻了`image_2`中冰箱贴所展示的光滑珐琅质感、抛光金属边框的视觉效果，以及整体柔和、均匀的布光环境和由此产生的自然阴影。"
- "[![]()
]{3600}")
-# 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}"
)
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py
index e7bb75fbe..7fae00c36 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py
+++ b/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[""]])
extra_token_pos["guidance"].append(token_count)
token_count += 1
+ if add_timestep_r_token:
+ token_seq.extend([self.special_token_map[""]])
+ 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,
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
index a1ad2c986..856e6f189 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
+++ b/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 )
if isinstance(image_size, str):
if image_size.startswith(" 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: , ![]()
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:
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/request_layout.py b/vllm_omni/diffusion/models/hunyuan_image3/request_layout.py
index d09c20c52..32353ed62 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/request_layout.py
+++ b/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: "