Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JingyaHuang wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions src/diffusers/modular_pipelines/flux2/encoders.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -330,8 +330,9 @@ def _get_qwen3_prompt_embeds(
all_input_ids.append(inputs["input_ids"])
all_attention_masks.append(inputs["attention_mask"])

input_ids = torch.cat(all_input_ids, dim=0).to(device)
attention_mask = torch.cat(all_attention_masks, dim=0).to(device)
model_device = text_encoder.device
input_ids = torch.cat(all_input_ids, dim=0).to(model_device)
attention_mask = torch.cat(all_attention_masks, dim=0).to(model_device)

# Forward pass through the model
output = text_encoder(
Expand Down Expand Up @@ -471,8 +472,9 @@ def _get_qwen3_prompt_embeds(
all_input_ids.append(inputs["input_ids"])
all_attention_masks.append(inputs["attention_mask"])

input_ids = torch.cat(all_input_ids, dim=0).to(device)
attention_mask = torch.cat(all_attention_masks, dim=0).to(device)
model_device = text_encoder.device
input_ids = torch.cat(all_input_ids, dim=0).to(model_device)
attention_mask = torch.cat(all_attention_masks, dim=0).to(model_device)

# Forward pass through the model
output = text_encoder(
Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@ def encode_image(components, image, device, num_images_per_prompt, output_hidden
if not isinstance(image, torch.Tensor):
image = components.feature_extractor(image, return_tensors="pt").pixel_values

image = image.to(device=device, dtype=dtype)
image = image.to(device=components.image_encoder.device, dtype=dtype)
if output_hidden_states:
image_enc_hidden_states = components.image_encoder(image, output_hidden_states=True).hidden_states[-2]
image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
Expand Down
14 changes: 8 additions & 6 deletions src/diffusers/pipelines/animatediff/pipeline_animatediff.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,17 +236,18 @@ def encode_prompt(
f" {self.tokenizer.model_max_length} tokens: {removed_text}"
)

model_device = self.text_encoder.device
if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
attention_mask = text_inputs.attention_mask.to(device)
attention_mask = text_inputs.attention_mask.to(model_device)
else:
attention_mask = None

if clip_skip is None:
prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)
prompt_embeds = self.text_encoder(text_input_ids.to(model_device), attention_mask=attention_mask)
prompt_embeds = prompt_embeds[0]
else:
prompt_embeds = self.text_encoder(
text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True
text_input_ids.to(model_device), attention_mask=attention_mask, output_hidden_states=True
)
# Access the `hidden_states` first, that contains a tuple of
# all the hidden states from the encoder layers. Then index into
Expand Down Expand Up @@ -306,13 +307,14 @@ def encode_prompt(
return_tensors="pt",
)

model_device = self.text_encoder.device
if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
attention_mask = uncond_input.attention_mask.to(device)
attention_mask = uncond_input.attention_mask.to(model_device)
else:
attention_mask = None

negative_prompt_embeds = self.text_encoder(
uncond_input.input_ids.to(device),
uncond_input.input_ids.to(model_device),
attention_mask=attention_mask,
)
negative_prompt_embeds = negative_prompt_embeds[0]
Expand Down Expand Up @@ -340,7 +342,7 @@ def encode_image(self, image, device, num_images_per_prompt, output_hidden_state
if not isinstance(image, torch.Tensor):
image = self.feature_extractor(image, return_tensors="pt").pixel_values

image = image.to(device=device, dtype=dtype)
image = image.to(device=self.image_encoder.device, dtype=dtype)
if output_hidden_states:
image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
Expand Down
14 changes: 8 additions & 6 deletions src/diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -281,17 +281,18 @@ def encode_prompt(
f" {self.tokenizer.model_max_length} tokens: {removed_text}"
)

model_device = self.text_encoder.device
if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
attention_mask = text_inputs.attention_mask.to(device)
attention_mask = text_inputs.attention_mask.to(model_device)
else:
attention_mask = None

if clip_skip is None:
prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)
prompt_embeds = self.text_encoder(text_input_ids.to(model_device), attention_mask=attention_mask)
prompt_embeds = prompt_embeds[0]
else:
prompt_embeds = self.text_encoder(
text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True
text_input_ids.to(model_device), attention_mask=attention_mask, output_hidden_states=True
)
# Access the `hidden_states` first, that contains a tuple of
# all the hidden states from the encoder layers. Then index into
Expand Down Expand Up @@ -351,13 +352,14 @@ def encode_prompt(
return_tensors="pt",
)

model_device = self.text_encoder.device
if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
attention_mask = uncond_input.attention_mask.to(device)
attention_mask = uncond_input.attention_mask.to(model_device)
else:
attention_mask = None

negative_prompt_embeds = self.text_encoder(
uncond_input.input_ids.to(device),
uncond_input.input_ids.to(model_device),
attention_mask=attention_mask,
)
negative_prompt_embeds = negative_prompt_embeds[0]
Expand Down Expand Up @@ -385,7 +387,7 @@ def encode_image(self, image, device, num_images_per_prompt, output_hidden_state
if not isinstance(image, torch.Tensor):
image = self.feature_extractor(image, return_tensors="pt").pixel_values

image = image.to(device=device, dtype=dtype)
image = image.to(device=self.image_encoder.device, dtype=dtype)
if output_hidden_states:
image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
Expand Down
6 changes: 3 additions & 3 deletions src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -446,7 +446,7 @@ def encode_prompt(
f" {tokenizer.model_max_length} tokens: {removed_text}"
)

prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
prompt_embeds = text_encoder(text_input_ids.to(text_encoder.device), output_hidden_states=True)

# We are only ALWAYS interested in the pooled output of the final text encoder
if pooled_prompt_embeds is None and prompt_embeds[0].ndim == 2:
Expand Down Expand Up @@ -507,7 +507,7 @@ def encode_prompt(
)

negative_prompt_embeds = text_encoder(
uncond_input.input_ids.to(device),
uncond_input.input_ids.to(text_encoder.device),
output_hidden_states=True,
)

Expand Down Expand Up @@ -569,7 +569,7 @@ def encode_image(self, image, device, num_images_per_prompt, output_hidden_state
if not isinstance(image, torch.Tensor):
image = self.feature_extractor(image, return_tensors="pt").pixel_values

image = image.to(device=device, dtype=dtype)
image = image.to(device=self.image_encoder.device, dtype=dtype)
if output_hidden_states:
image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
Expand Down
14 changes: 8 additions & 6 deletions src/diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -290,17 +290,18 @@ def encode_prompt(
f" {self.tokenizer.model_max_length} tokens: {removed_text}"
)

model_device = self.text_encoder.device
if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
attention_mask = text_inputs.attention_mask.to(device)
attention_mask = text_inputs.attention_mask.to(model_device)
else:
attention_mask = None

if clip_skip is None:
prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)
prompt_embeds = self.text_encoder(text_input_ids.to(model_device), attention_mask=attention_mask)
prompt_embeds = prompt_embeds[0]
else:
prompt_embeds = self.text_encoder(
text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True
text_input_ids.to(model_device), attention_mask=attention_mask, output_hidden_states=True
)
# Access the `hidden_states` first, that contains a tuple of
# all the hidden states from the encoder layers. Then index into
Expand Down Expand Up @@ -360,13 +361,14 @@ def encode_prompt(
return_tensors="pt",
)

model_device = self.text_encoder.device
if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
attention_mask = uncond_input.attention_mask.to(device)
attention_mask = uncond_input.attention_mask.to(model_device)
else:
attention_mask = None

negative_prompt_embeds = self.text_encoder(
uncond_input.input_ids.to(device),
uncond_input.input_ids.to(model_device),
attention_mask=attention_mask,
)
negative_prompt_embeds = negative_prompt_embeds[0]
Expand Down Expand Up @@ -394,7 +396,7 @@ def encode_image(self, image, device, num_images_per_prompt, output_hidden_state
if not isinstance(image, torch.Tensor):
image = self.feature_extractor(image, return_tensors="pt").pixel_values

image = image.to(device=device, dtype=dtype)
image = image.to(device=self.image_encoder.device, dtype=dtype)
if output_hidden_states:
image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -442,7 +442,7 @@ def encode_image(self, image, device, num_images_per_prompt, output_hidden_state
if not isinstance(image, torch.Tensor):
image = self.feature_extractor(image, return_tensors="pt").pixel_values

image = image.to(device=device, dtype=dtype)
image = image.to(device=self.image_encoder.device, dtype=dtype)
if output_hidden_states:
image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -473,7 +473,7 @@ def encode_image(self, image, device, num_images_per_prompt, output_hidden_state
if not isinstance(image, torch.Tensor):
image = self.feature_extractor(image, return_tensors="pt").pixel_values

image = image.to(device=device, dtype=dtype)
image = image.to(device=self.image_encoder.device, dtype=dtype)
if output_hidden_states:
image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
Expand Down
3 changes: 2 additions & 1 deletion src/diffusers/pipelines/anyflow/pipeline_anyflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,7 +161,8 @@ def _get_t5_prompt_embeds(
text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
seq_lens = mask.gt(0).sum(dim=1).long()

prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state
model_device = self.text_encoder.device
prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
prompt_embeds = torch.stack(
Expand Down
3 changes: 2 additions & 1 deletion src/diffusers/pipelines/anyflow/pipeline_anyflow_far.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,8 @@ def _get_t5_prompt_embeds(
text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
seq_lens = mask.gt(0).sum(dim=1).long()

prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state
model_device = self.text_encoder.device
prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
prompt_embeds = torch.stack(
Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/pipelines/chroma/pipeline_chroma.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -372,7 +372,7 @@ def encode_image(self, image, device, num_images_per_prompt):
if not isinstance(image, torch.Tensor):
image = self.feature_extractor(image, return_tensors="pt").pixel_values

image = image.to(device=device, dtype=dtype)
image = image.to(device=self.image_encoder.device, dtype=dtype)
image_embeds = self.image_encoder(image).image_embeds
image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
return image_embeds
Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -399,7 +399,7 @@ def encode_image(self, image, device, num_images_per_prompt):
if not isinstance(image, torch.Tensor):
image = self.feature_extractor(image, return_tensors="pt").pixel_values

image = image.to(device=device, dtype=dtype)
image = image.to(device=self.image_encoder.device, dtype=dtype)
image_embeds = self.image_encoder(image).image_embeds
image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
return image_embeds
Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/pipelines/chroma/pipeline_chroma_inpainting.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -401,7 +401,7 @@ def encode_image(self, image, device, num_images_per_prompt):
if not isinstance(image, torch.Tensor):
image = self.feature_extractor(image, return_tensors="pt").pixel_values

image = image.to(device=device, dtype=dtype)
image = image.to(device=self.image_encoder.device, dtype=dtype)
image_embeds = self.image_encoder(image).image_embeds
image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
return image_embeds
Expand Down
7 changes: 4 additions & 3 deletions src/diffusers/pipelines/chronoedit/pipeline_chronoedit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -210,7 +210,8 @@ def _get_t5_prompt_embeds(
text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
seq_lens = mask.gt(0).sum(dim=1).long()

prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state
model_device = self.text_encoder.device
prompt_embeds = self.text_encoder(text_input_ids.to(model_device), mask.to(model_device)).last_hidden_state
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
prompt_embeds = torch.stack(
Expand All @@ -231,9 +232,9 @@ def encode_image(
device: torch.device | None = None,
):
device = device or self._execution_device
image = self.image_processor(images=image, return_tensors="pt").to(device)
image = self.image_processor(images=image, return_tensors="pt").to(self.image_encoder.device)
image_embeds = self.image_encoder(**image, output_hidden_states=True)
return image_embeds.hidden_states[-2]
return image_embeds.hidden_states[-2].to(device)

# Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt
def encode_prompt(
Expand Down
18 changes: 11 additions & 7 deletions src/diffusers/pipelines/controlnet/pipeline_controlnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -380,17 +380,18 @@ def encode_prompt(
f" {self.tokenizer.model_max_length} tokens: {removed_text}"
)

model_device = self.text_encoder.device
if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
attention_mask = text_inputs.attention_mask.to(device)
attention_mask = text_inputs.attention_mask.to(model_device)
else:
attention_mask = None

if clip_skip is None:
prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)
prompt_embeds = self.text_encoder(text_input_ids.to(model_device), attention_mask=attention_mask)
prompt_embeds = prompt_embeds[0]
else:
prompt_embeds = self.text_encoder(
text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True
text_input_ids.to(model_device), attention_mask=attention_mask, output_hidden_states=True
)
# Access the `hidden_states` first, that contains a tuple of
# all the hidden states from the encoder layers. Then index into
Expand Down Expand Up @@ -450,13 +451,14 @@ def encode_prompt(
return_tensors="pt",
)

model_device = self.text_encoder.device
if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
attention_mask = uncond_input.attention_mask.to(device)
attention_mask = uncond_input.attention_mask.to(model_device)
else:
attention_mask = None

negative_prompt_embeds = self.text_encoder(
uncond_input.input_ids.to(device),
uncond_input.input_ids.to(model_device),
attention_mask=attention_mask,
)
negative_prompt_embeds = negative_prompt_embeds[0]
Expand Down Expand Up @@ -484,7 +486,7 @@ def encode_image(self, image, device, num_images_per_prompt, output_hidden_state
if not isinstance(image, torch.Tensor):
image = self.feature_extractor(image, return_tensors="pt").pixel_values

image = image.to(device=device, dtype=dtype)
image = image.to(device=self.image_encoder.device, dtype=dtype)
if output_hidden_states:
image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
Expand Down Expand Up @@ -557,7 +559,9 @@ def run_safety_checker(self, image, device, dtype):
feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")
else:
feature_extractor_input = self.image_processor.numpy_to_pil(image)
safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(
self.safety_checker.device
)
image, has_nsfw_concept = self.safety_checker(
images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
)
Expand Down
Loading
Loading