huggingface / huggingface/diffusers
UNet2DModel dtype property fails under nn.DataParallel with UnboundLocalError in get_parameter_dtype
- Dominant language
- Python
- Stars
- 34.5k
- Forks
- 7.3k
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 91
Description
### Describe the bug
## Description
`UNet2DModel` fails under `torch.nn.DataParallel` during the forward pass when it accesses `self.dtype` inside `UNet2DModel.forward`.
The error appears to come from `diffusers.models.modeling_utils.get_parameter_dtype`, where the nested function annotation uses `tuple`:
```python
def find_tensor_attributes(module: nn.Module) -> list[tuple[str, Tensor]]:
### Reproduction
import torch
from diffusers import UNet2DModel
print("torch:", torch.__version__)
import diffusers
print("diffusers:", diffusers.__version__)
model = UNet2DModel(
sample_size=128,
in_channels=10,
out_channels=5,
layers_per_block=2,
block_out_channels=(64, 128, 256, 256),
down_block_types=(
"DownBlock2D",
"DownBlock2D",
"AttnDownBlock2D",
"DownBlock2D",
),
up_block_types=(
"UpBlock2D",
"AttnUpBlock2D",
"UpBlock2D",
"UpBlock2D",
),
norm_num_groups=8,
)
model = model.cuda()
model = torch.nn.DataParallel(model, device_ids=[0, 1])
model.eval()
x = torch.randn(8, 10, 128, 128, device="cuda")
t = torch.rand(8, device="cuda") * 1000
with torch.no_grad():
y = model(x, t, return_dict=False)[0]
print(y.shape)
### Logs
```shell
File ".../diffusers/models/unets/unet_2d.py", line ..., in forward
t_emb = t_emb.to(dtype=self.dtype)
File ".../diffusers/models/modeling_utils.py", line ..., in dtype
return get_parameter_dtype(self)
File ".../diffusers/models/modeling_utils.py", line ..., in get_parameter_dtype
def find_tensor_attributes(module: nn.Module) -> list[tuple[str, Tensor]]:
UnboundLocalError: cannot access local variable 'tuple' where it is not associated with a value
```
### System Info
Python 3.11.13, torch 2.9.1+cu128, diffusers 0.38.0, CUDA Version: 12.9
### Who can help?
_No response_
Contributor guide
Research direction
Start in diffusers/models/modeling_utils.py at get_parameter_dtype and its find_tensor_attributes annotation, then follow the self.dtype access from diffusers/models/unets/unet_2d.py. Reproduce the failure with the provided UNet2DModel and torch.nn.DataParallel example, and add or run a regression check showing the forward pass completes without UnboundLocalError.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 74/100