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.14from dataclasses import dataclass15from typing import Any, Dict, List, Optional, Tuple, Union16 17import torch18from torch import nn19from torch.nn import functional as F20 21from ..configuration_utils import ConfigMixin, register_to_config22from ..utils import BaseOutput, logging23from .attention_processor import AttentionProcessor, AttnProcessor24from .embeddings import TimestepEmbedding, Timesteps25from .modeling_utils import ModelMixin26from .unet_2d_blocks import (27    CrossAttnDownBlock2D,28    DownBlock2D,29    UNetMidBlock2DCrossAttn,30    get_down_block,31)32from .unet_2d_condition import UNet2DConditionModel33 34 35logger = logging.get_logger(__name__)  # pylint: disable=invalid-name36 37 38@dataclass39class ControlNetOutput(BaseOutput):40    down_block_res_samples: Tuple[torch.Tensor]41    mid_block_res_sample: torch.Tensor42 43 44class ControlNetConditioningEmbedding(nn.Module):45    """46    Quoting from https://arxiv.org/abs/2302.05543: "Stable Diffusion uses a pre-processing method similar to VQ-GAN47    [11] to convert the entire dataset of 512 × 512 images into smaller 64 × 64 “latent images” for stabilized48    training. This requires ControlNets to convert image-based conditions to 64 × 64 feature space to match the49    convolution size. We use a tiny network E(·) of four convolution layers with 4 × 4 kernels and 2 × 2 strides50    (activated by ReLU, channels are 16, 32, 64, 128, initialized with Gaussian weights, trained jointly with the full51    model) to encode image-space conditions ... into feature maps ..."52    """53 54    def __init__(55        self,56        conditioning_embedding_channels: int,57        conditioning_channels: int = 3,58        block_out_channels: Tuple[int] = (16, 32, 96, 256),59    ):60        super().__init__()61 62        self.conv_in = nn.Conv2d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1)63 64        self.blocks = nn.ModuleList([])65 66        for i in range(len(block_out_channels) - 1):67            channel_in = block_out_channels[i]68            channel_out = block_out_channels[i + 1]69            self.blocks.append(nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1))70            self.blocks.append(nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1, stride=2))71 72        self.conv_out = zero_module(73            nn.Conv2d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1)74        )75 76    def forward(self, conditioning):77        embedding = self.conv_in(conditioning)78        embedding = F.silu(embedding)79 80        for block in self.blocks:81            embedding = block(embedding)82            embedding = F.silu(embedding)83 84        embedding = self.conv_out(embedding)85 86        return embedding87 88 89class ControlNetModel(ModelMixin, ConfigMixin):90    _supports_gradient_checkpointing = True91 92    @register_to_config93    def __init__(94        self,95        in_channels: int = 4,96        flip_sin_to_cos: bool = True,97        freq_shift: int = 0,98        down_block_types: Tuple[str] = (99            "CrossAttnDownBlock2D",100            "CrossAttnDownBlock2D",101            "CrossAttnDownBlock2D",102            "DownBlock2D",103        ),104        only_cross_attention: Union[bool, Tuple[bool]] = False,105        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),106        layers_per_block: int = 2,107        downsample_padding: int = 1,108        mid_block_scale_factor: float = 1,109        act_fn: str = "silu",110        norm_num_groups: Optional[int] = 32,111        norm_eps: float = 1e-5,112        cross_attention_dim: int = 1280,113        attention_head_dim: Union[int, Tuple[int]] = 8,114        use_linear_projection: bool = False,115        class_embed_type: Optional[str] = None,116        num_class_embeds: Optional[int] = None,117        upcast_attention: bool = False,118        resnet_time_scale_shift: str = "default",119        projection_class_embeddings_input_dim: Optional[int] = None,120        controlnet_conditioning_channel_order: str = "rgb",121        conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, 96, 256),122    ):123        super().__init__()124 125        # Check inputs126        if len(block_out_channels) != len(down_block_types):127            raise ValueError(128                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}."129            )130 131        if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):132            raise ValueError(133                f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."134            )135 136        if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):137            raise ValueError(138                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}."139            )140 141        # input142        conv_in_kernel = 3143        conv_in_padding = (conv_in_kernel - 1) // 2144        self.conv_in = nn.Conv2d(145            in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding146        )147 148        # time149        time_embed_dim = block_out_channels[0] * 4150 151        self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)152        timestep_input_dim = block_out_channels[0]153 154        self.time_embedding = TimestepEmbedding(155            timestep_input_dim,156            time_embed_dim,157            act_fn=act_fn,158        )159 160        # class embedding161        if class_embed_type is None and num_class_embeds is not None:162            self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)163        elif class_embed_type == "timestep":164            self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)165        elif class_embed_type == "identity":166            self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)167        elif class_embed_type == "projection":168            if projection_class_embeddings_input_dim is None:169                raise ValueError(170                    "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"171                )172            # The projection `class_embed_type` is the same as the timestep `class_embed_type` except173            # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings174            # 2. it projects from an arbitrary input dimension.175            #176            # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.177            # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.178            # As a result, `TimestepEmbedding` can be passed arbitrary vectors.179            self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)180        else:181            self.class_embedding = None182 183        # control net conditioning embedding184        self.controlnet_cond_embedding = ControlNetConditioningEmbedding(185            conditioning_embedding_channels=block_out_channels[0],186            block_out_channels=conditioning_embedding_out_channels,187        )188 189        self.down_blocks = nn.ModuleList([])190        self.controlnet_down_blocks = nn.ModuleList([])191 192        if isinstance(only_cross_attention, bool):193            only_cross_attention = [only_cross_attention] * len(down_block_types)194 195        if isinstance(attention_head_dim, int):196            attention_head_dim = (attention_head_dim,) * len(down_block_types)197 198        # down199        output_channel = block_out_channels[0]200 201        controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)202        controlnet_block = zero_module(controlnet_block)203        self.controlnet_down_blocks.append(controlnet_block)204 205        for i, down_block_type in enumerate(down_block_types):206            input_channel = output_channel207            output_channel = block_out_channels[i]208            is_final_block = i == len(block_out_channels) - 1209 210            down_block = get_down_block(211                down_block_type,212                num_layers=layers_per_block,213                in_channels=input_channel,214                out_channels=output_channel,215                temb_channels=time_embed_dim,216                add_downsample=not is_final_block,217                resnet_eps=norm_eps,218                resnet_act_fn=act_fn,219                resnet_groups=norm_num_groups,220                cross_attention_dim=cross_attention_dim,221                attn_num_head_channels=attention_head_dim[i],222                downsample_padding=downsample_padding,223                use_linear_projection=use_linear_projection,224                only_cross_attention=only_cross_attention[i],225                upcast_attention=upcast_attention,226                resnet_time_scale_shift=resnet_time_scale_shift,227            )228            self.down_blocks.append(down_block)229 230            for _ in range(layers_per_block):231                controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)232                controlnet_block = zero_module(controlnet_block)233                self.controlnet_down_blocks.append(controlnet_block)234 235            if not is_final_block:236                controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)237                controlnet_block = zero_module(controlnet_block)238                self.controlnet_down_blocks.append(controlnet_block)239 240        # mid241        mid_block_channel = block_out_channels[-1]242 243        controlnet_block = nn.Conv2d(mid_block_channel, mid_block_channel, kernel_size=1)244        controlnet_block = zero_module(controlnet_block)245        self.controlnet_mid_block = controlnet_block246 247        self.mid_block = UNetMidBlock2DCrossAttn(248            in_channels=mid_block_channel,249            temb_channels=time_embed_dim,250            resnet_eps=norm_eps,251            resnet_act_fn=act_fn,252            output_scale_factor=mid_block_scale_factor,253            resnet_time_scale_shift=resnet_time_scale_shift,254            cross_attention_dim=cross_attention_dim,255            attn_num_head_channels=attention_head_dim[-1],256            resnet_groups=norm_num_groups,257            use_linear_projection=use_linear_projection,258            upcast_attention=upcast_attention,259        )260 261    @classmethod262    def from_unet(263        cls,264        unet: UNet2DConditionModel,265        controlnet_conditioning_channel_order: str = "rgb",266        conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, 96, 256),267        load_weights_from_unet: bool = True,268    ):269        r"""270        Instantiate Controlnet class from UNet2DConditionModel.271 272        Parameters:273            unet (`UNet2DConditionModel`):274                UNet model which weights are copied to the ControlNet. Note that all configuration options are also275                copied where applicable.276        """277        controlnet = cls(278            in_channels=unet.config.in_channels,279            flip_sin_to_cos=unet.config.flip_sin_to_cos,280            freq_shift=unet.config.freq_shift,281            down_block_types=unet.config.down_block_types,282            only_cross_attention=unet.config.only_cross_attention,283            block_out_channels=unet.config.block_out_channels,284            layers_per_block=unet.config.layers_per_block,285            downsample_padding=unet.config.downsample_padding,286            mid_block_scale_factor=unet.config.mid_block_scale_factor,287            act_fn=unet.config.act_fn,288            norm_num_groups=unet.config.norm_num_groups,289            norm_eps=unet.config.norm_eps,290            cross_attention_dim=unet.config.cross_attention_dim,291            attention_head_dim=unet.config.attention_head_dim,292            use_linear_projection=unet.config.use_linear_projection,293            class_embed_type=unet.config.class_embed_type,294            num_class_embeds=unet.config.num_class_embeds,295            upcast_attention=unet.config.upcast_attention,296            resnet_time_scale_shift=unet.config.resnet_time_scale_shift,297            projection_class_embeddings_input_dim=unet.config.projection_class_embeddings_input_dim,298            controlnet_conditioning_channel_order=controlnet_conditioning_channel_order,299            conditioning_embedding_out_channels=conditioning_embedding_out_channels,300        )301 302        if load_weights_from_unet:303            controlnet.conv_in.load_state_dict(unet.conv_in.state_dict())304            controlnet.time_proj.load_state_dict(unet.time_proj.state_dict())305            controlnet.time_embedding.load_state_dict(unet.time_embedding.state_dict())306 307            if controlnet.class_embedding:308                controlnet.class_embedding.load_state_dict(unet.class_embedding.state_dict())309 310            controlnet.down_blocks.load_state_dict(unet.down_blocks.state_dict())311            controlnet.mid_block.load_state_dict(unet.mid_block.state_dict())312 313        return controlnet314 315    @property316    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors317    def attn_processors(self) -> Dict[str, AttentionProcessor]:318        r"""319        Returns:320            `dict` of attention processors: A dictionary containing all attention processors used in the model with321            indexed by its weight name.322        """323        # set recursively324        processors = {}325 326        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):327            if hasattr(module, "set_processor"):328                processors[f"{name}.processor"] = module.processor329 330            for sub_name, child in module.named_children():331                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)332 333            return processors334 335        for name, module in self.named_children():336            fn_recursive_add_processors(name, module, processors)337 338        return processors339 340    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor341    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):342        r"""343        Parameters:344            `processor (`dict` of `AttentionProcessor` or `AttentionProcessor`):345                The instantiated processor class or a dictionary of processor classes that will be set as the processor346                of **all** `Attention` layers.347            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.:348 349        """350        count = len(self.attn_processors.keys())351 352        if isinstance(processor, dict) and len(processor) != count:353            raise ValueError(354                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"355                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."356            )357 358        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):359            if hasattr(module, "set_processor"):360                if not isinstance(processor, dict):361                    module.set_processor(processor)362                else:363                    module.set_processor(processor.pop(f"{name}.processor"))364 365            for sub_name, child in module.named_children():366                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)367 368        for name, module in self.named_children():369            fn_recursive_attn_processor(name, module, processor)370 371    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor372    def set_default_attn_processor(self):373        """374        Disables custom attention processors and sets the default attention implementation.375        """376        self.set_attn_processor(AttnProcessor())377 378    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attention_slice379    def set_attention_slice(self, slice_size):380        r"""381        Enable sliced attention computation.382 383        When this option is enabled, the attention module will split the input tensor in slices, to compute attention384        in several steps. This is useful to save some memory in exchange for a small speed decrease.385 386        Args:387            slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):388                When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If389                `"max"`, maximum amount of memory will be saved by running only one slice at a time. If a number is390                provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`391                must be a multiple of `slice_size`.392        """393        sliceable_head_dims = []394 395        def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):396            if hasattr(module, "set_attention_slice"):397                sliceable_head_dims.append(module.sliceable_head_dim)398 399            for child in module.children():400                fn_recursive_retrieve_sliceable_dims(child)401 402        # retrieve number of attention layers403        for module in self.children():404            fn_recursive_retrieve_sliceable_dims(module)405 406        num_sliceable_layers = len(sliceable_head_dims)407 408        if slice_size == "auto":409            # half the attention head size is usually a good trade-off between410            # speed and memory411            slice_size = [dim // 2 for dim in sliceable_head_dims]412        elif slice_size == "max":413            # make smallest slice possible414            slice_size = num_sliceable_layers * [1]415 416        slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size417 418        if len(slice_size) != len(sliceable_head_dims):419            raise ValueError(420                f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"421                f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."422            )423 424        for i in range(len(slice_size)):425            size = slice_size[i]426            dim = sliceable_head_dims[i]427            if size is not None and size > dim:428                raise ValueError(f"size {size} has to be smaller or equal to {dim}.")429 430        # Recursively walk through all the children.431        # Any children which exposes the set_attention_slice method432        # gets the message433        def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):434            if hasattr(module, "set_attention_slice"):435                module.set_attention_slice(slice_size.pop())436 437            for child in module.children():438                fn_recursive_set_attention_slice(child, slice_size)439 440        reversed_slice_size = list(reversed(slice_size))441        for module in self.children():442            fn_recursive_set_attention_slice(module, reversed_slice_size)443 444    def _set_gradient_checkpointing(self, module, value=False):445        if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D)):446            module.gradient_checkpointing = value447 448    def forward(449        self,450        sample: torch.FloatTensor,451        timestep: Union[torch.Tensor, float, int],452        encoder_hidden_states: torch.Tensor,453        controlnet_cond: torch.FloatTensor,454        conditioning_scale: float = 1.0,455        class_labels: Optional[torch.Tensor] = None,456        timestep_cond: Optional[torch.Tensor] = None,457        attention_mask: Optional[torch.Tensor] = None,458        cross_attention_kwargs: Optional[Dict[str, Any]] = None,459        return_dict: bool = True,460    ) -> Union[ControlNetOutput, Tuple]:461        # check channel order462        channel_order = self.config.controlnet_conditioning_channel_order463 464        if channel_order == "rgb":465            # in rgb order by default466            ...467        elif channel_order == "bgr":468            controlnet_cond = torch.flip(controlnet_cond, dims=[1])469        else:470            raise ValueError(f"unknown `controlnet_conditioning_channel_order`: {channel_order}")471 472        # prepare attention_mask473        if attention_mask is not None:474            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0475            attention_mask = attention_mask.unsqueeze(1)476 477        # 1. time478        timesteps = timestep479        if not torch.is_tensor(timesteps):480            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can481            # This would be a good case for the `match` statement (Python 3.10+)482            is_mps = sample.device.type == "mps"483            if isinstance(timestep, float):484                dtype = torch.float32 if is_mps else torch.float64485            else:486                dtype = torch.int32 if is_mps else torch.int64487            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)488        elif len(timesteps.shape) == 0:489            timesteps = timesteps[None].to(sample.device)490 491        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML492        timesteps = timesteps.expand(sample.shape[0])493 494        t_emb = self.time_proj(timesteps)495 496        # timesteps does not contain any weights and will always return f32 tensors497        # but time_embedding might actually be running in fp16. so we need to cast here.498        # there might be better ways to encapsulate this.499        t_emb = t_emb.to(dtype=self.dtype)500 501        emb = self.time_embedding(t_emb, timestep_cond)502 503        if self.class_embedding is not None:504            if class_labels is None:505                raise ValueError("class_labels should be provided when num_class_embeds > 0")506 507            if self.config.class_embed_type == "timestep":508                class_labels = self.time_proj(class_labels)509 510            class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)511            emb = emb + class_emb512 513        # 2. pre-process514        sample = self.conv_in(sample)515 516        controlnet_cond = self.controlnet_cond_embedding(controlnet_cond)517 518        sample += controlnet_cond519 520        # 3. down521        down_block_res_samples = (sample,)522        for downsample_block in self.down_blocks:523            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:524                sample, res_samples = downsample_block(525                    hidden_states=sample,526                    temb=emb,527                    encoder_hidden_states=encoder_hidden_states,528                    attention_mask=attention_mask,529                    cross_attention_kwargs=cross_attention_kwargs,530                )531            else:532                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)533 534            down_block_res_samples += res_samples535 536        # 4. mid537        if self.mid_block is not None:538            sample = self.mid_block(539                sample,540                emb,541                encoder_hidden_states=encoder_hidden_states,542                attention_mask=attention_mask,543                cross_attention_kwargs=cross_attention_kwargs,544            )545 546        # 5. Control net blocks547 548        controlnet_down_block_res_samples = ()549 550        for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks):551            down_block_res_sample = controlnet_block(down_block_res_sample)552            controlnet_down_block_res_samples += (down_block_res_sample,)553 554        down_block_res_samples = controlnet_down_block_res_samples555 556        mid_block_res_sample = self.controlnet_mid_block(sample)557 558        # 6. scaling559        down_block_res_samples = [sample * conditioning_scale for sample in down_block_res_samples]560        mid_block_res_sample *= conditioning_scale561 562        if not return_dict:563            return (down_block_res_samples, mid_block_res_sample)564 565        return ControlNetOutput(566            down_block_res_samples=down_block_res_samples, mid_block_res_sample=mid_block_res_sample567        )568 569 570def zero_module(module):571    for p in module.parameters():572        nn.init.zeros_(p)573    return module574