svjack/MotionClone-Text-to-Video

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# 15#  Changes were made to this source code by Yuwei Guo.16from dataclasses import dataclass17from typing import Any, Dict, List, Optional, Tuple, Union18 19import torch20from torch import nn21from torch.nn import functional as F22 23from diffusers.configuration_utils import ConfigMixin, register_to_config24from diffusers.utils import BaseOutput, logging25from diffusers.models.embeddings import TimestepEmbedding, Timesteps26from diffusers.models.modeling_utils import ModelMixin27 28 29from .unet_blocks import (30    CrossAttnDownBlock3D,31    DownBlock3D,32    UNetMidBlock3DCrossAttn,33    get_down_block,34)35from einops import repeat, rearrange36from .resnet import InflatedConv3d37 38from diffusers.models.unet_2d_condition import UNet2DConditionModel39 40logger = logging.get_logger(__name__)  # pylint: disable=invalid-name41 42 43@dataclass44class SparseControlNetOutput(BaseOutput):45    down_block_res_samples: Tuple[torch.Tensor]46    mid_block_res_sample: torch.Tensor47 48 49class SparseControlNetConditioningEmbedding(nn.Module):50    def __init__(51        self,52        conditioning_embedding_channels: int,53        conditioning_channels: int = 3,54        block_out_channels: Tuple[int] = (16, 32, 96, 256),55    ):56        super().__init__()57 58        self.conv_in = InflatedConv3d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1)59 60        self.blocks = nn.ModuleList([])61 62        for i in range(len(block_out_channels) - 1):63            channel_in = block_out_channels[i]64            channel_out = block_out_channels[i + 1]65            self.blocks.append(InflatedConv3d(channel_in, channel_in, kernel_size=3, padding=1))66            self.blocks.append(InflatedConv3d(channel_in, channel_out, kernel_size=3, padding=1, stride=2))67 68        self.conv_out = zero_module(69            InflatedConv3d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1)70        )71 72    def forward(self, conditioning):73        embedding = self.conv_in(conditioning)74        embedding = F.silu(embedding)75 76        for block in self.blocks:77            embedding = block(embedding)78            embedding = F.silu(embedding)79 80        embedding = self.conv_out(embedding)81 82        return embedding83 84 85class SparseControlNetModel(ModelMixin, ConfigMixin):86    _supports_gradient_checkpointing = True87 88    @register_to_config89    def __init__(90        self,91        in_channels: int = 4,92        conditioning_channels: int = 3,93        flip_sin_to_cos: bool = True,94        freq_shift: int = 0,95        down_block_types: Tuple[str] = (96            "CrossAttnDownBlock2D",97            "CrossAttnDownBlock2D",98            "CrossAttnDownBlock2D",99            "DownBlock2D",100        ),101        only_cross_attention: Union[bool, Tuple[bool]] = False,102        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),103        layers_per_block: int = 2,104        downsample_padding: int = 1,105        mid_block_scale_factor: float = 1,106        act_fn: str = "silu",107        norm_num_groups: Optional[int] = 32,108        norm_eps: float = 1e-5,109        cross_attention_dim: int = 1280,110        attention_head_dim: Union[int, Tuple[int]] = 8,111        num_attention_heads: Optional[Union[int, Tuple[int]]] = None,112        use_linear_projection: bool = False,113        class_embed_type: Optional[str] = None,114        num_class_embeds: Optional[int] = None,115        upcast_attention: bool = False,116        resnet_time_scale_shift: str = "default",117        projection_class_embeddings_input_dim: Optional[int] = None,118        controlnet_conditioning_channel_order: str = "rgb",119        conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, 96, 256),120        global_pool_conditions: bool = False,121 122        use_motion_module         = True,123        motion_module_resolutions = ( 1,2,4,8 ),124        motion_module_mid_block   = False,125        motion_module_type        = "Vanilla",126        motion_module_kwargs      = {127            "num_attention_heads": 8,128            "num_transformer_block": 1,129            "attention_block_types": ["Temporal_Self"],130            "temporal_position_encoding": True,131            "temporal_position_encoding_max_len": 32,132            "temporal_attention_dim_div": 1,133            "causal_temporal_attention": False,134        },135 136        concate_conditioning_mask: bool = True,137        use_simplified_condition_embedding:  bool = False,138 139        set_noisy_sample_input_to_zero: bool = False,140    ):141        super().__init__()142 143        # If `num_attention_heads` is not defined (which is the case for most models)144        # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.145        # The reason for this behavior is to correct for incorrectly named variables that were introduced146        # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131147        # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking148        # which is why we correct for the naming here.149        num_attention_heads = num_attention_heads or attention_head_dim150 151        # Check inputs152        if len(block_out_channels) != len(down_block_types):153            raise ValueError(154                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}."155            )156 157        if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):158            raise ValueError(159                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}."160            )161 162        if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):163            raise ValueError(164                f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."165            )166 167        # input168        self.set_noisy_sample_input_to_zero  = set_noisy_sample_input_to_zero169 170        conv_in_kernel = 3171        conv_in_padding = (conv_in_kernel - 1) // 2172        self.conv_in = InflatedConv3d(173            in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding174        )175 176        if concate_conditioning_mask:177            conditioning_channels = conditioning_channels + 1178        self.concate_conditioning_mask = concate_conditioning_mask179 180        # control net conditioning embedding181        if use_simplified_condition_embedding:182            self.controlnet_cond_embedding = zero_module(183                InflatedConv3d(conditioning_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding)184            ).to(torch.float16)185        else:186            self.controlnet_cond_embedding = SparseControlNetConditioningEmbedding(187                conditioning_embedding_channels=block_out_channels[0],188                block_out_channels=conditioning_embedding_out_channels,189                conditioning_channels=conditioning_channels,190            ).to(torch.float16)191        self.use_simplified_condition_embedding = use_simplified_condition_embedding192 193        # time194        time_embed_dim = block_out_channels[0] * 4195 196        self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)197        timestep_input_dim = block_out_channels[0]198 199        self.time_embedding = TimestepEmbedding(200            timestep_input_dim,201            time_embed_dim,202            act_fn=act_fn,203        )204 205        # class embedding206        if class_embed_type is None and num_class_embeds is not None:207            self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)208        elif class_embed_type == "timestep":209            self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)210        elif class_embed_type == "identity":211            self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)212        elif class_embed_type == "projection":213            if projection_class_embeddings_input_dim is None:214                raise ValueError(215                    "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"216                )217            # The projection `class_embed_type` is the same as the timestep `class_embed_type` except218            # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings219            # 2. it projects from an arbitrary input dimension.220            #221            # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.222            # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.223            # As a result, `TimestepEmbedding` can be passed arbitrary vectors.224            self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)225        else:226            self.class_embedding = None227 228 229        self.down_blocks = nn.ModuleList([])230        self.controlnet_down_blocks = nn.ModuleList([])231 232        if isinstance(only_cross_attention, bool):233            only_cross_attention = [only_cross_attention] * len(down_block_types)234 235        if isinstance(attention_head_dim, int):236            attention_head_dim = (attention_head_dim,) * len(down_block_types)237 238        if isinstance(num_attention_heads, int):239            num_attention_heads = (num_attention_heads,) * len(down_block_types)240 241        # down242        output_channel = block_out_channels[0]243 244        controlnet_block = InflatedConv3d(output_channel, output_channel, kernel_size=1)245        controlnet_block = zero_module(controlnet_block)246        self.controlnet_down_blocks.append(controlnet_block)247 248        for i, down_block_type in enumerate(down_block_types):249            res = 2 ** i250            input_channel = output_channel251            output_channel = block_out_channels[i]252            is_final_block = i == len(block_out_channels) - 1253 254            down_block = get_down_block(255                down_block_type,256                num_layers=layers_per_block,257                in_channels=input_channel,258                out_channels=output_channel,259                temb_channels=time_embed_dim,260                add_downsample=not is_final_block,261                resnet_eps=norm_eps,262                resnet_act_fn=act_fn,263                resnet_groups=norm_num_groups,264                cross_attention_dim=cross_attention_dim,265                attn_num_head_channels=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,266                downsample_padding=downsample_padding,267                use_linear_projection=use_linear_projection,268                only_cross_attention=only_cross_attention[i],269                upcast_attention=upcast_attention,270                resnet_time_scale_shift=resnet_time_scale_shift,271 272                use_inflated_groupnorm=True,273 274                use_motion_module=use_motion_module and (res in motion_module_resolutions),275                motion_module_type=motion_module_type,276                motion_module_kwargs=motion_module_kwargs,277            )278            self.down_blocks.append(down_block)279 280            for _ in range(layers_per_block):281                controlnet_block = InflatedConv3d(output_channel, output_channel, kernel_size=1)282                controlnet_block = zero_module(controlnet_block)283                self.controlnet_down_blocks.append(controlnet_block)284 285            if not is_final_block:286                controlnet_block = InflatedConv3d(output_channel, output_channel, kernel_size=1)287                controlnet_block = zero_module(controlnet_block)288                self.controlnet_down_blocks.append(controlnet_block)289 290        # mid291        mid_block_channel = block_out_channels[-1]292 293        controlnet_block = InflatedConv3d(mid_block_channel, mid_block_channel, kernel_size=1)294        controlnet_block = zero_module(controlnet_block)295        self.controlnet_mid_block = controlnet_block296 297        self.mid_block = UNetMidBlock3DCrossAttn(298            in_channels=mid_block_channel,299            temb_channels=time_embed_dim,300            resnet_eps=norm_eps,301            resnet_act_fn=act_fn,302            output_scale_factor=mid_block_scale_factor,303            resnet_time_scale_shift=resnet_time_scale_shift,304            cross_attention_dim=cross_attention_dim,305            attn_num_head_channels=num_attention_heads[-1],306            resnet_groups=norm_num_groups,307            use_linear_projection=use_linear_projection,308            upcast_attention=upcast_attention,309 310            use_inflated_groupnorm=True,311            use_motion_module=use_motion_module and motion_module_mid_block,312            motion_module_type=motion_module_type,313            motion_module_kwargs=motion_module_kwargs,314        )315 316    @classmethod317    def from_unet(318        cls,319        unet: UNet2DConditionModel,320        controlnet_conditioning_channel_order: str = "rgb",321        conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, 96, 256),322        load_weights_from_unet: bool = True,323 324        controlnet_additional_kwargs: dict = {},325    ):326        controlnet = cls(327            in_channels=unet.config.in_channels,328            flip_sin_to_cos=unet.config.flip_sin_to_cos,329            freq_shift=unet.config.freq_shift,330            down_block_types=unet.config.down_block_types,331            only_cross_attention=unet.config.only_cross_attention,332            block_out_channels=unet.config.block_out_channels,333            layers_per_block=unet.config.layers_per_block,334            downsample_padding=unet.config.downsample_padding,335            mid_block_scale_factor=unet.config.mid_block_scale_factor,336            act_fn=unet.config.act_fn,337            norm_num_groups=unet.config.norm_num_groups,338            norm_eps=unet.config.norm_eps,339            cross_attention_dim=unet.config.cross_attention_dim,340            attention_head_dim=unet.config.attention_head_dim,341            num_attention_heads=unet.config.num_attention_heads,342            use_linear_projection=unet.config.use_linear_projection,343            class_embed_type=unet.config.class_embed_type,344            num_class_embeds=unet.config.num_class_embeds,345            upcast_attention=unet.config.upcast_attention,346            resnet_time_scale_shift=unet.config.resnet_time_scale_shift,347            projection_class_embeddings_input_dim=unet.config.projection_class_embeddings_input_dim,348            controlnet_conditioning_channel_order=controlnet_conditioning_channel_order,349            conditioning_embedding_out_channels=conditioning_embedding_out_channels,350 351            **controlnet_additional_kwargs,352        )353 354        if load_weights_from_unet:355            m, u = controlnet.conv_in.load_state_dict(cls.image_layer_filter(unet.conv_in.state_dict()), strict=False)356            assert len(u) == 0357            m, u = controlnet.time_proj.load_state_dict(cls.image_layer_filter(unet.time_proj.state_dict()), strict=False)358            assert len(u) == 0359            m, u = controlnet.time_embedding.load_state_dict(cls.image_layer_filter(unet.time_embedding.state_dict()), strict=False)360            assert len(u) == 0361 362            if controlnet.class_embedding:363                m, u = controlnet.class_embedding.load_state_dict(cls.image_layer_filter(unet.class_embedding.state_dict()), strict=False)364                assert len(u) == 0365            m, u = controlnet.down_blocks.load_state_dict(cls.image_layer_filter(unet.down_blocks.state_dict()), strict=False)366            assert len(u) == 0367            m, u = controlnet.mid_block.load_state_dict(cls.image_layer_filter(unet.mid_block.state_dict()), strict=False)368            assert len(u) == 0369 370        return controlnet371 372    @staticmethod373    def image_layer_filter(state_dict):374        new_state_dict = {}375        for name, param in state_dict.items():376            if "motion_modules." in name or "lora" in name: continue377            new_state_dict[name] = param378        return new_state_dict379 380    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attention_slice381    def set_attention_slice(self, slice_size):382        r"""383        Enable sliced attention computation.384 385        When this option is enabled, the attention module splits the input tensor in slices to compute attention in386        several steps. This is useful for saving some memory in exchange for a small decrease in speed.387 388        Args:389            slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):390                When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If391                `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is392                provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`393                must be a multiple of `slice_size`.394        """395        sliceable_head_dims = []396 397        def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):398            if hasattr(module, "set_attention_slice"):399                sliceable_head_dims.append(module.sliceable_head_dim)400 401            for child in module.children():402                fn_recursive_retrieve_sliceable_dims(child)403 404        # retrieve number of attention layers405        for module in self.children():406            fn_recursive_retrieve_sliceable_dims(module)407 408        num_sliceable_layers = len(sliceable_head_dims)409 410        if slice_size == "auto":411            # half the attention head size is usually a good trade-off between412            # speed and memory413            slice_size = [dim // 2 for dim in sliceable_head_dims]414        elif slice_size == "max":415            # make smallest slice possible416            slice_size = num_sliceable_layers * [1]417 418        slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size419 420        if len(slice_size) != len(sliceable_head_dims):421            raise ValueError(422                f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"423                f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."424            )425 426        for i in range(len(slice_size)):427            size = slice_size[i]428            dim = sliceable_head_dims[i]429            if size is not None and size > dim:430                raise ValueError(f"size {size} has to be smaller or equal to {dim}.")431 432        # Recursively walk through all the children.433        # Any children which exposes the set_attention_slice method434        # gets the message435        def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):436            if hasattr(module, "set_attention_slice"):437                module.set_attention_slice(slice_size.pop())438 439            for child in module.children():440                fn_recursive_set_attention_slice(child, slice_size)441 442        reversed_slice_size = list(reversed(slice_size))443        for module in self.children():444            fn_recursive_set_attention_slice(module, reversed_slice_size)445 446    def _set_gradient_checkpointing(self, module, value=False):447        if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D)):448            module.gradient_checkpointing = value449 450    def forward(451        self,452        sample: torch.FloatTensor,453        timestep: Union[torch.Tensor, float, int],454        encoder_hidden_states: torch.Tensor,455 456        controlnet_cond: torch.FloatTensor,457        conditioning_mask: Optional[torch.FloatTensor] = None,458 459        conditioning_scale: float = 1.0,460        class_labels: Optional[torch.Tensor] = None,461        attention_mask: Optional[torch.Tensor] = None,462        cross_attention_kwargs: Optional[Dict[str, Any]] = None,463        guess_mode: bool = False,464        return_dict: bool = True,465    ) -> Union[SparseControlNetOutput, Tuple]:466 467        # set input noise to zero468        # if self.set_noisy_sample_input_to_zero:469        #     sample = torch.zeros_like(sample).to(sample.device)470 471        # prepare attention_mask472        if attention_mask is not None:473            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0474            attention_mask = attention_mask.unsqueeze(1)475 476        # 1. time477        timesteps = timestep478        if not torch.is_tensor(timesteps):479            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can480            # This would be a good case for the `match` statement (Python 3.10+)481            is_mps = sample.device.type == "mps"482            if isinstance(timestep, float):483                dtype = torch.float32 if is_mps else torch.float64484            else:485                dtype = torch.int32 if is_mps else torch.int64486            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)487        elif len(timesteps.shape) == 0:488            timesteps = timesteps[None].to(sample.device)489 490        timesteps             = timesteps.repeat(sample.shape[0] // timesteps.shape[0])491        encoder_hidden_states = encoder_hidden_states.repeat(sample.shape[0] // encoder_hidden_states.shape[0], 1, 1)492 493        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML494        timesteps = timesteps.expand(sample.shape[0])495 496        t_emb = self.time_proj(timesteps)497 498        # timesteps does not contain any weights and will always return f32 tensors499        # but time_embedding might actually be running in fp16. so we need to cast here.500        # there might be better ways to encapsulate this.501        t_emb = t_emb.to(dtype=self.dtype)502        emb = self.time_embedding(t_emb)503 504        if self.class_embedding is not None:505            if class_labels is None:506                raise ValueError("class_labels should be provided when num_class_embeds > 0")507 508            if self.config.class_embed_type == "timestep":509                class_labels = self.time_proj(class_labels)510 511            class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)512            emb = emb + class_emb513 514        # 2. pre-process515        # equal to set input noise to zero516        if self.set_noisy_sample_input_to_zero:517            shape = sample.shape518            sample = self.conv_in.bias.reshape(1,-1,1,1,1).expand(shape[0],-1,shape[2],shape[3],shape[4])519        else:520            sample = self.conv_in(sample)521 522        if self.concate_conditioning_mask:523            controlnet_cond = torch.cat([controlnet_cond, conditioning_mask], dim=1).to(torch.float16)524        # import pdb; pdb.set_trace()525        controlnet_cond = self.controlnet_cond_embedding(controlnet_cond)526        527        sample = sample + controlnet_cond528 529        # 3. down530        down_block_res_samples = (sample,)531        for downsample_block in self.down_blocks:532            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:533                sample, res_samples = downsample_block(534                    hidden_states=sample,535                    temb=emb,536                    encoder_hidden_states=encoder_hidden_states,537                    attention_mask=attention_mask,538                    # cross_attention_kwargs=cross_attention_kwargs,539                )540            else: sample, res_samples = downsample_block(hidden_states=sample, temb=emb)541 542            down_block_res_samples += res_samples543 544        # 4. mid545        if self.mid_block is not None:546            sample = self.mid_block(547                sample,548                emb,549                encoder_hidden_states=encoder_hidden_states,550                attention_mask=attention_mask,551                # cross_attention_kwargs=cross_attention_kwargs,552            )553 554        # 5. controlnet blocks555        controlnet_down_block_res_samples = ()556 557        for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks):558            down_block_res_sample = controlnet_block(down_block_res_sample)559            controlnet_down_block_res_samples = controlnet_down_block_res_samples + (down_block_res_sample,)560 561        down_block_res_samples = controlnet_down_block_res_samples562 563        mid_block_res_sample = self.controlnet_mid_block(sample)564 565        # 6. scaling566        if guess_mode and not self.config.global_pool_conditions:567            scales = torch.logspace(-1, 0, len(down_block_res_samples) + 1, device=sample.device)  # 0.1 to 1.0568 569            scales = scales * conditioning_scale570            down_block_res_samples = [sample * scale for sample, scale in zip(down_block_res_samples, scales)]571            mid_block_res_sample = mid_block_res_sample * scales[-1]  # last one572        else:573            down_block_res_samples = [sample * conditioning_scale for sample in down_block_res_samples]574            mid_block_res_sample = mid_block_res_sample * conditioning_scale575 576        if self.config.global_pool_conditions:577            down_block_res_samples = [578                torch.mean(sample, dim=(2, 3), keepdim=True) for sample in down_block_res_samples579            ]580            mid_block_res_sample = torch.mean(mid_block_res_sample, dim=(2, 3), keepdim=True)581 582        if not return_dict:583            return (down_block_res_samples, mid_block_res_sample)584 585        return SparseControlNetOutput(586            down_block_res_samples=down_block_res_samples, mid_block_res_sample=mid_block_res_sample587        )588 589 590def zero_module(module):591    for p in module.parameters():592        nn.init.zeros_(p)593    return module594