declare-lab/tango2
92
1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14from dataclasses import dataclass15from typing import Optional, Tuple, Union16 17import torch18import torch.nn as nn19 20from ..configuration_utils import ConfigMixin, register_to_config21from ..utils import BaseOutput22from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps23from .modeling_utils import ModelMixin24from .unet_2d_blocks import UNetMidBlock2D, get_down_block, get_up_block25 26 27@dataclass28class UNet2DOutput(BaseOutput):29 """30 Args:31 sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):32 Hidden states output. Output of last layer of model.33 """34 35 sample: torch.FloatTensor36 37 38class UNet2DModel(ModelMixin, ConfigMixin):39 r"""40 UNet2DModel is a 2D UNet model that takes in a noisy sample and a timestep and returns sample shaped output.41 42 This model inherits from [`ModelMixin`]. Check the superclass documentation for the generic methods the library43 implements for all the model (such as downloading or saving, etc.)44 45 Parameters:46 sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):47 Height and width of input/output sample.48 in_channels (`int`, *optional*, defaults to 3): Number of channels in the input image.49 out_channels (`int`, *optional*, defaults to 3): Number of channels in the output.50 center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.51 time_embedding_type (`str`, *optional*, defaults to `"positional"`): Type of time embedding to use.52 freq_shift (`int`, *optional*, defaults to 0): Frequency shift for fourier time embedding.53 flip_sin_to_cos (`bool`, *optional*, defaults to :54 obj:`True`): Whether to flip sin to cos for fourier time embedding.55 down_block_types (`Tuple[str]`, *optional*, defaults to :56 obj:`("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")`): Tuple of downsample block57 types.58 mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2D"`):59 The mid block type. Choose from `UNetMidBlock2D` or `UnCLIPUNetMidBlock2D`.60 up_block_types (`Tuple[str]`, *optional*, defaults to :61 obj:`("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")`): Tuple of upsample block types.62 block_out_channels (`Tuple[int]`, *optional*, defaults to :63 obj:`(224, 448, 672, 896)`): Tuple of block output channels.64 layers_per_block (`int`, *optional*, defaults to `2`): The number of layers per block.65 mid_block_scale_factor (`float`, *optional*, defaults to `1`): The scale factor for the mid block.66 downsample_padding (`int`, *optional*, defaults to `1`): The padding for the downsample convolution.67 act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.68 attention_head_dim (`int`, *optional*, defaults to `8`): The attention head dimension.69 norm_num_groups (`int`, *optional*, defaults to `32`): The number of groups for the normalization.70 norm_eps (`float`, *optional*, defaults to `1e-5`): The epsilon for the normalization.71 resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config72 for resnet blocks, see [`~models.resnet.ResnetBlock2D`]. Choose from `default` or `scale_shift`.73 class_embed_type (`str`, *optional*, defaults to None):74 The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,75 `"timestep"`, or `"identity"`.76 num_class_embeds (`int`, *optional*, defaults to None):77 Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing78 class conditioning with `class_embed_type` equal to `None`.79 """80 81 @register_to_config82 def __init__(83 self,84 sample_size: Optional[Union[int, Tuple[int, int]]] = None,85 in_channels: int = 3,86 out_channels: int = 3,87 center_input_sample: bool = False,88 time_embedding_type: str = "positional",89 freq_shift: int = 0,90 flip_sin_to_cos: bool = True,91 down_block_types: Tuple[str] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D"),92 up_block_types: Tuple[str] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D"),93 block_out_channels: Tuple[int] = (224, 448, 672, 896),94 layers_per_block: int = 2,95 mid_block_scale_factor: float = 1,96 downsample_padding: int = 1,97 act_fn: str = "silu",98 attention_head_dim: Optional[int] = 8,99 norm_num_groups: int = 32,100 norm_eps: float = 1e-5,101 resnet_time_scale_shift: str = "default",102 add_attention: bool = True,103 class_embed_type: Optional[str] = None,104 num_class_embeds: Optional[int] = None,105 ):106 super().__init__()107 108 self.sample_size = sample_size109 time_embed_dim = block_out_channels[0] * 4110 111 # Check inputs112 if len(down_block_types) != len(up_block_types):113 raise ValueError(114 f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."115 )116 117 if len(block_out_channels) != len(down_block_types):118 raise ValueError(119 f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."120 )121 122 # input123 self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1))124 125 # time126 if time_embedding_type == "fourier":127 self.time_proj = GaussianFourierProjection(embedding_size=block_out_channels[0], scale=16)128 timestep_input_dim = 2 * block_out_channels[0]129 elif time_embedding_type == "positional":130 self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)131 timestep_input_dim = block_out_channels[0]132 133 self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)134 135 # class embedding136 if class_embed_type is None and num_class_embeds is not None:137 self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)138 elif class_embed_type == "timestep":139 self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)140 elif class_embed_type == "identity":141 self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)142 else:143 self.class_embedding = None144 145 self.down_blocks = nn.ModuleList([])146 self.mid_block = None147 self.up_blocks = nn.ModuleList([])148 149 # down150 output_channel = block_out_channels[0]151 for i, down_block_type in enumerate(down_block_types):152 input_channel = output_channel153 output_channel = block_out_channels[i]154 is_final_block = i == len(block_out_channels) - 1155 156 down_block = get_down_block(157 down_block_type,158 num_layers=layers_per_block,159 in_channels=input_channel,160 out_channels=output_channel,161 temb_channels=time_embed_dim,162 add_downsample=not is_final_block,163 resnet_eps=norm_eps,164 resnet_act_fn=act_fn,165 resnet_groups=norm_num_groups,166 attn_num_head_channels=attention_head_dim,167 downsample_padding=downsample_padding,168 resnet_time_scale_shift=resnet_time_scale_shift,169 )170 self.down_blocks.append(down_block)171 172 # mid173 self.mid_block = UNetMidBlock2D(174 in_channels=block_out_channels[-1],175 temb_channels=time_embed_dim,176 resnet_eps=norm_eps,177 resnet_act_fn=act_fn,178 output_scale_factor=mid_block_scale_factor,179 resnet_time_scale_shift=resnet_time_scale_shift,180 attn_num_head_channels=attention_head_dim,181 resnet_groups=norm_num_groups,182 add_attention=add_attention,183 )184 185 # up186 reversed_block_out_channels = list(reversed(block_out_channels))187 output_channel = reversed_block_out_channels[0]188 for i, up_block_type in enumerate(up_block_types):189 prev_output_channel = output_channel190 output_channel = reversed_block_out_channels[i]191 input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]192 193 is_final_block = i == len(block_out_channels) - 1194 195 up_block = get_up_block(196 up_block_type,197 num_layers=layers_per_block + 1,198 in_channels=input_channel,199 out_channels=output_channel,200 prev_output_channel=prev_output_channel,201 temb_channels=time_embed_dim,202 add_upsample=not is_final_block,203 resnet_eps=norm_eps,204 resnet_act_fn=act_fn,205 resnet_groups=norm_num_groups,206 attn_num_head_channels=attention_head_dim,207 resnet_time_scale_shift=resnet_time_scale_shift,208 )209 self.up_blocks.append(up_block)210 prev_output_channel = output_channel211 212 # out213 num_groups_out = norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4, 32)214 self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=num_groups_out, eps=norm_eps)215 self.conv_act = nn.SiLU()216 self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=3, padding=1)217 218 def forward(219 self,220 sample: torch.FloatTensor,221 timestep: Union[torch.Tensor, float, int],222 class_labels: Optional[torch.Tensor] = None,223 return_dict: bool = True,224 ) -> Union[UNet2DOutput, Tuple]:225 r"""226 Args:227 sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor228 timestep (`torch.FloatTensor` or `float` or `int): (batch) timesteps229 class_labels (`torch.FloatTensor`, *optional*, defaults to `None`):230 Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.231 return_dict (`bool`, *optional*, defaults to `True`):232 Whether or not to return a [`~models.unet_2d.UNet2DOutput`] instead of a plain tuple.233 234 Returns:235 [`~models.unet_2d.UNet2DOutput`] or `tuple`: [`~models.unet_2d.UNet2DOutput`] if `return_dict` is True,236 otherwise a `tuple`. When returning a tuple, the first element is the sample tensor.237 """238 # 0. center input if necessary239 if self.config.center_input_sample:240 sample = 2 * sample - 1.0241 242 # 1. time243 timesteps = timestep244 if not torch.is_tensor(timesteps):245 timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device)246 elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0:247 timesteps = timesteps[None].to(sample.device)248 249 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML250 timesteps = timesteps * torch.ones(sample.shape[0], dtype=timesteps.dtype, device=timesteps.device)251 252 t_emb = self.time_proj(timesteps)253 254 # timesteps does not contain any weights and will always return f32 tensors255 # but time_embedding might actually be running in fp16. so we need to cast here.256 # there might be better ways to encapsulate this.257 t_emb = t_emb.to(dtype=self.dtype)258 emb = self.time_embedding(t_emb)259 260 if self.class_embedding is not None:261 if class_labels is None:262 raise ValueError("class_labels should be provided when doing class conditioning")263 264 if self.config.class_embed_type == "timestep":265 class_labels = self.time_proj(class_labels)266 267 class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)268 emb = emb + class_emb269 270 # 2. pre-process271 skip_sample = sample272 sample = self.conv_in(sample)273 274 # 3. down275 down_block_res_samples = (sample,)276 for downsample_block in self.down_blocks:277 if hasattr(downsample_block, "skip_conv"):278 sample, res_samples, skip_sample = downsample_block(279 hidden_states=sample, temb=emb, skip_sample=skip_sample280 )281 else:282 sample, res_samples = downsample_block(hidden_states=sample, temb=emb)283 284 down_block_res_samples += res_samples285 286 # 4. mid287 sample = self.mid_block(sample, emb)288 289 # 5. up290 skip_sample = None291 for upsample_block in self.up_blocks:292 res_samples = down_block_res_samples[-len(upsample_block.resnets) :]293 down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]294 295 if hasattr(upsample_block, "skip_conv"):296 sample, skip_sample = upsample_block(sample, res_samples, emb, skip_sample)297 else:298 sample = upsample_block(sample, res_samples, emb)299 300 # 6. post-process301 sample = self.conv_norm_out(sample)302 sample = self.conv_act(sample)303 sample = self.conv_out(sample)304 305 if skip_sample is not None:306 sample += skip_sample307 308 if self.config.time_embedding_type == "fourier":309 timesteps = timesteps.reshape((sample.shape[0], *([1] * len(sample.shape[1:]))))310 sample = sample / timesteps311 312 if not return_dict:313 return (sample,)314 315 return UNet2DOutput(sample=sample)316