huggingface / huggingface/diffusers

UNet2DModel dtype property fails under nn.DataParallel with UnboundLocalError in get_parameter_dtype

Open Beginner friendly
#13,789 0 comments 0 reactions 0 assignees View on GitHub
bug models
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.