diff --git a/configs/hunyuan_image3/hunyuan_image3_t2i.json b/configs/hunyuan_image3/hunyuan_image3_t2i_tp_sp_cfg_flashinfer.json similarity index 67% rename from configs/hunyuan_image3/hunyuan_image3_t2i.json rename to configs/hunyuan_image3/hunyuan_image3_t2i_tp_sp_cfg_flashinfer.json index f01450c9c..bbb889f09 100644 --- a/configs/hunyuan_image3/hunyuan_image3_t2i.json +++ b/configs/hunyuan_image3/hunyuan_image3_t2i_tp_sp_cfg_flashinfer.json @@ -6,12 +6,21 @@ "target_width": 1024, "feature_caching": "NoCaching", "moe_impl": "flashinfer", - "attn_impl": "sdpa", - "hunyuan_cfg_mode": "batch", - "flashinfer_autotune_mode": "tune", - "flashinfer_autotune_cache": "save_results/hunyuan_image3_flashinfer_autotune_t2i.json", + "attn_impl": "torch_sdpa", + "flashinfer_autotune_mode": "auto", + "flashinfer_autotune_cache": "save_results/hunyuan_image3_flashinfer_autotune_t2i_tp2_sp2_cfg2.json", "flashinfer_tune_max_num_tokens": 16384, - "flashinfer_tuning_buckets": [128, 256, 512, 1024, 2048, 4096, 8192, 12288, 16384], + "flashinfer_tuning_buckets": [ + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 12288, + 16384 + ], "flashinfer_autotune_round_up": true, "enable_kv_cache": true, "enable_text_kv_cache": true, @@ -24,7 +33,14 @@ "taylor_cache_tailing_enhance_steps": 1, "taylor_cache_low_freqs_order": 2, "taylor_cache_high_freqs_order": 2, - "pipeline_parallel": true, + "parallel": { + "pipeline_parallel": false, + "tensor_p_size": 2, + "seq_p_size": 2, + "cfg_p_size": 2, + "seq_p_attn_type": "ulysses", + "cfg_mode": "parallel" + }, "enable_cfg": true, "bot_task": "think_recaption", "use_system_prompt": "en_unified", diff --git a/configs/hunyuan_image3/hunyuan_image3_t2t_tp_sp_cfg_flashinfer.json b/configs/hunyuan_image3/hunyuan_image3_t2t_tp_sp_cfg_flashinfer.json new file mode 100644 index 000000000..18056a19d --- /dev/null +++ b/configs/hunyuan_image3/hunyuan_image3_t2t_tp_sp_cfg_flashinfer.json @@ -0,0 +1,39 @@ +{ + "feature_caching": "NoCaching", + "moe_impl": "flashinfer", + "attn_impl": "torch_sdpa", + "flashinfer_autotune_mode": "auto", + "flashinfer_autotune_cache": "save_results/hunyuan_image3_flashinfer_autotune_t2t_tp2_sp2_cfg1.json", + "flashinfer_tune_max_num_tokens": 16384, + "flashinfer_tuning_buckets": [ + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 12288, + 16384 + ], + "flashinfer_autotune_round_up": true, + "enable_kv_cache": true, + "enable_text_kv_cache": true, + "parallel": { + "pipeline_parallel": false, + "tensor_p_size": 2, + "seq_p_size": 2, + "cfg_p_size": 1, + "seq_p_attn_type": "ulysses", + "cfg_mode": "batch" + }, + "enable_cfg": false, + "bot_task": "auto", + "use_system_prompt": "None", + "max_new_tokens": 2048, + "text_do_sample": true, + "text_top_k": 1024, + "text_top_p": 0.95, + "text_temperature": 0.6, + "text_stream_output": true +} diff --git a/configs/hunyuan_image3/hunyuan_image3_ti2i.json b/configs/hunyuan_image3/hunyuan_image3_ti2i.json deleted file mode 100644 index ed4ca9615..000000000 --- a/configs/hunyuan_image3/hunyuan_image3_ti2i.json +++ /dev/null @@ -1,32 +0,0 @@ -{ - "infer_steps": 50, - "sample_guide_scale": 5.0, - "flow_shift": 1.0, - "image_size": "auto", - "feature_caching": "NoCaching", - "moe_impl": "eager", - "attn_impl": "torch_sdpa", - "hunyuan_cfg_mode": "batch", - "flashinfer_autotune_mode": "off", - "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_i2i.json b/configs/hunyuan_image3/hunyuan_image3_ti2i_tp_sp_cfg_flashinfer.json similarity index 69% rename from configs/hunyuan_image3/hunyuan_image3_i2i.json rename to configs/hunyuan_image3/hunyuan_image3_ti2i_tp_sp_cfg_flashinfer.json index d6641152b..40a7e4ddb 100644 --- a/configs/hunyuan_image3/hunyuan_image3_i2i.json +++ b/configs/hunyuan_image3/hunyuan_image3_ti2i_tp_sp_cfg_flashinfer.json @@ -6,11 +6,20 @@ "feature_caching": "NoCaching", "moe_impl": "flashinfer", "attn_impl": "torch_sdpa", - "hunyuan_cfg_mode": "batch", - "flashinfer_autotune_mode": "tune", - "flashinfer_autotune_cache": "save_results/hunyuan_image3_flashinfer_autotune_i2i.json", + "flashinfer_autotune_mode": "auto", + "flashinfer_autotune_cache": "save_results/hunyuan_image3_flashinfer_autotune_ti2i_tp2_sp2_cfg2.json", "flashinfer_tune_max_num_tokens": 16384, - "flashinfer_tuning_buckets": [128, 256, 512, 1024, 2048, 4096, 8192, 12288, 16384], + "flashinfer_tuning_buckets": [ + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 12288, + 16384 + ], "flashinfer_autotune_round_up": true, "enable_kv_cache": true, "enable_text_kv_cache": true, @@ -23,7 +32,14 @@ "taylor_cache_tailing_enhance_steps": 1, "taylor_cache_low_freqs_order": 2, "taylor_cache_high_freqs_order": 2, - "pipeline_parallel": true, + "parallel": { + "pipeline_parallel": false, + "tensor_p_size": 2, + "seq_p_size": 2, + "cfg_p_size": 2, + "seq_p_attn_type": "ulysses", + "cfg_mode": "parallel" + }, "enable_cfg": true, "infer_align_image_size": true, "bot_task": "think_recaption", diff --git a/configs/hunyuan_image3/hunyuan_image3_ti2t_tp_sp_cfg_flashinfer.json b/configs/hunyuan_image3/hunyuan_image3_ti2t_tp_sp_cfg_flashinfer.json new file mode 100644 index 000000000..7e3b9f6d3 --- /dev/null +++ b/configs/hunyuan_image3/hunyuan_image3_ti2t_tp_sp_cfg_flashinfer.json @@ -0,0 +1,40 @@ +{ + "feature_caching": "NoCaching", + "moe_impl": "flashinfer", + "attn_impl": "torch_sdpa", + "flashinfer_autotune_mode": "auto", + "flashinfer_autotune_cache": "save_results/hunyuan_image3_flashinfer_autotune_ti2t_tp2_sp2_cfg1.json", + "flashinfer_tune_max_num_tokens": 16384, + "flashinfer_tuning_buckets": [ + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 12288, + 16384 + ], + "flashinfer_autotune_round_up": true, + "enable_kv_cache": true, + "enable_text_kv_cache": true, + "parallel": { + "pipeline_parallel": false, + "tensor_p_size": 2, + "seq_p_size": 2, + "cfg_p_size": 1, + "seq_p_attn_type": "ulysses", + "cfg_mode": "batch" + }, + "enable_cfg": false, + "infer_align_image_size": true, + "bot_task": "auto", + "use_system_prompt": "None", + "max_new_tokens": 2048, + "text_do_sample": true, + "text_top_k": 1024, + "text_top_p": 0.95, + "text_temperature": 0.6, + "text_stream_output": true +} diff --git a/lightx2v/common/ops/mm/mm_weight.py b/lightx2v/common/ops/mm/mm_weight.py index be61890e9..51ed835a8 100755 --- a/lightx2v/common/ops/mm/mm_weight.py +++ b/lightx2v/common/ops/mm/mm_weight.py @@ -2444,6 +2444,7 @@ def __init__( is_post_adapter=False, lora_prefix="diffusion_model.blocks", lora_path="", + reduce_output=True, ): super().__init__( weight_name, @@ -2460,6 +2461,7 @@ def __init__( self.tp_rank = tp_rank self.tp_size = tp_size self.split_dim = split_dim # "col" for column split, "row" for row split + self.reduce_output = reduce_output assert split_dim in ["col", "row"], f"split_dim must be 'col' or 'row', got {split_dim}" self._mm = MM_WEIGHT_REGISTER.get(mm_type, MMWeight)( @@ -2497,7 +2499,7 @@ def apply(self, input_tensor): output = self._mm.apply(input_tensor) # For row split, need all-reduce to combine results from all ranks - if self.split_dim == "row" and self.tp_size > 1 and self.tp_group is not None: + if self.split_dim == "row" and self.reduce_output and self.tp_size > 1 and self.tp_group is not None: dist.all_reduce(output, op=dist.ReduceOp.SUM, group=self.tp_group) # Add bias after all-reduce (bias is not split for row split) if self._row_split_bias is not None: diff --git a/lightx2v/disagg/examples/run_controller.py b/lightx2v/disagg/examples/run_controller.py index c206af60d..a3df941b3 100644 --- a/lightx2v/disagg/examples/run_controller.py +++ b/lightx2v/disagg/examples/run_controller.py @@ -32,12 +32,9 @@ def _parallel_world_size_from_config(config: dict[str, Any] | None) -> int: return 1 tensor_p_size = int(parallel.get("tensor_p_size", 1) or 1) - if tensor_p_size > 1: - return tensor_p_size - cfg_p_size = int(parallel.get("cfg_p_size", 1) or 1) seq_p_size = int(parallel.get("seq_p_size", 1) or 1) - return max(1, cfg_p_size * seq_p_size) + return max(1, tensor_p_size * cfg_p_size * seq_p_size) def _load_base_config_json(path: str) -> dict[str, Any]: diff --git a/lightx2v/infer.py b/lightx2v/infer.py index 61464c59f..6d85827ec 100755 --- a/lightx2v/infer.py +++ b/lightx2v/infer.py @@ -138,7 +138,10 @@ def main(): ) parser.add_argument( - "--task", type=str, choices=["t2v", "i2v", "t2i", "i2i", "flf2v", "vace", "animate", "s2v", "rs2v", "t2av", "i2av", "i2va", "v2av", "ltx2_s2v", "sr", "recon", "i23d"], default="t2v" + "--task", + type=str, + choices=["t2v", "i2v", "t2t", "t2i", "ti2t", "ti2i", "i2i", "flf2v", "vace", "animate", "s2v", "rs2v", "t2av", "i2av", "i2va", "v2av", "ltx2_s2v", "sr", "recon", "i23d"], + default="t2v", ) 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) @@ -152,7 +155,7 @@ def main(): "--image_path", type=str, default="", - help="The path to input image file(s) for image-to-video (i2v) or image-to-audio-video (i2av) task. Multiple paths should be comma-separated. Example: 'path1.jpg,path2.jpg'", + help="The path to input image file(s), including HunyuanImage3 ti2t/ti2i. Multiple paths should be comma-separated. Example: 'path1.jpg,path2.jpg'", ) parser.add_argument("--state_path", type=str, default="", help="The path to input robot state file for robot i2v/i2va inference.") parser.add_argument("--last_frame_path", type=str, default="", help="The path to last frame file for first-last-frame-to-video (flf2v) task") diff --git a/lightx2v/models/networks/hunyuan_image3/infer/post_infer.py b/lightx2v/models/networks/hunyuan_image3/infer/post_infer.py index 8161fedcb..aaecca30f 100644 --- a/lightx2v/models/networks/hunyuan_image3/infer/post_infer.py +++ b/lightx2v/models/networks/hunyuan_image3/infer/post_infer.py @@ -1,4 +1,5 @@ import torch +import torch.distributed as dist 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 @@ -7,6 +8,12 @@ class HunyuanImage3PostInfer: def __init__(self, config): self.config = config + if config.get("tensor_parallel", False): + self.tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p") + self.tp_size = dist.get_world_size(self.tp_group) + else: + self.tp_group = None + self.tp_size = 1 def set_scheduler(self, scheduler): self.scheduler = scheduler @@ -52,6 +59,15 @@ def infer(self, weights, hidden_states, pre_infer_out): 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) + # Autoregressive generation only consumes the last position. Limiting + # the TP vocabulary projection to that token avoids materializing and + # gathering prompt-length logits on every rank. + logits_hidden_states = hidden_states[:, -1:, :] if self.tp_size > 1 else hidden_states + normed = weights.final_norm.apply(logits_hidden_states) logits = apply_linear(weights.lm_head, normed.reshape(-1, normed.shape[-1])).reshape(*normed.shape[:-1], -1) + if self.tp_size > 1: + local_logits = logits.contiguous() + gathered_logits = [torch.empty_like(local_logits) for _ in range(self.tp_size)] + dist.all_gather(gathered_logits, local_logits, group=self.tp_group) + logits = torch.cat(gathered_logits, dim=-1) return {"logits": logits} diff --git a/lightx2v/models/networks/hunyuan_image3/infer/transformer_infer.py b/lightx2v/models/networks/hunyuan_image3/infer/transformer_infer.py index a9440522a..b36cd9988 100644 --- a/lightx2v/models/networks/hunyuan_image3/infer/transformer_infer.py +++ b/lightx2v/models/networks/hunyuan_image3/infer/transformer_infer.py @@ -37,10 +37,20 @@ def __init__(self, config): self.config = config self.num_layers = int(config.get("num_layers") or config["num_hidden_layers"]) self.hidden_size = config["hidden_size"] - self.num_heads = int(config.get("num_attention_heads") or config["num_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.global_num_heads = int(config.get("num_attention_heads") or config["num_heads"]) + self.global_num_key_value_heads = int(config.get("num_key_value_heads") or self.global_num_heads) + self.num_key_value_groups = self.global_num_heads // self.global_num_key_value_heads + self.head_dim = config.get("attention_head_dim", self.hidden_size // self.global_num_heads) + if config.get("tensor_parallel", False): + self.tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p") + self.tp_rank = dist.get_rank(self.tp_group) + self.tp_size = dist.get_world_size(self.tp_group) + else: + self.tp_group = None + self.tp_rank = 0 + self.tp_size = 1 + self.num_heads = self.global_num_heads // self.tp_size + self.num_key_value_heads = self.global_num_key_value_heads // self.tp_size 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")) @@ -551,7 +561,7 @@ def _sequence_parallel_gather_kv(self, key_states, value_states, state): 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.") + raise ValueError("HunyuanImage3 Ulysses expects batch size 1; use parallel.cfg_mode='serial' when cfg_p_size=1 (including TP+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) @@ -577,11 +587,15 @@ def infer_mlp(self, phase, hidden_states): 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": + output = self._infer_mlp_flashinfer(moe, hidden_states) + elif moe_impl == "eager": + output = self._infer_mlp_eager(moe, hidden_states) + else: raise ValueError(f"Unsupported HunyuanImage3 moe_impl={moe_impl!r}. Expected 'eager' or 'flashinfer'.") - return self._infer_mlp_eager(moe, hidden_states) + if self.tp_size > 1: + dist.all_reduce(output, op=dist.ReduceOp.SUM, group=self.tp_group) + return output def _moe_easy_topk(self, moe, hidden_states): flat = hidden_states.reshape(-1, hidden_states.shape[-1]) @@ -634,6 +648,8 @@ def _infer_mlp_flashinfer(self, moe, hidden_states): compute_dtype, output=combined_output, quant_scales=None, + tp_size=self.tp_size, + tp_rank=self.tp_rank, tune_max_num_tokens=self.flashinfer_tune_max_num_tokens, ) output = combined_output.reshape_as(hidden_states).to(original_dtype) diff --git a/lightx2v/models/networks/hunyuan_image3/infer/utils.py b/lightx2v/models/networks/hunyuan_image3/infer/utils.py index 9a70b395f..7387fd1f7 100644 --- a/lightx2v/models/networks/hunyuan_image3/infer/utils.py +++ b/lightx2v/models/networks/hunyuan_image3/infer/utils.py @@ -18,7 +18,13 @@ 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) + tensor = getattr(linear, f"pin_{attr}", None) + if tensor is not None: + return tensor + wrapped_linear = getattr(linear, "_mm", None) + if wrapped_linear is not None and wrapped_linear is not linear: + return _linear_tensor_attr(wrapped_linear, attr) + return None def match_linear_dtype(input_tensor, linear): @@ -67,8 +73,8 @@ def repeat_kv(hidden_states, n_rep): def weight_device(module): - for attr in ("weight", "pin_weight", "bias", "pin_bias"): - tensor = getattr(module, attr, None) + for attr in ("weight", "bias"): + tensor = _linear_tensor_attr(module, attr) if tensor is not None: return tensor.device return None diff --git a/lightx2v/models/networks/hunyuan_image3/model.py b/lightx2v/models/networks/hunyuan_image3/model.py index 9c19b9039..0c86a2a3f 100644 --- a/lightx2v/models/networks/hunyuan_image3/model.py +++ b/lightx2v/models/networks/hunyuan_image3/model.py @@ -25,7 +25,7 @@ def _normalize_device_name(device): if device.isdigit(): return f"cuda:{device}" if device == "cuda": - return "cuda:0" + return f"cuda:{torch.cuda.current_device()}" return device @@ -42,9 +42,10 @@ def _resolve_sequence_parallel_pipeline_lane(config, devices): if len(devices) % parallel_world_size: raise ValueError(f"HunyuanImage3 sequence parallel requires pipeline device count ({len(devices)}) to be divisible by cfg_p_size * seq_p_size ({parallel_world_size}).") - layout = str(config.get("hunyuan_image3_pipeline_layout", "interleaved")).strip().lower() + parallel = config.get("parallel") or {} + layout = str(parallel.get("pipeline_layout", config.get("hunyuan_image3_pipeline_layout", "interleaved"))).strip().lower() if layout != "interleaved": - raise ValueError(f"HunyuanImage3 sequence parallel uses Wan-style rank/device initialization and currently requires hunyuan_image3_pipeline_layout='interleaved'; got {layout!r}.") + raise ValueError(f"HunyuanImage3 sequence parallel uses Wan-style rank/device initialization and currently requires parallel.pipeline_layout='interleaved'; got {layout!r}.") lane = devices[global_rank::parallel_world_size] expected_device = f"cuda:{torch.cuda.current_device()}" @@ -54,7 +55,8 @@ def _resolve_sequence_parallel_pipeline_lane(config, devices): def resolve_pipeline_devices(config, fallback_device): - configured = config.get("pipeline_parallel_devices") or config.get("hunyuan_image3_pipeline_devices") + parallel = config.get("parallel") or {} + configured = parallel.get("pipeline_devices") or 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()] @@ -65,7 +67,8 @@ def resolve_pipeline_devices(config, fallback_device): return _resolve_sequence_parallel_pipeline_lane(config, devices) return devices - if config.get("pipeline_parallel", True) and torch.cuda.is_available(): + pipeline_parallel = parallel.get("pipeline_parallel", config.get("pipeline_parallel", True)) + if pipeline_parallel and torch.cuda.is_available(): device_count = torch.cuda.device_count() if device_count > 0: if config.get("seq_parallel", False): @@ -114,10 +117,12 @@ class HunyuanImage3Model(BaseTransformerModel): 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._init_tensor_parallel() 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_tensor_parallel_config() self._validate_sequence_parallel_config() if self.lazy_load: self.remove_keys.extend(["model.layers."]) @@ -145,7 +150,17 @@ def __init__(self, model_path, config, device, lora_path=None, lora_strength=1.0 self._init_weights() self._init_infer() - if self.config.get("seq_parallel", False): + if self.config.get("seq_parallel", False) and self.tensor_parallel: + logger.info( + "HunyuanImage3 tensor + sequence parallel initialized: " + f"tp_rank={self.tp_rank}, " + f"tp_size={self.tp_size}, " + f"sp_rank={dist.get_rank(self.seq_p_group)}, " + f"sp_size={dist.get_world_size(self.seq_p_group)}, " + f"attention={self.sequence_parallel_attn_type}, " + f"device={self.pipeline_devices[0]}" + ) + elif self.config.get("seq_parallel", False): logger.info( "HunyuanImage3 sequence parallel initialized: " f"rank={dist.get_rank(self.seq_p_group)}, " @@ -153,6 +168,110 @@ def __init__(self, model_path, config, device, lora_path=None, lora_strength=1.0 f"attention={self.sequence_parallel_attn_type}, " f"pipeline_devices={self.pipeline_devices}" ) + elif self.tensor_parallel: + logger.info(f"HunyuanImage3 tensor parallel initialized: rank={self.tp_rank}, size={self.tp_size}, device={self.pipeline_devices[0]}") + + def _init_tensor_parallel(self): + self.tensor_parallel = bool(self.config.get("tensor_parallel", False)) + if self.tensor_parallel: + self.tp_group = self.config["device_mesh"].get_group(mesh_dim="tensor_p") + self.tp_rank = dist.get_rank(self.tp_group) + self.tp_size = dist.get_world_size(self.tp_group) + else: + self.tp_group = None + self.tp_rank = 0 + self.tp_size = 1 + + @staticmethod + def _iter_config_ints(value): + if isinstance(value, list): + return [int(item) for item in value] + if value is None: + return [] + return [int(value)] + + def _validate_tensor_parallel_config(self): + if not self.tensor_parallel: + return + if self.pipeline_parallel: + raise ValueError("HunyuanImage3 tensor parallel cannot be combined with pipeline parallel; set parallel.pipeline_parallel=false.") + if self.config.get("cpu_offload", False): + raise NotImplementedError("HunyuanImage3 tensor parallel does not support cpu_offload.") + if self.config.get("lazy_load", False): + raise NotImplementedError("HunyuanImage3 tensor parallel does not support lazy_load.") + if self.config.get("dit_quantized", False): + raise NotImplementedError("HunyuanImage3 tensor parallel currently supports the unquantized checkpoint only.") + if self.config.get("load_from_rank0", False): + raise NotImplementedError("HunyuanImage3 tensor parallel loads and shards safetensors locally on each rank; set load_from_rank0=false.") + if self.lora_path is not None: + raise NotImplementedError("HunyuanImage3 tensor parallel does not support LoRA weight loading yet.") + moe_impl = str(self.config.get("moe_impl", "eager")).strip().lower() + if moe_impl not in ("eager", "flashinfer"): + raise NotImplementedError("HunyuanImage3 tensor parallel supports moe_impl='eager' or 'flashinfer'.") + + divisibility_checks = { + "num_attention_heads": [int(self.config.get("num_attention_heads") or self.config["num_heads"])], + "num_key_value_heads": [int(self.config.get("num_key_value_heads") or self.config.get("num_attention_heads") or self.config["num_heads"])], + "intermediate_size": self._iter_config_ints(self.config.get("intermediate_size")), + "moe_intermediate_size": self._iter_config_ints(self.config.get("moe_intermediate_size")), + "vocab_size": self._iter_config_ints(self.config.get("vocab_size")), + } + shared_experts = self._iter_config_ints(self.config.get("num_shared_expert")) + moe_intermediate = self._iter_config_ints(self.config.get("moe_intermediate_size")) + if shared_experts and moe_intermediate: + if len(shared_experts) == 1: + shared_experts *= len(moe_intermediate) + if len(moe_intermediate) == 1: + moe_intermediate *= len(shared_experts) + divisibility_checks["shared_mlp_intermediate_size"] = [experts * intermediate for experts, intermediate in zip(shared_experts, moe_intermediate)] + + for name, values in divisibility_checks.items(): + invalid = sorted({value for value in values if value % self.tp_size}) + if invalid: + raise ValueError(f"HunyuanImage3 TP size {self.tp_size} must divide every {name}; invalid values: {invalid}.") + + @staticmethod + def _tp_split_type(key): + if not key.startswith("model.layers.") and not key.startswith("lm_head."): + return None + if ".self_attn.qkv_proj." in key: + return "qkv_col" + if ".self_attn.o_proj." in key: + return "row" + if ".gate_and_up_proj." in key: + return "gate_up_col" + if ".down_proj." in key: + return "row" + if key.startswith("lm_head."): + return "col" + return None + + def _select_tensor_parallel_shard(self, key, tensor): + split_type = self._tp_split_type(key) + if split_type is None: + return tensor + + if split_type == "row": + # Row-parallel biases are replicated and added after the reduction. + if tensor.ndim == 1: + return tensor + if tensor.ndim != 2 or tensor.shape[1] % self.tp_size: + raise ValueError(f"Cannot row-shard HunyuanImage3 tensor {key} with shape {tuple(tensor.shape)} across TP size {self.tp_size}.") + return torch.chunk(tensor, self.tp_size, dim=1)[self.tp_rank].contiguous() + + if split_type == "gate_up_col": + if tensor.shape[0] % 2: + raise ValueError(f"HunyuanImage3 fused gate/up tensor {key} has an odd output dimension: {tuple(tensor.shape)}.") + gate, up = tensor.chunk(2, dim=0) + if gate.shape[0] % self.tp_size: + raise ValueError(f"Cannot shard HunyuanImage3 fused gate/up tensor {key} with shape {tuple(tensor.shape)} across TP size {self.tp_size}.") + gate_shard = torch.chunk(gate, self.tp_size, dim=0)[self.tp_rank] + up_shard = torch.chunk(up, self.tp_size, dim=0)[self.tp_rank] + return torch.cat((gate_shard, up_shard), dim=0).contiguous() + + if tensor.shape[0] % self.tp_size: + raise ValueError(f"Cannot column-shard HunyuanImage3 tensor {key} with shape {tuple(tensor.shape)} across TP size {self.tp_size}.") + return torch.chunk(tensor, self.tp_size, dim=0)[self.tp_rank].contiguous() def _resolve_sequence_parallel_attn_type(self): if not self.config.get("seq_parallel", False): @@ -174,22 +293,29 @@ def _validate_sequence_parallel_config(self): 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() + parallel = self.config.get("parallel") or {} + cfg_mode = str(parallel.get("cfg_mode", 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'.") + raise ValueError("HunyuanImage3 CFG+SP requires parallel.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.") + raise ValueError("HunyuanImage3 sequence parallel requires parallel.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.get("num_attention_heads") or self.config["num_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(f"HunyuanImage3 Ulysses requires seq_p_size to divide both Q and KV heads: Q={q_heads}, KV={kv_heads}, seq_p_size={world_size}.") + global_q_heads = int(self.config.get("num_attention_heads") or self.config["num_heads"]) + global_kv_heads = int(self.config.get("num_key_value_heads") or global_q_heads) + local_q_heads = global_q_heads // self.tp_size + local_kv_heads = global_kv_heads // self.tp_size + if local_q_heads % world_size or local_kv_heads % world_size: + raise ValueError( + "HunyuanImage3 Ulysses requires seq_p_size to divide TP-local Q and KV heads: " + f"global_Q={global_q_heads}, global_KV={global_kv_heads}, tp_size={self.tp_size}, " + f"local_Q={local_q_heads}, local_KV={local_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) @@ -197,6 +323,8 @@ def _tensor_target_device(self, key): 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"): + if self.tensor_parallel: + raise NotImplementedError("HunyuanImage3 tensor parallel requires safetensors checkpoints.") return super()._load_safetensor_to_dict(file_path, unified_dtype, sensitive_layer) remove_keys = self.remove_keys if hasattr(self, "remove_keys") else [] @@ -210,6 +338,8 @@ def _load_safetensor_to_dict(self, file_path, unified_dtype, sensitive_layer): continue tensor = f.get_tensor(key) + if self.tensor_parallel: + tensor = self._select_tensor_parallel_shard(key, tensor) 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: @@ -302,7 +432,9 @@ def _infer_cond_uncond(self, inputs, infer_condition=True): @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.") + raise ValueError( + "HunyuanImage3 sequence parallel expects batch size 1 per transformer forward; use parallel.cfg_mode='serial' when cfg_p_size=1 (including TP+SP), or 'parallel' for CFG+SP." + ) world_size = dist.get_world_size(self.seq_p_group) rank = dist.get_rank(self.seq_p_group) diff --git a/lightx2v/models/networks/hunyuan_image3/weights/common.py b/lightx2v/models/networks/hunyuan_image3/weights/common.py index f3bc2ca6b..a631b703d 100644 --- a/lightx2v/models/networks/hunyuan_image3/weights/common.py +++ b/lightx2v/models/networks/hunyuan_image3/weights/common.py @@ -1,6 +1,7 @@ import math import torch +import torch.distributed as dist import torch.nn.functional as F import lightx2v.common.ops # noqa: F401 @@ -34,6 +35,50 @@ def _patch_embed_hidden_dim(config): return 1024 if hidden_dim is None else int(hidden_dim) +def hunyuan_image3_mm_weight( + config, + mm_type, + weight_name, + bias_name=None, + split_dim=None, + reduce_output=True, + create_cuda_buffer=False, + create_cpu_buffer=False, + lazy_load=False, + lazy_load_file=None, + lora_prefix="diffusion_model.blocks", + lora_path=None, +): + if config.get("tensor_parallel", False) and split_dim is not None: + tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p") + return MM_WEIGHT_REGISTER["TensorParallel"]( + weight_name=weight_name, + bias_name=bias_name, + mm_type=mm_type, + tp_group=tp_group, + tp_rank=dist.get_rank(tp_group), + tp_size=dist.get_world_size(tp_group), + split_dim=split_dim, + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=lazy_load, + lazy_load_file=lazy_load_file, + lora_prefix=lora_prefix, + lora_path=lora_path, + reduce_output=reduce_output, + ) + return MM_WEIGHT_REGISTER[mm_type]( + weight_name, + bias_name, + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + lora_prefix=lora_prefix, + lora_path=lora_path, + ) + + class TensorPairWeight: def __init__(self, weight_name, bias_name=None): self.weight_name = weight_name @@ -245,8 +290,10 @@ class HunyuanImage3DenseMLPWeights(WeightModule): def __init__( self, prefix, + config, mm_type, mlp_bias=False, + defer_down_reduce=False, create_cuda_buffer=False, create_cpu_buffer=False, lazy_load=False, @@ -258,26 +305,33 @@ def __init__( down_bias = f"{prefix}.down_proj.bias" if mlp_bias else None self.add_module( "gate_and_up_proj", - MM_WEIGHT_REGISTER[mm_type]( + hunyuan_image3_mm_weight( + config, + mm_type, f"{prefix}.gate_and_up_proj.weight", gate_and_up_bias, - create_cuda_buffer, - create_cpu_buffer, - lazy_load, - lazy_load_file, + split_dim="col", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=lazy_load, + lazy_load_file=lazy_load_file, lora_prefix=prefix, lora_path=lora_path, ), ) self.add_module( "down_proj", - MM_WEIGHT_REGISTER[mm_type]( + hunyuan_image3_mm_weight( + config, + mm_type, f"{prefix}.down_proj.weight", down_bias, - create_cuda_buffer, - create_cpu_buffer, - lazy_load, - lazy_load_file, + split_dim="row", + reduce_output=not defer_down_reduce, + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=lazy_load, + lazy_load_file=lazy_load_file, lora_prefix=prefix, lora_path=lora_path, ), @@ -324,8 +378,10 @@ def __init__( "shared_mlp", HunyuanImage3DenseMLPWeights( f"{prefix}.shared_mlp", + config, mm_type, config.get("mlp_bias", False), + True, create_cuda_buffer, create_cpu_buffer, lazy_load, @@ -339,8 +395,10 @@ def __init__( [ HunyuanImage3DenseMLPWeights( f"{prefix}.experts.{i}", + config, mm_type, config.get("mlp_bias", False), + True, create_cuda_buffer, create_cpu_buffer, lazy_load, @@ -352,12 +410,30 @@ def __init__( ) self.add_module("experts", self.experts) + @staticmethod + def _unwrap_linear_weight(linear): + """Return the concrete MM weight hidden behind TP wrappers.""" + current = linear + visited = set() + while getattr(current, "_mm", None) is not None: + if id(current) in visited: + raise RuntimeError("Detected a cycle while unwrapping a HunyuanImage3 MM weight.") + visited.add(id(current)) + current = current._mm + return current + 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.") + linear = self._unwrap_linear_weight(linear) + 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) @@ -373,8 +449,9 @@ def _linear_weight_for_flashinfer(self, linear, device, dtype): # 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): + @classmethod + def _release_linear_weight(cls, linear, device, dtype): + linear = cls._unwrap_linear_weight(linear) empty = torch.empty(0, device=device, dtype=dtype) for attr in ("weight", "pin_weight", "bias", "pin_bias"): if hasattr(linear, attr): @@ -444,26 +521,32 @@ def __init__( attn_bias = ".bias" if config.get("attention_bias", False) else None self.add_module( "qkv_proj", - MM_WEIGHT_REGISTER[mm_type]( + hunyuan_image3_mm_weight( + config, + 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, + split_dim="col", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=lazy_load, + lazy_load_file=lazy_load_file, lora_prefix=prefix, lora_path=lora_path, ), ) self.add_module( "o_proj", - MM_WEIGHT_REGISTER[mm_type]( + hunyuan_image3_mm_weight( + config, + 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, + split_dim="row", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=lazy_load, + lazy_load_file=lazy_load_file, lora_prefix=prefix, lora_path=lora_path, ), @@ -533,8 +616,10 @@ def __init__( "dense_mlp", HunyuanImage3DenseMLPWeights( f"{prefix}.mlp", + config, mm_type, config.get("mlp_bias", False), + False, create_cuda_buffer, create_cpu_buffer, lazy_load, diff --git a/lightx2v/models/networks/hunyuan_image3/weights/post_weights.py b/lightx2v/models/networks/hunyuan_image3/weights/post_weights.py index f86f0fe3c..60b64465d 100644 --- a/lightx2v/models/networks/hunyuan_image3/weights/post_weights.py +++ b/lightx2v/models/networks/hunyuan_image3/weights/post_weights.py @@ -2,8 +2,9 @@ from lightx2v.models.networks.hunyuan_image3.weights.common import ( HunyuanImage3TimestepEmbedderWeights, HunyuanImage3UNetUpWeights, + hunyuan_image3_mm_weight, ) -from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import RMS_WEIGHT_REGISTER class HunyuanImage3PostWeights(WeightModule): @@ -18,7 +19,15 @@ def __init__(self, config): eps=config.get("rms_norm_eps", 1e-5), ), ) - self.add_module("lm_head", MM_WEIGHT_REGISTER["Default"]("lm_head.weight")) + self.add_module( + "lm_head", + hunyuan_image3_mm_weight( + config, + "Default", + "lm_head.weight", + split_dim="col", + ), + ) self.add_module("final_layer", HunyuanImage3UNetUpWeights("final_layer", config)) def to_cpu(self, non_blocking=True): diff --git a/lightx2v/models/runners/hunyuan_image3/flashinfer_autotune.py b/lightx2v/models/runners/hunyuan_image3/flashinfer_autotune.py index 8058275a5..28c26adb9 100644 --- a/lightx2v/models/runners/hunyuan_image3/flashinfer_autotune.py +++ b/lightx2v/models/runners/hunyuan_image3/flashinfer_autotune.py @@ -18,10 +18,11 @@ flashinfer_autotune = None -FlashInferAutotuneMode = Literal["off", "tune", "load"] +FlashInferAutotuneMode = Literal["off", "auto", "tune", "load"] AutotuneFactory = Callable[..., AbstractContextManager] TunerProvider = Callable[[], Any] StatusDeviceResolver = Callable[[], torch.device] +LIGHTX2V_ROOT = Path(__file__).resolve().parents[4] @dataclass(frozen=True) @@ -42,6 +43,13 @@ def enabled(self) -> bool: def tune_mode(self) -> bool: return self.mode == "tune" + def effective_tune_mode(self) -> bool: + if self.mode == "tune": + return True + if self.mode == "auto": + return self.cache_path is None or not self.cache_path.is_file() + return False + @classmethod def from_config(cls, config: Mapping[str, Any]) -> FlashInferAutotuneOptions: mode = cls.resolve_mode(config) @@ -78,8 +86,8 @@ def resolve_mode(cls, config: Mapping[str, Any]) -> FlashInferAutotuneMode: "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.") + if normalized not in ("off", "auto", "tune", "load"): + raise ValueError("flashinfer_autotune_mode must be one of: off, auto, tune, load.") return normalized @staticmethod @@ -102,7 +110,10 @@ def _parse_optional_bool(cls, value: Any) -> bool | None: def _resolve_cache_path(value: Any) -> Path | None: if value in (None, ""): return None - return Path(value).expanduser() + path = Path(value).expanduser() + if not path.is_absolute(): + path = LIGHTX2V_ROOT / path + return path.resolve() @staticmethod def _parse_tuning_buckets(value: Any) -> tuple[int, ...] | None: @@ -179,16 +190,17 @@ def context(self) -> AbstractContextManager: autotune_factory = self._resolve_autotune_factory() self._prepare_cache_path() + effective_tune_mode = self.options.effective_tune_mode() logger.info( "HunyuanImage3 FlashInfer autotune enabled: " - f"mode={self.options.mode}, cache={self.options.cache_path}, " + f"mode={self.options.mode}, effective_tune_mode={effective_tune_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) + return self._distributed_autotune_context(autotune_factory, effective_tune_mode) + return self._local_autotune_context(autotune_factory, effective_tune_mode) def _resolve_autotune_factory(self) -> AutotuneFactory: factory = self._autotune_factory or flashinfer_autotune @@ -205,29 +217,28 @@ def _resolve_tuner(self) -> Any: 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.") + if cache_path is None: + raise ValueError(f"flashinfer_autotune_mode={self.options.mode!r} requires flashinfer_autotune_cache.") + cache_path.parent.mkdir(parents=True, exist_ok=True) - def _local_autotune_context(self, autotune_factory: AutotuneFactory) -> AbstractContextManager: + def _local_autotune_context(self, autotune_factory: AutotuneFactory, effective_tune_mode: bool) -> 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, + tune_mode=effective_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): + def _distributed_autotune_context(self, autotune_factory: AutotuneFactory, effective_tune_mode: bool): 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.") + raise RuntimeError("Distributed FlashInfer autotune requires an initialized tensor, CFG, or sequence parallel group.") if dist.get_backend(group) == "nccl": if runtime.status_device_resolver is None: @@ -245,14 +256,14 @@ def _distributed_autotune_context(self, autotune_factory: AutotuneFactory): raise RuntimeError(f"Distributed FlashInfer autotune cache preflight failed for {cache_path}.") from load_error with autotune_factory( - tune_mode=self.options.tune_mode, + tune_mode=effective_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_error = self._save_distributed_cache(tuner, cache_path, effective_tune_mode) 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) @@ -273,11 +284,15 @@ def _load_distributed_cache(self, tuner: Any, cache_path: str | None) -> Excepti 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: + def _save_distributed_cache(self, tuner: Any, cache_path: str | None, effective_tune_mode: bool) -> Exception | None: + if not effective_tune_mode or cache_path is None or not self.distributed_context.is_cache_writer: return None try: tuner.save_configs(cache_path) + cache_file = Path(cache_path) + if not cache_file.is_file() or cache_file.stat().st_size == 0: + raise RuntimeError(f"FlashInfer autotune did not create a non-empty cache file: {cache_path}") + logger.info(f"HunyuanImage3 FlashInfer autotune cache saved to: {cache_path}") except Exception as error: logger.exception(f"Failed to persist distributed FlashInfer autotune cache: {cache_path}") return error diff --git a/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py b/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py index 59a4ea7f2..fc41ce186 100644 --- a/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py +++ b/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py @@ -2,6 +2,7 @@ import os import sys from dataclasses import dataclass +from functools import partial from pathlib import Path import torch @@ -127,23 +128,34 @@ def _ensure_pipeline_modules(self): model_version=self.hunyuan_config.model_version, ) self.hunyuan_image_processor = modules["HunyuanImage3ImageProcessor"](self.hunyuan_config) + # Newer transformers versions alias Siglip2ImageProcessorFast to the + # slow processor, whose default output is a Python list. The upstream + # Hunyuan processor expects tensors and calls .squeeze() immediately. + vit_processor = self.hunyuan_image_processor.vit_info.processor + self.hunyuan_image_processor.vit_info.processor = partial(vit_processor, return_tensors="pt") 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()) + # Text-only tasks do not need a VAE. Image paths load it on demand so + # T2T can run without paying the extra memory and initialization cost. + self.vae_decoder = None + self.vision_model = None + self.vision_aligner = None + self._hunyuan_pipeline_modules_ready = True + + def _ensure_vae(self): + if self.vae_decoder is not None: + return + + task = str(self.config.get("task", "t2i")).lower() + load_vae_on_this_rank = not (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) @@ -345,11 +357,13 @@ def _hunyuan_taylor_cache_enabled(self): def _build_flashinfer_autotune_controller(self): sequence_parallel_active = self._sequence_parallel_enabled() + tensor_parallel_active = self._tensor_parallel_enabled() + distributed_parallel_active = sequence_parallel_active or tensor_parallel_active 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, + coordination_required=distributed_parallel_active, + process_group=self._parallel_control_group() if distributed_parallel_active else None, + status_device_resolver=self._pipeline_latent_device if distributed_parallel_active else None, + is_cache_writer=self._is_output_rank() if distributed_parallel_active else False, ) return FlashInferAutotuneController.from_config( config=self.config, @@ -359,23 +373,28 @@ def _build_flashinfer_autotune_controller(self): 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 _tensor_parallel_enabled(self): + return bool(self.config.get("tensor_parallel", False) and dist.is_available() and dist.is_initialized() and getattr(self.model, "tp_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. + exchange remains scoped to ``cfg_p_group``. Tensor-parallel runs keep + sampled tokens and initial latents synchronized through ``tensor_p``. + Hybrid TP/SP/CFG runs use WORLD so runner state spans every active + mesh dimension. Pure CFG deliberately returns no runner control group. """ - if not self._sequence_parallel_enabled(): - return None - if self._cfg_parallel_enabled(): + tensor_parallel_active = self._tensor_parallel_enabled() + sequence_parallel_active = self._sequence_parallel_enabled() + cfg_parallel_active = self._cfg_parallel_enabled() + if sum((tensor_parallel_active, sequence_parallel_active, cfg_parallel_active)) > 1: return dist.group.WORLD - return self.model.seq_p_group + if tensor_parallel_active: + return self.model.tp_group + if sequence_parallel_active: + return self.model.seq_p_group + return None def _is_output_rank(self): return not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0 @@ -520,11 +539,12 @@ 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() + parallel = self.config.get("parallel") or {} + mode = str(parallel.get("cfg_mode", 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}.") + raise ValueError(f"HunyuanImage3 parallel.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.") + raise ValueError("HunyuanImage3 parallel.cfg_mode='parallel' requires enable_cfg=true and parallel.cfg_p_size=2.") return mode def _prepare_cfg_parallel_branch_inputs(self, prepared_inputs, cfg_p_rank, mark_parallel_branch=True): @@ -628,7 +648,7 @@ def _attach_tokenizer_cond_masks(self, cond_inputs, tokenizer_output, device): 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" + return str(bot_task or "image").strip().lower() 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")) @@ -641,6 +661,12 @@ def _resolve_system_prompt(self, bot_task): def _resolve_text_generation_plan(self, bot_task): tokenizer = self.hunyuan_tokenizer + if bot_task == "auto": + return HunyuanImage3TextGenerationPlan( + first_bot_task="auto", + stage_transitions=[], + final_stop_tokens=list(tokenizer.conversation.stop_token_ids), + ) if bot_task == "recaption": return HunyuanImage3TextGenerationPlan( first_bot_task="recaption", @@ -682,13 +708,14 @@ def _build_text_model_inputs( build_attention_mask=True, past_key_values=None, use_cache=False, + rope_image_info=None, ): 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) + custom_pos_emb = self.hunyuan_cached_rope(max_position, device, rope_image_info=rope_image_info, 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 @@ -724,35 +751,38 @@ def _prepare_text_generation_inputs(self, prompt, bot_task, system_prompt, batch 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, + image_base_size=None if bot_task == "auto" else 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"] + rope_image_info = self._build_batch_rope_image_info(output, tokenizer_out["sections"]) return { "input_ids": output.tokens.to(latent_device), "tokenizer_output": output, "cond_inputs": self._attach_tokenizer_cond_masks(cond_inputs, output, latent_device), + "rope_image_info": rope_image_info, } - def _sample_text_token(self, logits, generator): - do_sample = bool(self.config.get("text_do_sample", getattr(self.hunyuan_generation_config, "do_sample", False))) + def _sample_text_token(self, logits, generator, generation_options=None): + generation_options = generation_options or {} + do_sample = bool(generation_options.get("text_do_sample", 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))) + temperature = float(generation_options.get("text_temperature", 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) + top_k = int(generation_options.get("text_top_k", 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))) + top_p = float(generation_options.get("text_top_p", 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) @@ -777,11 +807,22 @@ def _print_text_generation_chunk(self, text="", end=""): print(text, end=end, flush=True) @torch.no_grad() - def _generate_text_tokens(self, input_ids, tokenizer_output, plan, stream_output=False, cond_inputs=None): + def _generate_text_tokens( + self, + input_ids, + tokenizer_output, + plan, + stream_output=False, + cond_inputs=None, + rope_image_info=None, + stream_callback=None, + generation_options=None, + ): + generation_options = generation_options or {} device = input_ids.device - seed = int(self.config.get("text_seed", self.config.get("seed", 42))) + seed = int(generation_options.get("text_seed", 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))) + max_new_tokens = int(generation_options.get("max_new_tokens", 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 = [] @@ -813,23 +854,28 @@ def _generate_text_tokens(self, input_ids, tokenizer_output, plan, stream_output build_attention_mask=first_cache_step, past_key_values=kv_cache, use_cache=True, + rope_image_info=rope_image_info if first_cache_step else None, ) cache_filled_length = input_ids.shape[1] else: if cond_inputs is None: - model_inputs = self._build_text_model_inputs(input_ids, tokenizer_output) + model_inputs = self._build_text_model_inputs(input_ids, tokenizer_output, rope_image_info=rope_image_info) else: - model_inputs = self._build_text_model_inputs(input_ids, tokenizer_output, cond_inputs=cond_inputs) + model_inputs = self._build_text_model_inputs(input_ids, tokenizer_output, cond_inputs=cond_inputs, rope_image_info=rope_image_info) 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 = self._sample_text_token(logits, generator, generation_options=generation_options) 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)) + text_chunk = self._decode_generated_text_token(next_token_id) + if stream_callback is not None and self._is_output_rank(): + stream_callback(text_chunk) + else: + self._print_text_generation_chunk(text_chunk) 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: @@ -841,6 +887,20 @@ def _generate_text_tokens(self, input_ids, tokenizer_output, plan, stream_output return generated + def _decode_text_result(self, generated_tokens, plan): + tokens = list(generated_tokens) + while tokens and tokens[-1] in plan.final_stop_tokens: + tokens.pop() + if not tokens: + return "" + + generated_text = self.hunyuan_tokenizer.decode(torch.tensor(tokens, dtype=torch.long)) + if plan.first_bot_task == "think": + return self.hunyuan_tokenizer.think_token + generated_text + if plan.first_bot_task == "recaption": + return self.hunyuan_tokenizer.recaption_token + generated_text + return generated_text + def _decode_cot_text(self, generated_tokens, plan): if not generated_tokens: return None @@ -850,6 +910,67 @@ def _decode_cot_text(self, generated_tokens, plan): return self.hunyuan_tokenizer.think_token + generated_text return self.hunyuan_tokenizer.recaption_token + generated_text + def _input_text_generation_options(self, input_info): + options = {} + for key in ("max_new_tokens", "text_do_sample", "text_temperature", "text_top_k", "text_top_p"): + value = getattr(input_info, key, None) + if value is not None: + options[key] = value + seed = getattr(input_info, "seed", None) + if seed is not None: + options["text_seed"] = seed + return options + + def _stream_text_prefix(self, plan, stream_callback=None): + if plan.first_bot_task == "think": + prefix = self.hunyuan_tokenizer.think_token + elif plan.first_bot_task == "recaption": + prefix = self.hunyuan_tokenizer.recaption_token + else: + prefix = "" + if not prefix: + return + if stream_callback is not None and self._is_output_rank(): + stream_callback(prefix) + else: + self._print_text_generation_chunk(prefix) + + def _generate_text(self, input_info, batch_cond_images=None, cond_inputs=None): + bot_task = str(getattr(input_info, "bot_task", None) or self._resolve_bot_task()).strip().lower() + if bot_task == "image": + raise ValueError("HunyuanImage3 text generation requires bot_task='auto', 'think', 'recaption', or 'think_recaption'.") + + plan = self._resolve_text_generation_plan(bot_task) + system_prompt = getattr(input_info, "system_prompt", None) + if system_prompt is None: + system_prompt = self._resolve_system_prompt(plan.first_bot_task) + prepared = self._prepare_text_generation_inputs( + getattr(input_info, "prompt", ""), + plan.first_bot_task, + system_prompt, + batch_cond_images=batch_cond_images, + cond_inputs=cond_inputs, + ) + + stream_callback = getattr(input_info, "stream_callback", None) + stream_output = self._text_stream_enabled() or stream_callback is not None + if stream_output: + self._stream_text_prefix(plan, stream_callback=stream_callback) + with self._build_flashinfer_autotune_controller().context(): + generated_tokens = self._generate_text_tokens( + prepared["input_ids"], + prepared["tokenizer_output"], + plan, + stream_output=stream_output, + cond_inputs=prepared.get("cond_inputs"), + rope_image_info=prepared.get("rope_image_info"), + stream_callback=stream_callback, + generation_options=self._input_text_generation_options(input_info), + ) + if stream_output and stream_callback is None: + self._print_text_generation_chunk("\n") + return self._decode_text_result(generated_tokens, plan) + 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": @@ -871,13 +992,15 @@ def _generate_cot_text(self, prompt, image_size, batch_cond_images=None, cond_in 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"), - ) + with self._build_flashinfer_autotune_controller().context(): + generated_tokens = self._generate_text_tokens( + prepared["input_ids"], + prepared["tokenizer_output"], + plan, + stream_output=stream_output, + cond_inputs=prepared.get("cond_inputs"), + rope_image_info=prepared.get("rope_image_info"), + ) if stream_output: self._print_text_generation_chunk("\n") cot_text = self._decode_cot_text(generated_tokens, plan) @@ -969,6 +1092,7 @@ def _resolve_ti2i_image_size(self, requested_image_size, batch_cond_images): 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): + self._ensure_vae() vae = self.vae_decoder vae_device = self._vae_device() image_tensor = image_tensor.unsqueeze(0).to(vae_device) @@ -1064,6 +1188,19 @@ def _prepare_cond_inputs(self, batch_cond_images, cfg_factor, seed): "cond_vit_embeds": cond_vit_embeds, } + def _repeat_cond_inputs(self, cond_inputs, factor): + if factor == 1: + return cond_inputs + repeated = {} + for key, value in cond_inputs.items(): + if isinstance(value, torch.Tensor): + repeated[key] = value.repeat(factor, *([1] * (value.ndim - 1))) + elif isinstance(value, list): + repeated[key] = value * factor + else: + repeated[key] = value + return repeated + def _prepare_latents(self, batch_size, image_size, generator): latent_device = self._pipeline_latent_device() downsample = self.hunyuan_config.vae_downsample_factor @@ -1099,7 +1236,7 @@ def _denoise_latents(self, prepared_inputs, image_size): 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.") + raise NotImplementedError("HunyuanImage3 Taylor cache currently supports parallel.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: @@ -1239,6 +1376,7 @@ def _denoise_latents(self, prepared_inputs, image_size): return latents def _decode_latents(self, latents, generator): + self._ensure_vae() vae = self.vae_decoder vae_device = self._vae_device() latents = latents.to(vae_device) @@ -1258,18 +1396,38 @@ def _decode_latents(self, latents, generator): 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_t2t(self, input_info): + self._ensure_pipeline_modules() + return self._generate_text(input_info) + + @torch.no_grad() + def generate_ti2t(self, input_info): + self._ensure_pipeline_modules() + seed = getattr(input_info, "seed", None) + if seed is None: + seed = 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 = getattr(input_info, "infer_align_image_size", None) + if infer_align_image_size is None: + infer_align_image_size = self.config.get("infer_align_image_size", False) + + batch_cond_images = self._build_batch_cond_images(image_paths, bool(infer_align_image_size)) + cond_inputs = self._prepare_cond_inputs(batch_cond_images, cfg_factor=1, seed=seed) + return self._generate_text(input_info, batch_cond_images=batch_cond_images, cond_inputs=cond_inputs) + @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) + 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 + cot_text=cot_text, ) latents = self._denoise_latents(prepared_inputs, image_size) if not self._is_output_rank(): @@ -1277,27 +1435,30 @@ def generate_t2i(self, input_info): return self._decode_latents(latents, prepared_inputs["generator"]) @torch.no_grad() - def generate_i2i(self, input_info): + def generate_ti2i(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))) + infer_align_image_size = getattr(input_info, "infer_align_image_size", None) + if infer_align_image_size is None: + infer_align_image_size = self.config.get("infer_align_image_size", False) + infer_align_image_size = bool(infer_align_image_size) 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) + 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) + gen_cond_inputs = self._repeat_cond_inputs(text_cond_inputs, factor=2 if do_cfg else 1) prepared_inputs = self._prepare_text_to_image_inputs( prompt, image_size, seed, - # cot_text=cot_text, + cot_text=cot_text, batch_cond_images=batch_cond_images, cond_inputs=gen_cond_inputs, ) @@ -1311,18 +1472,37 @@ def generate_i2i(self, input_info): infer_align_image_size=infer_align_image_size, ) + def generate_i2i(self, input_info): + """Compatibility alias for the pre-TI2I task name.""" + return self.generate_ti2i(input_info) + def run_pipeline(self, input_info): self.input_info = input_info task = self.config.get("task") - if task == "t2i": + if task == "t2t": + text = self.generate_t2t(input_info) + elif task == "ti2t": + text = self.generate_ti2t(input_info) + elif task == "t2i": images = self.generate_t2i(input_info) elif task in ("i2i", "ti2i"): - images = self.generate_i2i(input_info) + images = self.generate_ti2i(input_info) else: - raise NotImplementedError("HunyuanImage3 native runner currently supports task=t2i/i2i.") + raise NotImplementedError("HunyuanImage3 native runner supports task=t2t/t2i/ti2t/ti2i (i2i is a compatibility alias).") if not self._is_output_rank(): - return {"image": None} + return {"text": None} if task in ("t2t", "ti2t") else {"image": None} + + if task in ("t2t", "ti2t"): + if getattr(input_info, "return_result_tensor", False): + return {"text": text} + 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) + Path(save_result_path).write_text(text, encoding="utf-8") + logger.info(f"HunyuanImage3 text saved successfully to: {save_result_path}") + return {"text": None} + if getattr(input_info, "return_result_tensor", False): return {"image": images} diff --git a/lightx2v/utils/input_info.py b/lightx2v/utils/input_info.py index dd343efbd..14eed780c 100755 --- a/lightx2v/utils/input_info.py +++ b/lightx2v/utils/input_info.py @@ -208,6 +208,28 @@ class T2IInputInfo: aspect_ratio: str = field(default_factory=str) +@dataclass +class T2TInputInfo: + seed: int = field(default_factory=int) + prompt: str = field(default_factory=str) + save_result_path: str = field(default_factory=str) + return_result_tensor: bool = field(default_factory=lambda: False) + max_new_tokens: Optional[int] = None + text_do_sample: Optional[bool] = None + text_temperature: Optional[float] = None + text_top_k: Optional[int] = None + text_top_p: Optional[float] = None + bot_task: Optional[str] = None + system_prompt: Optional[str] = None + stream_callback: Any = None + + +@dataclass +class TI2TInputInfo(T2TInputInfo): + image_path: str = field(default_factory=str) + infer_align_image_size: Optional[bool] = None + + @dataclass class I2IInputInfo: seed: int = field(default_factory=int) @@ -231,6 +253,11 @@ class I2IInputInfo: aspect_ratio: str = field(default_factory=str) +@dataclass +class TI2IInputInfo(I2IInputInfo): + infer_align_image_size: Optional[bool] = None + + @dataclass class T2AVInputInfo: seed: int = field(default_factory=int) @@ -452,7 +479,10 @@ class WorldPlayT2VInputInfo: "s2v": S2VInputInfo, "rs2v": RS2VInputInfo, "animate": AnimateInputInfo, + "t2t": T2TInputInfo, "t2i": T2IInputInfo, + "ti2t": TI2TInputInfo, + "ti2i": TI2IInputInfo, "i2i": I2IInputInfo, "t2av": T2AVInputInfo, "i2av": I2AVInputInfo, diff --git a/lightx2v/utils/set_config.py b/lightx2v/utils/set_config.py index b1eafe58c..8497133b1 100755 --- a/lightx2v/utils/set_config.py +++ b/lightx2v/utils/set_config.py @@ -237,6 +237,23 @@ def auto_calc_config(config): config["infer_steps"] = config["num_inference_steps"] if config["model_cls"] == "hunyuan_image3": + task = str(config.get("task", "t2i")).strip().lower() + supported_tasks = {"t2t", "t2i", "ti2t", "ti2i", "i2i"} + if task not in supported_tasks: + raise ValueError(f"HunyuanImage3 task must be one of {sorted(supported_tasks)}, got {task!r}.") + + bot_task = str(config.get("bot_task", "image")).strip().lower() + supported_bot_tasks = {"image", "auto", "think", "recaption", "think_recaption"} + if bot_task not in supported_bot_tasks: + raise ValueError(f"HunyuanImage3 bot_task must be one of {sorted(supported_bot_tasks)}, got {bot_task!r}.") + config["bot_task"] = bot_task + + if task in {"t2t", "ti2t"}: + if config.get("enable_cfg", False): + raise ValueError(f"HunyuanImage3 task={task} does not support diffusion CFG; set enable_cfg=false.") + if bot_task == "image": + raise ValueError(f"HunyuanImage3 task={task} requires a text bot_task such as 'auto' or 'think_recaption'.") + 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: @@ -244,17 +261,52 @@ def auto_calc_config(config): parallel_config = config.get("parallel") if isinstance(parallel_config, dict): parallel_config = dict(parallel_config) + tensor_p_size = int(parallel_config.get("tensor_p_size", 1)) cfg_p_size = int(parallel_config.get("cfg_p_size", 1)) seq_p_size = int(parallel_config.get("seq_p_size", 1)) + + nested_pipeline_parallel = parallel_config.get("pipeline_parallel") + legacy_pipeline_parallel = config.get("pipeline_parallel") + if nested_pipeline_parallel is not None and legacy_pipeline_parallel is not None and nested_pipeline_parallel != legacy_pipeline_parallel: + raise ValueError(f"Conflicting HunyuanImage3 pipeline settings: parallel.pipeline_parallel={nested_pipeline_parallel!r}, pipeline_parallel={legacy_pipeline_parallel!r}.") + pipeline_parallel = nested_pipeline_parallel if nested_pipeline_parallel is not None else legacy_pipeline_parallel + if pipeline_parallel is None: + pipeline_parallel = True + if not isinstance(pipeline_parallel, bool): + raise ValueError(f"HunyuanImage3 parallel.pipeline_parallel must be a boolean, got {pipeline_parallel!r}.") + + nested_cfg_mode = parallel_config.get("cfg_mode") + legacy_cfg_mode = config.get("hunyuan_cfg_mode") + if nested_cfg_mode is not None and legacy_cfg_mode is not None: + if str(nested_cfg_mode).strip().lower() != str(legacy_cfg_mode).strip().lower(): + raise ValueError(f"Conflicting HunyuanImage3 CFG modes: parallel.cfg_mode={nested_cfg_mode!r}, hunyuan_cfg_mode={legacy_cfg_mode!r}.") + cfg_mode = nested_cfg_mode if nested_cfg_mode is not None else legacy_cfg_mode + cfg_mode = str(cfg_mode or "batch").strip().lower() + if cfg_mode not in ("batch", "serial", "parallel"): + raise ValueError(f"HunyuanImage3 parallel.cfg_mode must be one of batch/serial/parallel, got {cfg_mode!r}.") + + if tensor_p_size < 1: + raise ValueError(f"HunyuanImage3 parallel.tensor_p_size must be >= 1, got {tensor_p_size}.") 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 tensor_p_size > 1: + if pipeline_parallel: + raise ValueError("HunyuanImage3 tensor parallel requires parallel.pipeline_parallel=false.") + moe_impl = str(config.get("moe_impl", "eager")).strip().lower() + if moe_impl not in ("eager", "flashinfer"): + raise ValueError("HunyuanImage3 tensor parallel supports moe_impl='eager' or 'flashinfer'.") if cfg_p_size == 2 and not config.get("enable_cfg", False): raise ValueError("HunyuanImage3 parallel.cfg_p_size=2 requires enable_cfg=true.") + if task in {"t2t", "ti2t"} and cfg_p_size != 1: + raise ValueError(f"HunyuanImage3 task={task} requires parallel.cfg_p_size=1.") + parallel_config["tensor_p_size"] = tensor_p_size parallel_config["cfg_p_size"] = cfg_p_size parallel_config["seq_p_size"] = seq_p_size + parallel_config["pipeline_parallel"] = pipeline_parallel + parallel_config["cfg_mode"] = cfg_mode if seq_p_size > 1: attn_type = parallel_config.get("seq_p_attn_type", "kv_all_gather") @@ -266,18 +318,24 @@ def auto_calc_config(config): parallel_config["seq_p_attn_type"] = attn_type config["parallel"] = parallel_config + # Runtime aliases keep existing model code and legacy callers + # compatible while JSON files use the canonical parallel branch. + config["pipeline_parallel"] = pipeline_parallel + config["hunyuan_cfg_mode"] = cfg_mode - 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'.") + raise ValueError("HunyuanImage3 parallel.cfg_p_size=2 requires parallel.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'.") + raise ValueError("HunyuanImage3 sequence parallel with cfg_p_size=1 requires parallel.cfg_mode='serial'.") if seq_p_size > 1 and parallel_config["seq_p_attn_type"] == "ulysses": q_heads = int(config.get("num_attention_heads") or config["num_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(f"HunyuanImage3 Ulysses requires seq_p_size to divide Q and KV heads: Q={q_heads}, KV={kv_heads}, seq_p_size={seq_p_size}.") + combined_head_parallel_size = tensor_p_size * seq_p_size + if q_heads % combined_head_parallel_size or kv_heads % combined_head_parallel_size: + raise ValueError( + f"HunyuanImage3 Ulysses requires tensor_p_size * seq_p_size to divide Q and KV heads: Q={q_heads}, KV={kv_heads}, tensor_p_size={tensor_p_size}, seq_p_size={seq_p_size}." + ) if config["model_cls"] == "lingbot_va" and "target_video_length" not in config: ar_config = config.get("ar_config", {}) @@ -335,28 +393,45 @@ def set_config(args): def set_parallel_config(config): if config["parallel"]: - tensor_p_size = config["parallel"].get("tensor_p_size", 1) + tensor_p_size = int(config["parallel"].get("tensor_p_size", 1)) + cfg_p_size = int(config["parallel"].get("cfg_p_size", 1)) + seq_p_size = int(config["parallel"].get("seq_p_size", 1)) + world_size = dist.get_world_size() + expected_world_size = tensor_p_size * cfg_p_size * seq_p_size + if expected_world_size != world_size: + raise ValueError( + f"Parallel sizes must match the distributed world size: tensor_p_size ({tensor_p_size}) * cfg_p_size ({cfg_p_size}) * seq_p_size ({seq_p_size}) != world_size ({world_size})." + ) if tensor_p_size > 1: - # Tensor parallel only: 1D mesh - assert tensor_p_size == dist.get_world_size(), f"tensor_p_size ({tensor_p_size}) must be equal to world_size ({dist.get_world_size()})" - config["device_mesh"] = init_device_mesh(AI_DEVICE, (tensor_p_size,), mesh_dim_names=("tensor_p",)) + # Tensor parallel is the innermost dimension. Optional CFG and + # sequence dimensions are prepended so ranks with the same + # non-TP coordinates form contiguous TP groups. For TP+SP+CFG: + # mesh shape/names = [cfg_p, seq_p, tensor_p]. + mesh_shape = [] + mesh_dim_names = [] + if cfg_p_size > 1: + mesh_shape.append(cfg_p_size) + mesh_dim_names.append("cfg_p") + if seq_p_size > 1: + mesh_shape.append(seq_p_size) + mesh_dim_names.append("seq_p") + mesh_shape.append(tensor_p_size) + mesh_dim_names.append("tensor_p") + config["device_mesh"] = init_device_mesh( + AI_DEVICE, + tuple(mesh_shape), + mesh_dim_names=tuple(mesh_dim_names), + ) config["tensor_parallel"] = True - config["seq_parallel"] = False - config["cfg_parallel"] = False + config["seq_parallel"] = seq_p_size > 1 + config["cfg_parallel"] = bool(config.get("enable_cfg", False) and cfg_p_size > 1) else: # Original 2D mesh for cfg_p and seq_p - cfg_p_size = config["parallel"].get("cfg_p_size", 1) - seq_p_size = config["parallel"].get("seq_p_size", 1) - assert cfg_p_size * seq_p_size == dist.get_world_size(), f"cfg_p_size ({cfg_p_size}) * seq_p_size ({seq_p_size}) must be equal to world_size ({dist.get_world_size()})" config["device_mesh"] = init_device_mesh(AI_DEVICE, (cfg_p_size, seq_p_size), mesh_dim_names=("cfg_p", "seq_p")) config["tensor_parallel"] = False - - if config["parallel"] and config["parallel"].get("seq_p_size", False) and config["parallel"]["seq_p_size"] > 1: - config["seq_parallel"] = True - - if config.get("enable_cfg", False) and config["parallel"] and config["parallel"].get("cfg_p_size", False) and config["parallel"]["cfg_p_size"] > 1: - config["cfg_parallel"] = True + config["seq_parallel"] = seq_p_size > 1 + config["cfg_parallel"] = bool(config.get("enable_cfg", False) and cfg_p_size > 1) # warmup dist if AI_DEVICE == "cuda": diff --git a/scripts/hunyuan_image3/run_hunyuan_image3_i2i.sh b/scripts/hunyuan_image3/run_hunyuan_image3_i2i.sh deleted file mode 100755 index c6777dc71..000000000 --- a/scripts/hunyuan_image3/run_hunyuan_image3_i2i.sh +++ /dev/null @@ -1,27 +0,0 @@ -#!/bin/bash - -export lightx2v_path=${lightx2v_path:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)} -export model_path=${model_path:-/path/to/HunyuanImage-3-Instruct} -export image_path=${image_path:-/path/to/input.png} - -if [ -z "${CUDA_VISIBLE_DEVICES:-}" ] && command -v nvidia-smi >/dev/null 2>&1; then - gpu_count=$(nvidia-smi -L | wc -l) - if [ "$gpu_count" -gt 0 ]; then - export CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((gpu_count - 1))) - fi -fi - -source ${lightx2v_path}/scripts/base/base.sh - -config_json="${lightx2v_path}/configs/hunyuan_image3/hunyuan_image3_i2i.json" -echo "config_json=${config_json}" - -python -m lightx2v.infer \ - --model_cls hunyuan_image3 \ - --task i2i \ - --model_path $model_path \ - --config_json "$config_json" \ - --prompt "新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。" \ - --image_path "${image_path}" \ - --save_result_path ${lightx2v_path}/save_results/hunyuan_image3_i2i.png \ - --seed 42 diff --git a/scripts/hunyuan_image3/run_hunyuan_image3_t2i.sh b/scripts/hunyuan_image3/run_hunyuan_image3_t2i.sh deleted file mode 100755 index c0de39a54..000000000 --- a/scripts/hunyuan_image3/run_hunyuan_image3_t2i.sh +++ /dev/null @@ -1,43 +0,0 @@ -#!/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 -# HunyuanImage3 的模型和性能参数统一在 config_json 中设置。 - -GPU_IDS="${1:-${CUDA_VISIBLE_DEVICES:-}}" - -export lightx2v_path="${lightx2v_path:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" -export model_path="/data/nvme0/lhd_codes/HunyuanImage-3.0-instruct/HunyuanImage-3-Instruct" -export HUNYUAN_IMAGE3_REPO_PATH="/data/nvme0/lhd_codes/HunyuanImage-3.0" -export PYTHONPATH="${HUNYUAN_IMAGE3_REPO_PATH}:${PYTHONPATH:-}" - -if [ -n "$GPU_IDS" ]; then - export CUDA_VISIBLE_DEVICES="$GPU_IDS" -elif command -v nvidia-smi >/dev/null 2>&1; then - gpu_count=$(nvidia-smi -L | wc -l) - if [ "$gpu_count" -gt 0 ]; then - export CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((gpu_count - 1))) - fi -fi - -echo "Using CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-not set}" - -source ${lightx2v_path}/scripts/base/base.sh - -config_json="${lightx2v_path}/configs/hunyuan_image3/hunyuan_image3_t2i.json" -echo "config_json=${config_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 diff --git a/scripts/hunyuan_image3/run_hunyuan_image3_t2i_tp_sp_cfg_flashinfer.sh b/scripts/hunyuan_image3/run_hunyuan_image3_t2i_tp_sp_cfg_flashinfer.sh new file mode 100755 index 000000000..ba32c202f --- /dev/null +++ b/scripts/hunyuan_image3/run_hunyuan_image3_t2i_tp_sp_cfg_flashinfer.sh @@ -0,0 +1,133 @@ +#!/bin/bash +set -euo pipefail + +# Usage: +# bash scripts/hunyuan_image3/run_hunyuan_image3_t2i_tp_sp_cfg_flashinfer.sh +# bash scripts/hunyuan_image3/run_hunyuan_image3_t2i_tp_sp_cfg_flashinfer.sh 0,1,2,3,4,5,6,7 + +GPU_IDS="${1:-${CUDA_VISIBLE_DEVICES:-}}" + +export lightx2v_path="${lightx2v_path:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" +export model_path="${model_path:-/data/nvme0/lhd_codes/HunyuanImage-3.0-instruct/HunyuanImage-3-Instruct}" +export HUNYUAN_IMAGE3_REPO_PATH="${HUNYUAN_IMAGE3_REPO_PATH:-/data/nvme0/lhd_codes/HunyuanImage-3.0}" +export PYTHONPATH="${HUNYUAN_IMAGE3_REPO_PATH}:${PYTHONPATH:-}" + +config_json="${CONFIG_JSON:-${lightx2v_path}/configs/hunyuan_image3/hunyuan_image3_t2i_tp_sp_cfg_flashinfer.json}" + +parallel_values=$(python - "$config_json" "$model_path" "t2i" <<'PY' +import json +import sys +from pathlib import Path + +config_path = Path(sys.argv[1]) +model_path = Path(sys.argv[2]) +task = sys.argv[3] + +with config_path.open("r", encoding="utf-8") as f: + config = json.load(f) + +parallel = config.get("parallel") or {} +tp_size = int(parallel.get("tensor_p_size", 1)) +sp_size = int(parallel.get("seq_p_size", 1)) +cfg_size = int(parallel.get("cfg_p_size", 1)) +attn_type = str(parallel.get("seq_p_attn_type", "kv_all_gather")).strip().lower().replace("-", "_") +cfg_mode = str(parallel.get("cfg_mode", config.get("hunyuan_cfg_mode", "batch"))).strip().lower() +pipeline_parallel = parallel.get("pipeline_parallel", config.get("pipeline_parallel", True)) +moe_impl = str(config.get("moe_impl", "eager")).strip().lower() + +if attn_type in ("kv_allgather", "kv_gather"): + attn_type = "kv_all_gather" +if tp_size <= 1 or sp_size <= 1: + raise SystemExit(f"{task} TP+SP launcher requires tensor_p_size > 1 and seq_p_size > 1, got TP={tp_size}, SP={sp_size}.") +if pipeline_parallel: + raise SystemExit(f"{task} TP+SP launcher requires parallel.pipeline_parallel=false.") +if moe_impl != "flashinfer": + raise SystemExit(f"{task} FlashInfer launcher requires moe_impl='flashinfer', got {moe_impl!r}.") +if attn_type not in ("kv_all_gather", "ulysses"): + raise SystemExit(f"Unsupported seq_p_attn_type: {attn_type}.") +if not config.get("enable_cfg", False) or cfg_size != 2 or cfg_mode != "parallel": + raise SystemExit( + f"{task} CFG launcher requires enable_cfg=true, cfg_p_size=2, and cfg_mode='parallel'; " + f"got enable_cfg={config.get('enable_cfg', False)}, cfg_p_size={cfg_size}, cfg_mode={cfg_mode!r}." + ) + +if attn_type == "ulysses": + model_config_path = model_path / "config.json" + if not model_config_path.is_file(): + raise SystemExit(f"Cannot validate Ulysses heads because model config does not exist: {model_config_path}.") + with model_config_path.open("r", encoding="utf-8") as f: + model_config = json.load(f) + q_heads = int(model_config.get("num_attention_heads") or model_config["num_heads"]) + kv_heads = int(model_config.get("num_key_value_heads") or q_heads) + combined_size = tp_size * sp_size + if q_heads % combined_size or kv_heads % combined_size: + raise SystemExit( + "Ulysses requires TP_SIZE * SP_SIZE to divide Q and KV heads: " + f"Q={q_heads}, KV={kv_heads}, TP={tp_size}, SP={sp_size}." + ) + +print(tp_size, sp_size, cfg_size, attn_type) +PY +) +read -r TP_SIZE SP_SIZE CFG_SIZE SP_ATTN_TYPE <<< "$parallel_values" +WORLD_SIZE=$((TP_SIZE * SP_SIZE * CFG_SIZE)) + +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" -ge "$WORLD_SIZE" ]; then + export CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((WORLD_SIZE - 1))) + fi +fi + +if [ -z "${CUDA_VISIBLE_DEVICES:-}" ]; then + echo "No CUDA devices were selected; ${WORLD_SIZE} visible GPUs are required." >&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 + echo "CUDA_VISIBLE_DEVICES contains an empty GPU id." >&2 + exit 2 + 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 + +if [ "${#visible_gpu_ids[@]}" -ne "$WORLD_SIZE" ]; then + echo "HunyuanImage3 ${TP_SIZE}x${SP_SIZE}x${CFG_SIZE} parallelism requires exactly ${WORLD_SIZE} visible GPUs; CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" >&2 + exit 2 +fi +export CUDA_VISIBLE_DEVICES=$(IFS=,; echo "${visible_gpu_ids[*]}") + +prompt="${prompt:-生成图片：一辆汽车行驶在高速公路上，驾驶员在打电话，副驾驶坐着一只狗}" +seed="${seed:-42}" +save_result_path="${SAVE_RESULT_PATH:-${save_result_path:-${lightx2v_path}/save_results/hunyuan_image3_t2i_tp2_sp2_cfg2_flashinfer.png}}" + +echo "Using CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" +echo "Running HunyuanImage3 t2i with TP_SIZE=${TP_SIZE}, SP_SIZE=${SP_SIZE}, CFG_SIZE=${CFG_SIZE}, SP_ATTN_TYPE=${SP_ATTN_TYPE}" +echo "config_json=${config_json}" + +source "${lightx2v_path}/scripts/base/base.sh" + +torchrun \ + --standalone \ + --nproc_per_node="$WORLD_SIZE" \ + --max_restarts=0 \ + -m lightx2v.infer \ + --model_cls hunyuan_image3 \ + --task t2i \ + --model_path "$model_path" \ + --config_json "$config_json" \ + --prompt "$prompt" \ + --save_result_path "$save_result_path" \ + --seed "$seed" diff --git a/scripts/hunyuan_image3/run_hunyuan_image3_t2t_tp_sp_cfg_flashinfer.sh b/scripts/hunyuan_image3/run_hunyuan_image3_t2t_tp_sp_cfg_flashinfer.sh new file mode 100755 index 000000000..c5ecc212b --- /dev/null +++ b/scripts/hunyuan_image3/run_hunyuan_image3_t2t_tp_sp_cfg_flashinfer.sh @@ -0,0 +1,138 @@ +#!/bin/bash +set -euo pipefail + +# Usage: +# bash scripts/hunyuan_image3/run_hunyuan_image3_t2t_tp_sp_cfg_flashinfer.sh +# bash scripts/hunyuan_image3/run_hunyuan_image3_t2t_tp_sp_cfg_flashinfer.sh 0,1,2,3 +# +# T2T has no diffusion CFG branch. The cfg_p_size=1 dimension is retained in +# the config and filename so this launcher follows the common naming scheme. + +GPU_IDS="${1:-${CUDA_VISIBLE_DEVICES:-}}" + +export lightx2v_path="${lightx2v_path:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" +export model_path="${model_path:-/data/nvme0/lhd_codes/HunyuanImage-3.0-instruct/HunyuanImage-3-Instruct}" +export HUNYUAN_IMAGE3_REPO_PATH="${HUNYUAN_IMAGE3_REPO_PATH:-/data/nvme0/lhd_codes/HunyuanImage-3.0}" +export PYTHONPATH="${HUNYUAN_IMAGE3_REPO_PATH}:${PYTHONPATH:-}" + +config_json="${CONFIG_JSON:-${lightx2v_path}/configs/hunyuan_image3/hunyuan_image3_t2t_tp_sp_cfg_flashinfer.json}" + +parallel_values=$(python - "$config_json" "$model_path" "t2t" <<'PY' +import json +import sys +from pathlib import Path + +config_path = Path(sys.argv[1]) +model_path = Path(sys.argv[2]) +task = sys.argv[3] + +with config_path.open("r", encoding="utf-8") as f: + config = json.load(f) + +parallel = config.get("parallel") or {} +tp_size = int(parallel.get("tensor_p_size", 1)) +sp_size = int(parallel.get("seq_p_size", 1)) +cfg_size = int(parallel.get("cfg_p_size", 1)) +attn_type = str(parallel.get("seq_p_attn_type", "kv_all_gather")).strip().lower().replace("-", "_") +cfg_mode = str(parallel.get("cfg_mode", config.get("hunyuan_cfg_mode", "batch"))).strip().lower() +pipeline_parallel = parallel.get("pipeline_parallel", config.get("pipeline_parallel", True)) +moe_impl = str(config.get("moe_impl", "eager")).strip().lower() + +if attn_type in ("kv_allgather", "kv_gather"): + attn_type = "kv_all_gather" +if tp_size <= 1 or sp_size <= 1: + raise SystemExit(f"{task} TP+SP launcher requires tensor_p_size > 1 and seq_p_size > 1, got TP={tp_size}, SP={sp_size}.") +if pipeline_parallel: + raise SystemExit(f"{task} TP+SP launcher requires parallel.pipeline_parallel=false.") +if moe_impl != "flashinfer": + raise SystemExit(f"{task} FlashInfer launcher requires moe_impl='flashinfer', got {moe_impl!r}.") +if attn_type not in ("kv_all_gather", "ulysses"): + raise SystemExit(f"Unsupported seq_p_attn_type: {attn_type}.") +if config.get("enable_cfg", False) or cfg_size != 1: + raise SystemExit( + f"{task} does not support diffusion CFG and requires enable_cfg=false and cfg_p_size=1; " + f"got enable_cfg={config.get('enable_cfg', False)}, cfg_p_size={cfg_size}." + ) +if cfg_mode not in ("batch", "serial"): + raise SystemExit(f"{task} requires a non-parallel CFG mode, got cfg_mode={cfg_mode!r}.") + +if attn_type == "ulysses": + model_config_path = model_path / "config.json" + if not model_config_path.is_file(): + raise SystemExit(f"Cannot validate Ulysses heads because model config does not exist: {model_config_path}.") + with model_config_path.open("r", encoding="utf-8") as f: + model_config = json.load(f) + q_heads = int(model_config.get("num_attention_heads") or model_config["num_heads"]) + kv_heads = int(model_config.get("num_key_value_heads") or q_heads) + combined_size = tp_size * sp_size + if q_heads % combined_size or kv_heads % combined_size: + raise SystemExit( + "Ulysses requires TP_SIZE * SP_SIZE to divide Q and KV heads: " + f"Q={q_heads}, KV={kv_heads}, TP={tp_size}, SP={sp_size}." + ) + +print(tp_size, sp_size, cfg_size, attn_type) +PY +) +read -r TP_SIZE SP_SIZE CFG_SIZE SP_ATTN_TYPE <<< "$parallel_values" +WORLD_SIZE=$((TP_SIZE * SP_SIZE * CFG_SIZE)) + +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" -ge "$WORLD_SIZE" ]; then + export CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((WORLD_SIZE - 1))) + fi +fi + +if [ -z "${CUDA_VISIBLE_DEVICES:-}" ]; then + echo "No CUDA devices were selected; ${WORLD_SIZE} visible GPUs are required." >&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 + echo "CUDA_VISIBLE_DEVICES contains an empty GPU id." >&2 + exit 2 + 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 + +if [ "${#visible_gpu_ids[@]}" -ne "$WORLD_SIZE" ]; then + echo "HunyuanImage3 ${TP_SIZE}x${SP_SIZE}x${CFG_SIZE} parallelism requires exactly ${WORLD_SIZE} visible GPUs; CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" >&2 + exit 2 +fi +export CUDA_VISIBLE_DEVICES=$(IFS=,; echo "${visible_gpu_ids[*]}") + +prompt="${prompt:-请简要介绍这张图像生成模型的主要能力。}" +seed="${seed:-42}" +save_result_path="${SAVE_RESULT_PATH:-${save_result_path:-${lightx2v_path}/save_results/hunyuan_image3_t2t_tp2_sp2_cfg1_flashinfer.txt}}" + +echo "Using CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" +echo "Running HunyuanImage3 t2t with TP_SIZE=${TP_SIZE}, SP_SIZE=${SP_SIZE}, CFG_SIZE=${CFG_SIZE}, SP_ATTN_TYPE=${SP_ATTN_TYPE}" +echo "config_json=${config_json}" + +source "${lightx2v_path}/scripts/base/base.sh" + +torchrun \ + --standalone \ + --nproc_per_node="$WORLD_SIZE" \ + --max_restarts=0 \ + -m lightx2v.infer \ + --model_cls hunyuan_image3 \ + --task t2t \ + --model_path "$model_path" \ + --config_json "$config_json" \ + --prompt "$prompt" \ + --save_result_path "$save_result_path" \ + --seed "$seed" diff --git a/scripts/hunyuan_image3/run_hunyuan_image3_ti2i.sh b/scripts/hunyuan_image3/run_hunyuan_image3_ti2i.sh deleted file mode 100755 index 960dd55bb..000000000 --- a/scripts/hunyuan_image3/run_hunyuan_image3_ti2i.sh +++ /dev/null @@ -1,69 +0,0 @@ -#!/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}" - -source "${lightx2v_path}/scripts/base/base.sh" - -config_json="${lightx2v_path}/configs/hunyuan_image3/hunyuan_image3_ti2i.json" -echo "config_json=${config_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" diff --git a/scripts/hunyuan_image3/run_hunyuan_image3_ti2i_tp_sp_cfg_flashinfer.sh b/scripts/hunyuan_image3/run_hunyuan_image3_ti2i_tp_sp_cfg_flashinfer.sh new file mode 100755 index 000000000..8f2e5cbd7 --- /dev/null +++ b/scripts/hunyuan_image3/run_hunyuan_image3_ti2i_tp_sp_cfg_flashinfer.sh @@ -0,0 +1,169 @@ +#!/bin/bash +set -euo pipefail + +# Usage: +# bash scripts/hunyuan_image3/run_hunyuan_image3_ti2i_tp_sp_cfg_flashinfer.sh +# bash scripts/hunyuan_image3/run_hunyuan_image3_ti2i_tp_sp_cfg_flashinfer.sh ti2iv2 +# bash scripts/hunyuan_image3/run_hunyuan_image3_ti2i_tp_sp_cfg_flashinfer.sh 0,1,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 + +if [ "$DEMO" != "ti2i" ] && [ "$DEMO" != "ti2iv2" ]; then + echo "DEMO must be either ti2i or ti2iv2, got: $DEMO" >&2 + exit 2 +fi + +export lightx2v_path="${lightx2v_path:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" +export model_path="${model_path:-/data/nvme0/lhd_codes/HunyuanImage-3.0-instruct/HunyuanImage-3-Instruct}" +export HUNYUAN_IMAGE3_REPO_PATH="${HUNYUAN_IMAGE3_REPO_PATH:-/data/nvme0/lhd_codes/HunyuanImage-3.0}" +export PYTHONPATH="${HUNYUAN_IMAGE3_REPO_PATH}:${PYTHONPATH:-}" + +config_json="${CONFIG_JSON:-${lightx2v_path}/configs/hunyuan_image3/hunyuan_image3_ti2i_tp_sp_cfg_flashinfer.json}" + +parallel_values=$(python - "$config_json" "$model_path" "ti2i" <<'PY' +import json +import sys +from pathlib import Path + +config_path = Path(sys.argv[1]) +model_path = Path(sys.argv[2]) +task = sys.argv[3] + +with config_path.open("r", encoding="utf-8") as f: + config = json.load(f) + +parallel = config.get("parallel") or {} +tp_size = int(parallel.get("tensor_p_size", 1)) +sp_size = int(parallel.get("seq_p_size", 1)) +cfg_size = int(parallel.get("cfg_p_size", 1)) +attn_type = str(parallel.get("seq_p_attn_type", "kv_all_gather")).strip().lower().replace("-", "_") +cfg_mode = str(parallel.get("cfg_mode", config.get("hunyuan_cfg_mode", "batch"))).strip().lower() +pipeline_parallel = parallel.get("pipeline_parallel", config.get("pipeline_parallel", True)) +moe_impl = str(config.get("moe_impl", "eager")).strip().lower() + +if attn_type in ("kv_allgather", "kv_gather"): + attn_type = "kv_all_gather" +if tp_size <= 1 or sp_size <= 1: + raise SystemExit(f"{task} TP+SP launcher requires tensor_p_size > 1 and seq_p_size > 1, got TP={tp_size}, SP={sp_size}.") +if pipeline_parallel: + raise SystemExit(f"{task} TP+SP launcher requires parallel.pipeline_parallel=false.") +if moe_impl != "flashinfer": + raise SystemExit(f"{task} FlashInfer launcher requires moe_impl='flashinfer', got {moe_impl!r}.") +if attn_type not in ("kv_all_gather", "ulysses"): + raise SystemExit(f"Unsupported seq_p_attn_type: {attn_type}.") +if not config.get("enable_cfg", False) or cfg_size != 2 or cfg_mode != "parallel": + raise SystemExit( + f"{task} CFG launcher requires enable_cfg=true, cfg_p_size=2, and cfg_mode='parallel'; " + f"got enable_cfg={config.get('enable_cfg', False)}, cfg_p_size={cfg_size}, cfg_mode={cfg_mode!r}." + ) + +if attn_type == "ulysses": + model_config_path = model_path / "config.json" + if not model_config_path.is_file(): + raise SystemExit(f"Cannot validate Ulysses heads because model config does not exist: {model_config_path}.") + with model_config_path.open("r", encoding="utf-8") as f: + model_config = json.load(f) + q_heads = int(model_config.get("num_attention_heads") or model_config["num_heads"]) + kv_heads = int(model_config.get("num_key_value_heads") or q_heads) + combined_size = tp_size * sp_size + if q_heads % combined_size or kv_heads % combined_size: + raise SystemExit( + "Ulysses requires TP_SIZE * SP_SIZE to divide Q and KV heads: " + f"Q={q_heads}, KV={kv_heads}, TP={tp_size}, SP={sp_size}." + ) + +print(tp_size, sp_size, cfg_size, attn_type) +PY +) +read -r TP_SIZE SP_SIZE CFG_SIZE SP_ATTN_TYPE <<< "$parallel_values" +WORLD_SIZE=$((TP_SIZE * SP_SIZE * CFG_SIZE)) + +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" -ge "$WORLD_SIZE" ]; then + export CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((WORLD_SIZE - 1))) + fi +fi + +if [ -z "${CUDA_VISIBLE_DEVICES:-}" ]; then + echo "No CUDA devices were selected; ${WORLD_SIZE} visible GPUs are required." >&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 + echo "CUDA_VISIBLE_DEVICES contains an empty GPU id." >&2 + exit 2 + 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 + +if [ "${#visible_gpu_ids[@]}" -ne "$WORLD_SIZE" ]; then + echo "HunyuanImage3 ${TP_SIZE}x${SP_SIZE}x${CFG_SIZE} parallelism requires exactly ${WORLD_SIZE} visible GPUs; CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" >&2 + exit 2 +fi +export CUDA_VISIBLE_DEVICES=$(IFS=,; echo "${visible_gpu_ids[*]}") + +if [ "$DEMO" = "ti2iv2" ]; then + default_prompt="让图1的猫咪与图2的猫咪自拍，图1的猫咪说:“妈妈，我在乡下遇到了好朋喵”，背景为图3。" + default_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" + default_seed=43 + default_save_result_path="${lightx2v_path}/save_results/hunyuan_image3_ti2iv2_tp2_sp2_cfg2_flashinfer.png" +else + default_prompt="新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。 宝丽莱相纸，超现实主义，写实主义，胶片摄影，打印颗粒感肌理。肌理，超写实，复古感。" + default_image_path="${HUNYUAN_IMAGE3_REPO_PATH}/assets/demo_instruct_imgs/input_0_0.png" + default_seed=42 + default_save_result_path="${lightx2v_path}/save_results/hunyuan_image3_ti2i_tp2_sp2_cfg2_flashinfer.png" +fi + +prompt="${prompt:-$default_prompt}" +image_path="${image_path:-$default_image_path}" +seed="${seed:-$default_seed}" +save_result_path="${SAVE_RESULT_PATH:-${save_result_path:-$default_save_result_path}}" + +echo "Using CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" +echo "Running HunyuanImage3 ${DEMO} with TP_SIZE=${TP_SIZE}, SP_SIZE=${SP_SIZE}, CFG_SIZE=${CFG_SIZE}, SP_ATTN_TYPE=${SP_ATTN_TYPE}" +echo "config_json=${config_json}" + +source "${lightx2v_path}/scripts/base/base.sh" + +torchrun \ + --standalone \ + --nproc_per_node="$WORLD_SIZE" \ + --max_restarts=0 \ + -m lightx2v.infer \ + --model_cls hunyuan_image3 \ + --task ti2i \ + --model_path "$model_path" \ + --config_json "$config_json" \ + --prompt "$prompt" \ + --image_path "$image_path" \ + --save_result_path "$save_result_path" \ + --seed "$seed" diff --git a/scripts/hunyuan_image3/run_hunyuan_image3_ti2t_tp_sp_cfg_flashinfer.sh b/scripts/hunyuan_image3/run_hunyuan_image3_ti2t_tp_sp_cfg_flashinfer.sh new file mode 100755 index 000000000..eabb42381 --- /dev/null +++ b/scripts/hunyuan_image3/run_hunyuan_image3_ti2t_tp_sp_cfg_flashinfer.sh @@ -0,0 +1,140 @@ +#!/bin/bash +set -euo pipefail + +# Usage: +# bash scripts/hunyuan_image3/run_hunyuan_image3_ti2t_tp_sp_cfg_flashinfer.sh +# bash scripts/hunyuan_image3/run_hunyuan_image3_ti2t_tp_sp_cfg_flashinfer.sh 0,1,2,3 +# +# TI2T has no diffusion CFG branch. The cfg_p_size=1 dimension is retained in +# the config and filename so this launcher follows the common naming scheme. + +GPU_IDS="${1:-${CUDA_VISIBLE_DEVICES:-}}" + +export lightx2v_path="${lightx2v_path:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}" +export model_path="${model_path:-/data/nvme0/lhd_codes/HunyuanImage-3.0-instruct/HunyuanImage-3-Instruct}" +export HUNYUAN_IMAGE3_REPO_PATH="${HUNYUAN_IMAGE3_REPO_PATH:-/data/nvme0/lhd_codes/HunyuanImage-3.0}" +export PYTHONPATH="${HUNYUAN_IMAGE3_REPO_PATH}:${PYTHONPATH:-}" + +config_json="${CONFIG_JSON:-${lightx2v_path}/configs/hunyuan_image3/hunyuan_image3_ti2t_tp_sp_cfg_flashinfer.json}" + +parallel_values=$(python - "$config_json" "$model_path" "ti2t" <<'PY' +import json +import sys +from pathlib import Path + +config_path = Path(sys.argv[1]) +model_path = Path(sys.argv[2]) +task = sys.argv[3] + +with config_path.open("r", encoding="utf-8") as f: + config = json.load(f) + +parallel = config.get("parallel") or {} +tp_size = int(parallel.get("tensor_p_size", 1)) +sp_size = int(parallel.get("seq_p_size", 1)) +cfg_size = int(parallel.get("cfg_p_size", 1)) +attn_type = str(parallel.get("seq_p_attn_type", "kv_all_gather")).strip().lower().replace("-", "_") +cfg_mode = str(parallel.get("cfg_mode", config.get("hunyuan_cfg_mode", "batch"))).strip().lower() +pipeline_parallel = parallel.get("pipeline_parallel", config.get("pipeline_parallel", True)) +moe_impl = str(config.get("moe_impl", "eager")).strip().lower() + +if attn_type in ("kv_allgather", "kv_gather"): + attn_type = "kv_all_gather" +if tp_size <= 1 or sp_size <= 1: + raise SystemExit(f"{task} TP+SP launcher requires tensor_p_size > 1 and seq_p_size > 1, got TP={tp_size}, SP={sp_size}.") +if pipeline_parallel: + raise SystemExit(f"{task} TP+SP launcher requires parallel.pipeline_parallel=false.") +if moe_impl != "flashinfer": + raise SystemExit(f"{task} FlashInfer launcher requires moe_impl='flashinfer', got {moe_impl!r}.") +if attn_type not in ("kv_all_gather", "ulysses"): + raise SystemExit(f"Unsupported seq_p_attn_type: {attn_type}.") +if config.get("enable_cfg", False) or cfg_size != 1: + raise SystemExit( + f"{task} does not support diffusion CFG and requires enable_cfg=false and cfg_p_size=1; " + f"got enable_cfg={config.get('enable_cfg', False)}, cfg_p_size={cfg_size}." + ) +if cfg_mode not in ("batch", "serial"): + raise SystemExit(f"{task} requires a non-parallel CFG mode, got cfg_mode={cfg_mode!r}.") + +if attn_type == "ulysses": + model_config_path = model_path / "config.json" + if not model_config_path.is_file(): + raise SystemExit(f"Cannot validate Ulysses heads because model config does not exist: {model_config_path}.") + with model_config_path.open("r", encoding="utf-8") as f: + model_config = json.load(f) + q_heads = int(model_config.get("num_attention_heads") or model_config["num_heads"]) + kv_heads = int(model_config.get("num_key_value_heads") or q_heads) + combined_size = tp_size * sp_size + if q_heads % combined_size or kv_heads % combined_size: + raise SystemExit( + "Ulysses requires TP_SIZE * SP_SIZE to divide Q and KV heads: " + f"Q={q_heads}, KV={kv_heads}, TP={tp_size}, SP={sp_size}." + ) + +print(tp_size, sp_size, cfg_size, attn_type) +PY +) +read -r TP_SIZE SP_SIZE CFG_SIZE SP_ATTN_TYPE <<< "$parallel_values" +WORLD_SIZE=$((TP_SIZE * SP_SIZE * CFG_SIZE)) + +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" -ge "$WORLD_SIZE" ]; then + export CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((WORLD_SIZE - 1))) + fi +fi + +if [ -z "${CUDA_VISIBLE_DEVICES:-}" ]; then + echo "No CUDA devices were selected; ${WORLD_SIZE} visible GPUs are required." >&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 + echo "CUDA_VISIBLE_DEVICES contains an empty GPU id." >&2 + exit 2 + 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 + +if [ "${#visible_gpu_ids[@]}" -ne "$WORLD_SIZE" ]; then + echo "HunyuanImage3 ${TP_SIZE}x${SP_SIZE}x${CFG_SIZE} parallelism requires exactly ${WORLD_SIZE} visible GPUs; CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" >&2 + exit 2 +fi +export CUDA_VISIBLE_DEVICES=$(IFS=,; echo "${visible_gpu_ids[*]}") + +image_path="${image_path:-${HUNYUAN_IMAGE3_REPO_PATH}/assets/demo_instruct_imgs/input_0_0.png}" +prompt="${prompt:-请描述图像中的主要内容和视觉风格。}" +seed="${seed:-42}" +save_result_path="${SAVE_RESULT_PATH:-${save_result_path:-${lightx2v_path}/save_results/hunyuan_image3_ti2t_tp2_sp2_cfg1_flashinfer.txt}}" + +echo "Using CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" +echo "Running HunyuanImage3 ti2t with TP_SIZE=${TP_SIZE}, SP_SIZE=${SP_SIZE}, CFG_SIZE=${CFG_SIZE}, SP_ATTN_TYPE=${SP_ATTN_TYPE}" +echo "config_json=${config_json}" + +source "${lightx2v_path}/scripts/base/base.sh" + +torchrun \ + --standalone \ + --nproc_per_node="$WORLD_SIZE" \ + --max_restarts=0 \ + -m lightx2v.infer \ + --model_cls hunyuan_image3 \ + --task ti2t \ + --model_path "$model_path" \ + --config_json "$config_json" \ + --prompt "$prompt" \ + --image_path "$image_path" \ + --save_result_path "$save_result_path" \ + --seed "$seed"