o pi@sddlmZmZmZddlZddlZddlmZddl mm Z ddl m Z mZddlmZddlmZddlmZdd lmZdd lmZdd lmZdd lmZdd lmZddlm Z m!Z!e"e#Z$Gdddej%Z&Gdddej'Z(Gdddej'Z)Gdddej'Z*Gdddej'Z+Gdddej'Z,Gdddej'Z-Gdddej'Z.Gd d!d!ej'Z/Gd"d#d#ee eZ0dS)$)OptionalTupleUnionN) ConfigMixinregister_to_config)FromOriginalModelMixin)logging)apply_forward_hook)get_activation)CogVideoXDownsample3D)AutoencoderKLOutput) ModelMixin)CogVideoXUpsample3D) DecoderOutputDiagonalGaussianDistributioncs.eZdZdZdejdejffdd ZZS)CogVideoXSafeConv3dzq A 3D convolution layer that splits the input tensor into smaller parts to avoid OOM in CogVideoX Model. inputreturncstt|jdd}|dkrY|jdt|dd}tj||dddkr@dgfddtdt Dg}D] }| t |qDtj |dd}|St |S)Nr i@rrdimcsFg|]}tj|ddddd ddf|fddqS)rNr r)torchcat).0i input_chunks kernel_sizet/home/ubuntu/SoloSpeech/.venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py 5s8z/CogVideoXSafeConv3d.forward..)rprodtensorshapeitemrintchunkrangelenappendsuperforwardr)selfr memory_countpart_num output_chunks input_chunkoutput __class__rr!r-+s   zCogVideoXSafeConv3d.forward)__name__ __module__ __qualname____doc__rTensorr- __classcell__r r r4r!r&s"rc seZdZdZ   ddededeeeeeeffdeded ef fd d Zd e j d e j fddZ ddZ d e j d e j fddZ ZS)CogVideoXCausalConv3da&A 3D causal convolution layer that pads the input tensor to ensure causality in CogVideoX Model. Args: in_channels (int): Number of channels in the input tensor. out_channels (int): Number of output channels. kernel_size (Union[int, Tuple[int, int, int]]): Size of the convolutional kernel. stride (int, optional): Stride of the convolution. Default is 1. dilation (int, optional): Dilation rate of the convolution. Default is 1. pad_mode (str, optional): Padding mode. Default is "constant". rconstant in_channels out_channelsrstridedilationpad_modec stt|tr|fd}|\}}} ||_||dd|} |d} | d} | |_| |_| |_| | | | | df|_d|_ ||_ |ddf}|ddf}t |||||d|_ d|_ dS)Nrrr r)r>r?rr@rA)r,__init__ isinstancer'rB height_pad width_padtime_padtime_causal_padding temporal_dimtime_kernel_sizerconv conv_cache) r.r>r?rr@rArBrJheight_kernel_sizewidth_kernel_sizerGrErFr4r r!rCOs0      zCogVideoXCausalConv3d.__init__inputsrcCs|j}|j}|dkr |S|d|}|jdur*tj|jd||j|gdd}ntj|ddg|d|gdd}|d|}|S)Nrrr) rIrJ transposerLrrtodevice contiguous)r.rOrrr r r!fake_context_parallel_forwardxs  &&z3CogVideoXCausalConv3d.fake_context_parallel_forwardcCs|`d|_dSN)rL)r.r r r!"_clear_fake_context_parallel_caches z8CogVideoXCausalConv3d._clear_fake_context_parallel_cachecCs~||}||dddd|j ddf|_|j|j|j |j f}t j ||ddd}| |}|}|S)Nrr=r)modevalue) rTrVrJrSdetachclonecpurLrFrEFpadrK)r.rOinput_parallel padding_2doutput_parallelr3r r r!r-s 4 zCogVideoXCausalConv3d.forward)rrr=)r6r7r8r9r'rrstrrCrr:rTrVr-r;r r r4r!r<Cs*)r<csNeZdZdZ d dededeffdd Zdejd ejd ejfd d ZZ S)CogVideoXSpatialNorm3Da Spatially conditioned normalization as defined in https://arxiv.org/abs/2209.09002. This implementation is specific to 3D-video like data. CogVideoXSafeConv3d is used instead of nn.Conv3d to avoid OOM in CogVideoX Model. Args: f_channels (`int`): The number of channels for input to group normalization layer, and output of the spatial norm layer. zq_channels (`int`): The number of channels for the quantized vector as described in the paper. f_channels zq_channelsgroupscsFttj||ddd|_t||ddd|_t||ddd|_dS)Nư>T) num_channels num_groupsepsaffiner)rr@)r,rCnn GroupNorm norm_layerr<conv_yconv_b)r.rdrerfr4r r!rCs zCogVideoXSpatialNorm3D.__init__fzqrc Cs|jddkrm|jdddkrm|ddddddf|ddddddf}}|jdd|jdd}}|ddddddf|ddddddf}}tj||d}tj||d}tj||gdd}n tj||jddd}||} | ||||} | S)Nr r)sizer)r%r\ interpolaterrrnrorp) r.rqrrf_firstf_rest f_first_size f_rest_sizez_firstz_restnorm_fnew_fr r r!r-s 66 zCogVideoXSpatialNorm3D.forward)rc) r6r7r8r9r'rCrr:r-r;r r r4r!rbs$ rbcseZdZdZ         dd ed eed ed edededededeedeffdd Z  dde j dee j dee j de j fddZ Z S)CogVideoXResnetBlock3Da A 3D ResNet block used in the CogVideoX model. Args: in_channels (int): Number of input channels. out_channels (Optional[int], optional): Number of output channels. If None, defaults to `in_channels`. Default is None. dropout (float, optional): Dropout rate. Default is 0.0. temb_channels (int, optional): Number of time embedding channels. Default is 512. groups (int, optional): Number of groups for group normalization. Default is 32. eps (float, optional): Epsilon value for normalization layers. Default is 1e-6. non_linearity (str, optional): Activation function to use. Default is "swish". conv_shortcut (bool, optional): If True, use a convolutional shortcut. Default is False. spatial_norm_dim (Optional[int], optional): Dimension of the spatial normalization. Default is None. pad_mode (str, optional): Padding mode. Default is "first". NrcrgswishFfirstr>r?dropout temb_channelsrfrj non_linearity conv_shortcutspatial_norm_dimrBc st|p|}||_||_t||_||_| dur.tj|||d|_ tj|||d|_ nt || |d|_ t || |d|_ t ||d| d|_ |dkrStj||d|_t||_t ||d| d|_|j|jkr|jrvt ||d| d|_dSt||dddd|_dSdS) N)rhrirj)rdrerfr)r>r?rrBr) in_features out_featuresr)r>r?rr@padding)r,rCr>r?r nonlinearityuse_conv_shortcutrlrmnorm1norm2rbr<conv1Linear temb_projDropoutrconv2rr) r.r>r?rrrfrjrrrrBr4r r!rCsJ       zCogVideoXResnetBlock3D.__init__rOtembrrrcCs|}|dur |||}n||}||}||}|dur5||||dddddddf}|dur@|||}n||}||}||}||}|j|jkr_| |}||}|SrU) rrrrrrrr>r?r)r.rOrrr hidden_statesr r r!r-s$   *      zCogVideoXResnetBlock3D.forward) NrrrcrgrFNrNN)r6r7r8r9r'rfloatraboolrCrr:r-r;r r r4r!r~sV   >r~cseZdZdZdZ        d d ed ed edededededededededeffdd Z  d!de j de e j de e j de j fddZ ZS)"CogVideoXDownBlock3Da1 A downsampling block used in the CogVideoX model. Args: in_channels (int): Number of input channels. out_channels (int): Number of output channels. temb_channels (int): Number of time embedding channels. dropout (float, optional): Dropout rate. Default is 0.0. num_layers (int, optional): Number of layers in the block. Default is 1. resnet_eps (float, optional): Epsilon value for the ResNet layers. Default is 1e-6. resnet_act_fn (str, optional): Activation function for the ResNet layers. Default is "swish". resnet_groups (int, optional): Number of groups for group normalization in the ResNet layers. Default is 32. add_downsample (bool, optional): If True, add a downsampling layer at the end of the block. Default is True. downsample_padding (int, optional): Padding for the downsampling layer. Default is 0. compress_time (bool, optional): If True, apply temporal compression. Default is False. pad_mode (str, optional): Padding mode. Default is "first". TrrrgrrcrFrr>r?rr num_layers resnet_eps resnet_act_fn resnet_groupsadd_downsampledownsample_padding compress_timerBc  stg} t|D]}|dkr|n|}| t|||||||| dq t| |_d|_| r=tt ||| | dg|_d|_ dS)Nr)r>r?rrrfrjrrBrrF) r,rCr)r+r~rl ModuleListresnets downsamplersr gradient_checkpointing)r.r>r?rrrrrrrrrrBrr in_channelr4r r!rCHs4    zCogVideoXDownBlock3D.__init__NrrrrrcCsh|jD]}|jr|jrdd}tjj|||||}q||||}q|jdur2|jD]}||}q+|S)Ncfdd}|S)Nc|SrUr rOmoduler r!create_forwardzSCogVideoXDownBlock3D.forward..create_custom_forward..create_forwardr rrr rr!create_custom_forward z;CogVideoXDownBlock3D.forward..create_custom_forward)rtrainingrrutils checkpointr)r.rrrrresnetr downsamplerr r r!r-ws      zCogVideoXDownBlock3D.forward) rrrgrrcTrFrr)r6r7r8r9 _supports_gradient_checkpointingr'rrarrCrr:rr-r;r r r4r!r3s`     2rcseZdZdZdZ       dd ed ed ed ededededeedeffdd Z  dde j dee j dee j de j fddZ Z S)CogVideoXMidBlock3Da/ A middle block used in the CogVideoX model. Args: in_channels (int): Number of input channels. temb_channels (int): Number of time embedding channels. dropout (float, optional): Dropout rate. Default is 0.0. num_layers (int, optional): Number of layers in the block. Default is 1. resnet_eps (float, optional): Epsilon value for the ResNet layers. Default is 1e-6. resnet_act_fn (str, optional): Activation function for the ResNet layers. Default is "swish". resnet_groups (int, optional): Number of groups for group normalization in the ResNet layers. Default is 32. spatial_norm_dim (Optional[int], optional): Dimension of the spatial normalization. Default is None. pad_mode (str, optional): Padding mode. Default is "first". TrrrgrrcNrr>rrrrrrrrBc sRtg} t|D]} | t||||||||| d q t| |_d|_dS)N) r>r?rrrfrjrrrBF) r,rCr)r+r~rlrrr) r.r>rrrrrrrrBr_r4r r!rCs$   zCogVideoXMidBlock3D.__init__rrrrrcCsJ|jD]}|jr|jrdd}tjj|||||}q||||}q|S)Ncr)NcrrUr rrr r!rrzRCogVideoXMidBlock3D.forward..create_custom_forward..create_forwardr rr rr!rrz:CogVideoXMidBlock3D.forward..create_custom_forward)rrrrrr)r.rrrrrrr r r!r-s   zCogVideoXMidBlock3D.forward)rrrgrrcNrr)r6r7r8r9rr'rrarrCrr:r-r;r r r4r!rsP  $rcseZdZdZ         d!d ed ed ededededededededededeffdd Z  d"de j de e j de e j de j fdd Z Z S)#CogVideoXUpBlock3Da An upsampling block used in the CogVideoX model. Args: in_channels (int): Number of input channels. out_channels (int): Number of output channels. temb_channels (int): Number of time embedding channels. dropout (float, optional): Dropout rate. Default is 0.0. num_layers (int, optional): Number of layers in the block. Default is 1. resnet_eps (float, optional): Epsilon value for the ResNet layers. Default is 1e-6. resnet_act_fn (str, optional): Activation function for the ResNet layers. Default is "swish". resnet_groups (int, optional): Number of groups for group normalization in the ResNet layers. Default is 32. spatial_norm_dim (int, optional): Dimension of the spatial normalization. Default is 16. add_upsample (bool, optional): If True, add an upsampling layer at the end of the block. Default is True. upsample_padding (int, optional): Padding for the upsampling layer. Default is 1. compress_time (bool, optional): If True, apply temporal compression. Default is False. pad_mode (str, optional): Padding mode. Default is "first". rrrgrrcTFrr>r?rrrrrrr add_upsampleupsample_paddingrrBcstg}t|D]}|dkr|n|}|t|||||||| | d q t||_d|_| r>tt ||| | dg|_d|_ dS)Nr) r>r?rrrfrjrrrBrF) r,rCr)r+r~rlrr upsamplersrr)r.r>r?rrrrrrrrrrrBrrrr4r r!rCs6    zCogVideoXUpBlock3D.__init__NrrrrrcCsh|jD]}|jr|jrdd}tjj|||||}q||||}q|jdur2|jD]}||}q+|S)z1Forward method of the `CogVideoXUpBlock3D` class.cr)NcrrUr rrr r!r.rzQCogVideoXUpBlock3D.forward..create_custom_forward..create_forwardr rr rr!r-rz9CogVideoXUpBlock3D.forward..create_custom_forwardN)rrrrrrr)r.rrrrrr upsamplerr r r!r-#s      zCogVideoXUpBlock3D.forward) rrrgrrcrTrFrr)r6r7r8r9r'rrarrCrr:rr-r;r r r4r!rsd     4rcseZdZdZdZ        d!d ededeedfdeedfdededededededeffdd Z d"de j de e j de j fdd Z ZS)#CogVideoXEncoder3Da The `CogVideoXEncoder3D` layer of a variational autoencoder that encodes its input into a latent representation. Args: in_channels (`int`, *optional*, defaults to 3): The number of input channels. out_channels (`int`, *optional*, defaults to 3): The number of output channels. down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`): The types of down blocks to use. See `~diffusers.models.unet_2d_blocks.get_down_block` for available options. block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): The number of output channels for each block. layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. norm_num_groups (`int`, *optional*, defaults to 32): The number of groups for normalization. act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. See `~diffusers.models.activations.get_activation` for available options. double_z (`bool`, *optional*, defaults to `True`): Whether to double the number of output channels for the last block. Trrrrrrrrsilurgrcrrr>r?down_block_types.block_out_channelslayers_per_blockact_fnnorm_epsnorm_num_groupsrrBtemporal_compression_ratioc  sttt| } t||dd| d|_tg|_ |d} t |D]4\}}| }||} |t |dk}|| k}|dkrOt || d| ||||| |d }nt d|j |q%t|dd| d |||| d |_tj||dd d |_t|_t|dd |d| d|_d |_dS)NrrrrBrr) r>r?rrrrrrrrzEInvalid `down_block_type` encountered. Must be `CogVideoXDownBlock3D`r )r>rrrrrrrBrg)rjF)r,rCr'nplog2r<conv_inrlr down_blocks enumerater*r ValueErrorr+r mid_blockrmnorm_outSiLUconv_actconv_outr)r.r>r?rrrrrrrrBrtemporal_compress_leveloutput_channelrdown_block_type input_channelis_final_blockr down_blockr4r r!rCZsR     zCogVideoXEncoder3D.__init__NsamplerrcCs||}|jr/|jr/dd}|jD]}tjj||||d}qtjj||j||d}n|jD]}|||d}q2|||d}||}| |}| |}|S)z5The forward method of the `CogVideoXEncoder3D` class.cr)NcrrUr rrr r!custom_forwardrzQCogVideoXEncoder3D.forward..create_custom_forward..custom_forwardr rrr rr!rrz9CogVideoXEncoder3D.forward..create_custom_forwardN) rrrrrrrrrrr)r.rrrrrr r r!r-s"        zCogVideoXEncoder3D.forward) rrrrrrrgrcrrrrUr6r7r8r9rr'rrarrCrr:rr-r;r r r4r!r@sL     *IrcseZdZdZdZ        d!d ededeedfdeedfdededededededeffdd Z d"de j de e j de j fdd Z ZS)#CogVideoXDecoder3Da The `CogVideoXDecoder3D` layer of a variational autoencoder that decodes its latent representation into an output sample. Args: in_channels (`int`, *optional*, defaults to 3): The number of input channels. out_channels (`int`, *optional*, defaults to 3): The number of output channels. up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`): The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options. block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): The number of output channels for each block. layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. norm_num_groups (`int`, *optional*, defaults to 32): The number of groups for normalization. act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. See `~diffusers.models.activations.get_activation` for available options. norm_type (`str`, *optional*, defaults to `"group"`): The normalization type to use. Can be either `"group"` or `"spatial"`. Trrrrrrrrrgrcrrrr>r?up_block_types.rrrrrrrBrc s(ttt|} t|| dd| d|_t| ddd||||| d|_t g|_ | d} t t | }t|D]:\}}| }| |} |t|dk}||k}|dkrjt|| d| |d||||| || d }| }ntd |j |q:t| d ||d |_t|_t| d |d| d|_d |_dS) Nrrrr )r>rrrrrrrBrr) r>r?rrrrrrrrrrBzAInvalid `up_block_type` encountered. Must be `CogVideoXUpBlock3D`r)rfF)r,rClistreversedr<rrrrlr up_blocksr'rrrr*rrr+rbrrrrr)r.r>r?rrrrrrrrBrreversed_block_out_channelsrrr up_block_typeprev_output_channelrrup_blockr4r r!rCs^       zCogVideoXDecoder3D.__init__NrrrcCs||}|jr/|jr/dd}tjj||j|||}|jD]}tjj|||||}qn||||}|jD]}||||}q9|||}| |}| |}|S)z5The forward method of the `CogVideoXDecoder3D` class.cr)NcrrUr rrr r!r9rzQCogVideoXDecoder3D.forward..create_custom_forward..custom_forwardr rr rr!r8rz9CogVideoXDecoder3D.forward..create_custom_forward) rrrrrrrrrrr)r.rrrrrr r r!r-2s$        zCogVideoXDecoder3D.forward) rrrrrrrgrcrrrrUrr r r4r!rsL     *Prc(sdeZdZdZdZdgZe              d7dededee dee deededede de dede dede de e de ee d e ee d!e d"e d#e f&fd$d% Z d8d&d'Zd(d)Ze d9d*ejd+e d,eeeeffd-d.Zed9d/ejd+e d,eeejffd0d1Z   d:d2ejd3e d+e d4e ejd,eejejff d5d6ZZS);AutoencoderKLCogVideoXa A VAE model with KL loss for encoding images into latents and decoding latent representations into images. Used in [CogVideoX](https://github.com/THUDM/CogVideo). This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented for all models (such as downloading or saving). Parameters: in_channels (int, *optional*, defaults to 3): Number of channels in the input image. out_channels (int, *optional*, defaults to 3): Number of channels in the output. down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): Tuple of downsample block types. up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): Tuple of upsample block types. block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`): Tuple of block output channels. act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. sample_size (`int`, *optional*, defaults to `32`): Sample input size. scaling_factor (`float`, *optional*, defaults to 0.18215): The component-wise standard deviation of the trained latent space computed using the first batch of the training set. This is used to scale the latent space to have unit variance when training the diffusion model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1 / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper. force_upcast (`bool`, *optional*, default to `True`): If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE can be fine-tuned / trained to a lower range without loosing too much precision in which case `force_upcast` can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix Tr~rrrrrrrgrcrrYi1p?NFr>r?rrrlatent_channelsrrrrr sample_sizescaling_factor shift_factor latents_mean latents_std force_upcastuse_quant_convuse_post_quant_convc stt||||||| | | d |_t||||||| | | d |_|r-td|d|dnd|_|r8t||dnd|_d|_ d|_ |j j |_ t|j j ttfrU|j j dn|j j } t| dt|j jd|_d|_dS)N) r>r?rrrrrrr) r>r?rrrrrrrr rFrg?)r,rCrencoderrdecoderr quant_convpost_quant_conv use_slicing use_tilingconfigrtile_sample_min_sizerDrtupler'r*rtile_latent_min_sizetile_overlap_factor)r.r>r?rrrrrrrrrrrrrrrrrr4r r!rCzsD !    zAutoencoderKLCogVideoX.__init__cCst|ttfr ||_dSdSrU)rDrrr)r.rrXr r r!_set_gradient_checkpointings z2AutoencoderKLCogVideoX._set_gradient_checkpointingcCs8|D]\}}t|trtd||qdS)Nz0Clearing fake Context Parallel cache for layer: ) named_modulesrDr<loggerdebugrV)r.namerr r r!!clear_fake_context_parallel_caches  z8AutoencoderKLCogVideoX.clear_fake_context_parallel_cachex return_dictrcCs:||}|jdur||}t|}|s|fSt|dS)a Encode a batch of images into latents. Args: x (`torch.Tensor`): Input batch of images. return_dict (`bool`, *optional*, defaults to `True`): Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. Returns: The latent representations of the encoded images. If `return_dict` is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. N) latent_dist)rrrr)r.r r h posteriorr r r!encodes    zAutoencoderKLCogVideoX.encodezcCs2|jdur ||}||}|s|fSt|dS)a Decode a batch of images. Args: z (`torch.Tensor`): Input batch of latent vectors. return_dict (`bool`, *optional*, defaults to `True`): Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. Returns: [`~models.vae.DecoderOutput`] or `tuple`: If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is returned. N)r)rrr)r.rr decr r r!decodes    zAutoencoderKLCogVideoX.decodersample_posterior generatorc CsB|}||j}|r|j|d}n|}||}|s|fS|S)N)r)rr rrWr) r.rrr rr rrrr r r!r-s  zAutoencoderKLCogVideoX.forward)rrrrrrrrrgrcrrrNNNTFF)F)T)FTN)r6r7r8r9r_no_split_modulesrr'rrarrrrCrr r rr:rrrr FloatTensorrr Generatorr-r;r r r4r!rWs    G &r)1typingrrrnumpyrrtorch.nnrltorch.nn.functional functionalr\configuration_utilsrrloaders.single_file_modelrrr utils.accelerate_utilsr activationsr downsamplingr modeling_outputsrmodeling_utilsr upsamplingrvaerr get_loggerr6rConv3drModuler<rbr~rrrrrrr r r r!s8          W)p`Kb