From f62d4ad4a794bd0706340d8dcc9d92e67672af1c Mon Sep 17 00:00:00 2001 From: hyw <109567717+yuweih205@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:15:32 +0800 Subject: [PATCH] [Model][ERNIE-Image] Delay AdaLN modulation broadcast (#7171) Signed-off-by: hyw --- .../models/ernie_image/ernie_image_transformer.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py index 244aeeb95..bd96630e8 100644 --- a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py +++ b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_ernie_image.py from collections.abc import Iterable @@ -599,17 +599,12 @@ class ErnieImageTransformer2DModel(nn.Module): hidden_states, freqs_cos, freqs_sin, attention_mask = self.unified_prepare(hidden_states, text_bth, text_lens) attention_mask = self._slice_attention_mask_for_ring(attention_mask) rotary_pos_emb = (freqs_cos, freqs_sin) - S = hidden_states.shape[0] - sample = self.time_proj(timestep) sample = sample.to(dtype=dtype) c = self.time_embedding(sample) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = [ - t.unsqueeze(0).expand(S, -1, -1).contiguous() for t in self.adaLN_modulation(c).chunk(6, dim=-1) - ] + temb = self.adaLN_modulation(c).chunk(6, dim=-1) for layer in self.layers: - temb = [shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp] if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func( layer,