o Gi @sddlZddlZddlmZddlmmZddlm Z m Z ddl m Z ddl mZddlmZddlmZdd lmZdd lmZd d lmZmZmZe eZ d)d ededejdej dedej!f ddZ"Gdddej#Z$Gdddej#Z%Gdddej#Z&Gdddej#Z'Gdddej#Z(Gdd d ej#Z)Gd!d"d"ej#Z*Gd#d$d$ej#Z+Gd%d&d&ej#Z,Gd'd(d(eee Z-dS)*N) ConfigMixinregister_to_config)logging)apply_forward_hook)get_activation) Attention)AutoencoderKLOutput) ModelMixin)AutoencoderMixin DecoderOutputDiagonalGaussianDistribution num_frames height_widthdtypedevice batch_sizereturnc Csvtjd|dtj|d}||}tj||dd\}}t||kdtd j|d} |dur9| d |dd} | S) Nr )rrxy)indexingrinf)r) torcharangeint32repeat_interleavemeshgridwherefloatto unsqueezeexpand) rrrrrindicesindices_blocksxymaskr)n/home/ubuntu/.local/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.pyprepare_causal_attention_mask!s  r+cseZdZ      ddededeeeeefBd eeeeefBd eeeeefBd eeeeefBd ed eddffdd Zdej dej fddZ Z S)HunyuanVideoCausalConv3drr rT replicate in_channels out_channels kernel_sizestridepaddingdilationbiaspad_moderNc stt|tr|||fn|}||_|dd|dd|dd|dd|dddf|_tj|||||||d|_dS)Nrrr r4) super__init__ isinstanceintr5time_causal_paddingnnConv3dconv) selfr.r/r0r1r2r3r4r5 __class__r)r*r8/s       z!HunyuanVideoCausalConv3d.__init__ hidden_statescCstj||j|jd}||S)N)mode)Fpadr;r5r>r?rBr)r)r*forwardJs z HunyuanVideoCausalConv3d.forward)rr rr Tr-) __name__ __module__ __qualname__r:tupleboolstrr8rTensorrG __classcell__r)r)r@r*r,.s6  r,csjeZdZ     ddededBded ed ed eeeefd dffd d Zdej d ej fddZ Z S)HunyuanVideoUpsampleCausal3DNrr Trrrr.r/r0r1r4upsample_factorrcs0t|p|}||_t|||||d|_dSNr6)r7r8rRr,r>)r?r.r/r0r1r4rRr@r)r*r8Ps z%HunyuanVideoUpsampleCausal3D.__init__rBcCs|d}|jd|dfdd\}}tj|d|jddddd}|dkr@|}tj||jdd}tj ||fdd}n|}| |}|S)Nrr dimnearest) scale_factorrC) sizesplitrD interpolatesqueezerRr" contiguousrcatr>)r?rBr first_frame other_framesr)r)r*rG`s  z$HunyuanVideoUpsampleCausal3D.forward)Nrr TrQ) rHrIrJr:rLrKr r8rrNrGrOr)r)r@r*rPOs, rPc s\eZdZ     ddededBded ed ed df fd d Zdejd ejfddZZ S)HunyuanVideoDownsampleCausal3DNr rTrchannelsr/r2r0r4rcs,t|p|}t||||||d|_dSrS)r7r8r,r>)r?rar/r2r0r4r1r@r)r*r8zs z'HunyuanVideoDownsampleCausal3D.__init__rBcCs||}|SN)r>rFr)r)r*rGs z&HunyuanVideoDownsampleCausal3D.forward)Nr rTr) rHrIrJr:rLr8rrNrGrOr)r)r@r*r`ys(r`cs`eZdZ     ddededBded ed ed ed dffd d Zdejd ejfddZ Z S)HunyuanVideoResnetBlockCausal3DN ư>swishr.r/dropoutgroupseps non_linearityrcst|p|}t||_tj|||dd|_t||ddd|_tj|||dd|_ t ||_ t||ddd|_ d|_ ||krLt||ddd|_ dSdS)NT)rjaffinerr r)r7r8r nonlinearityr< GroupNormnorm1r,conv1norm2Dropoutrhconv2 conv_shortcut)r?r.r/rhrirjrkr@r)r*r8s   z(HunyuanVideoResnetBlockCausal3D.__init__rBcCsr|}|}||}||}||}||}||}||}||}|jdur3||}||}|Srb)r\rormrprqrhrsrt)r?rBresidualr)r)r*rGs         z'HunyuanVideoResnetBlockCausal3D.forward)Nrdrerfrg) rHrIrJr:r rMr8rrNrGrOr)r)r@r*rcs,rccsheZdZ       ddeded ed ed ed ed ededdffdd Zdej dej fddZ Z S)HunyuanVideoMidBlock3Drdr rfrgreTr.rh num_layers resnet_eps resnet_act_fn resnet_groups add_attentionattention_head_dimrNc st|dur |nt|dd}||_t||||||dg} g} t|D]*} |jr>| t||||||ddddd n| d| t||||||dq&t | |_ t | |_ d|_ dS)Nrer.r/rjrirhrkT)headsdim_headrjnorm_num_groupsresidual_connectionr4upcast_softmax_from_deprecated_attn_blockF) r7r8minr{rcrangeappendr r< ModuleList attentionsresnetsgradient_checkpointing) r?r.rhrwrxryrzr{r|rr_r@r)r*r8sT     zHunyuanVideoMidBlock3D.__init__rBc Csztrd|jrd||jd|}t|j|jddD]F\}}|dur[|j\}}}}}|ddddd dd}t ||||j |j |d} ||| d}| d|||fddddd}|||}q|S|jd|}t|j|jddD]D\}}|dur|j\}}}}}|ddddd dd}t ||||j |j |d} ||| d}| d|||fddddd}||}qv|S)Nrr rrr})r)attention_mask)ris_grad_enabledr_gradient_checkpointing_funcrziprshapepermuteflattenr+rr unflatten) r?rBattnresnetr num_channelsrheightwidthrr)r)r*rGs2     zHunyuanVideoMidBlock3D.forward)rdr rfrgreTr rHrIrJr:r rMrLr8rrNrGrOr)r)r@r*rvs8  >rvcsreZdZ        dded ed ed ed ed edededededdffdd Zdej dej fddZ Z S)HunyuanVideoDownBlock3Drdr rfrgreTrr.r/rhrwrxryrzadd_downsampledownsample_stridedownsample_paddingrNc stg} t|D]} | dkr|n|}| t||||||dq t| |_|r9tt||| | dg|_ nd|_ d|_ dS)Nrr~)r/r2r1F) r7r8rrrcr<rrr` downsamplersr) r?r.r/rhrwrxryrzrrrrir@r)r*r8s6    z HunyuanVideoDownBlock3D.__init__rBcC^tr|jr|jD]}|||}q n |jD]}||}q|jdur-|jD]}||}q&|Srb)rrrrrr)r?rBr downsamplerr)r)r*rGLs      zHunyuanVideoDownBlock3D.forward)rdr rfrgreTrr rr)r)r@r*rsB    /rcsveZdZ       dded ed ed ed ed edededeeeefddffdd Zde j de j fddZ Z S)HunyuanVideoUpBlock3Drdr rfrgreTrQr.r/rhrwrxryrz add_upsampleupsample_scale_factorrNc stg} t|D]} | dkr|n|} | t| |||||dq t| |_|r8tt||| dg|_ nd|_ d|_ dS)Nrr~)r/rRF) r7r8rrrcr<rrrP upsamplersr) r?r.r/rhrwrxryrzrrrrinput_channelsr@r)r*r8\s4    zHunyuanVideoUpBlock3D.__init__rBcCrrb)rrrrrr)r?rBr upsamplerr)r)r*rGs      zHunyuanVideoUpBlock3D.forward)rdr rfrgreTrQ) rHrIrJr:r rMrLrKr8rrNrGrOr)r)r@r*r[s<   .rcseZdZdZ          dd ed ed eedfdeedfdededededededdffdd Zde j de j fddZ Z S)HunyuanVideoEncoder3Dzx Causal encoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603). rrrrrrrresiluTr}r.r/down_block_types.block_out_channelslayers_per_blockract_fndouble_ztemporal_compression_ratiospatial_compression_ratiorNc  stt||dddd|_d|_tg|_|d} t|D]\} }|dkr0t d|| }|| } | t |dk}t t | }t t | }| dkrft| |k}t| t |d|koc| }n| dkrwt| |k}t| |k}nt d | |rd nd }|rd nd }t||}t||| t|p|d|||dd }|j|q!t|dd||d|| d|_tj|d|dd|_t|_|rd|n|}t|d|dd|_d|_dS)Nrrr r0r1rzUnsupported down_block_type: r}r$Unsupported time_compression_ratio: rrr r rr rf) rwr.r/rrxryrzrrrr.rxryr|rzr{r num_groupsrjrr0F)r7r8r,conv_in mid_blockr<r down_blocks enumerate ValueErrorlenr:nplog2rLrKrrrvrn conv_norm_outSiLUconv_actconv_outr)r?r.r/rrrrrrmid_block_add_attentionrroutput_channelrdown_block_type input_channelis_final_blocknum_spatial_downsample_layersnum_time_downsample_layersadd_spatial_downsampleadd_time_downsampledownsample_stride_HWdownsample_stride_Tr down_blockconv_out_channelsr@r)r*r8sd           zHunyuanVideoEncoder3D.__init__rBcCs||}tr |jr |jD]}|||}q||j|}n|jD]}||}q#||}||}||}| |}|Srb) rrrrrrrrrr)r?rBrr)r)r*rGs        zHunyuanVideoEncoder3D.forward) rrrrrrerTTr}r) rHrIrJ__doc__r:rKrMrLr8rrNrGrOr)r)r@r*rsJ     RrcseZdZdZ         dd ed ed eedfdeedfdededededeffdd Zdej dej fddZ Z S)HunyuanVideoDecoder3Dzx Causal decoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603). rrrrrrrrerTr}rr.r/up_block_types.rrrrtime_compression_ratiorc  st||_t||dddd|_tg|_t|dd||d||d|_ t t |} | d} t |D]y\} }|dkrEt d || }| | } | t|dk}tt| }tt| }| d kr{t| |k}t| t|d|kox| }nt d | |rd nd }|rdnd}t||}t|jd|| t|p||d||d}|j|| }q6tj|d|dd|_t|_t|d|dd|_d|_dS)Nrrr rrfrrrzUnsupported up_block_type: r}rrrrr)rwr.r/rrrxryrzrrF)r7r8rr,rr<r up_blocksrvrlistreversedrrrr:rrrLrKrrrnrrrrr)r?r.r/rrrrrrrrreversed_block_out_channelsrr up_block_typeprev_output_channelrnum_spatial_upsample_layersnum_time_upsample_layersadd_spatial_upsampleadd_time_upsampleupsample_scale_factor_HWupsample_scale_factor_Trup_blockr@r)r*r8 s^           zHunyuanVideoDecoder3D.__init__rBrcCs||}tr |jr ||j|}|jD]}|||}qn||}|jD]}||}q(||}||}| |}|Srb) rrrrrrrrrr)r?rBrr)r)r*rG[s        zHunyuanVideoDecoder3D.forward) rrrrrrerTr}r) rHrIrJrr:rKrMr8rrNrGrOr)r)r@r*rs@     Prcs\eZdZdZdZe        dJdedededeedfdeedfdeedededede dedede ddffdd Z      dKd edBd!edBd"edBd#e dBd$e dBd%e dBddfd&d'Z d(e jde jfd)d*Ze dLd(e jd+e deeeBfd,d-ZdLd.e jd+e dee jBfd/d0ZedLd.e jd+e dee jBfd1d2Zd3e jd4e jd5ede jfd6d7Zd3e jd4e jd5ede jfd8d9Zd3e jd4e jd5ede jfd:d;Zd(e jdefdd?Zd(e jdefd@dAZdLd.e jd+e dee jBfdBdCZ D  dMdEe jdFe d+e dGe jdBdee jBf dHdIZZ S)NAutoencoderKLHunyuanVideoaj A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Introduced in [HunyuanVideo](https://huggingface.co/papers/2412.03603). This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented for all models (such as downloading or saving). Trrrrrrre>I?rr}r.r/latent_channelsr.rrrrrscaling_factorrrrrNc st| |_t|||||| |d| | | d |_t|||||| || | | d |_tjd|d|dd|_ tj||dd|_ | |_ | |_ d|_ d|_d|_d|_d|_d|_d |_d |_d |_d |_dS) NT) r.r/rrrrrrrrr) r.r/rrrrrrrrrr rFrr )r7r8rrencoderrdecoderr<r= quant_convpost_quant_convrr use_slicing use_tilinguse_framewise_encodinguse_framewise_decodingtile_sample_min_heighttile_sample_min_widthtile_sample_min_num_framestile_sample_stride_heighttile_sample_stride_widthtile_sample_stride_num_frames)r?r.r/rrrrrrrrrrrr@r)r*r8|sR   z"AutoencoderKLHunyuanVideo.__init__rrrrrrcCsRd|_|p|j|_|p |j|_|p|j|_|p|j|_|p|j|_|p%|j|_dS)a 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_min_num_frames (`int`, *optional*): The minimum number of frames required for a sample to be separated into tiles across the frame 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. tile_sample_stride_num_frames (`int`, *optional*): The stride between two consecutive frame tiles. This is to ensure that there are no tiling artifacts produced across the frame dimension. TN)rrrrrrr)r?rrrrrrr)r)r* enable_tilings     z'AutoencoderKLHunyuanVideo.enable_tilingr&cCsf|j\}}}}}|jr||jkr||S|jr'||jks"||jkr'||S||}| |}|Srb) rrr_temporal_tiled_encoderrr tiled_encoderr)r?r&rrrrrencr)r)r*_encodes    z!AutoencoderKLHunyuanVideo._encode return_dictcs^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|]}|qSr))r).0x_slicer?r)r* sz4AutoencoderKLHunyuanVideo.encode..) latent_dist)rrrYrr]rrr )r?r&rencoded_slicesh posteriorr)r r*encode s   z AutoencoderKLHunyuanVideo.encodezc Cs|j\}}}}}|j|j}|j|j} |j|j} |jr(|| kr(|j||dS|jr:|| ks3||kr:|j ||dS| |}| |} |sI| fSt | dS)Nrsample) rrrrrrr_temporal_tiled_decoder tiled_decoderrr) r?rrrrrrrtile_latent_min_heighttile_latent_min_widthtile_latent_min_num_framesdecr)r)r*_decode%s      z!AutoencoderKLHunyuanVideo._decodecsXjr|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|]}|jqSr))rr)rz_slicer r)r*r Isz4AutoencoderKLHunyuanVideo.decode..r)rrrYrr]rrr)r?rrdecoded_slicesdecodedr)r r*decode9s   z AutoencoderKLHunyuanVideo.decodeab blend_extentc Cst|jd|jd|}t|D]@}|dddddd| |ddfd|||dddddd|ddf|||dddddd|ddf<q|S)Nr rrr)r?r r!r"r'r)r)r*blend_vS  R&z!AutoencoderKLHunyuanVideo.blend_vcCst|jd|jd|}t|D]@}|dddddddd| |fd|||dddddddd|f|||dddddddd|f<q|S)Nrr r$r?r r!r"r&r)r)r*blend_h[r&z!AutoencoderKLHunyuanVideo.blend_hc Cst|jd|jd|}t|D]@}|dddd| |ddddfd|||dddd|ddddf|||dddd|ddddf<q|S)Nr r$r'r)r)r*blend_tcr&z!AutoencoderKLHunyuanVideo.blend_tc Cs|j\}}}}}||j}||j}|j|j} |j|j} |j|j} |j|j} | | } | | }g}td||jD];}g}td||jD]*}|dddddd|||j|||jf}||}||}| |qF| |q;g}t |D]O\}}g}t |D]:\}}|dkr| ||d||| }|dkr| ||d||}| |ddddddd| d| fq| t j|ddq}t j|ddddddddd|d|f}|S)zEncode a batch of images using a tiled encoder. Args: x (`torch.Tensor`): Input batch of videos. Returns: `torch.Tensor`: The latent representation of the encoded videos. rNr r}rTr)rrrrrrrrrrrr%r(rr])r?r&rrrrr latent_height latent_widthrrtile_latent_stride_heighttile_latent_stride_width blend_height blend_widthrowsrrowjtile result_rows result_rowrr)r)r*rks<      2    .0z&AutoencoderKLHunyuanVideo.tiled_encodec Cs|j\}}}}}||j}||j} |j|j} |j|j} |j|j} |j|j} |j|j}|j|j}g}td|| D]8}g}td|| D](}|dddddd||| ||| f}||}||}| |qH| |q>g}t |D]Q\}}g}t |D]<\}}|dkr| ||d|||}|dkr| ||d||}| |ddddddd|jd|jfq| t j|ddq}t j|ddddddddd|d| f}|s|fSt|dS)a Decode a batch of images using a tiled decoder. Args: z (`torch.Tensor`): Input batch of latent vectors. return_dict (`bool`, *optional*, defaults to `True`): Whether or not 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. rNr rrTrr)rrrrrrrrrrrr%r(rr]r)r?rrrrrrr sample_height sample_widthrrr-r.r/r0r1rr2r3r4rr5r6rr)r)r*rs@        .    20 z&AutoencoderKLHunyuanVideo.tiled_decodec Cs|j\}}}}}|d|jd}|j|j}|j|j} || } g} td||jD]S} |dddd| | |jdddddf} |jrW||jksQ||jkrW|| } n | | } | | } | dkrx| ddddddddddf} | | q*g}t | D]B\} } | dkr| | | d| | } | | ddddd| ddddfq| | ddddd| dddddfqtj|ddddddd|f}|S)Nr rrrT)rrrrrrrrrrrrrr*rr])r?r&rrrrrlatent_num_framesrtile_latent_stride_num_framesblend_num_framesr2rr4r6rr)r)r*rs.  0   & .2$z0AutoencoderKLHunyuanVideo._temporal_tiled_encodec Cs|j\}}}}}|d|jd}|j|j} |j|j} |j|j} |j|j} |j|j} g}td|| D]Y}|dddd||| dddddf}|jrj|jd| ksa|jd| krj|j |ddj }n | |}| |}|dkr|ddddddddddf}| |q7g}t|D]D\}}|dkr|||d|| }| |ddddd|jddddfq| |ddddd|jdddddfqtj|ddddddd|f}|s|fSt|d S) Nr rrr#TrrrTr)rrrrrrrrrrrrrrrr*rr]r)r?rrrrrrrnum_sample_framesrrrr:r;r2rr4rr6rr)r)r*rs6     ."  & 04$ z0AutoencoderKLHunyuanVideo._temporal_tiled_decodeFrsample_posterior generatorc Cs<|}||j}|r|j|d}n|}|j||d}|S)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)rr rrCr) r?rr=rr>r&rrrr)r)r*rGs z!AutoencoderKLHunyuanVideo.forward) rrrrrrrrrerrr}T)NNNNNN)T)FTN)!rHrIrJr _supports_gradient_checkpointingrr:rKrMr rLr8rrrNrrr rrrrrr%r(r*rrrr GeneratorrGrOr)r)r@r*rqs  Y (   "    2: 'rrb).numpyrrtorch.nnr<torch.nn.functional functionalrDconfiguration_utilsrrutilsrutils.accelerate_utilsr activationsrattention_processorr modeling_outputsr modeling_utilsr vaer rr get_loggerrHloggerr:rrrNr+Moduler,rPr`rcrvrrrrrr)r)r)r*sJ          !*.a??lk