declare-lab/tango2

1# Copyright 2023 Alibaba DAMO-VILAB and The HuggingFace Team. All rights reserved.2# Copyright 2023 The ModelScope Team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15from dataclasses import dataclass16from typing import Any, Dict, List, Optional, Tuple, Union17 18import torch19import torch.nn as nn20import torch.utils.checkpoint21 22from ..configuration_utils import ConfigMixin, register_to_config23from ..utils import BaseOutput, logging24from .attention_processor import AttentionProcessor, AttnProcessor25from .embeddings import TimestepEmbedding, Timesteps26from .modeling_utils import ModelMixin27from .transformer_temporal import TransformerTemporalModel28from .unet_3d_blocks import (29    CrossAttnDownBlock3D,30    CrossAttnUpBlock3D,31    DownBlock3D,32    UNetMidBlock3DCrossAttn,33    UpBlock3D,34    get_down_block,35    get_up_block,36)37 38 39logger = logging.get_logger(__name__)  # pylint: disable=invalid-name40 41 42@dataclass43class UNet3DConditionOutput(BaseOutput):44    """45    Args:46        sample (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`):47            Hidden states conditioned on `encoder_hidden_states` input. Output of last layer of model.48    """49 50    sample: torch.FloatTensor51 52 53class UNet3DConditionModel(ModelMixin, ConfigMixin):54    r"""55    UNet3DConditionModel is a conditional 2D UNet model that takes in a noisy sample, conditional state, and a timestep56    and returns sample shaped output.57 58    This model inherits from [`ModelMixin`]. Check the superclass documentation for the generic methods the library59    implements for all the models (such as downloading or saving, etc.)60 61    Parameters:62        sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):63            Height and width of input/output sample.64        in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample.65        out_channels (`int`, *optional*, defaults to 4): The number of channels in the output.66        down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):67            The tuple of downsample blocks to use.68        up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D",)`):69            The tuple of upsample blocks to use.70        block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):71            The tuple of output channels for each block.72        layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.73        downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.74        mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block.75        act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.76        norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.77            If `None`, it will skip the normalization and activation layers in post-processing78        norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.79        cross_attention_dim (`int`, *optional*, defaults to 1280): The dimension of the cross attention features.80        attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.81    """82 83    _supports_gradient_checkpointing = False84 85    @register_to_config86    def __init__(87        self,88        sample_size: Optional[int] = None,89        in_channels: int = 4,90        out_channels: int = 4,91        down_block_types: Tuple[str] = (92            "CrossAttnDownBlock3D",93            "CrossAttnDownBlock3D",94            "CrossAttnDownBlock3D",95            "DownBlock3D",96        ),97        up_block_types: Tuple[str] = ("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D"),98        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),99        layers_per_block: int = 2,100        downsample_padding: int = 1,101        mid_block_scale_factor: float = 1,102        act_fn: str = "silu",103        norm_num_groups: Optional[int] = 32,104        norm_eps: float = 1e-5,105        cross_attention_dim: int = 1024,106        attention_head_dim: Union[int, Tuple[int]] = 64,107    ):108        super().__init__()109 110        self.sample_size = sample_size111 112        # Check inputs113        if len(down_block_types) != len(up_block_types):114            raise ValueError(115                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}."116            )117 118        if len(block_out_channels) != len(down_block_types):119            raise ValueError(120                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}."121            )122 123        if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):124            raise ValueError(125                f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}."126            )127 128        # input129        conv_in_kernel = 3130        conv_out_kernel = 3131        conv_in_padding = (conv_in_kernel - 1) // 2132        self.conv_in = nn.Conv2d(133            in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding134        )135 136        # time137        time_embed_dim = block_out_channels[0] * 4138        self.time_proj = Timesteps(block_out_channels[0], True, 0)139        timestep_input_dim = block_out_channels[0]140 141        self.time_embedding = TimestepEmbedding(142            timestep_input_dim,143            time_embed_dim,144            act_fn=act_fn,145        )146 147        self.transformer_in = TransformerTemporalModel(148            num_attention_heads=8,149            attention_head_dim=attention_head_dim,150            in_channels=block_out_channels[0],151            num_layers=1,152        )153 154        # class embedding155        self.down_blocks = nn.ModuleList([])156        self.up_blocks = nn.ModuleList([])157 158        if isinstance(attention_head_dim, int):159            attention_head_dim = (attention_head_dim,) * len(down_block_types)160 161        # down162        output_channel = block_out_channels[0]163        for i, down_block_type in enumerate(down_block_types):164            input_channel = output_channel165            output_channel = block_out_channels[i]166            is_final_block = i == len(block_out_channels) - 1167 168            down_block = get_down_block(169                down_block_type,170                num_layers=layers_per_block,171                in_channels=input_channel,172                out_channels=output_channel,173                temb_channels=time_embed_dim,174                add_downsample=not is_final_block,175                resnet_eps=norm_eps,176                resnet_act_fn=act_fn,177                resnet_groups=norm_num_groups,178                cross_attention_dim=cross_attention_dim,179                attn_num_head_channels=attention_head_dim[i],180                downsample_padding=downsample_padding,181                dual_cross_attention=False,182            )183            self.down_blocks.append(down_block)184 185        # mid186        self.mid_block = UNetMidBlock3DCrossAttn(187            in_channels=block_out_channels[-1],188            temb_channels=time_embed_dim,189            resnet_eps=norm_eps,190            resnet_act_fn=act_fn,191            output_scale_factor=mid_block_scale_factor,192            cross_attention_dim=cross_attention_dim,193            attn_num_head_channels=attention_head_dim[-1],194            resnet_groups=norm_num_groups,195            dual_cross_attention=False,196        )197 198        # count how many layers upsample the images199        self.num_upsamplers = 0200 201        # up202        reversed_block_out_channels = list(reversed(block_out_channels))203        reversed_attention_head_dim = list(reversed(attention_head_dim))204 205        output_channel = reversed_block_out_channels[0]206        for i, up_block_type in enumerate(up_block_types):207            is_final_block = i == len(block_out_channels) - 1208 209            prev_output_channel = output_channel210            output_channel = reversed_block_out_channels[i]211            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]212 213            # add upsample block for all BUT final layer214            if not is_final_block:215                add_upsample = True216                self.num_upsamplers += 1217            else:218                add_upsample = False219 220            up_block = get_up_block(221                up_block_type,222                num_layers=layers_per_block + 1,223                in_channels=input_channel,224                out_channels=output_channel,225                prev_output_channel=prev_output_channel,226                temb_channels=time_embed_dim,227                add_upsample=add_upsample,228                resnet_eps=norm_eps,229                resnet_act_fn=act_fn,230                resnet_groups=norm_num_groups,231                cross_attention_dim=cross_attention_dim,232                attn_num_head_channels=reversed_attention_head_dim[i],233                dual_cross_attention=False,234            )235            self.up_blocks.append(up_block)236            prev_output_channel = output_channel237 238        # out239        if norm_num_groups is not None:240            self.conv_norm_out = nn.GroupNorm(241                num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps242            )243            self.conv_act = nn.SiLU()244        else:245            self.conv_norm_out = None246            self.conv_act = None247 248        conv_out_padding = (conv_out_kernel - 1) // 2249        self.conv_out = nn.Conv2d(250            block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding251        )252 253    @property254    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors255    def attn_processors(self) -> Dict[str, AttentionProcessor]:256        r"""257        Returns:258            `dict` of attention processors: A dictionary containing all attention processors used in the model with259            indexed by its weight name.260        """261        # set recursively262        processors = {}263 264        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):265            if hasattr(module, "set_processor"):266                processors[f"{name}.processor"] = module.processor267 268            for sub_name, child in module.named_children():269                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)270 271            return processors272 273        for name, module in self.named_children():274            fn_recursive_add_processors(name, module, processors)275 276        return processors277 278    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attention_slice279    def set_attention_slice(self, slice_size):280        r"""281        Enable sliced attention computation.282 283        When this option is enabled, the attention module will split the input tensor in slices, to compute attention284        in several steps. This is useful to save some memory in exchange for a small speed decrease.285 286        Args:287            slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):288                When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If289                `"max"`, maximum amount of memory will be saved by running only one slice at a time. If a number is290                provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`291                must be a multiple of `slice_size`.292        """293        sliceable_head_dims = []294 295        def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):296            if hasattr(module, "set_attention_slice"):297                sliceable_head_dims.append(module.sliceable_head_dim)298 299            for child in module.children():300                fn_recursive_retrieve_sliceable_dims(child)301 302        # retrieve number of attention layers303        for module in self.children():304            fn_recursive_retrieve_sliceable_dims(module)305 306        num_sliceable_layers = len(sliceable_head_dims)307 308        if slice_size == "auto":309            # half the attention head size is usually a good trade-off between310            # speed and memory311            slice_size = [dim // 2 for dim in sliceable_head_dims]312        elif slice_size == "max":313            # make smallest slice possible314            slice_size = num_sliceable_layers * [1]315 316        slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size317 318        if len(slice_size) != len(sliceable_head_dims):319            raise ValueError(320                f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"321                f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."322            )323 324        for i in range(len(slice_size)):325            size = slice_size[i]326            dim = sliceable_head_dims[i]327            if size is not None and size > dim:328                raise ValueError(f"size {size} has to be smaller or equal to {dim}.")329 330        # Recursively walk through all the children.331        # Any children which exposes the set_attention_slice method332        # gets the message333        def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):334            if hasattr(module, "set_attention_slice"):335                module.set_attention_slice(slice_size.pop())336 337            for child in module.children():338                fn_recursive_set_attention_slice(child, slice_size)339 340        reversed_slice_size = list(reversed(slice_size))341        for module in self.children():342            fn_recursive_set_attention_slice(module, reversed_slice_size)343 344    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor345    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):346        r"""347        Parameters:348            `processor (`dict` of `AttentionProcessor` or `AttentionProcessor`):349                The instantiated processor class or a dictionary of processor classes that will be set as the processor350                of **all** `Attention` layers.351            In case `processor` is a dict, the key needs to define the path to the corresponding cross attention processor. This is strongly recommended when setting trainable attention processors.:352 353        """354        count = len(self.attn_processors.keys())355 356        if isinstance(processor, dict) and len(processor) != count:357            raise ValueError(358                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"359                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."360            )361 362        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):363            if hasattr(module, "set_processor"):364                if not isinstance(processor, dict):365                    module.set_processor(processor)366                else:367                    module.set_processor(processor.pop(f"{name}.processor"))368 369            for sub_name, child in module.named_children():370                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)371 372        for name, module in self.named_children():373            fn_recursive_attn_processor(name, module, processor)374 375    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor376    def set_default_attn_processor(self):377        """378        Disables custom attention processors and sets the default attention implementation.379        """380        self.set_attn_processor(AttnProcessor())381 382    def _set_gradient_checkpointing(self, module, value=False):383        if isinstance(module, (CrossAttnDownBlock3D, DownBlock3D, CrossAttnUpBlock3D, UpBlock3D)):384            module.gradient_checkpointing = value385 386    def forward(387        self,388        sample: torch.FloatTensor,389        timestep: Union[torch.Tensor, float, int],390        encoder_hidden_states: torch.Tensor,391        class_labels: Optional[torch.Tensor] = None,392        timestep_cond: Optional[torch.Tensor] = None,393        attention_mask: Optional[torch.Tensor] = None,394        cross_attention_kwargs: Optional[Dict[str, Any]] = None,395        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,396        mid_block_additional_residual: Optional[torch.Tensor] = None,397        return_dict: bool = True,398    ) -> Union[UNet3DConditionOutput, Tuple]:399        r"""400        Args:401            sample (`torch.FloatTensor`): (batch, num_frames, channel, height, width) noisy inputs tensor402            timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps403            encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states404            return_dict (`bool`, *optional*, defaults to `True`):405                Whether or not to return a [`models.unet_2d_condition.UNet3DConditionOutput`] instead of a plain tuple.406            cross_attention_kwargs (`dict`, *optional*):407                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under408                `self.processor` in409                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).410 411        Returns:412            [`~models.unet_2d_condition.UNet3DConditionOutput`] or `tuple`:413            [`~models.unet_2d_condition.UNet3DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When414            returning a tuple, the first element is the sample tensor.415        """416        # By default samples have to be AT least a multiple of the overall upsampling factor.417        # The overall upsampling factor is equal to 2 ** (# num of upsampling layears).418        # However, the upsampling interpolation output size can be forced to fit any upsampling size419        # on the fly if necessary.420        default_overall_up_factor = 2**self.num_upsamplers421 422        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`423        forward_upsample_size = False424        upsample_size = None425 426        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):427            logger.info("Forward upsample size to force interpolation output size.")428            forward_upsample_size = True429 430        # prepare attention_mask431        if attention_mask is not None:432            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0433            attention_mask = attention_mask.unsqueeze(1)434 435        # 1. time436        timesteps = timestep437        if not torch.is_tensor(timesteps):438            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can439            # This would be a good case for the `match` statement (Python 3.10+)440            is_mps = sample.device.type == "mps"441            if isinstance(timestep, float):442                dtype = torch.float32 if is_mps else torch.float64443            else:444                dtype = torch.int32 if is_mps else torch.int64445            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)446        elif len(timesteps.shape) == 0:447            timesteps = timesteps[None].to(sample.device)448 449        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML450        num_frames = sample.shape[2]451        timesteps = timesteps.expand(sample.shape[0])452 453        t_emb = self.time_proj(timesteps)454 455        # timesteps does not contain any weights and will always return f32 tensors456        # but time_embedding might actually be running in fp16. so we need to cast here.457        # there might be better ways to encapsulate this.458        t_emb = t_emb.to(dtype=self.dtype)459 460        emb = self.time_embedding(t_emb, timestep_cond)461        emb = emb.repeat_interleave(repeats=num_frames, dim=0)462        encoder_hidden_states = encoder_hidden_states.repeat_interleave(repeats=num_frames, dim=0)463 464        # 2. pre-process465        sample = sample.permute(0, 2, 1, 3, 4).reshape((sample.shape[0] * num_frames, -1) + sample.shape[3:])466        sample = self.conv_in(sample)467 468        sample = self.transformer_in(sample, num_frames=num_frames).sample469 470        # 3. down471        down_block_res_samples = (sample,)472        for downsample_block in self.down_blocks:473            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:474                sample, res_samples = downsample_block(475                    hidden_states=sample,476                    temb=emb,477                    encoder_hidden_states=encoder_hidden_states,478                    attention_mask=attention_mask,479                    num_frames=num_frames,480                    cross_attention_kwargs=cross_attention_kwargs,481                )482            else:483                sample, res_samples = downsample_block(hidden_states=sample, temb=emb, num_frames=num_frames)484 485            down_block_res_samples += res_samples486 487        if down_block_additional_residuals is not None:488            new_down_block_res_samples = ()489 490            for down_block_res_sample, down_block_additional_residual in zip(491                down_block_res_samples, down_block_additional_residuals492            ):493                down_block_res_sample = down_block_res_sample + down_block_additional_residual494                new_down_block_res_samples += (down_block_res_sample,)495 496            down_block_res_samples = new_down_block_res_samples497 498        # 4. mid499        if self.mid_block is not None:500            sample = self.mid_block(501                sample,502                emb,503                encoder_hidden_states=encoder_hidden_states,504                attention_mask=attention_mask,505                num_frames=num_frames,506                cross_attention_kwargs=cross_attention_kwargs,507            )508 509        if mid_block_additional_residual is not None:510            sample = sample + mid_block_additional_residual511 512        # 5. up513        for i, upsample_block in enumerate(self.up_blocks):514            is_final_block = i == len(self.up_blocks) - 1515 516            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]517            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]518 519            # if we have not reached the final block and need to forward the520            # upsample size, we do it here521            if not is_final_block and forward_upsample_size:522                upsample_size = down_block_res_samples[-1].shape[2:]523 524            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:525                sample = upsample_block(526                    hidden_states=sample,527                    temb=emb,528                    res_hidden_states_tuple=res_samples,529                    encoder_hidden_states=encoder_hidden_states,530                    upsample_size=upsample_size,531                    attention_mask=attention_mask,532                    num_frames=num_frames,533                    cross_attention_kwargs=cross_attention_kwargs,534                )535            else:536                sample = upsample_block(537                    hidden_states=sample,538                    temb=emb,539                    res_hidden_states_tuple=res_samples,540                    upsample_size=upsample_size,541                    num_frames=num_frames,542                )543 544        # 6. post-process545        if self.conv_norm_out:546            sample = self.conv_norm_out(sample)547            sample = self.conv_act(sample)548 549        sample = self.conv_out(sample)550 551        # reshape to (batch, channel, framerate, width, height)552        sample = sample[None, :].reshape((-1, num_frames) + sample.shape[1:]).permute(0, 2, 1, 3, 4)553 554        if not return_dict:555            return (sample,)556 557        return UNet3DConditionOutput(sample=sample)558