diff --git a/configs/dist_infer/hunyuan_image3_t2i_dist_cfg.json b/configs/dist_infer/hunyuan_image3_t2i_dist_cfg.json new file mode 100644 index 000000000..b6a26bd35 --- /dev/null +++ b/configs/dist_infer/hunyuan_image3_t2i_dist_cfg.json @@ -0,0 +1,32 @@ +{ + "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, + "hunyuan_cfg_mode": "parallel", + "parallel": { + "cfg_p_size": 2 + }, + "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/dist_infer/hunyuan_image3_t2i_dist_cfg_kv_all_gather.json b/configs/dist_infer/hunyuan_image3_t2i_dist_cfg_kv_all_gather.json new file mode 100644 index 000000000..687702243 --- /dev/null +++ b/configs/dist_infer/hunyuan_image3_t2i_dist_cfg_kv_all_gather.json @@ -0,0 +1,30 @@ +{ + "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, + "pipeline_parallel": true, + "hunyuan_image3_pipeline_layout": "interleaved", + "enable_cfg": true, + "hunyuan_cfg_mode": "parallel", + "moe_impl": "flashinfer", + "attn_impl": "flash_attention_2", + "flashinfer_autotune_mode": "tune", + "parallel": { + "cfg_p_size": 2, + "seq_p_size": 2, + "seq_p_attn_type": "kv_all_gather" + }, + "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/dist_infer/hunyuan_image3_t2i_dist_cfg_ulysses.json b/configs/dist_infer/hunyuan_image3_t2i_dist_cfg_ulysses.json new file mode 100644 index 000000000..e19d7068e --- /dev/null +++ b/configs/dist_infer/hunyuan_image3_t2i_dist_cfg_ulysses.json @@ -0,0 +1,30 @@ +{ + "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, + "pipeline_parallel": true, + "hunyuan_image3_pipeline_layout": "interleaved", + "enable_cfg": true, + "hunyuan_cfg_mode": "parallel", + "moe_impl": "flashinfer", + "attn_impl": "sdpa", + "flashinfer_autotune_mode": "tune", + "parallel": { + "cfg_p_size": 2, + "seq_p_size": 2, + "seq_p_attn_type": "ulysses" + }, + "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/dist_infer/hunyuan_image3_t2i_dist_kv_all_gather.json b/configs/dist_infer/hunyuan_image3_t2i_dist_kv_all_gather.json new file mode 100644 index 000000000..0c605e531 --- /dev/null +++ b/configs/dist_infer/hunyuan_image3_t2i_dist_kv_all_gather.json @@ -0,0 +1,29 @@ +{ + "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, + "pipeline_parallel": true, + "enable_cfg": true, + "hunyuan_cfg_mode": "serial", + "moe_impl": "flashinfer", + "attn_impl": "sdpa", + "flashinfer_autotune_mode": "tune", + "parallel": { + "cfg_p_size": 1, + "seq_p_size": 2, + "seq_p_attn_type": "kv_all_gather" + }, + "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/dist_infer/hunyuan_image3_t2i_dist_ulysses.json b/configs/dist_infer/hunyuan_image3_t2i_dist_ulysses.json new file mode 100644 index 000000000..7c0ab31e3 --- /dev/null +++ b/configs/dist_infer/hunyuan_image3_t2i_dist_ulysses.json @@ -0,0 +1,29 @@ +{ + "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, + "pipeline_parallel": true, + "enable_cfg": true, + "hunyuan_cfg_mode": "serial", + "moe_impl": "flashinfer", + "attn_impl": "sdpa", + "flashinfer_autotune_mode": "tune", + "parallel": { + "cfg_p_size": 1, + "seq_p_size": 2, + "seq_p_attn_type": "ulysses" + }, + "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_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/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 f1e2cc40f..fc3834ad4 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.lingbot_video.lingbot_video_runner import LingBotVideoRunner # noqa: F401 @@ -49,6 +50,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) @@ -91,6 +103,7 @@ def main(): "wan2.2_s2v", "hunyuan_video_1.5", "hunyuan_video_1.5_distill", + "hunyuan_image3", "hunyuan3d", "worldplay_distill", "worldplay_ar", @@ -120,6 +133,110 @@ 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) + hunyuan_native_args.add_argument( + "--hunyuan_cfg_mode", + "--hunyuan-cfg-mode", + dest="hunyuan_cfg_mode", + type=str, + choices=["batch", "serial", "parallel"], + default=None, + help="HunyuanImage3 CFG execution mode: batch packs cond/uncond, serial runs two batch=1 forwards, parallel splits cond/uncond by cfg rank.", + ) + hunyuan_native_args.add_argument( + "--attn_impl", + "--attn-impl", + dest="attn_impl", + type=str, + choices=[ + "eager", + "sdpa", + "torch_sdpa", + "flash_attention_2", + "flash_attn2", + "flash_attention_3", + "flash_attn3", + "sage_attn2", + "sage_attn3", + ], + default=None, + ) + hunyuan_sp_args = parser.add_argument_group("HunyuanImage3 sequence parallel options") + hunyuan_sp_args.add_argument( + "--hunyuan_sp_size", + "--hunyuan-sp-size", + dest="hunyuan_sp_size", + type=int, + default=None, + ) + hunyuan_sp_args.add_argument( + "--hunyuan_sp_attn_type", + "--hunyuan-sp-attn-type", + dest="hunyuan_sp_attn_type", + type=str, + choices=["kv_all_gather", "kv-all-gather", "ulysses"], + default=None, + ) + hunyuan_sp_args.add_argument( + "--hunyuan_pipeline_layout", + "--hunyuan-pipeline-layout", + dest="hunyuan_image3_pipeline_layout", + type=str, + choices=["interleaved", "contiguous"], + default=None, + ) + hunyuan_flashinfer_args = parser.add_argument_group("HunyuanImage3 FlashInfer autotune options") + hunyuan_flashinfer_args.add_argument("--flashinfer_autotune", "--flashinfer-autotune", dest="flashinfer_autotune", type=_str2bool, nargs="?", const=True, default=None) + hunyuan_flashinfer_args.add_argument( + "--flashinfer_autotune_mode", + "--flashinfer-autotune-mode", + dest="flashinfer_autotune_mode", + type=str, + choices=["off", "tune", "load"], + default=None, + help="'tune' profiles missing FlashInfer fused MoE tactics and saves cache; 'load' only loads cache; 'off' disables autotune.", + ) + hunyuan_flashinfer_args.add_argument("--flashinfer_autotune_cache", "--flashinfer-autotune-cache", dest="flashinfer_autotune_cache", type=str, default=None) + hunyuan_flashinfer_args.add_argument("--flashinfer_tuning_buckets", "--flashinfer-tuning-buckets", dest="flashinfer_tuning_buckets", type=str, default=None) + hunyuan_flashinfer_args.add_argument( + "--flashinfer_autotune_round_up", + "--flashinfer-autotune-round-up", + dest="flashinfer_autotune_round_up", + type=_str2bool, + nargs="?", + const=True, + default=None, + ) + hunyuan_flashinfer_args.add_argument("--flashinfer_tune_max_num_tokens", "--flashinfer-tune-max-num-tokens", dest="flashinfer_tune_max_num_tokens", type=int, default=None) parser.add_argument("--use_prompt_enhancer", action="store_true") parser.add_argument("--warmup", action="store_true", help="Warm up the model before inference. Disabled by default.") parser.add_argument("--prompt", type=str, default="", help="The input prompt for text-to-video generation") 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..0730417bb --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/infer/kv_cache.py @@ -0,0 +1,130 @@ +import math +from dataclasses import dataclass + +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..a384b3e76 --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/infer/module_io.py @@ -0,0 +1,32 @@ +from dataclasses import dataclass + +import torch + + +@dataclass +class HunyuanImage3SequenceParallelState: + attn_type: str + original_seq_len: int + padded_seq_len: int + local_seq_len: int + local_start: int + valid_local_seq_len: int + global_position_ids: torch.Tensor | None = None + global_attention_mask: torch.Tensor | None = None + + +@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 + full_attn_slices: list | None = None + sequence_parallel_state: HunyuanImage3SequenceParallelState | None = None + attention_segment_specs: list | 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..4fe0e153e --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/infer/pre_infer.py @@ -0,0 +1,229 @@ +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"), + full_attn_slices=inputs.get("full_attn_slices"), + ) 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..7a2705835 --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/infer/transformer_infer.py @@ -0,0 +1,655 @@ +import torch +import torch.distributed as dist +import torch.nn.functional as F +from loguru import logger + +from lightx2v.common.ops.attn import * # noqa: F403,F401 - registers LightX2V attention kernels +from lightx2v.common.ops.attn.utils.all2all import all2all_head2seq, all2all_seq2head +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 +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER + +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): + ATTENTION_IMPL_ALIASES = { + "eager": "torch_sdpa", + "sdpa": "torch_sdpa", + "torch_sdpa": "torch_sdpa", + "flash_attention_2": "flash_attn2", + "flash_attn2": "flash_attn2", + "flash_attention_3": "flash_attn3", + "flash_attn3": "flash_attn3", + "sage_attn2": "sage_attn2", + "sage_attn3": "sage_attn3", + } + + 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") + self.flashinfer_tune_max_num_tokens = int(config.get("flashinfer_tune_max_num_tokens", 8192)) + self.attn_impl = self._normalize_attention_impl(config.get("attn_impl", "torch_sdpa")) + self.attn_kernel = None if self.attn_impl == "torch_sdpa" else self._build_attention_kernel(self.attn_impl) + self._attn_cu_seqlens_cache = {} + self._attn_fallback_warnings = set() + self._sp_gather_buffers = {} + self._pre_infer_device_cache = {} + if config.get("seq_parallel", False): + self.seq_p_group = config.get("device_mesh").get_group(mesh_dim="seq_p") + self.sequence_parallel_attn_type = str(config["parallel"].get("seq_p_attn_type", "kv_all_gather")).strip().lower().replace("-", "_") + if self.sequence_parallel_attn_type in ("kv_allgather", "kv_gather"): + self.sequence_parallel_attn_type = "kv_all_gather" + elif self.sequence_parallel_attn_type == "ulysses_sp": + self.sequence_parallel_attn_type = "ulysses" + else: + self.seq_p_group = None + self.sequence_parallel_attn_type = None + + def set_scheduler(self, scheduler): + self.scheduler = scheduler + + @classmethod + def _normalize_attention_impl(cls, attn_impl): + attn_impl = str(attn_impl or "torch_sdpa") + if attn_impl not in cls.ATTENTION_IMPL_ALIASES: + supported = ", ".join(sorted(cls.ATTENTION_IMPL_ALIASES)) + raise ValueError(f"Unsupported HunyuanImage3 attn_impl={attn_impl!r}. Supported values: {supported}.") + return cls.ATTENTION_IMPL_ALIASES[attn_impl] + + def _build_attention_kernel(self, attn_impl): + if attn_impl == "flash_attn2": + from lightx2v.common.ops.attn.flash_attn import flash_attn_func_v2, flash_attn_varlen_func_v2 + + if flash_attn_func_v2 is None or flash_attn_varlen_func_v2 is None: + raise ImportError("HunyuanImage3 attn_impl='flash_attn2' requires flash-attn v2.") + elif attn_impl == "flash_attn3": + from lightx2v.common.ops.attn.flash_attn import flash_attn_func_v3, flash_attn_varlen_func_v3 + + if flash_attn_func_v3 is None or flash_attn_varlen_func_v3 is None: + raise ImportError("HunyuanImage3 attn_impl='flash_attn3' requires flash-attn v3 / flash_attn_interface.") + elif attn_impl == "sage_attn2": + from lightx2v.common.ops.attn.sage_attn import sageattn + + if sageattn is None: + raise ImportError("HunyuanImage3 attn_impl='sage_attn2' requires sageattention.") + elif attn_impl == "sage_attn3": + from lightx2v.common.ops.attn.sage_attn import sageattn3_blackwell + + if sageattn3_blackwell is None: + raise ImportError("HunyuanImage3 attn_impl='sage_attn3' requires sageattention3.") + if attn_impl not in ATTN_WEIGHT_REGISTER: + raise ValueError(f"HunyuanImage3 attn_impl={attn_impl!r} is not registered in LightX2V ATTN_WEIGHT_REGISTER.") + return ATTN_WEIGHT_REGISTER[attn_impl]() + + def _normalize_attention_dtype(self, tensor): + if tensor.dtype in (torch.float16, torch.bfloat16): + return tensor + if tensor.device.type == "cuda": + return tensor.to(torch.bfloat16) + return tensor.to(torch.float32) + + def _get_cu_seqlens(self, name, batch, seq_len, device): + key = (name, batch, seq_len, device.type, device.index) + cu_seqlens = self._attn_cu_seqlens_cache.get(key) + if cu_seqlens is None: + cu_seqlens = torch.arange(0, batch * seq_len + 1, seq_len, dtype=torch.int32) + if self.attn_impl in ("flash_attn2", "flash_attn3"): + cu_seqlens = cu_seqlens.to(device, non_blocking=True) + self._attn_cu_seqlens_cache[key] = cu_seqlens + return cu_seqlens + + def _attention_mask_mode(self, attention_mask, q_len, kv_len): + if attention_mask is None: + return "none" + if attention_mask.dtype != torch.bool or attention_mask.dim() != 4: + return "custom" + if attention_mask.shape[-2] != q_len or attention_mask.shape[-1] != kv_len: + return "custom" + if attention_mask.shape[1] != 1: + return "custom" + + mask = attention_mask[:, 0] + if torch.all(mask): + return "full" + if q_len == kv_len: + causal_mask = torch.ones((q_len, kv_len), device=attention_mask.device, dtype=torch.bool).tril() + if torch.equal(mask, causal_mask.expand_as(mask)): + return "causal" + return "custom" + + def _warn_attention_fallback_once(self, mask_mode): + key = (self.attn_impl, mask_mode) + if key in self._attn_fallback_warnings: + return + self._attn_fallback_warnings.add(key) + logger.warning( + "HunyuanImage3 attn_impl='{}' does not support {} attention masks in the low-intrusion path; " + "falling back to PyTorch SDPA for this attention call.", + self.attn_impl, + mask_mode, + ) + + def _sdpa_attention(self, query_states, key_states, value_states, attention_mask): + 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() + return F.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask=attention_mask, dropout_p=0.0) + + def _apply_registered_attention_kernel(self, query_states, key_states, value_states, causal): + batch, query_heads, q_len, _ = query_states.shape + kv_len = key_states.shape[2] + original_dtype = query_states.dtype + q = self._normalize_attention_dtype(query_states.transpose(1, 2)).contiguous() + k = self._normalize_attention_dtype(key_states.transpose(1, 2)).contiguous() + v = self._normalize_attention_dtype(value_states.transpose(1, 2)).contiguous() + cu_seqlens_q = self._get_cu_seqlens("q", batch, q_len, q.device) + cu_seqlens_kv = self._get_cu_seqlens("kv", batch, kv_len, k.device) + attn_output = self.attn_kernel.apply( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=q_len, + max_seqlen_kv=kv_len, + causal=causal, + ) + if attn_output.dim() == 2: + attn_output = attn_output.reshape(batch, q_len, query_heads, self.head_dim) + elif attn_output.dim() == 3: + attn_output = attn_output.reshape(batch, q_len, query_heads, self.head_dim) + else: + raise RuntimeError(f"HunyuanImage3 attn_impl={self.attn_impl!r} returned unexpected shape {tuple(attn_output.shape)}.") + return attn_output.to(original_dtype).transpose(1, 2) + + def _normalize_full_attn_slices(self, full_attn_slices, batch): + if full_attn_slices is None: + return [[] for _ in range(batch)] + if len(full_attn_slices) == 0: + return [[] for _ in range(batch)] + + first_item = full_attn_slices[0] + if isinstance(first_item, slice) or ( + isinstance(first_item, (list, tuple)) and len(first_item) == 2 and all(isinstance(v, int) for v in first_item) + ): + full_attn_slices = [full_attn_slices for _ in range(batch)] + elif len(full_attn_slices) == 1 and batch > 1: + full_attn_slices = list(full_attn_slices) * batch + + normalized = [] + for sample_slices in full_attn_slices: + sample = [] + for item in sample_slices: + if isinstance(item, slice): + start = 0 if item.start is None else int(item.start) + stop = int(item.stop) + else: + start, stop = int(item[0]), int(item[1]) + if stop > start: + sample.append((start, stop)) + normalized.append(sorted(sample)) + if len(normalized) != batch: + return None + return normalized + + @staticmethod + def _find_full_attn_slice(full_slices, position): + for start, stop in full_slices: + if start <= position < stop: + return start, stop + return None + + def _build_segment_specs(self, position_ids, full_slices, kv_len): + if torch.is_tensor(position_ids): + positions = [int(position) for position in position_ids.detach().cpu().reshape(-1).tolist()] + else: + positions = [int(position) for position in position_ids] + segments = [] + local_start = 0 + while local_start < len(positions): + pos = positions[local_start] + full_slice = self._find_full_attn_slice(full_slices, pos) + causal = full_slice is None + kv_end = pos + 1 if causal else full_slice[1] + if kv_end <= 0 or kv_end > kv_len: + return None + + local_end = local_start + 1 + previous_pos = pos + while local_end < len(positions): + next_pos = positions[local_end] + next_full_slice = self._find_full_attn_slice(full_slices, next_pos) + next_causal = next_full_slice is None + if next_causal != causal or next_pos != previous_pos + 1: + break + if causal: + next_kv_end = next_pos + 1 + else: + if next_full_slice != full_slice: + break + next_kv_end = full_slice[1] + if next_kv_end <= 0 or next_kv_end > kv_len: + return None + previous_pos = next_pos + kv_end = next_kv_end + local_end += 1 + + segments.append((local_start, local_end, kv_end, causal)) + local_start = local_end + return segments + + def _segmented_flash_attention(self, query_states, key_states, value_states, position_ids, full_attn_slices, segment_specs=None): + if position_ids is None: + return None + batch, _, q_len, _ = query_states.shape + kv_len = key_states.shape[2] + if position_ids.shape != (batch, q_len): + return None + + if segment_specs is None: + if full_attn_slices is None: + return None + batch_full_slices = self._normalize_full_attn_slices(full_attn_slices, batch) + if batch_full_slices is None or not any(batch_full_slices): + return None + segment_specs = [ + self._build_segment_specs(position_ids[batch_idx], batch_full_slices[batch_idx], kv_len) + for batch_idx in range(batch) + ] + if len(segment_specs) != batch or any(specs is None for specs in segment_specs): + return None + + output = torch.empty_like(query_states) + for batch_idx in range(batch): + for q_start, q_stop, kv_stop, causal in segment_specs[batch_idx]: + segment_output = self._apply_registered_attention_kernel( + query_states[batch_idx : batch_idx + 1, :, q_start:q_stop], + key_states[batch_idx : batch_idx + 1, :, :kv_stop], + value_states[batch_idx : batch_idx + 1, :, :kv_stop], + causal=causal, + ) + output[batch_idx : batch_idx + 1, :, q_start:q_stop] = segment_output + return output + + def _registered_attention( + self, + query_states, + key_states, + value_states, + attention_mask, + position_ids=None, + full_attn_slices=None, + segment_specs=None, + ): + if self.attn_impl in ("flash_attn2", "flash_attn3") and segment_specs is not None: + segmented_output = self._segmented_flash_attention( + query_states, + key_states, + value_states, + position_ids, + full_attn_slices, + segment_specs=segment_specs, + ) + if segmented_output is not None: + return segmented_output + raise RuntimeError("HunyuanImage3 precomputed segmented attention plan does not match the current Q/KV layout.") + + batch, _, q_len, _ = query_states.shape + kv_len = key_states.shape[2] + mask_mode = self._attention_mask_mode(attention_mask, q_len, kv_len) + + if self.attn_impl == "torch_sdpa": + return self._sdpa_attention(query_states, key_states, value_states, attention_mask) + if self.attn_impl in ("flash_attn2", "flash_attn3"): + if mask_mode not in ("none", "full", "causal"): + segmented_output = self._segmented_flash_attention( + query_states, + key_states, + value_states, + position_ids, + full_attn_slices, + segment_specs=segment_specs, + ) + if segmented_output is not None: + return segmented_output + self._warn_attention_fallback_once(mask_mode) + return self._sdpa_attention(query_states, key_states, value_states, attention_mask) + causal = mask_mode == "causal" + elif self.attn_impl in ("sage_attn2", "sage_attn3"): + if mask_mode not in ("none", "full"): + self._warn_attention_fallback_once(mask_mode) + return self._sdpa_attention(query_states, key_states, value_states, attention_mask) + causal = False + else: + raise ValueError(f"Unsupported HunyuanImage3 normalized attn_impl={self.attn_impl!r}.") + + return self._apply_registered_attention_kernel(query_states, key_states, value_states, causal=causal) + + @torch.no_grad() + def infer(self, weights, pre_infer_out): + self._pre_infer_device_cache = {} + pre_infer_out.attention_segment_specs = self._prepare_attention_segment_specs(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) + if device is not None and device.type == "cuda" and device.index is not None: + torch.cuda.set_device(device.index) + hidden_states = to_device(hidden_states, device) + use_segment_specs = self.attn_impl in ("flash_attn2", "flash_attn3") and pre_infer_out.attention_segment_specs is not None + attention_mask = None if use_segment_specs else self._cached_pre_infer_to_device("attention_mask", pre_infer_out.attention_mask, device) + position_ids = self._cached_pre_infer_to_device("position_ids", pre_infer_out.position_ids, device) + custom_pos_emb = self._cached_pre_infer_to_device("custom_pos_emb", 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.full_attn_slices, + pre_infer_out.past_key_values if pre_infer_out.use_cache else None, + pre_infer_out.sequence_parallel_state, + pre_infer_out.attention_segment_specs, + ) + 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, + full_attn_slices=None, + past_key_values=None, + sequence_parallel_state=None, + segment_specs=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) + + cache_position_ids = position_ids + if sequence_parallel_state is not None: + if sequence_parallel_state.attn_type == "kv_all_gather": + key_states, value_states = self._sequence_parallel_gather_kv( + key_states, + value_states, + sequence_parallel_state, + ) + cache_position_ids = self._cached_pre_infer_to_device( + "sp_global_position_ids", + sequence_parallel_state.global_position_ids, + key_states.device, + ) + elif sequence_parallel_state.attn_type == "ulysses": + query_states, key_states, value_states = self._sequence_parallel_ulysses_seq_to_head( + query_states, + key_states, + value_states, + sequence_parallel_state, + ) + cache_position_ids = self._cached_pre_infer_to_device( + "sp_global_position_ids", + sequence_parallel_state.global_position_ids, + key_states.device, + ) + position_ids = cache_position_ids + if not (self.attn_impl in ("flash_attn2", "flash_attn3") and segment_specs is not None): + attention_mask = self._cached_pre_infer_to_device( + "sp_global_attention_mask", + sequence_parallel_state.global_attention_mask, + key_states.device, + ) + else: + raise ValueError(f"Unsupported HunyuanImage3 sequence parallel attention type: {sequence_parallel_state.attn_type!r}.") + + if past_key_values is not None: + if cache_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, cache_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 sequence_parallel_state is not None and sequence_parallel_state.attn_type == "kv_all_gather": + valid_q_len = sequence_parallel_state.valid_local_seq_len + attn_output = torch.zeros_like(query_states) + if valid_q_len: + attn_output[:, :, :valid_q_len] = self._registered_attention( + query_states[:, :, :valid_q_len], + key_states, + value_states, + None if attention_mask is None else attention_mask[:, :, :valid_q_len], + position_ids=None if position_ids is None else position_ids[:, :valid_q_len], + full_attn_slices=full_attn_slices, + segment_specs=segment_specs, + ) + else: + attn_output = self._registered_attention( + query_states, + key_states, + value_states, + attention_mask, + position_ids=position_ids, + full_attn_slices=full_attn_slices, + segment_specs=segment_specs, + ) + + if sequence_parallel_state is not None and sequence_parallel_state.attn_type == "ulysses": + attn_output = self._sequence_parallel_ulysses_head_to_seq(attn_output, sequence_parallel_state) + + 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 _prepare_attention_segment_specs(self, pre_infer_out): + if self.attn_impl not in ("flash_attn2", "flash_attn3"): + return None + + state = pre_infer_out.sequence_parallel_state + if state is not None and state.attn_type == "ulysses": + position_ids = state.global_position_ids + attention_mask = state.global_attention_mask + else: + position_ids = pre_infer_out.position_ids + attention_mask = pre_infer_out.attention_mask + if state is not None: + position_ids = position_ids[:, : state.valid_local_seq_len] + if attention_mask is not None: + attention_mask = attention_mask[:, :, : state.valid_local_seq_len] + + if position_ids is None or attention_mask is None: + return None + if pre_infer_out.full_attn_slices is None: + # A generic custom mask has no lossless causal/full-slice description. + # Keep the dense mask and let registered attention fall back to SDPA. + return None + batch, q_len = position_ids.shape + kv_len = attention_mask.shape[-1] + if self._attention_mask_mode(attention_mask, q_len, kv_len) != "custom": + return None + batch_full_slices = self._normalize_full_attn_slices(pre_infer_out.full_attn_slices, batch) + if batch_full_slices is None or not any(batch_full_slices): + return None + segment_specs = [ + self._build_segment_specs(position_ids[batch_idx], batch_full_slices[batch_idx], kv_len) + for batch_idx in range(batch) + ] + return None if any(specs is None for specs in segment_specs) else segment_specs + + def _cached_pre_infer_to_device(self, name, value, device): + if value is None: + return None + cache = getattr(self, "_pre_infer_device_cache", None) + if cache is None: + self._pre_infer_device_cache = {} + cache = self._pre_infer_device_cache + key = (name, device, id(value)) + cached = cache.get(key) + if cached is None: + cached = to_device(value, device) + cache[key] = cached + return cached + + def _sequence_parallel_gather_kv(self, key_states, value_states, state): + world_size = dist.get_world_size(self.seq_p_group) + local = torch.stack((key_states, value_states), dim=2).permute(3, 0, 2, 1, 4).contiguous() + output_shape = (local.shape[0] * world_size, *local.shape[1:]) + buffer_key = ("kv", local.device, local.dtype, output_shape) + gathered = self._sp_gather_buffers.get(buffer_key) + if gathered is None or gathered.shape != output_shape: + gathered = torch.empty(output_shape, device=local.device, dtype=local.dtype) + self._sp_gather_buffers[buffer_key] = gathered + dist.all_gather_into_tensor(gathered, local, group=self.seq_p_group) + valid = gathered[: state.original_seq_len] + key_value = valid.permute(2, 1, 3, 0, 4).contiguous() + return key_value[0], key_value[1] + + def _sequence_parallel_ulysses_seq_to_head(self, query_states, key_states, value_states, state): + if query_states.shape[0] != 1: + raise ValueError( + "HunyuanImage3 Ulysses expects batch size 1; use hunyuan_cfg_mode='serial' " + "for pure SP or 'parallel' for CFG+SP." + ) + query = all2all_seq2head(query_states[0].transpose(0, 1).contiguous(), group=self.seq_p_group) + key = all2all_seq2head(key_states[0].transpose(0, 1).contiguous(), group=self.seq_p_group) + value = all2all_seq2head(value_states[0].transpose(0, 1).contiguous(), group=self.seq_p_group) + original_seq_len = state.original_seq_len + query = query[:original_seq_len].transpose(0, 1).unsqueeze(0).contiguous() + key = key[:original_seq_len].transpose(0, 1).unsqueeze(0).contiguous() + value = value[:original_seq_len].transpose(0, 1).unsqueeze(0).contiguous() + return query, key, value + + def _sequence_parallel_ulysses_head_to_seq(self, attn_output, state): + output = attn_output[0].transpose(0, 1).contiguous() + padding_size = state.padded_seq_len - state.original_seq_len + if padding_size: + padding = output.new_zeros(padding_size, output.shape[1], output.shape[2]) + output = torch.cat((output, padding), dim=0) + output = all2all_head2seq(output, group=self.seq_p_group) + return output.unsqueeze(0).transpose(1, 2).contiguous() + + 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, + tune_max_num_tokens=self.flashinfer_tune_max_num_tokens, + ) + 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..a716d355e --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/model.py @@ -0,0 +1,475 @@ +import os +import re + +import torch +import torch.distributed as dist +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.module_io import HunyuanImage3SequenceParallelState +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_sequence_parallel_pipeline_lane(config, devices): + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("HunyuanImage3 sequence-parallel pipeline split requires torch.distributed initialization.") + + # In pure SP, global rank/world are identical to seq rank/world. In a + # CFG+SP mesh, however, using seq rank would make both CFG rows select the + # same physical pipeline devices. Allocate one disjoint lane per global + # process instead; collectives themselves still use the seq_p group. + parallel_world_size = dist.get_world_size() + global_rank = dist.get_rank() + if len(devices) % parallel_world_size: + raise ValueError( + f"HunyuanImage3 sequence parallel requires pipeline device count ({len(devices)}) " + f"to be divisible by cfg_p_size * seq_p_size ({parallel_world_size})." + ) + + layout = str(config.get("hunyuan_image3_pipeline_layout", "interleaved")).strip().lower() + if layout != "interleaved": + raise ValueError( + "HunyuanImage3 sequence parallel uses Wan-style rank/device initialization and currently requires " + f"hunyuan_image3_pipeline_layout='interleaved'; got {layout!r}." + ) + + lane = devices[global_rank::parallel_world_size] + expected_device = f"cuda:{torch.cuda.current_device()}" + if lane[0] != expected_device: + raise RuntimeError( + "HunyuanImage3 interleaved pipeline lane must start on the CUDA device selected by " + f"init_parallel_env: lane={lane}, current_device={expected_device}." + ) + return lane + + +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) + devices = [_normalize_device_name(device) for device in devices] + if config.get("seq_parallel", False): + return _resolve_sequence_parallel_pipeline_lane(config, devices) + return devices + + if config.get("pipeline_parallel", False) and torch.cuda.is_available(): + device_count = torch.cuda.device_count() + if device_count > 0: + if config.get("seq_parallel", False): + devices = [f"cuda:{idx}" for idx in range(device_count)] + return _resolve_sequence_parallel_pipeline_lane(config, devices) + if config.get("cfg_parallel", False): + cfg_p_size = int(config.get("parallel", {}).get("cfg_p_size", 1)) + if cfg_p_size > 1: + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("HunyuanImage3 cfg_parallel pipeline split requires an initialized distributed process group.") + if device_count % cfg_p_size != 0: + raise ValueError( + f"HunyuanImage3 cfg_parallel requires visible CUDA device count ({device_count}) " + f"to be divisible by cfg_p_size ({cfg_p_size})." + ) + cfg_p_group = config["device_mesh"].get_group(mesh_dim="cfg_p") + cfg_p_rank = dist.get_rank(cfg_p_group) + devices_per_cfg_rank = device_count // cfg_p_size + start = cfg_p_rank * devices_per_cfg_rank + stop = start + devices_per_cfg_rank + return [f"cuda:{idx}" for idx in range(start, stop)] + 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 + self.sequence_parallel_attn_type = self._resolve_sequence_parallel_attn_type() + self._sp_gather_buffers = {} + self._validate_sequence_parallel_config() + 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() + + if self.config.get("seq_parallel", False): + logger.info( + "HunyuanImage3 sequence parallel initialized: " + f"rank={dist.get_rank(self.seq_p_group)}, " + f"size={dist.get_world_size(self.seq_p_group)}, " + f"attention={self.sequence_parallel_attn_type}, " + f"pipeline_devices={self.pipeline_devices}" + ) + + def _resolve_sequence_parallel_attn_type(self): + if not self.config.get("seq_parallel", False): + return None + parallel = self.config.get("parallel") or {} + attn_type = str(parallel.get("seq_p_attn_type", "kv_all_gather")).strip().lower().replace("-", "_") + aliases = { + "kv_allgather": "kv_all_gather", + "kv_gather": "kv_all_gather", + "ulysses_sp": "ulysses", + } + attn_type = aliases.get(attn_type, attn_type) + if attn_type not in ("kv_all_gather", "ulysses"): + raise ValueError( + "HunyuanImage3 parallel.seq_p_attn_type must be one of: " + f"kv_all_gather, ulysses; got {attn_type!r}." + ) + return attn_type + + def _validate_sequence_parallel_config(self): + if not self.config.get("seq_parallel", False): + return + if self.seq_p_group is None: + raise RuntimeError("HunyuanImage3 sequence parallel requires an initialized seq_p process group.") + cfg_mode = str(self.config.get("hunyuan_cfg_mode", "batch")).strip().lower() + if self.config.get("cfg_parallel", False): + if not self.config.get("enable_cfg", False): + raise ValueError("HunyuanImage3 cfg_parallel requires enable_cfg=true.") + if cfg_mode != "parallel": + raise ValueError("HunyuanImage3 CFG+SP requires hunyuan_cfg_mode='parallel'.") + elif self.config.get("enable_cfg", False) and cfg_mode != "serial": + raise ValueError( + "HunyuanImage3 sequence parallel requires hunyuan_cfg_mode='serial' so every transformer forward has batch size 1." + ) + if self.config.get("use_taylor_cache", False) and self.config.get("enable_kv_cache", False): + raise ValueError("HunyuanImage3 sequence parallel does not support enabling Taylor cache and KV cache together.") + if self.sequence_parallel_attn_type == "ulysses": + world_size = dist.get_world_size(self.seq_p_group) + q_heads = int(self.config["num_attention_heads"]) + kv_heads = int(self.config.get("num_key_value_heads") or q_heads) + if q_heads % world_size or kv_heads % world_size: + raise ValueError( + "HunyuanImage3 Ulysses requires seq_p_size to divide both Q and KV heads: " + f"Q={q_heads}, KV={kv_heads}, seq_p_size={world_size}." + ) + + 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, pre_infer_out.sequence_parallel_state) + 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) + # for seq_parallel, the _seq_parallel_post_process function is run in _infer_transformer + 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): + if pre_infer_out.hidden_states.shape[0] != 1: + raise ValueError( + "HunyuanImage3 sequence parallel expects batch size 1 per transformer forward; " + "use hunyuan_cfg_mode='serial' for pure SP or 'parallel' for CFG+SP." + ) + + world_size = dist.get_world_size(self.seq_p_group) + rank = dist.get_rank(self.seq_p_group) + original_seq_len = int(pre_infer_out.hidden_states.shape[1]) + padding_size = (-original_seq_len) % world_size + padded_seq_len = original_seq_len + padding_size + local_seq_len = padded_seq_len // world_size + local_start = rank * local_seq_len + valid_local_seq_len = max(0, min(local_seq_len, original_seq_len - local_start)) + + global_position_ids = pre_infer_out.position_ids + global_attention_mask = pre_infer_out.attention_mask if self.sequence_parallel_attn_type == "ulysses" else None + + pre_infer_out.hidden_states = self._pad_and_slice_sequence( + pre_infer_out.hidden_states, 1, padding_size, local_start, local_seq_len, value=0 + ) + pre_infer_out.position_ids = self._pad_and_slice_sequence( + pre_infer_out.position_ids, 1, padding_size, local_start, local_seq_len, value=0 + ) + if pre_infer_out.custom_pos_emb is not None: + cos, sin = pre_infer_out.custom_pos_emb + pre_infer_out.custom_pos_emb = ( + self._pad_and_slice_sequence(cos, 1, padding_size, local_start, local_seq_len, value=1), + self._pad_and_slice_sequence(sin, 1, padding_size, local_start, local_seq_len, value=0), + ) + if pre_infer_out.attention_mask is not None: + local_attention_mask = self._pad_and_slice_sequence( + pre_infer_out.attention_mask, + -2, + padding_size, + local_start, + local_seq_len, + value=False, + ) + if valid_local_seq_len < local_seq_len: + local_attention_mask = local_attention_mask.clone() + local_attention_mask[..., valid_local_seq_len:, 0] = True + pre_infer_out.attention_mask = local_attention_mask + + pre_infer_out.sequence_parallel_state = HunyuanImage3SequenceParallelState( + attn_type=self.sequence_parallel_attn_type, + original_seq_len=original_seq_len, + padded_seq_len=padded_seq_len, + local_seq_len=local_seq_len, + local_start=local_start, + valid_local_seq_len=valid_local_seq_len, + global_position_ids=global_position_ids, + global_attention_mask=global_attention_mask, + ) + return pre_infer_out + + @torch.no_grad() + def _seq_parallel_post_process(self, x, sequence_parallel_state): + if sequence_parallel_state is None: + raise RuntimeError("HunyuanImage3 sequence parallel post-process is missing sequence metadata.") + world_size = dist.get_world_size(self.seq_p_group) + local = x.transpose(0, 1).contiguous() + output_shape = (local.shape[0] * world_size, *local.shape[1:]) + key = ("hidden", local.device, local.dtype, output_shape) + gathered = self._sp_gather_buffers.get(key) + if gathered is None or gathered.shape != output_shape: + gathered = torch.empty(output_shape, device=local.device, dtype=local.dtype) + self._sp_gather_buffers[key] = gathered + dist.all_gather_into_tensor(gathered, local, group=self.seq_p_group) + return gathered[: sequence_parallel_state.original_seq_len].transpose(0, 1).contiguous() + + @staticmethod + def _pad_and_slice_sequence(tensor, sequence_dim, padding_size, start, length, value): + if tensor is None: + return None + sequence_dim %= tensor.ndim + if padding_size: + pad_shape = list(tensor.shape) + pad_shape[sequence_dim] = padding_size + padding = tensor.new_full(pad_shape, value) + tensor = torch.cat((tensor, padding), dim=sequence_dim) + return tensor.narrow(sequence_dim, start, length).contiguous() + + def _guidance_scale(self): + return float(self.config.get("sample_guide_scale", self.config.get("diff_guidance_scale", 1.0))) + + def _set_cfg_scheduler_predictions(self, noise_pred_cond, noise_pred_uncond, noise_pred_guided): + if not hasattr(self, "scheduler"): + return + self.scheduler.noise_pred_cond = noise_pred_cond + self.scheduler.noise_pred_uncond = noise_pred_uncond + self.scheduler.noise_pred_guided = noise_pred_guided + self.scheduler.noise_pred = noise_pred_guided + + def combine_cfg_predictions(self, noise_pred_cond, noise_pred_uncond): + return noise_pred_uncond + self._guidance_scale() * (noise_pred_cond - noise_pred_uncond) + + @torch.no_grad() + def infer_branch(self, inputs, infer_condition=True): + output = self._infer_cond_uncond(inputs, infer_condition=infer_condition) + if hasattr(self, "scheduler") and "diffusion_prediction" in output: + self.scheduler.noise_pred = output["diffusion_prediction"] + return output + + @torch.no_grad() + def infer_cfg_serial(self, cond_inputs, uncond_inputs): + cond_output = self._infer_cond_uncond(cond_inputs, infer_condition=True) + uncond_output = self._infer_cond_uncond(uncond_inputs, infer_condition=False) + noise_pred_cond = cond_output["diffusion_prediction"] + noise_pred_uncond = uncond_output["diffusion_prediction"] + noise_pred_guided = self.combine_cfg_predictions(noise_pred_cond, noise_pred_uncond) + output = dict(cond_output) + output["diffusion_prediction"] = noise_pred_guided + self._set_cfg_scheduler_predictions(noise_pred_cond, noise_pred_uncond, noise_pred_guided) + return output + + @torch.no_grad() + def infer(self, inputs): + cfg_parallel_branch = bool(inputs.get("_cfg_parallel_branch", False)) + infer_condition = True + cfg_p_group = None + # 当当前 forward 属于 CFG parallel 时，找到当前进程对应的 CFG 通信组，并判断它负责 conditional 还是 unconditional 分支。 + if cfg_parallel_branch: + if not self.config.get("cfg_parallel", False): + raise RuntimeError("HunyuanImage3 received a cfg-parallel branch input, but config['cfg_parallel'] is not enabled.") + cfg_p_group = self.config["device_mesh"].get_group(mesh_dim="cfg_p") + assert dist.get_world_size(cfg_p_group) == 2, "cfg_p_world_size must be equal to 2" + infer_condition = dist.get_rank(cfg_p_group) == 0 + + output = self._infer_cond_uncond(inputs, infer_condition=infer_condition) + + if cfg_parallel_branch and "diffusion_prediction" in output: + # Keep the CFG collective on the device that produced the + # diffusion prediction (the last device of each pipeline lane). + # This matches the proven pure-CFG path and avoids an extra + # cross-device copy immediately before the NCCL collective. + noise_pred = output["diffusion_prediction"].contiguous() + if noise_pred.device.type == "cuda": + torch.cuda.set_device(noise_pred.device) + noise_pred_list = [torch.empty_like(noise_pred) for _ in range(2)] + dist.all_gather(noise_pred_list, noise_pred, group=cfg_p_group) + noise_pred_cond = noise_pred_list[0] + noise_pred_uncond = noise_pred_list[1] + noise_pred_guided = self.combine_cfg_predictions(noise_pred_cond, noise_pred_uncond) + output["diffusion_prediction"] = noise_pred_guided + self._set_cfg_scheduler_predictions(noise_pred_cond, noise_pred_uncond, noise_pred_guided) + elif 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/flashinfer_autotune.py b/lightx2v/models/runners/hunyuan_image3/flashinfer_autotune.py new file mode 100644 index 000000000..8058275a5 --- /dev/null +++ b/lightx2v/models/runners/hunyuan_image3/flashinfer_autotune.py @@ -0,0 +1,284 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping +from contextlib import AbstractContextManager, contextmanager, nullcontext +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import torch +import torch.distributed as dist +from loguru import logger + +try: + from flashinfer.autotuner import AutoTuner as FlashInferAutoTuner + from flashinfer.autotuner import autotune as flashinfer_autotune +except Exception: + FlashInferAutoTuner = None + flashinfer_autotune = None + + +FlashInferAutotuneMode = Literal["off", "tune", "load"] +AutotuneFactory = Callable[..., AbstractContextManager] +TunerProvider = Callable[[], Any] +StatusDeviceResolver = Callable[[], torch.device] + + +@dataclass(frozen=True) +class FlashInferAutotuneOptions: + """Normalized, side-effect-free FlashInfer autotune configuration.""" + + mode: FlashInferAutotuneMode + cache_path: Path | None = None + tuning_buckets: tuple[int, ...] | None = None + round_up: bool | None = None + tune_max_num_tokens: int = 8192 + + @property + def enabled(self) -> bool: + return self.mode != "off" + + @property + def tune_mode(self) -> bool: + return self.mode == "tune" + + @classmethod + def from_config(cls, config: Mapping[str, Any]) -> FlashInferAutotuneOptions: + mode = cls.resolve_mode(config) + if mode == "off": + return cls(mode=mode) + return cls( + mode=mode, + cache_path=cls._resolve_cache_path(config.get("flashinfer_autotune_cache")), + tuning_buckets=cls._parse_tuning_buckets(config.get("flashinfer_tuning_buckets")), + round_up=cls._parse_optional_bool(config.get("flashinfer_autotune_round_up")), + tune_max_num_tokens=int(config.get("flashinfer_tune_max_num_tokens", 8192)), + ) + + @classmethod + def resolve_mode(cls, config: Mapping[str, Any]) -> FlashInferAutotuneMode: + mode = config.get("flashinfer_autotune_mode") + if mode is None and "flashinfer_autotune" in config: + mode = "tune" if cls._parse_bool(config.get("flashinfer_autotune")) else "off" + if mode is None: + return "off" + if isinstance(mode, bool): + return "tune" if mode else "off" + + normalized = str(mode).strip().lower() + bool_aliases = { + "1": "tune", + "true": "tune", + "yes": "tune", + "y": "tune", + "on": "tune", + "0": "off", + "false": "off", + "no": "off", + "n": "off", + } + normalized = bool_aliases.get(normalized, normalized) + if normalized not in ("off", "tune", "load"): + raise ValueError("flashinfer_autotune_mode must be one of: off, tune, load.") + return normalized + + @staticmethod + def _parse_bool(value: Any) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in ("1", "true", "yes", "y", "on"): + return True + if normalized in ("0", "false", "no", "n", "off"): + return False + return bool(value) + + @classmethod + def _parse_optional_bool(cls, value: Any) -> bool | None: + return None if value is None else cls._parse_bool(value) + + @staticmethod + def _resolve_cache_path(value: Any) -> Path | None: + if value in (None, ""): + return None + return Path(value).expanduser() + + @staticmethod + def _parse_tuning_buckets(value: Any) -> tuple[int, ...] | None: + if value in (None, ""): + return None + if isinstance(value, int): + return (value,) + if isinstance(value, (list, tuple)): + buckets = tuple(int(item) for item in value if str(item).strip()) + else: + buckets = tuple(int(item.strip()) for item in str(value).split(",") if item.strip()) + if not buckets: + raise ValueError("flashinfer_tuning_buckets must contain at least one integer when provided.") + return tuple(sorted(set(buckets))) + + +@dataclass(frozen=True) +class DistributedAutotuneContext: + """Distributed resources selected by the runner's parallel topology.""" + + coordination_required: bool + process_group: dist.ProcessGroup | None = None + status_device_resolver: StatusDeviceResolver | None = None + is_cache_writer: bool = False + + +class FlashInferAutotuneController: + """Own the local or distributed FlashInfer autotune lifecycle.""" + + def __init__( + self, + options: FlashInferAutotuneOptions, + distributed_context: DistributedAutotuneContext, + moe_impl: str, + *, + autotune_factory: AutotuneFactory | None = None, + tuner_provider: TunerProvider | None = None, + ): + self.options = options + self.distributed_context = distributed_context + self.moe_impl = moe_impl + self._autotune_factory = autotune_factory + self._tuner_provider = tuner_provider + + @classmethod + def from_config( + cls, + config: Mapping[str, Any], + distributed_context: DistributedAutotuneContext, + *, + autotune_factory: AutotuneFactory | None = None, + tuner_provider: TunerProvider | None = None, + ) -> FlashInferAutotuneController: + """Build a controller without validating inactive autotune settings.""" + moe_impl = config.get("moe_impl", "eager") + mode = FlashInferAutotuneOptions.resolve_mode(config) + options = FlashInferAutotuneOptions(mode=mode) + if mode != "off" and moe_impl == "flashinfer": + options = FlashInferAutotuneOptions.from_config(config) + return cls( + options=options, + distributed_context=distributed_context, + moe_impl=moe_impl, + autotune_factory=autotune_factory, + tuner_provider=tuner_provider, + ) + + def context(self) -> AbstractContextManager: + if not self.options.enabled: + return nullcontext() + if self.moe_impl != "flashinfer": + logger.warning("flashinfer_autotune_mode is set but moe_impl is not 'flashinfer'; autotune is disabled for this run.") + return nullcontext() + + autotune_factory = self._resolve_autotune_factory() + self._prepare_cache_path() + logger.info( + "HunyuanImage3 FlashInfer autotune enabled: " + f"mode={self.options.mode}, cache={self.options.cache_path}, " + f"buckets={self.options.tuning_buckets}, round_up={self.options.round_up}, " + f"tune_max_num_tokens={self.options.tune_max_num_tokens}" + ) + + if self.distributed_context.coordination_required: + return self._distributed_autotune_context(autotune_factory) + return self._local_autotune_context(autotune_factory) + + def _resolve_autotune_factory(self) -> AutotuneFactory: + factory = self._autotune_factory or flashinfer_autotune + if factory is None: + raise ImportError("flashinfer_autotune_mode requires flashinfer.autotuner.autotune.") + return factory + + def _resolve_tuner(self) -> Any: + if self._tuner_provider is not None: + return self._tuner_provider() + if FlashInferAutoTuner is None: + raise ImportError("Distributed FlashInfer autotune requires flashinfer.autotuner.AutoTuner.") + return FlashInferAutoTuner.get() + + def _prepare_cache_path(self) -> None: + cache_path = self.options.cache_path + if cache_path is not None: + cache_path.parent.mkdir(parents=True, exist_ok=True) + if self.options.mode == "load" and cache_path is None: + raise ValueError("flashinfer_autotune_mode='load' requires flashinfer_autotune_cache.") + + def _local_autotune_context(self, autotune_factory: AutotuneFactory) -> AbstractContextManager: + cache_path = self.options.cache_path + if self.options.mode == "load" and (cache_path is None or not cache_path.is_file()): + raise FileNotFoundError(f"FlashInfer autotune cache does not exist for load mode: {cache_path}") + return autotune_factory( + tune_mode=self.options.tune_mode, + cache=None if cache_path is None else str(cache_path), + tuning_buckets=self.options.tuning_buckets, + round_up=self.options.round_up, + ) + + @contextmanager + def _distributed_autotune_context(self, autotune_factory: AutotuneFactory): + tuner = self._resolve_tuner() + runtime = self.distributed_context + group = runtime.process_group + if group is None: + raise RuntimeError("Distributed FlashInfer autotune requires an initialized CFG or sequence parallel group.") + + if dist.get_backend(group) == "nccl": + if runtime.status_device_resolver is None: + raise RuntimeError("Distributed FlashInfer autotune requires an NCCL status device resolver.") + status_device = torch.device(runtime.status_device_resolver()) + torch.cuda.set_device(status_device) + else: + status_device = torch.device("cpu") + + cache_path = None if self.options.cache_path is None else str(self.options.cache_path) + load_error = self._load_distributed_cache(tuner, cache_path) + load_succeeded = torch.tensor([0 if load_error is not None else 1], device=status_device, dtype=torch.int32) + dist.all_reduce(load_succeeded, op=dist.ReduceOp.MIN, group=group) + if not bool(load_succeeded.item()): + raise RuntimeError(f"Distributed FlashInfer autotune cache preflight failed for {cache_path}.") from load_error + + with autotune_factory( + tune_mode=self.options.tune_mode, + cache=None, + tuning_buckets=self.options.tuning_buckets, + round_up=self.options.round_up, + ): + yield + + save_error = self._save_distributed_cache(tuner, cache_path) + save_succeeded = torch.tensor([0 if save_error is not None else 1], device=status_device, dtype=torch.int32) + source_rank = dist.get_global_rank(group, 0) + dist.broadcast(save_succeeded, src=source_rank, group=group) + if not bool(save_succeeded.item()): + raise RuntimeError(f"Distributed FlashInfer autotune cache could not be saved to {cache_path}.") from save_error + + def _load_distributed_cache(self, tuner: Any, cache_path: str | None) -> Exception | None: + if cache_path is not None: + if Path(cache_path).is_file(): + try: + if not tuner.load_configs(cache_path): + return RuntimeError("FlashInfer autotune cache metadata does not match this runtime; choose a new cache path or remove the incompatible cache.") + except Exception as error: + return error + elif self.options.mode == "load": + return FileNotFoundError(f"FlashInfer autotune cache does not exist for load mode: {cache_path}") + elif self.options.mode == "load": + return ValueError("FlashInfer load mode requires a cache path.") + return None + + def _save_distributed_cache(self, tuner: Any, cache_path: str | None) -> Exception | None: + if not self.options.tune_mode or cache_path is None or not self.distributed_context.is_cache_writer: + return None + try: + tuner.save_configs(cache_path) + except Exception as error: + logger.exception(f"Failed to persist distributed FlashInfer autotune cache: {cache_path}") + return error + return None 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..8d90732b3 --- /dev/null +++ b/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py @@ -0,0 +1,1361 @@ +import importlib +import os +import sys +from dataclasses import dataclass +from pathlib import Path + +import torch +import torch.distributed as dist +from PIL import Image +from loguru import logger +from safetensors import safe_open +from tqdm.auto import tqdm +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.runners.hunyuan_image3.flashinfer_autotune import ( + DistributedAutotuneContext, + FlashInferAutotuneController, +) +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"] + + # self.vae_decoder is a complete VAE model (encoder + decoder). T2I only + # decodes on rank 0, so other SP ranks do not need this extra allocation. + load_vae_on_this_rank = not ( + self.config.get("task") == "t2i" + and self._sequence_parallel_enabled() + and not self._is_output_rank() + ) + if load_vae_on_this_rank: + self.vae_decoder = self._load_vae_decoder() + else: + self.vae_decoder = None + logger.info("Skipping HunyuanImage3 VAE load on non-output SP rank for T2I.") + 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) + full_attn_slices = self._build_full_attn_slices(tokenizer_output, batch, seq_len=seq_len) + for batch_idx in range(batch): + for start, stop in full_attn_slices[batch_idx]: + attention_mask[batch_idx, start:stop, start:stop] = True + return attention_mask.unsqueeze(1) + + def _build_full_attn_slices(self, tokenizer_output, batch, seq_len=None): + batch_slices = [] + for batch_idx in range(batch): + sample_slices = [] + for image_slice in self.hunyuan_image_processor.prepare_full_attn_slices(tokenizer_output, batch_idx): + start = 0 if image_slice.start is None else int(image_slice.start) + stop = seq_len if image_slice.stop is None else int(image_slice.stop) + if seq_len is not None: + start = max(0, min(start, seq_len)) + stop = max(0, min(stop, seq_len)) + if stop > start: + sample_slices.append((start, stop)) + sample_slices.sort() + batch_slices.append(sample_slices) + return batch_slices + + 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_flashinfer_autotune_controller(self): + sequence_parallel_active = self._sequence_parallel_enabled() + distributed_context = DistributedAutotuneContext( + coordination_required=sequence_parallel_active, + process_group=self._parallel_control_group() if sequence_parallel_active else None, + status_device_resolver=self._pipeline_latent_device if sequence_parallel_active else None, + is_cache_writer=self._is_output_rank() if sequence_parallel_active else False, + ) + return FlashInferAutotuneController.from_config( + config=self.config, + distributed_context=distributed_context, + ) + + def _sequence_parallel_enabled(self): + return bool( + self.config.get("seq_parallel", False) + and dist.is_available() + and dist.is_initialized() + and getattr(self.model, "seq_p_group", None) is not None + ) + + def _parallel_control_group(self): + """Return the group used to keep replicated runner state identical. + + Attention collectives remain scoped to ``seq_p_group`` and CFG prediction + exchange remains scoped to ``cfg_p_group``. Runner state such as sampled + tokens and latents is synchronized only for sequence-parallel runs. + Hybrid runs use WORLD so the state also spans the CFG dimension. Pure + CFG deliberately returns no runner control group: its only collective + is the prediction all-gather in the model, matching the known-good CFG + implementation and avoiding cfg_p collectives on both ends of a + multi-device pipeline. + """ + if not self._sequence_parallel_enabled(): + return None + if self._cfg_parallel_enabled(): + return dist.group.WORLD + return self.model.seq_p_group + + def _is_output_rank(self): + return not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0 + + def _broadcast_parallel_tensor(self, tensor): + group = self._parallel_control_group() + if group is None: + return tensor + source_rank = dist.get_global_rank(group, 0) + dist.broadcast(tensor, src=source_rank, group=group) + return tensor + + def _parallel_barrier(self): + group = self._parallel_control_group() + if group is None: + return + if dist.get_backend(group) == "nccl": + barrier_device = self._pipeline_latent_device() + torch.cuda.set_device(barrier_device) + device_index = barrier_device.index + if device_index is None: + device_index = torch.cuda.current_device() + dist.barrier(group=group, device_ids=[device_index]) + else: + dist.barrier(group=group) + + 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 _cfg_parallel_enabled(self): + return bool(self.config.get("cfg_parallel", False)) and bool(self.config.get("enable_cfg", False)) + + def _is_distributed_nonzero_rank(self): + return dist.is_available() and dist.is_initialized() and dist.get_rank() != 0 + + def _cfg_parallel_rank(self): + if not self._cfg_parallel_enabled(): + return 0 + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("HunyuanImage3 cfg_parallel requires torch.distributed to be initialized.") + cfg_p_group = self.config["device_mesh"].get_group(mesh_dim="cfg_p") + assert dist.get_world_size(cfg_p_group) == 2, "cfg_p_world_size must be equal to 2" + return dist.get_rank(cfg_p_group) + + def _slice_cfg_branch_value(self, value, branch_idx, cfg_batch_size): + if isinstance(value, torch.Tensor): + if value.ndim > 0 and value.shape[0] == cfg_batch_size: + return value[branch_idx : branch_idx + 1] + return value + if isinstance(value, list): + if len(value) == cfg_batch_size: + return [value[branch_idx]] + return [self._slice_cfg_branch_value(item, branch_idx, cfg_batch_size) for item in value] + if isinstance(value, tuple): + if len(value) == cfg_batch_size: + return (value[branch_idx],) + return tuple(self._slice_cfg_branch_value(item, branch_idx, cfg_batch_size) for item in value) + if isinstance(value, dict): + return {k: self._slice_cfg_branch_value(v, branch_idx, cfg_batch_size) for k, v in value.items()} + return value + + def _resolve_denoise_cfg_mode(self, prepared_inputs): + if not prepared_inputs.get("do_cfg", False): + return "none" + + mode = str(self.config.get("hunyuan_cfg_mode", "batch")).lower() + if mode not in ("batch", "serial", "parallel"): + raise ValueError(f"HunyuanImage3 hunyuan_cfg_mode must be one of batch/serial/parallel, got {mode!r}.") + if mode == "parallel" and not self._cfg_parallel_enabled(): + raise ValueError( + "HunyuanImage3 hunyuan_cfg_mode='parallel' requires enable_cfg=true and " + "a distributed config with parallel.cfg_p_size=2." + ) + return mode + + def _prepare_cfg_parallel_branch_inputs(self, prepared_inputs, cfg_p_rank, mark_parallel_branch=True): + cfg_batch_size = int(prepared_inputs["input_ids"].shape[0]) + if cfg_batch_size != 2: + raise ValueError(f"HunyuanImage3 CFG branch inputs expect packed cond/uncond batch size 2, got {cfg_batch_size}.") + + branch_inputs = {} + for key, value in prepared_inputs.items(): + if key in ("custom_pos_emb",): + continue + branch_inputs[key] = self._slice_cfg_branch_value(value, cfg_p_rank, cfg_batch_size) + + branch_inputs["batch_size"] = 1 + branch_inputs["do_cfg"] = False + if mark_parallel_branch: + branch_inputs["_cfg_parallel_branch"] = True + branch_inputs["custom_pos_emb"] = self._build_custom_pos_emb(branch_inputs["position_ids"], branch_inputs.get("rope_image_info")) + return branch_inputs + + def _build_denoise_kv_state(self, prepared_inputs, use_kv_cache): + if not use_kv_cache: + return { + "kv_cache": None, + "cache_position_ids": None, + "cache_local_inputs": None, + } + cache_position_ids = self._build_denoise_cache_position_ids(prepared_inputs) + return { + "kv_cache": HunyuanImage3StaticKVCache( + num_layers=self._hunyuan_num_layers(), + max_cache_len=prepared_inputs["input_ids"].shape[1], + ), + "cache_position_ids": cache_position_ids, + "cache_local_inputs": self._build_denoise_cache_local_indices(prepared_inputs, cache_position_ids), + } + + def _build_denoise_model_inputs(self, prepared_inputs, latent_model_input, timestep, step_index, use_kv_cache, kv_state, guidance_scale): + 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"], + "full_attn_slices": prepared_inputs.get("full_attn_slices"), + "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: + cache_position_ids = kv_state["cache_position_ids"] + cache_local_inputs = kv_state["cache_local_inputs"] + model_inputs = { + "input_ids": None, + "attention_mask": self._slice_denoise_cache_attention_mask(prepared_inputs["attention_mask"], cache_position_ids), + "full_attn_slices": prepared_inputs.get("full_attn_slices"), + "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 prepared_inputs.get("_cfg_parallel_branch", False): + model_inputs["_cfg_parallel_branch"] = True + if use_kv_cache: + model_inputs["past_key_values"] = kv_state["kv_cache"] + model_inputs["use_cache"] = True + if self.config.get("cfg_distilled", False): + model_inputs["guidance"] = torch.tensor([1000.0 * guidance_scale], device=latent_model_input.device, dtype=torch.bfloat16) + if self.config.get("use_meanflow", False): + raise NotImplementedError("HunyuanImage3 native meanflow sampling is not implemented yet.") + return model_inputs + + 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) + full_attn_slices = self._build_full_attn_slices(tokenizer_output, input_ids.shape[0], seq_len=input_ids.shape[1]) if attention_mask is not None else None + model_inputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + "custom_pos_emb": custom_pos_emb, + "full_attn_slices": full_attn_slices, + } + 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=""): + if self._is_output_rank(): + 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 = self._broadcast_parallel_tensor(next_token) + 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) + full_attn_slices = self._build_full_attn_slices(output, input_ids.shape[0], seq_len=input_ids.shape[1]) + + 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), + "full_attn_slices": full_attn_slices, + "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): + cfg_mode = self._resolve_denoise_cfg_mode(prepared_inputs) + serial_branch_inputs = None + if cfg_mode == "parallel": + prepared_inputs = self._prepare_cfg_parallel_branch_inputs(prepared_inputs, self._cfg_parallel_rank(), mark_parallel_branch=True) + elif cfg_mode == "serial": + serial_branch_inputs = [ + self._prepare_cfg_parallel_branch_inputs(prepared_inputs, 0, mark_parallel_branch=False), + self._prepare_cfg_parallel_branch_inputs(prepared_inputs, 1, mark_parallel_branch=False), + ] + + latents = self._prepare_latents(prepared_inputs["batch_size"], image_size, prepared_inputs["generator"]) + latents = self._broadcast_parallel_tensor(latents) + 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 cfg_mode in ("serial", "parallel"): + raise NotImplementedError("HunyuanImage3 Taylor cache currently supports hunyuan_cfg_mode='batch' only.") + 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}") + + if cfg_mode == "serial": + kv_states = [self._build_denoise_kv_state(branch_inputs, use_kv_cache) for branch_inputs in serial_branch_inputs] + else: + kv_states = [self._build_denoise_kv_state(prepared_inputs, use_kv_cache)] + if hasattr(self.model, "transformer_infer") and hasattr(self.model.transformer_infer, "reset_moe_profile"): + self.model.transformer_infer.reset_moe_profile() + + denoise_steps = tqdm( + self.scheduler.timesteps, + total=len(self.scheduler.timesteps), + desc="HunyuanImage3 denoise", + dynamic_ncols=True, + disable=bool(self.config.get("disable_progress_bar", False)) or not self._is_output_rank(), + ) + autotune_controller = self._build_flashinfer_autotune_controller() + with autotune_controller.context(): + for step_index, timestep in enumerate(denoise_steps): + # ==================== CFG serial Processing ==================== + if cfg_mode == "serial": + cond_model_inputs = self._build_denoise_model_inputs( + serial_branch_inputs[0], + latents, + timestep, + step_index, + use_kv_cache, + kv_states[0], + guidance_scale, + ) + uncond_model_inputs = self._build_denoise_model_inputs( + serial_branch_inputs[1], + latents, + timestep, + step_index, + use_kv_cache, + kv_states[1], + guidance_scale, + ) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=latents.device.type == "cuda"): + prediction = self.model.infer_cfg_serial(cond_model_inputs, uncond_model_inputs)["diffusion_prediction"].to(latents.device) + + elif cfg_mode == "batch": + # ==================== CFG Batch Processing ==================== + # prepared_inputs 中是 packed CFG batch： + # batch[0] = conditional + # batch[1] = unconditional + # + # latents 原本是 batch=1，需要复制为 batch=2， + # 让 cond/uncond 在一次 transformer forward 中完成。 + latent_model_input = torch.cat([latents, latents], dim=0) + + model_inputs = self._build_denoise_model_inputs( + prepared_inputs, + latent_model_input, + timestep, + step_index, + use_kv_cache, + kv_states[0], + guidance_scale, + ) + + if taylor_cache_dic is not None: + taylor_cache_dic["current_step"] = step_index + model_inputs["cache_dic"] = taylor_cache_dic + + 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) + + # Transformer 返回 batch=2 prediction， + # 在 runner 中拆分并完成 CFG guidance。 + pred_cond, pred_uncond = prediction.chunk(2, dim=0) + + if hasattr(self.model, "combine_cfg_predictions"): + prediction = self.model.combine_cfg_predictions( + pred_cond, + pred_uncond, + ) + else: + prediction = pred_uncond + guidance_scale * ( + pred_cond - pred_uncond + ) + + else: + # ==================== CFG Parallel Processing ==================== + # cfg_mode == "parallel" 时： + # prepared_inputs 已经按照 cfg_p_rank 切成 batch=1 + # cfg rank 0 负责 conditional + # cfg rank 1 负责 unconditional + # + # cfg_mode == "none" 也会复用此单 batch forward 路径。 + latent_model_input = latents + + model_inputs = self._build_denoise_model_inputs( + prepared_inputs, + latent_model_input, + timestep, + step_index, + use_kv_cache, + kv_states[0], + guidance_scale, + ) + + # parallel 模式前面已经禁止 Taylor cache； + # 这里保留判断是为了兼容 cfg_mode == "none"。 + if taylor_cache_dic is not None: + taylor_cache_dic["current_step"] = step_index + model_inputs["cache_dic"] = taylor_cache_dic + + 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) + + # CFG parallel 不需要在 runner 中 chunk/combine。 + # model.infer() 内部已经通过 cfg_p_group： + # 1. all-gather cond/uncond prediction + # 2. 完成 CFG guidance + # 最终返回的 prediction 已经是 guided prediction。 + + prediction = self._broadcast_parallel_tensor(prediction) + 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) + + self._parallel_barrier() + if hasattr(self.model, "transformer_infer") and hasattr(self.model.transformer_infer, "print_moe_profile"): + self.model.transformer_infer.print_moe_profile(reset=True) + + 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) + if not self._is_output_rank(): + return [] + 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) + if not self._is_output_rank(): + return [] + 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 not self._is_output_rank(): + return {"image": None} + 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 97155fa10..746f7acb7 100755 --- a/lightx2v/utils/set_config.py +++ b/lightx2v/utils/set_config.py @@ -39,6 +39,40 @@ 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", + "hunyuan_cfg_mode", + "attn_impl", + "flashinfer_autotune", + "flashinfer_autotune_mode", + "flashinfer_autotune_cache", + "flashinfer_tuning_buckets", + "flashinfer_autotune_round_up", + "flashinfer_tune_max_num_tokens", + "hunyuan_sp_size", + "hunyuan_sp_attn_type", + "hunyuan_image3_pipeline_layout", + ) + 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. @@ -228,6 +262,123 @@ 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("attn_impl", "torch_sdpa") + config.setdefault("hunyuan_cfg_mode", "batch") + 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 + + sp_size_override = config.pop("hunyuan_sp_size", None) + sp_attn_override = config.pop("hunyuan_sp_attn_type", None) + parallel_config = config.get("parallel") + if sp_size_override is not None: + sp_size = int(sp_size_override) + if sp_size < 1: + raise ValueError(f"hunyuan_sp_size must be >= 1, got {sp_size}.") + if sp_size == 1: + # Only collapse the sequence-parallel dimension. A hybrid + # CFG+SP config may still need its cfg_p_size=2 process mesh + # when SP is disabled from the CLI. + if isinstance(parallel_config, dict) and int(parallel_config.get("cfg_p_size", 1)) > 1: + parallel_config = dict(parallel_config) + parallel_config["seq_p_size"] = 1 + parallel_config.pop("seq_p_attn_type", None) + config["parallel"] = parallel_config + else: + config["parallel"] = False + parallel_config = None + else: + parallel_config = dict(parallel_config) if isinstance(parallel_config, dict) else {} + parallel_config["seq_p_size"] = sp_size + config["parallel"] = parallel_config + + if isinstance(parallel_config, dict): + parallel_config = dict(parallel_config) + cfg_p_size = int(parallel_config.get("cfg_p_size", 1)) + seq_p_size = int(parallel_config.get("seq_p_size", 1)) + if cfg_p_size not in (1, 2): + raise ValueError(f"HunyuanImage3 parallel.cfg_p_size must be 1 or 2, got {cfg_p_size}.") + if seq_p_size < 1: + raise ValueError(f"HunyuanImage3 parallel.seq_p_size must be >= 1, got {seq_p_size}.") + if cfg_p_size == 2 and not config.get("enable_cfg", False): + raise ValueError("HunyuanImage3 parallel.cfg_p_size=2 requires enable_cfg=true.") + + parallel_config["cfg_p_size"] = cfg_p_size + parallel_config["seq_p_size"] = seq_p_size + + if seq_p_size > 1: + attn_type = sp_attn_override or parallel_config.get("seq_p_attn_type", "kv_all_gather") + attn_type = str(attn_type).strip().lower().replace("-", "_") + if attn_type in ("kv_allgather", "kv_gather"): + attn_type = "kv_all_gather" + if attn_type not in ("kv_all_gather", "ulysses"): + raise ValueError( + "HunyuanImage3 sequence parallel attention must be 'kv_all_gather' or 'ulysses', " + f"got {attn_type!r}." + ) + parallel_config["seq_p_attn_type"] = attn_type + + config["parallel"] = parallel_config + + cfg_mode = str(config.get("hunyuan_cfg_mode", "batch")).strip().lower() + if cfg_p_size == 2 and cfg_mode != "parallel": + raise ValueError("HunyuanImage3 parallel.cfg_p_size=2 requires hunyuan_cfg_mode='parallel'.") + if cfg_p_size == 1 and seq_p_size > 1 and config.get("enable_cfg", False) and cfg_mode != "serial": + raise ValueError("HunyuanImage3 sequence parallel with cfg_p_size=1 requires hunyuan_cfg_mode='serial'.") + + if seq_p_size > 1 and parallel_config["seq_p_attn_type"] == "ulysses": + q_heads = int(config["num_attention_heads"]) + kv_heads = int(config.get("num_key_value_heads") or q_heads) + if q_heads % seq_p_size or kv_heads % seq_p_size: + raise ValueError( + "HunyuanImage3 Ulysses requires seq_p_size to divide Q and KV heads: " + f"Q={q_heads}, KV={kv_heads}, seq_p_size={seq_p_size}." + ) + + config.pop("_hunyuan_image3_cli_cache_snapshot", None) + config.pop("_hunyuan_image3_cli_native_snapshot", None) + if config["task"] in ["i2v", "t2av", "i2av", "i2va", "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.") @@ -295,7 +446,11 @@ def set_parallel_config(config): config["cfg_parallel"] = True # warmup dist - _a = torch.zeros([1]).to(f"{AI_DEVICE}:{dist.get_rank()}") + if AI_DEVICE == "cuda": + warmup_device = f"{AI_DEVICE}:{torch.cuda.current_device()}" + else: + warmup_device = AI_DEVICE + _a = torch.zeros([1], device=warmup_device) dist.all_reduce(_a) diff --git a/scripts/dist_infer/run_hunyuan_image3_t2i_dist_cfg.sh b/scripts/dist_infer/run_hunyuan_image3_t2i_dist_cfg.sh new file mode 100755 index 000000000..9badf1642 --- /dev/null +++ b/scripts/dist_infer/run_hunyuan_image3_t2i_dist_cfg.sh @@ -0,0 +1,66 @@ +#!/bin/bash +set -e + +# 用法： +# bash run_hunyuan_image3_t2i_dist_cfg.sh 0,1,2,3 +# +# 该脚本是 HunyuanImage3 t2i 的 CFG parallel 入口，结构参考 WAN: +# 普通 t2i 使用 run_hunyuan_image3_t2i.sh；CFG parallel 使用本脚本和 configs/dist_infer 下的 dist cfg。 + +GPU_IDS="${1:-${CUDA_VISIBLE_DEVICES:-}}" + +export lightx2v_path="${lightx2v_path:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" +export model_path="${model_path:-/path/to/HunyuanImage-3.0-model}" +export HUNYUAN_IMAGE3_REPO_PATH="${HUNYUAN_IMAGE3_REPO_PATH:-/path/to/HunyuanImage-3.0-repo}" +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 "nproc_per_node=${nproc_per_node:-2}" +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 "hunyuan_cfg_mode=${hunyuan_cfg_mode:-parallel}" +echo "moe_impl=${moe_impl:-flashinfer}" +echo "flashinfer_autotune_mode=${flashinfer_autotune_mode:-tune}" +echo "flashinfer_autotune_cache=${flashinfer_autotune_cache:-${lightx2v_path}/save_results/hunyuan_image3_flashinfer_autotune_t2i_dist_cfg_parallel.json}" + +source ${lightx2v_path}/scripts/base/base.sh + +config_json="${lightx2v_path}/configs/dist_infer/hunyuan_image3_t2i_dist_cfg.json" + +torchrun --nproc_per_node="${nproc_per_node:-2}" -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_cfg_parallel.png" \ + --seed 42 \ + --hunyuan_cfg_mode "${hunyuan_cfg_mode:-parallel}" \ + --moe_impl "${moe_impl:-flashinfer}" \ + --attn_impl "${attn_impl:-sdpa}" \ + --flashinfer_autotune_mode "${flashinfer_autotune_mode:-tune}" \ + --flashinfer_autotune_cache "${flashinfer_autotune_cache:-${lightx2v_path}/save_results/hunyuan_image3_flashinfer_autotune_t2i_dist_cfg_parallel.json}" \ + --flashinfer_tune_max_num_tokens "${flashinfer_tune_max_num_tokens:-16384}" \ + --flashinfer_tuning_buckets "${flashinfer_tuning_buckets:-128,256,512,1024,2048,4096,8192,12288,16384}" \ + --flashinfer_autotune_round_up "${flashinfer_autotune_round_up:-true}" \ + --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/dist_infer/run_hunyuan_image3_t2i_dist_cfg_kv_all_gather.sh b/scripts/dist_infer/run_hunyuan_image3_t2i_dist_cfg_kv_all_gather.sh new file mode 100755 index 000000000..aff41921b --- /dev/null +++ b/scripts/dist_infer/run_hunyuan_image3_t2i_dist_cfg_kv_all_gather.sh @@ -0,0 +1,10 @@ +#!/bin/bash +set -euo pipefail + +# Example (CFG=2, SP=2, PP=2): +# SP_SIZE=2 bash scripts/dist_infer/run_hunyuan_image3_t2i_dist_cfg_kv_all_gather.sh 0,1,2,3,4,5,6,7 + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +export CFG_SIZE=2 +export SP_ATTN_TYPE=kv_all_gather +exec bash "${SCRIPT_DIR}/run_hunyuan_image3_t2i_dist_cfg_sp.sh" "$@" diff --git a/scripts/dist_infer/run_hunyuan_image3_t2i_dist_cfg_sp.sh b/scripts/dist_infer/run_hunyuan_image3_t2i_dist_cfg_sp.sh new file mode 100755 index 000000000..9f38168c8 --- /dev/null +++ b/scripts/dist_infer/run_hunyuan_image3_t2i_dist_cfg_sp.sh @@ -0,0 +1,140 @@ +#!/bin/bash +set -euo pipefail + +# Common HunyuanImage3 T2I launcher for CFG parallel + sequence parallel. +# A logical torchrun rank owns one interleaved pipeline lane: +# world_size = CFG_SIZE * SP_SIZE +# pipeline_size = visible_gpu_count / world_size + +GPU_IDS="${1:-${CUDA_VISIBLE_DEVICES:-}}" +CFG_SIZE="${CFG_SIZE:-2}" +SP_SIZE="${SP_SIZE:-2}" +SP_ATTN_TYPE="${SP_ATTN_TYPE:-ulysses}" + +export lightx2v_path="${lightx2v_path:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" +export model_path="${model_path:-/path/to/HunyuanImage-3.0-model}" +export HUNYUAN_IMAGE3_REPO_PATH="${HUNYUAN_IMAGE3_REPO_PATH:-/path/to/HunyuanImage-3.0-repo}" +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 [ "$CFG_SIZE" != "2" ]; then + echo "HunyuanImage3 CFG parallel has exactly two cond/uncond branches; CFG_SIZE must be 2, got: $CFG_SIZE" >&2 + exit 2 +fi +if ! [[ "$SP_SIZE" =~ ^[1-9][0-9]*$ ]] || [ "$SP_SIZE" -le 1 ]; then + echo "Hybrid CFG+SP requires SP_SIZE to be an integer greater than 1, got: $SP_SIZE" >&2 + exit 2 +fi +if [ "$SP_ATTN_TYPE" != "kv_all_gather" ] && [ "$SP_ATTN_TYPE" != "ulysses" ]; then + echo "SP_ATTN_TYPE must be kv_all_gather or ulysses, got: $SP_ATTN_TYPE" >&2 + exit 2 +fi + +IFS=',' read -r -a raw_visible_gpu_ids <<< "${CUDA_VISIBLE_DEVICES:-}" +visible_gpu_ids=() +declare -A seen_gpu_ids=() +for gpu_id in "${raw_visible_gpu_ids[@]}"; do + gpu_id="${gpu_id//[[:space:]]/}" + if [ -z "$gpu_id" ]; then + continue + fi + if [ -n "${seen_gpu_ids[$gpu_id]:-}" ]; then + echo "Duplicate GPU id in CUDA_VISIBLE_DEVICES: $gpu_id" >&2 + exit 2 + fi + seen_gpu_ids[$gpu_id]=1 + visible_gpu_ids+=("$gpu_id") +done +visible_gpu_count=${#visible_gpu_ids[@]} +if [ "$visible_gpu_count" -eq 0 ]; then + echo "No visible CUDA devices were found." >&2 + exit 2 +fi +export CUDA_VISIBLE_DEVICES=$(IFS=,; echo "${visible_gpu_ids[*]}") + +world_size=$((CFG_SIZE * SP_SIZE)) +if [ $((visible_gpu_count % world_size)) -ne 0 ]; then + echo "Visible GPU count ($visible_gpu_count) must be divisible by CFG_SIZE * SP_SIZE ($CFG_SIZE * $SP_SIZE = $world_size)." >&2 + exit 2 +fi +if [ "$SP_ATTN_TYPE" = "ulysses" ] && { [ $((32 % SP_SIZE)) -ne 0 ] || [ $((8 % SP_SIZE)) -ne 0 ]; }; then + echo "Ulysses SP_SIZE must divide both 32 Q heads and 8 KV heads, got: $SP_SIZE" >&2 + exit 2 +fi + +pipeline_gpus_per_lane=$((visible_gpu_count / world_size)) +if [ "$SP_ATTN_TYPE" = "ulysses" ]; then + config_json="${lightx2v_path}/configs/dist_infer/hunyuan_image3_t2i_dist_cfg_ulysses.json" + default_attn_impl="sdpa" +else + config_json="${lightx2v_path}/configs/dist_infer/hunyuan_image3_t2i_dist_cfg_kv_all_gather.json" + default_attn_impl="flash_attention_2" +fi +resolved_attn_impl="${attn_impl:-$default_attn_impl}" +default_autotune_cache="${lightx2v_path}/save_results/hunyuan_image3_flashinfer_${SP_ATTN_TYPE}_cfg${CFG_SIZE}_sp${SP_SIZE}_pp${pipeline_gpus_per_lane}_t2i.json" +resolved_autotune_cache="${flashinfer_autotune_cache:-$default_autotune_cache}" +if [ -n "${flashinfer_autotune_mode:-}" ]; then + resolved_autotune_mode="$flashinfer_autotune_mode" +elif [ "${flashinfer_autotune:-true}" = "false" ] || [ "${flashinfer_autotune:-true}" = "0" ]; then + resolved_autotune_mode="off" +else + resolved_autotune_mode="tune" +fi + +echo "Using CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" +echo "CFG_SIZE=${CFG_SIZE}, SP_SIZE=${SP_SIZE}, SP_ATTN_TYPE=${SP_ATTN_TYPE}, WORLD_SIZE=${world_size}" +echo "pipeline_gpus_per_lane=${pipeline_gpus_per_lane}, layout=interleaved" +for ((rank=0; rank/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 ! [[ "$SP_SIZE" =~ ^[1-9][0-9]*$ ]]; then + echo "SP_SIZE must be a positive integer, got: $SP_SIZE" >&2 + exit 2 +fi +if [ "$SP_ATTN_TYPE" != "kv_all_gather" ] && [ "$SP_ATTN_TYPE" != "ulysses" ]; then + echo "SP_ATTN_TYPE must be kv_all_gather or ulysses, got: $SP_ATTN_TYPE" >&2 + exit 2 +fi +IFS=',' read -r -a visible_gpu_ids <<< "${CUDA_VISIBLE_DEVICES:-}" +visible_gpu_count=${#visible_gpu_ids[@]} +if [ "$visible_gpu_count" -eq 0 ] || [ -z "${visible_gpu_ids[0]}" ]; then + echo "No visible CUDA devices were found." >&2 + exit 2 +fi +declare -A seen_gpu_ids=() +for gpu_id in "${visible_gpu_ids[@]}"; do + gpu_id="${gpu_id//[[:space:]]/}" + if [ -n "${seen_gpu_ids[$gpu_id]:-}" ]; then + echo "Duplicate GPU id in CUDA_VISIBLE_DEVICES: $gpu_id" >&2 + exit 2 + fi + seen_gpu_ids[$gpu_id]=1 +done +if [ $((visible_gpu_count % SP_SIZE)) -ne 0 ]; then + echo "Visible GPU count ($visible_gpu_count) must be divisible by SP_SIZE ($SP_SIZE)." >&2 + exit 2 +fi +if [ "$SP_ATTN_TYPE" = "ulysses" ] && { [ $((32 % SP_SIZE)) -ne 0 ] || [ $((8 % SP_SIZE)) -ne 0 ]; }; then + echo "Ulysses SP_SIZE must divide both 32 Q heads and 8 KV heads, got: $SP_SIZE" >&2 + exit 2 +fi + +pipeline_gpus_per_lane=$((visible_gpu_count / SP_SIZE)) +min_pipeline_gpus_per_lane="${MIN_PIPELINE_GPUS_PER_LANE:-3}" +if [ "$pipeline_gpus_per_lane" -lt "$min_pipeline_gpus_per_lane" ] && [ "${ALLOW_UNSAFE_SP_MEMORY_LAYOUT:-false}" != "true" ]; then + echo "Each SP lane has only ${pipeline_gpus_per_lane} GPU(s); this BF16 checkpoint expects at least ${min_pipeline_gpus_per_lane}." >&2 + echo "Set ALLOW_UNSAFE_SP_MEMORY_LAYOUT=true only for a separately validated quantized/offload setup." >&2 + exit 2 +fi + +min_free_gpu_memory_mib="${MIN_FREE_GPU_MEMORY_MIB:-50000}" +if command -v nvidia-smi >/dev/null 2>&1; then + for gpu_id in "${visible_gpu_ids[@]}"; do + free_mib=$(nvidia-smi -i "$gpu_id" --query-gpu=memory.free --format=csv,noheader,nounits 2>/dev/null | head -n 1 | tr -d '[:space:]' || true) + if [[ "$free_mib" =~ ^[0-9]+$ ]]; then + echo "GPU ${gpu_id}: free_memory=${free_mib} MiB" + if [ "$free_mib" -lt "$min_free_gpu_memory_mib" ] && [ "${ALLOW_LOW_FREE_MEMORY:-false}" != "true" ]; then + echo "GPU ${gpu_id} has only ${free_mib} MiB free; require at least ${min_free_gpu_memory_mib} MiB for this launcher." >&2 + echo "Choose an idle GPU or set ALLOW_LOW_FREE_MEMORY=true after validating memory manually." >&2 + exit 2 + fi + fi + done +fi + +if [ "$SP_ATTN_TYPE" = "ulysses" ]; then + config_json="${lightx2v_path}/configs/dist_infer/hunyuan_image3_t2i_dist_ulysses.json" +else + config_json="${lightx2v_path}/configs/dist_infer/hunyuan_image3_t2i_dist_kv_all_gather.json" +fi +default_autotune_cache="${lightx2v_path}/save_results/hunyuan_image3_flashinfer_${SP_ATTN_TYPE}_sp${SP_SIZE}_t2i.json" +resolved_autotune_cache="${flashinfer_autotune_cache:-$default_autotune_cache}" +if [ -n "${flashinfer_autotune_mode:-}" ]; then + resolved_autotune_mode="$flashinfer_autotune_mode" +elif [ "${flashinfer_autotune:-true}" = "false" ] || [ "${flashinfer_autotune:-true}" = "0" ]; then + resolved_autotune_mode="off" +else + resolved_autotune_mode="tune" +fi + +echo "Using CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" +echo "SP_SIZE=${SP_SIZE}, SP_ATTN_TYPE=${SP_ATTN_TYPE}, layout=interleaved" +for ((rank=0; rank/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 "HUNYUAN_CFG_MODE=${HUNYUAN_CFG_MODE}" +echo "moe_impl=${moe_impl:-flashinfer}" +echo "ATTN_IMPL=${ATTN_IMPL}" +echo "FLASHINFER_AUTOTUNE_MODE=${FLASHINFER_AUTOTUNE_MODE}" +echo "FLASHINFER_AUTOTUNE_CACHE=${FLASHINFER_AUTOTUNE_CACHE}" + +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 \ + --hunyuan_cfg_mode "${HUNYUAN_CFG_MODE}" \ + --moe_impl "${moe_impl:-flashinfer}" \ + --attn_impl "${ATTN_IMPL}" \ + --flashinfer_autotune_mode "${FLASHINFER_AUTOTUNE_MODE}" \ + --flashinfer_autotune_cache "${FLASHINFER_AUTOTUNE_CACHE}" \ + --flashinfer_tune_max_num_tokens "${FLASHINFER_TUNE_MAX_NUM_TOKENS}" \ + --flashinfer_tuning_buckets "${FLASHINFER_TUNING_BUCKETS}" \ + --flashinfer_autotune_round_up "${FLASHINFER_AUTOTUNE_ROUND_UP}" diff --git a/scripts/hunyuan_image3/run_hunyuan_image3_t2i.sh b/scripts/hunyuan_image3/run_hunyuan_image3_t2i.sh new file mode 100755 index 000000000..cf3da34cd --- /dev/null +++ b/scripts/hunyuan_image3/run_hunyuan_image3_t2i.sh @@ -0,0 +1,69 @@ +#!/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="${lightx2v_path:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" +export model_path="${model_path:-/path/to/HunyuanImage-3.0-model}" +export HUNYUAN_IMAGE3_REPO_PATH="${HUNYUAN_IMAGE3_REPO_PATH:-/path/to/HunyuanImage-3.0-repo}" +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 "hunyuan_cfg_mode=${hunyuan_cfg_mode:-batch}" +echo "moe_impl=${moe_impl:-flashinfer}" +echo "attn_impl=${attn_impl:-torch_sdpa}" +echo "flashinfer_autotune_mode=${flashinfer_autotune_mode:-off}" +echo "flashinfer_autotune_cache=${flashinfer_autotune_cache:-${lightx2v_path}/save_results/hunyuan_image3_flashinfer_autotune_t2i.json}" + +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 \ + --hunyuan_cfg_mode "${hunyuan_cfg_mode:-batch}" \ + --moe_impl "${moe_impl:-flashinfer}" \ + --attn_impl "${attn_impl:-sdpa}" \ + --flashinfer_autotune_mode "${flashinfer_autotune_mode:-tune}" \ + --flashinfer_autotune_cache "${flashinfer_autotune_cache:-${lightx2v_path}/save_results/hunyuan_image3_flashinfer_autotune_t2i.json}" \ + --flashinfer_tune_max_num_tokens "${flashinfer_tune_max_num_tokens:-16384}" \ + --flashinfer_tuning_buckets "${flashinfer_tuning_buckets:-128,256,512,1024,2048,4096,8192,12288,16384}" \ + --flashinfer_autotune_round_up "${flashinfer_autotune_round_up:-true}" \ + --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_t2i_dist_cfg.sh b/scripts/hunyuan_image3/run_hunyuan_image3_t2i_dist_cfg.sh new file mode 100755 index 000000000..5fcc4d45c --- /dev/null +++ b/scripts/hunyuan_image3/run_hunyuan_image3_t2i_dist_cfg.sh @@ -0,0 +1,66 @@ +#!/bin/bash +set -e + +# 用法： +# bash run_hunyuan_image3_t2i_dist_cfg.sh 0,1,2,3 +# +# 该脚本是 HunyuanImage3 t2i 的 CFG parallel 入口，结构参考 WAN: +# 普通 t2i 使用 run_hunyuan_image3_t2i.sh；CFG parallel 使用本脚本和 configs/dist_infer 下的 dist cfg。 + +GPU_IDS="${1:-${CUDA_VISIBLE_DEVICES:-}}" + +export lightx2v_path="${lightx2v_path:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" +export model_path="${model_path:-/path/to/HunyuanImage-3.0-model}" +export HUNYUAN_IMAGE3_REPO_PATH="${HUNYUAN_IMAGE3_REPO_PATH:-/path/to/HunyuanImage-3.0-repo}" +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 "nproc_per_node=${nproc_per_node:-2}" +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 "hunyuan_cfg_mode=${hunyuan_cfg_mode:-parallel}" +echo "moe_impl=${moe_impl:-flashinfer}" +echo "flashinfer_autotune_mode=${flashinfer_autotune_mode:-tune}" +echo "flashinfer_autotune_cache=${flashinfer_autotune_cache:-${lightx2v_path}/save_results/hunyuan_image3_flashinfer_autotune_t2i_dist_cfg.json}" + +source ${lightx2v_path}/scripts/base/base.sh + +config_json="${lightx2v_path}/configs/dist_infer/hunyuan_image3_t2i_dist_cfg.json" + +torchrun --nproc_per_node="${nproc_per_node:-2}" -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 \ + --hunyuan_cfg_mode "${hunyuan_cfg_mode:-parallel}" \ + --moe_impl "${moe_impl:-flashinfer}" \ + --attn_impl "${attn_impl:-sdpa}" \ + --flashinfer_autotune_mode "${flashinfer_autotune_mode:-tune}" \ + --flashinfer_autotune_cache "${flashinfer_autotune_cache:-${lightx2v_path}/save_results/hunyuan_image3_flashinfer_autotune_t2i.json}" \ + --flashinfer_tune_max_num_tokens "${flashinfer_tune_max_num_tokens:-16384}" \ + --flashinfer_tuning_buckets "${flashinfer_tuning_buckets:-128,256,512,1024,2048,4096,8192,12288,16384}" \ + --flashinfer_autotune_round_up "${flashinfer_autotune_round_up:-true}" \ + --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..5fea47023 --- /dev/null +++ b/scripts/hunyuan_image3/run_hunyuan_image3_ti2i.sh @@ -0,0 +1,95 @@ +#!/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:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" +export model_path="${model_path:-/path/to/HunyuanImage-3-Instruct}" +export HUNYUAN_IMAGE3_REPO_PATH="${HUNYUAN_IMAGE3_REPO_PATH:-/path/to/HunyuanImage-3.0-repo}" +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 "hunyuan_cfg_mode=${hunyuan_cfg_mode:-batch}" +echo "moe_impl=${moe_impl:-flashinfer}" +echo "attn_impl=${attn_impl:-torch_sdpa}" +echo "flashinfer_autotune_mode=${flashinfer_autotune_mode:-off}" +echo "flashinfer_autotune_cache=${flashinfer_autotune_cache:-${lightx2v_path}/save_results/hunyuan_image3_flashinfer_autotune_${DEMO}.json}" + +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" \ + --hunyuan_cfg_mode "${hunyuan_cfg_mode:-batch}" \ + --moe_impl "${moe_impl:-eager}" \ + --attn_impl "${attn_impl:-torch_sdpa}" \ + --flashinfer_autotune_mode "${flashinfer_autotune_mode:-tune}" \ + --flashinfer_autotune_cache "${flashinfer_autotune_cache:-${lightx2v_path}/save_results/hunyuan_image3_flashinfer_autotune_${DEMO}.json}" \ + --flashinfer_tune_max_num_tokens "${flashinfer_tune_max_num_tokens:-16384}" \ + --flashinfer_tuning_buckets "${flashinfer_tuning_buckets:-128,256,512,1024,2048,4096,8192,12288,16384}" \ + --flashinfer_autotune_round_up "${flashinfer_autotune_round_up:-true}" \ + --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}"