SIGN IN SIGN UP

[core] Fix component device placement for split-device pipelines (#14739)

* fix: component device placement for split-device pipelines

DiffusionPipeline.device/`_execution_device` (#14383) already prefer any
non-CPU, non-meta component so a pipeline with components split across
devices (e.g. text encoder on CPU, denoising backbone on an accelerator)
reports the right execution device. But several component call sites
still moved a tensor to that execution `device` before handing it to a
specific submodule, instead of the submodule's own device -- so encoding
a prompt, encoding an IP-Adapter image, or decoding latents would crash
with a device-mismatch error as soon as that submodule wasn't already
sitting on the same device as everything else.

Fixes three component types across the major pipeline families
(StableDiffusionPipeline, StableDiffusionXLPipeline,
StableDiffusion3Pipeline, FluxPipeline, Flux2(Klein)Pipeline, WanPipeline,
WanImageToVideoPipeline):

- Text encoder(s): `encode_prompt` / `_get_t5_prompt_embeds` /
  `_get_clip_prompt_embeds` now read the encoder's own `.device` instead
  of trusting the passed-in `device`.
- VAE: the `vae.encode`/`vae.decode` calls in each `__call__`, plus
  `Flux2Pipeline.prepare_image_latents` (whose output also wasn't being
  moved back to the execution device).
- image_encoder (IP-Adapter): `encode_image`, plus
  `run_safety_checker`'s feature-extractor input.

Also fixes a cross-device bug in `WanImageToVideoPipeline.prepare_latents`
and `Flux2KleinInpaintPipeline.prepare_latents`, where a vae-encoded
tensor and the freshly-sampled noise tensor could end up on different
devices before being combined by the scheduler.

Fixed the canonical/root pipelines by hand, then ran `make fix-copies` to
propagate the text-encoder/image-encoder fixes to every `# Copied from`
derivative.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* review:reduce scope

* fix: fix not failing test while offloading

---------

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
J
Jingya HUANG committed
9f1246971270c84dcbe71233edb7a519596a5d02
Parent: 80c7ed2
Committed by GitHub <noreply@github.com> on 9/21/2026, 7:17:32 AM