declare-lab/tango2

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.14 15from dataclasses import dataclass16from typing import Optional, Tuple, Union17 18import torch19import torch.nn as nn20 21from ..configuration_utils import ConfigMixin, register_to_config22from ..utils import BaseOutput23from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps24from .modeling_utils import ModelMixin25from .unet_1d_blocks import get_down_block, get_mid_block, get_out_block, get_up_block26 27 28@dataclass29class UNet1DOutput(BaseOutput):30    """31    Args:32        sample (`torch.FloatTensor` of shape `(batch_size, num_channels, sample_size)`):33            Hidden states output. Output of last layer of model.34    """35 36    sample: torch.FloatTensor37 38 39class UNet1DModel(ModelMixin, ConfigMixin):40    r"""41    UNet1DModel is a 1D UNet model that takes in a noisy sample and a timestep and returns sample shaped output.42 43    This model inherits from [`ModelMixin`]. Check the superclass documentation for the generic methods the library44    implements for all the model (such as downloading or saving, etc.)45 46    Parameters:47        sample_size (`int`, *optional*): Default length of sample. Should be adaptable at runtime.48        in_channels (`int`, *optional*, defaults to 2): Number of channels in the input sample.49        out_channels (`int`, *optional*, defaults to 2): Number of channels in the output.50        extra_in_channels (`int`, *optional*, defaults to 0):51            Number of additional channels to be added to the input of the first down block. Useful for cases where the52            input data has more channels than what the model is initially designed for.53        time_embedding_type (`str`, *optional*, defaults to `"fourier"`): Type of time embedding to use.54        freq_shift (`float`, *optional*, defaults to 0.0): Frequency shift for fourier time embedding.55        flip_sin_to_cos (`bool`, *optional*, defaults to :56            obj:`False`): Whether to flip sin to cos for fourier time embedding.57        down_block_types (`Tuple[str]`, *optional*, defaults to :58            obj:`("DownBlock1D", "DownBlock1DNoSkip", "AttnDownBlock1D")`): Tuple of downsample block types.59        up_block_types (`Tuple[str]`, *optional*, defaults to :60            obj:`("UpBlock1D", "UpBlock1DNoSkip", "AttnUpBlock1D")`): Tuple of upsample block types.61        block_out_channels (`Tuple[int]`, *optional*, defaults to :62            obj:`(32, 32, 64)`): Tuple of block output channels.63        mid_block_type (`str`, *optional*, defaults to "UNetMidBlock1D"): block type for middle of UNet.64        out_block_type (`str`, *optional*, defaults to `None`): optional output processing of UNet.65        act_fn (`str`, *optional*, defaults to None): optional activation function in UNet blocks.66        norm_num_groups (`int`, *optional*, defaults to 8): group norm member count in UNet blocks.67        layers_per_block (`int`, *optional*, defaults to 1): added number of layers in a UNet block.68        downsample_each_block (`int`, *optional*, defaults to False:69            experimental feature for using a UNet without upsampling.70    """71 72    @register_to_config73    def __init__(74        self,75        sample_size: int = 65536,76        sample_rate: Optional[int] = None,77        in_channels: int = 2,78        out_channels: int = 2,79        extra_in_channels: int = 0,80        time_embedding_type: str = "fourier",81        flip_sin_to_cos: bool = True,82        use_timestep_embedding: bool = False,83        freq_shift: float = 0.0,84        down_block_types: Tuple[str] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D"),85        up_block_types: Tuple[str] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip"),86        mid_block_type: Tuple[str] = "UNetMidBlock1D",87        out_block_type: str = None,88        block_out_channels: Tuple[int] = (32, 32, 64),89        act_fn: str = None,90        norm_num_groups: int = 8,91        layers_per_block: int = 1,92        downsample_each_block: bool = False,93    ):94        super().__init__()95        self.sample_size = sample_size96 97        # time98        if time_embedding_type == "fourier":99            self.time_proj = GaussianFourierProjection(100                embedding_size=8, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos101            )102            timestep_input_dim = 2 * block_out_channels[0]103        elif time_embedding_type == "positional":104            self.time_proj = Timesteps(105                block_out_channels[0], flip_sin_to_cos=flip_sin_to_cos, downscale_freq_shift=freq_shift106            )107            timestep_input_dim = block_out_channels[0]108 109        if use_timestep_embedding:110            time_embed_dim = block_out_channels[0] * 4111            self.time_mlp = TimestepEmbedding(112                in_channels=timestep_input_dim,113                time_embed_dim=time_embed_dim,114                act_fn=act_fn,115                out_dim=block_out_channels[0],116            )117 118        self.down_blocks = nn.ModuleList([])119        self.mid_block = None120        self.up_blocks = nn.ModuleList([])121        self.out_block = None122 123        # down124        output_channel = in_channels125        for i, down_block_type in enumerate(down_block_types):126            input_channel = output_channel127            output_channel = block_out_channels[i]128 129            if i == 0:130                input_channel += extra_in_channels131 132            is_final_block = i == len(block_out_channels) - 1133 134            down_block = get_down_block(135                down_block_type,136                num_layers=layers_per_block,137                in_channels=input_channel,138                out_channels=output_channel,139                temb_channels=block_out_channels[0],140                add_downsample=not is_final_block or downsample_each_block,141            )142            self.down_blocks.append(down_block)143 144        # mid145        self.mid_block = get_mid_block(146            mid_block_type,147            in_channels=block_out_channels[-1],148            mid_channels=block_out_channels[-1],149            out_channels=block_out_channels[-1],150            embed_dim=block_out_channels[0],151            num_layers=layers_per_block,152            add_downsample=downsample_each_block,153        )154 155        # up156        reversed_block_out_channels = list(reversed(block_out_channels))157        output_channel = reversed_block_out_channels[0]158        if out_block_type is None:159            final_upsample_channels = out_channels160        else:161            final_upsample_channels = block_out_channels[0]162 163        for i, up_block_type in enumerate(up_block_types):164            prev_output_channel = output_channel165            output_channel = (166                reversed_block_out_channels[i + 1] if i < len(up_block_types) - 1 else final_upsample_channels167            )168 169            is_final_block = i == len(block_out_channels) - 1170 171            up_block = get_up_block(172                up_block_type,173                num_layers=layers_per_block,174                in_channels=prev_output_channel,175                out_channels=output_channel,176                temb_channels=block_out_channels[0],177                add_upsample=not is_final_block,178            )179            self.up_blocks.append(up_block)180            prev_output_channel = output_channel181 182        # out183        num_groups_out = norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4, 32)184        self.out_block = get_out_block(185            out_block_type=out_block_type,186            num_groups_out=num_groups_out,187            embed_dim=block_out_channels[0],188            out_channels=out_channels,189            act_fn=act_fn,190            fc_dim=block_out_channels[-1] // 4,191        )192 193    def forward(194        self,195        sample: torch.FloatTensor,196        timestep: Union[torch.Tensor, float, int],197        return_dict: bool = True,198    ) -> Union[UNet1DOutput, Tuple]:199        r"""200        Args:201            sample (`torch.FloatTensor`): `(batch_size, num_channels, sample_size)` noisy inputs tensor202            timestep (`torch.FloatTensor` or `float` or `int): (batch) timesteps203            return_dict (`bool`, *optional*, defaults to `True`):204                Whether or not to return a [`~models.unet_1d.UNet1DOutput`] instead of a plain tuple.205 206        Returns:207            [`~models.unet_1d.UNet1DOutput`] or `tuple`: [`~models.unet_1d.UNet1DOutput`] if `return_dict` is True,208            otherwise a `tuple`. When returning a tuple, the first element is the sample tensor.209        """210 211        # 1. time212        timesteps = timestep213        if not torch.is_tensor(timesteps):214            timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device)215        elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0:216            timesteps = timesteps[None].to(sample.device)217 218        timestep_embed = self.time_proj(timesteps)219        if self.config.use_timestep_embedding:220            timestep_embed = self.time_mlp(timestep_embed)221        else:222            timestep_embed = timestep_embed[..., None]223            timestep_embed = timestep_embed.repeat([1, 1, sample.shape[2]]).to(sample.dtype)224            timestep_embed = timestep_embed.broadcast_to((sample.shape[:1] + timestep_embed.shape[1:]))225 226        # 2. down227        down_block_res_samples = ()228        for downsample_block in self.down_blocks:229            sample, res_samples = downsample_block(hidden_states=sample, temb=timestep_embed)230            down_block_res_samples += res_samples231 232        # 3. mid233        if self.mid_block:234            sample = self.mid_block(sample, timestep_embed)235 236        # 4. up237        for i, upsample_block in enumerate(self.up_blocks):238            res_samples = down_block_res_samples[-1:]239            down_block_res_samples = down_block_res_samples[:-1]240            sample = upsample_block(sample, res_hidden_states_tuple=res_samples, temb=timestep_embed)241 242        # 5. post-process243        if self.out_block:244            sample = self.out_block(sample, timestep_embed)245 246        if not return_dict:247            return (sample,)248 249        return UNet1DOutput(sample=sample)250