From bb56997d4b7e87f0743f26a612f49ec4e7ce7213 Mon Sep 17 00:00:00 2001 From: "Wang, Yi" Date: Thu, 13 Aug 2026 15:36:08 +0800 Subject: [PATCH] =?UTF-8?q?Fix=20AuraFlow=20model=20parallelism=20device?= =?UTF-8?q?=20mismatch=20and=20update=20XPU=20IP-Adap=E2=80=A6=20(#14273)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Fix AuraFlow model parallelism device mismatch and update XPU IP-Adapter expectations Description This PR fixes two test issues currently reflected in the working tree. Changes Fixed AuraFlow model parallelism in auraflow_transformer_2d.py. Under device_map="auto", context_embedder can place encoder_hidden_states on a different device from the top-level register_tokens parameter. The forward pass now moves register_tokens to the projected encoder hidden states device and dtype before concatenation, avoiding cross-device torch.cat failures. Updated XPU-specific expected slices in test_ip_adapter_stable_diffusion.py. The IP-Adapter SD integration expectation now distinguishes XPU backend versions 3 and 5, since they produce slightly different deterministic slices. Signed-off-by: Wang, Yi A * update Signed-off-by: Wang, Yi A * update Signed-off-by: Wang, Yi A * fmt Signed-off-by: Wang, Yi A --------- Signed-off-by: Wang, Yi A --- .../models/transformers/auraflow_transformer_2d.py | 8 ++++++-- .../ip_adapters/test_ip_adapter_stable_diffusion.py | 5 +---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/src/diffusers/models/transformers/auraflow_transformer_2d.py b/src/diffusers/models/transformers/auraflow_transformer_2d.py index ff6c0c78a5..cdc76c5ec0 100644 --- a/src/diffusers/models/transformers/auraflow_transformer_2d.py +++ b/src/diffusers/models/transformers/auraflow_transformer_2d.py @@ -22,7 +22,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import apply_lora_scale, logging -from ...utils.torch_utils import maybe_allow_in_graph +from ...utils.torch_utils import maybe_adjust_dtype_for_device, maybe_allow_in_graph from ..attention import AttentionMixin from ..attention_processor import ( Attention, @@ -435,8 +435,12 @@ class AuraFlowTransformer2DModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAd temb = self.time_step_embed(timestep).to(dtype=next(self.parameters()).dtype) temb = self.time_step_proj(temb) encoder_hidden_states = self.context_embedder(encoder_hidden_states) + register_tokens_dtype = maybe_adjust_dtype_for_device( + encoder_hidden_states.dtype, encoder_hidden_states.device + ) + register_tokens = self.register_tokens.to(device=encoder_hidden_states.device, dtype=register_tokens_dtype) encoder_hidden_states = torch.cat( - [self.register_tokens.repeat(encoder_hidden_states.size(0), 1, 1), encoder_hidden_states], dim=1 + [register_tokens.repeat(encoder_hidden_states.size(0), 1, 1), encoder_hidden_states], dim=1 ) # MMDiT blocks. diff --git a/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py b/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py index 056521c7f5..86519fea0d 100644 --- a/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py +++ b/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py @@ -223,13 +223,10 @@ class IPAdapterSDIntegrationTests(IPAdapterNightlyTestsMixin): with self.get_fixed_randn_tensor_patch(pipeline): images = pipeline(**inputs).images image_slice = images[0, :3, :3, -1].flatten() - expected_slice = Expectations( { ("cuda", None): np.array([0.1238, 0.0579, 0.0312, 0.0493, 0.0010, 0.0, 0.0188, 0.0, 0.0]), - ("xpu", None): np.array( - [0.11938477, 0.05249023, 0.02490234, 0.04370117, 0.0, 0.0, 0.01342773, 0.0, 0.0] - ), + ("xpu", 5): np.array([0.1228, 0.0566, 0.0286, 0.0479, 0.0, 0.0, 0.0173, 0.0, 0.0]), (None, None): np.array([0.1238, 0.0579, 0.0312, 0.0493, 0.0010, 0.0, 0.0188, 0.0, 0.0]), } ).get_expectation()