diff --git "a/src/unet_block_hacked_garmnet.py" "b/src/unet_block_hacked_garmnet.py"
deleted file mode 100644--- "a/src/unet_block_hacked_garmnet.py"
+++ /dev/null
@@ -1,3579 +0,0 @@
-# Copyright 2023 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import Any, Dict, Optional, Tuple, Union
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-from diffusers.utils import is_torch_version, logging
-from diffusers.utils.torch_utils import apply_freeu
-from diffusers.models.activations import get_activation
-from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor, AttnAddedKVProcessor2_0
-from diffusers.models.dual_transformer_2d import DualTransformer2DModel
-from diffusers.models.normalization import AdaGroupNorm
-from diffusers.models.resnet import Downsample2D, FirDownsample2D, FirUpsample2D, KDownsample2D, KUpsample2D, ResnetBlock2D, Upsample2D
-from src.transformerhacked_garmnet import Transformer2DModel
-from einops import rearrange
-
-logger = logging.get_logger(__name__) # pylint: disable=invalid-name
-
-
-def get_down_block(
- down_block_type: str,
- num_layers: int,
- in_channels: int,
- out_channels: int,
- temb_channels: int,
- add_downsample: bool,
- resnet_eps: float,
- resnet_act_fn: str,
- transformer_layers_per_block: int = 1,
- num_attention_heads: Optional[int] = None,
- resnet_groups: Optional[int] = None,
- cross_attention_dim: Optional[int] = None,
- downsample_padding: Optional[int] = None,
- dual_cross_attention: bool = False,
- use_linear_projection: bool = False,
- only_cross_attention: bool = False,
- upcast_attention: bool = False,
- resnet_time_scale_shift: str = "default",
- attention_type: str = "default",
- resnet_skip_time_act: bool = False,
- resnet_out_scale_factor: float = 1.0,
- cross_attention_norm: Optional[str] = None,
- attention_head_dim: Optional[int] = None,
- downsample_type: Optional[str] = None,
- dropout: float = 0.0,
-):
- # If attn head dim is not defined, we default it to the number of heads
- if attention_head_dim is None:
- logger.warn(
- f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
- )
- attention_head_dim = num_attention_heads
-
- down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
- if down_block_type == "DownBlock2D":
- return DownBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- dropout=dropout,
- add_downsample=add_downsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_groups=resnet_groups,
- downsample_padding=downsample_padding,
- resnet_time_scale_shift=resnet_time_scale_shift,
- )
- elif down_block_type == "ResnetDownsampleBlock2D":
- return ResnetDownsampleBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- dropout=dropout,
- add_downsample=add_downsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_groups=resnet_groups,
- resnet_time_scale_shift=resnet_time_scale_shift,
- skip_time_act=resnet_skip_time_act,
- output_scale_factor=resnet_out_scale_factor,
- )
- elif down_block_type == "AttnDownBlock2D":
- if add_downsample is False:
- downsample_type = None
- else:
- downsample_type = downsample_type or "conv" # default to 'conv'
- return AttnDownBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- dropout=dropout,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_groups=resnet_groups,
- downsample_padding=downsample_padding,
- attention_head_dim=attention_head_dim,
- resnet_time_scale_shift=resnet_time_scale_shift,
- downsample_type=downsample_type,
- )
- elif down_block_type == "CrossAttnDownBlock2D":
- if cross_attention_dim is None:
- raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock2D")
- return CrossAttnDownBlock2D(
- num_layers=num_layers,
- transformer_layers_per_block=transformer_layers_per_block,
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- dropout=dropout,
- add_downsample=add_downsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_groups=resnet_groups,
- downsample_padding=downsample_padding,
- cross_attention_dim=cross_attention_dim,
- num_attention_heads=num_attention_heads,
- dual_cross_attention=dual_cross_attention,
- use_linear_projection=use_linear_projection,
- only_cross_attention=only_cross_attention,
- upcast_attention=upcast_attention,
- resnet_time_scale_shift=resnet_time_scale_shift,
- attention_type=attention_type,
- )
- elif down_block_type == "SimpleCrossAttnDownBlock2D":
- if cross_attention_dim is None:
- raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnDownBlock2D")
- return SimpleCrossAttnDownBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- dropout=dropout,
- add_downsample=add_downsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_groups=resnet_groups,
- cross_attention_dim=cross_attention_dim,
- attention_head_dim=attention_head_dim,
- resnet_time_scale_shift=resnet_time_scale_shift,
- skip_time_act=resnet_skip_time_act,
- output_scale_factor=resnet_out_scale_factor,
- only_cross_attention=only_cross_attention,
- cross_attention_norm=cross_attention_norm,
- )
- elif down_block_type == "SkipDownBlock2D":
- return SkipDownBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- dropout=dropout,
- add_downsample=add_downsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- downsample_padding=downsample_padding,
- resnet_time_scale_shift=resnet_time_scale_shift,
- )
- elif down_block_type == "AttnSkipDownBlock2D":
- return AttnSkipDownBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- dropout=dropout,
- add_downsample=add_downsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- attention_head_dim=attention_head_dim,
- resnet_time_scale_shift=resnet_time_scale_shift,
- )
- elif down_block_type == "DownEncoderBlock2D":
- return DownEncoderBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- dropout=dropout,
- add_downsample=add_downsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_groups=resnet_groups,
- downsample_padding=downsample_padding,
- resnet_time_scale_shift=resnet_time_scale_shift,
- )
- elif down_block_type == "AttnDownEncoderBlock2D":
- return AttnDownEncoderBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- dropout=dropout,
- add_downsample=add_downsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_groups=resnet_groups,
- downsample_padding=downsample_padding,
- attention_head_dim=attention_head_dim,
- resnet_time_scale_shift=resnet_time_scale_shift,
- )
- elif down_block_type == "KDownBlock2D":
- return KDownBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- dropout=dropout,
- add_downsample=add_downsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- )
- elif down_block_type == "KCrossAttnDownBlock2D":
- return KCrossAttnDownBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- dropout=dropout,
- add_downsample=add_downsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- cross_attention_dim=cross_attention_dim,
- attention_head_dim=attention_head_dim,
- add_self_attention=True if not add_downsample else False,
- )
- raise ValueError(f"{down_block_type} does not exist.")
-
-
-def get_up_block(
- up_block_type: str,
- num_layers: int,
- in_channels: int,
- out_channels: int,
- prev_output_channel: int,
- temb_channels: int,
- add_upsample: bool,
- resnet_eps: float,
- resnet_act_fn: str,
- resolution_idx: Optional[int] = None,
- transformer_layers_per_block: int = 1,
- num_attention_heads: Optional[int] = None,
- resnet_groups: Optional[int] = None,
- cross_attention_dim: Optional[int] = None,
- dual_cross_attention: bool = False,
- use_linear_projection: bool = False,
- only_cross_attention: bool = False,
- upcast_attention: bool = False,
- resnet_time_scale_shift: str = "default",
- attention_type: str = "default",
- resnet_skip_time_act: bool = False,
- resnet_out_scale_factor: float = 1.0,
- cross_attention_norm: Optional[str] = None,
- attention_head_dim: Optional[int] = None,
- upsample_type: Optional[str] = None,
- dropout: float = 0.0,
-) -> nn.Module:
- # If attn head dim is not defined, we default it to the number of heads
- if attention_head_dim is None:
- logger.warn(
- f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
- )
- attention_head_dim = num_attention_heads
-
- up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
- if up_block_type == "UpBlock2D":
- return UpBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- prev_output_channel=prev_output_channel,
- temb_channels=temb_channels,
- resolution_idx=resolution_idx,
- dropout=dropout,
- add_upsample=add_upsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_groups=resnet_groups,
- resnet_time_scale_shift=resnet_time_scale_shift,
- )
- elif up_block_type == "ResnetUpsampleBlock2D":
- return ResnetUpsampleBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- prev_output_channel=prev_output_channel,
- temb_channels=temb_channels,
- resolution_idx=resolution_idx,
- dropout=dropout,
- add_upsample=add_upsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_groups=resnet_groups,
- resnet_time_scale_shift=resnet_time_scale_shift,
- skip_time_act=resnet_skip_time_act,
- output_scale_factor=resnet_out_scale_factor,
- )
- elif up_block_type == "CrossAttnUpBlock2D":
- if cross_attention_dim is None:
- raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D")
- return CrossAttnUpBlock2D(
- num_layers=num_layers,
- transformer_layers_per_block=transformer_layers_per_block,
- in_channels=in_channels,
- out_channels=out_channels,
- prev_output_channel=prev_output_channel,
- temb_channels=temb_channels,
- resolution_idx=resolution_idx,
- dropout=dropout,
- add_upsample=add_upsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_groups=resnet_groups,
- cross_attention_dim=cross_attention_dim,
- num_attention_heads=num_attention_heads,
- dual_cross_attention=dual_cross_attention,
- use_linear_projection=use_linear_projection,
- only_cross_attention=only_cross_attention,
- upcast_attention=upcast_attention,
- resnet_time_scale_shift=resnet_time_scale_shift,
- attention_type=attention_type,
- )
- elif up_block_type == "SimpleCrossAttnUpBlock2D":
- if cross_attention_dim is None:
- raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnUpBlock2D")
- return SimpleCrossAttnUpBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- prev_output_channel=prev_output_channel,
- temb_channels=temb_channels,
- resolution_idx=resolution_idx,
- dropout=dropout,
- add_upsample=add_upsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_groups=resnet_groups,
- cross_attention_dim=cross_attention_dim,
- attention_head_dim=attention_head_dim,
- resnet_time_scale_shift=resnet_time_scale_shift,
- skip_time_act=resnet_skip_time_act,
- output_scale_factor=resnet_out_scale_factor,
- only_cross_attention=only_cross_attention,
- cross_attention_norm=cross_attention_norm,
- )
- elif up_block_type == "AttnUpBlock2D":
- if add_upsample is False:
- upsample_type = None
- else:
- upsample_type = upsample_type or "conv" # default to 'conv'
-
- return AttnUpBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- prev_output_channel=prev_output_channel,
- temb_channels=temb_channels,
- resolution_idx=resolution_idx,
- dropout=dropout,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_groups=resnet_groups,
- attention_head_dim=attention_head_dim,
- resnet_time_scale_shift=resnet_time_scale_shift,
- upsample_type=upsample_type,
- )
- elif up_block_type == "SkipUpBlock2D":
- return SkipUpBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- prev_output_channel=prev_output_channel,
- temb_channels=temb_channels,
- resolution_idx=resolution_idx,
- dropout=dropout,
- add_upsample=add_upsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_time_scale_shift=resnet_time_scale_shift,
- )
- elif up_block_type == "AttnSkipUpBlock2D":
- return AttnSkipUpBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- prev_output_channel=prev_output_channel,
- temb_channels=temb_channels,
- resolution_idx=resolution_idx,
- dropout=dropout,
- add_upsample=add_upsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- attention_head_dim=attention_head_dim,
- resnet_time_scale_shift=resnet_time_scale_shift,
- )
- elif up_block_type == "UpDecoderBlock2D":
- return UpDecoderBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- resolution_idx=resolution_idx,
- dropout=dropout,
- add_upsample=add_upsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_groups=resnet_groups,
- resnet_time_scale_shift=resnet_time_scale_shift,
- temb_channels=temb_channels,
- )
- elif up_block_type == "AttnUpDecoderBlock2D":
- return AttnUpDecoderBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- resolution_idx=resolution_idx,
- dropout=dropout,
- add_upsample=add_upsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- resnet_groups=resnet_groups,
- attention_head_dim=attention_head_dim,
- resnet_time_scale_shift=resnet_time_scale_shift,
- temb_channels=temb_channels,
- )
- elif up_block_type == "KUpBlock2D":
- return KUpBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- resolution_idx=resolution_idx,
- dropout=dropout,
- add_upsample=add_upsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- )
- elif up_block_type == "KCrossAttnUpBlock2D":
- return KCrossAttnUpBlock2D(
- num_layers=num_layers,
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- resolution_idx=resolution_idx,
- dropout=dropout,
- add_upsample=add_upsample,
- resnet_eps=resnet_eps,
- resnet_act_fn=resnet_act_fn,
- cross_attention_dim=cross_attention_dim,
- attention_head_dim=attention_head_dim,
- )
-
- raise ValueError(f"{up_block_type} does not exist.")
-
-
-class AutoencoderTinyBlock(nn.Module):
- """
- Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU
- blocks.
-
- Args:
- in_channels (`int`): The number of input channels.
- out_channels (`int`): The number of output channels.
- act_fn (`str`):
- ` The activation function to use. Supported values are `"swish"`, `"mish"`, `"gelu"`, and `"relu"`.
-
- Returns:
- `torch.FloatTensor`: A tensor with the same shape as the input tensor, but with the number of channels equal to
- `out_channels`.
- """
-
- def __init__(self, in_channels: int, out_channels: int, act_fn: str):
- super().__init__()
- act_fn = get_activation(act_fn)
- self.conv = nn.Sequential(
- nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
- act_fn,
- nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
- act_fn,
- nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
- )
- self.skip = (
- nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
- if in_channels != out_channels
- else nn.Identity()
- )
- self.fuse = nn.ReLU()
-
- def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
- return self.fuse(self.conv(x) + self.skip(x))
-
-
-class UNetMidBlock2D(nn.Module):
- """
- A 2D UNet mid-block [`UNetMidBlock2D`] with multiple residual blocks and optional attention blocks.
-
- Args:
- in_channels (`int`): The number of input channels.
- temb_channels (`int`): The number of temporal embedding channels.
- dropout (`float`, *optional*, defaults to 0.0): The dropout rate.
- num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.
- resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.
- resnet_time_scale_shift (`str`, *optional*, defaults to `default`):
- The type of normalization to apply to the time embeddings. This can help to improve the performance of the
- model on tasks with long-range temporal dependencies.
- resnet_act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks.
- resnet_groups (`int`, *optional*, defaults to 32):
- The number of groups to use in the group normalization layers of the resnet blocks.
- attn_groups (`Optional[int]`, *optional*, defaults to None): The number of groups for the attention blocks.
- resnet_pre_norm (`bool`, *optional*, defaults to `True`):
- Whether to use pre-normalization for the resnet blocks.
- add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks.
- attention_head_dim (`int`, *optional*, defaults to 1):
- Dimension of a single attention head. The number of attention heads is determined based on this value and
- the number of input channels.
- output_scale_factor (`float`, *optional*, defaults to 1.0): The output scale factor.
-
- Returns:
- `torch.FloatTensor`: The output of the last residual block, which is a tensor of shape `(batch_size,
- in_channels, height, width)`.
-
- """
-
- def __init__(
- self,
- in_channels: int,
- temb_channels: int,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default", # default, spatial
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- attn_groups: Optional[int] = None,
- resnet_pre_norm: bool = True,
- add_attention: bool = True,
- attention_head_dim: int = 1,
- output_scale_factor: float = 1.0,
- ):
- super().__init__()
- resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
- self.add_attention = add_attention
-
- if attn_groups is None:
- attn_groups = resnet_groups if resnet_time_scale_shift == "default" else None
-
- # there is always at least one resnet
- resnets = [
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=in_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- ]
- attentions = []
-
- if attention_head_dim is None:
- logger.warn(
- f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}."
- )
- attention_head_dim = in_channels
-
- for _ in range(num_layers):
- if self.add_attention:
- attentions.append(
- Attention(
- in_channels,
- heads=in_channels // attention_head_dim,
- dim_head=attention_head_dim,
- rescale_output_factor=output_scale_factor,
- eps=resnet_eps,
- norm_num_groups=attn_groups,
- spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None,
- residual_connection=True,
- bias=True,
- upcast_softmax=True,
- _from_deprecated_attn_block=True,
- )
- )
- else:
- attentions.append(None)
-
- resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=in_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
-
- self.attentions = nn.ModuleList(attentions)
- self.resnets = nn.ModuleList(resnets)
-
- def forward(self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None) -> torch.FloatTensor:
- hidden_states = self.resnets[0](hidden_states, temb)
- for attn, resnet in zip(self.attentions, self.resnets[1:]):
- if attn is not None:
- hidden_states = attn(hidden_states, temb=temb)
- hidden_states = resnet(hidden_states, temb)
-
- return hidden_states
-
-
-class UNetMidBlock2DCrossAttn(nn.Module):
- def __init__(
- self,
- in_channels: int,
- temb_channels: int,
- dropout: float = 0.0,
- num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- num_attention_heads: int = 1,
- output_scale_factor: float = 1.0,
- cross_attention_dim: int = 1280,
- dual_cross_attention: bool = False,
- use_linear_projection: bool = False,
- upcast_attention: bool = False,
- attention_type: str = "default",
- ):
- super().__init__()
-
- 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)
-
- # support for variable transformer layers per block
- if isinstance(transformer_layers_per_block, int):
- transformer_layers_per_block = [transformer_layers_per_block] * num_layers
-
- # there is always at least one resnet
- resnets = [
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=in_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- ]
- attentions = []
-
- for i in range(num_layers):
- if not dual_cross_attention:
- attentions.append(
- Transformer2DModel(
- num_attention_heads,
- in_channels // num_attention_heads,
- in_channels=in_channels,
- num_layers=transformer_layers_per_block[i],
- cross_attention_dim=cross_attention_dim,
- norm_num_groups=resnet_groups,
- use_linear_projection=use_linear_projection,
- upcast_attention=upcast_attention,
- attention_type=attention_type,
- )
- )
- else:
- attentions.append(
- DualTransformer2DModel(
- num_attention_heads,
- in_channels // num_attention_heads,
- in_channels=in_channels,
- num_layers=1,
- cross_attention_dim=cross_attention_dim,
- norm_num_groups=resnet_groups,
- )
- )
- resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=in_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
-
- self.attentions = nn.ModuleList(attentions)
- self.resnets = nn.ModuleList(resnets)
-
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- temb: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- ) -> torch.FloatTensor:
- lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
- hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale)
- garment_features = []
- for attn, resnet in zip(self.attentions, self.resnets[1:]):
- if self.training and self.gradient_checkpointing:
-
- def create_custom_forward(module, return_dict=None):
- def custom_forward(*inputs):
- if return_dict is not None:
- return module(*inputs, return_dict=return_dict)
- else:
- return module(*inputs)
-
- return custom_forward
-
- ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
- # hidden_states = attn(
- hidden_states,out_garment_feat = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- cross_attention_kwargs=cross_attention_kwargs,
- attention_mask=attention_mask,
- encoder_attention_mask=encoder_attention_mask,
- return_dict=False,
- )
- hidden_states=hidden_states[0]
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet),
- hidden_states,
- temb,
- **ckpt_kwargs,
- )
- else:
- # hidden_states= attn(
- hidden_states,out_garment_feat = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- cross_attention_kwargs=cross_attention_kwargs,
- attention_mask=attention_mask,
- encoder_attention_mask=encoder_attention_mask,
- return_dict=False,
- )
- hidden_states=hidden_states[0]
- hidden_states = resnet(hidden_states, temb, scale=lora_scale)
- garment_features += out_garment_feat
- return hidden_states,garment_features
- # return hidden_states
-
-
-class UNetMidBlock2DSimpleCrossAttn(nn.Module):
- def __init__(
- self,
- in_channels: int,
- temb_channels: int,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- attention_head_dim: int = 1,
- output_scale_factor: float = 1.0,
- cross_attention_dim: int = 1280,
- skip_time_act: bool = False,
- only_cross_attention: bool = False,
- cross_attention_norm: Optional[str] = None,
- ):
- super().__init__()
-
- self.has_cross_attention = True
-
- self.attention_head_dim = attention_head_dim
- resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
-
- self.num_heads = in_channels // self.attention_head_dim
-
- # there is always at least one resnet
- resnets = [
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=in_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- skip_time_act=skip_time_act,
- )
- ]
- attentions = []
-
- for _ in range(num_layers):
- processor = (
- AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor()
- )
-
- attentions.append(
- Attention(
- query_dim=in_channels,
- cross_attention_dim=in_channels,
- heads=self.num_heads,
- dim_head=self.attention_head_dim,
- added_kv_proj_dim=cross_attention_dim,
- norm_num_groups=resnet_groups,
- bias=True,
- upcast_softmax=True,
- only_cross_attention=only_cross_attention,
- cross_attention_norm=cross_attention_norm,
- processor=processor,
- )
- )
- resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=in_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- skip_time_act=skip_time_act,
- )
- )
-
- self.attentions = nn.ModuleList(attentions)
- self.resnets = nn.ModuleList(resnets)
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- temb: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- ) -> torch.FloatTensor:
- cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
- lora_scale = cross_attention_kwargs.get("scale", 1.0)
-
- if attention_mask is None:
- # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask.
- mask = None if encoder_hidden_states is None else encoder_attention_mask
- else:
- # when attention_mask is defined: we don't even check for encoder_attention_mask.
- # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks.
- # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask.
- # then we can simplify this whole if/else block to:
- # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask
- mask = attention_mask
-
- hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale)
- for attn, resnet in zip(self.attentions, self.resnets[1:]):
- # attn
- hidden_states = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- attention_mask=mask,
- **cross_attention_kwargs,
- )
-
- # resnet
- hidden_states = resnet(hidden_states, temb, scale=lora_scale)
-
- return hidden_states
-
-
-class AttnDownBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- temb_channels: int,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- attention_head_dim: int = 1,
- output_scale_factor: float = 1.0,
- downsample_padding: int = 1,
- downsample_type: str = "conv",
- ):
- super().__init__()
- resnets = []
- attentions = []
- self.downsample_type = downsample_type
-
- if attention_head_dim is None:
- logger.warn(
- f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}."
- )
- attention_head_dim = out_channels
-
- for i in range(num_layers):
- in_channels = in_channels if i == 0 else out_channels
- resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
- attentions.append(
- Attention(
- out_channels,
- heads=out_channels // attention_head_dim,
- dim_head=attention_head_dim,
- rescale_output_factor=output_scale_factor,
- eps=resnet_eps,
- norm_num_groups=resnet_groups,
- residual_connection=True,
- bias=True,
- upcast_softmax=True,
- _from_deprecated_attn_block=True,
- )
- )
-
- self.attentions = nn.ModuleList(attentions)
- self.resnets = nn.ModuleList(resnets)
-
- if downsample_type == "conv":
- self.downsamplers = nn.ModuleList(
- [
- Downsample2D(
- out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
- )
- ]
- )
- elif downsample_type == "resnet":
- self.downsamplers = nn.ModuleList(
- [
- ResnetBlock2D(
- in_channels=out_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- down=True,
- )
- ]
- )
- else:
- self.downsamplers = None
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- temb: Optional[torch.FloatTensor] = None,
- upsample_size: Optional[int] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
- cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
-
- lora_scale = cross_attention_kwargs.get("scale", 1.0)
-
- output_states = ()
-
- for resnet, attn in zip(self.resnets, self.attentions):
- cross_attention_kwargs.update({"scale": lora_scale})
- hidden_states = resnet(hidden_states, temb, scale=lora_scale)
- hidden_states = attn(hidden_states, **cross_attention_kwargs)
- output_states = output_states + (hidden_states,)
-
- if self.downsamplers is not None:
- for downsampler in self.downsamplers:
- if self.downsample_type == "resnet":
- hidden_states = downsampler(hidden_states, temb=temb, scale=lora_scale)
- else:
- hidden_states = downsampler(hidden_states, scale=lora_scale)
-
- output_states += (hidden_states,)
-
- return hidden_states, output_states
-
-
-class CrossAttnDownBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- temb_channels: int,
- dropout: float = 0.0,
- num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- num_attention_heads: int = 1,
- cross_attention_dim: int = 1280,
- output_scale_factor: float = 1.0,
- downsample_padding: int = 1,
- add_downsample: bool = True,
- dual_cross_attention: bool = False,
- use_linear_projection: bool = False,
- only_cross_attention: bool = False,
- upcast_attention: bool = False,
- attention_type: str = "default",
- ):
- super().__init__()
- resnets = []
- attentions = []
-
- self.has_cross_attention = True
- self.num_attention_heads = num_attention_heads
- if isinstance(transformer_layers_per_block, int):
- transformer_layers_per_block = [transformer_layers_per_block] * num_layers
-
- for i in range(num_layers):
- in_channels = in_channels if i == 0 else out_channels
- resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
- if not dual_cross_attention:
- attentions.append(
- Transformer2DModel(
- num_attention_heads,
- out_channels // num_attention_heads,
- in_channels=out_channels,
- num_layers=transformer_layers_per_block[i],
- cross_attention_dim=cross_attention_dim,
- norm_num_groups=resnet_groups,
- use_linear_projection=use_linear_projection,
- only_cross_attention=only_cross_attention,
- upcast_attention=upcast_attention,
- attention_type=attention_type,
- )
- )
- else:
- attentions.append(
- DualTransformer2DModel(
- num_attention_heads,
- out_channels // num_attention_heads,
- in_channels=out_channels,
- num_layers=1,
- cross_attention_dim=cross_attention_dim,
- norm_num_groups=resnet_groups,
- )
- )
- self.attentions = nn.ModuleList(attentions)
- self.resnets = nn.ModuleList(resnets)
-
- if add_downsample:
- self.downsamplers = nn.ModuleList(
- [
- Downsample2D(
- out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
- )
- ]
- )
- else:
- self.downsamplers = None
-
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- temb: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- additional_residuals: Optional[torch.FloatTensor] = None,
- ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
- output_states = ()
-
- lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
-
- blocks = list(zip(self.resnets, self.attentions))
- garment_features = []
- for i, (resnet, attn) in enumerate(blocks):
- if self.training and self.gradient_checkpointing:
-
- def create_custom_forward(module, return_dict=None):
- def custom_forward(*inputs):
- if return_dict is not None:
- return module(*inputs, return_dict=return_dict)
- else:
- return module(*inputs)
-
- return custom_forward
-
- ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet),
- hidden_states,
- temb,
- **ckpt_kwargs,
- )
- hidden_states,out_garment_feat = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- cross_attention_kwargs=cross_attention_kwargs,
- attention_mask=attention_mask,
- encoder_attention_mask=encoder_attention_mask,
- return_dict=False,
- )
- hidden_states=hidden_states[0]
- else:
- hidden_states = resnet(hidden_states, temb, scale=lora_scale)
- hidden_states,out_garment_feat = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- cross_attention_kwargs=cross_attention_kwargs,
- attention_mask=attention_mask,
- encoder_attention_mask=encoder_attention_mask,
- return_dict=False,
- )
- hidden_states=hidden_states[0]
- garment_features += out_garment_feat
- # apply additional residuals to the output of the last pair of resnet and attention blocks
- if i == len(blocks) - 1 and additional_residuals is not None:
- hidden_states = hidden_states + additional_residuals
-
- output_states = output_states + (hidden_states,)
-
- if self.downsamplers is not None:
- for downsampler in self.downsamplers:
- hidden_states = downsampler(hidden_states, scale=lora_scale)
-
- output_states = output_states + (hidden_states,)
-
- return hidden_states, output_states,garment_features
-
-
-class DownBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- temb_channels: int,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- output_scale_factor: float = 1.0,
- add_downsample: bool = True,
- downsample_padding: int = 1,
- ):
- super().__init__()
- resnets = []
-
- for i in range(num_layers):
- in_channels = in_channels if i == 0 else out_channels
- resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
-
- self.resnets = nn.ModuleList(resnets)
-
- if add_downsample:
- self.downsamplers = nn.ModuleList(
- [
- Downsample2D(
- out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
- )
- ]
- )
- else:
- self.downsamplers = None
-
- self.gradient_checkpointing = False
-
- def forward(
- self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
- ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
- output_states = ()
-
- for resnet in self.resnets:
- if self.training and self.gradient_checkpointing:
-
- def create_custom_forward(module):
- def custom_forward(*inputs):
- return module(*inputs)
-
- return custom_forward
-
- if is_torch_version(">=", "1.11.0"):
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
- )
- else:
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet), hidden_states, temb
- )
- else:
- hidden_states = resnet(hidden_states, temb, scale=scale)
-
- output_states = output_states + (hidden_states,)
-
- if self.downsamplers is not None:
- for downsampler in self.downsamplers:
- hidden_states = downsampler(hidden_states, scale=scale)
-
- output_states = output_states + (hidden_states,)
-
- return hidden_states, output_states
-
-
-class DownEncoderBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- output_scale_factor: float = 1.0,
- add_downsample: bool = True,
- downsample_padding: int = 1,
- ):
- super().__init__()
- resnets = []
-
- for i in range(num_layers):
- in_channels = in_channels if i == 0 else out_channels
- resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=None,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
-
- self.resnets = nn.ModuleList(resnets)
-
- if add_downsample:
- self.downsamplers = nn.ModuleList(
- [
- Downsample2D(
- out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
- )
- ]
- )
- else:
- self.downsamplers = None
-
- def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
- for resnet in self.resnets:
- hidden_states = resnet(hidden_states, temb=None, scale=scale)
-
- if self.downsamplers is not None:
- for downsampler in self.downsamplers:
- hidden_states = downsampler(hidden_states, scale)
-
- return hidden_states
-
-
-class AttnDownEncoderBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- attention_head_dim: int = 1,
- output_scale_factor: float = 1.0,
- add_downsample: bool = True,
- downsample_padding: int = 1,
- ):
- super().__init__()
- resnets = []
- attentions = []
-
- if attention_head_dim is None:
- logger.warn(
- f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}."
- )
- attention_head_dim = out_channels
-
- for i in range(num_layers):
- in_channels = in_channels if i == 0 else out_channels
- resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=None,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
- attentions.append(
- Attention(
- out_channels,
- heads=out_channels // attention_head_dim,
- dim_head=attention_head_dim,
- rescale_output_factor=output_scale_factor,
- eps=resnet_eps,
- norm_num_groups=resnet_groups,
- residual_connection=True,
- bias=True,
- upcast_softmax=True,
- _from_deprecated_attn_block=True,
- )
- )
-
- self.attentions = nn.ModuleList(attentions)
- self.resnets = nn.ModuleList(resnets)
-
- if add_downsample:
- self.downsamplers = nn.ModuleList(
- [
- Downsample2D(
- out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
- )
- ]
- )
- else:
- self.downsamplers = None
-
- def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
- for resnet, attn in zip(self.resnets, self.attentions):
- hidden_states = resnet(hidden_states, temb=None, scale=scale)
- cross_attention_kwargs = {"scale": scale}
- hidden_states = attn(hidden_states, **cross_attention_kwargs)
-
- if self.downsamplers is not None:
- for downsampler in self.downsamplers:
- hidden_states = downsampler(hidden_states, scale)
-
- return hidden_states
-
-
-class AttnSkipDownBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- temb_channels: int,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_pre_norm: bool = True,
- attention_head_dim: int = 1,
- output_scale_factor: float = np.sqrt(2.0),
- add_downsample: bool = True,
- ):
- super().__init__()
- self.attentions = nn.ModuleList([])
- self.resnets = nn.ModuleList([])
-
- if attention_head_dim is None:
- logger.warn(
- f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}."
- )
- attention_head_dim = out_channels
-
- for i in range(num_layers):
- in_channels = in_channels if i == 0 else out_channels
- self.resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=min(in_channels // 4, 32),
- groups_out=min(out_channels // 4, 32),
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
- self.attentions.append(
- Attention(
- out_channels,
- heads=out_channels // attention_head_dim,
- dim_head=attention_head_dim,
- rescale_output_factor=output_scale_factor,
- eps=resnet_eps,
- norm_num_groups=32,
- residual_connection=True,
- bias=True,
- upcast_softmax=True,
- _from_deprecated_attn_block=True,
- )
- )
-
- if add_downsample:
- self.resnet_down = ResnetBlock2D(
- in_channels=out_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=min(out_channels // 4, 32),
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- use_in_shortcut=True,
- down=True,
- kernel="fir",
- )
- self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)])
- self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1))
- else:
- self.resnet_down = None
- self.downsamplers = None
- self.skip_conv = None
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- temb: Optional[torch.FloatTensor] = None,
- skip_sample: Optional[torch.FloatTensor] = None,
- scale: float = 1.0,
- ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...], torch.FloatTensor]:
- output_states = ()
-
- for resnet, attn in zip(self.resnets, self.attentions):
- hidden_states = resnet(hidden_states, temb, scale=scale)
- cross_attention_kwargs = {"scale": scale}
- hidden_states = attn(hidden_states, **cross_attention_kwargs)
- output_states += (hidden_states,)
-
- if self.downsamplers is not None:
- hidden_states = self.resnet_down(hidden_states, temb, scale=scale)
- for downsampler in self.downsamplers:
- skip_sample = downsampler(skip_sample)
-
- hidden_states = self.skip_conv(skip_sample) + hidden_states
-
- output_states += (hidden_states,)
-
- return hidden_states, output_states, skip_sample
-
-
-class SkipDownBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- temb_channels: int,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_pre_norm: bool = True,
- output_scale_factor: float = np.sqrt(2.0),
- add_downsample: bool = True,
- downsample_padding: int = 1,
- ):
- super().__init__()
- self.resnets = nn.ModuleList([])
-
- for i in range(num_layers):
- in_channels = in_channels if i == 0 else out_channels
- self.resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=min(in_channels // 4, 32),
- groups_out=min(out_channels // 4, 32),
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
-
- if add_downsample:
- self.resnet_down = ResnetBlock2D(
- in_channels=out_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=min(out_channels // 4, 32),
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- use_in_shortcut=True,
- down=True,
- kernel="fir",
- )
- self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)])
- self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1))
- else:
- self.resnet_down = None
- self.downsamplers = None
- self.skip_conv = None
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- temb: Optional[torch.FloatTensor] = None,
- skip_sample: Optional[torch.FloatTensor] = None,
- scale: float = 1.0,
- ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...], torch.FloatTensor]:
- output_states = ()
-
- for resnet in self.resnets:
- hidden_states = resnet(hidden_states, temb, scale)
- output_states += (hidden_states,)
-
- if self.downsamplers is not None:
- hidden_states = self.resnet_down(hidden_states, temb, scale)
- for downsampler in self.downsamplers:
- skip_sample = downsampler(skip_sample)
-
- hidden_states = self.skip_conv(skip_sample) + hidden_states
-
- output_states += (hidden_states,)
-
- return hidden_states, output_states, skip_sample
-
-
-class ResnetDownsampleBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- temb_channels: int,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- output_scale_factor: float = 1.0,
- add_downsample: bool = True,
- skip_time_act: bool = False,
- ):
- super().__init__()
- resnets = []
-
- for i in range(num_layers):
- in_channels = in_channels if i == 0 else out_channels
- resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- skip_time_act=skip_time_act,
- )
- )
-
- self.resnets = nn.ModuleList(resnets)
-
- if add_downsample:
- self.downsamplers = nn.ModuleList(
- [
- ResnetBlock2D(
- in_channels=out_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- skip_time_act=skip_time_act,
- down=True,
- )
- ]
- )
- else:
- self.downsamplers = None
-
- self.gradient_checkpointing = False
-
- def forward(
- self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
- ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
- output_states = ()
-
- for resnet in self.resnets:
- if self.training and self.gradient_checkpointing:
-
- def create_custom_forward(module):
- def custom_forward(*inputs):
- return module(*inputs)
-
- return custom_forward
-
- if is_torch_version(">=", "1.11.0"):
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
- )
- else:
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet), hidden_states, temb
- )
- else:
- hidden_states = resnet(hidden_states, temb, scale)
-
- output_states = output_states + (hidden_states,)
-
- if self.downsamplers is not None:
- for downsampler in self.downsamplers:
- hidden_states = downsampler(hidden_states, temb, scale)
-
- output_states = output_states + (hidden_states,)
-
- return hidden_states, output_states
-
-
-class SimpleCrossAttnDownBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- temb_channels: int,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- attention_head_dim: int = 1,
- cross_attention_dim: int = 1280,
- output_scale_factor: float = 1.0,
- add_downsample: bool = True,
- skip_time_act: bool = False,
- only_cross_attention: bool = False,
- cross_attention_norm: Optional[str] = None,
- ):
- super().__init__()
-
- self.has_cross_attention = True
-
- resnets = []
- attentions = []
-
- self.attention_head_dim = attention_head_dim
- self.num_heads = out_channels // self.attention_head_dim
-
- for i in range(num_layers):
- in_channels = in_channels if i == 0 else out_channels
- resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- skip_time_act=skip_time_act,
- )
- )
-
- processor = (
- AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor()
- )
-
- attentions.append(
- Attention(
- query_dim=out_channels,
- cross_attention_dim=out_channels,
- heads=self.num_heads,
- dim_head=attention_head_dim,
- added_kv_proj_dim=cross_attention_dim,
- norm_num_groups=resnet_groups,
- bias=True,
- upcast_softmax=True,
- only_cross_attention=only_cross_attention,
- cross_attention_norm=cross_attention_norm,
- processor=processor,
- )
- )
- self.attentions = nn.ModuleList(attentions)
- self.resnets = nn.ModuleList(resnets)
-
- if add_downsample:
- self.downsamplers = nn.ModuleList(
- [
- ResnetBlock2D(
- in_channels=out_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- skip_time_act=skip_time_act,
- down=True,
- )
- ]
- )
- else:
- self.downsamplers = None
-
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- temb: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
- output_states = ()
- cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
-
- lora_scale = cross_attention_kwargs.get("scale", 1.0)
-
- if attention_mask is None:
- # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask.
- mask = None if encoder_hidden_states is None else encoder_attention_mask
- else:
- # when attention_mask is defined: we don't even check for encoder_attention_mask.
- # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks.
- # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask.
- # then we can simplify this whole if/else block to:
- # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask
- mask = attention_mask
-
- for resnet, attn in zip(self.resnets, self.attentions):
- if self.training and self.gradient_checkpointing:
-
- def create_custom_forward(module, return_dict=None):
- def custom_forward(*inputs):
- if return_dict is not None:
- return module(*inputs, return_dict=return_dict)
- else:
- return module(*inputs)
-
- return custom_forward
-
- hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
- hidden_states = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- attention_mask=mask,
- **cross_attention_kwargs,
- )
- else:
- hidden_states = resnet(hidden_states, temb, scale=lora_scale)
-
- hidden_states = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- attention_mask=mask,
- **cross_attention_kwargs,
- )
-
- output_states = output_states + (hidden_states,)
-
- if self.downsamplers is not None:
- for downsampler in self.downsamplers:
- hidden_states = downsampler(hidden_states, temb, scale=lora_scale)
-
- output_states = output_states + (hidden_states,)
-
- return hidden_states, output_states
-
-
-class KDownBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- temb_channels: int,
- dropout: float = 0.0,
- num_layers: int = 4,
- resnet_eps: float = 1e-5,
- resnet_act_fn: str = "gelu",
- resnet_group_size: int = 32,
- add_downsample: bool = False,
- ):
- super().__init__()
- resnets = []
-
- for i in range(num_layers):
- in_channels = in_channels if i == 0 else out_channels
- groups = in_channels // resnet_group_size
- groups_out = out_channels // resnet_group_size
-
- resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=out_channels,
- dropout=dropout,
- temb_channels=temb_channels,
- groups=groups,
- groups_out=groups_out,
- eps=resnet_eps,
- non_linearity=resnet_act_fn,
- time_embedding_norm="ada_group",
- conv_shortcut_bias=False,
- )
- )
-
- self.resnets = nn.ModuleList(resnets)
-
- if add_downsample:
- # YiYi's comments- might be able to use FirDownsample2D, look into details later
- self.downsamplers = nn.ModuleList([KDownsample2D()])
- else:
- self.downsamplers = None
-
- self.gradient_checkpointing = False
-
- def forward(
- self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
- ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
- output_states = ()
-
- for resnet in self.resnets:
- if self.training and self.gradient_checkpointing:
-
- def create_custom_forward(module):
- def custom_forward(*inputs):
- return module(*inputs)
-
- return custom_forward
-
- if is_torch_version(">=", "1.11.0"):
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
- )
- else:
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet), hidden_states, temb
- )
- else:
- hidden_states = resnet(hidden_states, temb, scale)
-
- output_states += (hidden_states,)
-
- if self.downsamplers is not None:
- for downsampler in self.downsamplers:
- hidden_states = downsampler(hidden_states)
-
- return hidden_states, output_states
-
-
-class KCrossAttnDownBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- temb_channels: int,
- cross_attention_dim: int,
- dropout: float = 0.0,
- num_layers: int = 4,
- resnet_group_size: int = 32,
- add_downsample: bool = True,
- attention_head_dim: int = 64,
- add_self_attention: bool = False,
- resnet_eps: float = 1e-5,
- resnet_act_fn: str = "gelu",
- ):
- super().__init__()
- resnets = []
- attentions = []
-
- self.has_cross_attention = True
-
- for i in range(num_layers):
- in_channels = in_channels if i == 0 else out_channels
- groups = in_channels // resnet_group_size
- groups_out = out_channels // resnet_group_size
-
- resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=out_channels,
- dropout=dropout,
- temb_channels=temb_channels,
- groups=groups,
- groups_out=groups_out,
- eps=resnet_eps,
- non_linearity=resnet_act_fn,
- time_embedding_norm="ada_group",
- conv_shortcut_bias=False,
- )
- )
- attentions.append(
- KAttentionBlock(
- out_channels,
- out_channels // attention_head_dim,
- attention_head_dim,
- cross_attention_dim=cross_attention_dim,
- temb_channels=temb_channels,
- attention_bias=True,
- add_self_attention=add_self_attention,
- cross_attention_norm="layer_norm",
- group_size=resnet_group_size,
- )
- )
-
- self.resnets = nn.ModuleList(resnets)
- self.attentions = nn.ModuleList(attentions)
-
- if add_downsample:
- self.downsamplers = nn.ModuleList([KDownsample2D()])
- else:
- self.downsamplers = None
-
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- temb: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
- output_states = ()
- lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
-
- for resnet, attn in zip(self.resnets, self.attentions):
- if self.training and self.gradient_checkpointing:
-
- def create_custom_forward(module, return_dict=None):
- def custom_forward(*inputs):
- if return_dict is not None:
- return module(*inputs, return_dict=return_dict)
- else:
- return module(*inputs)
-
- return custom_forward
-
- ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet),
- hidden_states,
- temb,
- **ckpt_kwargs,
- )
- hidden_states = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- emb=temb,
- attention_mask=attention_mask,
- cross_attention_kwargs=cross_attention_kwargs,
- encoder_attention_mask=encoder_attention_mask,
- )
- else:
- hidden_states = resnet(hidden_states, temb, scale=lora_scale)
- hidden_states = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- emb=temb,
- attention_mask=attention_mask,
- cross_attention_kwargs=cross_attention_kwargs,
- encoder_attention_mask=encoder_attention_mask,
- )
-
- if self.downsamplers is None:
- output_states += (None,)
- else:
- output_states += (hidden_states,)
-
- if self.downsamplers is not None:
- for downsampler in self.downsamplers:
- hidden_states = downsampler(hidden_states)
-
- return hidden_states, output_states
-
-
-class AttnUpBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- prev_output_channel: int,
- out_channels: int,
- temb_channels: int,
- resolution_idx: int = None,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- attention_head_dim: int = 1,
- output_scale_factor: float = 1.0,
- upsample_type: str = "conv",
- ):
- super().__init__()
- resnets = []
- attentions = []
-
- self.upsample_type = upsample_type
-
- if attention_head_dim is None:
- logger.warn(
- f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}."
- )
- attention_head_dim = out_channels
-
- for i in range(num_layers):
- res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
- resnet_in_channels = prev_output_channel if i == 0 else out_channels
-
- resnets.append(
- ResnetBlock2D(
- in_channels=resnet_in_channels + res_skip_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
- attentions.append(
- Attention(
- out_channels,
- heads=out_channels // attention_head_dim,
- dim_head=attention_head_dim,
- rescale_output_factor=output_scale_factor,
- eps=resnet_eps,
- norm_num_groups=resnet_groups,
- residual_connection=True,
- bias=True,
- upcast_softmax=True,
- _from_deprecated_attn_block=True,
- )
- )
-
- self.attentions = nn.ModuleList(attentions)
- self.resnets = nn.ModuleList(resnets)
-
- if upsample_type == "conv":
- self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
- elif upsample_type == "resnet":
- self.upsamplers = nn.ModuleList(
- [
- ResnetBlock2D(
- in_channels=out_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- up=True,
- )
- ]
- )
- else:
- self.upsamplers = None
-
- self.resolution_idx = resolution_idx
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
- temb: Optional[torch.FloatTensor] = None,
- upsample_size: Optional[int] = None,
- scale: float = 1.0,
- ) -> torch.FloatTensor:
- for resnet, attn in zip(self.resnets, self.attentions):
- # pop res hidden states
- res_hidden_states = res_hidden_states_tuple[-1]
- res_hidden_states_tuple = res_hidden_states_tuple[:-1]
- hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
-
- hidden_states = resnet(hidden_states, temb, scale=scale)
- cross_attention_kwargs = {"scale": scale}
- hidden_states = attn(hidden_states, **cross_attention_kwargs)
-
- if self.upsamplers is not None:
- for upsampler in self.upsamplers:
- if self.upsample_type == "resnet":
- hidden_states = upsampler(hidden_states, temb=temb, scale=scale)
- else:
- hidden_states = upsampler(hidden_states, scale=scale)
-
- return hidden_states
-
-
-class CrossAttnUpBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- prev_output_channel: int,
- temb_channels: int,
- resolution_idx: Optional[int] = None,
- dropout: float = 0.0,
- num_layers: int = 1,
- transformer_layers_per_block: Union[int, Tuple[int]] = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- num_attention_heads: int = 1,
- cross_attention_dim: int = 1280,
- output_scale_factor: float = 1.0,
- add_upsample: bool = True,
- dual_cross_attention: bool = False,
- use_linear_projection: bool = False,
- only_cross_attention: bool = False,
- upcast_attention: bool = False,
- attention_type: str = "default",
- ):
- super().__init__()
- resnets = []
- attentions = []
-
- self.has_cross_attention = True
- self.num_attention_heads = num_attention_heads
-
- if isinstance(transformer_layers_per_block, int):
- transformer_layers_per_block = [transformer_layers_per_block] * num_layers
-
- for i in range(num_layers):
- res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
- resnet_in_channels = prev_output_channel if i == 0 else out_channels
-
- resnets.append(
- ResnetBlock2D(
- in_channels=resnet_in_channels + res_skip_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
- if not dual_cross_attention:
- attentions.append(
- Transformer2DModel(
- num_attention_heads,
- out_channels // num_attention_heads,
- in_channels=out_channels,
- num_layers=transformer_layers_per_block[i],
- cross_attention_dim=cross_attention_dim,
- norm_num_groups=resnet_groups,
- use_linear_projection=use_linear_projection,
- only_cross_attention=only_cross_attention,
- upcast_attention=upcast_attention,
- attention_type=attention_type,
- )
- )
- else:
- attentions.append(
- DualTransformer2DModel(
- num_attention_heads,
- out_channels // num_attention_heads,
- in_channels=out_channels,
- num_layers=1,
- cross_attention_dim=cross_attention_dim,
- norm_num_groups=resnet_groups,
- )
- )
- self.attentions = nn.ModuleList(attentions)
- self.resnets = nn.ModuleList(resnets)
-
- if add_upsample:
- self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
- else:
- self.upsamplers = None
-
- self.gradient_checkpointing = False
- self.resolution_idx = resolution_idx
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
- temb: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- upsample_size: Optional[int] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- ) -> torch.FloatTensor:
- lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
- is_freeu_enabled = (
- getattr(self, "s1", None)
- and getattr(self, "s2", None)
- and getattr(self, "b1", None)
- and getattr(self, "b2", None)
- )
- garment_features = []
- for resnet, attn in zip(self.resnets, self.attentions):
- # pop res hidden states
- # print("h.shape")
- # print(h.shape)
- # print("hidden_states.shape)
- # print(hidden_states.shape)
- # print("attn_block")
- # print(attn)
-
- res_hidden_states = res_hidden_states_tuple[-1]
- res_hidden_states_tuple = res_hidden_states_tuple[:-1]
-
- # FreeU: Only operate on the first two stages
- if is_freeu_enabled:
- hidden_states, res_hidden_states = apply_freeu(
- self.resolution_idx,
- hidden_states,
- res_hidden_states,
- s1=self.s1,
- s2=self.s2,
- b1=self.b1,
- b2=self.b2,
- )
-
- hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
- # print(hidden_states.shape)
- # print(encoder_hidden_states.shape)
- if self.training and self.gradient_checkpointing:
-
- def create_custom_forward(module, return_dict=None):
- def custom_forward(*inputs):
- if return_dict is not None:
- return module(*inputs, return_dict=return_dict)
- else:
- return module(*inputs)
-
- return custom_forward
-
- ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet),
- hidden_states,
- temb,
- **ckpt_kwargs,
- )
- hidden_states,out_garment_feat = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- cross_attention_kwargs=cross_attention_kwargs,
- attention_mask=attention_mask,
- encoder_attention_mask=encoder_attention_mask,
- return_dict=False,
- )
- hidden_states=hidden_states[0]
- else:
- hidden_states = resnet(hidden_states, temb, scale=lora_scale)
- hidden_states,out_garment_feat = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- cross_attention_kwargs=cross_attention_kwargs,
- attention_mask=attention_mask,
- encoder_attention_mask=encoder_attention_mask,
- return_dict=False,
- )
- hidden_states=hidden_states[0]
- garment_features += out_garment_feat
- if self.upsamplers is not None:
- for upsampler in self.upsamplers:
- hidden_states = upsampler(hidden_states, upsample_size, scale=lora_scale)
-
- return hidden_states,garment_features
-
-
-class UpBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- prev_output_channel: int,
- out_channels: int,
- temb_channels: int,
- resolution_idx: Optional[int] = None,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- output_scale_factor: float = 1.0,
- add_upsample: bool = True,
- ):
- super().__init__()
- resnets = []
-
- for i in range(num_layers):
- res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
- resnet_in_channels = prev_output_channel if i == 0 else out_channels
-
- resnets.append(
- ResnetBlock2D(
- in_channels=resnet_in_channels + res_skip_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
-
- self.resnets = nn.ModuleList(resnets)
-
- if add_upsample:
- self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
- else:
- self.upsamplers = None
-
- self.gradient_checkpointing = False
- self.resolution_idx = resolution_idx
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
- temb: Optional[torch.FloatTensor] = None,
- upsample_size: Optional[int] = None,
- scale: float = 1.0,
- ) -> torch.FloatTensor:
- is_freeu_enabled = (
- getattr(self, "s1", None)
- and getattr(self, "s2", None)
- and getattr(self, "b1", None)
- and getattr(self, "b2", None)
- )
-
- for resnet in self.resnets:
- # pop res hidden states
- res_hidden_states = res_hidden_states_tuple[-1]
- res_hidden_states_tuple = res_hidden_states_tuple[:-1]
-
- # FreeU: Only operate on the first two stages
- if is_freeu_enabled:
- hidden_states, res_hidden_states = apply_freeu(
- self.resolution_idx,
- hidden_states,
- res_hidden_states,
- s1=self.s1,
- s2=self.s2,
- b1=self.b1,
- b2=self.b2,
- )
-
- hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
-
- if self.training and self.gradient_checkpointing:
-
- def create_custom_forward(module):
- def custom_forward(*inputs):
- return module(*inputs)
-
- return custom_forward
-
- if is_torch_version(">=", "1.11.0"):
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
- )
- else:
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet), hidden_states, temb
- )
- else:
- hidden_states = resnet(hidden_states, temb, scale=scale)
-
- if self.upsamplers is not None:
- for upsampler in self.upsamplers:
- hidden_states = upsampler(hidden_states, upsample_size, scale=scale)
-
- return hidden_states
- # def forward(
- # self,
- # hidden_states: torch.FloatTensor,
- # res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
- # temb: Optional[torch.FloatTensor] = None,
- # upsample_size: Optional[int] = None,
- # scale: float = 1.0,
- # zero_block=None,
- # hint=None,
- # ) -> torch.FloatTensor:
- # is_freeu_enabled = (
- # getattr(self, "s1", None)
- # and getattr(self, "s2", None)
- # and getattr(self, "b1", None)
- # and getattr(self, "b2", None)
- # )
-
- # # print(len(self.resnets))
- # # print(len(zero_block))
- # # print(len(hint))
- # # for resnet in self.resnets:
- # for resnet, zero,h in zip(self.resnets,zero_block,hint):
-
- # # pop res hidden states
- # res_hidden_states = res_hidden_states_tuple[-1]
- # res_hidden_states_tuple = res_hidden_states_tuple[:-1]
-
- # res_hidden_states = res_hidden_states + zero(h)
- # # FreeU: Only operate on the first two stages
- # if is_freeu_enabled:
- # hidden_states, res_hidden_states = apply_freeu(
- # self.resolution_idx,
- # hidden_states,
- # res_hidden_states,
- # s1=self.s1,
- # s2=self.s2,
- # b1=self.b1,
- # b2=self.b2,
- # )
-
- # # print(hidden_states.shape)
- # # # print(h.shape)
- # # print(res_hidden_states.shape)
- # hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
- # # print(hidden_states.shape)
-
- # if self.training and self.gradient_checkpointing:
-
- # def create_custom_forward(module):
- # def custom_forward(*inputs):
- # return module(*inputs)
-
- # return custom_forward
-
- # if is_torch_version(">=", "1.11.0"):
- # hidden_states = torch.utils.checkpoint.checkpoint(
- # create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
- # )
- # else:
- # hidden_states = torch.utils.checkpoint.checkpoint(
- # create_custom_forward(resnet), hidden_states, temb
- # )
- # else:
- # hidden_states = resnet(hidden_states, temb, scale=scale)
-
- # if self.upsamplers is not None:
- # for upsampler in self.upsamplers:
- # hidden_states = upsampler(hidden_states, upsample_size, scale=scale)
-
- # return hidden_states
-
-
-class UpDecoderBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- resolution_idx: Optional[int] = None,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default", # default, spatial
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- output_scale_factor: float = 1.0,
- add_upsample: bool = True,
- temb_channels: Optional[int] = None,
- ):
- super().__init__()
- resnets = []
-
- for i in range(num_layers):
- input_channels = in_channels if i == 0 else out_channels
-
- resnets.append(
- ResnetBlock2D(
- in_channels=input_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
-
- self.resnets = nn.ModuleList(resnets)
-
- if add_upsample:
- self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
- else:
- self.upsamplers = None
-
- self.resolution_idx = resolution_idx
-
- def forward(
- self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
- ) -> torch.FloatTensor:
- for resnet in self.resnets:
- hidden_states = resnet(hidden_states, temb=temb, scale=scale)
-
- if self.upsamplers is not None:
- for upsampler in self.upsamplers:
- hidden_states = upsampler(hidden_states)
-
- return hidden_states
-
-
-class AttnUpDecoderBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- resolution_idx: Optional[int] = None,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- attention_head_dim: int = 1,
- output_scale_factor: float = 1.0,
- add_upsample: bool = True,
- temb_channels: Optional[int] = None,
- ):
- super().__init__()
- resnets = []
- attentions = []
-
- if attention_head_dim is None:
- logger.warn(
- f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `out_channels`: {out_channels}."
- )
- attention_head_dim = out_channels
-
- for i in range(num_layers):
- input_channels = in_channels if i == 0 else out_channels
-
- resnets.append(
- ResnetBlock2D(
- in_channels=input_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
- attentions.append(
- Attention(
- out_channels,
- heads=out_channels // attention_head_dim,
- dim_head=attention_head_dim,
- rescale_output_factor=output_scale_factor,
- eps=resnet_eps,
- norm_num_groups=resnet_groups if resnet_time_scale_shift != "spatial" else None,
- spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None,
- residual_connection=True,
- bias=True,
- upcast_softmax=True,
- _from_deprecated_attn_block=True,
- )
- )
-
- self.attentions = nn.ModuleList(attentions)
- self.resnets = nn.ModuleList(resnets)
-
- if add_upsample:
- self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
- else:
- self.upsamplers = None
-
- self.resolution_idx = resolution_idx
-
- def forward(
- self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
- ) -> torch.FloatTensor:
- for resnet, attn in zip(self.resnets, self.attentions):
- hidden_states = resnet(hidden_states, temb=temb, scale=scale)
- cross_attention_kwargs = {"scale": scale}
- hidden_states = attn(hidden_states, temb=temb, **cross_attention_kwargs)
-
- if self.upsamplers is not None:
- for upsampler in self.upsamplers:
- hidden_states = upsampler(hidden_states, scale=scale)
-
- return hidden_states
-
-
-class AttnSkipUpBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- prev_output_channel: int,
- out_channels: int,
- temb_channels: int,
- resolution_idx: Optional[int] = None,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_pre_norm: bool = True,
- attention_head_dim: int = 1,
- output_scale_factor: float = np.sqrt(2.0),
- add_upsample: bool = True,
- ):
- super().__init__()
- self.attentions = nn.ModuleList([])
- self.resnets = nn.ModuleList([])
-
- for i in range(num_layers):
- res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
- resnet_in_channels = prev_output_channel if i == 0 else out_channels
-
- self.resnets.append(
- ResnetBlock2D(
- in_channels=resnet_in_channels + res_skip_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=min(resnet_in_channels + res_skip_channels // 4, 32),
- groups_out=min(out_channels // 4, 32),
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
-
- if attention_head_dim is None:
- logger.warn(
- f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `out_channels`: {out_channels}."
- )
- attention_head_dim = out_channels
-
- self.attentions.append(
- Attention(
- out_channels,
- heads=out_channels // attention_head_dim,
- dim_head=attention_head_dim,
- rescale_output_factor=output_scale_factor,
- eps=resnet_eps,
- norm_num_groups=32,
- residual_connection=True,
- bias=True,
- upcast_softmax=True,
- _from_deprecated_attn_block=True,
- )
- )
-
- self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels)
- if add_upsample:
- self.resnet_up = ResnetBlock2D(
- in_channels=out_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=min(out_channels // 4, 32),
- groups_out=min(out_channels // 4, 32),
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- use_in_shortcut=True,
- up=True,
- kernel="fir",
- )
- self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
- self.skip_norm = torch.nn.GroupNorm(
- num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True
- )
- self.act = nn.SiLU()
- else:
- self.resnet_up = None
- self.skip_conv = None
- self.skip_norm = None
- self.act = None
-
- self.resolution_idx = resolution_idx
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
- temb: Optional[torch.FloatTensor] = None,
- skip_sample=None,
- scale: float = 1.0,
- ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
- for resnet in self.resnets:
- # pop res hidden states
- res_hidden_states = res_hidden_states_tuple[-1]
- res_hidden_states_tuple = res_hidden_states_tuple[:-1]
- hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
-
- hidden_states = resnet(hidden_states, temb, scale=scale)
-
- cross_attention_kwargs = {"scale": scale}
- hidden_states = self.attentions[0](hidden_states, **cross_attention_kwargs)
-
- if skip_sample is not None:
- skip_sample = self.upsampler(skip_sample)
- else:
- skip_sample = 0
-
- if self.resnet_up is not None:
- skip_sample_states = self.skip_norm(hidden_states)
- skip_sample_states = self.act(skip_sample_states)
- skip_sample_states = self.skip_conv(skip_sample_states)
-
- skip_sample = skip_sample + skip_sample_states
-
- hidden_states = self.resnet_up(hidden_states, temb, scale=scale)
-
- return hidden_states, skip_sample
-
-
-class SkipUpBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- prev_output_channel: int,
- out_channels: int,
- temb_channels: int,
- resolution_idx: Optional[int] = None,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_pre_norm: bool = True,
- output_scale_factor: float = np.sqrt(2.0),
- add_upsample: bool = True,
- upsample_padding: int = 1,
- ):
- super().__init__()
- self.resnets = nn.ModuleList([])
-
- for i in range(num_layers):
- res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
- resnet_in_channels = prev_output_channel if i == 0 else out_channels
-
- self.resnets.append(
- ResnetBlock2D(
- in_channels=resnet_in_channels + res_skip_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=min((resnet_in_channels + res_skip_channels) // 4, 32),
- groups_out=min(out_channels // 4, 32),
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- )
- )
-
- self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels)
- if add_upsample:
- self.resnet_up = ResnetBlock2D(
- in_channels=out_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=min(out_channels // 4, 32),
- groups_out=min(out_channels // 4, 32),
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- use_in_shortcut=True,
- up=True,
- kernel="fir",
- )
- self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
- self.skip_norm = torch.nn.GroupNorm(
- num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True
- )
- self.act = nn.SiLU()
- else:
- self.resnet_up = None
- self.skip_conv = None
- self.skip_norm = None
- self.act = None
-
- self.resolution_idx = resolution_idx
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
- temb: Optional[torch.FloatTensor] = None,
- skip_sample=None,
- scale: float = 1.0,
- ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
- for resnet in self.resnets:
- # pop res hidden states
- res_hidden_states = res_hidden_states_tuple[-1]
- res_hidden_states_tuple = res_hidden_states_tuple[:-1]
- hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
-
- hidden_states = resnet(hidden_states, temb, scale=scale)
-
- if skip_sample is not None:
- skip_sample = self.upsampler(skip_sample)
- else:
- skip_sample = 0
-
- if self.resnet_up is not None:
- skip_sample_states = self.skip_norm(hidden_states)
- skip_sample_states = self.act(skip_sample_states)
- skip_sample_states = self.skip_conv(skip_sample_states)
-
- skip_sample = skip_sample + skip_sample_states
-
- hidden_states = self.resnet_up(hidden_states, temb, scale=scale)
-
- return hidden_states, skip_sample
-
-
-class ResnetUpsampleBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- prev_output_channel: int,
- out_channels: int,
- temb_channels: int,
- resolution_idx: Optional[int] = None,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- output_scale_factor: float = 1.0,
- add_upsample: bool = True,
- skip_time_act: bool = False,
- ):
- super().__init__()
- resnets = []
-
- for i in range(num_layers):
- res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
- resnet_in_channels = prev_output_channel if i == 0 else out_channels
-
- resnets.append(
- ResnetBlock2D(
- in_channels=resnet_in_channels + res_skip_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- skip_time_act=skip_time_act,
- )
- )
-
- self.resnets = nn.ModuleList(resnets)
-
- if add_upsample:
- self.upsamplers = nn.ModuleList(
- [
- ResnetBlock2D(
- in_channels=out_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- skip_time_act=skip_time_act,
- up=True,
- )
- ]
- )
- else:
- self.upsamplers = None
-
- self.gradient_checkpointing = False
- self.resolution_idx = resolution_idx
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
- temb: Optional[torch.FloatTensor] = None,
- upsample_size: Optional[int] = None,
- scale: float = 1.0,
- ) -> torch.FloatTensor:
- for resnet in self.resnets:
- # pop res hidden states
- res_hidden_states = res_hidden_states_tuple[-1]
- res_hidden_states_tuple = res_hidden_states_tuple[:-1]
- hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
-
- if self.training and self.gradient_checkpointing:
-
- def create_custom_forward(module):
- def custom_forward(*inputs):
- return module(*inputs)
-
- return custom_forward
-
- if is_torch_version(">=", "1.11.0"):
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
- )
- else:
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet), hidden_states, temb
- )
- else:
- hidden_states = resnet(hidden_states, temb, scale=scale)
-
- if self.upsamplers is not None:
- for upsampler in self.upsamplers:
- hidden_states = upsampler(hidden_states, temb, scale=scale)
-
- return hidden_states
-
-
-class SimpleCrossAttnUpBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- prev_output_channel: int,
- temb_channels: int,
- resolution_idx: Optional[int] = None,
- dropout: float = 0.0,
- num_layers: int = 1,
- resnet_eps: float = 1e-6,
- resnet_time_scale_shift: str = "default",
- resnet_act_fn: str = "swish",
- resnet_groups: int = 32,
- resnet_pre_norm: bool = True,
- attention_head_dim: int = 1,
- cross_attention_dim: int = 1280,
- output_scale_factor: float = 1.0,
- add_upsample: bool = True,
- skip_time_act: bool = False,
- only_cross_attention: bool = False,
- cross_attention_norm: Optional[str] = None,
- ):
- super().__init__()
- resnets = []
- attentions = []
-
- self.has_cross_attention = True
- self.attention_head_dim = attention_head_dim
-
- self.num_heads = out_channels // self.attention_head_dim
-
- for i in range(num_layers):
- res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
- resnet_in_channels = prev_output_channel if i == 0 else out_channels
-
- resnets.append(
- ResnetBlock2D(
- in_channels=resnet_in_channels + res_skip_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- skip_time_act=skip_time_act,
- )
- )
-
- processor = (
- AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor()
- )
-
- attentions.append(
- Attention(
- query_dim=out_channels,
- cross_attention_dim=out_channels,
- heads=self.num_heads,
- dim_head=self.attention_head_dim,
- added_kv_proj_dim=cross_attention_dim,
- norm_num_groups=resnet_groups,
- bias=True,
- upcast_softmax=True,
- only_cross_attention=only_cross_attention,
- cross_attention_norm=cross_attention_norm,
- processor=processor,
- )
- )
- self.attentions = nn.ModuleList(attentions)
- self.resnets = nn.ModuleList(resnets)
-
- if add_upsample:
- self.upsamplers = nn.ModuleList(
- [
- ResnetBlock2D(
- in_channels=out_channels,
- out_channels=out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=resnet_groups,
- dropout=dropout,
- time_embedding_norm=resnet_time_scale_shift,
- non_linearity=resnet_act_fn,
- output_scale_factor=output_scale_factor,
- pre_norm=resnet_pre_norm,
- skip_time_act=skip_time_act,
- up=True,
- )
- ]
- )
- else:
- self.upsamplers = None
-
- self.gradient_checkpointing = False
- self.resolution_idx = resolution_idx
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
- temb: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- upsample_size: Optional[int] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- ) -> torch.FloatTensor:
- cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
-
- lora_scale = cross_attention_kwargs.get("scale", 1.0)
- if attention_mask is None:
- # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask.
- mask = None if encoder_hidden_states is None else encoder_attention_mask
- else:
- # when attention_mask is defined: we don't even check for encoder_attention_mask.
- # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks.
- # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask.
- # then we can simplify this whole if/else block to:
- # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask
- mask = attention_mask
-
- for resnet, attn in zip(self.resnets, self.attentions):
- # resnet
- # pop res hidden states
- res_hidden_states = res_hidden_states_tuple[-1]
- res_hidden_states_tuple = res_hidden_states_tuple[:-1]
- hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
-
- if self.training and self.gradient_checkpointing:
-
- def create_custom_forward(module, return_dict=None):
- def custom_forward(*inputs):
- if return_dict is not None:
- return module(*inputs, return_dict=return_dict)
- else:
- return module(*inputs)
-
- return custom_forward
-
- hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
- hidden_states = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- attention_mask=mask,
- **cross_attention_kwargs,
- )
- else:
- hidden_states = resnet(hidden_states, temb, scale=lora_scale)
-
- hidden_states = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- attention_mask=mask,
- **cross_attention_kwargs,
- )
-
- if self.upsamplers is not None:
- for upsampler in self.upsamplers:
- hidden_states = upsampler(hidden_states, temb, scale=lora_scale)
-
- return hidden_states
-
-
-class KUpBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- temb_channels: int,
- resolution_idx: int,
- dropout: float = 0.0,
- num_layers: int = 5,
- resnet_eps: float = 1e-5,
- resnet_act_fn: str = "gelu",
- resnet_group_size: Optional[int] = 32,
- add_upsample: bool = True,
- ):
- super().__init__()
- resnets = []
- k_in_channels = 2 * out_channels
- k_out_channels = in_channels
- num_layers = num_layers - 1
-
- for i in range(num_layers):
- in_channels = k_in_channels if i == 0 else out_channels
- groups = in_channels // resnet_group_size
- groups_out = out_channels // resnet_group_size
-
- resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=k_out_channels if (i == num_layers - 1) else out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=groups,
- groups_out=groups_out,
- dropout=dropout,
- non_linearity=resnet_act_fn,
- time_embedding_norm="ada_group",
- conv_shortcut_bias=False,
- )
- )
-
- self.resnets = nn.ModuleList(resnets)
-
- if add_upsample:
- self.upsamplers = nn.ModuleList([KUpsample2D()])
- else:
- self.upsamplers = None
-
- self.gradient_checkpointing = False
- self.resolution_idx = resolution_idx
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
- temb: Optional[torch.FloatTensor] = None,
- upsample_size: Optional[int] = None,
- scale: float = 1.0,
- ) -> torch.FloatTensor:
- res_hidden_states_tuple = res_hidden_states_tuple[-1]
- if res_hidden_states_tuple is not None:
- hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], dim=1)
-
- for resnet in self.resnets:
- if self.training and self.gradient_checkpointing:
-
- def create_custom_forward(module):
- def custom_forward(*inputs):
- return module(*inputs)
-
- return custom_forward
-
- if is_torch_version(">=", "1.11.0"):
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
- )
- else:
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet), hidden_states, temb
- )
- else:
- hidden_states = resnet(hidden_states, temb, scale=scale)
-
- if self.upsamplers is not None:
- for upsampler in self.upsamplers:
- hidden_states = upsampler(hidden_states)
-
- return hidden_states
-
-
-class KCrossAttnUpBlock2D(nn.Module):
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- temb_channels: int,
- resolution_idx: int,
- dropout: float = 0.0,
- num_layers: int = 4,
- resnet_eps: float = 1e-5,
- resnet_act_fn: str = "gelu",
- resnet_group_size: int = 32,
- attention_head_dim: int = 1, # attention dim_head
- cross_attention_dim: int = 768,
- add_upsample: bool = True,
- upcast_attention: bool = False,
- ):
- super().__init__()
- resnets = []
- attentions = []
-
- is_first_block = in_channels == out_channels == temb_channels
- is_middle_block = in_channels != out_channels
- add_self_attention = True if is_first_block else False
-
- self.has_cross_attention = True
- self.attention_head_dim = attention_head_dim
-
- # in_channels, and out_channels for the block (k-unet)
- k_in_channels = out_channels if is_first_block else 2 * out_channels
- k_out_channels = in_channels
-
- num_layers = num_layers - 1
-
- for i in range(num_layers):
- in_channels = k_in_channels if i == 0 else out_channels
- groups = in_channels // resnet_group_size
- groups_out = out_channels // resnet_group_size
-
- if is_middle_block and (i == num_layers - 1):
- conv_2d_out_channels = k_out_channels
- else:
- conv_2d_out_channels = None
-
- resnets.append(
- ResnetBlock2D(
- in_channels=in_channels,
- out_channels=out_channels,
- conv_2d_out_channels=conv_2d_out_channels,
- temb_channels=temb_channels,
- eps=resnet_eps,
- groups=groups,
- groups_out=groups_out,
- dropout=dropout,
- non_linearity=resnet_act_fn,
- time_embedding_norm="ada_group",
- conv_shortcut_bias=False,
- )
- )
- attentions.append(
- KAttentionBlock(
- k_out_channels if (i == num_layers - 1) else out_channels,
- k_out_channels // attention_head_dim
- if (i == num_layers - 1)
- else out_channels // attention_head_dim,
- attention_head_dim,
- cross_attention_dim=cross_attention_dim,
- temb_channels=temb_channels,
- attention_bias=True,
- add_self_attention=add_self_attention,
- cross_attention_norm="layer_norm",
- upcast_attention=upcast_attention,
- )
- )
-
- self.resnets = nn.ModuleList(resnets)
- self.attentions = nn.ModuleList(attentions)
-
- if add_upsample:
- self.upsamplers = nn.ModuleList([KUpsample2D()])
- else:
- self.upsamplers = None
-
- self.gradient_checkpointing = False
- self.resolution_idx = resolution_idx
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
- temb: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- upsample_size: Optional[int] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- ) -> torch.FloatTensor:
- res_hidden_states_tuple = res_hidden_states_tuple[-1]
- if res_hidden_states_tuple is not None:
- hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], dim=1)
-
- lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
- for resnet, attn in zip(self.resnets, self.attentions):
- if self.training and self.gradient_checkpointing:
-
- def create_custom_forward(module, return_dict=None):
- def custom_forward(*inputs):
- if return_dict is not None:
- return module(*inputs, return_dict=return_dict)
- else:
- return module(*inputs)
-
- return custom_forward
-
- ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(resnet),
- hidden_states,
- temb,
- **ckpt_kwargs,
- )
- hidden_states = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- emb=temb,
- attention_mask=attention_mask,
- cross_attention_kwargs=cross_attention_kwargs,
- encoder_attention_mask=encoder_attention_mask,
- )
- else:
- hidden_states = resnet(hidden_states, temb, scale=lora_scale)
- hidden_states = attn(
- hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- emb=temb,
- attention_mask=attention_mask,
- cross_attention_kwargs=cross_attention_kwargs,
- encoder_attention_mask=encoder_attention_mask,
- )
-
- if self.upsamplers is not None:
- for upsampler in self.upsamplers:
- hidden_states = upsampler(hidden_states)
-
- return hidden_states
-
-
-# can potentially later be renamed to `No-feed-forward` attention
-class KAttentionBlock(nn.Module):
- r"""
- A basic Transformer block.
-
- Parameters:
- dim (`int`): The number of channels in the input and output.
- num_attention_heads (`int`): The number of heads to use for multi-head attention.
- attention_head_dim (`int`): The number of channels in each head.
- dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
- cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
- attention_bias (`bool`, *optional*, defaults to `False`):
- Configure if the attention layers should contain a bias parameter.
- upcast_attention (`bool`, *optional*, defaults to `False`):
- Set to `True` to upcast the attention computation to `float32`.
- temb_channels (`int`, *optional*, defaults to 768):
- The number of channels in the token embedding.
- add_self_attention (`bool`, *optional*, defaults to `False`):
- Set to `True` to add self-attention to the block.
- cross_attention_norm (`str`, *optional*, defaults to `None`):
- The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`.
- group_size (`int`, *optional*, defaults to 32):
- The number of groups to separate the channels into for group normalization.
- """
-
- def __init__(
- self,
- dim: int,
- num_attention_heads: int,
- attention_head_dim: int,
- dropout: float = 0.0,
- cross_attention_dim: Optional[int] = None,
- attention_bias: bool = False,
- upcast_attention: bool = False,
- temb_channels: int = 768, # for ada_group_norm
- add_self_attention: bool = False,
- cross_attention_norm: Optional[str] = None,
- group_size: int = 32,
- ):
- super().__init__()
- self.add_self_attention = add_self_attention
-
- # 1. Self-Attn
- if add_self_attention:
- self.norm1 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size))
- self.attn1 = Attention(
- query_dim=dim,
- heads=num_attention_heads,
- dim_head=attention_head_dim,
- dropout=dropout,
- bias=attention_bias,
- cross_attention_dim=None,
- cross_attention_norm=None,
- )
-
- # 2. Cross-Attn
- self.norm2 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size))
- self.attn2 = Attention(
- query_dim=dim,
- cross_attention_dim=cross_attention_dim,
- heads=num_attention_heads,
- dim_head=attention_head_dim,
- dropout=dropout,
- bias=attention_bias,
- upcast_attention=upcast_attention,
- cross_attention_norm=cross_attention_norm,
- )
-
- def _to_3d(self, hidden_states: torch.FloatTensor, height: int, weight: int) -> torch.FloatTensor:
- return hidden_states.permute(0, 2, 3, 1).reshape(hidden_states.shape[0], height * weight, -1)
-
- def _to_4d(self, hidden_states: torch.FloatTensor, height: int, weight: int) -> torch.FloatTensor:
- return hidden_states.permute(0, 2, 1).reshape(hidden_states.shape[0], -1, height, weight)
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- # TODO: mark emb as non-optional (self.norm2 requires it).
- # requires assessing impact of change to positional param interface.
- emb: Optional[torch.FloatTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- ) -> torch.FloatTensor:
- cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
-
- # 1. Self-Attention
- if self.add_self_attention:
- norm_hidden_states = self.norm1(hidden_states, emb)
-
- height, weight = norm_hidden_states.shape[2:]
- norm_hidden_states = self._to_3d(norm_hidden_states, height, weight)
-
- attn_output = self.attn1(
- norm_hidden_states,
- encoder_hidden_states=None,
- attention_mask=attention_mask,
- **cross_attention_kwargs,
- )
- attn_output = self._to_4d(attn_output, height, weight)
-
- hidden_states = attn_output + hidden_states
-
- # 2. Cross-Attention/None
- norm_hidden_states = self.norm2(hidden_states, emb)
-
- height, weight = norm_hidden_states.shape[2:]
- norm_hidden_states = self._to_3d(norm_hidden_states, height, weight)
- attn_output = self.attn2(
- norm_hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- attention_mask=attention_mask if encoder_hidden_states is None else encoder_attention_mask,
- **cross_attention_kwargs,
- )
- attn_output = self._to_4d(attn_output, height, weight)
-
- hidden_states = attn_output + hidden_states
-
- return hidden_states