Skip to content

[WIP]correct the attn naming for UNet3DConditionModel #6873

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

Closed
wants to merge 2 commits into from
Closed

[WIP]correct the attn naming for UNet3DConditionModel #6873

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b71f35b
draft
Feb 6, 2024
fd7f9c7
fix
Feb 6, 2024
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
107 changes: 83 additions & 24 deletions src/diffusers/models/unets/unet_3d_blocks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
import torch
from torch import nn

from ...utils import is_torch_version
from ...utils import deprecate, is_torch_version
from ...utils.torch_utils import apply_freeu
from ..attention import Attention
from ..resnet import (
Expand All @@ -44,7 +44,8 @@ def get_down_block(
add_downsample: bool,
resnet_eps: float,
resnet_act_fn: str,
num_attention_heads: int,
num_attention_heads: Optional[int] = None,
attention_head_dim: Optional[int] = None,
resnet_groups: Optional[int] = None,
cross_attention_dim: Optional[int] = None,
downsample_padding: Optional[int] = None,
Expand Down Expand Up @@ -80,6 +81,16 @@ def get_down_block(
elif down_block_type == "CrossAttnDownBlock3D":
if cross_attention_dim is None:
raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock3D")
if num_attention_heads is not None:
deprecation_message = (
" passing `num`_attention_heads` to `unet_3d_blocks.get_down_block` for CrossAttnDownBlock3D is deprecated. "
" Please use `attention_head_dim` instead."
)
deprecate("num_attention_heads not None", "1.0.0", deprecation_message, standard_warn=False)
if attention_head_dim is None:
attention_head_dim = num_attention_heads
if attention_head_dim is None:
raise ValueError("`attention_head_dim` must be specified for CrossAttnDownBlock3D")
return CrossAttnDownBlock3D(
num_layers=num_layers,
in_channels=in_channels,
Expand All @@ -91,7 +102,8 @@ def get_down_block(
resnet_groups=resnet_groups,
downsample_padding=downsample_padding,
cross_attention_dim=cross_attention_dim,
num_attention_heads=num_attention_heads,
num_attention_heads=None,
attention_head_dim=attention_head_dim,
dual_cross_attention=dual_cross_attention,
use_linear_projection=use_linear_projection,
only_cross_attention=only_cross_attention,
Expand Down Expand Up @@ -173,7 +185,8 @@ def get_up_block(
add_upsample: bool,
resnet_eps: float,
resnet_act_fn: str,
num_attention_heads: int,
num_attention_heads: Optional[int] = None,
attention_head_dim: Optional[int] = None,
resolution_idx: Optional[int] = None,
resnet_groups: Optional[int] = None,
cross_attention_dim: Optional[int] = None,
Expand Down Expand Up @@ -212,6 +225,16 @@ def get_up_block(
elif up_block_type == "CrossAttnUpBlock3D":
if cross_attention_dim is None:
raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock3D")
if num_attention_heads is not None:
deprecation_message = (
" passing `num`_attention_heads` to `unet_3d_blocks.get_up_block` for CrossAttnUpBlock3D is deprecated. "
" Please use `attention_head_dim` instead."
)
deprecate("num_attention_heads not None", "1.0.0", deprecation_message, standard_warn=False)
if attention_head_dim is None:
attention_head_dim = num_attention_heads
if attention_head_dim is None:
raise ValueError("`attention_head_dim` must be specified for CrossAttnUpBlock3D")
return CrossAttnUpBlock3D(
num_layers=num_layers,
in_channels=in_channels,
Expand All @@ -223,7 +246,8 @@ def get_up_block(
resnet_act_fn=resnet_act_fn,
resnet_groups=resnet_groups,
cross_attention_dim=cross_attention_dim,
num_attention_heads=num_attention_heads,
num_attention_heads=None,
attention_head_dim=attention_head_dim,
dual_cross_attention=dual_cross_attention,
use_linear_projection=use_linear_projection,
only_cross_attention=only_cross_attention,
Expand Down Expand Up @@ -314,17 +338,28 @@ def __init__(
resnet_act_fn: str = "swish",
resnet_groups: int = 32,
resnet_pre_norm: bool = True,
num_attention_heads: int = 1,
num_attention_heads: Optional[int] = 1,
attention_head_dim: Optional[int] = None,
output_scale_factor: float = 1.0,
cross_attention_dim: int = 1280,
dual_cross_attention: bool = False,
use_linear_projection: bool = True,
upcast_attention: bool = False,
):
super().__init__()

if num_attention_heads is not None:
deprecation_message = (
" passing `num`_attention_heads` to `unet_3d_blocks.UNetMidBlock3DCrossAttn` is deprecated. "
" Please use `attention_head_dim` instead."
)
deprecate("num_attention_heads not None", "1.0.0", deprecation_message, standard_warn=False)
if attention_head_dim is None:
attention_head_dim = num_attention_heads
self.num_attention_heads = num_attention_heads
if attention_head_dim is None:
raise ValueError("`attention_head_dim` must be specified for UNetMidBlock3DCrossAttn")
self.has_cross_attention = True
self.num_attention_heads = num_attention_heads

resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)

# there is always at least one resnet
Expand Down Expand Up @@ -356,8 +391,8 @@ def __init__(
for _ in range(num_layers):
attentions.append(
Transformer2DModel(
in_channels // num_attention_heads,
num_attention_heads,
in_channels // attention_head_dim,
attention_head_dim,
in_channels=in_channels,
num_layers=1,
cross_attention_dim=cross_attention_dim,
Expand All @@ -368,8 +403,8 @@ def __init__(
)
temp_attentions.append(
TransformerTemporalModel(
in_channels // num_attention_heads,
num_attention_heads,
in_channels // attention_head_dim,
attention_head_dim,
in_channels=in_channels,
num_layers=1,
cross_attention_dim=cross_attention_dim,
Expand Down Expand Up @@ -449,7 +484,8 @@ def __init__(
resnet_act_fn: str = "swish",
resnet_groups: int = 32,
resnet_pre_norm: bool = True,
num_attention_heads: int = 1,
num_attention_heads: Optional[int] = 1,
attention_head_dim: Optional[int] = None,
cross_attention_dim: int = 1280,
output_scale_factor: float = 1.0,
downsample_padding: int = 1,
Expand All @@ -460,13 +496,23 @@ def __init__(
upcast_attention: bool = False,
):
super().__init__()
if num_attention_heads is not None:
deprecation_message = (
" passing `num`_attention_heads` to `unet_3d_blocks.CrossAttnDownBlock3D` is deprecated. "
" Please use `attention_head_dim` instead."
)
deprecate("num_attention_heads not None", "1.0.0", deprecation_message, standard_warn=False)
if attention_head_dim is None:
attention_head_dim = num_attention_heads
self.num_attention_heads = num_attention_heads
if attention_head_dim is None:
raise ValueError("`attention_head_dim` must be specified for CrossAttnDownBlock3D")
resnets = []
attentions = []
temp_attentions = []
temp_convs = []

self.has_cross_attention = True
self.num_attention_heads = num_attention_heads

for i in range(num_layers):
in_channels = in_channels if i == 0 else out_channels
Expand Down Expand Up @@ -494,8 +540,8 @@ def __init__(
)
attentions.append(
Transformer2DModel(
out_channels // num_attention_heads,
num_attention_heads,
out_channels // attention_head_dim,
attention_head_dim,
in_channels=out_channels,
num_layers=1,
cross_attention_dim=cross_attention_dim,
Expand All @@ -507,8 +553,8 @@ def __init__(
)
temp_attentions.append(
TransformerTemporalModel(
out_channels // num_attention_heads,
num_attention_heads,
out_channels // attention_head_dim,
attention_head_dim,
in_channels=out_channels,
num_layers=1,
cross_attention_dim=cross_attention_dim,
Expand Down Expand Up @@ -681,7 +727,8 @@ def __init__(
resnet_act_fn: str = "swish",
resnet_groups: int = 32,
resnet_pre_norm: bool = True,
num_attention_heads: int = 1,
num_attention_heads: Optional[int] = 1,
attention_head_dim: Optional[int] = None,
cross_attention_dim: int = 1280,
output_scale_factor: float = 1.0,
add_upsample: bool = True,
Expand All @@ -692,13 +739,25 @@ def __init__(
resolution_idx: Optional[int] = None,
):
super().__init__()
if num_attention_heads is not None:
deprecation_message = (
" passing `num`_attention_heads` to `unet_3d_blocks.CrossAttnUpBlock3D` is deprecated. "
" Please use `attention_head_dim` instead."
)
deprecate("num_attention_heads not None", "1.0.0", deprecation_message, standard_warn=False)
if attention_head_dim is None:
attention_head_dim = num_attention_heads
self.num_attention_heads = num_attention_heads

if attention_head_dim is None:
raise ValueError("`attention_head_dim` must be specified for CrossAttnUpBlock3D")

resnets = []
temp_convs = []
attentions = []
temp_attentions = []

self.has_cross_attention = True
self.num_attention_heads = num_attention_heads

for i in range(num_layers):
res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
Expand Down Expand Up @@ -728,8 +787,8 @@ def __init__(
)
attentions.append(
Transformer2DModel(
out_channels // num_attention_heads,
num_attention_heads,
out_channels // attention_head_dim,
attention_head_dim,
in_channels=out_channels,
num_layers=1,
cross_attention_dim=cross_attention_dim,
Expand All @@ -741,8 +800,8 @@ def __init__(
)
temp_attentions.append(
TransformerTemporalModel(
out_channels // num_attention_heads,
num_attention_heads,
out_channels // attention_head_dim,
attention_head_dim,
in_channels=out_channels,
num_layers=1,
cross_attention_dim=cross_attention_dim,
Expand Down
25 changes: 9 additions & 16 deletions src/diffusers/models/unets/unet_3d_condition.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,14 +132,6 @@ def __init__(
"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."
)

# 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

# Check inputs
if len(down_block_types) != len(up_block_types):
raise ValueError(
Expand All @@ -151,9 +143,9 @@ def __init__(
f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
)

if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):
raise ValueError(
f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}."
)

# input
Expand Down Expand Up @@ -187,8 +179,8 @@ def __init__(
self.down_blocks = nn.ModuleList([])
self.up_blocks = nn.ModuleList([])

if isinstance(num_attention_heads, int):
num_attention_heads = (num_attention_heads,) * len(down_block_types)
if isinstance(attention_head_dim, int):
attention_head_dim = (attention_head_dim,) * len(down_block_types)

# down
output_channel = block_out_channels[0]
Expand All @@ -208,7 +200,7 @@ def __init__(
resnet_act_fn=act_fn,
resnet_groups=norm_num_groups,
cross_attention_dim=cross_attention_dim,
num_attention_heads=num_attention_heads[i],
attention_head_dim=attention_head_dim[i],
downsample_padding=downsample_padding,
dual_cross_attention=False,
)
Expand All @@ -222,7 +214,8 @@ def __init__(
resnet_act_fn=act_fn,
output_scale_factor=mid_block_scale_factor,
cross_attention_dim=cross_attention_dim,
num_attention_heads=num_attention_heads[-1],
num_attention_heads=None,
attention_head_dim=attention_head_dim[-1],
resnet_groups=norm_num_groups,
dual_cross_attention=False,
)
Expand All @@ -232,7 +225,7 @@ def __init__(

# up
reversed_block_out_channels = list(reversed(block_out_channels))
reversed_num_attention_heads = list(reversed(num_attention_heads))
reversed_attention_head_dim = list(reversed(attention_head_dim))

output_channel = reversed_block_out_channels[0]
for i, up_block_type in enumerate(up_block_types):
Expand Down Expand Up @@ -261,7 +254,7 @@ def __init__(
resnet_act_fn=act_fn,
resnet_groups=norm_num_groups,
cross_attention_dim=cross_attention_dim,
num_attention_heads=reversed_num_attention_heads[i],
attention_head_dim=reversed_attention_head_dim[i],
dual_cross_attention=False,
resolution_idx=i,
)
Expand Down