AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-23 09:09:01 +03:00
Code Activity

Fix AuraFlow model parallelism device mismatch and update XPU IP-Adap… (#14273)

Browse Source 
* 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 <yi.a.wang@intel.com>

* update

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

* update

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

* fmt

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

---------

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
This commit is contained in:
Wang, Yi
2026-08-13 15:36:08 +08:00
committed by GitHub
parent 9558d918ba
commit bb56997d4b
2 changed files with 7 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/diffusers/models/transformers/auraflow_transformer_2d.py
View File

@@ -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.

5
tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py
View File

@@ -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()