From c3376a466b67db11d5ac4abd9bc19f3f53eef145 Mon Sep 17 00:00:00 2001 From: Ding Zuhao Date: Wed, 1 Apr 2026 16:33:02 +0800 Subject: [PATCH] [Feature][HunyuanImage3.0] Add cfgP to HunyuanImage3.0 (#1751) Signed-off-by: Ding Zuhao --- .../hunyuan_image_3_transformer.py | 103 ++++++++++++++++-- .../pipeline_hunyuan_image_3.py | 5 +- 2 files changed, 96 insertions(+), 12 deletions(-) diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py index d18913723..3d670809b 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py @@ -61,6 +61,9 @@ from vllm_omni.diffusion.attention.backends.abstract import ( ) from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.distributed.parallel_state import ( + get_cfg_group, + get_classifier_free_guidance_rank, + get_classifier_free_guidance_world_size, get_pp_group, get_sequence_parallel_rank, get_sequence_parallel_world_size, @@ -2535,6 +2538,61 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): def set_scheduler(self, new_scheduler): self.register_modules(scheduler=new_scheduler) + @staticmethod + def _split_model_kwargs_for_cfg_parallel(model_kwargs: dict[str, Any], batch_size: int, cfg_rank: int) -> None: + """Split batch-doubled model_kwargs in-place for CFG parallel. + + The tokenizer produces inputs with cfg_factor=2, so all batch-dim + tensors have shape [batch_size*2, ...]. This method slices them + so that rank 0 gets the conditioned half and rank 1 gets the + unconditioned half. + """ + s = slice(cfg_rank * batch_size, (cfg_rank + 1) * batch_size) + + # Tensor fields with leading batch dimension + tensor_keys = [ + "position_ids", + "image_mask", + "gen_timestep_scatter_index", + "cond_vae_image_mask", + "cond_vit_image_mask", + "cond_timestep_scatter_index", + ] + for key in tensor_keys: + if key in model_kwargs and model_kwargs[key] is not None: + model_kwargs[key] = model_kwargs[key][s] + + # custom_pos_emb: tuple of (cos, sin) + if "custom_pos_emb" in model_kwargs and model_kwargs["custom_pos_emb"] is not None: + cos, sin = model_kwargs["custom_pos_emb"] + model_kwargs["custom_pos_emb"] = (cos[s], sin[s]) + + # cond_vae_images: tensor or list + if model_kwargs.get("cond_vae_images") is not None: + v = model_kwargs["cond_vae_images"] + if isinstance(v, torch.Tensor): + model_kwargs["cond_vae_images"] = v[s] + elif isinstance(v, list): + model_kwargs["cond_vae_images"] = v[s.start : s.stop] + + # cond_timestep: tensor or list + if model_kwargs.get("cond_timestep") is not None: + v = model_kwargs["cond_timestep"] + if isinstance(v, torch.Tensor): + model_kwargs["cond_timestep"] = v[s] + elif isinstance(v, list): + model_kwargs["cond_timestep"] = v[s.start : s.stop] + + # cond_vit_images: list of tensors + if model_kwargs.get("cond_vit_images") is not None: + model_kwargs["cond_vit_images"] = model_kwargs["cond_vit_images"][s.start : s.stop] + + # vit_kwargs: dict of lists + if model_kwargs.get("vit_kwargs") is not None: + model_kwargs["vit_kwargs"] = { + k: v[s.start : s.stop] if isinstance(v, list) else v[s] for k, v in model_kwargs["vit_kwargs"].items() + } + @torch.no_grad() def __call__( self, @@ -2621,7 +2679,8 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): self._guidance_scale = guidance_scale self._guidance_rescale = guidance_rescale - cfg_factor = 1 + self.do_classifier_free_guidance + # Detect CFG parallel configuration (only 2-branch layout is supported) + cfg_parallel_ready = self.do_classifier_free_guidance and get_classifier_free_guidance_world_size() == 2 # Define call parameters device = self._execution_device @@ -2649,13 +2708,33 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): # Prepare extra step kwargs. _scheduler_step_extra_kwargs = self.prepare_extra_func_kwargs(self.scheduler.step, {"generator": generator}) - # Prepare model kwargs + # Prepare model kwargs — attention mask is built from the full + # (cfg_factor=2) batch before any splitting so that each rank's + # slice is correct. input_ids = model_kwargs.pop("input_ids") attention_mask = self.model._prepare_attention_mask_for_generation( # noqa input_ids, self.model.generation_config, model_kwargs=model_kwargs, ) + + # Split inputs for CFG parallel: each rank processes only its branch. + if cfg_parallel_ready: + cfg_group = get_cfg_group() + cfg_rank = get_classifier_free_guidance_rank() + + # Ensure all ranks start with the same latents + latents = latents.contiguous() + cfg_group.broadcast(latents, src=0) + + # Split batch-doubled tensors: rank 0 → conditioned, rank 1 → unconditioned + s = slice(cfg_rank * batch_size, (cfg_rank + 1) * batch_size) + input_ids = input_ids[s] + attention_mask = attention_mask[s] + self._split_model_kwargs_for_cfg_parallel(model_kwargs, batch_size, cfg_rank) + else: + cfg_factor = 1 + self.do_classifier_free_guidance + b, _, q_len1, seq_len = attention_mask.shape query_lens = [q_len1] * b seq_lens = [seq_len] * b @@ -2678,9 +2757,12 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * cfg_factor) - # latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + if cfg_parallel_ready: + # CFG parallel: each rank forwards its own branch (no batch doubling) + latent_model_input = latents + else: + # Sequential CFG: double the batch + latent_model_input = torch.cat([latents] * cfg_factor) t_expand = t.repeat(latent_model_input.shape[0]) @@ -2721,14 +2803,17 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): # TeaCache fast path: reuse previous prediction pred = tc_prev_pred - # perform guidance - if self.do_classifier_free_guidance: + # Perform guidance + if cfg_parallel_ready: + # CFG parallel: all_gather → all ranks combine locally (no broadcast needed) + gathered = cfg_group.all_gather(pred, separate_tensors=True) + pred = self.cfg_operator(gathered[0], gathered[1], self.guidance_scale, step=i) + elif self.do_classifier_free_guidance: pred_cond, pred_uncond = pred.chunk(2) pred = self.cfg_operator(pred_cond, pred_uncond, self.guidance_scale, step=i) - # compute the previous noisy sample x_t -> x_t-1 + # Scheduler step (all ranks compute locally in CFG parallel) latents = self.scheduler.step(pred, t, latents, **_scheduler_step_extra_kwargs, return_dict=False)[0] - if i != len(timesteps) - 1 and should_compute: model_kwargs = self.model._update_model_kwargs_for_generation( # noqa model_output, diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py b/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py index c19e8a65a..ba24818dc 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py +++ b/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py @@ -149,13 +149,12 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin, Diffu "time_embed_2", "final_layer.model", ] - - device_str = f"{get_local_device()}" + device = get_local_device() named_modules = dict(self.named_modules()) for prefix in non_model_layer_prefixes: mod = named_modules.get(prefix) if mod: - mod.to(device_str) + mod.to(device) unexpected_keywords = [ "guidance_emb",