+hUSSKJrJrJrJr SSKrSSKJs Jr SSKJr SSK J r J r SSK Jr SSKJr SS KJrJr SS KJrJrJr SS KJr SS KJr SS KJrJrJr \R@"\!5r""SS\RF5r$"SS\RF5r%"SS5r&\"SS\RF55r'"SS\\ 5r(g))ListOptionalTupleUnionN)nn) ConfigMixinregister_to_config)logging)maybe_allow_in_graph) Attention FeedForward)TimestepEmbedding Timestepsget_3d_rotary_pos_embed)Transformer2DModelOutput) ModelMixin) AdaLayerNorm FP32LayerNormRMSNormc^\rSrSrSS\S\S\S\S\S\SS 4U4S jjjrS \ RS \ RS \ RS\ \ R\ R\ R\ R44Sjr Sr U=r$)EasyAnimateLayerNormZero#conditioning_dim embedding_dimelementwise_affineepsbias norm_typereturnNc>[TU]5 [R"5Ul[R "USU-US9UlUS:Xa[R"X#US9UlgUS:Xa[X#US9Ulg[SUS35e)N)r layer_normrrfp32_layer_normzUnsupported `norm_type` (z@) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'.) super__init__rSiLUsiluLinearlinear LayerNormnormr ValueError)selfrrrrrr __class__s o/home/james-whalen/.local/lib/python3.13/site-packages/diffusers/models/transformers/transformer_easyanimate.pyr(!EasyAnimateLayerNormZero.__init__$s GGI ii 0!m2C$O  $ ]_bcDI + +%m`cdDI+I;6vw  hidden_statesencoder_hidden_statestembcNURURU55RSSS9upEpgpURU5SUR S5--UR S5-nURU5SUR S5--UR S5-nXXi4$)Nr#dim)r,r*chunkr. unsqueeze) r0r5r6r7shiftscalegate enc_shift enc_scaleenc_gates r2forward EasyAnimateLayerNormZero.forward;s>B[[SW=Y=_=_`agh=_=i:dy -0A8J4JKeoo^_N`` $ *? @A H[H[\]H^D^ _bkbubu c ! TCCr4)r,r.r*)Th㈵>Tr&)__name__ __module__ __qualname____firstlineno__intboolfloatstrr(torchTensorrrD__static_attributes__ __classcell__r1s@r2rr#s $(*!      .D"\\DBG,,DV[VbVbD u||U\\5<<E FDDr4rc~^\rSrSrS\S\\SS4U4SjjrSrS\RS\R4S jr S r U=r $) EasyAnimateRotaryPosEmbedF patch_sizerope_dimr!Nc:>[TU]5 XlX lgN)r'r(rWrX)r0rWrXr1s r2r("EasyAnimateRotaryPosEmbed.__init__Gs $ r4cUnUnUupgXg- nXU- :aUn [[XV- U-55n OUn [[XG- U-55n [[XY- S- 55n [[XJ- S- 55n X4X-X-44$)Ng@)rKround) r0src tgt_width tgt_heighttwthhwr resize_height resize_widthcrop_top crop_lefts r2get_resize_crop_region_for_grid9EasyAnimateRotaryPosEmbed.get_resize_crop_region_for_gridMs   E R=MuRVaZ01LLbfqj 12Mub0C789r0C789 $x'?AY&ZZZr4r5cUR5up#pEnXPR-nX`R-nSUR-nSUR-nURXV4Xx5n [URU XV4URS5SS9n U $)NZ<r T) grid_size temporal_sizeuse_real)sizerWrjrrX) r0r5bsc num_frames grid_height grid_widthbase_size_widthbase_size_heightgrid_crops_coordsimage_rotary_embs r2rD!EasyAnimateRotaryPosEmbed.forward^s5B5G5G5I2z !__4 ??2 /0 @@  % 3 MM "/',,Q/   r4)rWrX) rGrHrIrJrKrr(rjrOrPrDrQrRrSs@r2rUrUFsG!3!$s)!! [" U\\ ell  r4rUc \rSrSrSrSrS S\S\RS\RS\ \RS \ \RS \R4 S jjr S r g)EasyAnimateAttnProcessor2_0rz Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is used in the EasyAnimateTransformer3DModel model. cD[[S5(d [S5eg)Nscaled_dot_product_attentionzaEasyAnimateAttnProcessor2_0 requires PyTorch 2.0 or above. To use it, please install PyTorch 2.0.)hasattrF ImportError)r0s r2r($EasyAnimateAttnProcessor2_0.__init__xs%q899s :r4Nattnr5r6attention_maskr{r!c 0URcUb[R"X2/SS9nURU5nUR U5nUR U5nUR SURS45RSS5nUR SURS45RSS5nUR SURS45RSS5nURbURU5nURbURU5nURGb?UGb;URU5n URU5n URU5n U R SURS45RSS5n U R SURS45RSS5n U R SURS45RSS5n URbURU 5n URbURU 5n [R"X/SS9n[R"X/SS9n[R"X/SS9nUbSSKJn U "USS2SS2UR"SS24U5USS2SS2UR"SS24'UR$(d;U "USS2SS2UR"SS24U5USS2SS2UR"SS24'[&R("XgXSSS9nURSS5R+SS 5nUR-UR.5nUbUSS2SUR"S24USS2UR"SS24p#[1US S5b*UR2S "U5nUR2S"U5n[1US S5bUR5U5nX#4$[1US S5b*UR2S "U5nUR2S"U5nX#4$) Nr9r:r )apply_rotary_embF) attn_mask dropout_p is_causalrto_outr to_add_out) add_q_projrOcatto_qto_kto_v unflattenheads transposenorm_qnorm_k add_k_proj add_v_proj norm_added_q norm_added_k embeddingsrshapeis_cross_attentionrrflattentodtypegetattrrr) r0rr5r6rr{querykeyvalue encoder_query encoder_key encoder_valuers r2__call__$EasyAnimateAttnProcessor2_0.__call__~s ?? "'<'H!II'<&LRSTM -(ii & -(DJJ#34>>q!DmmA B/0::1a@DJJ#34>>q!D ;; "KK&E ;; "++c"C ?? &+@+L OO,ABM//*?@K OO,ABM)33A B7GHRRSTVWXM%//DJJ3CDNNqRSTK)33A B7GHRRSTVWXM  , $ 1 1- @   ,"// < II}4!N1399!<>>?AQ?Aq/55a8::; 66 3RW &//15==aC %((5  ! ,a!A#8#>#>q#A!AABa!6![TU]5 [XAXxUSS9Ul[ UUUU (aSOSSSSU(aUOSU(aSOS[ 5S9 Ul[XAXxUSS9Ul[UUUU U U S9Ul SUl U(a[UUUU U U S9Ul SUl U (a[XUS9Ul gg) NT)r rr$ư>F) query_dimdim_headrrrradded_proj_biasadded_kv_proj_dimcontext_pre_only processor)rrr inner_dimrr%) r'r(rnorm1rr~attn1norm2rfftxt_ffnorm3r)r0r;rrrrrrrrrrrrr rr1s r2r($EasyAnimateTransformerBlock.__init__s$ . !8i^b  '%$+L %3c&4U$13   . !8i^b   ''"   %++& DK &s\deDJ r4r5r6r7r{r!cURXU5upVpxURUUUS9upXRS5U --nX(RS5U --nURXU5upVpURboUR UR U55nUR b!UR UR U55nObUR UR U55nOAUR U5nUR bUR U5nOUR U5nXRS5U--nX,RS5U--nX4$)N)r5r6r{r9)rrr=rrrr) r0r5r6r7r{norm_hidden_statesnorm_encoder_hidden_statesgate_msa enc_gate_msaattn_hidden_statesattn_encoder_hidden_statesgate_ff enc_gate_ffs r2rD#EasyAnimateTransformerBlock.forwardsoRVQ[Q[ $R N:>,"<-:D: 6 &(:(:1(=@R(RR 58N8Nq8QTn8n nPTzz $P L :: !!%DGG4F,G!H {{&-1ZZ D^8_-`*-1ZZ@Z8[-\*!%);!< {{&-1[[9S-T*-1WW5O-P*%(9(9!(a` A Transformer model for video-like data in [EasyAnimate](https://github.com/aigc-apps/EasyAnimate). Parameters: num_attention_heads (`int`, defaults to `48`): The number of heads to use for multi-head attention. attention_head_dim (`int`, defaults to `64`): The number of channels in each head. in_channels (`int`, defaults to `16`): The number of channels in the input. out_channels (`int`, *optional*, defaults to `16`): The number of channels in the output. patch_size (`int`, defaults to `2`): The size of the patches to use in the patch embedding layer. sample_width (`int`, defaults to `90`): The width of the input latents. sample_height (`int`, defaults to `60`): The height of the input latents. activation_fn (`str`, defaults to `"gelu-approximate"`): Activation function to use in feed-forward. timestep_activation_fn (`str`, defaults to `"silu"`): Activation function to use when generating the timestep embeddings. num_layers (`int`, defaults to `30`): The number of layers of Transformer blocks to use. mmdit_layers (`int`, defaults to `1000`): The number of layers of Multi Modal Transformer blocks to use. dropout (`float`, defaults to `0.0`): The dropout probability to use. time_embed_dim (`int`, defaults to `512`): Output dimension of timestep embeddings. text_embed_dim (`int`, defaults to `4096`): Input dimension of text embeddings from the text encoder. norm_eps (`float`, defaults to `1e-5`): The epsilon value to use in normalization layers. norm_elementwise_affine (`bool`, defaults to `True`): Whether to use elementwise affine in normalization layers. flip_sin_to_cos (`bool`, defaults to `True`): Whether to flip the sin to cos in the time embedding. time_position_encoding_type (`str`, defaults to `3d_rope`): Type of time position encoding. after_norm (`bool`, defaults to `False`): Flag to apply normalization after. resize_inpaint_mask_directly (`bool`, defaults to `True`): Flag to resize inpaint mask directly. enable_text_attention_mask (`bool`, defaults to `True`): Flag to enable text attention mask. add_noise_in_inpaint_model (`bool`, defaults to `False`): Flag to add noise in inpaint model. Tr)z^proj$r.z ^proj_out$rr in_channels out_channelsrW sample_width sample_heightrtimestep_activation_fn freq_shift num_layers mmdit_layersrradd_norm_text_encodertext_embed_dimtext_embed_dim_t5rrflip_sin_to_costime_position_encoding_typeresize_inpaint_mask_directlyenable_text_attention_maskadd_noise_in_inpaint_modelc>[TU]5 X-n[UUU 5Ul[ UX5Ul[ XR5Ul[R"UUXU4USS9Ul SUl SUl U(d<[R"UU5Ul Ub[R"UU5Ul Ow[R"[USSS9[R"UU55Ul Ub:[R"[USSS9[R"UU55Ul [R "[#U 5Vs/sHn[%UUUUU UUUUUU :aSOSS9 PM sn5Ul[R("UUU5Ul[-USU-UUSS 9Ul[R"UXU-U-5UlSUlgs snf) NT) kernel_sizestriderr)rF) r;rrrrrrrrrr r9)r output_dimrr chunk_dim)r'r(r time_projrtime_embeddingrUrope_embeddingrConv2dproj text_proj text_proj_t5r+ Sequentialr ModuleListrangertransformer_blocksr- norm_finalrnorm_outproj_outgradient_checkpointing)r0rrrrrWrrrrrrrrrrrrrrrrrrrrr_r1s r2r(&EasyAnimateTransformer3DModel.__init__us: '< #9ozJ/ >b7 WII 0HQ[bf   $YY~yADN ,$&II.?$K!]]F R`bkHlDN!,$&MMNDTJBIIVgirLs%! #%--z* +A,!(;'9#1#"/,C%)+,|+;4 + # ",,y(s^0d(,$67'G$$&"$%)&*$(/&,!&+"!%(, $+4-1+/+/5W, W, W,c] W, sm W, SM W,W,W,W,!$W,W,W,W,W,W, $!W,"#W,$%W,&'W,("&)W,*+W,,&)-W,0'+1W,2%)3W,4%)5W,W,z158<;?2626 @7||@7,,@7  - @7 ( 5 @7 #+5<<"8 @7"%,,/@7"%,,/@7@7 uU\\"$<< =@7@7r4r))typingrrrrrOtorch.nn.functionalr functionalrconfiguration_utilsr r utilsr utils.torch_utilsr attentionrrrrrrmodeling_outputsrmodeling_utilsr normalizationrrr get_loggerrGloggerModulerrUr~rrrr4r2r2s 0/ B5.NN7'@@   H % Dryy DF)  ) X]4]4@h4"))h4h4VQ7J Q7r4