ó
    oM‡j§>  ã                   ó¼  • S SK r S SKJr  S SKrS SKJr  S SKJs  Jr  S SK	J
r
  SSKJrJr  SSKJr  SSKJr  SSKJr  S	S
KJr  SSKJr  S,S jrS r " S S\R4                  5      rS-S\S\S\R4                  4S jjrS\R>                  S\S\R>                  4S jr  " S S\R4                  5      r! " S S\R4                  5      r" " S S\R4                  5      r# " S S\R4                  5      r$ " S  S!\R4                  5      r% " S" S#\R4                  5      r& " S$ S%\R4                  5      r'\ " S& S'\5      5       r(\ " S( S)\5      5       r) " S* S+\\\5      r*g).é    N)Ú	dataclass)Úweight_normé   )ÚConfigMixinÚregister_to_config)Ú
BaseOutput)Úapply_forward_hook)Úrandn_tensoré   )Ú
ModelMixiné   )ÚAutoencoderMixinc                 óB   • [        [        R                  " XX#XTUS95      $ )N)Úbias)r   ÚnnÚConv1d)Úin_channelsÚout_channelsÚkernel_sizeÚstrideÚdilationÚpaddingr   s          ÚÉC:\Users\johnn\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.13_qbz5n2kfra8p0\LocalCache\local-packages\Python313\site-packages\diffusers/models/autoencoders/autoencoder_longcat_audio_dit.pyÚ
_wn_conv1dr   "   s   € Ü”r—y’y ¸KÐQXÐimÑnÓoÐoó    c                  ó@   • [        [        R                  " U 0 UD65      $ ©N)r   r   ÚConvTranspose1d)ÚargsÚkwargss     r   Ú_wn_conv_transpose1dr!   &   s   € Ü”r×)Ò)¨4Ð:°6Ñ:Ó;Ð;r   c                   ór   ^ • \ rS rSrS	S\S\4U 4S jjjrS\R                  S\R                  4S jr	Sr
U =r$ )
ÚSnake1dé*   ÚchannelsÚalpha_logscalec                 óê   >• [         TU ]  5         X l        [        R                  " [
        R                  " U5      5      U l        [        R                  " [
        R                  " U5      5      U l        g r   )	ÚsuperÚ__init__r&   r   Ú	ParameterÚtorchÚzerosÚalphaÚbeta)Úselfr%   r&   Ú	__class__s      €r   r)   ÚSnake1d.__init__+   sF   ø€ Ü‰ÑÔØ,ÔÜ—\’\¤%§+¢+¨hÓ"7Ó8ˆŒ
Ü—L’L¤§¢¨XÓ!6Ó7ˆ�	r   Úhidden_statesÚreturnc                 ó2  • U R                   S S S 2S 4   nU R                  S S S 2S 4   nU R                  (       a,  [        R                  " U5      n[        R                  " U5      nUSUS-   -  [        R
                  " X-  5      R                  S5      -  -   $ )Ng      ð?g•Ö&è.>r   )r-   r.   r&   r+   ÚexpÚsinÚpow)r/   r2   r-   r.   s       r   ÚforwardÚSnake1d.forward1   s   € Ø—
‘
˜4¢ D˜=Ñ)ˆØ�y‰y˜šq $˜Ñ'ˆØ××Ü—I’I˜eÓ$ˆEÜ—9’9˜T“?ˆDØ  t¨d¡{Ñ 3´u·y²yÀÑAVÓ7W×7[Ñ7[Ð\]Ó7^Ñ^Ñ^Ð^r   )r-   r&   r.   ©T)Ú__name__Ú
__module__Ú__qualname__Ú__firstlineno__ÚintÚboolr)   r+   ÚTensorr8   Ú__static_attributes__Ú__classcell__©r0   s   @r   r#   r#   *   s?   ø† ñ8 ð 8°d÷ 8ð 8ð_ U§\¡\ð _°e·l±l÷ _ò _r   r#   Únamer%   r3   c                 ó~   • U S:X  a  [         R                  " 5       nU$ U S:X  a  [        U5      nU$ [        SU  35      e)NÚeluÚsnakezUnknown activation: )r   ÚELUr#   Ú
ValueError)rE   r%   Úacts      r   Ú_get_vae_activationrL   :   sJ   € Øˆuƒ}Ü�fŠf‹hˆð
 €Jð	 
�‹Ü�hÓˆð €Jô Ð/°¨vÐ6Ó7Ð7r   r2   Úfactorc                 ó¼   • U R                  5       u  p#nU R                  X#U-  X5      R                  SSSS5      R                  5       R                  X#U-  XA-  5      $ )Nr   r   r   r   )ÚsizeÚviewÚpermuteÚ
contiguous)r2   rM   Úbatchr%   Úwidths        r   Ú_pixel_shuffle_1drU   D   sY   € Ø*×/Ñ/Ó1Ñ€E�Uà×Ñ˜5¨fÑ"4°fÓDß	‰��A�q˜!Ó	ß	‰‹ß	‰ˆe Ñ'¨©Ó	8ð	r   c                   ór   ^ • \ rS rSrS\S\S\4U 4S jjrS\R                  S\R                  4S jrS	r	U =r
$ )
ÚDownsampleShortcutéN   r   r   rM   c                 óR   >• [         TU ]  5         X0l        X-  U-  U l        X l        g r   )r(   r)   rM   Ú
group_sizer   ©r/   r   r   rM   r0   s       €r   r)   ÚDownsampleShortcut.__init__O   s(   ø€ Ü‰ÑÔØŒØ%Ñ.°,Ñ>ˆŒØ(Õr   r2   r3   c                 ó„  • UR                   u  p#nUR                  X#X@R                  -  U R                  5      R                  SSSS5      R	                  5       R                  X#U R                  -  X@R                  -  5      nUR                  X R
                  U R                  X@R                  -  5      R                  SS9$ )Nr   r   r   r   ©Údim)ÚshaperP   rM   rQ   rR   r   rZ   Úmean)r/   r2   rS   r%   rT   s        r   r8   ÚDownsampleShortcut.forwardU   s    € Ø!.×!4Ñ!4Ñˆ˜à×Ñ˜u°¿¹Ñ0DÀdÇkÁkÓRß‰W�Q˜˜1˜aÓ ß‰Z‹\ß‰T�% D§K¡KÑ/°¿+¹+Ñ1EÓFð	 	ð ×!Ñ! %×):Ñ):¸D¿O¹OÈU×VaÑVaÑMaÓb×gÑgÐlmÐgÐnÐnr   )rM   rZ   r   ©r;   r<   r=   r>   r?   r)   r+   rA   r8   rB   rC   rD   s   @r   rW   rW   N   sA   ø† ð) Cð )°sð )ÀC÷ )ðo U§\¡\ð o°e·l±l÷ oò or   rW   c                   ór   ^ • \ rS rSrS\S\S\4U 4S jjrS\R                  S\R                  4S jrS	r	U =r
$ )
ÚUpsampleShortcuté`   r   r   rM   c                 óF   >• [         TU ]  5         X0l        X#-  U-  U l        g r   )r(   r)   rM   Úrepeatsr[   s       €r   r)   ÚUpsampleShortcut.__init__a   s"   ø€ Ü‰ÑÔØŒØ#Ñ,°Ñ;ˆ�r   r2   r3   c                 ó`   • UR                  U R                  SS9n[        XR                  5      $ )Nr   r^   )Úrepeat_interleaverh   rU   rM   ©r/   r2   s     r   r8   ÚUpsampleShortcut.forwardf   s*   € Ø%×7Ñ7¸¿¹È!Ð7ÐLˆÜ  ·±Ó<Ð<r   )rM   rh   rc   rD   s   @r   re   re   `   s=   ø† ð< Cð <°sð <ÀC÷ <ð
= U§\¡\ð =°e·l±l÷ =ò =r   re   c                   ó€   ^ • \ rS rSr SS\S\S\S\S\4
U 4S jjjrS\R                  S	\R                  4S
 jr	Sr
U =r$ )ÚVaeResidualUniték   r   r   r   r   Úact_fnc                 ó°   >• [         TU ]  5         X4S-
  -  S-  n[        R                  " [	        XRS9[        XXCUS9[	        XRS9[        X"SS95      U l        g )Nr   r   ©r%   )r   r   )r   )r(   r)   r   Ú
SequentialrL   r   Úlayers)r/   r   r   r   r   rq   r   r0   s          €r   r)   ÚVaeResidualUnit.__init__l   sX   ø€ ô 	‰ÑÔØ¨Q™Ñ/°AÑ5ˆÜ—m’mÜ Ñ>Ü�{°+ÐZaÑbÜ Ñ>Ü�|¸qÑAó	
ˆ�r   r2   r3   c                 ó(   • XR                  U5      -   $ r   ©ru   rl   s     r   r8   ÚVaeResidualUnit.forwardx   s   € ØŸ{™{¨=Ó9Ñ9Ð9r   rx   )é   rH   ©r;   r<   r=   r>   r?   Ústrr)   r+   rA   r8   rB   rC   rD   s   @r   ro   ro   k   sX   ø† àfmñ

Øð

Ø.1ð

Ø=@ð

ØORð

Ø`c÷

ð 

ð: U§\¡\ð :°e·l±l÷ :ò :r   ro   c                   ó‚   ^ • \ rS rSr  SS\S\S\S\S\4
U 4S jjjrS\R                  S	\R                  4S
 jr	Sr
U =r$ )ÚVaeEncoderBlocké|   r   r   r   rq   Údownsample_shortcutc                 ój  >• [         TU ]  5         [        XSUS9[        XSUS9[        XSUS9/nUR                  [	        XAS95        UR                  [        XSU-  U[        R                  " US-  5      S95        [        R                  " U6 U l
        US:X  a  [        XU5      U l        g S U l        g )	Nr   ©r   rq   r   é	   rs   r   ©r   r   r   Ú	averaging)r(   r)   ro   ÚappendrL   r   ÚmathÚceilr   rt   ru   rW   Úresidual)r/   r   r   r   rq   r€   ru   r0   s          €r   r)   ÚVaeEncoderBlock.__init__}   s·   ø€ ô 	‰ÑÔä˜K¸qÈÑPÜ˜K¸qÈÑPÜ˜K¸qÈÑPð
ˆð
 	�‰Ô)¨&ÑGÔHØ�‰Ü�{¸aÀ&¹jÐQWÔae×ajÒajÐkqÐtuÑkuÓavÑwô	
ô —m’m VÐ,ˆŒàEXÐ\gÓEgÔ˜{¸&ÓAð 	�Ømqð 	�r   r2   r3   c                 ól   • U R                  U5      nU R                  b  U R                  U5      nX#-   nU$ r   ©ru   r‰   ©r/   r2   Úoutput_hidden_statesr‰   s       r   r8   ÚVaeEncoderBlock.forward”   ó8   € Ø#Ÿ{™{¨=Ó9ÐØ�=‰=Ñ$Ø—}‘} ]Ó3ˆHØ#7Ñ#BÐ Ø#Ð#r   rŒ   ©rH   Únoner{   rD   s   @r   r~   r~   |   se   ø† ð Ø#)ñ
àð
ð ð
ð ð	
ð
 ð
ð !÷
ð 
ð.$ U§\¡\ð $°e·l±l÷ $ò $r   r~   c                   ó‚   ^ • \ rS rSr  SS\S\S\S\S\4
U 4S jjjrS\R                  S	\R                  4S
 jr	Sr
U =r$ )ÚVaeDecoderBlockéœ   r   r   r   rq   Úupsample_shortcutc                 ó*  >• [         TU ]  5         [        XAS9[        XSU-  U[        R
                  " US-  5      S9[        X"SUS9[        X"SUS9[        X"SUS9/n[        R                  " U6 U l	        US:X  a  [        XU5      U l        g S U l        g )	Nrs   r   r„   r   r‚   r   rƒ   Úduplicating)r(   r)   rL   r!   r‡   rˆ   ro   r   rt   ru   re   r‰   )r/   r   r   r   rq   r–   ru   r0   s          €r   r)   ÚVaeDecoderBlock.__init__�   s£   ø€ ô 	‰ÑÔä Ñ=Ü Ø°q¸6±zÈ&ÔZ^×ZcÒZcÐdjÐmnÑdnÓZoñô ˜LÀÈ6ÑRÜ˜LÀÈ6ÑRÜ˜LÀÈ6ÑRð
ˆô —m’m VÐ,ˆŒàCTÐXeÓCeÔ˜[¸Ó?ð 	�Økoð 	�r   r2   r3   c                 ól   • U R                  U5      nU R                  b  U R                  U5      nX#-   nU$ r   rŒ   r�   s       r   r8   ÚVaeDecoderBlock.forward´   r�   r   rŒ   r‘   r{   rD   s   @r   r”   r”   œ   se   ø† ð Ø!'ñ
àð
ð ð
ð ð	
ð
 ð
ð ÷
ð 
ð.$ U§\¡\ð $°e·l±l÷ $ò $r   r”   c                   ó¸   ^ • \ rS rSr         SS\S\S\\   S-  S\\   S-  S\S\S	\S
\S\4U 4S jjjrS\R                  S\R                  4S jr
SrU =r$ )ÚAudioDiTVaeEncoderé¼   Nr   r%   Úc_multsÚstridesÚ
latent_dimÚencoder_latent_dimrq   r€   Úout_shortcutc
                 óÐ  >• [         TU ]  5         S/U=(       d    / SQ-   n[        U=(       d    S/[        U5      S-
  -  5      n[        U5      [        U5      S-
  :  a:  UR	                  U(       a  US   OS/[        U5      S-
  [        U5      -
  -  5        OUS [        U5      S-
   nUn
[        XS   U
-  SSS9/n[        [        U5      S-
  5       H.  nUR                  [        X<   U
-  X<S-      U
-  XL   UUS	95        M0     UR                  [        US   U
-  USSS95        [        R                  " U6 U l        U	S
:X  a  [        US   U
-  US5      U l        g S U l        g )Nr   ©r   r   é   é   é   r   éÿÿÿÿr   rz   r   ©r   r   )rq   r€   r…   )r(   r)   ÚlistÚlenÚextendr   Úranger†   r~   r   rt   ru   rW   Úshortcut)r/   r   r%   rŸ   r    r¡   r¢   rq   r€   r£   Úchannels_baseru   Úidxr0   s                €r   r)   ÚAudioDiTVaeEncoder.__init__½   sy  ø€ ô 	‰ÑÔØ�#˜×4Ò$4Ñ5ˆÜ�w×: 1 #¬¨W«¸Ñ)9Ñ":Ó;ˆÜˆw‹<œ#˜g›,¨Ñ*Ó*Ø�N‰N®7˜G BšK¸Ð:¼cÀ'»lÈQÑ>NÔQTÐU\ÓQ]Ñ>]Ñ^Õ_àÐ0¤ G£¨qÑ 0Ð1ˆGØ ˆÜ˜[°!©*°}Ñ*DÐRSÐ]^Ñ_Ð`ˆÜœ˜W›¨Ñ)Ö*ˆCØ�M‰MÜØ‘L =Ñ0Ø !™GÑ$ }Ñ4Ø‘LØ!Ø(;ñöñ +ð 	�‰”j ¨¡¨}Ñ!<Ð>PÐ^_ÐijÑkÔlÜ—m’m VÐ,ˆŒð ˜{Ó*ô ˜w r™{¨]Ñ:Ð<NÐPQÓRð 	�ð ð 	�r   r2   r3   c                 óž   • U R                   S S " U5      nU R                   S   " U5      nU R                  b  U R                  U5      nX#-   nU$ )Nr©   ©ru   r¯   )r/   r2   rŽ   r¯   s       r   r8   ÚAudioDiTVaeEncoder.forwardä   sQ   € ØŸ™ C RÑ(¨Ó7ˆØ#Ÿ{™{¨2š¨}Ó=ÐØ�=‰=Ñ$Ø—}‘} ]Ó3ˆHØ#7Ñ#BÐ Ø#Ð#r   r´   )	r   é€   NNé@   r¶   rH   r…   r…   )r;   r<   r=   r>   r?   r«   r|   r)   r+   rA   r8   rB   rC   rD   s   @r   r�   r�   ¼   sµ   ø† ð ØØ$(Ø$(ØØ"%ØØ#.Ø'ñ%
àð%
ð ð%
ð �c‘˜TÑ!ð	%
ð
 �c‘˜TÑ!ð%
ð ð%
ð  ð%
ð ð%
ð !ð%
ð ÷%
ð %
ðN$ U§\¡\ð $°e·l±l÷ $ò $r   r�   c                   ó¸   ^ • \ rS rSr         SS\S\S\\   S-  S\\   S-  S\S\S	\S
\S\4U 4S jjjrS\	R                  S\	R                  4S jrSrU =r$ )ÚAudioDiTVaeDecoderéí   Nr   r%   rŸ   r    r¡   rq   Úin_shortcutÚ
final_tanhr–   c
                 óŠ  >• [         TU ]  5         S/U=(       d    / SQ-   n[        U=(       d    S/[        U5      S-
  -  5      n[        U5      [        U5      S-
  :  a:  UR	                  U(       a  US   OS/[        U5      S-
  [        U5      -
  -  5        OUS [        U5      S-
   nUn
US:X  a  [        XSS   U
-  S5      OS U l        [        XSS   U
-  SSS9/n[        [        U5      S-
  S	S5       H1  nUR                  [        X<   U
-  X<S-
     U
-  XLS-
     UU	S
95        M3     UR                  [        XcS	   U
-  S95        UR                  [        US	   U
-  USSSS95        UR                  U(       a  [        R                  " 5       O[        R                  " 5       5        [        R                  " U6 U l        g )Nr   r¥   r   r©   r˜   rz   r   rª   r   )rq   r–   rs   F)r   r   r   )r(   r)   r«   r¬   r­   re   r¯   r   r®   r†   r”   rL   r   ÚTanhÚIdentityrt   ru   )r/   r   r%   rŸ   r    r¡   rq   r»   r¼   r–   r°   ru   r±   r0   s                €r   r)   ÚAudioDiTVaeDecoder.__init__î   s±  ø€ ô 	‰ÑÔØ�#˜×4Ò$4Ñ5ˆÜ�w×: 1 #¬¨W«¸Ñ)9Ñ":Ó;ˆÜˆw‹<œ#˜g›,¨Ñ*Ó*Ø�N‰N®7˜G BšK¸Ð:¼cÀ'»lÈQÑ>NÔQTÐU\ÓQ]Ñ>]Ñ^Õ_àÐ0¤ G£¨qÑ 0Ð1ˆGØ ˆð MXÐ[hÓLhÔ˜Z°©°}Ñ)DÀaÔHÐnrð 	Œô ˜Z°©°}Ñ)DÐRSÐ]^Ñ_Ð`ˆÜœ˜W›¨Ñ)¨1¨bÖ1ˆCØ�M‰MÜØ‘L =Ñ0Ø !™GÑ$ }Ñ4Ø !™GÑ$Ø!Ø&7ñöñ 2ð 	�‰Ô)¨&À1¹:ÈÑ;UÑVÔWØ�‰”j ¨¡¨mÑ!;¸[ÐVWÐabÐinÑoÔpØ�‰¦:”b—g’g”i´2·;²;³=ÔAÜ—m’m VÐ,ˆ�r   r2   r3   c                 ó¶   • U R                   c  U R                  U5      $ U R                  U5      U R                  S   " U5      -   nU R                  SS  " U5      $ )Nr   r   )r¯   ru   rl   s     r   r8   ÚAudioDiTVaeDecoder.forward  sP   € Ø�=‰=Ñ Ø—;‘;˜}Ó-Ð-ØŸ™ mÓ4°t·{±{À1²~ÀmÓ7TÑTˆØ�{‰{˜1˜2‰˜}Ó-Ð-r   r´   )	r   r¶   NNr·   rH   r˜   Fr˜   )r;   r<   r=   r>   r?   r«   r|   r@   r)   r+   rA   r8   rB   rC   rD   s   @r   r¹   r¹   í   sµ   ø† ð ØØ$(Ø$(ØØØ(Ø Ø!.ñ'-àð'-ð ð'-ð �c‘˜TÑ!ð	'-ð
 �c‘˜TÑ!ð'-ð ð'-ð ð'-ð ð'-ð ð'-ð ÷'-ð '-ðR. U§\¡\ð .°e·l±l÷ .ò .r   r¹   c                   ó4   • \ rS rSr% \R
                  \S'   Srg)ÚLongCatAudioDiTVaeEncoderOutputi  Úlatents© N©r;   r<   r=   r>   r+   rA   Ú__annotations__rB   rÆ   r   r   rÄ   rÄ     s   ‡ à�\‰\Ör   rÄ   c                   ó4   • \ rS rSr% \R
                  \S'   Srg)ÚLongCatAudioDiTVaeDecoderOutputi#  ÚsamplerÆ   NrÇ   rÆ   r   r   rÊ   rÊ   #  s   ‡ à�L‰LÖr   rÊ   c            "       ó  ^ • \ rS rSrSr\                SS\S\S\\   S-  S\\   S-  S\S	\S
\S-  S\	S-  S\S\S\S\S\	S\S\S\
4 U 4S jjj5       r\   S S\R                  S\	S\	S\R                  S-  S\\\R                     -  4
S jj5       r\ S!S\R                  S\	S\\\R                     -  4S jj5       r   S"S\R                  S\	S\	S\R                  S-  S\\\R                     -  4
S jjrSrU =r$ )#ÚLongCatAudioDiTVaei(  FNr   r%   rŸ   r    r¡   r¢   rq   Ú	use_snaker€   r–   r£   r»   r¼   Údownsampling_ratioÚsample_rateÚscalec                 ó¢   >• [         TU ]  5         Uc  Uc  SnOU(       a  SOSn[        UUUUUUUU	US9	U l        [	        UUUUUUUUU
S9	U l        g )NrH   rG   )	r   r%   rŸ   r    r¡   r¢   rq   r€   r£   )	r   r%   rŸ   r    r¡   rq   r»   r¼   r–   )r(   r)   r�   Úencoderr¹   Údecoder)r/   r   r%   rŸ   r    r¡   r¢   rq   rÎ   r€   r–   r£   r»   r¼   rÏ   rÐ   rÑ   r0   s                    €r   r)   ÚLongCatAudioDiTVae.__init__+  sy   ø€ ô( 	‰ÑÔØ‰>ØÑ Ø ‘æ$-™°5�Ü)Ø#ØØØØ!Ø1ØØ 3Ø%ñ

ˆŒô *Ø#ØØØØ!ØØ#Ø!Ø/ñ

ˆ�r   rË   Úsample_posteriorÚreturn_dictÚ	generatorr3   c                 ó(  • [        U R                  R                  5       5      R                  nUR                  U:w  a  UR	                  U5      nU R                  U5      nUR                  SSS9u  px[        R                  " U5      S-   n	U(       a1  [        UR                  XGR                  UR                  S9n
XyU
-  -   nOUnX°R                  R                  -  nU[        R                  :w  a  UR                  5       nU(       d  U4$ [!        US9$ )Nr   r   r^   g-Cëâ6?)rØ   ÚdeviceÚdtype)rÅ   )ÚnextrÓ   Ú
parametersrÛ   ÚtoÚchunkÚFÚsoftplusr
   r`   rÚ   ÚconfigrÑ   r+   Úfloat32ÚfloatrÄ   )r/   rË   rÖ   r×   rØ   Úencoder_dtypeÚencodedra   Úscale_paramÚstdÚnoiserÅ   s               r   ÚencodeÚLongCatAudioDiTVae.encode\  sá   € ô ˜TŸ\™\×4Ñ4Ó6Ó7×=Ñ=ˆØ�<‰<˜=Ó(Ø—Y‘Y˜}Ó-ˆFØ—,‘,˜vÓ&ˆØ#ŸM™M¨!°˜MÐ3ÑˆÜ�jŠj˜Ó%¨Ñ,ˆÞÜ  §¡°yÏÉÐ\`×\fÑ\fÑgˆEØ 5™[Ñ(‰GàˆGØŸK™K×-Ñ-Ñ-ˆØœEŸM™MÓ)Ø—m‘m“oˆGÞØ�:ÐÜ.°wÑ?Ð?r   rÅ   c                 ó^  • [        U R                  R                  5       5      R                  nXR                  R
                  -  nUR                  U:w  a  UR                  U5      nU R                  U5      nU[        R                  :w  a  UR                  5       nU(       d  U4$ [        US9$ )N©rË   )rÜ   rÔ   rÝ   rÛ   râ   rÑ   rÞ   r+   rã   rä   rÊ   )r/   rÅ   r×   Údecoder_dtypeÚdecodeds        r   ÚdecodeÚLongCatAudioDiTVae.decodev  sˆ   € ô ˜TŸ\™\×4Ñ4Ó6Ó7×=Ñ=ˆØŸK™K×-Ñ-Ñ-ˆØ�=‰=˜MÓ)Ø—j‘j Ó/ˆGØ—,‘,˜wÓ'ˆØœEŸM™MÓ)Ø—m‘m“oˆGÞØ�:ÐÜ.°gÑ>Ð>r   c                 ó’   • U R                  XSUS9R                  nU R                  USS9R                  nU(       d  U4$ [	        US9$ )a²  
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 [`LongCatAudioDiTVaeDecoderOutput`] instead of a plain tuple.
    generator (`torch.Generator`, *optional*):
        A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
        deterministic.

Returns:
    [`LongCatAudioDiTVaeDecoderOutput`] or `tuple`:
        If `return_dict` is True, a [`LongCatAudioDiTVaeDecoderOutput`] is returned, otherwise a plain `tuple`
        is returned.
T)rÖ   r×   rØ   )r×   rí   )rê   rÅ   rð   rË   rÊ   )r/   rË   rÖ   r×   rØ   rÅ   rï   s          r   r8   ÚLongCatAudioDiTVae.forward…  sN   € ð. —+‘+˜fÐUYÐen�+Ðo×wÑwˆØ—+‘+˜g°4�+Ð8×?Ñ?ˆÞØ�:ÐÜ.°gÑ>Ð>r   )rÔ   rÓ   )r   r¶   NNr·   r¶   NNr…   r˜   r…   r˜   Fi   iÀ]  g¸…ëQ¸æ?)TTNr:   )FTN)r;   r<   r=   r>   Ú_supports_group_offloadingr   r?   r«   r|   r@   rä   r)   r	   r+   rA   Ú	GeneratorrÄ   Útuplerê   rÊ   rð   r8   rB   rC   rD   s   @r   rÍ   rÍ   (  s  ø† Ø!&Ðàð ØØ$(Ø$(ØØ"%Ø!Ø!%Ø#.Ø!.Ø'Ø(Ø Ø"&Ø Øñ#.
àð.
ð ð.
ð �c‘˜TÑ!ð	.
ð
 �c‘˜TÑ!ð.
ð ð.
ð  ð.
ð �d‘
ð.
ð ˜$‘;ð.
ð !ð.
ð ð.
ð ð.
ð ð.
ð ð.
ð  ð.
ð  ð!.
ð" ÷#.
ó ð.
ð` ð "&Ø Ø,0ñ@à—‘ð@ð ð@ð ð	@ð
 —?‘? TÑ)ð@ð 
)¨5°·±Ñ+>Ñ	>ô@ó ð@ð2 à9=ñ?Ø—|‘|ð?Ø26ð?à	(¨5°·±Ñ+>Ñ	>ô?ó ð?ð" "'Ø Ø,0ñ?à—‘ð?ð ð?ð ð	?ð
 —?‘? TÑ)ð?ð 
)¨5°·±Ñ+>Ñ	>÷?ó ?r   rÍ   )r   r   r   T)r   )+r‡   Údataclassesr   r+   Útorch.nnr   Útorch.nn.functionalÚ
functionalrà   Útorch.nn.utilsr   Úconfiguration_utilsr   r   Úutilsr   Úutils.accelerate_utilsr	   Úutils.torch_utilsr
   Úmodeling_utilsr   Úvaer   r   r!   ÚModuler#   r|   r?   rL   rA   rU   rW   re   ro   r~   r”   r�   r¹   rÄ   rÊ   rÍ   rÆ   r   r   Ú<module>r     sJ  ðó$ Ý !ã Ý ß Ð Ý &ç BÝ Ý 8Ý -Ý 'Ý !ôpò<ô_ˆb�i‰iô _ñ ˜cð ¨Sð ¸¿¹õ ð U§\¡\ð ¸3ð À5Ç<Á<ô ôo˜Ÿ™ô oô$=�r—y‘yô =ô:�b—i‘iô :ô"$�b—i‘iô $ô@$�b—i‘iô $ô@.$˜Ÿ™ô .$ôb..˜Ÿ™ô ..ðb ô jó ó ðð ô jó ó ðôx?˜Ð%5°{õ x?r   