From 91400aa02ef923df162befce0ed525b52505e291 Mon Sep 17 00:00:00 2001 From: liuhongda Date: Tue, 7 Jul 2026 03:13:30 +0000 Subject: [PATCH] after add cache, this version support flashinfer --- .../hunyuan_image3/hunyuan_image3_i2i.json | 28 + .../hunyuan_image3/hunyuan_image3_t2i.json | 28 + ...26-07-02-hunyuan-image3-think-recaption.md | 45 + lightx2v/common/ops/norm/rms_norm_weight.py | 5 +- lightx2v/infer.py | 43 + .../networks/hunyuan_image3/__init__.py | 3 + .../networks/hunyuan_image3/infer/kv_cache.py | 130 ++ .../hunyuan_image3/infer/module_io.py | 17 + .../hunyuan_image3/infer/post_infer.py | 57 + .../hunyuan_image3/infer/pre_infer.py | 228 +++ .../hunyuan_image3/infer/transformer_infer.py | 177 +++ .../networks/hunyuan_image3/infer/utils.py | 114 ++ .../models/networks/hunyuan_image3/model.py | 223 +++ .../hunyuan_image3/weights/__init__.py | 9 + .../networks/hunyuan_image3/weights/common.py | 534 +++++++ .../hunyuan_image3/weights/post_weights.py | 32 + .../hunyuan_image3/weights/pre_weights.py | 30 + .../weights/transformer_weights.py | 162 ++ .../models/runners/hunyuan_image3/__init__.py | 1 + .../hunyuan_image3/hunyuan_image3_runner.py | 1037 +++++++++++++ .../schedulers/hunyuan_image3/__init__.py | 3 + .../schedulers/hunyuan_image3/scheduler.py | 40 + lightx2v/utils/set_config.py | 71 + .../hunyuan_image3/run_hunyuan_image3_i2i.sh | 74 + .../hunyuan_image3/run_hunyuan_image3_t2i.sh | 58 + .../hunyuan_image3/run_hunyuan_image3_ti2i.sh | 84 ++ scripts/wan/run_wan_i2v.sh | 4 +- scripts/wan/run_wan_t2v.sh | 4 +- tests/test_hunyuan_image3_integration.py | 1297 +++++++++++++++++ 29 files changed, 4533 insertions(+), 5 deletions(-) create mode 100644 configs/hunyuan_image3/hunyuan_image3_i2i.json create mode 100644 configs/hunyuan_image3/hunyuan_image3_t2i.json create mode 100644 docs/superpowers/plans/2026-07-02-hunyuan-image3-think-recaption.md create mode 100644 lightx2v/models/networks/hunyuan_image3/__init__.py create mode 100644 lightx2v/models/networks/hunyuan_image3/infer/kv_cache.py create mode 100644 lightx2v/models/networks/hunyuan_image3/infer/module_io.py create mode 100644 lightx2v/models/networks/hunyuan_image3/infer/post_infer.py create mode 100644 lightx2v/models/networks/hunyuan_image3/infer/pre_infer.py create mode 100644 lightx2v/models/networks/hunyuan_image3/infer/transformer_infer.py create mode 100644 lightx2v/models/networks/hunyuan_image3/infer/utils.py create mode 100644 lightx2v/models/networks/hunyuan_image3/model.py create mode 100644 lightx2v/models/networks/hunyuan_image3/weights/__init__.py create mode 100644 lightx2v/models/networks/hunyuan_image3/weights/common.py create mode 100644 lightx2v/models/networks/hunyuan_image3/weights/post_weights.py create mode 100644 lightx2v/models/networks/hunyuan_image3/weights/pre_weights.py create mode 100644 lightx2v/models/networks/hunyuan_image3/weights/transformer_weights.py create mode 100644 lightx2v/models/runners/hunyuan_image3/__init__.py create mode 100644 lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py create mode 100644 lightx2v/models/schedulers/hunyuan_image3/__init__.py create mode 100644 lightx2v/models/schedulers/hunyuan_image3/scheduler.py create mode 100644 scripts/hunyuan_image3/run_hunyuan_image3_i2i.sh create mode 100644 scripts/hunyuan_image3/run_hunyuan_image3_t2i.sh create mode 100755 scripts/hunyuan_image3/run_hunyuan_image3_ti2i.sh create mode 100644 tests/test_hunyuan_image3_integration.py diff --git a/configs/hunyuan_image3/hunyuan_image3_i2i.json b/configs/hunyuan_image3/hunyuan_image3_i2i.json new file mode 100644 index 000000000..0022a7584 --- /dev/null +++ b/configs/hunyuan_image3/hunyuan_image3_i2i.json @@ -0,0 +1,28 @@ +{ + "infer_steps": 50, + "sample_guide_scale": 5.0, + "flow_shift": 1.0, + "image_size": "auto", + "feature_caching": "NoCaching", + "enable_kv_cache": true, + "enable_text_kv_cache": true, + "use_taylor_cache": false, + "taylor_cache_interval": 5, + "taylor_cache_order": 2, + "taylor_cache_enable_first_enhance": false, + "taylor_cache_first_enhance_steps": 3, + "taylor_cache_enable_tailing_enhance": false, + "taylor_cache_tailing_enhance_steps": 1, + "taylor_cache_low_freqs_order": 2, + "taylor_cache_high_freqs_order": 2, + "pipeline_parallel": true, + "enable_cfg": true, + "infer_align_image_size": true, + "bot_task": "think_recaption", + "use_system_prompt": "en_unified", + "max_new_tokens": 2048, + "text_do_sample": true, + "text_top_k": 1024, + "text_top_p": 0.95, + "text_temperature": 0.6 +} diff --git a/configs/hunyuan_image3/hunyuan_image3_t2i.json b/configs/hunyuan_image3/hunyuan_image3_t2i.json new file mode 100644 index 000000000..da7abe891 --- /dev/null +++ b/configs/hunyuan_image3/hunyuan_image3_t2i.json @@ -0,0 +1,28 @@ +{ + "infer_steps": 50, + "sample_guide_scale": 5.0, + "flow_shift": 1.0, + "target_height": 1024, + "target_width": 1024, + "feature_caching": "NoCaching", + "enable_kv_cache": true, + "enable_text_kv_cache": true, + "use_taylor_cache": false, + "taylor_cache_interval": 5, + "taylor_cache_order": 2, + "taylor_cache_enable_first_enhance": false, + "taylor_cache_first_enhance_steps": 3, + "taylor_cache_enable_tailing_enhance": false, + "taylor_cache_tailing_enhance_steps": 1, + "taylor_cache_low_freqs_order": 2, + "taylor_cache_high_freqs_order": 2, + "pipeline_parallel": true, + "enable_cfg": true, + "bot_task": "think_recaption", + "use_system_prompt": "en_unified", + "max_new_tokens": 2048, + "text_do_sample": true, + "text_top_k": 1024, + "text_top_p": 0.95, + "text_temperature": 0.6 +} diff --git a/docs/superpowers/plans/2026-07-02-hunyuan-image3-think-recaption.md b/docs/superpowers/plans/2026-07-02-hunyuan-image3-think-recaption.md new file mode 100644 index 000000000..087c53fc9 --- /dev/null +++ b/docs/superpowers/plans/2026-07-02-hunyuan-image3-think-recaption.md @@ -0,0 +1,45 @@ +# Hunyuan Image3 Think Recaption Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add a native LightX2V text-generation pre-stage for Hunyuan Image3 `recaption` and `think_recaption`, then feed generated COT text into the existing native image generation path. + +**Architecture:** Keep the DiT path native. Add small runner-level helpers for text input preparation, autoregressive token generation without KV cache, stage transitions from `` to ``, and generated COT injection via tokenizer `batch_cot_text` for `gen_image`. + +**Tech Stack:** Python, PyTorch, Hunyuan tokenizer/image processor helpers, LightX2V native HunyuanImage3 model/infer/weights. + +--- + +### Task 1: Regression Tests + +**Files:** +- Modify: `tests/test_hunyuan_image3_integration.py` + +- [ ] Add tests for resolving bot task, generating COT before image inputs, and stage transition token forcing. +- [ ] Run targeted pytest and verify failures before implementation. + +### Task 2: Runner Text Generation Helpers + +**Files:** +- Modify: `lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py` + +- [ ] Add helpers for system prompt resolution, text template creation, next-token selection, and COT text generation. +- [ ] Preserve existing `bot_task=image` behavior. +- [ ] For `think_recaption`, force `` after `` and stop after ``. + +### Task 3: Feed COT Into Image Generation + +**Files:** +- Modify: `lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py` +- Modify: `configs/hunyuan_image3/hunyuan_image3_t2i.json` + +- [ ] Pass `batch_cot_text` into `apply_chat_template(..., mode="gen_image")`. +- [ ] Add config keys for `bot_task`, `use_system_prompt`, `max_new_tokens`, and deterministic text generation. + +### Task 4: Verification + +**Files:** +- Test: `tests/test_hunyuan_image3_integration.py` + +- [ ] Run full Hunyuan integration tests. +- [ ] Run ruff and py_compile on touched files. diff --git a/lightx2v/common/ops/norm/rms_norm_weight.py b/lightx2v/common/ops/norm/rms_norm_weight.py index 340b747aa..54d500e59 100755 --- a/lightx2v/common/ops/norm/rms_norm_weight.py +++ b/lightx2v/common/ops/norm/rms_norm_weight.py @@ -66,14 +66,17 @@ def _get_lora_attr_mapping(self): self.weight_diff = torch.tensor(0.0, dtype=GET_DTYPE(), device=AI_DEVICE) def _get_actual_weight(self): - if not hasattr(self, "weight_diff"): + if not hasattr(self, "weight_diff") or not self.has_diff: return self.weight + if self.weight_diff.device != self.weight.device or self.weight_diff.dtype != self.weight.dtype: + self.weight_diff = self.weight_diff.to(device=self.weight.device, dtype=self.weight.dtype) return self.weight + self.weight_diff def register_diff(self, weight_dict): if not self.lazy_load or self.create_cuda_buffer or self.create_cpu_buffer: if self.weight_diff_name in weight_dict: self.weight_diff = weight_dict[self.weight_diff_name] + self.has_diff = True logger.debug(f"Register Diff to {self.weight_name}") def load(self, weight_dict): diff --git a/lightx2v/infer.py b/lightx2v/infer.py index caa12bce8..3b6ab46b9 100755 --- a/lightx2v/infer.py +++ b/lightx2v/infer.py @@ -12,6 +12,7 @@ from lightx2v.models.runners.flux2.flux2_runner import Flux2DevRunner, Flux2KleinRunner # noqa: F401 from lightx2v.models.runners.hidream_o1_image.hidream_o1_image_runner import HidreamO1ImageRunner # noqa: F401 from lightx2v.models.runners.hunyuan3d.hunyuan3d_shape_runner import Hunyuan3DShapeRunner # noqa: F401 +from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner # noqa: F401 from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_distill_runner import HunyuanVideo15DistillRunner # noqa: F401 from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner import HunyuanVideo15Runner # noqa: F401 from lightx2v.models.runners.longcat_image.longcat_image_runner import LongCatImageRunner # noqa: F401 @@ -47,6 +48,17 @@ from lightx2v_platform.registry_factory import PLATFORM_DEVICE_REGISTER +def _str2bool(value): + if isinstance(value, bool): + return value + lowered = value.lower() + if lowered in ("1", "true", "yes", "y", "on"): + return True + if lowered in ("0", "false", "no", "n", "off"): + return False + raise argparse.ArgumentTypeError(f"Expected a boolean value, got {value!r}") + + def init_runner(config): torch.set_grad_enabled(False) runner = RUNNER_REGISTER[config["model_cls"]](config) @@ -88,6 +100,7 @@ def main(): "wan2.2_s2v", "hunyuan_video_1.5", "hunyuan_video_1.5_distill", + "hunyuan_image3", "hunyuan3d", "worldplay_distill", "worldplay_ar", @@ -116,6 +129,36 @@ def main(): parser.add_argument("--support_tasks", type=str, nargs="+", default=[], help="Set supported tasks for the model") parser.add_argument("--model_path", type=str, required=True) parser.add_argument("--config_json", type=str, required=True) + hunyuan_cache_args = parser.add_argument_group("HunyuanImage3 cache options") + hunyuan_cache_args.add_argument("--enable_kv_cache", "--enable-kv-cache", dest="enable_kv_cache", type=_str2bool, nargs="?", const=True, default=None) + hunyuan_cache_args.add_argument("--enable_text_kv_cache", "--enable-text-kv-cache", dest="enable_text_kv_cache", type=_str2bool, nargs="?", const=True, default=None) + hunyuan_cache_args.add_argument("--use_taylor_cache", "--use-taylor-cache", dest="use_taylor_cache", type=_str2bool, nargs="?", const=True, default=None) + hunyuan_cache_args.add_argument("--taylor_cache_interval", "--taylor-cache-interval", dest="taylor_cache_interval", type=int, default=None) + hunyuan_cache_args.add_argument("--taylor_cache_order", "--taylor-cache-order", dest="taylor_cache_order", type=int, default=None) + hunyuan_cache_args.add_argument( + "--taylor_cache_enable_first_enhance", + "--taylor-cache-enable-first-enhance", + dest="taylor_cache_enable_first_enhance", + type=_str2bool, + nargs="?", + const=True, + default=None, + ) + hunyuan_cache_args.add_argument("--taylor_cache_first_enhance_steps", "--taylor-cache-first-enhance-steps", dest="taylor_cache_first_enhance_steps", type=int, default=None) + hunyuan_cache_args.add_argument( + "--taylor_cache_enable_tailing_enhance", + "--taylor-cache-enable-tailing-enhance", + dest="taylor_cache_enable_tailing_enhance", + type=_str2bool, + nargs="?", + const=True, + default=None, + ) + hunyuan_cache_args.add_argument("--taylor_cache_tailing_enhance_steps", "--taylor-cache-tailing-enhance-steps", dest="taylor_cache_tailing_enhance_steps", type=int, default=None) + hunyuan_cache_args.add_argument("--taylor_cache_low_freqs_order", "--taylor-cache-low-freqs-order", dest="taylor_cache_low_freqs_order", type=int, default=None) + hunyuan_cache_args.add_argument("--taylor_cache_high_freqs_order", "--taylor-cache-high-freqs-order", dest="taylor_cache_high_freqs_order", type=int, default=None) + hunyuan_native_args = parser.add_argument_group("HunyuanImage3 native implementation options") + hunyuan_native_args.add_argument("--moe_impl", "--moe-impl", dest="moe_impl", type=str, choices=["eager", "flashinfer"], default=None) parser.add_argument("--use_prompt_enhancer", action="store_true") parser.add_argument("--prompt", type=str, default="", help="The input prompt for text-to-video generation") parser.add_argument("--negative_prompt", type=str, default="") diff --git a/lightx2v/models/networks/hunyuan_image3/__init__.py b/lightx2v/models/networks/hunyuan_image3/__init__.py new file mode 100644 index 000000000..05842eaa6 --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/__init__.py @@ -0,0 +1,3 @@ +from .model import HunyuanImage3Model + +__all__ = ["HunyuanImage3Model"] diff --git a/lightx2v/models/networks/hunyuan_image3/infer/kv_cache.py b/lightx2v/models/networks/hunyuan_image3/infer/kv_cache.py new file mode 100644 index 000000000..c57e306b8 --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/infer/kv_cache.py @@ -0,0 +1,130 @@ +from dataclasses import dataclass +import math + +import torch + + +@dataclass +class HunyuanImage3KVCacheLayer: + key: torch.Tensor | None = None + value: torch.Tensor | None = None + + +class HunyuanImage3StaticKVCache: + """Per-layer KV cache matching HunyuanImage3 gen_text/gen_image inference.""" + + def __init__(self, num_layers, max_cache_len, dynamic=False): + self.num_layers = int(num_layers) + self.max_cache_len = int(max_cache_len) + self.dynamic = bool(dynamic) + self.layers = [HunyuanImage3KVCacheLayer() for _ in range(self.num_layers)] + + def _ensure_layer(self, layer_idx, key_states, value_states): + layer = self.layers[layer_idx] + if layer.key is None: + key_shape = (*key_states.shape[:2], self.max_cache_len, key_states.shape[-1]) + value_shape = (*value_states.shape[:2], self.max_cache_len, value_states.shape[-1]) + layer.key = torch.zeros(key_shape, device=key_states.device, dtype=key_states.dtype) + layer.value = torch.zeros(value_shape, device=value_states.device, dtype=value_states.dtype) + return layer + + def update(self, key_states, value_states, layer_idx, cache_position=None): + layer = self._ensure_layer(layer_idx, key_states, value_states) + if cache_position is None: + layer.key[:, :, : key_states.shape[2]].copy_(key_states) + layer.value[:, :, : value_states.shape[2]].copy_(value_states) + return self._slice_dynamic(layer, key_states.shape[2]) + + cache_position = cache_position.to(device=key_states.device, dtype=torch.long) + if cache_position.dim() == 1: + layer.key.index_copy_(2, cache_position, key_states) + layer.value.index_copy_(2, cache_position, value_states) + return self._slice_dynamic(layer, int(cache_position[-1].item()) + 1) + + if cache_position.dim() != 2: + raise ValueError(f"HunyuanImage3 cache_position must be 1D or 2D, got {cache_position.shape}.") + if cache_position.shape[0] != key_states.shape[0]: + raise ValueError( + f"HunyuanImage3 cache batch mismatch: cache_position={cache_position.shape}, key_states={key_states.shape}." + ) + + for batch_idx in range(cache_position.shape[0]): + layer.key[batch_idx].index_copy_(1, cache_position[batch_idx], key_states[batch_idx]) + layer.value[batch_idx].index_copy_(1, cache_position[batch_idx], value_states[batch_idx]) + return self._slice_dynamic(layer, int(cache_position.max().item()) + 1) + + def _slice_dynamic(self, layer, end): + if not self.dynamic: + return layer.key, layer.value + end = min(int(end), self.max_cache_len) + return layer.key[:, :, :end], layer.value[:, :, :end] + + +def _decompose_freqs(x, cutoff_ratio=0.1): + original_dtype = x.dtype + x_fp32 = x.float() + freq = torch.fft.fft(x_fp32, dim=1) + freqs = torch.fft.fftfreq(x_fp32.shape[1], d=1.0, device=x.device) + cutoff = cutoff_ratio * freqs.abs().max() + low_mask = (freqs.abs() <= cutoff)[None, :, None] + high_mask = ~low_mask + low = torch.fft.ifft(freq * low_mask, dim=1).real.to(dtype=original_dtype) + high = torch.fft.ifft(freq * high_mask, dim=1).real.to(dtype=original_dtype) + return low, high + + +class HunyuanImage3TaylorCache: + """Frequency-split Taylor hidden-state cache used by HunyuanImage3 sampling.""" + + def __init__(self, max_order): + self.max_order = int(max_order) + self.low_derivatives = [None for _ in range(self.max_order + 1)] + self.high_derivatives = [None for _ in range(self.max_order + 1)] + self.last_past_key_values = None + + def taylor_formula(self, distance): + low_output = None + high_output = None + for order, derivative in enumerate(self.low_derivatives): + if derivative is None: + break + term = (distance**order / math.factorial(order)) * derivative + low_output = term if low_output is None else low_output + term + for order, derivative in enumerate(self.high_derivatives): + if derivative is None: + break + term = (distance**order / math.factorial(order)) * derivative + high_output = term if high_output is None else high_output + term + if low_output is None and high_output is None: + raise RuntimeError("HunyuanImage3 Taylor cache has no derivatives to extrapolate from.") + if low_output is None: + return high_output + if high_output is None: + return low_output + return low_output + high_output + + def derivatives_computation(self, hidden_states, distance, low_freqs_order, high_freqs_order): + low, high = _decompose_freqs(hidden_states) + new_low = [None for _ in range(self.max_order + 1)] + new_high = [None for _ in range(self.max_order + 1)] + new_low[0] = low + new_high[0] = high + safe_distance = max(int(distance), 1) + + for order in range(min(int(low_freqs_order), self.max_order)): + if self.low_derivatives[order] is None: + break + new_low[order + 1] = (new_low[order] - self.low_derivatives[order]) / safe_distance + + for order in range(min(int(high_freqs_order), self.max_order)): + if self.high_derivatives[order] is None: + break + new_high[order + 1] = (new_high[order] - self.high_derivatives[order]) / safe_distance + + self.low_derivatives = new_low + self.high_derivatives = new_high + + def clear_derivatives(self): + self.low_derivatives = [None for _ in range(self.max_order + 1)] + self.high_derivatives = [None for _ in range(self.max_order + 1)] + self.last_past_key_values = None diff --git a/lightx2v/models/networks/hunyuan_image3/infer/module_io.py b/lightx2v/models/networks/hunyuan_image3/infer/module_io.py new file mode 100644 index 000000000..cadb6cc98 --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/infer/module_io.py @@ -0,0 +1,17 @@ +from dataclasses import dataclass + +import torch + + +@dataclass +class HunyuanImage3PreInferOutput: + hidden_states: torch.Tensor + attention_mask: torch.Tensor | None = None + position_ids: torch.Tensor | None = None + custom_pos_emb: tuple[torch.Tensor, torch.Tensor] | None = None + past_key_values: object | None = None + use_cache: bool = False + image_mask: torch.Tensor | None = None + timesteps: torch.Tensor | None = None + token_hw: tuple[int, int] | None = None + first_step: bool | None = None diff --git a/lightx2v/models/networks/hunyuan_image3/infer/post_infer.py b/lightx2v/models/networks/hunyuan_image3/infer/post_infer.py new file mode 100644 index 000000000..8161fedcb --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/infer/post_infer.py @@ -0,0 +1,57 @@ +import torch +import torch.nn.functional as F + +from lightx2v.models.networks.hunyuan_image3.infer.utils import apply_linear, apply_timestep_embedder, first_weight_device, to_device + + +class HunyuanImage3PostInfer: + def __init__(self, config): + self.config = config + + def set_scheduler(self, scheduler): + self.scheduler = scheduler + + def _apply_resblock(self, block, x, emb): + h = block.in_norm.apply_group_norm(x) + h = F.silu(h) + h = block.in_conv.apply(h) + emb_out = apply_linear(block.emb_proj, F.silu(emb)) + while emb_out.ndim < h.ndim: + emb_out = emb_out[..., None] + scale, shift = emb_out.chunk(2, dim=1) + h = block.out_norm.apply_group_norm(h) * (1.0 + scale) + shift + h = F.silu(h) + h = block.out_conv.apply(h) + if block.skip_connection is None: + return x + h + return block.skip_connection.apply(x) + h + + def _final_layer(self, weights, image_output, timesteps, token_h, token_w): + t_emb = apply_timestep_embedder(weights.time_embed_2, timesteps) + x = image_output.reshape(image_output.shape[0], token_h, token_w, image_output.shape[-1]) + x = x.permute(0, 3, 1, 2).contiguous() + for block in weights.final_layer.blocks: + x = self._apply_resblock(block, x, t_emb) + x = weights.final_layer.out_norm.apply_group_norm(x) + x = F.silu(x) + return weights.final_layer.output_conv.apply(x) + + @torch.no_grad() + def infer(self, weights, hidden_states, pre_infer_out): + device = first_weight_device(weights) + hidden_states = to_device(hidden_states, device) + if pre_infer_out.image_mask is not None and pre_infer_out.timesteps is not None and pre_infer_out.token_hw is not None: + token_h, token_w = pre_infer_out.token_hw + image_mask = to_device(pre_infer_out.image_mask, device) + timesteps = to_device(pre_infer_out.timesteps, device) + if pre_infer_out.first_step is False: + special_tokens = 1 + int(self.config.get("cfg_distilled", False)) + int(self.config.get("use_meanflow", False)) + image_output = hidden_states[:, special_tokens:, :] + else: + hidden = hidden_states.shape[-1] + image_output = hidden_states.masked_select(image_mask.unsqueeze(-1).bool()).reshape(-1, token_h * token_w, hidden) + return {"diffusion_prediction": self._final_layer(weights, image_output, timesteps, token_h, token_w)} + + normed = weights.final_norm.apply(hidden_states) + logits = apply_linear(weights.lm_head, normed.reshape(-1, normed.shape[-1])).reshape(*normed.shape[:-1], -1) + return {"logits": logits} diff --git a/lightx2v/models/networks/hunyuan_image3/infer/pre_infer.py b/lightx2v/models/networks/hunyuan_image3/infer/pre_infer.py new file mode 100644 index 000000000..b2da2da73 --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/infer/pre_infer.py @@ -0,0 +1,228 @@ +import torch +import torch.nn.functional as F + +from lightx2v.models.networks.hunyuan_image3.infer.module_io import HunyuanImage3PreInferOutput +from lightx2v.models.networks.hunyuan_image3.infer.utils import apply_linear, apply_timestep_embedder, first_weight_device, to_device + + +class HunyuanImage3PreInfer: + def __init__(self, config): + self.config = config + + def set_scheduler(self, scheduler): + self.scheduler = scheduler + + def _apply_resblock(self, block, x, emb): + h = block.in_norm.apply_group_norm(x) + h = F.silu(h) + h = block.in_conv.apply(h) + emb_out = apply_linear(block.emb_proj, F.silu(emb)) + while emb_out.ndim < h.ndim: + emb_out = emb_out[..., None] + scale, shift = emb_out.chunk(2, dim=1) + h = block.out_norm.apply_group_norm(h) * (1.0 + scale) + shift + h = F.silu(h) + h = block.out_conv.apply(h) + if block.skip_connection is None: + return x + h + return block.skip_connection.apply(x) + h + + def _patch_embed(self, weights, images, timesteps): + t_emb = apply_timestep_embedder(weights.time_embed, timesteps) + x = weights.patch_embed.input_conv.apply(images) + for block in weights.patch_embed.blocks: + x = self._apply_resblock(block, x, t_emb) + token_h, token_w = x.shape[-2:] + x = x.flatten(2).transpose(1, 2).contiguous() + return x, token_h, token_w + + def _instantiate_image_tokens(self, weights, hidden_states, images, image_mask, timesteps): + if hidden_states is None: + image_seq, token_h, token_w = self._patch_embed(weights, images, timesteps) + if image_mask is not None: + batch, _, hidden = image_seq.shape + image_mask = image_mask.to(device=image_seq.device) + out = image_seq.new_zeros(batch, image_mask.shape[1], hidden) + index = torch.arange(image_mask.shape[1], device=image_seq.device).unsqueeze(0).repeat(batch, 1) + image_scatter_index = index.masked_select(image_mask.bool()).reshape(batch, -1) + if image_scatter_index.shape[1] != image_seq.shape[1]: + raise ValueError( + f"HunyuanImage3 image_mask selects {image_scatter_index.shape[1]} tokens, " + f"but image patch embed produced {image_seq.shape[1]} tokens." + ) + out.scatter_( + dim=1, + index=image_scatter_index.unsqueeze(-1).repeat(1, 1, hidden), + src=image_seq, + ) + return out, (token_h, token_w) + timestep_emb = apply_timestep_embedder(weights.timestep_emb, timesteps).reshape(images.size(0), -1, self.config["hidden_size"]) + return torch.cat([timestep_emb, image_seq], dim=1), (token_h, token_w) + + batch, seqlen, hidden = hidden_states.shape + index = torch.arange(seqlen, device=hidden_states.device).unsqueeze(0).repeat(batch, 1) + if isinstance(images, torch.Tensor): + image_seq, token_h, token_w = self._patch_embed(weights, images, timesteps) + image_scatter_index = index.masked_select(image_mask.bool()).reshape(batch, -1) + hidden_states.scatter_( + dim=1, + index=image_scatter_index.unsqueeze(-1).repeat(1, 1, hidden), + src=image_seq, + ) + return hidden_states, (token_h, token_w) + + token_hw = None + for batch_idx, (image, timestep) in enumerate(zip(images, timesteps)): + timestep = timestep.to(hidden_states.device) + if isinstance(image, torch.Tensor): + image = image.to(hidden_states.device) + image_seq, token_h, token_w = self._patch_embed(weights, image, timestep) + token_hw = (token_h, token_w) + elif isinstance(image, list): + image_seq_parts = [] + for image_idx, image_item in enumerate(image): + image_item = image_item.unsqueeze(0).to(hidden_states.device) + image_seq_item, token_h, token_w = self._patch_embed(weights, image_item, timestep[image_idx : image_idx + 1]) + token_hw = (token_h, token_w) + image_seq_parts.append(image_seq_item) + image_seq = torch.cat(image_seq_parts, dim=1) + else: + raise TypeError(f"HunyuanImage3 image item should be a tensor or list, got {type(image)}") + + image_scatter_index = index[batch_idx : batch_idx + 1].masked_select(image_mask[batch_idx : batch_idx + 1].bool()).reshape(1, -1) + hidden_states[batch_idx : batch_idx + 1].scatter_( + dim=1, + index=image_scatter_index.unsqueeze(-1).repeat(1, 1, hidden), + src=image_seq.reshape(1, -1, hidden), + ) + return hidden_states, token_hw + + def _instantiate_vit_image_tokens(self, hidden_states, image_embeds, image_mask): + if image_embeds is None or image_mask is None: + return hidden_states + + batch, seqlen, hidden = hidden_states.shape + index = torch.arange(seqlen, device=hidden_states.device).unsqueeze(0).repeat(batch, 1) + if isinstance(image_embeds, torch.Tensor): + image_embeds = image_embeds.to(device=hidden_states.device, dtype=hidden_states.dtype) + image_scatter_index = index.masked_select(image_mask.bool()).reshape(batch, -1) + hidden_states.scatter_( + dim=1, + index=image_scatter_index.unsqueeze(-1).repeat(1, 1, hidden), + src=image_embeds.reshape(batch, -1, hidden), + ) + return hidden_states + + for batch_idx, embeds in enumerate(image_embeds): + embeds = embeds.to(device=hidden_states.device, dtype=hidden_states.dtype).reshape(1, -1, hidden) + image_scatter_index = index[batch_idx : batch_idx + 1].masked_select(image_mask[batch_idx : batch_idx + 1].bool()).reshape(1, -1) + hidden_states[batch_idx : batch_idx + 1].scatter_( + dim=1, + index=image_scatter_index.unsqueeze(-1).repeat(1, 1, hidden), + src=embeds, + ) + return hidden_states + + def _instantiate_continuous_tokens(self, hidden_states, embedder_weights, values, scatter_index): + if values is None or scatter_index is None: + return hidden_states + + batch, _, hidden = hidden_states.shape + if isinstance(values, list): + for batch_idx, value in enumerate(values): + token_emb = apply_timestep_embedder(embedder_weights, value) + hidden_states[batch_idx : batch_idx + 1].scatter_( + dim=1, + index=scatter_index[batch_idx].unsqueeze(0).unsqueeze(-1).repeat(1, 1, hidden), + src=token_emb.reshape(1, -1, hidden), + ) + return hidden_states + + token_emb = apply_timestep_embedder(embedder_weights, values.reshape(-1)) + hidden_states.scatter_( + dim=1, + index=scatter_index.unsqueeze(-1).repeat(1, 1, hidden), + src=token_emb.reshape(batch, -1, hidden), + ) + return hidden_states + + @torch.no_grad() + def infer(self, weights, inputs): + device = first_weight_device(weights) + input_ids = inputs.get("input_ids") + input_ids = to_device(input_ids, device) + hidden_states = inputs.get("inputs_embeds") + hidden_states = to_device(hidden_states, device) + if hidden_states is None and input_ids is not None: + hidden_states = weights.token_embedding.apply(input_ids) + + token_hw = None + images = inputs.get("images") + images = to_device(images, device) + if images is not None: + hidden_states, token_hw = self._instantiate_image_tokens( + weights, + hidden_states, + images, + to_device(inputs.get("image_mask"), device), + to_device(inputs.get("timesteps"), device), + ) + + cond_vae_images = to_device(inputs.get("cond_vae_images"), device) + if cond_vae_images is not None: + hidden_states, _ = self._instantiate_image_tokens( + weights, + hidden_states, + cond_vae_images, + to_device(inputs.get("cond_vae_image_mask"), device), + to_device(inputs.get("cond_timesteps"), device), + ) + + hidden_states = self._instantiate_vit_image_tokens( + hidden_states, + to_device(inputs.get("cond_vit_embeds"), device), + to_device(inputs.get("cond_vit_image_mask"), device), + ) + + hidden_states = self._instantiate_continuous_tokens( + hidden_states, + weights.timestep_emb, + to_device(inputs.get("timesteps"), device), + to_device(inputs.get("timesteps_index"), device), + ) + if self.config.get("cfg_distilled", False): + hidden_states = self._instantiate_continuous_tokens( + hidden_states, + weights.guidance_emb, + to_device(inputs.get("guidance"), device), + to_device(inputs.get("guidance_index"), device), + ) + if self.config.get("use_meanflow", False): + hidden_states = self._instantiate_continuous_tokens( + hidden_states, + weights.timestep_r_emb, + to_device(inputs.get("timesteps_r"), device), + to_device(inputs.get("timesteps_r_index"), device), + ) + hidden_states = self._instantiate_continuous_tokens( + hidden_states, + weights.timestep_emb, + to_device(inputs.get("cond_timesteps"), device), + to_device(inputs.get("cond_timestep_index"), device), + ) + + if hidden_states is None: + raise ValueError("HunyuanImage3 pre_infer requires input_ids, inputs_embeds, or images.") + + return HunyuanImage3PreInferOutput( + hidden_states=hidden_states, + attention_mask=to_device(inputs.get("attention_mask"), device), + position_ids=to_device(inputs.get("position_ids"), device), + custom_pos_emb=to_device(inputs.get("custom_pos_emb"), device), + past_key_values=inputs.get("past_key_values"), + use_cache=bool(inputs.get("use_cache", False)), + image_mask=to_device(inputs.get("image_mask"), device), + timesteps=to_device(inputs.get("timesteps"), device), + token_hw=token_hw, + first_step=inputs.get("first_step"), + ) diff --git a/lightx2v/models/networks/hunyuan_image3/infer/transformer_infer.py b/lightx2v/models/networks/hunyuan_image3/infer/transformer_infer.py new file mode 100644 index 000000000..12dcb9456 --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/infer/transformer_infer.py @@ -0,0 +1,177 @@ +import torch +import torch.nn.functional as F + +from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer +from lightx2v.models.networks.hunyuan_image3.infer.utils import apply_linear, apply_mlp, apply_rotary_pos_emb, first_weight_device, repeat_kv, to_device + +try: + from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe +except Exception: + try: + import flashinfer + + flashinfer_cutlass_fused_moe = flashinfer.fused_moe.cutlass_fused_moe + except Exception: + flashinfer_cutlass_fused_moe = None + + +class HunyuanImage3TransformerInfer(BaseTransformerInfer): + def __init__(self, config): + self.config = config + self.num_layers = config["num_layers"] + self.hidden_size = config["hidden_size"] + self.num_heads = config["num_attention_heads"] + self.num_key_value_heads = config.get("num_key_value_heads") or self.num_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.head_dim = config.get("attention_head_dim", self.hidden_size // self.num_heads) + self.hidden_act = config.get("hidden_act", "silu") + + def set_scheduler(self, scheduler): + self.scheduler = scheduler + + @torch.no_grad() + def infer(self, weights, pre_infer_out): + hidden_states = pre_infer_out.hidden_states + for block_idx, block in enumerate(weights.blocks): + hidden_states = self.infer_block(block_idx, block, hidden_states, pre_infer_out) + return hidden_states + + def infer_block(self, block_idx, block, hidden_states, pre_infer_out): + attention_phase = block.compute_phases[0] + mlp_phase = block.compute_phases[1] + device = first_weight_device(attention_phase) + hidden_states = to_device(hidden_states, device) + attention_mask = to_device(pre_infer_out.attention_mask, device) + position_ids = to_device(pre_infer_out.position_ids, device) + custom_pos_emb = to_device(pre_infer_out.custom_pos_emb, device) + + residual = hidden_states + normed = attention_phase.input_layernorm.apply(hidden_states) + attn_out = self.infer_attention( + block_idx, + attention_phase, + normed, + attention_mask, + position_ids, + custom_pos_emb, + pre_infer_out.past_key_values if pre_infer_out.use_cache else None, + ) + hidden_states = residual + attn_out + + residual = hidden_states + normed = mlp_phase.post_attention_layernorm.apply(hidden_states) + mlp_out = self.infer_mlp(mlp_phase, normed) + return residual + mlp_out + + def infer_attention(self, block_idx, phase, hidden_states, attention_mask, position_ids, custom_pos_emb, past_key_values=None): + batch, q_len, _ = hidden_states.shape + qkv_states = apply_linear(phase.qkv_proj, hidden_states.reshape(-1, hidden_states.shape[-1])) + qkv_states = qkv_states.reshape( + batch, + q_len, + self.num_key_value_heads, + self.num_key_value_groups + 2, + self.head_dim, + ) + query_states, key_states, value_states = torch.split(qkv_states, [self.num_key_value_groups, 1, 1], dim=3) + query_states = query_states.reshape(batch, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.reshape(batch, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.reshape(batch, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + if custom_pos_emb is not None: + cos, sin = custom_pos_emb + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if getattr(phase, "query_layernorm", None) is not None: + query_states = phase.query_layernorm.apply(query_states) + key_states = phase.key_layernorm.apply(key_states) + + query_states = query_states.to(value_states.dtype) + key_states = key_states.to(value_states.dtype) + if past_key_values is not None: + if position_ids is None: + raise ValueError("HunyuanImage3 KV cache requires position_ids.") + key_states, value_states = past_key_values.update(key_states, value_states, block_idx, position_ids) + query_states = query_states.to(key_states.dtype) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + if query_states.device.type == "cuda" and attention_mask is not None: + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask=attention_mask, dropout_p=0.0) + attn_output = attn_output.transpose(1, 2).contiguous().reshape(batch, q_len, -1) + attn_output = apply_linear(phase.o_proj, attn_output.reshape(-1, attn_output.shape[-1])) + return attn_output.reshape(batch, q_len, -1) + + def infer_mlp(self, phase, hidden_states): + if not phase.is_moe: + return apply_mlp(phase.gate_and_up_proj, phase.down_proj, hidden_states, self.hidden_act) + + moe = phase.moe + moe_impl = getattr(moe, "moe_impl", self.config.get("moe_impl", "eager")) + if moe_impl == "flashinfer": + return self._infer_mlp_flashinfer(moe, hidden_states) + if moe_impl != "eager": + raise ValueError(f"Unsupported HunyuanImage3 moe_impl={moe_impl!r}. Expected 'eager' or 'flashinfer'.") + + return self._infer_mlp_eager(moe, hidden_states) + + def _moe_easy_topk(self, moe, hidden_states): + flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + logits = apply_linear(moe.gate, flat) + topk_weight, topk_idx = torch.topk(torch.softmax(logits, dim=-1), moe.moe_topk, dim=-1) + topk_weight = topk_weight / torch.clamp(topk_weight.sum(dim=-1, keepdim=True), min=1e-8) + return flat, topk_weight, topk_idx + + def _infer_mlp_eager(self, moe, hidden_states): + flat, topk_weight, topk_idx = self._moe_easy_topk(moe, hidden_states) + repeated = flat.repeat_interleave(moe.moe_topk, dim=0) + expert_outputs = torch.zeros_like(repeated) + flat_topk_idx = topk_idx.reshape(-1) + for expert_idx, expert in enumerate(moe.experts): + mask = flat_topk_idx == expert_idx + if not torch.any(mask): + continue + expert_out = apply_mlp(expert.gate_and_up_proj, expert.down_proj, repeated[mask], self.hidden_act) + expert_outputs[mask] = expert_out.to(expert_outputs.dtype) + combined = (expert_outputs.reshape(flat.shape[0], moe.moe_topk, -1) * topk_weight.to(expert_outputs.dtype).unsqueeze(-1)).sum(dim=1) + output = combined.reshape_as(hidden_states) + if getattr(moe, "shared_mlp", None) is not None: + shared_out = apply_mlp(moe.shared_mlp.gate_and_up_proj, moe.shared_mlp.down_proj, hidden_states, self.hidden_act) + output = output + shared_out.to(output.dtype) + return output + + def _infer_mlp_flashinfer(self, moe, hidden_states): + if flashinfer_cutlass_fused_moe is None: + raise ImportError("HunyuanImage3 moe_impl='flashinfer' requires flashinfer.fused_moe.cutlass_fused_moe.") + if self.hidden_act != "silu": + raise NotImplementedError("HunyuanImage3 moe_impl='flashinfer' currently supports only silu/SwiGLU experts.") + if not hasattr(moe, "ensure_flashinfer_weights"): + raise RuntimeError("HunyuanImage3 moe_impl='flashinfer' requires HunyuanImage3MoEWeights.") + + if hidden_states.device.type == "cuda" and hidden_states.device.index is not None: + torch.cuda.set_device(hidden_states.device.index) + + original_dtype = hidden_states.dtype + compute_dtype = original_dtype if original_dtype in (torch.float16, torch.bfloat16) else torch.bfloat16 + flat, topk_weight, topk_idx = self._moe_easy_topk(moe, hidden_states) + fused_input = flat.to(dtype=compute_dtype).contiguous() + moe_weight, moe_weight_2 = moe.ensure_flashinfer_weights(fused_input.device, compute_dtype) + combined_output = torch.zeros_like(fused_input) + flashinfer_cutlass_fused_moe( + fused_input, + topk_idx.to(torch.int32).contiguous(), + topk_weight.to(torch.float32).contiguous(), + moe_weight, + moe_weight_2, + compute_dtype, + output=combined_output, + quant_scales=None, + ) + output = combined_output.reshape_as(hidden_states).to(original_dtype) + if getattr(moe, "shared_mlp", None) is not None: + shared_out = apply_mlp(moe.shared_mlp.gate_and_up_proj, moe.shared_mlp.down_proj, hidden_states, self.hidden_act) + output = output + shared_out.to(output.dtype) + return output diff --git a/lightx2v/models/networks/hunyuan_image3/infer/utils.py b/lightx2v/models/networks/hunyuan_image3/infer/utils.py new file mode 100644 index 000000000..9a70b395f --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/infer/utils.py @@ -0,0 +1,114 @@ +import math + +import torch +import torch.nn.functional as F + + +def timestep_embedding(t, dim, max_period=10000): + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(0, half, dtype=torch.float32, device=t.device) / half) + args = t.float()[:, None] * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = F.pad(embedding, (0, 1)) + return embedding + + +def _linear_tensor_attr(linear, attr): + tensor = getattr(linear, attr, None) + if tensor is not None: + return tensor + return getattr(linear, f"pin_{attr}", None) + + +def match_linear_dtype(input_tensor, linear): + weight = _linear_tensor_attr(linear, "weight") + bias = _linear_tensor_attr(linear, "bias") + ref = weight if weight is not None else bias + if ref is None: + return input_tensor + return input_tensor.to(device=ref.device, dtype=ref.dtype) + + +def apply_linear(linear, input_tensor): + return linear.apply(match_linear_dtype(input_tensor, linear)) + + +def apply_timestep_embedder(weights, t): + emb = timestep_embedding(t.reshape(-1), 256) + emb = apply_linear(weights.linear_1, emb) + emb = F.gelu(emb) + return apply_linear(weights.linear_2, emb) + + +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + if position_ids is not None: + cos = cos[position_ids] + sin = sin[position_ids] + cos = cos.to(device=q.device, dtype=q.dtype) + sin = sin.to(device=q.device, dtype=q.dtype) + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin) + + +def repeat_kv(hidden_states, n_rep): + if n_rep == 1: + return hidden_states + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def weight_device(module): + for attr in ("weight", "pin_weight", "bias", "pin_bias"): + tensor = getattr(module, attr, None) + if tensor is not None: + return tensor.device + return None + + +def first_weight_device(module): + device = weight_device(module) + if device is not None: + return device + for child in getattr(module, "_modules", {}).values(): + device = first_weight_device(child) + if device is not None: + return device + return None + + +def to_device(value, device): + if value is None or device is None: + return value + if torch.is_tensor(value): + return value.to(device) + if isinstance(value, tuple): + return tuple(to_device(item, device) for item in value) + if isinstance(value, list): + return [to_device(item, device) for item in value] + if isinstance(value, dict): + return {key: to_device(item, device) for key, item in value.items()} + return value + + +def apply_mlp(gate_and_up_proj, down_proj, hidden_states, hidden_act="silu"): + original_shape = hidden_states.shape + flat = hidden_states.reshape(-1, original_shape[-1]) + projected = apply_linear(gate_and_up_proj, flat) + if hidden_act == "silu": + x1, x2 = projected.chunk(2, dim=-1) + hidden = x1 * F.silu(x2) + elif hidden_act == "gelu": + hidden = F.gelu(projected) + else: + raise NotImplementedError(f"Unsupported HunyuanImage3 hidden_act: {hidden_act}") + out = apply_linear(down_proj, hidden) + return out.reshape(*original_shape) diff --git a/lightx2v/models/networks/hunyuan_image3/model.py b/lightx2v/models/networks/hunyuan_image3/model.py new file mode 100644 index 000000000..8b3fa17db --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/model.py @@ -0,0 +1,223 @@ +import os +import re + +import torch +from loguru import logger +from safetensors import safe_open + +from lightx2v.models.networks.base_model import BaseTransformerModel +from lightx2v.models.networks.hunyuan_image3.infer.kv_cache import HunyuanImage3TaylorCache +from lightx2v.models.networks.hunyuan_image3.infer.post_infer import HunyuanImage3PostInfer +from lightx2v.models.networks.hunyuan_image3.infer.pre_infer import HunyuanImage3PreInfer +from lightx2v.models.networks.hunyuan_image3.infer.transformer_infer import HunyuanImage3TransformerInfer +from lightx2v.models.networks.hunyuan_image3.weights.post_weights import HunyuanImage3PostWeights +from lightx2v.models.networks.hunyuan_image3.weights.pre_weights import HunyuanImage3PreWeights +from lightx2v.models.networks.hunyuan_image3.weights.transformer_weights import HunyuanImage3TransformerWeights +from lightx2v.utils.envs import GET_DTYPE, GET_SENSITIVE_DTYPE + + +def _normalize_device_name(device): + if isinstance(device, int): + return f"cuda:{device}" + device = str(device) + if device.isdigit(): + return f"cuda:{device}" + if device == "cuda": + return "cuda:0" + return device + + +def resolve_pipeline_devices(config, fallback_device): + configured = config.get("pipeline_parallel_devices") or config.get("hunyuan_image3_pipeline_devices") + if configured: + if isinstance(configured, str): + devices = [item.strip() for item in configured.split(",") if item.strip()] + else: + devices = list(configured) + return [_normalize_device_name(device) for device in devices] + + if config.get("pipeline_parallel", False) and torch.cuda.is_available(): + device_count = torch.cuda.device_count() + if device_count > 0: + return [f"cuda:{idx}" for idx in range(device_count)] + + return [_normalize_device_name(fallback_device)] + + +def resolve_pipeline_device_for_key(key, config, devices): + devices = [_normalize_device_name(device) for device in devices] + if len(devices) == 1: + return devices[0] + + layer_match = re.match(r"model\.layers\.(\d+)\.", key) + if layer_match is not None: + layer_idx = int(layer_match.group(1)) + num_layers = int(config["num_layers"]) + stage_idx = min(layer_idx * len(devices) // num_layers, len(devices) - 1) + return devices[stage_idx] + + if key.startswith(("model.ln_f.", "lm_head.", "final_layer.", "time_embed_2.")): + return devices[-1] + + return devices[0] + + +class HunyuanImage3Model(BaseTransformerModel): + pre_weight_class = HunyuanImage3PreWeights + transformer_weight_class = HunyuanImage3TransformerWeights + post_weight_class = HunyuanImage3PostWeights + + def __init__(self, model_path, config, device, lora_path=None, lora_strength=1.0): + super().__init__(model_path, config, device, "hunyuan_image3", lora_path, lora_strength) + self.pipeline_devices = resolve_pipeline_devices(config, device) + self.pipeline_parallel = len(set(self.pipeline_devices)) > 1 + if self.lazy_load: + self.remove_keys.extend(["model.layers."]) + self.preserved_keys = [ + "model.", + "lm_head.", + "patch_embed.", + "final_layer.", + "time_embed.", + "time_embed_2.", + "timestep_emb.", + "guidance_emb.", + "timestep_r_emb.", + ] + self.sensitive_layer = { + "norm", + "layernorm", + "time_embed", + "timestep_emb", + "guidance_emb", + "timestep_r_emb", + "gate.wg", + } + self._init_infer_class() + self._init_weights() + self._init_infer() + + def _tensor_target_device(self, key): + return resolve_pipeline_device_for_key(key, self.config, self.pipeline_devices) + + def _load_safetensor_to_dict(self, file_path, unified_dtype, sensitive_layer): + ext = os.path.splitext(file_path)[-1] + if ext in (".pt", ".pth", ".tar"): + return super()._load_safetensor_to_dict(file_path, unified_dtype, sensitive_layer) + + remove_keys = self.remove_keys if hasattr(self, "remove_keys") else [] + preserve_keys = self.preserved_keys if hasattr(self, "preserved_keys") else None + weight_dict = {} + with safe_open(file_path, framework="pt", device="cpu") as f: + for key in f.keys(): + if any(remove_key in key for remove_key in remove_keys): + continue + if preserve_keys is not None and not any(preserve_key in key for preserve_key in preserve_keys): + continue + + tensor = f.get_tensor(key) + if tensor.dtype.is_floating_point: + dtype = GET_DTYPE() if unified_dtype or all(s not in key for s in sensitive_layer) else GET_SENSITIVE_DTYPE() + else: + dtype = tensor.dtype + target_device = self._tensor_target_device(key) + weight_dict[key] = tensor.to(device=target_device, dtype=dtype) + + if self.pipeline_parallel: + logger.debug(f"HunyuanImage3 loaded {len(weight_dict)} tensors from {file_path} across {self.pipeline_devices}") + return weight_dict + + def _init_infer_class(self): + if self.config["feature_caching"] != "NoCaching": + raise NotImplementedError("HunyuanImage3 native feature caching is not implemented yet.") + self.pre_infer_class = HunyuanImage3PreInfer + self.transformer_infer_class = HunyuanImage3TransformerInfer + self.post_infer_class = HunyuanImage3PostInfer + + def _init_infer(self): + self.pre_infer = self.pre_infer_class(self.config) + self.transformer_infer = self.transformer_infer_class(self.config) + self.post_infer = self.post_infer_class(self.config) + self.reset_taylor_cache() + if hasattr(self.transformer_infer, "offload_manager"): + self._init_offload_manager() + + def reset_taylor_cache(self): + self.taylor_cache = None + self.taylor_counter = 0 + self.taylor_last_full_computation_step = 0 + + def _should_taylor_full_compute(self, cache_dic): + current_step = int(cache_dic["current_step"]) + taylor_counter = getattr(self, "taylor_counter", 0) + return ( + current_step == 0 + or taylor_counter == int(cache_dic["cache_interval"]) - 1 + or (cache_dic["enable_first_enhance"] and current_step < int(cache_dic["first_enhance_steps"])) + or ( + cache_dic["enable_tailing_enhance"] + and current_step >= int(cache_dic["num_steps"]) - int(cache_dic["tailing_enhance_steps"]) + ) + ) + + def _infer_transformer(self, pre_infer_out): + hidden_states = self.transformer_infer.infer(self.transformer_weights, pre_infer_out) + if self.config["seq_parallel"]: + hidden_states = self._seq_parallel_post_process(hidden_states) + return hidden_states + + def _infer_transformer_with_taylor_cache(self, pre_infer_out, cache_dic): + if not hasattr(self, "taylor_counter"): + self.reset_taylor_cache() + if self.taylor_cache is None: + self.taylor_cache = HunyuanImage3TaylorCache(cache_dic["max_order"]) + + current_step = int(cache_dic["current_step"]) + if self._should_taylor_full_compute(cache_dic): + self.taylor_counter = 0 + hidden_states = self._infer_transformer(pre_infer_out) + if not (cache_dic["enable_first_enhance"] and current_step < int(cache_dic["first_enhance_steps"]) - 1): + self.taylor_cache.derivatives_computation( + hidden_states, + distance=current_step - self.taylor_last_full_computation_step, + low_freqs_order=cache_dic["low_freqs_order"], + high_freqs_order=cache_dic["high_freqs_order"], + ) + self.taylor_last_full_computation_step = current_step + self.taylor_cache.last_past_key_values = pre_infer_out.past_key_values + else: + self.taylor_counter += 1 + hidden_states = self.taylor_cache.taylor_formula(distance=self.taylor_counter) + + if current_step == int(cache_dic["num_steps"]) - 1: + self.taylor_cache.clear_derivatives() + return hidden_states + + @torch.no_grad() + def _infer_cond_uncond(self, inputs, infer_condition=True): + if hasattr(self, "scheduler"): + self.scheduler.infer_condition = infer_condition + pre_infer_out = self.pre_infer.infer(self.pre_weight, inputs) + # in default setting seq_parallel is False + if self.config["seq_parallel"]: + pre_infer_out = self._seq_parallel_pre_process(pre_infer_out) + if inputs.get("cache_dic") is not None: + hidden_states = self._infer_transformer_with_taylor_cache(pre_infer_out, inputs["cache_dic"]) + else: + hidden_states = self._infer_transformer(pre_infer_out) + return self.post_infer.infer(self.post_weight, hidden_states, pre_infer_out) + + @torch.no_grad() + def _seq_parallel_pre_process(self, pre_infer_out): + raise NotImplementedError("HunyuanImage3 native sequence parallel is not implemented yet.") + + @torch.no_grad() + def _seq_parallel_post_process(self, x): + raise NotImplementedError("HunyuanImage3 native sequence parallel is not implemented yet.") + + @torch.no_grad() + def infer(self, inputs): + output = self._infer_cond_uncond(inputs, infer_condition=True) + if hasattr(self, "scheduler") and "diffusion_prediction" in output: + self.scheduler.noise_pred = output["diffusion_prediction"] + return output diff --git a/lightx2v/models/networks/hunyuan_image3/weights/__init__.py b/lightx2v/models/networks/hunyuan_image3/weights/__init__.py new file mode 100644 index 000000000..7a9edf8b3 --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/weights/__init__.py @@ -0,0 +1,9 @@ +from .post_weights import HunyuanImage3PostWeights +from .pre_weights import HunyuanImage3PreWeights +from .transformer_weights import HunyuanImage3TransformerWeights + +__all__ = [ + "HunyuanImage3PostWeights", + "HunyuanImage3PreWeights", + "HunyuanImage3TransformerWeights", +] diff --git a/lightx2v/models/networks/hunyuan_image3/weights/common.py b/lightx2v/models/networks/hunyuan_image3/weights/common.py new file mode 100644 index 000000000..2cc652258 --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/weights/common.py @@ -0,0 +1,534 @@ +import math + +import torch +import torch.nn.functional as F + +import lightx2v.common.ops # noqa: F401 +from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList +from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER +from lightx2v_platform.base.global_var import AI_DEVICE + + +def _as_list(value, length): + if isinstance(value, list): + return value + return [value for _ in range(length)] + + +def _moe_value(config, name, block_index, default=None): + value = config.get(name, default) + if isinstance(value, list): + return value[block_index] + return value + + +class TensorPairWeight: + def __init__(self, weight_name, bias_name=None): + self.weight_name = weight_name + self.bias_name = bias_name + self.weight = None + self.bias = None + + def load(self, weight_dict): + self.weight = weight_dict[self.weight_name] + self.bias = weight_dict[self.bias_name] if self.bias_name is not None else None + + def apply_group_norm(self, x, eps=1e-5, groups=32): + channels = x.shape[1] + groups = math.gcd(groups, channels) + return F.group_norm(x, groups, self.weight, self.bias, eps) + + def to_cpu(self, non_blocking=False): + self.weight = self.weight.cpu(non_blocking=non_blocking) + if self.bias is not None: + self.bias = self.bias.cpu(non_blocking=non_blocking) + + def to_cuda(self, non_blocking=False): + self.weight = self.weight.to(AI_DEVICE, non_blocking=non_blocking) + if self.bias is not None: + self.bias = self.bias.to(AI_DEVICE, non_blocking=non_blocking) + + def state_dict(self, destination=None): + if destination is None: + destination = {} + destination[self.weight_name] = self.weight.detach().cpu().clone() + if self.bias is not None: + destination[self.bias_name] = self.bias.detach().cpu().clone() + return destination + + +class HunyuanImage3Conv2dWeight: + def __init__(self, weight_name, bias_name, stride=1, padding=0, dilation=1, groups=1): + self.weight_name = weight_name + self.bias_name = bias_name + self.stride = stride + self.padding = padding + self.dilation = dilation + self.groups = groups + self.weight = None + self.bias = None + + def load(self, weight_dict): + self.weight = weight_dict[self.weight_name] + self.bias = weight_dict[self.bias_name] if self.bias_name is not None else None + + def apply(self, input_tensor): + input_tensor = input_tensor.to(device=self.weight.device, dtype=self.weight.dtype) + return torch.nn.functional.conv2d( + input_tensor, + weight=self.weight, + bias=self.bias, + stride=self.stride, + padding=self.padding, + dilation=self.dilation, + groups=self.groups, + ) + + def to_cpu(self, non_blocking=False): + self.weight = self.weight.cpu(non_blocking=non_blocking) + if self.bias is not None: + self.bias = self.bias.cpu(non_blocking=non_blocking) + + def to_cuda(self, non_blocking=False): + self.weight = self.weight.to(AI_DEVICE, non_blocking=non_blocking) + if self.bias is not None: + self.bias = self.bias.to(AI_DEVICE, non_blocking=non_blocking) + + def state_dict(self, destination=None): + if destination is None: + destination = {} + destination[self.weight_name] = self.weight.detach().cpu().clone() + if self.bias is not None: + destination[self.bias_name] = self.bias.detach().cpu().clone() + return destination + + +class HunyuanImage3TimestepEmbedderWeights(WeightModule): + def __init__(self, prefix): + super().__init__() + self.add_module("linear_1", MM_WEIGHT_REGISTER["Default"](f"{prefix}.mlp.0.weight", f"{prefix}.mlp.0.bias")) + self.add_module("linear_2", MM_WEIGHT_REGISTER["Default"](f"{prefix}.mlp.2.weight", f"{prefix}.mlp.2.bias")) + + +class HunyuanImage3ResBlockWeights(WeightModule): + def __init__( + self, + prefix, + in_channels, + out_channels, + create_cuda_buffer=False, + create_cpu_buffer=False, + lazy_load=False, + lazy_load_file=None, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.add_module("in_norm", TensorPairWeight(f"{prefix}.in_layers.0.weight", f"{prefix}.in_layers.0.bias")) + self.add_module( + "in_conv", + HunyuanImage3Conv2dWeight( + f"{prefix}.in_layers.2.weight", + f"{prefix}.in_layers.2.bias", + stride=1, + padding=1, + ), + ) + self.add_module( + "emb_proj", + MM_WEIGHT_REGISTER["Default"]( + f"{prefix}.emb_layers.1.weight", + f"{prefix}.emb_layers.1.bias", + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + ), + ) + self.add_module("out_norm", TensorPairWeight(f"{prefix}.out_layers.0.weight", f"{prefix}.out_layers.0.bias")) + self.add_module( + "out_conv", + HunyuanImage3Conv2dWeight( + f"{prefix}.out_layers.3.weight", + f"{prefix}.out_layers.3.bias", + stride=1, + padding=1, + ), + ) + if in_channels != out_channels: + self.add_module( + "skip_connection", + HunyuanImage3Conv2dWeight( + f"{prefix}.skip_connection.weight", + f"{prefix}.skip_connection.bias", + stride=1, + padding=0, + ), + ) + else: + self.skip_connection = None + + +class HunyuanImage3UNetDownWeights(WeightModule): + def __init__(self, prefix, config): + super().__init__() + patch_size = int(config.get("patch_size", 1) or 1) + hidden_channels = int(config["patch_embed_hidden_dim"]) + out_channels = int(config["hidden_size"]) + in_channels = int(config.get("vae", {}).get("latent_channels", config.get("latent_channels", 32))) + self.patch_size = patch_size + self.add_module( + "input_conv", + HunyuanImage3Conv2dWeight( + f"{prefix}.model.0.weight", + f"{prefix}.model.0.bias", + stride=1, + padding=1, + ), + ) + + block_count = 1 if patch_size == 1 else patch_size // 2 + blocks = [] + for i in range(block_count): + block_in = hidden_channels + block_out = out_channels if patch_size == 1 or (i + 1) * 2 == patch_size else hidden_channels + blocks.append(HunyuanImage3ResBlockWeights(f"{prefix}.model.{i + 1}", block_in, block_out)) + self.blocks = WeightModuleList(blocks) + self.add_module("blocks", self.blocks) + self.in_channels = in_channels + + +class HunyuanImage3UNetUpWeights(WeightModule): + def __init__(self, prefix, config): + super().__init__() + patch_size = int(config.get("patch_size", 1) or 1) + hidden_channels = int(config["patch_embed_hidden_dim"]) + in_channels = int(config["hidden_size"]) + out_channels = int(config.get("vae", {}).get("latent_channels", config.get("latent_channels", 32))) + self.patch_size = patch_size + + block_count = 1 if patch_size == 1 else patch_size // 2 + blocks = [] + for i in range(block_count): + block_in = in_channels if i == 0 else hidden_channels + blocks.append(HunyuanImage3ResBlockWeights(f"{prefix}.model.{i}", block_in, hidden_channels)) + self.blocks = WeightModuleList(blocks) + self.add_module("blocks", self.blocks) + + last_index = block_count + self.add_module("out_norm", TensorPairWeight(f"{prefix}.model.{last_index}.0.weight", f"{prefix}.model.{last_index}.0.bias")) + self.add_module( + "output_conv", + HunyuanImage3Conv2dWeight( + f"{prefix}.model.{last_index}.2.weight", + f"{prefix}.model.{last_index}.2.bias", + stride=1, + padding=1, + ), + ) + self.out_channels = out_channels + + +class HunyuanImage3DenseMLPWeights(WeightModule): + def __init__( + self, + prefix, + mm_type, + mlp_bias=False, + create_cuda_buffer=False, + create_cpu_buffer=False, + lazy_load=False, + lazy_load_file=None, + lora_path=None, + ): + super().__init__() + gate_and_up_bias = f"{prefix}.gate_and_up_proj.bias" if mlp_bias else None + down_bias = f"{prefix}.down_proj.bias" if mlp_bias else None + self.add_module( + "gate_and_up_proj", + MM_WEIGHT_REGISTER[mm_type]( + f"{prefix}.gate_and_up_proj.weight", + gate_and_up_bias, + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + lora_prefix=prefix, + lora_path=lora_path, + ), + ) + self.add_module( + "down_proj", + MM_WEIGHT_REGISTER[mm_type]( + f"{prefix}.down_proj.weight", + down_bias, + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + lora_prefix=prefix, + lora_path=lora_path, + ), + ) + + +class HunyuanImage3MoEWeights(WeightModule): + def __init__( + self, + prefix, + block_index, + config, + mm_type, + create_cuda_buffer=False, + create_cpu_buffer=False, + lazy_load=False, + lazy_load_file=None, + lora_path=None, + ): + super().__init__() + self.num_experts = int(_moe_value(config, "num_experts", block_index, 1)) + self.moe_topk = int(_moe_value(config, "moe_topk", block_index, 1)) + self.moe_impl = config.get("moe_impl", "eager") + self.moe_weight = None + self.moe_weight_2 = None + self._flashinfer_weights_initialized = False + self._flashinfer_weight_device = None + self._flashinfer_weight_dtype = None + self.add_module( + "gate", + MM_WEIGHT_REGISTER["Default-ForceFp32"]( + f"{prefix}.gate.wg.weight", + None, + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + lora_prefix=prefix, + lora_path=lora_path, + ), + ) + if config.get("use_mixed_mlp_moe", False): + self.add_module( + "shared_mlp", + HunyuanImage3DenseMLPWeights( + f"{prefix}.shared_mlp", + mm_type, + config.get("mlp_bias", False), + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + lora_path, + ), + ) + else: + self.shared_mlp = None + self.experts = WeightModuleList( + [ + HunyuanImage3DenseMLPWeights( + f"{prefix}.experts.{i}", + mm_type, + config.get("mlp_bias", False), + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + lora_path, + ) + for i in range(self.num_experts) + ] + ) + self.add_module("experts", self.experts) + + def _linear_weight_for_flashinfer(self, linear, device, dtype): + if getattr(linear, "has_lora_branch", False): + raise RuntimeError("HunyuanImage3 moe_impl='flashinfer' does not support dynamic LoRA branches.") + if hasattr(linear, "weight_diff"): + raise RuntimeError("HunyuanImage3 moe_impl='flashinfer' does not support runtime diff weights.") + + weight = getattr(linear, "weight", None) + if weight is None: + weight = getattr(linear, "pin_weight", None) + if weight is None or weight.numel() == 0: + raise RuntimeError("HunyuanImage3 FlashInfer MoE weights were already released and cannot be rebuilt.") + + bias = getattr(linear, "bias", None) + pin_bias = getattr(linear, "pin_bias", None) + if (bias is not None and bias.numel() > 0) or (pin_bias is not None and pin_bias.numel() > 0): + raise RuntimeError("HunyuanImage3 moe_impl='flashinfer' expects bias-free expert MLP weights.") + + # LightX2V MM weights are stored as [in, out] for torch.mm(input, weight). + # FlashInfer follows torch.nn.Linear/checkpoint layout [out, in]. + return weight.t().to(device=device, dtype=dtype).contiguous() + + @staticmethod + def _release_linear_weight(linear, device, dtype): + empty = torch.empty(0, device=device, dtype=dtype) + for attr in ("weight", "pin_weight", "bias", "pin_bias"): + if hasattr(linear, attr): + setattr(linear, attr, empty) + + def ensure_flashinfer_weights(self, device, dtype): + device = torch.device(device) + if dtype not in (torch.float16, torch.bfloat16): + dtype = torch.bfloat16 + + if self._flashinfer_weights_initialized: + if self.moe_weight.device != device: + self.moe_weight = self.moe_weight.to(device=device) + self.moe_weight_2 = self.moe_weight_2.to(device=device) + if self.moe_weight.dtype != dtype: + self.moe_weight = self.moe_weight.to(dtype=dtype) + self.moe_weight_2 = self.moe_weight_2.to(dtype=dtype) + self._flashinfer_weight_device = self.moe_weight.device + self._flashinfer_weight_dtype = self.moe_weight.dtype + return self.moe_weight, self.moe_weight_2 + + expert_weights_gate_up = [] + expert_weights_down = [] + for expert in self.experts: + expert_weights_gate_up.append(self._linear_weight_for_flashinfer(expert.gate_and_up_proj, device, dtype)) + expert_weights_down.append(self._linear_weight_for_flashinfer(expert.down_proj, device, dtype)) + + self.moe_weight = torch.stack(expert_weights_gate_up, dim=0).contiguous() + self.moe_weight_2 = torch.stack(expert_weights_down, dim=0).contiguous() + + for expert in self.experts: + self._release_linear_weight(expert.gate_and_up_proj, self.moe_weight.device, self.moe_weight.dtype) + self._release_linear_weight(expert.down_proj, self.moe_weight_2.device, self.moe_weight_2.dtype) + + self._flashinfer_weights_initialized = True + self._flashinfer_weight_device = self.moe_weight.device + self._flashinfer_weight_dtype = self.moe_weight.dtype + return self.moe_weight, self.moe_weight_2 + + +class HunyuanImage3AttentionWeights(WeightModule): + def __init__( + self, + block_prefix, + block_index, + config, + mm_type, + create_cuda_buffer=False, + create_cpu_buffer=False, + lazy_load=False, + lazy_load_file=None, + lora_path=None, + ): + super().__init__() + prefix = f"{block_prefix}.{block_index}" + rms_type = config.get("rms_norm_type", "fp32_variance") + self.heads = int(config["num_attention_heads"]) + self.kv_heads = int(config.get("num_key_value_heads") or self.heads) + self.head_dim = int(config.get("attention_head_dim", config["hidden_size"] // self.heads)) + self.add_module( + "input_layernorm", + RMS_WEIGHT_REGISTER[config.get("rms_norm_type", "fp32_variance")]( + f"{prefix}.input_layernorm.weight", + eps=config.get("rms_norm_eps", 1e-5), + ), + ) + attn_bias = ".bias" if config.get("attention_bias", False) else None + self.add_module( + "qkv_proj", + MM_WEIGHT_REGISTER[mm_type]( + f"{prefix}.self_attn.qkv_proj.weight", + attn_bias and f"{prefix}.self_attn.qkv_proj{attn_bias}", + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + lora_prefix=prefix, + lora_path=lora_path, + ), + ) + self.add_module( + "o_proj", + MM_WEIGHT_REGISTER[mm_type]( + f"{prefix}.self_attn.o_proj.weight", + attn_bias and f"{prefix}.self_attn.o_proj{attn_bias}", + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + lora_prefix=prefix, + lora_path=lora_path, + ), + ) + if config.get("use_qk_norm", True): + self.add_module( + "query_layernorm", + RMS_WEIGHT_REGISTER[config.get("rms_norm_type", "fp32_variance")]( + f"{prefix}.self_attn.query_layernorm.weight", + eps=config.get("rms_norm_eps", 1e-5), + ), + ) + self.add_module( + "key_layernorm", + RMS_WEIGHT_REGISTER[config.get("rms_norm_type", "fp32_variance")]( + f"{prefix}.self_attn.key_layernorm.weight", + eps=config.get("rms_norm_eps", 1e-5), + ), + ) + else: + self.query_layernorm = None + self.key_layernorm = None + + +class HunyuanImage3MLPPhaseWeights(WeightModule): + def __init__( + self, + block_prefix, + block_index, + config, + mm_type, + create_cuda_buffer=False, + create_cpu_buffer=False, + lazy_load=False, + lazy_load_file=None, + lora_path=None, + ): + super().__init__() + prefix = f"{block_prefix}.{block_index}" + self.add_module( + "post_attention_layernorm", + RMS_WEIGHT_REGISTER[config.get("rms_norm_type", "fp32_variance")]( + f"{prefix}.post_attention_layernorm.weight", + eps=config.get("rms_norm_eps", 1e-5), + ), + ) + is_moe = int(_moe_value(config, "num_experts", block_index, 1)) > 1 and block_index >= int(config.get("moe_layer_num_skipped", 0)) + self.is_moe = is_moe + if is_moe: + self.add_module( + "moe", + HunyuanImage3MoEWeights( + f"{prefix}.mlp", + block_index, + config, + mm_type, + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + lora_path, + ), + ) + self.experts = self.moe.experts + else: + self.add_module( + "dense_mlp", + HunyuanImage3DenseMLPWeights( + f"{prefix}.mlp", + mm_type, + config.get("mlp_bias", False), + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + lora_path, + ), + ) + self.gate_and_up_proj = self.dense_mlp.gate_and_up_proj + self.down_proj = self.dense_mlp.down_proj diff --git a/lightx2v/models/networks/hunyuan_image3/weights/post_weights.py b/lightx2v/models/networks/hunyuan_image3/weights/post_weights.py new file mode 100644 index 000000000..f86f0fe3c --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/weights/post_weights.py @@ -0,0 +1,32 @@ +from lightx2v.common.modules.weight_module import WeightModule +from lightx2v.models.networks.hunyuan_image3.weights.common import ( + HunyuanImage3TimestepEmbedderWeights, + HunyuanImage3UNetUpWeights, +) +from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER + + +class HunyuanImage3PostWeights(WeightModule): + def __init__(self, config): + super().__init__() + self.config = config + self.add_module("time_embed_2", HunyuanImage3TimestepEmbedderWeights("time_embed_2")) + self.add_module( + "final_norm", + RMS_WEIGHT_REGISTER[config.get("rms_norm_type", "fp32_variance")]( + "model.ln_f.weight", + eps=config.get("rms_norm_eps", 1e-5), + ), + ) + self.add_module("lm_head", MM_WEIGHT_REGISTER["Default"]("lm_head.weight")) + self.add_module("final_layer", HunyuanImage3UNetUpWeights("final_layer", config)) + + def to_cpu(self, non_blocking=True): + for module in self._modules.values(): + if module is not None and hasattr(module, "to_cpu"): + module.to_cpu(non_blocking=non_blocking) + + def to_cuda(self, non_blocking=True): + for module in self._modules.values(): + if module is not None and hasattr(module, "to_cuda"): + module.to_cuda(non_blocking=non_blocking) diff --git a/lightx2v/models/networks/hunyuan_image3/weights/pre_weights.py b/lightx2v/models/networks/hunyuan_image3/weights/pre_weights.py new file mode 100644 index 000000000..e366b77e8 --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/weights/pre_weights.py @@ -0,0 +1,30 @@ +from lightx2v.common.modules.weight_module import WeightModule +from lightx2v.models.networks.hunyuan_image3.weights.common import ( + HunyuanImage3TimestepEmbedderWeights, + HunyuanImage3UNetDownWeights, +) +from lightx2v.utils.registry_factory import EMBEDDING_WEIGHT_REGISTER + + +class HunyuanImage3PreWeights(WeightModule): + def __init__(self, config): + super().__init__() + self.config = config + self.add_module("token_embedding", EMBEDDING_WEIGHT_REGISTER["Default"]("model.wte.weight")) + self.add_module("timestep_emb", HunyuanImage3TimestepEmbedderWeights("timestep_emb")) + self.add_module("time_embed", HunyuanImage3TimestepEmbedderWeights("time_embed")) + if config.get("cfg_distilled", False): + self.add_module("guidance_emb", HunyuanImage3TimestepEmbedderWeights("guidance_emb")) + if config.get("use_meanflow", False): + self.add_module("timestep_r_emb", HunyuanImage3TimestepEmbedderWeights("timestep_r_emb")) + self.add_module("patch_embed", HunyuanImage3UNetDownWeights("patch_embed", config)) + + def to_cpu(self, non_blocking=True): + for module in self._modules.values(): + if module is not None and hasattr(module, "to_cpu"): + module.to_cpu(non_blocking=non_blocking) + + def to_cuda(self, non_blocking=True): + for module in self._modules.values(): + if module is not None and hasattr(module, "to_cuda"): + module.to_cuda(non_blocking=non_blocking) diff --git a/lightx2v/models/networks/hunyuan_image3/weights/transformer_weights.py b/lightx2v/models/networks/hunyuan_image3/weights/transformer_weights.py new file mode 100644 index 000000000..f37e13278 --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/weights/transformer_weights.py @@ -0,0 +1,162 @@ +from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList +from lightx2v.models.networks.hunyuan_image3.weights.common import ( + HunyuanImage3AttentionWeights, + HunyuanImage3MLPPhaseWeights, +) + + +class HunyuanImage3TransformerWeights(WeightModule): + def __init__(self, config, lazy_load_path=None, lora_path=None): + super().__init__() + self.config = config + self.blocks_num = config["num_layers"] + self.mm_type = config.get("dit_quant_scheme", "Default") + if self.mm_type != "Default": + assert config.get("dit_quantized") is True + self.lazy_load = config.get("lazy_load", False) + self.blocks = WeightModuleList( + [ + HunyuanImage3TransformerBlock( + block_index=i, + config=config, + mm_type=self.mm_type, + create_cuda_buffer=False, + create_cpu_buffer=False, + block_prefix="model.layers", + lazy_load=self.lazy_load, + lazy_load_path=lazy_load_path, + lora_path=lora_path, + ) + for i in range(self.blocks_num) + ] + ) + self.register_offload_buffers(config, lazy_load_path, lora_path) + self.add_module("blocks", self.blocks) + + def register_offload_buffers(self, config, lazy_load_path, lora_path): + self.offload_block_cuda_buffers = None + self.offload_phase_cuda_buffers = None + self.offload_block_cpu_buffers = None + self.offload_phase_cpu_buffers = None + if not config.get("cpu_offload", False): + return + + if config.get("offload_granularity", "block") == "block": + self.offload_blocks_num = 2 + self.offload_block_cuda_buffers = WeightModuleList( + [ + HunyuanImage3TransformerBlock( + i, + config, + self.mm_type, + True, + False, + "model.layers", + self.lazy_load, + lazy_load_path, + lora_path, + ) + for i in range(self.offload_blocks_num) + ] + ) + self.add_module("offload_block_cuda_buffers", self.offload_block_cuda_buffers) + if self.lazy_load: + self.offload_block_cpu_buffers = WeightModuleList( + [ + HunyuanImage3TransformerBlock( + i, + config, + self.mm_type, + False, + True, + "model.layers", + self.lazy_load, + lazy_load_path, + lora_path, + ) + for i in range(self.offload_blocks_num) + ] + ) + self.add_module("offload_block_cpu_buffers", self.offload_block_cpu_buffers) + elif config.get("offload_granularity") == "phase": + self.offload_phase_cuda_buffers = HunyuanImage3TransformerBlock( + 0, + config, + self.mm_type, + True, + False, + "model.layers", + self.lazy_load, + lazy_load_path, + lora_path, + ).compute_phases + self.add_module("offload_phase_cuda_buffers", self.offload_phase_cuda_buffers) + if self.lazy_load: + self.offload_phase_cpu_buffers = WeightModuleList( + [ + HunyuanImage3TransformerBlock( + i, + config, + self.mm_type, + False, + True, + "model.layers", + self.lazy_load, + lazy_load_path, + lora_path, + ).compute_phases + for i in range(2) + ] + ) + self.add_module("offload_phase_cpu_buffers", self.offload_phase_cpu_buffers) + + def non_block_weights_to_cuda(self): + pass + + def non_block_weights_to_cpu(self): + pass + + +class HunyuanImage3TransformerBlock(WeightModule): + def __init__( + self, + block_index, + config, + mm_type, + create_cuda_buffer=False, + create_cpu_buffer=False, + block_prefix="model.layers", + lazy_load=False, + lazy_load_path=None, + lora_path=None, + ): + super().__init__() + self.block_index = block_index + lazy_load_file = lazy_load_path if lazy_load else None + self.compute_phases = WeightModuleList( + [ + HunyuanImage3AttentionWeights( + block_prefix, + block_index, + config, + mm_type, + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + lora_path, + ), + HunyuanImage3MLPPhaseWeights( + block_prefix, + block_index, + config, + mm_type, + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + lora_path, + ), + ] + ) + self.add_module("compute_phases", self.compute_phases) diff --git a/lightx2v/models/runners/hunyuan_image3/__init__.py b/lightx2v/models/runners/hunyuan_image3/__init__.py new file mode 100644 index 000000000..87697f260 --- /dev/null +++ b/lightx2v/models/runners/hunyuan_image3/__init__.py @@ -0,0 +1 @@ +from .hunyuan_image3_runner import HunyuanImage3Runner diff --git a/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py b/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py new file mode 100644 index 000000000..ddef99154 --- /dev/null +++ b/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py @@ -0,0 +1,1037 @@ +import importlib +import os +import sys +from dataclasses import dataclass +from pathlib import Path + +import torch +from PIL import Image +from loguru import logger +from safetensors import safe_open +from transformers import GenerationConfig + +from lightx2v.models.networks.hunyuan_image3.infer.kv_cache import HunyuanImage3StaticKVCache +from lightx2v.models.networks.hunyuan_image3.model import HunyuanImage3Model +from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.schedulers.hunyuan_image3.scheduler import HunyuanImage3Scheduler +from lightx2v.utils.registry_factory import RUNNER_REGISTER + + +@dataclass(frozen=True) +class HunyuanImage3TextGenerationPlan: + first_bot_task: str + stage_transitions: list[tuple[int, list[int]]] + final_stop_tokens: list[int] + + +@RUNNER_REGISTER("hunyuan_image3") +class HunyuanImage3Runner(DefaultRunner): + model_cpu_offload_seq = "transformer" + + def load_model(self): + self.model = self.load_transformer() + self.text_encoders = [] + self.image_encoder = None + self.vae_encoder = None + self.vae_decoder = None + + def load_transformer(self): + logger.info("Loading native HunyuanImage3 transformer weights") + return HunyuanImage3Model( + model_path=self.config["model_path"], + config=self.config, + device=self.init_device, + ) + + def load_text_encoder(self): + return [] + + def load_image_encoder(self): + return None + + def load_vae(self): + return None, None + + def init_scheduler(self): + super().init_scheduler() + self.scheduler = HunyuanImage3Scheduler(self.config) + + def init_modules(self): + super().init_modules() + self.run_dit = self._run_dit_local + + def _run_dit_local(self, total_steps=None): + if self.config.get("lazy_load", False) or self.config.get("unload_modules", False): + self.model = self.load_transformer() + self.model.set_scheduler(self.scheduler) + return self.run_main(total_steps=total_steps) + + def _resolve_upstream_repo_path(self): + candidates = [ + self.config.get("hunyuan_image3_repo_path"), + os.environ.get("HUNYUAN_IMAGE3_REPO_PATH"), + ] + model_path = Path(self.config["model_path"]).resolve() + candidates.extend(parent / "HunyuanImage-3.0" for parent in model_path.parents) + candidates.append(Path.cwd().parent / "HunyuanImage-3.0") + + for candidate in candidates: + if not candidate: + continue + candidate = Path(candidate).expanduser().resolve() + if (candidate / "hunyuan_image_3" / "__init__.py").exists(): + return str(candidate) + return None + + def _import_upstream_modules(self): + try: + importlib.import_module("hunyuan_image_3") + except ModuleNotFoundError: + repo_path = self._resolve_upstream_repo_path() + if repo_path is None: + raise ModuleNotFoundError( + "Cannot import hunyuan_image_3. Set HUNYUAN_IMAGE3_REPO_PATH or " + "`hunyuan_image3_repo_path` to the HunyuanImage-3.0 project path." + ) + if repo_path not in sys.path: + sys.path.insert(0, repo_path) + + return { + "AutoencoderKLConv3D": importlib.import_module("hunyuan_image_3.autoencoder_kl_3d").AutoencoderKLConv3D, + "CachedRoPE": importlib.import_module("hunyuan_image_3.modeling_hunyuan_image_3").CachedRoPE, + "HunyuanImage3Config": importlib.import_module("hunyuan_image_3.configuration_hunyuan_image_3").HunyuanImage3Config, + "HunyuanImage3ImageProcessor": importlib.import_module("hunyuan_image_3.image_processor").HunyuanImage3ImageProcessor, + "LightProjector": importlib.import_module("hunyuan_image_3.siglip2").LightProjector, + "Siglip2VisionTransformer": importlib.import_module("hunyuan_image_3.siglip2").Siglip2VisionTransformer, + "HunyuanImage3TokenizerFast": importlib.import_module("hunyuan_image_3.tokenization_hunyuan_image_3").HunyuanImage3TokenizerFast, + "get_system_prompt": importlib.import_module("hunyuan_image_3.system_prompt").get_system_prompt, + } + + def _ensure_pipeline_modules(self): + if getattr(self, "_hunyuan_pipeline_modules_ready", False): + return + + modules = self._import_upstream_modules() + model_path = self.config["model_path"] + self.hunyuan_config = modules["HunyuanImage3Config"].from_pretrained(model_path) + try: + self.hunyuan_generation_config = GenerationConfig.from_pretrained(model_path) + except OSError: + self.hunyuan_generation_config = GenerationConfig() + + self.hunyuan_tokenizer = modules["HunyuanImage3TokenizerFast"].from_pretrained( + model_path, + model_version=self.hunyuan_config.model_version, + ) + self.hunyuan_image_processor = modules["HunyuanImage3ImageProcessor"](self.hunyuan_config) + self.hunyuan_cached_rope = modules["CachedRoPE"](self.hunyuan_config) + self.hunyuan_vae_cls = modules["AutoencoderKLConv3D"] + self.hunyuan_vision_cls = modules["Siglip2VisionTransformer"] + self.hunyuan_vision_aligner_cls = modules["LightProjector"] + self.hunyuan_get_system_prompt = modules["get_system_prompt"] + + # todo: actually, self.vae_decoder is a complete VAE model, which includes both encoder and decoder. There is a name confusion here. + self.vae_decoder = self._load_vae_decoder() + self.vision_model = None + self.vision_aligner = None + self._hunyuan_pipeline_modules_ready = True + + def _pipeline_latent_device(self): + devices = getattr(self.model, "pipeline_devices", None) + if devices: + return torch.device(devices[0]) + return torch.device(self.init_device) + + def _vae_device(self): + configured = self.config.get("hunyuan_image3_vae_device") + if configured is not None: + return torch.device(str(configured)) + devices = getattr(self.model, "pipeline_devices", None) + if devices: + return torch.device(devices[-1]) + return torch.device(self.init_device) + + def _vision_device(self): + configured = self.config.get("hunyuan_image3_vision_device") + if configured is not None: + return torch.device(str(configured)) + return self._pipeline_latent_device() + + def _resolve_torch_dtype(self, value, default): + if value in (None, "auto"): + return default + if isinstance(value, torch.dtype): + return value + return getattr(torch, str(value), default) + + def _iter_vae_weight_files(self): + model_path = Path(self.config["model_path"]) + index_path = model_path / "model.safetensors.index.json" + if index_path.exists(): + import json + + weight_map = json.loads(index_path.read_text())["weight_map"] + files = sorted({filename for key, filename in weight_map.items() if key.startswith("vae.")}) + return [model_path / filename for filename in files] + return sorted(model_path.glob("*.safetensors")) + + def _iter_prefixed_weight_files(self, prefix): + model_path = Path(self.config["model_path"]) + index_path = model_path / "model.safetensors.index.json" + if index_path.exists(): + import json + + weight_map = json.loads(index_path.read_text())["weight_map"] + files = sorted({filename for key, filename in weight_map.items() if key.startswith(prefix)}) + return [model_path / filename for filename in files] + return sorted(model_path.glob("*.safetensors")) + + def _load_vae_decoder(self): + logger.info("Loading HunyuanImage3 VAE weights") + vae = self.hunyuan_vae_cls.from_config(self.hunyuan_config.vae) + vae.eval() + for param in vae.parameters(): + param.requires_grad = False + + loaded_keys = set() + for file_path in self._iter_vae_weight_files(): + state_dict = {} + with safe_open(str(file_path), framework="pt", device="cpu") as f: + for key in f.keys(): + if key.startswith("vae."): + state_key = key.removeprefix("vae.") + state_dict[state_key] = f.get_tensor(key) + loaded_keys.add(state_key) + if state_dict: + vae.load_state_dict(state_dict, strict=False) + + missing = sorted(set(vae.state_dict()) - loaded_keys) + if missing: + raise RuntimeError(f"HunyuanImage3 VAE weights are incomplete; missing {len(missing)} keys, first key: {missing[0]}") + + vae_dtype = getattr(torch, self.hunyuan_config.vae_dtype, torch.float32) + return vae.to(device=self._vae_device(), dtype=vae_dtype) + + def _load_prefixed_module_weights(self, module, prefix): + loaded_keys = set() + for file_path in self._iter_prefixed_weight_files(prefix): + state_dict = {} + with safe_open(str(file_path), framework="pt", device="cpu") as f: + for key in f.keys(): + if key.startswith(prefix): + state_key = key.removeprefix(prefix) + state_dict[state_key] = f.get_tensor(key) + loaded_keys.add(state_key) + if state_dict: + module.load_state_dict(state_dict, strict=False) + + missing = sorted(set(module.state_dict()) - loaded_keys) + if missing: + raise RuntimeError(f"HunyuanImage3 {prefix.rstrip('.')} weights are incomplete; missing {len(missing)} keys, first key: {missing[0]}") + + def _ensure_vision_encoder(self): + if self.vision_model is not None and self.vision_aligner is not None: + return + + logger.info("Loading HunyuanImage3 conditional vision encoder weights") + vision_model = self.hunyuan_vision_cls(self.hunyuan_config.vit) + vision_aligner = self.hunyuan_vision_aligner_cls(self.hunyuan_config.vit_aligner) + vision_model.eval() + vision_aligner.eval() + for module in (vision_model, vision_aligner): + for param in module.parameters(): + param.requires_grad = False + + self._load_prefixed_module_weights(vision_model, "vision_model.") + self._load_prefixed_module_weights(vision_aligner, "vision_aligner.") + + vision_dtype = self._resolve_torch_dtype(self.config.get("hunyuan_image3_vision_dtype"), torch.bfloat16) + vision_device = self._vision_device() + self.vision_model = vision_model.to(device=vision_device, dtype=vision_dtype) + self.vision_aligner = vision_aligner.to(device=vision_device, dtype=vision_dtype) + + def _resolve_image_size(self, input_info): + if self.config.get("image_size"): + return self.config["image_size"] + target_shape = getattr(input_info, "target_shape", None) or self.config.get("target_shape") + if target_shape: + if len(target_shape) >= 2: + return int(target_shape[-2]), int(target_shape[-1]) + return int(self.config.get("target_height", 1024)), int(self.config.get("target_width", 1024)) + + def _build_batch_rope_image_info(self, output, sections): + if self.hunyuan_config.rope_type == "default": + return None + if self.hunyuan_config.rope_type != "2d": + raise NotImplementedError(f"HunyuanImage3 rope type {self.hunyuan_config.rope_type} is not supported.") + + rope_image_info = [] + for image_slices, sample_sections in zip(output.all_image_slices, sections): + image_idx = 0 + sample_info = [] + for section in sample_sections: + if section["type"] in ("gen_image", "cond_vae_image", "cond_vit_image"): + sample_info.append((image_slices[image_idx], (section["token_height"], section["token_width"]))) + image_idx += 1 + elif section["type"] == "cond_joint_image": + if self.hunyuan_image_processor.cond_token_attn_type in ("full", "joint_full"): + sample_info.extend( + [ + (image_slices[image_idx], (section["token_height"][0], section["token_width"][0])), + (image_slices[image_idx + 1], (section["token_height"][1], section["token_width"][1])), + ] + ) + elif self.hunyuan_image_processor.cond_token_attn_type == "full_causal": + sample_info.append((image_slices[image_idx], (section["token_height"][0], section["token_width"][0]))) + elif self.hunyuan_image_processor.cond_token_attn_type == "causal": + pass + else: + raise NotImplementedError( + f"HunyuanImage3 cond_token_attn_type={self.hunyuan_image_processor.cond_token_attn_type!r} is not supported." + ) + image_idx += 2 + rope_image_info.append(sample_info) + return rope_image_info + + def _build_attention_mask(self, input_ids, tokenizer_output): + batch, seq_len = input_ids.shape + attention_mask = torch.ones(seq_len, seq_len, dtype=torch.bool, device=input_ids.device).tril(diagonal=0).repeat(batch, 1, 1) + for batch_idx in range(batch): + for image_slice in self.hunyuan_image_processor.prepare_full_attn_slices(tokenizer_output, batch_idx): + attention_mask[batch_idx, image_slice, image_slice] = True + return attention_mask.unsqueeze(1) + + def _hunyuan_kv_cache_enabled(self): + return bool(self.config.get("enable_kv_cache", True)) + + def _hunyuan_text_kv_cache_enabled(self): + if "enable_text_kv_cache" in self.config: + return bool(self.config["enable_text_kv_cache"]) + if "enable_kv_cache" in self.config: + return bool(self.config["enable_kv_cache"]) + return hasattr(self, "hunyuan_config") + + def _hunyuan_num_layers(self): + hunyuan_config = getattr(self, "hunyuan_config", None) + return int(self.config.get("num_layers", getattr(hunyuan_config, "num_hidden_layers", 1))) + + def _hunyuan_taylor_cache_enabled(self): + return bool(self.config.get("use_taylor_cache", False)) + + def _build_taylor_cache_dic(self, num_steps): + return { + "counter": 0, + "current_step": 0, + "cache_interval": int(self.config.get("taylor_cache_interval", 5)), + "max_order": int(self.config.get("taylor_cache_order", 2)), + "num_steps": int(num_steps), + "enable_first_enhance": bool(self.config.get("taylor_cache_enable_first_enhance", False)), + "first_enhance_steps": int(self.config.get("taylor_cache_first_enhance_steps", 3)), + "enable_tailing_enhance": bool(self.config.get("taylor_cache_enable_tailing_enhance", False)), + "tailing_enhance_steps": int(self.config.get("taylor_cache_tailing_enhance_steps", 1)), + "low_freqs_order": int(self.config.get("taylor_cache_low_freqs_order", 2)), + "high_freqs_order": int(self.config.get("taylor_cache_high_freqs_order", 2)), + "enable_force_control": False, + "force_compute": False, + } + + def _build_denoise_cache_position_ids(self, prepared_inputs): + image_mask = prepared_inputs["image_mask"] + batch, seq_len = image_mask.shape + device = image_mask.device + batch_index = torch.arange(batch, device=device) + token_index = torch.arange(seq_len, device=device).unsqueeze(0).repeat(batch, 1) + position_parts = [] + + timesteps_index = prepared_inputs.get("timesteps_index") + if timesteps_index is not None: + position_parts.append(token_index[batch_index, timesteps_index[:, -1]].unsqueeze(-1)) + + guidance_index = prepared_inputs.get("guidance_index") + if guidance_index is not None: + position_parts.append(token_index[batch_index, guidance_index[:, -1]].unsqueeze(-1)) + + timesteps_r_index = prepared_inputs.get("timesteps_r_index") + if timesteps_r_index is not None: + position_parts.append(token_index[batch_index, timesteps_r_index[:, -1]].unsqueeze(-1)) + + image_position_ids = token_index.masked_select(image_mask.bool()).reshape(batch, -1) + position_parts.append(image_position_ids) + return torch.cat(position_parts, dim=1) + + def _slice_denoise_cache_attention_mask(self, attention_mask, position_ids): + if attention_mask is None: + return None + mask_parts = [] + for sample_attention_mask, sample_position_ids in zip(attention_mask, position_ids): + mask_parts.append(sample_attention_mask.index_select(dim=1, index=sample_position_ids.reshape(-1))) + return torch.stack(mask_parts, dim=0) + + def _build_denoise_cache_local_indices(self, prepared_inputs, position_ids): + image_tokens = int(prepared_inputs["image_mask"][0].sum().item()) + batch, local_seq_len = position_ids.shape + device = position_ids.device + special_tokens = local_seq_len - image_tokens + local_image_mask = torch.zeros((batch, local_seq_len), dtype=torch.bool, device=device) + local_image_mask[:, special_tokens:] = True + + local_index = 0 + local_inputs = { + "image_mask": local_image_mask, + "timesteps_index": None, + "guidance_index": None, + "timesteps_r_index": None, + } + if prepared_inputs.get("timesteps_index") is not None: + local_inputs["timesteps_index"] = torch.full((batch, 1), local_index, dtype=torch.long, device=device) + local_index += 1 + if prepared_inputs.get("guidance_index") is not None: + local_inputs["guidance_index"] = torch.full((batch, 1), local_index, dtype=torch.long, device=device) + local_index += 1 + if prepared_inputs.get("timesteps_r_index") is not None: + local_inputs["timesteps_r_index"] = torch.full((batch, 1), local_index, dtype=torch.long, device=device) + return local_inputs + + def _build_custom_pos_emb(self, position_ids, rope_image_info): + return self.hunyuan_cached_rope( + self.config.get("max_position_embeddings", self.hunyuan_config.max_position_embeddings), + position_ids.device, + rope_image_info=rope_image_info, + position_ids=position_ids, + ) + + def _attach_tokenizer_cond_masks(self, cond_inputs, tokenizer_output, device): + if not cond_inputs: + return cond_inputs + cond_inputs = dict(cond_inputs) + cond_inputs["cond_vae_image_mask"] = None if getattr(tokenizer_output, "vae_image_mask", None) is None else tokenizer_output.vae_image_mask.to(device) + cond_inputs["cond_vit_image_mask"] = None if getattr(tokenizer_output, "vit_image_mask", None) is None else tokenizer_output.vit_image_mask.to(device) + cond_inputs["cond_timestep_index"] = ( + None + if getattr(tokenizer_output, "cond_timestep_scatter_index", None) is None + else tokenizer_output.cond_timestep_scatter_index.to(device) + ) + return cond_inputs + + def _resolve_bot_task(self): + bot_task = self.config.get("bot_task", getattr(self.hunyuan_generation_config, "bot_task", "image")) + return bot_task or "image" + + def _resolve_system_prompt(self, bot_task): + system_prompt_type = self.config.get("use_system_prompt", getattr(self.hunyuan_generation_config, "use_system_prompt", "None")) + system_prompt = self.config.get("system_prompt") + if hasattr(self, "hunyuan_get_system_prompt"): + return self.hunyuan_get_system_prompt(system_prompt_type, bot_task, system_prompt) + if system_prompt_type == "custom": + return system_prompt + return system_prompt if system_prompt_type not in (None, "None") else None + + def _resolve_text_generation_plan(self, bot_task): + tokenizer = self.hunyuan_tokenizer + if bot_task == "recaption": + return HunyuanImage3TextGenerationPlan( + first_bot_task="recaption", + stage_transitions=[], + final_stop_tokens=[tokenizer.end_of_recaption_token_id], + ) + if bot_task == "think_recaption": + recaption_token_id = tokenizer.convert_tokens_to_ids(tokenizer.recaption_token) + return HunyuanImage3TextGenerationPlan( + first_bot_task="think", + stage_transitions=[(tokenizer.end_of_think_token_id, [recaption_token_id])], + final_stop_tokens=[tokenizer.end_of_recaption_token_id], + ) + if bot_task == "think": + return HunyuanImage3TextGenerationPlan( + first_bot_task="think", + stage_transitions=[], + final_stop_tokens=[tokenizer.end_of_think_token_id, tokenizer.end_of_recaption_token_id], + ) + raise NotImplementedError(f"HunyuanImage3 native runner does not support bot_task={bot_task!r}.") + + def _expand_sequence_mask(self, mask, seq_len): + if mask is None: + return None + if mask.shape[1] == seq_len: + return mask + if mask.shape[1] > seq_len: + return mask[:, :seq_len] + pad = torch.zeros((mask.shape[0], seq_len - mask.shape[1]), dtype=mask.dtype, device=mask.device) + return torch.cat([mask, pad], dim=1) + + def _build_text_model_inputs( + self, + input_ids, + tokenizer_output, + cond_inputs=None, + position_ids=None, + attention_mask=None, + build_attention_mask=True, + past_key_values=None, + use_cache=False, + ): + device = input_ids.device + if position_ids is None: + position_ids = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=device)[None].expand(input_ids.shape[0], -1) + position_max = int(position_ids.max().item()) + 1 if position_ids.numel() > 0 else input_ids.shape[1] + max_position = max(int(self.config.get("max_position_embeddings", self.hunyuan_config.max_position_embeddings)), position_max) + custom_pos_emb = self.hunyuan_cached_rope(max_position, device, rope_image_info=None, position_ids=position_ids) + if build_attention_mask and attention_mask is None: + attention_mask = self._build_attention_mask(input_ids, tokenizer_output) + model_inputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + "custom_pos_emb": custom_pos_emb, + } + if use_cache: + model_inputs["past_key_values"] = past_key_values + model_inputs["use_cache"] = True + if cond_inputs: + model_inputs.update( + { + "cond_vae_images": cond_inputs.get("cond_vae_images"), + "cond_vae_image_mask": self._expand_sequence_mask(cond_inputs.get("cond_vae_image_mask"), input_ids.shape[1]), + "cond_timesteps": cond_inputs.get("cond_timesteps"), + "cond_timestep_index": cond_inputs.get("cond_timestep_index"), + "cond_vit_embeds": cond_inputs.get("cond_vit_embeds"), + "cond_vit_image_mask": self._expand_sequence_mask(cond_inputs.get("cond_vit_image_mask"), input_ids.shape[1]), + } + ) + return model_inputs + + def _prepare_text_generation_inputs(self, prompt, bot_task, system_prompt, batch_cond_images=None, cond_inputs=None): + generation_config = self.hunyuan_generation_config + tokenizer_out = self.hunyuan_tokenizer.apply_chat_template( + batch_prompt=[prompt], + mode="gen_text", + batch_cond_images=batch_cond_images, + batch_system_prompt=[system_prompt], + max_length=getattr(generation_config, "max_length", self.config.get("max_position_embeddings")), + bot_task=bot_task, + image_base_size=self.hunyuan_image_processor.vae_reso_group.base_size, + sequence_template=getattr(generation_config, "sequence_template", "pretrain"), + cfg_factor=1, + drop_think=getattr(generation_config, "drop_think", False), + ) + latent_device = self._pipeline_latent_device() + output = tokenizer_out["output"] + return { + "input_ids": output.tokens.to(latent_device), + "tokenizer_output": output, + "cond_inputs": self._attach_tokenizer_cond_masks(cond_inputs, output, latent_device), + } + + def _sample_text_token(self, logits, generator): + do_sample = bool(self.config.get("text_do_sample", getattr(self.hunyuan_generation_config, "do_sample", False))) + logits = logits.float() + if not do_sample: + return torch.argmax(logits, dim=-1) + + temperature = float(self.config.get("text_temperature", getattr(self.hunyuan_generation_config, "temperature", 1.0))) + if temperature > 0: + logits = logits / temperature + + top_k = int(self.config.get("text_top_k", getattr(self.hunyuan_generation_config, "top_k", 0)) or 0) + if top_k > 0 and top_k < logits.shape[-1]: + kth_values = torch.topk(logits, top_k, dim=-1).values[..., -1, None] + logits = logits.masked_fill(logits < kth_values, torch.finfo(logits.dtype).min) + + top_p = float(self.config.get("text_top_p", getattr(self.hunyuan_generation_config, "top_p", 1.0))) + if 0.0 < top_p < 1.0: + sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1) + cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1) + sorted_remove = cumulative_probs > top_p + sorted_remove[..., 1:] = sorted_remove[..., :-1].clone() + sorted_remove[..., 0] = False + remove = torch.zeros_like(sorted_remove).scatter(dim=-1, index=sorted_indices, src=sorted_remove) + logits = logits.masked_fill(remove, torch.finfo(logits.dtype).min) + + probs = torch.softmax(logits, dim=-1) + return torch.multinomial(probs, num_samples=1, generator=generator).squeeze(-1) + + def _text_stream_enabled(self): + return bool(self.config.get("text_stream_output", True)) + + def _decode_generated_text_token(self, token_id): + token_tensor = torch.tensor([token_id], dtype=torch.long) + return self.hunyuan_tokenizer.decode(token_tensor) + + def _print_text_generation_chunk(self, text="", end=""): + print(text, end=end, flush=True) + + @torch.no_grad() + def _generate_text_tokens(self, input_ids, tokenizer_output, plan, stream_output=False, cond_inputs=None): + device = input_ids.device + seed = int(self.config.get("text_seed", self.config.get("seed", 42))) + generator = torch.Generator(device=device).manual_seed(seed) + max_new_tokens = int(self.config.get("max_new_tokens", getattr(self.hunyuan_generation_config, "max_new_tokens", 2048))) + transition_map = {stop_id: list(append_ids) for stop_id, append_ids in plan.stage_transitions} + completed_transitions = set() + pending_tokens = [] + generated = [] + use_kv_cache = self._hunyuan_text_kv_cache_enabled() + kv_cache = None + cache_filled_length = 0 + if use_kv_cache: + kv_cache = HunyuanImage3StaticKVCache( + num_layers=self._hunyuan_num_layers(), + max_cache_len=input_ids.shape[1] + max_new_tokens + sum(len(tokens) for tokens in transition_map.values()), + dynamic=True, + ) + + for _ in range(max_new_tokens): + if pending_tokens: + next_token_id = pending_tokens.pop(0) + next_token = torch.tensor([next_token_id], device=device, dtype=input_ids.dtype) + else: + if use_kv_cache: + model_input_ids = input_ids[:, cache_filled_length:] + position_ids = torch.arange(cache_filled_length, input_ids.shape[1], dtype=torch.long, device=device)[None].expand( + input_ids.shape[0], -1 + ) + first_cache_step = cache_filled_length == 0 + model_inputs = self._build_text_model_inputs( + model_input_ids, + tokenizer_output, + cond_inputs=cond_inputs if first_cache_step else None, + position_ids=position_ids, + build_attention_mask=first_cache_step, + past_key_values=kv_cache, + use_cache=True, + ) + cache_filled_length = input_ids.shape[1] + else: + if cond_inputs is None: + model_inputs = self._build_text_model_inputs(input_ids, tokenizer_output) + else: + model_inputs = self._build_text_model_inputs(input_ids, tokenizer_output, cond_inputs=cond_inputs) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"): + logits = self.model.infer(model_inputs)["logits"][:, -1, :] + next_token = self._sample_text_token(logits, generator) + next_token = next_token.to(device=device, dtype=input_ids.dtype) + next_token_id = int(next_token.item()) + + generated.append(next_token_id) + if stream_output: + self._print_text_generation_chunk(self._decode_generated_text_token(next_token_id)) + input_ids = torch.cat([input_ids, next_token.reshape(1, 1)], dim=1) + + if next_token_id in transition_map and next_token_id not in completed_transitions: + completed_transitions.add(next_token_id) + pending_tokens.extend(transition_map[next_token_id]) + continue + if next_token_id in plan.final_stop_tokens and not pending_tokens: + break + + return generated + + def _decode_cot_text(self, generated_tokens, plan): + if not generated_tokens: + return None + token_tensor = torch.tensor(generated_tokens, dtype=torch.long) + generated_text = self.hunyuan_tokenizer.decode(token_tensor) + if plan.first_bot_task == "think": + return self.hunyuan_tokenizer.think_token + generated_text + return self.hunyuan_tokenizer.recaption_token + generated_text + + def _generate_cot_text(self, prompt, image_size, batch_cond_images=None, cond_inputs=None): + bot_task = self._resolve_bot_task() + if bot_task == "image": + return None + # Generate COT text for bot_task in ("think", "recaption", "think_recaption"), default is "think_recaption" + plan = self._resolve_text_generation_plan(bot_task) + system_prompt = self._resolve_system_prompt(plan.first_bot_task) + if batch_cond_images is None and cond_inputs is None: + prepared = self._prepare_text_generation_inputs(prompt, plan.first_bot_task, system_prompt) + else: + prepared = self._prepare_text_generation_inputs( + prompt, + plan.first_bot_task, + system_prompt, + batch_cond_images=batch_cond_images, + cond_inputs=cond_inputs, + ) + stream_output = self._text_stream_enabled() + if stream_output: + prefix = self.hunyuan_tokenizer.think_token if plan.first_bot_task == "think" else self.hunyuan_tokenizer.recaption_token + self._print_text_generation_chunk(prefix) + generated_tokens = self._generate_text_tokens( + prepared["input_ids"], + prepared["tokenizer_output"], + plan, + stream_output=stream_output, + cond_inputs=prepared.get("cond_inputs"), + ) + if stream_output: + self._print_text_generation_chunk("\n") + cot_text = self._decode_cot_text(generated_tokens, plan) + if cot_text: + logger.info(f"HunyuanImage3 generated COT text for bot_task={bot_task}: {cot_text}") + return cot_text + + def _prepare_text_to_image_inputs(self, prompt, image_size, seed, cot_text=None, batch_cond_images=None, cond_inputs=None): + do_cfg = bool(self.config.get("enable_cfg", False)) and not bool(self.config.get("cfg_distilled", False)) + cfg_factor = 2 if do_cfg else 1 + gen_image_info = self.hunyuan_image_processor.build_gen_image_info( + image_size, + add_guidance_token=bool(self.config.get("cfg_distilled", False)), + add_timestep_r_token=bool(self.config.get("use_meanflow", False)), + ) + generation_config = self.hunyuan_generation_config + tokenizer_out = self.hunyuan_tokenizer.apply_chat_template( + batch_prompt=[prompt], + mode="gen_image", + batch_gen_image_info=[gen_image_info], + batch_cond_images=batch_cond_images, + batch_system_prompt=[self._resolve_system_prompt("image")], + batch_cot_text=[cot_text] if cot_text else None, + max_length=getattr(generation_config, "max_length", self.config.get("max_position_embeddings")), + bot_task="image", + image_base_size=self.hunyuan_image_processor.vae_reso_group.base_size, + sequence_template=getattr(generation_config, "sequence_template", "pretrain"), + cfg_factor=cfg_factor, + drop_think=getattr(generation_config, "drop_think", False), + ) + + latent_device = self._pipeline_latent_device() + output = tokenizer_out["output"] + input_ids = output.tokens.to(latent_device) + position_ids = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=latent_device)[None].expand(input_ids.shape[0], -1) + rope_image_info = self._build_batch_rope_image_info(output, tokenizer_out["sections"]) + custom_pos_emb = self._build_custom_pos_emb(position_ids, rope_image_info) + + generator = torch.Generator(device=latent_device).manual_seed(int(seed)) + prepared_inputs = { + "input_ids": input_ids, + "position_ids": position_ids, + "attention_mask": self._build_attention_mask(input_ids, output), + "custom_pos_emb": custom_pos_emb, + "rope_image_info": rope_image_info, + "image_mask": output.gen_image_mask.to(latent_device), + "timesteps_index": None if output.gen_timestep_scatter_index is None else output.gen_timestep_scatter_index.to(latent_device), + "guidance_index": None if output.guidance_scatter_index is None else output.guidance_scatter_index.to(latent_device), + "timesteps_r_index": None if output.gen_timestep_r_scatter_index is None else output.gen_timestep_r_scatter_index.to(latent_device), + "generator": generator, + "do_cfg": do_cfg, + "batch_size": 1, + } + prepared_inputs.update(self._attach_tokenizer_cond_masks(cond_inputs, output, latent_device) or {}) + return prepared_inputs + + def _split_image_paths(self, image_path): + if not image_path: + raise ValueError("HunyuanImage3 ti2i requires --image_path with one or more comma-separated images.") + if isinstance(image_path, (list, tuple)): + return [str(path) for path in image_path if str(path)] + return [path.strip() for path in str(image_path).split(",") if path.strip()] + + def _build_batch_cond_images(self, image_paths, infer_align_image_size): + return [ + self.hunyuan_image_processor.build_cond_images( + image_list=image_paths, + infer_align_image_size=infer_align_image_size, + ) + ] + + def _cond_vae_image(self, cond_image): + return cond_image.vae_image if getattr(cond_image, "section_type", None) == "cond_joint_image" else cond_image + + def _cond_vit_image(self, cond_image): + return cond_image.vit_image if getattr(cond_image, "section_type", None) == "cond_joint_image" else cond_image + + def _resolve_ti2i_image_size(self, requested_image_size, batch_cond_images): + if requested_image_size != "auto": + return requested_image_size + if not batch_cond_images or not batch_cond_images[0]: + return int(self.config.get("target_height", 1024)), int(self.config.get("target_width", 1024)) + first_cond_image = batch_cond_images[0][0] + vae_image = self._cond_vae_image(first_cond_image) + if hasattr(vae_image, "i"): + return int(vae_image.i.image_height), int(vae_image.i.image_width) + return int(self.config.get("target_height", 1024)), int(self.config.get("target_width", 1024)) + + def _vae_encode_cond_tensor(self, image_tensor, generator=None): + vae = self.vae_decoder + vae_device = self._vae_device() + image_tensor = image_tensor.unsqueeze(0).to(vae_device) + autocast_dtype = getattr(torch, self.hunyuan_config.vae_autocast_dtype, torch.float16) + with torch.autocast(device_type="cuda", dtype=autocast_dtype, enabled=vae_device.type == "cuda" and autocast_dtype != torch.float32): + encoded = vae.encode(image_tensor) + latents = encoded if torch.is_tensor(encoded) else encoded.latent_dist.sample(generator) + if hasattr(vae.config, "shift_factor") and vae.config.shift_factor: + latents = latents - vae.config.shift_factor + if hasattr(vae.config, "scaling_factor") and vae.config.scaling_factor: + latents = latents * vae.config.scaling_factor + if hasattr(vae, "ffactor_temporal"): + latents = latents.squeeze(2) + return latents.squeeze(0).to(self._pipeline_latent_device(), dtype=torch.bfloat16) + + def _encode_cond_vae_images(self, batch_cond_images, cfg_factor, seed): + if not batch_cond_images or not batch_cond_images[0]: + return None, None + first_image = batch_cond_images[0][0] + if getattr(first_image, "section_type", None) not in ("cond_vae_image", "cond_joint_image"): + return None, None + + vae_device = self._vae_device() + generator = torch.Generator(device=vae_device).manual_seed(int(seed)) + batch_cond_vae_images = [] + batch_cond_timesteps = [] + for cond_images in batch_cond_images: + image_items = [] + timestep_items = [] + for cond_image in cond_images: + vae_image = self._cond_vae_image(cond_image) + latent = self._vae_encode_cond_tensor(vae_image, generator=generator) + image_items.append(latent) + timestep_items.append(torch.zeros(1, device=self._pipeline_latent_device(), dtype=torch.bfloat16)) + batch_cond_vae_images.append(image_items) + batch_cond_timesteps.append(timestep_items) + + if all(len(items) == 1 for items in batch_cond_vae_images) and all( + items[0].shape == batch_cond_vae_images[0][0].shape for items in batch_cond_vae_images + ): + cond_vae_images = torch.stack([items[0] for items in batch_cond_vae_images], dim=0) + cond_timesteps = torch.cat([items[0] for items in batch_cond_timesteps], dim=0) + if cfg_factor > 1: + cond_vae_images = cond_vae_images.repeat(cfg_factor, 1, 1, 1) + cond_timesteps = cond_timesteps.repeat(cfg_factor) + return cond_vae_images, cond_timesteps + + cond_vae_images = [] + cond_timesteps = [torch.cat(items, dim=0) for items in batch_cond_timesteps] + for items in batch_cond_vae_images: + if all(items[0].shape == item.shape for item in items): + cond_vae_images.append(torch.stack(items, dim=0)) + else: + cond_vae_images.append(items) + if cfg_factor > 1: + cond_vae_images = cond_vae_images * cfg_factor + cond_timesteps = cond_timesteps * cfg_factor + return cond_vae_images, cond_timesteps + + def _encode_cond_vit_embeds(self, batch_cond_images, cfg_factor): + if not batch_cond_images or not batch_cond_images[0]: + return None + first_image = batch_cond_images[0][0] + if getattr(first_image, "section_type", None) not in ("cond_vit_image", "cond_joint_image"): + return None + + self._ensure_vision_encoder() + vision_device = self._vision_device() + cond_vit_embeds = [] + for cond_images in batch_cond_images: + vit_images = [self._cond_vit_image(cond_image) for cond_image in cond_images] + pixel_values = torch.stack(vit_images, dim=0).to(vision_device) + image_kwargs = {} + first_vit_image = vit_images[0] + if hasattr(first_vit_image, "vision_encoder_kwargs") and first_vit_image.vision_encoder_kwargs: + image_kwargs = { + "spatial_shapes": torch.stack([vit_image.vision_encoder_kwargs["spatial_shapes"] for vit_image in vit_images], dim=0).to(vision_device), + "attention_mask": torch.stack([vit_image.vision_encoder_kwargs["pixel_attention_mask"] for vit_image in vit_images], dim=0).to(vision_device), + } + with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=vision_device.type == "cuda"): + image_embeds = self.vision_model(pixel_values, **image_kwargs).last_hidden_state + image_embeds = self.vision_aligner(image_embeds) + cond_vit_embeds.append(image_embeds.to(self._pipeline_latent_device(), dtype=torch.bfloat16)) + + if cfg_factor > 1: + cond_vit_embeds = cond_vit_embeds * cfg_factor + return cond_vit_embeds + + def _prepare_cond_inputs(self, batch_cond_images, cfg_factor, seed): + cond_vae_images, cond_timesteps = self._encode_cond_vae_images(batch_cond_images, cfg_factor, seed) + cond_vit_embeds = self._encode_cond_vit_embeds(batch_cond_images, cfg_factor) + return { + "cond_vae_images": cond_vae_images, + "cond_timesteps": cond_timesteps, + "cond_vit_embeds": cond_vit_embeds, + } + + def _prepare_latents(self, batch_size, image_size, generator): + latent_device = self._pipeline_latent_device() + downsample = self.hunyuan_config.vae_downsample_factor + if isinstance(downsample, int): + downsample_h = downsample_w = downsample + else: + downsample_h, downsample_w = downsample[:2] + height, width = image_size + shape = ( + batch_size, + self.hunyuan_config.vae["latent_channels"], + height // downsample_h, + width // downsample_w, + ) + return torch.randn(shape, generator=generator, device=latent_device, dtype=torch.bfloat16) + + def _denoise_latents(self, prepared_inputs, image_size): + latents = self._prepare_latents(prepared_inputs["batch_size"], image_size, prepared_inputs["generator"]) + num_steps = int(self.config.get("infer_steps", self.config.get("diff_infer_steps", 50))) + self.scheduler.set_timesteps(num_steps, device=latents.device) + guidance_scale = float(self.config.get("sample_guide_scale", self.config.get("diff_guidance_scale", 1.0))) + use_kv_cache = self._hunyuan_kv_cache_enabled() + taylor_cache_dic = self._build_taylor_cache_dic(num_steps) if self._hunyuan_taylor_cache_enabled() else None + if taylor_cache_dic is not None and hasattr(self.model, "reset_taylor_cache"): + self.model.reset_taylor_cache() + if taylor_cache_dic is not None: + logger.info(f"HunyuanImage3 Taylor cache enabled: {taylor_cache_dic}") + kv_cache = None + cache_position_ids = None + cache_local_inputs = None + if use_kv_cache: + kv_cache = HunyuanImage3StaticKVCache( + num_layers=self._hunyuan_num_layers(), + max_cache_len=prepared_inputs["input_ids"].shape[1], + ) + cache_position_ids = self._build_denoise_cache_position_ids(prepared_inputs) + cache_local_inputs = self._build_denoise_cache_local_indices(prepared_inputs, cache_position_ids) + + for step_index, timestep in enumerate(self.scheduler.timesteps): + cfg_factor = 2 if prepared_inputs["do_cfg"] else 1 + latent_model_input = torch.cat([latents] * cfg_factor) if cfg_factor > 1 else latents + timestep_input = timestep.repeat(latent_model_input.shape[0]) + first_step = step_index == 0 or not use_kv_cache + if first_step: + model_inputs = { + "input_ids": prepared_inputs["input_ids"], + "attention_mask": prepared_inputs["attention_mask"], + "position_ids": prepared_inputs["position_ids"], + "custom_pos_emb": prepared_inputs["custom_pos_emb"], + "images": latent_model_input, + "image_mask": prepared_inputs["image_mask"], + "timesteps": timestep_input, + "timesteps_index": prepared_inputs["timesteps_index"], + "guidance_index": prepared_inputs["guidance_index"], + "timesteps_r_index": prepared_inputs["timesteps_r_index"], + "first_step": True, + } + if prepared_inputs.get("cond_vae_images") is not None: + model_inputs["cond_vae_images"] = prepared_inputs["cond_vae_images"] + model_inputs["cond_vae_image_mask"] = prepared_inputs.get("cond_vae_image_mask") + model_inputs["cond_timesteps"] = prepared_inputs.get("cond_timesteps") + model_inputs["cond_timestep_index"] = prepared_inputs.get("cond_timestep_index") + if prepared_inputs.get("cond_vit_embeds") is not None: + model_inputs["cond_vit_embeds"] = prepared_inputs["cond_vit_embeds"] + model_inputs["cond_vit_image_mask"] = prepared_inputs.get("cond_vit_image_mask") + else: + model_inputs = { + "input_ids": None, + "attention_mask": self._slice_denoise_cache_attention_mask(prepared_inputs["attention_mask"], cache_position_ids), + "position_ids": cache_position_ids, + "custom_pos_emb": self._build_custom_pos_emb(cache_position_ids, prepared_inputs.get("rope_image_info")), + "images": latent_model_input, + "image_mask": cache_local_inputs["image_mask"], + "timesteps": timestep_input, + "timesteps_index": cache_local_inputs["timesteps_index"], + "guidance_index": cache_local_inputs["guidance_index"], + "timesteps_r_index": cache_local_inputs["timesteps_r_index"], + "first_step": False, + } + if use_kv_cache: + model_inputs["past_key_values"] = kv_cache + model_inputs["use_cache"] = True + if taylor_cache_dic is not None: + taylor_cache_dic["current_step"] = step_index + model_inputs["cache_dic"] = taylor_cache_dic + if self.config.get("cfg_distilled", False): + model_inputs["guidance"] = torch.tensor([1000.0 * guidance_scale], device=latents.device, dtype=torch.bfloat16) + if self.config.get("use_meanflow", False): + raise NotImplementedError("HunyuanImage3 native meanflow sampling is not implemented yet.") + + with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=latents.device.type == "cuda"): + prediction = self.model.infer(model_inputs)["diffusion_prediction"].to(latents.device) + + if prepared_inputs["do_cfg"]: + pred_cond, pred_uncond = prediction.chunk(2) + prediction = pred_uncond + guidance_scale * (pred_cond - pred_uncond) + + sigma = self.scheduler.sigmas[step_index].to(latents.device) + sigma_next = self.scheduler.sigmas[step_index + 1].to(latents.device) + latents = latents.float() + prediction.float() * (sigma_next - sigma) + + return latents + + def _decode_latents(self, latents, generator): + vae = self.vae_decoder + vae_device = self._vae_device() + latents = latents.to(vae_device) + if hasattr(vae.config, "scaling_factor") and vae.config.scaling_factor: + latents = latents / vae.config.scaling_factor + if hasattr(vae.config, "shift_factor") and vae.config.shift_factor: + latents = latents + vae.config.shift_factor + if hasattr(vae, "ffactor_temporal"): + latents = latents.unsqueeze(2) + + autocast_dtype = getattr(torch, self.hunyuan_config.vae_autocast_dtype, torch.float16) + with torch.autocast(device_type="cuda", dtype=autocast_dtype, enabled=vae_device.type == "cuda" and autocast_dtype != torch.float32): + image = vae.decode(latents, return_dict=False, generator=generator)[0] + if hasattr(vae, "ffactor_temporal"): + image = image.squeeze(2) + + image = (image / 2 + 0.5).clamp(0, 1).detach().cpu().permute(0, 2, 3, 1).float().numpy() + return [Image.fromarray((sample * 255.0).round().astype("uint8")) for sample in image] + + @torch.no_grad() + def generate_t2i(self, input_info): + self._ensure_pipeline_modules() + prompt = getattr(input_info, "prompt_enhanced", None) or getattr(input_info, "prompt", "") + image_size = self._resolve_image_size(input_info) + seed = getattr(input_info, "seed", None) or self.config.get("seed", 42) + cot_text = self._generate_cot_text(prompt, image_size) + prepared_inputs = self._prepare_text_to_image_inputs(prompt, image_size, seed, cot_text=cot_text) + latents = self._denoise_latents(prepared_inputs, image_size) + return self._decode_latents(latents, prepared_inputs["generator"]) + + @torch.no_grad() + def generate_i2i(self, input_info): + self._ensure_pipeline_modules() + prompt = getattr(input_info, "prompt_enhanced", None) or getattr(input_info, "prompt", "") + seed = getattr(input_info, "seed", None) or self.config.get("seed", 42) + image_paths = self._split_image_paths(getattr(input_info, "image_path", None) or self.config.get("image_path")) + infer_align_image_size = bool(getattr(input_info, "infer_align_image_size", self.config.get("infer_align_image_size", False))) + + batch_cond_images = self._build_batch_cond_images(image_paths, infer_align_image_size) + image_size = self._resolve_ti2i_image_size(self._resolve_image_size(input_info), batch_cond_images) + + text_cond_inputs = self._prepare_cond_inputs(batch_cond_images, cfg_factor=1, seed=seed) + cot_text = self._generate_cot_text(prompt, image_size, batch_cond_images=batch_cond_images, cond_inputs=text_cond_inputs) + + # in default, we enable CFG for i2i/ti2i, but disable CFG if cfg_distilled is True + do_cfg = bool(self.config.get("enable_cfg", False)) and not bool(self.config.get("cfg_distilled", False)) + gen_cond_inputs = self._prepare_cond_inputs(batch_cond_images, cfg_factor=2 if do_cfg else 1, seed=seed) + prepared_inputs = self._prepare_text_to_image_inputs( + prompt, + image_size, + seed, + cot_text=cot_text, + batch_cond_images=batch_cond_images, + cond_inputs=gen_cond_inputs, + ) + latents = self._denoise_latents(prepared_inputs, image_size) + images = self._decode_latents(latents, prepared_inputs["generator"]) + return self.hunyuan_image_processor.postprocess_outputs( + images, + batch_cond_images=batch_cond_images, + infer_align_image_size=infer_align_image_size, + ) + + def run_pipeline(self, input_info): + self.input_info = input_info + task = self.config.get("task") + if task == "t2i": + images = self.generate_t2i(input_info) + elif task in ("i2i", "ti2i"): + images = self.generate_i2i(input_info) + else: + raise NotImplementedError("HunyuanImage3 native runner currently supports task=t2i/i2i.") + + if getattr(input_info, "return_result_tensor", False): + return {"image": images} + + save_result_path = getattr(input_info, "save_result_path", None) + if save_result_path: + Path(save_result_path).parent.mkdir(parents=True, exist_ok=True) + images[0].save(save_result_path) + logger.info(f"HunyuanImage3 image saved successfully to: {save_result_path}") + return {"image": None} diff --git a/lightx2v/models/schedulers/hunyuan_image3/__init__.py b/lightx2v/models/schedulers/hunyuan_image3/__init__.py new file mode 100644 index 000000000..645f4e8c5 --- /dev/null +++ b/lightx2v/models/schedulers/hunyuan_image3/__init__.py @@ -0,0 +1,3 @@ +from .scheduler import HunyuanImage3Scheduler + +__all__ = ["HunyuanImage3Scheduler"] diff --git a/lightx2v/models/schedulers/hunyuan_image3/scheduler.py b/lightx2v/models/schedulers/hunyuan_image3/scheduler.py new file mode 100644 index 000000000..02e9e1a60 --- /dev/null +++ b/lightx2v/models/schedulers/hunyuan_image3/scheduler.py @@ -0,0 +1,40 @@ +import torch + +from lightx2v.models.schedulers.scheduler import BaseScheduler +from lightx2v_platform.base.global_var import AI_DEVICE + + +class HunyuanImage3Scheduler(BaseScheduler): + def __init__(self, config): + super().__init__(config) + self.sample_guide_scale = config.get("sample_guide_scale", config.get("diff_guidance_scale", 1.0)) + self.flow_shift = config.get("sample_shift", config.get("flow_shift", 1.0)) + self.timesteps = None + self.sigmas = None + self.noise_pred = None + + def prepare(self, input_info): + seed = getattr(input_info, "seed", None) or self.config.get("seed", 42) + self.generator = torch.Generator(device=AI_DEVICE).manual_seed(seed) + + def set_timesteps(self, num_inference_steps=None, device=None): + num_inference_steps = num_inference_steps or self.infer_steps + device = device or AI_DEVICE + sigmas = torch.linspace(1, 0, num_inference_steps + 1, device=device) + if self.flow_shift != 1.0: + sigmas = (self.flow_shift * sigmas) / (1 + (self.flow_shift - 1) * sigmas) + self.sigmas = sigmas + self.timesteps = (sigmas[:-1] * 1000).to(dtype=torch.float32) + + def step_pre(self, step_index): + super().step_pre(step_index) + if self.timesteps is None: + self.set_timesteps(self.infer_steps, device=self.latents.device if self.latents is not None else AI_DEVICE) + self.timestep_input = self.timesteps[step_index : step_index + 1] + + def step_post(self): + if self.latents is None or self.noise_pred is None: + return + sigma = self.sigmas[self.step_index] + sigma_next = self.sigmas[self.step_index + 1] + self.latents = self.latents.float() + self.noise_pred.float() * (sigma_next - sigma) diff --git a/lightx2v/utils/set_config.py b/lightx2v/utils/set_config.py index fd4795e18..338e791eb 100755 --- a/lightx2v/utils/set_config.py +++ b/lightx2v/utils/set_config.py @@ -39,6 +39,29 @@ def set_args2config(args): config = get_default_config() config.update({k: v for k, v in vars(args).items() if k not in ALL_INPUT_INFO_KEYS and v is not None}) + _hunyuan_image3_cli_cache_keys = ( + "enable_kv_cache", + "enable_text_kv_cache", + "use_taylor_cache", + "taylor_cache_interval", + "taylor_cache_order", + "taylor_cache_enable_first_enhance", + "taylor_cache_first_enhance_steps", + "taylor_cache_enable_tailing_enhance", + "taylor_cache_tailing_enhance_steps", + "taylor_cache_low_freqs_order", + "taylor_cache_high_freqs_order", + ) + config["_hunyuan_image3_cli_cache_snapshot"] = { + k: getattr(args, k, None) for k in _hunyuan_image3_cli_cache_keys if hasattr(args, k) and getattr(args, k, None) is not None + } + _hunyuan_image3_cli_native_keys = ( + "moe_impl", + ) + config["_hunyuan_image3_cli_native_snapshot"] = { + k: getattr(args, k, None) for k in _hunyuan_image3_cli_native_keys if hasattr(args, k) and getattr(args, k, None) is not None + } + # Snapshot worldmirror-specific CLI flags so they can win over JSON # defaults after auto_calc_config merges the JSON in. See the # ``worldmirror`` branch of ``auto_calc_config`` for the replay logic. @@ -215,6 +238,54 @@ def auto_calc_config(config): if "infer_steps" not in config and "num_inference_steps" in config: config["infer_steps"] = config["num_inference_steps"] + if config["model_cls"] == "hunyuan_image3": + if "num_layers" not in config and "num_hidden_layers" in config: + config["num_layers"] = config["num_hidden_layers"] + if "num_heads" not in config and "num_attention_heads" in config: + config["num_heads"] = config["num_attention_heads"] + if "num_attention_heads" not in config and "num_heads" in config: + config["num_attention_heads"] = config["num_heads"] + if "attention_head_dim" not in config and "hidden_size" in config and "num_attention_heads" in config: + config["attention_head_dim"] = config["hidden_size"] // config["num_attention_heads"] + if "infer_steps" not in config and "diff_infer_steps" in config: + config["infer_steps"] = config["diff_infer_steps"] + if "sample_guide_scale" not in config and "diff_guidance_scale" in config: + config["sample_guide_scale"] = config["diff_guidance_scale"] + if "sample_shift" not in config and "flow_shift" in config: + config["sample_shift"] = config["flow_shift"] + config.setdefault("img_proj_type", "unet") + if config.get("patch_size") is None or isinstance(config.get("patch_size"), (list, tuple)): + config["patch_size"] = 1 + if config.get("patch_embed_hidden_dim") is None: + config["patch_embed_hidden_dim"] = 1024 + config["cfg_distilled"] = bool(config.get("cfg_distilled", False)) + config["use_meanflow"] = bool(config.get("use_meanflow", False)) + config.setdefault("use_qk_norm", True) + config.setdefault("rms_norm_eps", 1e-5) + config.setdefault("rms_norm_type", "fp32_variance") + config.setdefault("hidden_act", "silu") + config.setdefault("mlp_bias", False) + config.setdefault("attention_bias", False) + config.setdefault("moe_layer_num_skipped", 0) + config.setdefault("use_mixed_mlp_moe", False) + config.setdefault("moe_impl", "eager") + config.setdefault("pipeline_parallel", True) + if "vae_scale_factor" not in config: + vae_downsample_factor = config.get("vae_downsample_factor") + if isinstance(vae_downsample_factor, list) and vae_downsample_factor: + config["vae_scale_factor"] = int(vae_downsample_factor[0]) + _hunyuan_image3_cli_cache_snapshot = config.pop("_hunyuan_image3_cli_cache_snapshot", None) + if isinstance(_hunyuan_image3_cli_cache_snapshot, dict): + for k, v in _hunyuan_image3_cli_cache_snapshot.items(): + config[k] = v + _hunyuan_image3_cli_native_snapshot = config.pop("_hunyuan_image3_cli_native_snapshot", None) + if isinstance(_hunyuan_image3_cli_native_snapshot, dict): + for k, v in _hunyuan_image3_cli_native_snapshot.items(): + config[k] = v + + config.pop("_hunyuan_image3_cli_cache_snapshot", None) + config.pop("_hunyuan_image3_cli_native_snapshot", None) + if config["task"] in ["i2v", "s2v", "rs2v", "ltx2_s2v", "v2av"]: if config["target_video_length"] % config["vae_stride"][0] != 1: logger.warning(f"`num_frames - 1` has to be divisible by {config['vae_stride'][0]}. Rounding to the nearest number.") diff --git a/scripts/hunyuan_image3/run_hunyuan_image3_i2i.sh b/scripts/hunyuan_image3/run_hunyuan_image3_i2i.sh new file mode 100644 index 000000000..585eff589 --- /dev/null +++ b/scripts/hunyuan_image3/run_hunyuan_image3_i2i.sh @@ -0,0 +1,74 @@ +#!/bin/bash + +export lightx2v_path=${lightx2v_path:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)} +export model_path=${model_path:-/path/to/HunyuanImage-3-Instruct} +export image_path=${image_path:-/path/to/input.png} +export ENABLE_KV_CACHE="${ENABLE_KV_CACHE:-true}" +export ENABLE_TEXT_KV_CACHE="${ENABLE_TEXT_KV_CACHE:-$ENABLE_KV_CACHE}" +export USE_TAYLOR_CACHE="${USE_TAYLOR_CACHE:-false}" +export TAYLOR_CACHE_INTERVAL="${TAYLOR_CACHE_INTERVAL:-5}" +export TAYLOR_CACHE_ORDER="${TAYLOR_CACHE_ORDER:-2}" +export TAYLOR_CACHE_ENABLE_FIRST_ENHANCE="${TAYLOR_CACHE_ENABLE_FIRST_ENHANCE:-false}" +export TAYLOR_CACHE_FIRST_ENHANCE_STEPS="${TAYLOR_CACHE_FIRST_ENHANCE_STEPS:-3}" +export TAYLOR_CACHE_ENABLE_TAILING_ENHANCE="${TAYLOR_CACHE_ENABLE_TAILING_ENHANCE:-false}" +export TAYLOR_CACHE_TAILING_ENHANCE_STEPS="${TAYLOR_CACHE_TAILING_ENHANCE_STEPS:-1}" +export TAYLOR_CACHE_LOW_FREQS_ORDER="${TAYLOR_CACHE_LOW_FREQS_ORDER:-2}" +export TAYLOR_CACHE_HIGH_FREQS_ORDER="${TAYLOR_CACHE_HIGH_FREQS_ORDER:-2}" + +if [ -z "${CUDA_VISIBLE_DEVICES:-}" ] && command -v nvidia-smi >/dev/null 2>&1; then + gpu_count=$(nvidia-smi -L | wc -l) + if [ "$gpu_count" -gt 0 ]; then + export CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((gpu_count - 1))) + fi +fi + +source ${lightx2v_path}/scripts/base/base.sh + +echo "ENABLE_KV_CACHE=${ENABLE_KV_CACHE}" +echo "ENABLE_TEXT_KV_CACHE=${ENABLE_TEXT_KV_CACHE}" +echo "USE_TAYLOR_CACHE=${USE_TAYLOR_CACHE}" +echo "moe_impl=${moe_impl:-flashinfer}" + +base_config="${lightx2v_path}/configs/hunyuan_image3/hunyuan_image3_i2i.json" +runtime_config="$(mktemp /tmp/hunyuan_image3_i2i_config.XXXXXX.json)" +trap 'rm -f "$runtime_config"' EXIT + +python - "$base_config" "$runtime_config" <<'PY' +import json +import os +import sys + +def env_bool(name, default=False): + return os.environ.get(name, str(default)).lower() in ("1", "true", "yes", "on") + +def env_int(name, default): + return int(os.environ.get(name, default)) + +base_config, runtime_config = sys.argv[1:3] +with open(base_config, "r") as f: + config = json.load(f) +config["enable_kv_cache"] = env_bool("ENABLE_KV_CACHE", True) +config["enable_text_kv_cache"] = env_bool("ENABLE_TEXT_KV_CACHE", config["enable_kv_cache"]) +config["use_taylor_cache"] = env_bool("USE_TAYLOR_CACHE", False) +config["taylor_cache_interval"] = env_int("TAYLOR_CACHE_INTERVAL", 5) +config["taylor_cache_order"] = env_int("TAYLOR_CACHE_ORDER", 2) +config["taylor_cache_enable_first_enhance"] = env_bool("TAYLOR_CACHE_ENABLE_FIRST_ENHANCE", False) +config["taylor_cache_first_enhance_steps"] = env_int("TAYLOR_CACHE_FIRST_ENHANCE_STEPS", 3) +config["taylor_cache_enable_tailing_enhance"] = env_bool("TAYLOR_CACHE_ENABLE_TAILING_ENHANCE", False) +config["taylor_cache_tailing_enhance_steps"] = env_int("TAYLOR_CACHE_TAILING_ENHANCE_STEPS", 1) +config["taylor_cache_low_freqs_order"] = env_int("TAYLOR_CACHE_LOW_FREQS_ORDER", 2) +config["taylor_cache_high_freqs_order"] = env_int("TAYLOR_CACHE_HIGH_FREQS_ORDER", 2) +with open(runtime_config, "w") as f: + json.dump(config, f, indent=4) +PY + +python -m lightx2v.infer \ + --model_cls hunyuan_image3 \ + --task i2i \ + --model_path $model_path \ + --config_json "$runtime_config" \ + --prompt "新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。" \ + --image_path "${image_path}" \ + --save_result_path ${lightx2v_path}/save_results/hunyuan_image3_i2i.png \ + --seed 42 \ + --moe_impl "${moe_impl:-flashinfer}" diff --git a/scripts/hunyuan_image3/run_hunyuan_image3_t2i.sh b/scripts/hunyuan_image3/run_hunyuan_image3_t2i.sh new file mode 100644 index 000000000..3f4b95502 --- /dev/null +++ b/scripts/hunyuan_image3/run_hunyuan_image3_t2i.sh @@ -0,0 +1,58 @@ +#!/bin/bash +set -e + +# 用法： +# bash run_hunyuan.sh 0 +# bash run_hunyuan.sh 1,2 +# bash run_hunyuan.sh 3,5,6 +# +# 如果不传参数，则使用已有的 CUDA_VISIBLE_DEVICES +# 如果也没有设置 CUDA_VISIBLE_DEVICES，则自动使用所有 GPU +# Cache 开关直接写在 python 入口参数里，测试时改入口处的 true/false/数字即可。 + +GPU_IDS="${1:-${CUDA_VISIBLE_DEVICES:-}}" + +export lightx2v_path="/data/nvme0/lhd_codes/LightX2V" +export model_path="/data/nvme0/lhd_codes/HunyuanImage-3.0-instruct/HunyuanImage-3-Instruct" +export HUNYUAN_IMAGE3_REPO_PATH="${HUNYUAN_IMAGE3_REPO_PATH:-/data/nvme0/lhd_codes/HunyuanImage-3.0}" +export PYTHONPATH="${HUNYUAN_IMAGE3_REPO_PATH}:${PYTHONPATH:-}" + +if [ -n "$GPU_IDS" ]; then + export CUDA_VISIBLE_DEVICES="$GPU_IDS" +elif command -v nvidia-smi >/dev/null 2>&1; then + gpu_count=$(nvidia-smi -L | wc -l) + if [ "$gpu_count" -gt 0 ]; then + export CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((gpu_count - 1))) + fi +fi + +echo "Using CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-not set}" +echo "enable_kv_cache=${enable_kv_cache:-true}" +echo "enable_text_kv_cache=${enable_text_kv_cache:-${enable_kv_cache:-true}}" +echo "use_taylor_cache=${use_taylor_cache:-false}" +echo "moe_impl=${moe_impl:-flashinfer}" + +source ${lightx2v_path}/scripts/base/base.sh + +config_json="${lightx2v_path}/configs/hunyuan_image3/hunyuan_image3_t2i.json" + +python -m lightx2v.infer \ + --model_cls hunyuan_image3 \ + --task t2i \ + --model_path "$model_path" \ + --config_json "$config_json" \ + --prompt "生成图片：一辆汽车行驶在高速公路上，驾驶员在打电话，副驾驶坐着一只狗" \ + --save_result_path "${lightx2v_path}/save_results/hunyuan_image3_t2i.png" \ + --seed 42 \ + --moe_impl "${moe_impl:-eager}" \ + --enable_kv_cache "${enable_kv_cache:-true}" \ + --enable_text_kv_cache "${enable_text_kv_cache:-${enable_kv_cache:-true}}" \ + --use_taylor_cache "${use_taylor_cache:-false}" \ + --taylor_cache_interval "${taylor_cache_interval:-5}" \ + --taylor_cache_order "${taylor_cache_order:-2}" \ + --taylor_cache_enable_first_enhance "${taylor_cache_enable_first_enhance:-false}" \ + --taylor_cache_first_enhance_steps "${taylor_cache_first_enhance_steps:-3}" \ + --taylor_cache_enable_tailing_enhance "${taylor_cache_enable_tailing_enhance:-false}" \ + --taylor_cache_tailing_enhance_steps "${taylor_cache_tailing_enhance_steps:-1}" \ + --taylor_cache_low_freqs_order "${taylor_cache_low_freqs_order:-2}" \ + --taylor_cache_high_freqs_order "${taylor_cache_high_freqs_order:-2}" diff --git a/scripts/hunyuan_image3/run_hunyuan_image3_ti2i.sh b/scripts/hunyuan_image3/run_hunyuan_image3_ti2i.sh new file mode 100755 index 000000000..6387bf9db --- /dev/null +++ b/scripts/hunyuan_image3/run_hunyuan_image3_ti2i.sh @@ -0,0 +1,84 @@ +#!/bin/bash +set -e + +# Usage: +# bash scripts/hunyuan_image3/run_hunyuan_image3_ti2i.sh +# bash scripts/hunyuan_image3/run_hunyuan_image3_ti2i.sh ti2iv2 +# bash scripts/hunyuan_image3/run_hunyuan_image3_ti2i.sh 0,2,3,4,5,6,7 ti2iv2 + +GPU_IDS="${CUDA_VISIBLE_DEVICES:-}" +DEMO="${DEMO:-ti2i}" + +if [ $# -gt 0 ]; then + case "$1" in + ti2i|ti2iv2) + DEMO="$1" + ;; + *) + GPU_IDS="$1" + ;; + esac +fi + +if [ $# -gt 1 ]; then + DEMO="$2" +fi + +export lightx2v_path="${lightx2v_path:-/data/nvme0/lhd_codes/LightX2V}" +export model_path="${model_path:-/data/nvme0/lhd_codes/HunyuanImage-3.0-instruct/HunyuanImage-3-Instruct}" +export HUNYUAN_IMAGE3_REPO_PATH="${HUNYUAN_IMAGE3_REPO_PATH:-/data/nvme0/lhd_codes/HunyuanImage-3.0}" +export PYTHONPATH="${HUNYUAN_IMAGE3_REPO_PATH}:${PYTHONPATH:-}" + +if [ -n "$GPU_IDS" ]; then + export CUDA_VISIBLE_DEVICES="$GPU_IDS" +elif command -v nvidia-smi >/dev/null 2>&1; then + gpu_count=$(nvidia-smi -L | wc -l) + if [ "$gpu_count" -gt 0 ]; then + export CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((gpu_count - 1))) + fi +fi + +if [ "$DEMO" = "ti2iv2" ]; then + prompt="让图1的猫咪与图2的猫咪自拍，图1的猫咪说:“妈妈，我在乡下遇到了好朋喵”，背景为图3。" + image_path="${HUNYUAN_IMAGE3_REPO_PATH}/assets/demo_instruct_imgs/input_2_0.png,${HUNYUAN_IMAGE3_REPO_PATH}/assets/demo_instruct_imgs/input_2_1.png,${HUNYUAN_IMAGE3_REPO_PATH}/assets/demo_instruct_imgs/input_2_2.png" + seed=43 + save_result_path="${lightx2v_path}/save_results/hunyuan_image3_ti2iv3.png" +else + prompt="新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。 宝丽莱相纸，超现实主义，写实主义，胶片摄影，打印颗粒感肌理。肌理，超写实，复古感。" + image_path="${HUNYUAN_IMAGE3_REPO_PATH}/assets/demo_instruct_imgs/input_0_0.png" + seed=42 + save_result_path="${lightx2v_path}/save_results/hunyuan_image3_ti2i.png" +fi + +echo "Using CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-not set}" +echo "Running HunyuanImage3 ${DEMO}" +echo "enable_kv_cache=${enable_kv_cache:-true}" +echo "enable_text_kv_cache=${enable_text_kv_cache:-${enable_kv_cache:-true}}" +echo "use_taylor_cache=${use_taylor_cache:-false}" +echo "moe_impl=${moe_impl:-flashinfer}" + +source "${lightx2v_path}/scripts/base/base.sh" + +config_json="${lightx2v_path}/configs/hunyuan_image3/hunyuan_image3_i2i.json" + +python -m lightx2v.infer \ + --model_cls hunyuan_image3 \ + --task i2i \ + --model_path "$model_path" \ + --config_json "$config_json" \ + --prompt "$prompt" \ + --image_path "$image_path" \ + --save_result_path "$save_result_path" \ + --seed "$seed" \ + --moe_impl "${moe_impl:-eager}" \ + --enable_kv_cache "${enable_kv_cache:-true}" \ + --enable_text_kv_cache "${enable_text_kv_cache:-${enable_kv_cache:-true}}" \ + --use_taylor_cache "${use_taylor_cache:-false}" \ + --taylor_cache_interval "${taylor_cache_interval:-5}" \ + --taylor_cache_order "${taylor_cache_order:-2}" \ + --taylor_cache_enable_first_enhance "${taylor_cache_enable_first_enhance:-false}" \ + --taylor_cache_first_enhance_steps "${taylor_cache_first_enhance_steps:-3}" \ + --taylor_cache_enable_tailing_enhance "${taylor_cache_enable_tailing_enhance:-false}" \ + --taylor_cache_tailing_enhance_steps "${taylor_cache_tailing_enhance_steps:-1}" \ + --taylor_cache_low_freqs_order "${taylor_cache_low_freqs_order:-2}" \ + --taylor_cache_high_freqs_order "${taylor_cache_high_freqs_order:-2}" diff --git a/scripts/wan/run_wan_i2v.sh b/scripts/wan/run_wan_i2v.sh index ea50e4e20..cc6df50ba 100755 --- a/scripts/wan/run_wan_i2v.sh +++ b/scripts/wan/run_wan_i2v.sh @@ -1,8 +1,8 @@ #!/bin/bash # set path firstly -lightx2v_path= -model_path= +lightx2v_path="/data/nvme0/lhd_codes/LightX2V" +model_path="/data/nvme0/models/Wan-AI/Wan2.1-I2V-14B-720P" export CUDA_VISIBLE_DEVICES=0 diff --git a/scripts/wan/run_wan_t2v.sh b/scripts/wan/run_wan_t2v.sh index a3436d47a..3d801e867 100755 --- a/scripts/wan/run_wan_t2v.sh +++ b/scripts/wan/run_wan_t2v.sh @@ -1,8 +1,8 @@ #!/bin/bash # set path firstly -lightx2v_path= -model_path= +lightx2v_path="/data/nvme0/lhd_codes/LightX2V" +model_path="/data/nvme0/models/Wan-AI/Wan2.1-T2V-1.3B" export CUDA_VISIBLE_DEVICES=0 diff --git a/tests/test_hunyuan_image3_integration.py b/tests/test_hunyuan_image3_integration.py new file mode 100644 index 000000000..2e94329ac --- /dev/null +++ b/tests/test_hunyuan_image3_integration.py @@ -0,0 +1,1297 @@ +import argparse +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from PIL import Image +from safetensors.torch import save_file + +from lightx2v.utils.set_config import set_config + + +def _tiny_hunyuan_config(): + return { + "model_cls": "hunyuan_image3", + "task": "t2i", + "seq_parallel": False, + "cpu_offload": False, + "feature_caching": "NoCaching", + "num_layers": 2, + "hidden_size": 16, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "intermediate_size": 12, + "moe_intermediate_size": 12, + "num_experts": 3, + "moe_topk": 2, + "moe_layer_num_skipped": 1, + "use_mixed_mlp_moe": False, + "hidden_act": "silu", + "mlp_bias": False, + "rms_norm_eps": 1e-5, + "patch_size": 1, + "patch_embed_hidden_dim": 8, + "vae": {"latent_channels": 4}, + "vocab_size": 32, + "cfg_distilled": False, + "use_meanflow": False, + } + + +def _tiny_weight_dict(config): + hidden = config["hidden_size"] + heads = config["num_attention_heads"] + kv_heads = config["num_key_value_heads"] + head_dim = hidden // heads + qkv_out = (heads + 2 * kv_heads) * head_dim + inter = config["intermediate_size"] + moe_inter = config["moe_intermediate_size"] + patch_hidden = config["patch_embed_hidden_dim"] + latent = config["vae"]["latent_channels"] + vocab = config["vocab_size"] + weights = { + "model.wte.weight": torch.randn(vocab, hidden), + "model.ln_f.weight": torch.randn(hidden), + "lm_head.weight": torch.randn(vocab, hidden), + "timestep_emb.mlp.0.weight": torch.randn(hidden, 256), + "timestep_emb.mlp.0.bias": torch.randn(hidden), + "timestep_emb.mlp.2.weight": torch.randn(hidden, hidden), + "timestep_emb.mlp.2.bias": torch.randn(hidden), + "time_embed.mlp.0.weight": torch.randn(hidden, 256), + "time_embed.mlp.0.bias": torch.randn(hidden), + "time_embed.mlp.2.weight": torch.randn(hidden, hidden), + "time_embed.mlp.2.bias": torch.randn(hidden), + "time_embed_2.mlp.0.weight": torch.randn(hidden, 256), + "time_embed_2.mlp.0.bias": torch.randn(hidden), + "time_embed_2.mlp.2.weight": torch.randn(hidden, hidden), + "time_embed_2.mlp.2.bias": torch.randn(hidden), + "patch_embed.model.0.weight": torch.randn(patch_hidden, latent, 3, 3), + "patch_embed.model.0.bias": torch.randn(patch_hidden), + "patch_embed.model.1.in_layers.0.weight": torch.randn(patch_hidden), + "patch_embed.model.1.in_layers.0.bias": torch.randn(patch_hidden), + "patch_embed.model.1.in_layers.2.weight": torch.randn(hidden, patch_hidden, 3, 3), + "patch_embed.model.1.in_layers.2.bias": torch.randn(hidden), + "patch_embed.model.1.emb_layers.1.weight": torch.randn(hidden * 2, hidden), + "patch_embed.model.1.emb_layers.1.bias": torch.randn(hidden * 2), + "patch_embed.model.1.out_layers.0.weight": torch.randn(hidden), + "patch_embed.model.1.out_layers.0.bias": torch.randn(hidden), + "patch_embed.model.1.out_layers.3.weight": torch.randn(hidden, hidden, 3, 3), + "patch_embed.model.1.out_layers.3.bias": torch.randn(hidden), + "patch_embed.model.1.skip_connection.weight": torch.randn(hidden, patch_hidden, 1, 1), + "patch_embed.model.1.skip_connection.bias": torch.randn(hidden), + "final_layer.model.0.in_layers.0.weight": torch.randn(hidden), + "final_layer.model.0.in_layers.0.bias": torch.randn(hidden), + "final_layer.model.0.in_layers.2.weight": torch.randn(patch_hidden, hidden, 3, 3), + "final_layer.model.0.in_layers.2.bias": torch.randn(patch_hidden), + "final_layer.model.0.emb_layers.1.weight": torch.randn(patch_hidden * 2, hidden), + "final_layer.model.0.emb_layers.1.bias": torch.randn(patch_hidden * 2), + "final_layer.model.0.out_layers.0.weight": torch.randn(patch_hidden), + "final_layer.model.0.out_layers.0.bias": torch.randn(patch_hidden), + "final_layer.model.0.out_layers.3.weight": torch.randn(patch_hidden, patch_hidden, 3, 3), + "final_layer.model.0.out_layers.3.bias": torch.randn(patch_hidden), + "final_layer.model.0.skip_connection.weight": torch.randn(patch_hidden, hidden, 1, 1), + "final_layer.model.0.skip_connection.bias": torch.randn(patch_hidden), + "final_layer.model.1.0.weight": torch.randn(patch_hidden), + "final_layer.model.1.0.bias": torch.randn(patch_hidden), + "final_layer.model.1.2.weight": torch.randn(latent, patch_hidden, 3, 3), + "final_layer.model.1.2.bias": torch.randn(latent), + } + for layer in range(config["num_layers"]): + prefix = f"model.layers.{layer}" + weights[f"{prefix}.input_layernorm.weight"] = torch.randn(hidden) + weights[f"{prefix}.post_attention_layernorm.weight"] = torch.randn(hidden) + weights[f"{prefix}.self_attn.qkv_proj.weight"] = torch.randn(qkv_out, hidden) + weights[f"{prefix}.self_attn.o_proj.weight"] = torch.randn(hidden, hidden) + weights[f"{prefix}.self_attn.query_layernorm.weight"] = torch.randn(head_dim) + weights[f"{prefix}.self_attn.key_layernorm.weight"] = torch.randn(head_dim) + if layer < config["moe_layer_num_skipped"]: + weights[f"{prefix}.mlp.gate_and_up_proj.weight"] = torch.randn(inter * 2, hidden) + weights[f"{prefix}.mlp.down_proj.weight"] = torch.randn(hidden, inter) + else: + weights[f"{prefix}.mlp.gate.wg.weight"] = torch.randn(config["num_experts"], hidden) + for expert in range(config["num_experts"]): + weights[f"{prefix}.mlp.experts.{expert}.gate_and_up_proj.weight"] = torch.randn(moe_inter * 2, hidden) + weights[f"{prefix}.mlp.experts.{expert}.down_proj.weight"] = torch.randn(hidden, moe_inter) + return weights + + +def _loaded_shape(weight): + tensor = getattr(weight, "weight", None) + if tensor is None: + tensor = getattr(weight, "pin_weight") + return tensor.shape + + +def test_infer_entrypoint_registers_hunyuan_image3_runner(): + infer_source = Path("lightx2v/infer.py").read_text() + + assert "HunyuanImage3Runner" in infer_source + assert '"hunyuan_image3"' in infer_source + + +def test_hunyuan_image3_config_merges_model_config(tmp_path): + model_path = tmp_path / "model" + model_path.mkdir() + (model_path / "config.json").write_text( + json.dumps( + { + "model_type": "hunyuan_image_3_moe", + "num_hidden_layers": 32, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "hidden_size": 4096, + "intermediate_size": 3072, + "diff_infer_steps": 8, + "diff_guidance_scale": 6.0, + "flow_shift": 7.0, + } + ) + ) + config_json = tmp_path / "runtime.json" + config_json.write_text( + json.dumps( + { + "infer_steps": 4, + "sample_guide_scale": 3.5, + "target_height": 1024, + "target_width": 1024, + } + ) + ) + + args = argparse.Namespace( + seed=123, + model_cls="hunyuan_image3", + task="t2i", + support_tasks=[], + model_path=str(model_path), + config_json=str(config_json), + ) + + config = set_config(args) + + assert config["num_layers"] == 32 + assert config["num_heads"] == 32 + assert config["num_key_value_heads"] == 8 + assert config["hidden_size"] == 4096 + assert config["infer_steps"] == 4 + assert config["sample_guide_scale"] == 3.5 + assert config["diff_infer_steps"] == 8 + assert config["diff_guidance_scale"] == 6.0 + assert config["flow_shift"] == 7.0 + + +def test_hunyuan_image3_cache_cli_args_override_config_json(tmp_path): + model_path = tmp_path / "model" + model_path.mkdir() + config_json = tmp_path / "runtime.json" + config_json.write_text( + json.dumps( + { + "enable_kv_cache": False, + "enable_text_kv_cache": False, + "use_taylor_cache": False, + "taylor_cache_interval": 5, + "taylor_cache_order": 2, + "taylor_cache_enable_first_enhance": False, + "taylor_cache_first_enhance_steps": 3, + "taylor_cache_enable_tailing_enhance": False, + "taylor_cache_tailing_enhance_steps": 1, + "taylor_cache_low_freqs_order": 2, + "taylor_cache_high_freqs_order": 2, + } + ) + ) + + args = argparse.Namespace( + seed=123, + model_cls="hunyuan_image3", + task="i2i", + support_tasks=[], + model_path=str(model_path), + config_json=str(config_json), + enable_kv_cache=True, + enable_text_kv_cache=True, + use_taylor_cache=True, + taylor_cache_interval=7, + taylor_cache_order=3, + taylor_cache_enable_first_enhance=True, + taylor_cache_first_enhance_steps=4, + taylor_cache_enable_tailing_enhance=True, + taylor_cache_tailing_enhance_steps=2, + taylor_cache_low_freqs_order=1, + taylor_cache_high_freqs_order=3, + ) + + config = set_config(args) + + assert config["enable_kv_cache"] is True + assert config["enable_text_kv_cache"] is True + assert config["use_taylor_cache"] is True + assert config["taylor_cache_interval"] == 7 + assert config["taylor_cache_order"] == 3 + assert config["taylor_cache_enable_first_enhance"] is True + assert config["taylor_cache_first_enhance_steps"] == 4 + assert config["taylor_cache_enable_tailing_enhance"] is True + assert config["taylor_cache_tailing_enhance_steps"] == 2 + assert config["taylor_cache_low_freqs_order"] == 1 + assert config["taylor_cache_high_freqs_order"] == 3 + + +def test_hunyuan_image3_moe_impl_cli_arg_overrides_config_json(tmp_path): + model_path = tmp_path / "model" + model_path.mkdir() + config_json = tmp_path / "runtime.json" + config_json.write_text(json.dumps({"moe_impl": "eager"})) + + args = argparse.Namespace( + seed=123, + model_cls="hunyuan_image3", + task="t2i", + support_tasks=[], + model_path=str(model_path), + config_json=str(config_json), + moe_impl="flashinfer", + ) + + config = set_config(args) + + assert config["moe_impl"] == "flashinfer" + + +def test_hunyuan_image3_ti2i_script_passes_cache_args_in_python_entrypoint(): + script_source = Path("scripts/hunyuan_image3/run_hunyuan_image3_ti2i.sh").read_text() + python_entry = script_source.split("python -m lightx2v.infer", maxsplit=1)[1] + + for cli_arg in [ + "--moe_impl", + "--enable_kv_cache", + "--enable_text_kv_cache", + "--use_taylor_cache", + "--taylor_cache_interval", + "--taylor_cache_order", + "--taylor_cache_enable_first_enhance", + "--taylor_cache_first_enhance_steps", + "--taylor_cache_enable_tailing_enhance", + "--taylor_cache_tailing_enhance_steps", + "--taylor_cache_low_freqs_order", + "--taylor_cache_high_freqs_order", + ]: + assert cli_arg in python_entry + + +def test_hunyuan_image3_t2i_script_passes_cache_args_in_python_entrypoint(): + script_source = Path("scripts/hunyuan_image3/run_hunyuan_image3_t2i.sh").read_text() + python_entry = script_source.split("python -m lightx2v.infer", maxsplit=1)[1] + + for cli_arg in [ + "--moe_impl", + "--enable_kv_cache", + "--enable_text_kv_cache", + "--use_taylor_cache", + "--taylor_cache_interval", + "--taylor_cache_order", + "--taylor_cache_enable_first_enhance", + "--taylor_cache_first_enhance_steps", + "--taylor_cache_enable_tailing_enhance", + "--taylor_cache_tailing_enhance_steps", + "--taylor_cache_low_freqs_order", + "--taylor_cache_high_freqs_order", + ]: + assert cli_arg in python_entry + + +def test_hunyuan_image3_i2i_script_passes_moe_impl_in_python_entrypoint(): + script_source = Path("scripts/hunyuan_image3/run_hunyuan_image3_i2i.sh").read_text() + python_entry = script_source.split("python -m lightx2v.infer", maxsplit=1)[1] + + assert "--moe_impl" in python_entry + + +def test_hunyuan_image3_runner_is_not_upstream_pipeline_wrapper(): + runner_source = Path("lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py").read_text() + + assert "HunyuanImage3ForCausalMM" not in runner_source + assert ".generate_image(" not in runner_source + assert "device_map" not in runner_source + + +def test_hunyuan_image3_weights_load_real_key_layout(): + from lightx2v.models.networks.hunyuan_image3.weights.post_weights import HunyuanImage3PostWeights + from lightx2v.models.networks.hunyuan_image3.weights.pre_weights import HunyuanImage3PreWeights + from lightx2v.models.networks.hunyuan_image3.weights.transformer_weights import HunyuanImage3TransformerWeights + + config = _tiny_hunyuan_config() + weight_dict = _tiny_weight_dict(config) + + pre = HunyuanImage3PreWeights(config) + transformer = HunyuanImage3TransformerWeights(config) + post = HunyuanImage3PostWeights(config) + + pre.load(weight_dict) + transformer.load(weight_dict) + post.load(weight_dict) + + assert _loaded_shape(pre.token_embedding) == (32, 16) + assert pre.patch_embed.input_conv.weight.shape == (8, 4, 3, 3) + assert len(transformer.blocks) == 2 + assert _loaded_shape(transformer.blocks[0].compute_phases[1].gate_and_up_proj) == (16, 24) + assert _loaded_shape(transformer.blocks[1].compute_phases[1].experts[2].down_proj) == (12, 16) + assert _loaded_shape(post.final_norm) == (16,) + assert post.final_layer.output_conv.weight.shape == (4, 8, 3, 3) + + +def test_hunyuan_image3_pipeline_parallel_maps_weights_across_devices(): + from lightx2v.models.networks.hunyuan_image3.model import resolve_pipeline_device_for_key + + config = _tiny_hunyuan_config() + config["num_layers"] = 8 + devices = ["cuda:0", "cuda:1", "cuda:2", "cuda:3"] + + assert resolve_pipeline_device_for_key("model.wte.weight", config, devices) == "cuda:0" + assert resolve_pipeline_device_for_key("patch_embed.model.0.weight", config, devices) == "cuda:0" + assert resolve_pipeline_device_for_key("time_embed.mlp.0.weight", config, devices) == "cuda:0" + + assert resolve_pipeline_device_for_key("model.layers.0.self_attn.qkv_proj.weight", config, devices) == "cuda:0" + assert resolve_pipeline_device_for_key("model.layers.1.mlp.experts.0.down_proj.weight", config, devices) == "cuda:0" + assert resolve_pipeline_device_for_key("model.layers.2.self_attn.qkv_proj.weight", config, devices) == "cuda:1" + assert resolve_pipeline_device_for_key("model.layers.5.self_attn.qkv_proj.weight", config, devices) == "cuda:2" + assert resolve_pipeline_device_for_key("model.layers.7.mlp.gate.wg.weight", config, devices) == "cuda:3" + + assert resolve_pipeline_device_for_key("model.ln_f.weight", config, devices) == "cuda:3" + assert resolve_pipeline_device_for_key("lm_head.weight", config, devices) == "cuda:3" + assert resolve_pipeline_device_for_key("final_layer.model.0.in_layers.0.weight", config, devices) == "cuda:3" + + +def test_hunyuan_image3_loader_streams_core_tensors_with_pipeline_mapping(tmp_path): + from lightx2v.models.networks.hunyuan_image3.model import HunyuanImage3Model + + ckpt = tmp_path / "tiny.safetensors" + save_file( + { + "model.wte.weight": torch.randn(4, 4), + "model.layers.1.self_attn.qkv_proj.weight": torch.randn(4, 4), + "vae.decoder.conv_in.weight": torch.randn(4, 4, 1, 1), + }, + ckpt, + ) + model = HunyuanImage3Model.__new__(HunyuanImage3Model) + model.config = {**_tiny_hunyuan_config(), "num_layers": 2} + model.pipeline_devices = ["cpu", "cpu"] + model.pipeline_parallel = True + model.remove_keys = [] + model.preserved_keys = ["model.", "lm_head.", "patch_embed.", "final_layer.", "time_embed.", "time_embed_2.", "timestep_emb."] + + loaded = model._load_safetensor_to_dict(str(ckpt), unified_dtype=True, sensitive_layer={}) + + assert sorted(loaded) == ["model.layers.1.self_attn.qkv_proj.weight", "model.wte.weight"] + assert loaded["model.wte.weight"].device.type == "cpu" + assert loaded["model.layers.1.self_attn.qkv_proj.weight"].shape == (4, 4) + + +def test_hunyuan_image3_runner_pipeline_saves_generated_image(tmp_path): + from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner + + save_path = tmp_path / "out.png" + runner = HunyuanImage3Runner.__new__(HunyuanImage3Runner) + runner.config = {"task": "t2i"} + runner._gc_frozen = True + runner.generate_t2i = lambda input_info: [Image.new("RGB", (2, 3), color=(10, 20, 30))] + + result = runner.run_pipeline( + SimpleNamespace( + prompt="hello", + save_result_path=str(save_path), + return_result_tensor=False, + ) + ) + + assert result == {"image": None} + assert save_path.exists() + assert Image.open(save_path).size == (2, 3) + + +def test_hunyuan_image3_runner_pipeline_saves_ti2i_image(tmp_path): + from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner + + save_path = tmp_path / "out.png" + runner = HunyuanImage3Runner.__new__(HunyuanImage3Runner) + runner.config = {"task": "i2i"} + runner._gc_frozen = True + runner.generate_i2i = lambda input_info: [Image.new("RGB", (4, 5), color=(10, 20, 30))] + + result = runner.run_pipeline( + SimpleNamespace( + prompt="edit", + image_path="ref.png", + save_result_path=str(save_path), + return_result_tensor=False, + ) + ) + + assert result == {"image": None} + assert save_path.exists() + assert Image.open(save_path).size == (4, 5) + + +def test_hunyuan_image3_resolves_text_generation_plan_for_think_recaption(): + from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner + + runner = HunyuanImage3Runner.__new__(HunyuanImage3Runner) + runner.hunyuan_tokenizer = SimpleNamespace( + end_of_think_token_id=10, + end_of_recaption_token_id=12, + recaption_token="", + convert_tokens_to_ids=lambda token: {"": 11}[token], + ) + + plan = runner._resolve_text_generation_plan("think_recaption") + + assert plan.first_bot_task == "think" + assert plan.stage_transitions == [(10, [11])] + assert plan.final_stop_tokens == [12] + + +def test_hunyuan_image3_prepare_image_inputs_passes_cot_text(): + from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner + + class FakeImageProcessor: + vae_reso_group = SimpleNamespace(base_size=1024) + + def build_gen_image_info(self, image_size, add_guidance_token=False, add_timestep_r_token=False): + return SimpleNamespace(image_size=image_size) + + def prepare_full_attn_slices(self, tokenizer_output, batch_idx): + return [] + + class FakeTokenizer: + def __init__(self): + self.kwargs = None + + def apply_chat_template(self, **kwargs): + self.kwargs = kwargs + output = SimpleNamespace( + tokens=torch.tensor([[1, 2, 3]]), + all_image_slices=[[]], + gen_image_mask=torch.tensor([[False, True, True]]), + gen_timestep_scatter_index=None, + guidance_scatter_index=None, + gen_timestep_r_scatter_index=None, + ) + return {"output": output, "sections": [[{"type": "text"}]]} + + runner = HunyuanImage3Runner.__new__(HunyuanImage3Runner) + runner.config = {"max_position_embeddings": 16} + runner.hunyuan_generation_config = SimpleNamespace(max_length=16, sequence_template="instruct", drop_think=True) + runner.hunyuan_image_processor = FakeImageProcessor() + runner.hunyuan_tokenizer = FakeTokenizer() + runner.hunyuan_config = SimpleNamespace(max_position_embeddings=16, rope_type="default") + runner.hunyuan_cached_rope = lambda *args, **kwargs: None + runner._pipeline_latent_device = lambda: torch.device("cpu") + + runner._prepare_text_to_image_inputs("prompt", (8, 8), 123, cot_text="xy") + + assert runner.hunyuan_tokenizer.kwargs["batch_cot_text"] == ["xy"] + assert runner.hunyuan_tokenizer.kwargs["bot_task"] == "image" + + +def test_hunyuan_image3_prepare_image_inputs_passes_cond_images(): + from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner + + class FakeImageProcessor: + vae_reso_group = SimpleNamespace(base_size=1024) + + def build_gen_image_info(self, image_size, add_guidance_token=False, add_timestep_r_token=False): + return SimpleNamespace(image_size=image_size) + + def prepare_full_attn_slices(self, tokenizer_output, batch_idx): + return [] + + class FakeTokenizer: + def __init__(self): + self.kwargs = None + + def apply_chat_template(self, **kwargs): + self.kwargs = kwargs + output = SimpleNamespace( + tokens=torch.tensor([[1, 2, 3], [1, 2, 3]]), + all_image_slices=[[], []], + gen_image_mask=torch.tensor([[False, True, True], [False, True, True]]), + gen_timestep_scatter_index=torch.tensor([[0], [0]]), + guidance_scatter_index=None, + gen_timestep_r_scatter_index=None, + vae_image_mask=torch.tensor([[True, False, False], [True, False, False]]), + vit_image_mask=torch.tensor([[False, True, False], [False, True, False]]), + cond_timestep_scatter_index=torch.tensor([[2], [2]]), + ) + return {"output": output, "sections": [[{"type": "cond_vae_image"}, {"type": "gen_image"}]]} + + runner = HunyuanImage3Runner.__new__(HunyuanImage3Runner) + runner.config = {"max_position_embeddings": 16, "enable_cfg": True} + runner.hunyuan_generation_config = SimpleNamespace(max_length=16, sequence_template="instruct", drop_think=True) + runner.hunyuan_image_processor = FakeImageProcessor() + runner.hunyuan_tokenizer = FakeTokenizer() + runner.hunyuan_config = SimpleNamespace(max_position_embeddings=16, rope_type="default") + runner.hunyuan_cached_rope = lambda *args, **kwargs: None + runner._pipeline_latent_device = lambda: torch.device("cpu") + + prepared = runner._prepare_text_to_image_inputs( + "prompt", + (8, 8), + 123, + cot_text="x", + batch_cond_images=[["cond"]], + cond_inputs={ + "cond_vae_images": [torch.zeros(1, 4, 2, 2), torch.zeros(1, 4, 2, 2)], + "cond_timesteps": [torch.zeros(1), torch.zeros(1)], + "cond_vit_embeds": [torch.zeros(1, 1, 16), torch.zeros(1, 1, 16)], + }, + ) + + assert runner.hunyuan_tokenizer.kwargs["batch_cond_images"] == [["cond"]] + assert prepared["cond_vae_images"][0].shape == (1, 4, 2, 2) + assert torch.equal(prepared["cond_vae_image_mask"], torch.tensor([[True, False, False], [True, False, False]])) + assert torch.equal(prepared["cond_timestep_index"], torch.tensor([[2], [2]])) + assert prepared["cond_vit_embeds"][0].shape == (1, 1, 16) + + +def test_hunyuan_image3_generate_t2i_uses_cot_for_think_recaption(): + from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner + + runner = HunyuanImage3Runner.__new__(HunyuanImage3Runner) + runner.config = {"task": "t2i", "bot_task": "think_recaption"} + runner._ensure_pipeline_modules = lambda: None + runner._resolve_image_size = lambda input_info: (8, 8) + runner._generate_cot_text = lambda prompt, image_size: "xy" + seen = {} + + def fake_prepare(prompt, image_size, seed, cot_text=None): + seen["cot_text"] = cot_text + return {"batch_size": 1, "generator": "generator"} + + runner._prepare_text_to_image_inputs = fake_prepare + runner._denoise_latents = lambda prepared_inputs, image_size: "latents" + runner._decode_latents = lambda latents, generator: ["image"] + + images = runner.generate_t2i(SimpleNamespace(prompt="prompt", seed=123)) + + assert images == ["image"] + assert seen["cot_text"] == "xy" + + +def test_hunyuan_image3_text_generation_applies_think_to_recaption_transition(): + from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner + + class FakeModel: + def __init__(self): + self.next_tokens = [10, 12] + self.calls = 0 + + def infer(self, inputs): + vocab_size = 16 + logits = torch.zeros(1, inputs["input_ids"].shape[1], vocab_size) + logits[:, -1, self.next_tokens[self.calls]] = 100.0 + self.calls += 1 + return {"logits": logits} + + runner = HunyuanImage3Runner.__new__(HunyuanImage3Runner) + runner.config = {"text_do_sample": False, "max_new_tokens": 5} + runner.hunyuan_generation_config = SimpleNamespace(max_new_tokens=5, do_sample=False) + runner.hunyuan_tokenizer = SimpleNamespace( + think_token="", + recaption_token="", + decode=lambda tokens: "".join({10: "", 11: "", 12: ""}[int(token)] for token in tokens), + ) + runner.model = FakeModel() + runner._build_text_model_inputs = lambda input_ids, tokenizer_output: {"input_ids": input_ids} + plan = SimpleNamespace(first_bot_task="think", stage_transitions=[(10, [11])], final_stop_tokens=[12]) + + generated_tokens = runner._generate_text_tokens(torch.tensor([[1, 2]]), SimpleNamespace(), plan) + cot_text = runner._decode_cot_text(generated_tokens, plan) + + assert generated_tokens == [10, 11, 12] + assert cot_text == "" + + +def test_hunyuan_image3_text_generation_moves_sampled_token_to_input_device(): + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + return + + from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner + + input_device = torch.device("cuda:0") + logits_device = torch.device("cuda:1") + + class FakeModel: + def infer(self, inputs): + logits = torch.zeros(1, inputs["input_ids"].shape[1], 8, device=logits_device) + logits[:, -1, 3] = 100.0 + return {"logits": logits} + + runner = HunyuanImage3Runner.__new__(HunyuanImage3Runner) + runner.config = {"text_do_sample": False, "max_new_tokens": 1} + runner.hunyuan_generation_config = SimpleNamespace(max_new_tokens=1, do_sample=False) + runner.model = FakeModel() + runner._build_text_model_inputs = lambda input_ids, tokenizer_output: {"input_ids": input_ids} + plan = SimpleNamespace(first_bot_task="recaption", stage_transitions=[], final_stop_tokens=[3]) + + generated_tokens = runner._generate_text_tokens(torch.tensor([[1, 2]], device=input_device), SimpleNamespace(), plan) + + assert generated_tokens == [3] + + +def test_hunyuan_image3_generate_cot_text_streams_process_to_stdout(capsys): + from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner + + class FakeModel: + def __init__(self): + self.next_tokens = [10, 12] + self.calls = 0 + + def infer(self, inputs): + vocab_size = 16 + logits = torch.zeros(1, inputs["input_ids"].shape[1], vocab_size) + logits[:, -1, self.next_tokens[self.calls]] = 100.0 + self.calls += 1 + return {"logits": logits} + + runner = HunyuanImage3Runner.__new__(HunyuanImage3Runner) + runner.config = {"bot_task": "think_recaption", "text_do_sample": False, "max_new_tokens": 5, "text_stream_output": True} + runner.hunyuan_generation_config = SimpleNamespace(max_new_tokens=5, do_sample=False) + runner.hunyuan_tokenizer = SimpleNamespace( + think_token="", + recaption_token="", + end_of_think_token_id=10, + end_of_recaption_token_id=12, + convert_tokens_to_ids=lambda token: {"": 11}[token], + decode=lambda tokens: "".join({10: "", 11: "", 12: ""}[int(token)] for token in tokens), + ) + runner.model = FakeModel() + runner._resolve_system_prompt = lambda bot_task: None + runner._prepare_text_generation_inputs = lambda prompt, bot_task, system_prompt: { + "input_ids": torch.tensor([[1, 2]]), + "tokenizer_output": SimpleNamespace(), + } + runner._build_text_model_inputs = lambda input_ids, tokenizer_output: {"input_ids": input_ids} + + cot_text = runner._generate_cot_text("prompt", (8, 8)) + + captured = capsys.readouterr() + assert cot_text == "" + assert "" in captured.out + + +def test_hunyuan_image3_pre_infer_instantiates_timestep_guidance_and_meanflow_tokens(monkeypatch): + from lightx2v.models.networks.hunyuan_image3.infer import pre_infer + + def fake_timestep_embedder(weights, values): + fill = {"timestep": 1.0, "guidance": 2.0, "timestep_r": 3.0}[weights] + return torch.full((values.numel(), 4), fill) + + monkeypatch.setattr(pre_infer, "apply_timestep_embedder", fake_timestep_embedder) + infer = pre_infer.HunyuanImage3PreInfer({"hidden_size": 4, "cfg_distilled": True, "use_meanflow": True}) + weights = SimpleNamespace(timestep_emb="timestep", guidance_emb="guidance", timestep_r_emb="timestep_r") + + out = infer.infer( + weights, + { + "inputs_embeds": torch.zeros(1, 5, 4), + "timesteps": torch.tensor([10.0]), + "timesteps_index": torch.tensor([[1]]), + "guidance": torch.tensor([5.0]), + "guidance_index": torch.tensor([[3]]), + "timesteps_r": torch.tensor([2.0]), + "timesteps_r_index": torch.tensor([[4]]), + }, + ) + + assert torch.equal(out.hidden_states[0, 0], torch.zeros(4)) + assert torch.equal(out.hidden_states[0, 1], torch.ones(4)) + assert torch.equal(out.hidden_states[0, 3], torch.full((4,), 2.0)) + assert torch.equal(out.hidden_states[0, 4], torch.full((4,), 3.0)) + + +def test_hunyuan_image3_pre_infer_instantiates_cond_vae_and_vit_tokens(monkeypatch): + from lightx2v.models.networks.hunyuan_image3.infer import pre_infer + + def fake_timestep_embedder(weights, values): + return torch.full((values.numel(), 4), 3.0) + + monkeypatch.setattr(pre_infer, "apply_timestep_embedder", fake_timestep_embedder) + infer = pre_infer.HunyuanImage3PreInfer({"hidden_size": 4, "cfg_distilled": False, "use_meanflow": False}) + infer._patch_embed = lambda weights, images, timesteps: (torch.full((1, 2, 4), 5.0), 1, 2) + weights = SimpleNamespace(timestep_emb="timestep") + + out = infer.infer( + weights, + { + "inputs_embeds": torch.zeros(1, 6, 4), + "cond_vae_images": torch.zeros(1, 4, 1, 2), + "cond_vae_image_mask": torch.tensor([[False, True, True, False, False, False]]), + "cond_timesteps": torch.tensor([0.0]), + "cond_timestep_index": torch.tensor([[3]]), + "cond_vit_embeds": torch.full((1, 2, 4), 7.0), + "cond_vit_image_mask": torch.tensor([[False, False, False, False, True, True]]), + }, + ) + + assert torch.equal(out.hidden_states[0, 1], torch.full((4,), 5.0)) + assert torch.equal(out.hidden_states[0, 2], torch.full((4,), 5.0)) + assert torch.equal(out.hidden_states[0, 3], torch.full((4,), 3.0)) + assert torch.equal(out.hidden_states[0, 4], torch.full((4,), 7.0)) + assert torch.equal(out.hidden_states[0, 5], torch.full((4,), 7.0)) + + +def test_hunyuan_image3_timestep_embedder_matches_weight_dtype(): + from lightx2v.models.networks.hunyuan_image3.infer.utils import apply_timestep_embedder + + class FakeLinear: + def __init__(self, in_features, out_features, dtype): + self.weight = torch.randn(in_features, out_features, dtype=dtype) + self.bias = torch.randn(out_features, dtype=dtype) + + def apply(self, input_tensor): + output = torch.empty(input_tensor.shape[0], self.weight.shape[1], dtype=input_tensor.dtype) + return torch.addmm(self.bias, input_tensor, self.weight, out=output) + + weights = SimpleNamespace( + linear_1=FakeLinear(256, 8, torch.bfloat16), + linear_2=FakeLinear(8, 4, torch.bfloat16), + ) + + output = apply_timestep_embedder(weights, torch.tensor([10.0])) + + assert output.dtype == torch.bfloat16 + assert output.shape == (1, 4) + + +def test_hunyuan_image3_apply_linear_matches_weight_dtype(): + from lightx2v.models.networks.hunyuan_image3.infer.utils import apply_linear + + class FakeLinear: + def __init__(self, in_features, out_features, dtype): + self.weight = torch.randn(in_features, out_features, dtype=dtype) + self.bias = None + + def apply(self, input_tensor): + output = torch.empty(input_tensor.shape[0], self.weight.shape[1], dtype=input_tensor.dtype) + return torch.mm(input_tensor, self.weight, out=output) + + linear = FakeLinear(4, 6, torch.bfloat16) + + output = apply_linear(linear, torch.randn(2, 4, dtype=torch.float32)) + + assert output.dtype == torch.bfloat16 + assert output.shape == (2, 6) + + +def test_hunyuan_image3_apply_mlp_matches_weight_dtype(): + from lightx2v.models.networks.hunyuan_image3.infer.utils import apply_mlp + + class FakeLinear: + def __init__(self, in_features, out_features, dtype): + self.weight = torch.randn(in_features, out_features, dtype=dtype) + self.bias = None + + def apply(self, input_tensor): + output = torch.empty(input_tensor.shape[0], self.weight.shape[1], dtype=input_tensor.dtype) + return torch.mm(input_tensor, self.weight, out=output) + + gate_and_up = FakeLinear(4, 8, torch.bfloat16) + down = FakeLinear(4, 4, torch.bfloat16) + + output = apply_mlp(gate_and_up, down, torch.randn(2, 3, 4, dtype=torch.float32)) + + assert output.dtype == torch.bfloat16 + assert output.shape == (2, 3, 4) + + +def test_hunyuan_image3_rotary_preserves_query_key_dtype(): + from lightx2v.models.networks.hunyuan_image3.infer.utils import apply_rotary_pos_emb + + q = torch.randn(1, 2, 3, 4, dtype=torch.bfloat16) + k = torch.randn(1, 1, 3, 4, dtype=torch.bfloat16) + cos = torch.randn(1, 3, 4, dtype=torch.float32) + sin = torch.randn(1, 3, 4, dtype=torch.float32) + + q_out, k_out = apply_rotary_pos_emb(q, k, cos, sin) + + assert q_out.dtype == torch.bfloat16 + assert k_out.dtype == torch.bfloat16 + + +def test_hunyuan_image3_conv2d_matches_weight_dtype(): + from lightx2v.models.networks.hunyuan_image3.weights.common import HunyuanImage3Conv2dWeight + + conv = HunyuanImage3Conv2dWeight("weight", "bias", padding=1) + conv.load( + { + "weight": torch.randn(2, 2, 3, 3, dtype=torch.bfloat16), + "bias": torch.randn(2, dtype=torch.bfloat16), + } + ) + + output = conv.apply(torch.randn(1, 2, 4, 4, dtype=torch.float32)) + + assert output.dtype == torch.bfloat16 + assert output.shape == (1, 2, 4, 4) + + +def test_hunyuan_image3_moe_scatter_matches_destination_dtype(): + from lightx2v.models.networks.hunyuan_image3.infer.transformer_infer import HunyuanImage3TransformerInfer + + class FakeLinear: + def __init__(self, in_features, out_features, dtype): + self.weight = torch.randn(in_features, out_features, dtype=dtype) + self.bias = None + + def apply(self, input_tensor): + output = torch.empty(input_tensor.shape[0], self.weight.shape[1], dtype=input_tensor.dtype) + return torch.mm(input_tensor, self.weight, out=output) + + class FakeExpert: + def __init__(self): + self.gate_and_up_proj = FakeLinear(4, 8, torch.bfloat16) + self.down_proj = FakeLinear(4, 4, torch.bfloat16) + + infer = HunyuanImage3TransformerInfer( + { + "num_layers": 1, + "hidden_size": 4, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "attention_head_dim": 4, + "hidden_act": "silu", + } + ) + phase = SimpleNamespace( + is_moe=True, + moe=SimpleNamespace( + gate=FakeLinear(4, 2, torch.float32), + moe_topk=1, + experts=[FakeExpert(), FakeExpert()], + shared_mlp=None, + ), + ) + + output = infer.infer_mlp(phase, torch.randn(1, 3, 4, dtype=torch.float32)) + + assert output.dtype == torch.float32 + assert output.shape == (1, 3, 4) + + +class _FakeHunyuanLinear: + def __init__(self, weight, bias=None): + self.weight = weight + self.bias = bias + self.has_lora_branch = False + + def _get_actual_weight(self): + return self.weight + + def apply(self, input_tensor): + input_tensor = input_tensor.to(device=self.weight.device, dtype=self.weight.dtype) + output = torch.mm(input_tensor, self.weight) + if self.bias is not None: + output = output + self.bias + return output + + +def _fake_hunyuan_moe_phase(moe_impl="eager"): + from lightx2v.models.networks.hunyuan_image3.weights.common import HunyuanImage3MoEWeights + + torch.manual_seed(7) + hidden = 4 + intermediate = 6 + num_experts = 3 + moe = HunyuanImage3MoEWeights.__new__(HunyuanImage3MoEWeights) + moe.num_experts = num_experts + moe.moe_topk = 2 + moe.moe_impl = moe_impl + moe.shared_mlp = None + moe.gate = _FakeHunyuanLinear(torch.randn(hidden, num_experts, dtype=torch.float32) * 0.2) + moe.experts = [] + for _ in range(num_experts): + moe.experts.append( + SimpleNamespace( + gate_and_up_proj=_FakeHunyuanLinear(torch.randn(hidden, intermediate * 2, dtype=torch.bfloat16) * 0.2), + down_proj=_FakeHunyuanLinear(torch.randn(intermediate, hidden, dtype=torch.bfloat16) * 0.2), + ) + ) + moe.moe_weight = None + moe.moe_weight_2 = None + moe._flashinfer_weights_initialized = False + moe._flashinfer_weight_device = None + moe._flashinfer_weight_dtype = None + return SimpleNamespace(is_moe=True, moe=moe) + + +def _fake_cutlass_fused_moe(input_tensor, token_selected_experts, token_final_scales, fc1_expert_weights, fc2_expert_weights, output_dtype, quant_scales, output=None, **kwargs): + result = torch.zeros(input_tensor.shape[0], input_tensor.shape[1], device=input_tensor.device, dtype=output_dtype) + for token_idx in range(input_tensor.shape[0]): + token = input_tensor[token_idx : token_idx + 1].to(output_dtype) + for topk_idx in range(token_selected_experts.shape[1]): + expert_idx = int(token_selected_experts[token_idx, topk_idx]) + projected = torch.mm(token, fc1_expert_weights[expert_idx].t()) + x1, x2 = projected.chunk(2, dim=-1) + activated = x1 * torch.nn.functional.silu(x2) + expert_out = torch.mm(activated, fc2_expert_weights[expert_idx].t()) + result[token_idx] += expert_out.squeeze(0) * token_final_scales[token_idx, topk_idx].to(output_dtype) + if output is not None: + output.copy_(result) + return output + return result + + +def test_hunyuan_image3_flashinfer_moe_packs_lightx2v_weights_like_torch_linear(): + phase = _fake_hunyuan_moe_phase("flashinfer") + moe = phase.moe + expert_gate_up = moe.experts[0].gate_and_up_proj.weight.clone() + expert_down = moe.experts[0].down_proj.weight.clone() + + moe_weight, moe_weight_2 = moe.ensure_flashinfer_weights(torch.device("cpu"), torch.bfloat16) + + assert moe_weight.shape == (3, 12, 4) + assert moe_weight_2.shape == (3, 4, 6) + torch.testing.assert_close(moe_weight[0], expert_gate_up.t().contiguous()) + torch.testing.assert_close(moe_weight_2[0], expert_down.t().contiguous()) + assert moe.experts[0].gate_and_up_proj.weight.numel() == 0 + assert moe.experts[0].down_proj.weight.numel() == 0 + + +def test_hunyuan_image3_flashinfer_moe_matches_eager_with_mocked_kernel(monkeypatch): + from lightx2v.models.networks.hunyuan_image3.infer import transformer_infer + from lightx2v.models.networks.hunyuan_image3.infer.transformer_infer import HunyuanImage3TransformerInfer + + monkeypatch.setattr(transformer_infer, "flashinfer_cutlass_fused_moe", _fake_cutlass_fused_moe) + infer = HunyuanImage3TransformerInfer( + { + "num_layers": 1, + "hidden_size": 4, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "attention_head_dim": 4, + "hidden_act": "silu", + "moe_impl": "flashinfer", + } + ) + torch.manual_seed(11) + hidden_states = torch.randn(2, 3, 4, dtype=torch.bfloat16) + + eager_output = infer.infer_mlp(_fake_hunyuan_moe_phase("eager"), hidden_states) + flashinfer_output = infer.infer_mlp(_fake_hunyuan_moe_phase("flashinfer"), hidden_states) + + torch.testing.assert_close(flashinfer_output.float(), eager_output.float(), atol=2e-2, rtol=2e-2) + + +def test_hunyuan_image3_flashinfer_moe_requires_flashinfer(monkeypatch): + from lightx2v.models.networks.hunyuan_image3.infer import transformer_infer + from lightx2v.models.networks.hunyuan_image3.infer.transformer_infer import HunyuanImage3TransformerInfer + + monkeypatch.setattr(transformer_infer, "flashinfer_cutlass_fused_moe", None) + infer = HunyuanImage3TransformerInfer( + { + "num_layers": 1, + "hidden_size": 4, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "attention_head_dim": 4, + "hidden_act": "silu", + "moe_impl": "flashinfer", + } + ) + + with pytest.raises(ImportError, match="flashinfer"): + infer.infer_mlp(_fake_hunyuan_moe_phase("flashinfer"), torch.randn(1, 2, 4)) + + +def test_hunyuan_image3_static_kv_cache_updates_requested_positions(): + from lightx2v.models.networks.hunyuan_image3.infer.kv_cache import HunyuanImage3StaticKVCache + + cache = HunyuanImage3StaticKVCache(num_layers=1, max_cache_len=5) + first_keys = torch.arange(20, dtype=torch.float32).reshape(1, 1, 5, 4) + first_values = first_keys + 100 + + key_cache, value_cache = cache.update(first_keys, first_values, layer_idx=0, cache_position=torch.arange(5)) + + assert key_cache.shape == (1, 1, 5, 4) + assert torch.equal(key_cache, first_keys) + assert torch.equal(value_cache, first_values) + + second_keys = torch.full((1, 1, 2, 4), -1.0) + second_values = torch.full((1, 1, 2, 4), -2.0) + key_cache, value_cache = cache.update(second_keys, second_values, layer_idx=0, cache_position=torch.tensor([[1, 4]])) + + assert torch.equal(key_cache[:, :, 0], first_keys[:, :, 0]) + assert torch.equal(key_cache[:, :, 1], second_keys[:, :, 0]) + assert torch.equal(key_cache[:, :, 4], second_keys[:, :, 1]) + assert torch.equal(value_cache[:, :, 1], second_values[:, :, 0]) + assert torch.equal(value_cache[:, :, 4], second_values[:, :, 1]) + + +def test_hunyuan_image3_dynamic_kv_cache_returns_filled_prefix(): + from lightx2v.models.networks.hunyuan_image3.infer.kv_cache import HunyuanImage3StaticKVCache + + cache = HunyuanImage3StaticKVCache(num_layers=1, max_cache_len=6, dynamic=True) + first_keys = torch.arange(12, dtype=torch.float32).reshape(1, 1, 3, 4) + first_values = first_keys + 100 + + key_cache, value_cache = cache.update(first_keys, first_values, layer_idx=0, cache_position=torch.arange(3)) + + assert key_cache.shape == (1, 1, 3, 4) + assert value_cache.shape == (1, 1, 3, 4) + + next_keys = torch.full((1, 1, 2, 4), 5.0) + next_values = torch.full((1, 1, 2, 4), 6.0) + key_cache, value_cache = cache.update(next_keys, next_values, layer_idx=0, cache_position=torch.tensor([[3, 4]])) + + assert key_cache.shape == (1, 1, 5, 4) + assert value_cache.shape == (1, 1, 5, 4) + assert torch.equal(key_cache[:, :, 3], next_keys[:, :, 0]) + assert torch.equal(key_cache[:, :, 4], next_keys[:, :, 1]) + + +def test_hunyuan_image3_runner_builds_cache_decode_positions_and_masks(): + from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner + + runner = HunyuanImage3Runner.__new__(HunyuanImage3Runner) + runner.config = {"cfg_distilled": True, "use_meanflow": False} + prepared_inputs = { + "image_mask": torch.tensor([[False, True, False, True], [False, True, False, True]]), + "timesteps_index": torch.tensor([[0], [0]]), + "guidance_index": torch.tensor([[2], [2]]), + "timesteps_r_index": None, + "attention_mask": torch.arange(2 * 1 * 4 * 4).reshape(2, 1, 4, 4), + } + + position_ids = runner._build_denoise_cache_position_ids(prepared_inputs) + attention_mask = runner._slice_denoise_cache_attention_mask(prepared_inputs["attention_mask"], position_ids) + local_inputs = runner._build_denoise_cache_local_indices(prepared_inputs, position_ids) + + assert torch.equal(position_ids, torch.tensor([[0, 2, 1, 3], [0, 2, 1, 3]])) + assert attention_mask.shape == (2, 1, 4, 4) + assert torch.equal(attention_mask[0], prepared_inputs["attention_mask"][0].index_select(1, position_ids[0])) + assert torch.equal(local_inputs["image_mask"], torch.tensor([[False, False, True, True], [False, False, True, True]])) + assert torch.equal(local_inputs["timesteps_index"], torch.tensor([[0], [0]])) + assert torch.equal(local_inputs["guidance_index"], torch.tensor([[1], [1]])) + assert local_inputs["timesteps_r_index"] is None + + +def test_hunyuan_image3_transformer_attention_uses_kv_cache(): + from lightx2v.models.networks.hunyuan_image3.infer.kv_cache import HunyuanImage3StaticKVCache + from lightx2v.models.networks.hunyuan_image3.infer.transformer_infer import HunyuanImage3TransformerInfer + + class FakeLinear: + def __init__(self, in_features, out_features): + self.weight = torch.randn(in_features, out_features) + self.bias = None + + def apply(self, input_tensor): + return torch.mm(input_tensor, self.weight) + + infer = HunyuanImage3TransformerInfer( + { + "num_layers": 1, + "hidden_size": 4, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "attention_head_dim": 4, + "hidden_act": "silu", + } + ) + phase = SimpleNamespace(qkv_proj=FakeLinear(4, 12), o_proj=FakeLinear(4, 4), query_layernorm=None, key_layernorm=None) + cache = HunyuanImage3StaticKVCache(num_layers=1, max_cache_len=5) + + first = infer.infer_attention( + 0, + phase, + torch.randn(1, 5, 4), + attention_mask=None, + position_ids=torch.arange(5).reshape(1, 5), + custom_pos_emb=None, + past_key_values=cache, + ) + second = infer.infer_attention( + 0, + phase, + torch.randn(1, 2, 4), + attention_mask=None, + position_ids=torch.tensor([[0, 3]]), + custom_pos_emb=None, + past_key_values=cache, + ) + + assert first.shape == (1, 5, 4) + assert second.shape == (1, 2, 4) + assert cache.layers[0].key.shape == (1, 1, 5, 4) + + +def test_hunyuan_image3_text_generation_uses_dynamic_kv_cache_for_uncached_tokens(): + from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner + + class FakeModel: + def __init__(self): + self.next_tokens = [10, 12] + self.calls = [] + + def infer(self, inputs): + self.calls.append( + { + "input_ids": inputs["input_ids"].clone(), + "position_ids": inputs["position_ids"].clone(), + "use_cache": inputs.get("use_cache"), + "past_key_values": inputs.get("past_key_values"), + } + ) + vocab_size = 16 + logits = torch.zeros(1, inputs["input_ids"].shape[1], vocab_size) + logits[:, -1, self.next_tokens[len(self.calls) - 1]] = 100.0 + return {"logits": logits} + + runner = HunyuanImage3Runner.__new__(HunyuanImage3Runner) + runner.config = {"text_do_sample": False, "max_new_tokens": 5, "enable_kv_cache": True} + runner.hunyuan_generation_config = SimpleNamespace(max_new_tokens=5, do_sample=False) + runner.hunyuan_config = SimpleNamespace(max_position_embeddings=16) + runner.hunyuan_cached_rope = lambda *args, **kwargs: None + runner.hunyuan_image_processor = SimpleNamespace(prepare_full_attn_slices=lambda tokenizer_output, batch_idx: []) + runner.model = FakeModel() + plan = SimpleNamespace(first_bot_task="think", stage_transitions=[(10, [11])], final_stop_tokens=[12]) + + generated_tokens = runner._generate_text_tokens(torch.tensor([[1, 2]]), SimpleNamespace(), plan) + + assert generated_tokens == [10, 11, 12] + assert len(runner.model.calls) == 2 + assert torch.equal(runner.model.calls[0]["input_ids"], torch.tensor([[1, 2]])) + assert torch.equal(runner.model.calls[0]["position_ids"], torch.tensor([[0, 1]])) + assert torch.equal(runner.model.calls[1]["input_ids"], torch.tensor([[10, 11]])) + assert torch.equal(runner.model.calls[1]["position_ids"], torch.tensor([[2, 3]])) + assert runner.model.calls[0]["past_key_values"] is runner.model.calls[1]["past_key_values"] + assert runner.model.calls[0]["use_cache"] is True + + +def test_hunyuan_image3_model_taylor_cache_skips_transformer_between_full_steps(): + from lightx2v.models.networks.hunyuan_image3.infer.module_io import HunyuanImage3PreInferOutput + from lightx2v.models.networks.hunyuan_image3.model import HunyuanImage3Model + + class FakePreInfer: + def infer(self, weights, inputs): + value = float(inputs["cache_dic"]["current_step"]) + return HunyuanImage3PreInferOutput(hidden_states=torch.full((1, 3, 4), value)) + + class FakeTransformerInfer: + def __init__(self): + self.calls = 0 + + def infer(self, weights, pre_infer_out): + self.calls += 1 + return pre_infer_out.hidden_states + 10.0 + + class FakePostInfer: + def infer(self, weights, hidden_states, pre_infer_out): + return {"diffusion_prediction": hidden_states} + + model = HunyuanImage3Model.__new__(HunyuanImage3Model) + model.config = {"seq_parallel": False} + model.pre_weight = None + model.transformer_weights = None + model.post_weight = None + model.pre_infer = FakePreInfer() + model.transformer_infer = FakeTransformerInfer() + model.post_infer = FakePostInfer() + + cache_dic = { + "current_step": 0, + "cache_interval": 2, + "max_order": 2, + "num_steps": 3, + "enable_first_enhance": False, + "first_enhance_steps": 3, + "enable_tailing_enhance": False, + "tailing_enhance_steps": 1, + "low_freqs_order": 2, + "high_freqs_order": 2, + } + + outputs = [] + for step in range(3): + cache_dic["current_step"] = step + outputs.append(model.infer({"cache_dic": cache_dic})["diffusion_prediction"]) + + assert model.transformer_infer.calls == 2 + assert torch.equal(outputs[0], torch.full((1, 3, 4), 10.0)) + assert torch.equal(outputs[1], outputs[0]) + assert torch.equal(outputs[2], torch.full((1, 3, 4), 12.0)) + + +def test_hunyuan_image3_runner_builds_taylor_cache_config_from_runtime_config(): + from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner + + runner = HunyuanImage3Runner.__new__(HunyuanImage3Runner) + runner.config = { + "use_taylor_cache": True, + "taylor_cache_interval": 5, + "taylor_cache_order": 2, + "taylor_cache_enable_first_enhance": True, + "taylor_cache_first_enhance_steps": 4, + "taylor_cache_enable_tailing_enhance": True, + "taylor_cache_tailing_enhance_steps": 2, + "taylor_cache_low_freqs_order": 1, + "taylor_cache_high_freqs_order": 2, + } + + cache_dic = runner._build_taylor_cache_dic(num_steps=50) + + assert cache_dic == { + "counter": 0, + "current_step": 0, + "cache_interval": 5, + "max_order": 2, + "num_steps": 50, + "enable_first_enhance": True, + "first_enhance_steps": 4, + "enable_tailing_enhance": True, + "tailing_enhance_steps": 2, + "low_freqs_order": 1, + "high_freqs_order": 2, + "enable_force_control": False, + "force_compute": False, + } + + +def test_hunyuan_image3_rms_norm_ignores_inactive_diff_tensor(): + from lightx2v.common.ops.norm.rms_norm_weight import RMSWeightFP32 + + rms = RMSWeightFP32("norm.weight") + rms.weight = torch.ones(4) + rms.weight_diff = torch.ones(5) + rms.has_diff = False + + output = rms.apply(torch.ones(1, 4)) + + assert output.shape == (1, 4) + + +def test_hunyuan_image3_rms_norm_aligns_active_diff_to_weight(): + from lightx2v.common.ops.norm.rms_norm_weight import RMSWeightFP32 + + rms = RMSWeightFP32("norm.weight") + rms.weight = torch.ones(4, device="meta", dtype=torch.bfloat16) + rms.weight_diff = torch.ones(4, dtype=torch.float32) + rms.has_diff = True + + actual_weight = rms._get_actual_weight() + + assert actual_weight.device == rms.weight.device + assert actual_weight.dtype == rms.weight.dtype + + +def test_hunyuan_image3_rms_norm_register_diff_marks_active(): + from lightx2v.common.ops.norm.rms_norm_weight import RMSWeightFP32 + + rms = RMSWeightFP32("norm.weight") + + rms.register_diff({rms.weight_diff_name: torch.ones(4)}) + + assert rms.has_diff is True