huggingface / huggingface/diffusers
StableDiffusionXLInpaintPipeline: image latent not transferred to correct device when using `device_map="balanced"` on VAE model
- Lingua principale
- Python
- Stelle
- 34.5k
- Fork
- 7.3k
- Merge medio
- 3g 3h
- PR unite (30g)
- 91
Descrizione
### Describe the bug
Running StableDiffusionXLInpaintPipeline on multiple GPUs with `device_map="balanced"` on the VAE model, I got the following error (irrelevant part excluded)
```
File "/home1/uname/.cache/pypoetry/virtualenvs/project-350heM32-py3.11/lib/python3.11/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py", line 1579, in __call__
latents_outputs = self.prepare_latents(
^^^^^^^^^^^^^^^^^^^^^
File "/home1/uname/.cache/pypoetry/virtualenvs/project-350heM32-py3.11/lib/python3.11/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py", line 928, in prepare_latents
latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home1/uname/.cache/pypoetry/virtualenvs/project-350heM32-py3.11/lib/python3.11/site-packages/diffusers/schedulers/scheduling_euler_discrete.py", line 631, in add_noise
noisy_samples = original_samples + noise * sigma
~~~~~~^~~~~~~
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:7!
```
It seems that the `image_latents` is on the same device with the VAE model while the `noise` is not.
https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py#L799-L812
I haven't debugged further, but for my use case I can simply force `image_latents.to(device)` or move the VAE to "cuda:0" to get past this.
Is this a bug or it's just me using `device_map` wrong?
### Reproduction
```python
import torch
from diffusers import (
AutoencoderKL,
StableDiffusionXLInpaintPipeline,
)
from PIL import Image
vae = AutoencoderKL.from_pretrained(
"madebyollin/sdxl-vae-fp16-fix",
torch_dtype=torch.float16,
device_map="balanced",
)
pipeline = StableDiffusionXLInpaintPipeline.from_pretrained(
"OzzyGT/RealVisXL_V4.0_inpainting",
torch_dtype=torch.float16,
variant="fp16",
vae=vae,
device_map="balanced",
)
base_image = Image.new("RGB", (1024, 1024))
mask = Image.new("L", base_image.size)
final_images = pipeline(
prompt="",
image=base_image,
mask_image=mask,
).images
```
### Logs
```shell
/home1/uname/.cache/pypoetry/virtualenvs/project-350heM32-py3.11/lib/python3.11/site-packages/diffusers/models/transformers/transformer_2d.py:34: FutureWarning: `Transformer2DModelOutput` is deprecated and will be removed in version 1.0.0. Importing `Transformer2DModelOutput` from `diffusers.models.transformer_2d` is deprecated and this will be removed in a future version. Please use `from diffusers.models.modeling_outputs import Transformer2DModelOutput`, instead.
deprecate("Transformer2DModelOutput", "1.0.0", deprecation_message)
The safetensors archive passed at /home1/uname/.cache/huggingface/hub/models--madebyollin--sdxl-vae-fp16-fix/snapshots/207b116dae70ace3637169f1ddd2434b91b3a8cd/diffusion_pytorch_model.safetensors does not contain metadata. Make sure to save your model with the `save_pretrained` method. Defaulting to 'pt' metadata.
Loading pipeline components...: 100%|███████████████████████████████████████████████████| 7/7 [00:05<00:00, 1.37it/s]
Traceback (most recent call last):
File "/home1/uname/yoshi/src/project/test.py", line 25, in
final_images = pipeline(
^^^^^^^^^
File "/home1/uname/.cache/pypoetry/virtualenvs/project-350heM32-py3.11/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/home1/uname/.cache/pypoetry/virtualenvs/project-350heM32-py3.11/lib/python3.11/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py", line 1579, in __call__
latents_outputs = self.prepare_latents(
^^^^^^^^^^^^^^^^^^^^^
File "/home1/uname/.cache/pypoetry/virtualenvs/project-350heM32-py3.11/lib/python3.11/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py", line 928, in prepare_latents
latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home1/uname/.cache/pypoetry/virtualenvs/project-350heM32-py3.11/lib/python3.11/site-packages/diffusers/schedulers/scheduling_euler_discrete.py", line 631, in add_noise
noisy_samples = original_samples + noise * sigma
~~~~~~^~~~~~~
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:7!
```
### System Info
- 🤗 Diffusers version: 0.28.2
- Platform: CentOS Linux 7 (Core) - Linux-3.10.0-1160.42.2.el7.x86_64-x86_64-with-glibc2.17
- Running on a notebook?: No
- Running on Google Colab?: No
- Python version: 3.11.4
- PyTorch version (GPU?): 2.2.2+cu121 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 0.23.3
- Transformers version: 4.41.2
- Accelerate version: 0.31.0
- PEFT version: not installed
- Bitsandbytes version: not installed
- Safetensors version: 0.4.3
- xFormers version: not installed
- Accelerator: Tesla V100-SXM2-32GB, 32768 MiB
Tesla V100-SXM2-32GB, 32768 MiB
Tesla V100-SXM2-32GB, 32768 MiB
Tesla V100-SXM2-32GB, 32768 MiB
Tesla V100-SXM2-32GB, 32768 MiB
Tesla V100-SXM2-32GB, 32768 MiB
Tesla V100-SXM2-32GB, 32768 MiB
Tesla V100-SXM2-32GB, 32768 MiB VRAM
- Using GPU in script?: yes (all 8)
- Using distributed or parallel set-up in script?: yes (`device_map="balanced"`)
### Who can help?
@yiyixuxu
Guida per i contributori
Apri la guida per i contributori
Direzione di ricerca
Start in src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py, especially prepare_latents and the image-latent handling around the referenced lines. Reproduce the reported device_map="balanced" setup and trace the devices of image_latents and noise before scheduler.add_noise; done means the pipeline no longer raises a cross-device tensor error.
Scritto dal modello di indicizzazione a partire dal testo della issue.
Valutazione
- Stack tecnologico
- python, pytorch
- Ambito
- distributed-systems, machine-learning
- Tipo di issue
- Bug
- Difficoltà
- 4/5
- Tempo stimato
- 3-5 giorni
- Stato di attività
- Ferma
- Chiarezza
- Abbastanza chiara
- Idoneità per principianti
- 35/100