ZeroDivisionError when performing forward pass with UNet3DConditionModel

[txz32102]

Describe the bug

ZeroDivisionError when performing forward pass with UNet3DConditionModel

I'm encountering a ZeroDivisionError when attempting to perform a forward pass with the UNet3DConditionModel. This seems to be related to the num_attention_heads parameter being None, which causes self.inner_dim to be 0.

Here's the code I'm using:

from diffusers import UNet3DConditionModel
import torch

model = UNet3DConditionModel(
    down_block_types=(
        "CrossAttnDownBlock3D",
        "CrossAttnDownBlock3D",
        "CrossAttnDownBlock3D",
        "DownBlock3D",
    ),
    up_block_types=(
        "UpBlock3D",
        "CrossAttnUpBlock3D",
        "CrossAttnUpBlock3D",
        "CrossAttnUpBlock3D",
    ),
    block_out_channels=(32, 64, 128, 128),
    norm_num_groups=4,
)

data = torch.randn(1, 4, 32, 32, 32)

model(data, timestep=3, encoder_hidden_states=torch.zeros(1, 4, 32, 32, 32))

The error traceback indicates that the issue occurs in the attention processing:

ZeroDivisionError: integer division or modulo by zero

This seems to be because num_attention_heads is None, leading to self.inner_dim = 0 in the transformer configuration.

I noticed that in the UNet3DConditionModel implementation, there's a check that raises an error if num_attention_heads is provided:

if num_attention_heads is not None:
    raise NotImplementedError(
        "At the moment it is not possible to define the number of attention heads via num_attention_heads because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 . Passing num_attention_heads will only be supported in diffusers v0.19."
    )

Given this limitation, I'm unsure how to properly configure the model to avoid this error. Could you provide guidance on:

1. How to correctly perform a forward pass with demo hidden states
2. What parameters I should adjust to ensure the model is properly configured
3. If there's a workaround for this issue in the current version of diffusers

Thank you for your assistance!

Reproduction

from diffusers import UNet3DConditionModel
import torch

model = UNet3DConditionModel(
    down_block_types=(
        "CrossAttnDownBlock3D",
        "CrossAttnDownBlock3D",
        "CrossAttnDownBlock3D",
        "DownBlock3D",
    ),
    up_block_types=(
        "UpBlock3D",
        "CrossAttnUpBlock3D",
        "CrossAttnUpBlock3D",
        "CrossAttnUpBlock3D",
    ),
    block_out_channels=(32, 64, 128, 128),
    norm_num_groups=4,
)

data = torch.randn(1, 4, 32, 32, 32)

model(data, timestep=3, encoder_hidden_states=torch.zeros(1, 4, 32, 32, 32))

Logs

System Info

Python 3.11.10
diffusers version 0.32.2
ubuntu 24.04

Who can help?

No response

[hlky]

Hi @txz32102.

attention_head_dim or num_attention_heads can be used, attention_head_dim is preferred, for details see:

diffusers/src/diffusers/models/unets/unet_3d_condition.py

Lines 137 to 143 in 733b44a

	# If `num_attention_heads` is not defined (which is the case for most models)
	# it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
	# The reason for this behavior is to correct for incorrectly named variables that were introduced
	# when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
	# Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
	# which is why we correct for the naming here.
	num_attention_heads = num_attention_heads or attention_head_dim

The default attention_head_dim is 64.

diffusers/src/diffusers/models/unets/unet_3d_condition.py

Line 124 in 733b44a

attention_head_dim: Union[int, Tuple[int]] = 64,

block_out_channels is passed to the block.

diffusers/src/diffusers/models/unets/unet_3d_condition.py

Line 200 in 733b44a

output_channel = block_out_channels[i]

diffusers/src/diffusers/models/unets/unet_3d_condition.py

Line 207 in 733b44a

out_channels=output_channel,

Transformer2DModel's num_attention_heads is set to out_channels // num_attention_heads.

diffusers/src/diffusers/models/unets/unet_3d_blocks.py

Line 466 in 733b44a

out_channels // num_attention_heads,

inner_dim uses Transformer2DModel's num_attention_heads

diffusers/src/diffusers/models/transformers/transformer_2d.py

Line 152 in 733b44a

self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim

out_channels = 32
num_attention_heads = 64
32 // 64 = 0

Also note that encoder_hidden_states is expected to be shape (batch, sequence_length, feature_dim)

diffusers/src/diffusers/models/unets/unet_3d_condition.py

Lines 565 to 566 in 733b44a

	encoder_hidden_states (`torch.Tensor`):
	The encoder hidden states with shape `(batch, sequence_length, feature_dim)`.

and the default feature_dim/cross_attention_dim is 1024.

diffusers/src/diffusers/models/unets/unet_3d_condition.py

Line 123 in 733b44a

cross_attention_dim: int = 1024,

Hope this helps, let us know if you have any questions.

[txz32102]

Thanks, my bad, very helpful