AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-24 20:59:27 +03:00
Code Activity

clean up unet

Browse Source
This commit is contained in:
patil-suraj
2023-11-27 10:03:40 +01:00
parent 760333d524
commit af85fb1bc1
4 changed files with 25 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/diffusers/models/transformer_temporal.py
View File

@@ -253,12 +253,10 @@ class TransformerSpatioTemporalModel(ModelMixin, ConfigMixin):
attention_type: str = "default",
merge_factor: float = 0.5,
merge_strategy: str = "learned_with_images",
max_time_embed_period: int = 10000,
):
super().__init__()
self.num_attention_heads = num_attention_heads
self.attention_head_dim = attention_head_dim
self.max_time_embed_period = max_time_embed_period
inner_dim = num_attention_heads * attention_head_dim
self.inner_dim = inner_dim

11
src/diffusers/models/unet_3d_blocks.py
View File

@@ -60,7 +60,6 @@ def get_down_block(
kernel_size_3d: Optional[torch.FloatTensor] = (3, 1, 1),
merge_factor: float = 0.5,
merge_strategy: str = "learned_with_images",
max_time_embed_period: int = 10000,
) -> Union[
"DownBlock3D",
"CrossAttnDownBlock3D",
@@ -183,7 +182,6 @@ def get_down_block(
kernel_size_3d=kernel_size_3d,
merge_factor=merge_factor,
merge_strategy=merge_strategy,
max_time_embed_period=max_time_embed_period,
)
raise ValueError(f"{down_block_type} does not exist.")
@@ -216,7 +214,6 @@ def get_up_block(
kernel_size_3d: Optional[torch.FloatTensor] = (3, 1, 1),
merge_factor: float = 0.5,
merge_strategy: str = "learned_with_images",
max_time_embed_period: int = 10000,
) -> Union[
"UpBlock3D",
"CrossAttnUpBlock3D",
@@ -345,7 +342,6 @@ def get_up_block(
kernel_size_3d=kernel_size_3d,
merge_factor=merge_factor,
merge_strategy=merge_strategy,
max_time_embed_period=max_time_embed_period,
)
raise ValueError(f"{up_block_type} does not exist.")
@@ -1919,7 +1915,6 @@ class UpBlockTemporalDecoder(nn.Module):
attention_type: str = "default",
merge_factor: float = 0.0,
merge_strategy: str = "learned",
max_time_embed_period: int = 10000,
transformer_layers_per_block: Union[int, Tuple[int]] = (1,),
switch_spatial_to_temporal_mix: bool = True,
):
@@ -2000,7 +1995,6 @@ class UNetMidBlockSpatioTemporal(nn.Module):
kernel_size_3d: Optional[torch.FloatTensor] = (3, 1, 1),
merge_factor: float = 0.5,
merge_strategy: str = "learned_with_images",
max_time_embed_period: int = 10000,
):
super().__init__()
@@ -2045,7 +2039,6 @@ class UNetMidBlockSpatioTemporal(nn.Module):
attention_type=attention_type,
merge_factor=merge_factor,
merge_strategy=merge_strategy,
max_time_embed_period=max_time_embed_period,
)
)
@@ -2286,7 +2279,6 @@ class CrossAttnDownBlockSpatioTemporal(nn.Module):
kernel_size_3d: Optional[torch.FloatTensor] = (3, 1, 1),
merge_factor: float = 0.5,
merge_strategy: str = "learned_with_images",
max_time_embed_period: int = 10000,
):
super().__init__()
resnets = []
@@ -2329,7 +2321,6 @@ class CrossAttnDownBlockSpatioTemporal(nn.Module):
attention_type=attention_type,
merge_factor=merge_factor,
merge_strategy=merge_strategy,
max_time_embed_period=max_time_embed_period,
)
)
@@ -2594,7 +2585,6 @@ class CrossAttnUpBlockSpatioTemporal(nn.Module):
kernel_size_3d: Optional[torch.FloatTensor] = (3, 1, 1),
merge_factor: float = 0.5,
merge_strategy: str = "learned_with_images",
max_time_embed_period: int = 10000,
):
super().__init__()
resnets = []
@@ -2640,7 +2630,6 @@ class CrossAttnUpBlockSpatioTemporal(nn.Module):
attention_type=attention_type,
merge_factor=merge_factor,
merge_strategy=merge_strategy,
max_time_embed_period=max_time_embed_period,
)
)

126
src/diffusers/models/unet_spatio_temporal_condition.py
View File

@@ -140,7 +140,6 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
out_channels: int = 4,
center_input_sample: bool = False,
flip_sin_to_cos: bool = True,
freq_shift: int = 0,
down_block_types: Tuple[str] = (
"CrossAttnDownBlockSpatioTemporal",
"CrossAttnDownBlockSpatioTemporal",
@@ -154,47 +153,26 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
"CrossAttnUpBlockSpatioTemporal",
"CrossAttnUpBlockSpatioTemporal",
),
only_cross_attention: Union[bool, Tuple[bool]] = False,
block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
projection_class_embeddings_input_dim: int = 768,
addition_time_embed_dim: int = 256,
layers_per_block: Union[int, Tuple[int]] = 2,
downsample_padding: int = 1,
mid_block_scale_factor: float = 1,
dropout: float = 0.0,
act_fn: str = "silu",
norm_num_groups: Optional[int] = 32,
norm_eps: float = 1e-5,
cross_attention_dim: Union[int, Tuple[int]] = 1280,
cross_attention_dim: Union[int, Tuple[int]] = 1024,
transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None,
encoder_hid_dim: Optional[int] = None,
encoder_hid_dim_type: Optional[str] = None,
attention_head_dim: Union[int, Tuple[int]] = 8,
num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
dual_cross_attention: bool = False,
use_linear_projection: bool = False,
class_embed_type: Optional[str] = None,
addition_embed_type: Optional[str] = None,
addition_time_embed_dim: Optional[int] = None,
num_class_embeds: Optional[int] = None,
upcast_attention: bool = False,
resnet_time_scale_shift: str = "default",
resnet_out_scale_factor: int = 1.0,
time_embedding_type: str = "positional",
time_embedding_dim: Optional[int] = None,
time_embedding_act_fn: Optional[str] = None,
timestep_post_act: Optional[str] = None,
time_cond_proj_dim: Optional[int] = None,
conv_in_kernel: int = 3,
conv_out_kernel: int = 3,
projection_class_embeddings_input_dim: Optional[int] = None,
class_embeddings_concat: bool = False,
mid_block_only_cross_attention: Optional[bool] = None,
cross_attention_norm: Optional[str] = None,
addition_embed_type_num_heads=64,
kernel_size_3d: Optional[torch.FloatTensor] = (3, 1, 1),
merge_factor: float = 0.5,
merge_strategy: str = "learned_with_images",
max_time_embed_period: int = 10000,
):
super().__init__()
@@ -243,10 +221,6 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
raise ValueError(
f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}."
)
if isinstance(transformer_layers_per_block, list) and reverse_transformer_layers_per_block is None:
for layer_number_per_block in transformer_layers_per_block:
if isinstance(layer_number_per_block, list):
raise ValueError("Must provide 'reverse_transformer_layers_per_block` if using asymmetrical UNet.")
# input
conv_in_padding = (conv_in_kernel - 1) // 2
@@ -257,35 +231,18 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
# time
time_embed_dim = time_embedding_dim or block_out_channels[0] * 4
self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, downscale_freq_shift=0)
timestep_input_dim = block_out_channels[0]
self.time_embedding = TimestepEmbedding(
timestep_input_dim,
time_embed_dim,
act_fn=act_fn,
post_act_fn=timestep_post_act,
cond_proj_dim=time_cond_proj_dim,
)
self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift)
self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, downscale_freq_shift=0)
self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
if encoder_hid_dim_type is None and encoder_hid_dim is not None:
encoder_hid_dim_type = "text_proj"
self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type)
logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.")
if encoder_hid_dim is None and encoder_hid_dim_type is not None:
raise ValueError(
f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."
)
if time_embedding_act_fn is None:
self.time_embed_act = None
else:
self.time_embed_act = get_activation(time_embedding_act_fn)
self.down_blocks = nn.ModuleList([])
self.up_blocks = nn.ModuleList([])
@@ -326,15 +283,12 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
resnet_groups=norm_num_groups,
cross_attention_dim=cross_attention_dim[i],
num_attention_heads=num_attention_heads[i],
downsample_padding=downsample_padding,
dual_cross_attention=dual_cross_attention,
downsample_padding=1,
upcast_attention=upcast_attention,
resnet_time_scale_shift=resnet_time_scale_shift,
dropout=dropout,
kernel_size_3d=kernel_size_3d,
merge_factor=merge_factor,
merge_strategy=merge_strategy,
max_time_embed_period=max_time_embed_period,
)
self.down_blocks.append(down_block)
@@ -345,18 +299,14 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
transformer_layers_per_block=transformer_layers_per_block[-1],
resnet_eps=norm_eps,
resnet_act_fn=act_fn,
output_scale_factor=mid_block_scale_factor,
resnet_time_scale_shift=resnet_time_scale_shift,
cross_attention_dim=cross_attention_dim[-1],
num_attention_heads=num_attention_heads[-1],
resnet_groups=norm_num_groups,
dual_cross_attention=dual_cross_attention,
upcast_attention=upcast_attention,
dropout=dropout,
kernel_size_3d=kernel_size_3d,
merge_factor=merge_factor,
merge_strategy=merge_strategy,
max_time_embed_period=max_time_embed_period,
)
# count how many layers upsample the images
@@ -367,12 +317,7 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
reversed_num_attention_heads = list(reversed(num_attention_heads))
reversed_layers_per_block = list(reversed(layers_per_block))
reversed_cross_attention_dim = list(reversed(cross_attention_dim))
reversed_transformer_layers_per_block = (
list(reversed(transformer_layers_per_block))
if reverse_transformer_layers_per_block is None
else reverse_transformer_layers_per_block
)
# only_cross_attention = list(reversed(only_cross_attention))
reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block))
output_channel = reversed_block_out_channels[0]
for i, up_block_type in enumerate(up_block_types):
@@ -404,29 +349,18 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
resnet_groups=norm_num_groups,
cross_attention_dim=reversed_cross_attention_dim[i],
num_attention_heads=reversed_num_attention_heads[i],
dual_cross_attention=dual_cross_attention,
upcast_attention=upcast_attention,
resnet_time_scale_shift=resnet_time_scale_shift,
dropout=dropout,
kernel_size_3d=kernel_size_3d,
merge_factor=merge_factor,
merge_strategy=merge_strategy,
max_time_embed_period=max_time_embed_period,
)
self.up_blocks.append(up_block)
prev_output_channel = output_channel
# out
if norm_num_groups is not None:
self.conv_norm_out = nn.GroupNorm(
num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps
)
self.conv_act = get_activation(act_fn)
else:
self.conv_norm_out = None
self.conv_act = None
self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps)
self.conv_act = get_activation(act_fn)
conv_out_padding = (conv_out_kernel - 1) // 2
self.conv_out = nn.Conv2d(
@@ -614,15 +548,10 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
sample: torch.FloatTensor,
timestep: Union[torch.Tensor, float, int],
encoder_hidden_states: torch.Tensor,
added_cond_kwargs: Dict[str, torch.Tensor],
image_only_indicator: Optional[torch.Tensor] = None,
class_labels: Optional[torch.Tensor] = None,
timestep_cond: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
mid_block_additional_residual: Optional[torch.Tensor] = None,
down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
attention_mask: Optional[torch.Tensor] = None,
encoder_attention_mask: Optional[torch.Tensor] = None,
return_dict: bool = True,
) -> Union[UNetSpatioTemporalConditionOutput, Tuple]:
@@ -634,12 +563,6 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
The noisy input tensor with the following shape `(batch, channel, height, width)`.
timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.
encoder_hidden_states (`torch.FloatTensor`):
The encoder hidden states with shape `(batch, sequence_length, feature_dim)`.
class_labels (`torch.Tensor`, *optional*, defaults to `None`):
Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`):
Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed
through the `self.time_embedding` layer to obtain the timestep embeddings.
attention_mask (`torch.Tensor`, *optional*, defaults to `None`):
An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
@@ -648,13 +571,9 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
added_cond_kwargs: (`dict`, *optional*):
added_cond_kwargs: (`dict`):
A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that
are passed along to the UNet blocks.
down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*):
A tuple of tensors that if specified are added to the residuals of down unet blocks.
mid_block_additional_residual: (`torch.Tensor`, *optional*):
A tensor that if specified is added to the residual of the middle unet block.
encoder_attention_mask (`torch.Tensor`):
A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If
`True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias,
@@ -664,17 +583,6 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
tuple.
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttnProcessor`].
added_cond_kwargs: (`dict`, *optional*):
A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that
are passed along to the UNet blocks.
down_block_additional_residuals (`tuple` of `torch.Tensor`, *optional*):
additional residuals to be added to UNet long skip connections from down blocks to up blocks for
example from ControlNet side model(s)
mid_block_additional_residual (`torch.Tensor`, *optional*):
additional residual to be added to UNet mid block output, for example from ControlNet side model
down_intrablock_additional_residuals (`tuple` of `torch.Tensor`, *optional*):
additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s)
Returns:
[`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise
@@ -746,7 +654,7 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
# there might be better ways to encapsulate this.
t_emb = t_emb.to(dtype=sample.dtype)
emb = self.time_embedding(t_emb, timestep_cond)
emb = self.time_embedding(t_emb)
time_ids = added_cond_kwargs.get("time_ids")
time_embeds = self.add_time_proj(time_ids.flatten())
@@ -755,9 +663,6 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
aug_emb = self.add_embedding(time_embeds)
emb = emb + aug_emb
if self.time_embed_act is not None:
emb = self.time_embed_act(emb)
# Flatten the batch and frames dimensions
# sample: [batch, frames, channels, height, width] -> [batch * frames, channels, height, width]
sample = sample.flatten(0, 1)
@@ -857,9 +762,8 @@ class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionL
)
# 6. post-process
if self.conv_norm_out:
sample = self.conv_norm_out(sample)
sample = self.conv_act(sample)
sample = self.conv_norm_out(sample)
sample = self.conv_act(sample)
sample = self.conv_out(sample)
# 7. Reshape back to original shape

42
tests/models/test_models_unet_spatiotemporal.py
View File

@@ -36,9 +36,7 @@ logger = logging.get_logger(__name__)
enable_full_determinism()
class UNetSpatioTemporalConditionModelTests(
ModelTesterMixin, UNetTesterMixin, unittest.TestCase
):
class UNetSpatioTemporalConditionModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase):
model_class = UNetSpatioTemporalConditionModel
main_input_name = "sample"
@@ -110,9 +108,7 @@ class UNetSpatioTemporalConditionModelTests(
model.enable_xformers_memory_efficient_attention()
assert (
model.mid_block.attentions[0]
.transformer_blocks[0]
.attn1.processor.__class__.__name__
model.mid_block.attentions[0].transformer_blocks[0].attn1.processor.__class__.__name__
== "XFormersAttnProcessor"
), "xformers is not enabled"
@@ -155,11 +151,7 @@ class UNetSpatioTemporalConditionModelTests(
named_params = dict(model.named_parameters())
named_params_2 = dict(model_2.named_parameters())
for name, param in named_params.items():
self.assertTrue(
torch_all_close(
param.grad.data, named_params_2[name].grad.data, atol=5e-5
)
)
self.assertTrue(torch_all_close(param.grad.data, named_params_2[name].grad.data, atol=5e-5))
def test_model_with_attention_head_dim_tuple(self):
init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
@@ -178,9 +170,7 @@ class UNetSpatioTemporalConditionModelTests(
self.assertIsNotNone(output)
expected_shape = inputs_dict["sample"].shape
self.assertEqual(
output.shape, expected_shape, "Input and output shapes do not match"
)
self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
def test_model_with_use_linear_projection(self):
init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
@@ -199,9 +189,7 @@ class UNetSpatioTemporalConditionModelTests(
self.assertIsNotNone(output)
expected_shape = inputs_dict["sample"].shape
self.assertEqual(
output.shape, expected_shape, "Input and output shapes do not match"
)
self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
def test_model_with_cross_attention_dim_tuple(self):
init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
@@ -220,9 +208,7 @@ class UNetSpatioTemporalConditionModelTests(
self.assertIsNotNone(output)
expected_shape = inputs_dict["sample"].shape
self.assertEqual(
output.shape, expected_shape, "Input and output shapes do not match"
)
self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
def test_model_with_simple_projection(self):
init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
@@ -232,9 +218,7 @@ class UNetSpatioTemporalConditionModelTests(
init_dict["class_embed_type"] = "simple_projection"
init_dict["projection_class_embeddings_input_dim"] = sample_size
inputs_dict["class_labels"] = floats_tensor((batch_size, sample_size)).to(
torch_device
)
inputs_dict["class_labels"] = floats_tensor((batch_size, sample_size)).to(torch_device)
model = self.model_class(**init_dict)
model.to(torch_device)
@@ -248,9 +232,7 @@ class UNetSpatioTemporalConditionModelTests(
self.assertIsNotNone(output)
expected_shape = inputs_dict["sample"].shape
self.assertEqual(
output.shape, expected_shape, "Input and output shapes do not match"
)
self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
def test_model_with_class_embeddings_concat(self):
init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
@@ -261,9 +243,7 @@ class UNetSpatioTemporalConditionModelTests(
init_dict["projection_class_embeddings_input_dim"] = sample_size
init_dict["class_embeddings_concat"] = True
inputs_dict["class_labels"] = floats_tensor((batch_size, sample_size)).to(
torch_device
)
inputs_dict["class_labels"] = floats_tensor((batch_size, sample_size)).to(torch_device)
model = self.model_class(**init_dict)
model.to(torch_device)
@@ -277,9 +257,7 @@ class UNetSpatioTemporalConditionModelTests(
self.assertIsNotNone(output)
expected_shape = inputs_dict["sample"].shape
self.assertEqual(
output.shape, expected_shape, "Input and output shapes do not match"
)
self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
def test_model_attention_slicing(self):
init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()