ó
    ÷“,jS+  ã                   óÐ   • S r SSKJr  SSKrSSKrSSKJrJr  SSK	J
r
Jr  \R                  " 5       S\S\4S	 j5       r\R                  " 5       SS\4S
 jj5       rS\\\4   4S jrg)zWConversion functions for 3rd part state-dicts and non-torch native checkpoint formats.
é    )ÚUnionNé   )ÚCLIPÚCustomTextCLIP)ÚTextTransformerÚTransformerÚmodelÚcheckpoint_pathc                 óÜ  ^^^^^	^
^• SSK Jm
Jm	  SS jm[        R                  " U5      mSmSmUUUU	U
U4S jnS[
        4UU4S	 jjmS[        4UUU4S
 jjnST;   a  SOSnU" U R                  R                  U S35        U" U R                  U S35        U R                  R                  T" TU S3   5      S   5        U R                  R                  T" TU S3   5      S   5        g)zÅLoad weights from .npz checkpoints for official Google big_vision image-text models

Currently, the SigLIP source models are supported and a CustomTextCLIP destination model
w/ timm image encoder.
r   )Úresample_patch_embedÚresample_abs_pos_embedNc                 óÜ  • Ub  X   n U R                   S:X  aN  U R                  S   U R                  S   s=:X  a  U R                  S   s=:X  a  S:X  a  O  OU R                  5       n U(       ak  U R                   S:X  a  U R                  / SQ5      n OGU R                   S:X  a  U R                  / SQ5      n O#U R                   S:X  a  U R                  SS/5      n [        R
                  " U 5      $ )Né   r   r   é   )é   r   r   r   r   )r   r   r   )ÚndimÚshapeÚflattenÚ	transposeÚtorchÚ
from_numpy)ÚwÚtÚidxs      ÚŸC:\Users\johnn\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.13_qbz5n2kfra8p0\LocalCache\local-packages\Python313\site-packages\open_clip/convert.pyÚ_n2pÚ%load_big_vision_weights.<locals>._n2p   s°   € Ø‰?Ø‘ˆAØ�6‰6�Q‹;˜1Ÿ7™7 1™:¨¯©°©ÕF°q·w±w¸q±zÕFÀQÖFØ—	‘	“ˆAÞØ�v‰v˜‹{Ø—K‘K¢Ó-‘Ø—‘˜1“Ø—K‘K¢	Ó*‘Ø—‘˜1“Ø—K‘K  A Ó'�Ü×Ò Ó"Ð"ó    ÚbilinearFc                 ó^  >• T" TU S3   5      nUR                   SS  U R                  R                  R                  R                   SS  :w  a5  T" UU R                  R                  R                  R                   SS  TTSS9nU R                  R                  R                  R	                  U5        U R                  R                  R
                  R	                  T" TU S3   5      5        U R                  b&  U R                  R	                  T" TU S3   SS95        T" TU S	3   SS9nUR                   U R                  R                   :w  a*   UR                    S
U R                  R                    35       eU R                  R	                  U5        Su  pVn[        U R                  R                  5       5       GHý  u  p‰U S3T;   a  U S3n
UnO
U SU S3n
S nU
SU S3-   nU	R                  R                  R	                  T" TU
 S3   US95        U	R                  R
                  R	                  T" TU
 S3   US95        U	R                  R                  R                  R	                  [         R"                  " S Vs/ s H,  nT" TU U S3   SUS9R%                  S5      R&                  PM.     sn5      5        U	R                  R                  R
                  R	                  [         R"                  " S Vs/ s H"  nT" TU U S3   SUS9R)                  S5      PM$     sn5      5        U	R                  R                  R                  R	                  T" TU S3   US9R%                  S5      5        U	R                  R                  R
                  R	                  T" TU S3   US95        U	R*                  R                  R	                  T" TU
 SU S3   US95        U	R*                  R
                  R	                  T" TU
 SU S3   US95        [-        S5       H�  n[        U	R.                  SUS-    35      R                  R	                  T" TU
 S U S!U S3   US95        [        U	R.                  SUS-    35      R
                  R	                  T" TU
 S U S!U S3   US95        M‘     GM      U R0                  R                  R	                  T" TU S"3   5      5        U R0                  R
                  R	                  T" TU S#3   5      5        U R2                  Gbj  U S$3n
U
S%-   nU R2                  R4                  R	                  T" TU
 S&3   SS95        U R2                  R6                  R                  R	                  T" TU S'3   SS9R%                  S5      R&                  5        U R2                  R6                  R
                  R	                  T" TU S(3   SS9R)                  S5      5        U R2                  R8                  R                  R	                  [         R"                  " S) Vs/ s H+  nT" TU U S3   SS9R%                  S5      R&                  PM-     sn5      5        U R2                  R8                  R
                  R	                  [         R"                  " S) Vs/ s H!  nT" TU U S3   SS9R)                  S5      PM#     sn5      5        U R2                  R                  R                  R	                  T" TU S3   5      R%                  S5      5        U R2                  R                  R
                  R	                  T" TU S3   5      5        U R2                  R0                  R                  R	                  T" TU
 S3   5      5        U R2                  R0                  R
                  R	                  T" TU
 S3   5      5        [-        S5       HŸ  n[        U R2                  R.                  SUS-    35      R                  R	                  T" TU
 S*U S3   5      5        [        U R2                  R.                  SUS-    35      R
                  R	                  T" TU
 S*U S3   5      5        M¡     g g s  snf s  snf s  snf s  snf )+Nzembedding/kerneléþÿÿÿT)ÚinterpolationÚ	antialiasÚverbosezembedding/biasÚclsF©r   Úpos_embeddingz, r   )r   r   r   z*Transformer/encoderblock/LayerNorm_0/scalezTransformer/encoderblock/zTransformer/encoderblock_Ú/ÚMultiHeadDotProductAttention_úLayerNorm_0/scale©r   úLayerNorm_0/bias©ÚqueryÚkeyÚvalueú/kernel©r   r   ú/biaséÿÿÿÿú
out/kernelúout/biasÚ
LayerNorm_z/scaler   ÚfcÚ	MlpBlock_z/Dense_zTransformer/encoder_norm/scalezTransformer/encoder_norm/biasz
MAPHead_0/úMultiHeadDotProductAttention_0/Úprobezquery/kernelz
query/bias)r/   r0   zMlpBlock_0/Dense_)r   Úpatch_embedÚprojÚweightÚcopy_ÚbiasÚ	cls_tokenÚ	pos_embedÚgetattrÚ	grid_sizeÚ	enumerateÚblocksÚchildrenÚnorm1ÚattnÚqkvr   Úcatr   ÚTÚreshapeÚnorm2ÚrangeÚmlpÚnormÚ	attn_poolÚlatentÚqÚkv)ÚmoduleÚprefixÚembed_conv_wÚpos_embed_wÚnum_prefix_tokensÚmha_subÚb_subÚln1_subÚiÚblockÚblock_prefixr   Ú
mha_prefixÚnÚrr   r#   r"   r   r   r   s                  €€€€€€r   Ú_convert_timm_imgÚ2load_big_vision_weights.<locals>._convert_timm_img'   s6  ø€ Ù˜A  Ð(8Ð9Ñ:Ó;ˆØ×Ñ˜b˜cÐ" f×&8Ñ&8×&=Ñ&=×&DÑ&D×&JÑ&JÈ2È3Ð&OÓOÙ/ØØ×"Ñ"×'Ñ'×.Ñ.×4Ñ4°R°SÐ9Ø+Ø#ØñˆLð 	×Ñ×Ñ×&Ñ&×,Ñ,¨\Ô:Ø×Ñ×Ñ×$Ñ$×*Ñ*©4°°V°H¸NÐ2KÑ0LÓ+MÔNà×ÑÑ'Ø×Ñ×"Ñ"¡4¨¨V¨H°C¨.Ñ(9¸UÑ#CÔDá˜1 ˜x }Ð5Ñ6¸%Ñ@ˆØ×Ñ × 0Ñ 0× 6Ñ 6Ó6ØJ˜[×.Ñ.Ð/¨r°&×2BÑ2B×2HÑ2HÐ1IÐJÓJ�5ð 	×Ñ×Ñ˜{Ô+à"+Ñˆ˜Ü! &§-¡-×"8Ñ"8Ó":×;‰HˆAØ�ÐCÐDÈÓIØ"( Ð)BÐC�Ø‘à"( Ð)BÀ1À#ÀQÐG�Ø�Ø%Ð*GÈÀyÐPQÐ(RÑRˆJØ�K‰K×Ñ×$Ñ$¡T¨!¨|¨nÐ<MÐ,NÑ*OÐUXÑ%YÔZØ�K‰K×Ñ×"Ñ"¡4¨¨\¨NÐ:JÐ*KÑ(LÐRUÑ#VÔWØ�J‰J�N‰N×!Ñ!×'Ñ'¬¯	ª	Ù[tó3vÚ[tÐVW‘�Q˜*˜ a S¨Ð0Ñ1°UÀÑD×LÑLÈQÓO×QÔQÑ[tñ3vó )wô xà�J‰J�N‰N×Ñ×%Ñ%¤e§i¢iÙXqó1sÚXqÐST‘�Q˜*˜ a S¨Ð.Ñ/°5¸cÑB×JÑJÈ2ÖNÑXqñ1só 'tô uà�J‰J�O‰O×"Ñ"×(Ñ(©¨a°:°,¸jÐ0IÑ.JÐPSÑ)T×)\Ñ)\Ð]^Ó)_Ô`Ø�J‰J�O‰O× Ñ ×&Ñ&¡t¨A°°¸HÐ.EÑ,FÈCÑ'PÔQØ�K‰K×Ñ×$Ñ$¡T¨!¨|¨n¸JÀwÀiÈvÐ,VÑ*WÐ]`Ñ%aÔbØ�K‰K×Ñ×"Ñ"¡4¨¨\¨N¸*ÀWÀIÈUÐ*SÑ(TÐZ]Ñ#^Ô_Ü˜1–X�Ü˜Ÿ	™	 R¨¨A© w <Ó0×7Ñ7×=Ñ=Ù˜˜l˜^¨9°U°G¸7À1À#ÀWÐMÑNÐTWÑXôZä˜Ÿ	™	 R¨¨A© w <Ó0×5Ñ5×;Ñ;Ù˜˜l˜^¨9°U°G¸7À1À#ÀUÐKÑLÐRUÑVöXô ñ% <ð0 	�‰×Ñ× Ñ ¡ a¨6¨(Ð2PÐ(QÑ&RÓ!SÔTØ�‰×Ñ×Ñ™t A¨¨Ð0MÐ&NÑ$OÓPÔQà×ÑÒ'Ø$˜X ZÐ0ˆLØ%Ð*IÑJˆJØ×Ñ×#Ñ#×)Ñ)©$¨q°L°>ÀÐ1GÑ/HÈEÑ*RÔSØ×Ñ×Ñ×%Ñ%×+Ñ+©D°°j°\ÀÐ3NÑ1OÐSXÑ,Y×,aÑ,aÐbcÓ,d×,fÑ,fÔgØ×Ñ×Ñ×#Ñ#×)Ñ)©$¨q°J°<¸zÐ1JÑ/KÈuÑ*U×*]Ñ*]Ð^`Ó*aÔbØ×Ñ×Ñ×&Ñ&×,Ñ,¬U¯YªYÙRbó8dÚRbÈQ‘�Q˜*˜ a S¨Ð0Ñ1°UÑ;×CÑCÀAÓF×HÔHÑRbñ8dó .eô fà×Ñ×Ñ×$Ñ$×*Ñ*¬5¯9ª9ÙO_ó6aÚO_È!‘�Q˜*˜ a S¨Ð.Ñ/°5Ñ9×AÑAÀ"ÖEÑO_ñ6aó ,bô cà×Ñ×!Ñ!×(Ñ(×.Ñ.©t°A¸¸ÀJÐ6OÑ4PÓ/Q×/YÑ/YÐZ[Ó/\Ô]Ø×Ñ×!Ñ!×&Ñ&×,Ñ,©T°!°z°lÀ(Ð4KÑ2LÓ-MÔNØ×Ñ×!Ñ!×(Ñ(×.Ñ.©t°A¸¸ÐFWÐ6XÑ4YÓ/ZÔ[Ø×Ñ×!Ñ!×&Ñ&×,Ñ,©T°!°|°nÐDTÐ4UÑ2VÓ-WÔXÜ˜1–X�Ü˜×(Ñ(×,Ñ,°°1°q±5°'¨lÓ;×BÑB×HÑHÉÈaÐS_ÐR`Ð`qÐrsÐqtÐt{ÐP|ÑN}ÓI~ÔÜ˜×(Ñ(×,Ñ,°°1°q±5°'¨lÓ;×@Ñ@×FÑFÁtÈAÐQ]ÐP^Ð^oÐpqÐorÐrwÐNxÑLyÓGzÖ{ò ð (ùò#3vùò1sùò*8dùò6as   Ê3bÌ )b Ù2b%Ú=(b*rV   c                 ó¼  >• [        U R                  R                  5       5       GH)  u  p#U S3T	;   a  U S3nUnO
U SU S3nS nUS-   nUR                  R                  R                  T" T	U S3   US95        UR                  R                  R                  T" T	U S3   US95        UR                  R                  R                  [        R                  " S	 Vs/ s H,  nT" T	U U S
3   SUS9R                  S5      R                  PM.     sn5      5        UR                  R                  R                  [        R                  " S	 Vs/ s H"  nT" T	U U S3   SUS9R                  S5      PM$     sn5      5        UR                  R                  R                  R                  T" T	U S3   US9R                  S5      5        UR                  R                  R                  R                  T" T	U S3   US95        UR                   R                  R                  T" T	U S3   US95        UR                   R                  R                  T" T	U S3   US95        UR"                  R$                  R                  R                  T" T	U S3   US95        UR"                  R$                  R                  R                  T" T	U S3   US95        UR"                  R&                  R                  R                  T" T	U S3   US95        UR"                  R&                  R                  R                  T" T	U S3   US95        GM,     g s  snf s  snf )Nzencoderblock/LayerNorm_0/scalezencoderblock/Úencoderblock_r(   r:   r*   r+   r,   r-   r1   Fr2   r   r3   r4   r5   r6   zLayerNorm_1/scalezLayerNorm_1/biaszMlpBlock_0/Dense_0/kernelzMlpBlock_0/Dense_0/biaszMlpBlock_0/Dense_1/kernelzMlpBlock_0/Dense_1/bias)rE   Ú	resblocksrG   Úln_1r>   r?   r@   rI   Úin_proj_weightr   rK   r   rL   Úin_proj_biasrM   Úout_projÚln_2rP   Úc_fcÚc_proj)
rV   rW   r^   r_   r`   r   ra   rb   r   r   s
           €€r   Ú_convert_openclip_transformerÚ>load_big_vision_weights.<locals>._convert_openclip_transformers   s(  ø€ Ü! &×"2Ñ"2×";Ñ";Ó"=×>‰HˆAØ�Ð7Ð8¸AÓ=Ø"( ¨Ð7�Ø‘à"( ¨°q°c¸Ð;�Ø�Ø%Ð*IÑJˆJØ�J‰J×Ñ×#Ñ#¡D¨¨l¨^Ð;LÐ+MÑ)NÐTWÑ$XÔYØ�J‰J�O‰O×!Ñ!¡$ q¨L¨>Ð9IÐ)JÑ'KÐQTÑ"UÔVØ�J‰J×%Ñ%×+Ñ+¬E¯IªIÙ[tó7vÚ[tÐVW‘�Q˜*˜ a S¨Ð0Ñ1°UÀÑD×LÑLÈQÓO×QÔQÑ[tñ7vó -wô xà�J‰J×#Ñ#×)Ñ)¬%¯)ª)ÙXqó5sÚXqÐST‘�Q˜*˜ a S¨Ð.Ñ/°5¸cÑB×JÑJÈ2ÖNÑXqñ5só +tô uà�J‰J×Ñ×&Ñ&×,Ñ,©T°!°z°lÀ*Ð4MÑ2NÐTWÑ-X×-`Ñ-`ÐabÓ-cÔdØ�J‰J×Ñ×$Ñ$×*Ñ*©4°°Z°LÀÐ2IÑ0JÐPSÑ+TÔUØ�J‰J×Ñ×#Ñ#¡D¨¨l¨^Ð;LÐ+MÑ)NÐTWÑ$XÔYØ�J‰J�O‰O×!Ñ!¡$ q¨L¨>Ð9IÐ)JÑ'KÐQTÑ"UÔVØ�I‰I�N‰N×!Ñ!×'Ñ'©¨Q°,°Ð?XÐ/YÑ-ZÐ`cÑ(dÔeØ�I‰I�N‰N×Ñ×%Ñ%¡d¨1°¨~Ð=TÐ-UÑ+VÐ\_Ñ&`ÔaØ�I‰I×Ñ×#Ñ#×)Ñ)©$¨q°L°>ÐAZÐ1[Ñ/\ÐbeÑ*fÔgØ�I‰I×Ñ×!Ñ!×'Ñ'©¨Q°,°Ð?VÐ/WÑ-XÐ^aÑ(b×cò+ ?ùò7vùò5ss   Ã3MÅ)Mc                 ó   >• U R                   R                  R                  T" TU S3   SS95        T" TU S3   SS9R                  S5      nU R                  R                  U5        T" U R
                  US-   S9  U R                  R                  R                  T" TU S3   5      5        U R                  R                  R                  T" TU S	3   5      5        U R                  bc  U R                  R                  R                  T" TU S
3   5      5        U R                  R                  R                  T" TU S3   5      5        g g )NzEmbed_0/embeddingFr&   r'   r   z
Encoder_0/)rW   zEncoder_0/encoder_norm/scalezEncoder_0/encoder_norm/biaszhead/kernelz	head/bias)	Útoken_embeddingr>   r?   ÚsqueezeÚpositional_embeddingÚtransformerÚln_finalr@   Útext_projection)rV   rW   rY   rp   r   r   s      €€€r   Ú_convert_openclip_txtÚ6load_big_vision_weights.<locals>._convert_openclip_txt‹   s9  ø€ Ø×Ñ×%Ñ%×+Ñ+©D°°f°XÐ=NÐ3OÑ1PÐTYÑ,ZÔ[Ù˜1 ˜x }Ð5Ñ6¸%Ñ@×HÑHÈÓKˆØ×#Ñ#×)Ñ)¨+Ô6Ù% f×&8Ñ&8ÀÈ,ÑAVÒWØ�‰×Ñ×$Ñ$¡T¨!¨v¨hÐ6RÐ,SÑ*TÓ%UÔVØ�‰×Ñ×"Ñ"¡4¨¨V¨HÐ4OÐ*PÑ(QÓ#RÔSØ×!Ñ!Ñ-Ø×"Ñ"×)Ñ)×/Ñ/±°Q¸&¸ÀÐ7MÑ5NÓ0OÔPØ×"Ñ"×'Ñ'×-Ñ-©d°1¸°x¸yÐ5IÑ3JÓ.KÕLð .r   zparams/bzparams/Ú zimg/ztxt/Úbr   )TN)Útimm.layersr   r   ÚnpÚloadr   r   ÚvisualÚtrunkÚtextÚ
logit_biasr?   Úlogit_scale)r	   r
   rd   ry   Úroot_prefixrp   r   r#   r"   r   r   r   s        @@@@@@@r   Úload_big_vision_weightsr†      sò   þ€ ÷ Iô#ô 	�Š�Ó €AØ€MØ€I÷J|ò J|ðXd¬k÷ dð dð0	M¤o÷ 	Mñ 	Mð  *¨Q›‘)°B€KÙ�e—l‘l×(Ñ(¨[¨M¸Ð*>Ô?Ù˜%Ÿ*™*¨¨°TÐ&:Ô;Ø	×Ñ×Ñ™4  [ M°Ð"3Ñ 4Ó5°aÑ8Ô9Ø	×Ñ×Ñ™D  k ]°!Ð#4Ñ!5Ó6°qÑ9Õ:r   c                 ó^   ^ ^• UU 4S jnSS jnU" U5      nU" U5      n0 UEUEnUS   US'   U$ )Nc                 óÆ   >• T(       a  SSK Jn  OSSKJn  U" U TR                  R                  5      nUR                  5        VVs0 s H  u  p4SU-   U_M     nnnU$ s  snnf )Nr   )Úcheckpoint_filter_fnzvisual.trunk.)Útimm.models.fastvitr‰   Ú%timm.models.vision_transformer_hybridr€   r�   Úitems)Ú
state_dictr‰   Útimm_state_dictÚkÚvÚfastvitr	   s        €€r   rd   Ú9convert_mobile_clip_state_dict.<locals>._convert_timm_img    sY   ø€ ÞÞ@åRÙ.¨z¸5¿<¹<×;MÑ;MÓNˆØ>M×>SÑ>SÔ>UÔVÒ>U±d°a˜?¨QÑ.°Ò1Ñ>UˆÑVØÐùó Ws   ÁAc                 óž  • 0 nU R                  5        GH5  u  p4UR                  U5      (       d  M  UR                  US5      nUR                  SS5      nUR                  SS5      nUR                  S5      (       a"  UR                  SS5      nUR                  5       nUR                  SS	5      nUR                  S
S5      nUR                  SS5      nUR                  SS5      nUR                  SS5      nUR                  SS5      nUR                  SS5      nUR                  SS5      nUR                  SS5      nXBSU-   '   GM8     U$ )Nr{   Úprojection_layerrx   Úembedding_layerrs   z(positional_embedding.pos_embed.pos_embedru   Úfinal_layer_normrw   zpre_norm_mha.0ri   zpre_norm_mha.1rI   zpre_norm_ffn.0rm   zpre_norm_ffn.1zmlp.c_fczpre_norm_ffn.4z
mlp.c_projzqkv_proj.weightrj   zqkv_proj.biasrk   ztransformer.ztransformer.resblocks.ztext.)rŒ   Ú
startswithÚreplacert   )r�   rW   Ú	text_dictr�   r�   s        r   ry   Ú=convert_mobile_clip_state_dict.<locals>._convert_openclip_txt©   s<  € Øˆ	Ø×$Ñ$×&‰DˆAØ—<‘< ×'Ñ'ÙØ—	‘	˜& "Ó%ˆAØ—	‘	Ð,Ð.?Ó@ˆAØ—	‘	Ð+Ð->Ó?ˆAØ�|‰|ÐF×GÑGØ—I‘IÐHÐJ`Óa�Ø—I‘I“K�Ø—	‘	Ð,¨jÓ9ˆAØ—	‘	Ð*¨FÓ3ˆAØ—	‘	Ð*¨FÓ3ˆAØ—	‘	Ð*¨FÓ3ˆAØ—	‘	Ð*¨JÓ7ˆAØ—	‘	Ð*¨LÓ9ˆAØ—	‘	Ð+Ð-=Ó>ˆAØ—	‘	˜/¨>Ó:ˆAØ—	‘	˜.Ð*BÓCˆAØ%&�g ‘kÔ"ñ% 'ð& Ðr   r„   )ztext_encoder.© )r	   r�   r‘   rd   ry   Ú
image_dictr™   Úout_dicts   ` `     r   Úconvert_mobile_clip_state_dictrž   �   sF   ù€ öôñ. # :Ó.€JÙ% jÓ1€IØ*�*Ð* 	Ð*€HØ(¨Ñ7€Hˆ]ÑØ€Or   c                 óH   • SU;   a  [        X5      nSU;   a
  [        XSS9nU$ )Nz8image_encoder.model.patch_embed.0.rbr_conv.0.conv.weightz1image_encoder.model.patch_emb.0.block.conv.weightF)r‘   )rž   )r	   r�   s     r   Úconvert_state_dictr    Ç   s.   € ØAÀZÓOä3°EÓFˆ
Ø:¸jÓHä3°EÈuÑUˆ
ØÐr   )T)Ú__doc__Útypingr   r   Únumpyr~   r	   r   r   rv   r   r   Úno_gradÚstrr†   rž   r    r›   r   r   Ú<module>r¦      s~   ðñå ã Û ç 'ß 5ð ‡‚ƒðM; >ð M;ÀCó M;ó ðM;ð` ‡‚ƒñ&¨.ô &ó ð&ðR˜e N°DÐ$8Ñ9õ r   