o Gi@sRddlZddlZddlmZddlmmZddlmZm Z ddl m Z ddl m Z ddlmZddlmZdd lmZd d lmZmZmZe eZGd d d ejZGdddejZGdddejZGdddejZ GdddejZ!GdddejZ"GdddejZ#GdddejZ$GdddejZ%GdddeeeZ&dS) N) ConfigMixinregister_to_config)logging)apply_forward_hook)get_activation)AutoencoderKLOutput) ModelMixin)AutoencoderMixin DecoderOutputDiagonalGaussianDistributioncseZdZ       ddededeeedfBd eeedfBd eeedfBd eeedfBd ed edeffdd ZddZde j de j ffdd Z Z S)EasyAnimateCausalConv3drr Tzeros in_channels out_channels kernel_size.stridepaddingdilationgroupsbias padding_modec  s~t|tr|n|fd}t|dksJd|dt|tr!|n|fd}t|dks4Jd|dt|tr;|n|fd}t|dksNJd|d|\} } } |\|_} }|\}}}| d|}|durt| d|d| d}t| d|d|d}nt|tr|}}ntsJ||_t| d|d|d|_ d|_ t j |||||d||f||| d dS) Nrz#Kernel size must be a 3-tuple, got z instead.zStride must be a 3-tuple, got z Dilation must be a 3-tuple, got r rr) rrrrrrrrr) isinstancetuplelent_stridemathceilintNotImplementedErrortemporal_paddingtemporal_padding_origin prev_featuressuper__init__)selfrrrrrrrrrt_ksh_ksw_ksh_stridew_stride t_dilation h_dilation w_dilationt_padh_padw_pad __class__g/home/ubuntu/.local/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_magvit.pyr&#s<          z EasyAnimateCausalConv3d.__init__cC|`d|_dSNr$r'r5r5r6_clear_conv_cache] z)EasyAnimateCausalConv3d._clear_conv_cache hidden_statesreturncs|j}|jdurttj|dddd|jdfdd}|j|d}||dddd|j df|_|d}g}d}||jd|krnt |dddd|||jdf}||j 7}| |||jd|ksEt |dS|j dkrt j|jdddd|jd df|gdd}n t j|j|gdd}|j|d}||dddd|j df|_|d}g}d}||jd|krt |dddd|||jdf}||j 7}| |||jd|kst |dS)Nr replicate)padmode)dtyperr dim)rBr$Fr@r"tor;clonesizer%forwardrappendtorchconcat)r'r=rB num_framesoutputsioutr3r5r6rIasH  $ ,    ( $ ,   zEasyAnimateCausalConv3d.forward)rr r r r Tr) __name__ __module__ __qualname__r rboolstrr&r;rKTensorrI __classcell__r5r5r3r6r"s:  :"rcsbeZdZ      ddeded ed ed ed ed edeffdd Zdej dej fddZ Z S)EasyAnimateResidualBlock3Dsilu ư>T?rr non_linearitynorm_num_groupsnorm_epsspatial_group_normdropoutoutput_scale_factorc st||_tj|||dd|_t||_t||dd|_ tj|||dd|_ t ||_ t||dd|_ ||krEtj||dd|_nt|_||_dS)NT) num_groups num_channelsepsaffinerrr )r%r&rcnn GroupNormnorm1r nonlinearityrconv1norm2Dropoutrbconv2Conv3dshortcutIdentityra) r'rrr^r_r`rarbrcr3r5r6r&s"     z#EasyAnimateResidualBlock3D.__init__r=r>cCs||}|jr/|d}|ddddddd}||}|d|dfddddd}n||}||}||}|jrh|d}|ddddddd}| |}|d|dfddddd}n| |}||}| |}| |}|||j SNrrr r) rrrarHpermuteflattenrk unflattenrlrmrnrbrprc)r'r=rr batch_sizer5r5r6rIs.              z"EasyAnimateResidualBlock3D.forward)rYrZr[Tr\r] rQrRrSr rUfloatrTr&rKrVrIrWr5r5r3r6rXs2 $rXc sFeZdZd dedededeffdd Zd ejd ejfd d ZZ S)EasyAnimateDownsampler3Drrrrrrrrcs"tt||||dd|_dS)Nr)rrrrr)r%r&rconv)r'rrrrr3r5r6r&s   z!EasyAnimateDownsampler3D.__init__r=r>cCst|d}||}|S)N)rr rr )rEr@rr'r=r5r5r6rIs  z EasyAnimateDownsampler3D.forward)rr~) rQrRrSr rr&rKrVrIrWr5r5r3r6r}s r}c sXeZdZ   ddededededef fd d Zd d Zd ejdejfddZ Z S)EasyAnimateUpsampler3DrFTrrrtemporal_upsampleracs8t|p|}||_||_t|||d|_d|_dS)N)rrr)r%r&rrarrr$)r'rrrrrar3r5r6r&s  zEasyAnimateUpsampler3D.__init__cCr7r8r9r:r5r5r6r;r<z(EasyAnimateUpsampler3D._clear_conv_cacher=r>cCsRtj|ddd}||}|jr'|jdur||_|Stj|d|js#dndd}|S)Nr rrnearest) scale_factorrA)rr r trilinear)rE interpolaterrr$rarr5r5r6rIs   zEasyAnimateUpsampler3D.forward)rFT) rQrRrSr rTr&r;rKrVrIrWr5r5r3r6rs"rcsteZdZ         dded ed ed ed ed edededededeffdd Zdej dej fddZ Z S)EasyAnimateDownBlock3Dr rYrZr[Tr\r]rr num_layersact_fnr_r`rarbrcadd_downsampleadd_temporal_downsamplec sttg|_t|D]} | dkr|n|}|jt|||||||| dq| r?| r?t||ddd|_ d|_ d|_ dS| rT| sTt||ddd|_ d|_ d|_ dSd|_ d|_ d|_ dS) Nrrrr^r_r`rarbrcrr~)rrrrr ) r%r&ri ModuleListconvsrangerJrXr} downsamplerspatial_downsample_factortemporal_downsample_factor) r'rrrrr_r`rarbrcrrrOr3r5r6r&s6       zEasyAnimateDownBlock3D.__init__r=r>cC,|jD]}||}q|jdur||}|Sr8)rrr'r=rr5r5r6rI=    zEasyAnimateDownBlock3D.forward) r rYrZr[Tr\r]TTr{r5r5r3r6rsD    -rcsteZdZ         dd ed ed ed ed ededededededeffdd Zdej dej fddZ Z S)EasyAnimateUpBlock3dr rYrZr[Fr\r]Trrrrr_r`rarbrc add_upsampleadd_temporal_upsamplec sxttg|_t|D]} | dkr|n|}|jt|||||||| dq| r7t||| |d|_ dSd|_ dS)Nrr)rra) r%r&rirrrrJrXr upsampler) r'rrrrr_r`rarbrcrrrOr3r5r6r&Fs0     zEasyAnimateUpBlock3d.__init__r=r>cCrr8)rrrr5r5r6rIprzEasyAnimateUpBlock3d.forward) r rYrZr[Fr\r]TTr{r5r5r3r6rEsD    *rcsdeZdZ       dded ed ed ed ed ededeffdd Zdej dej fddZ Z S)EasyAnimateMidBlock3dr rYrZr[Tr\r]rrrr_r`rarbrcc s|t|dur |nt|dd}tt||||||||dg|_t|dD]} |jt||||||||dq)dS)NrurZrr ) r%r&minrirrXrrrJ) r'rrrr_r`rarbrc_r3r5r6r&ys: zEasyAnimateMidBlock3d.__init__r=r>cCs.|jd|}|jddD]}||}q|S)Nrr )r)r'r=resnetr5r5r6rIs zEasyAnimateMidBlock3d.forward)r rYrZr[Tr\r]r{r5r5r3r6rxs4 ,rcseZdZdZdZdddgdddd dd f d ed ed eedfdeedfdededededeffdd Z de j de j fddZ Z S)EasyAnimateEncoderzp Causal encoder for 3D video-like data used in [EasyAnimate](https://huggingface.co/papers/2405.18991). TrSpatialDownBlock3DSpatialTemporalDownBlock3DrrrrrZrYFrrdown_block_types.block_out_channelslayers_per_blockr_rdouble_zrac  s8tt||ddd|_tg|_|d} t|D]E\} } | } || } | t|dk}| dkrBt | | |||d| | dd }n| d krUt | | |||d| | d d }nt d | |j |qt |d ||| |dddd |_ | |_tj|d |dd|_t||_|rd|n|}t|d |dd|_d|_dS)Nrrrhr rr[F) rrrrr_r`rarrrTUnknown up block type: rv)rrrrar_r`rbrcrerdrfr)r%r&rconv_inrir down_blocks enumeraterr ValueErrorrJr mid_blockrarj conv_norm_outrconv_actconv_outgradient_checkpointing)r'rrrrrr_rrraoutput_channelsrOdown_block_typeinput_channelsis_final_block down_blockconv_out_channelsr3r5r6r&sl       zEasyAnimateEncoder.__init__r=r>cCs||}|jD]}tr|jr|||}q||}q||}|jrL|d}| ddddd dd}| |}| d|df ddddd}n| |}| |}||}|Srt)rrrKis_grad_enabledr_gradient_checkpointing_funcrrarHrwrxrryrr)r'r=rrzr5r5r6rIs          zEasyAnimateEncoder.forwardrQrRrS__doc__ _supports_gradient_checkpointingr rrUrTr&rKrVrIrWr5r5r3r6rs@     TrcseZdZdZdZdddgdddd d fd ed ed eedfdeedfdedededeffdd Z de j de j fddZ Z S)EasyAnimateDecoderzp Causal decoder for 3D video-like data used in [EasyAnimate](https://huggingface.co/papers/2405.18991). TrrSpatialUpBlock3DSpatialTemporalUpBlock3DrrrrrZrYFrrup_block_types.rrr_rrac  s:tt||ddd|_t|d|||dddd|_tg|_t t |} | d} t |D]I\} } | } | | } | t |dk}| dkrXt | | |d||d|| d d }n| d krmt | | |d||d|| d d }ntd | |j|q1||_tj|d|dd|_t||_t|d|dd|_d |_dS)Nrvrrhr[rr )rrrr_r`rbrcrF) rrrrr_r`rarrrTrr)r%r&rrrrrir up_blockslistreversedrrrrrJrarjrrrrr)r'rrrrrr_rrareversed_block_out_channelsrrO up_block_typerrup_blockr3r5r6r&'sj      zEasyAnimateDecoder.__init__r=r>cCs||}tr|jr||j|}n||}|jD]}tr,|jr,|||}q||}q|jr[|d}| ddddd dd}| |}| d|df ddddd}n| |}| |}||}|Srt)rrKrrrrrrarHrwrxrryrr)r'r=rrzr5r5r6rI{s&          zEasyAnimateDecoder.forwardrr5r5r3r6r s:     Trcs(eZdZdZdZedddgdgdgddd d d df d ed ededeedfdeedfdeedfdededede de ffdd Z ddZ      dBdedBdedBdedBd e dBd!e dBd"e dBd#dfd$d%Z e dCd&ejd'e d#eeeBfd(d)Ze dCd&ejd'e d#eeeBfd*d+ZdCd,ejd'e d#eejBfd-d.ZedCd,ejd'e d#eejBfd/d0Zd1ejd2ejd3ed#ejfd4d5Zd1ejd2ejd3ed#ejfd6d7ZdCd&ejd'e d#efd8d9ZdCd,ejd'e d#eejBfd:d;Z <  dDd=ejd>e d'e d?ejdBd#eejBf d@dAZZS)EAutoencoderKLMagvitaq A VAE model with KL loss for encoding images into latents and decoding latent representations into images. This model is used in [EasyAnimate](https://huggingface.co/papers/2405.18991). This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented for all models (such as downloading or saving). TrrrrrrYrZg?rlatent_channelsrr.rrrrr_scaling_factorrac stt|||||| |d| d |_t|||||| || d|_tjd|d|dd|_tj||dd|_ dt |d|_ dt |d|_ d|_ d|_d|_d|_d|_d|_d |_d |_d|_d |_d |_d |_dS) NT) rrrrrr_rrra)rrrrrr_rrarr rhFrurir)r%r&rencoderrdecoderrirq quant_convpost_quant_convrspatial_compression_ratiotemporal_compression_ratio use_slicing use_tilinguse_framewise_encodinguse_framewise_decodingnum_sample_frames_batch_sizenum_latent_frames_batch_sizetile_sample_min_heighttile_sample_min_widthtile_sample_min_num_framestile_sample_stride_heighttile_sample_stride_widthtile_sample_stride_num_frames) r'rrrrrrrrr_rrar3r5r6r&sL    zAutoencoderKLMagvit.__init__cCs:|D]\}}t|tr|t|tr|qdSr8) named_modulesrrr;r)r'namemoduler5r5r6r;s  z%AutoencoderKLMagvit._clear_conv_cacheNrrrrrrr>cCs^d|_d|_d|_|p |j|_|p|j|_|p|j|_|p|j|_|p%|j|_|p+|j|_dS)aX Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow processing larger images. Args: tile_sample_min_height (`int`, *optional*): The minimum height required for a sample to be separated into tiles across the height dimension. tile_sample_min_width (`int`, *optional*): The minimum width required for a sample to be separated into tiles across the width dimension. tile_sample_stride_height (`int`, *optional*): The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are no tiling artifacts produced across the height dimension. tile_sample_stride_width (`int`, *optional*): The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling artifacts produced across the width dimension. TN) rrrrrrrrr)r'rrrrrrr5r5r6 enable_tilings     z!AutoencoderKLMagvit.enable_tilingx return_dictc Cs|jr|jd|jks|jd|jkr|j||dS||ddddddddddf}|g}td|jd|jD] }||dddd|||jddddf}||q=t j |dd}| |}| |S)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. rvrNr rrC) rshaperr tiled_encoderrrrJrKcatrr;)r'rr first_frameshrO next_framesmomentsr5r5r6_encode's&,2  zAutoencoderKLMagvit._encodecs^jr|jddkrfdd|dD}t|}n|}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 videos. If `return_dict` is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. rr csg|]}|qSr5)r).0x_slicer:r5r6 Vsz.AutoencoderKLMagvit.encode..) latent_dist)rrsplitrKrrrr )r'rrencoded_slicesr posteriorr5r:r6encodeEs   zAutoencoderKLMagvit.encodezc Cs |j\}}}}}|j|j}|j|j} |jr,|jd|ks%|jd| kr,|j||dS||}||ddddddddddf} | g} td|jd|j D] } ||dddd| | |j ddddf} | | qTt j | dd} |s| fSt | dS)Nrvrrr rrCsample)rrrrr tiled_decoderrrrrJrKrr )r'rrrzrerMheightwidthtile_latent_min_heighttile_latent_min_widthrdecrOrr5r5r6_decodeas  " ,2  zAutoencoderKLMagvit._decodecs`jr|jddkrfdd|dD}t|}n|j}|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. rr csg|]}|jqSr5)rr)rz_slicer:r5r6rsz.AutoencoderKLMagvit.decode..r) rrrrKrrrr;r )r'rrdecoded_slicesdecodedr5r:r6decode{s   zAutoencoderKLMagvit.decodeab blend_extentc Cst|jd|jd|}t|D]@}|dddddd| |ddfd|||dddddd|ddf|||dddddd|ddf<q|S)Nrr rrr)r'r r r yr5r5r6blend_v  R&zAutoencoderKLMagvit.blend_vcCst|jd|jd|}t|D]@}|dddddddd| |fd|||dddddddd|f|||dddddddd|f<q|S)Nrur r )r'r r r rr5r5r6blend_hrzAutoencoderKLMagvit.blend_hc Cs\|j\}}}}}||j}||j} |j|j} |j|j} |j|j} |j|j} | | }| | }g}td||jD]}g}td||jD]q}|dddddd|||j|||jf}||ddddddddddf}|g}td||jD] }||dddd|||jddddf}| |qt j |dd}| |}| | |qF| |q;g}t|D]O\}}g}t|D]:\}}|dkr|||d|||}|dkr|||d||}| |ddddddd|d| fq| t j |ddqt j |ddddddddd|d| f}|S)Nrr rrCrur)rrrrrrrrrrJrKrrr;rrr)r'rrrzrerMrr latent_height latent_widthrrtile_latent_stride_heighttile_latent_stride_width blend_height blend_widthrowsrOrowjtilertile_hkr result_rows result_rowrr5r5r6rsV        ,2    .0z AutoencoderKLMagvit.tiled_encodec Csr|j\}}}}}||j}||j} |j|j} |j|j} |j|j} |j|j} |j|j}|j|j}g}td|| D]}g}td|| D]o}|dddddd||| ||| f}||}||ddddddddddf}|g}td||j D] }||dddd|||j ddddf}| |qt j |dd}| | |qH| |q>g}t|D]Q\}}g}t|D]<\}}|dkr|||d|||}|dkr|||d||}| |ddddddd|jd|jfq| t j |ddqt j |ddddddddd|d| f}|s4|fSt|dS)Nrr rrCrurr)rrrrrrrrrrrJrKrr;rrrr )r'rrrzrerMrr sample_height sample_widthrrrrrrrrOrrrrtile_decrrrrrrr5r5r6rsZ           ,2   20 z AutoencoderKLMagvit.tiled_decodeFrsample_posterior generatorc CsJ|}||j}|r|j|d}n|}||j}|s |fSt|dS)aa Args: sample (`torch.Tensor`): Input sample. sample_posterior (`bool`, *optional*, defaults to `False`): Whether to sample from the posterior. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`DecoderOutput`] instead of a plain tuple. )r#r)rrrrArr ) r'rr"rr#rrrrr5r5r6rIs   zAutoencoderKLMagvit.forward)NNNNNN)T)FTN)rQrRrSrrrr rrUr|rTr&r;rrrKrVr rrrr rrrrrr GeneratorrIrWr5r5r3r6rs   V  $    "   4@r)'rrKtorch.nnritorch.nn.functional functionalrEconfiguration_utilsrrutilsrutils.accelerate_utilsr activationsrmodeling_outputsr modeling_utilsr vaer r r get_loggerrQloggerrqrModulerXr}rrrrrrrr5r5r5r6s,       qF(634ty