
    { jA'                    L   S SK Jr  S SKJrJrJrJr  S SKrS SK	r	\(       a  S SK	J
r
  SSKJrJrJrJr  \R                   " SSS	S
5       " S S\5      5       r\R                   " SSS	5       " S S\5      5       r " S S\5      r\R                   " S
5       " S S\\5      5       rg)    )annotations)AnyCallableIterableTYPE_CHECKINGN)Tensor   )MmprojModel	ModelBase	TextModelggufQwen2VLModelQwen2VLForConditionalGeneration"Qwen2_5_VLForConditionalGenerationQwen2_5OmniModelc                  t   ^  \ rS rSr\R
                  R                  rU 4S jrS r	\
SU 4S jj5       rSrU =r$ )r      c                "   > [         TU ]  5         g N)superset_gguf_parametersself	__class__s    -/mnt/workspace/llama.cpp/conversion/qwenvl.pyr    Qwen2VLModel.set_gguf_parameters   s    #%    c                .    U R                  5         g ! [         a     Of = fU R                  5       u  pnU R                  R	                  S5        U R                  R                  U5        U R                  R                  U5        U R                  R                  U5        [        R                  " U R                  SS9nS H4  nU R                  R                  US-   5      nUc  M&  XdR                  U'   M6     UR                  U R                  5        g )Ngpt2T)load_merges)boseos	_token_id)_set_vocab_sentencepieceFileNotFoundErrorget_vocab_basegguf_writeradd_tokenizer_modeladd_tokenizer_preadd_token_listadd_token_typesr   SpecialVocab	dir_modelhparamsgetspecial_token_idsadd_to_gguf)r   tokenstoktypestokprespecial_vocabkeytids          r   	set_vocabQwen2VLModel.set_vocab   s    	))+  		 $(#6#6#8 &,,V4**62''/((2))$..dK "C,,""3#45C7://4 "
 	!!$"2"23s    
  c                   > Uu  p#UR                  S5      (       a  UR                  SSS5      nUR                  S5      (       a  UR                  SS5      n[        TU ]  X#45      $ )Ntransformer. r	   thinker.
startswithreplacer   filter_tensorsclsitemnamegenr   s       r   rA   Qwen2VLModel.filter_tensors2   s^    	 ??>**<<A6D??:&&<<
B/Dw%tk22r    rD   z tuple[str, Callable[[], Tensor]]returnz'tuple[str, Callable[[], Tensor]] | None)__name__
__module____qualname____firstlineno__r   
MODEL_ARCHQWEN2VL
model_archr   r8   classmethodrA   __static_attributes____classcell__r   s   @r   r   r      s2     ((J&40 
3 
3r   c                  j   ^  \ rS rSrU 4S jrU 4S jrU 4S jr\SU 4S jj5       rS	U 4S jjr	Sr
U =r$ )
Qwen2VLVisionModel@   c                  > [         TU ]  " U0 UD6  U R                  c   eU R                  R                  SS5      U R                  S'   U R                  R                  S5      U R                  S'   U R                  R                  S5      U R                  S'   SU R                  ;   aQ  U R                  R                  S5      U R                  S	'   U R                  R                  S5      U R                  S'   g g )
N
image_sizei0  	num_headsnum_attention_headsdepthnum_hidden_layers	embed_dimhidden_sizeintermediate_size)r   __init__hparams_visionr/   r   argskwargsr   s      r   rb   Qwen2VLVisionModel.__init__B   s    $)&)""...,0,?,?,C,CLRU,VL)595H5H5L5L[5Y12373F3F3J3J73S/0$---7;7J7J7N7N}7]D 34151D1D1H1H1UD. .r   c                  > [         TU ]  5         U R                  c   eU R                  nU R                  S   nUS:X  a5  U R                  R                  [        R                  R                  5        GO)US:X  d  US:X  Ga  US:X  a4  U R                  R                  [        R                  R                  5        O3U R                  R                  [        R                  R                  5        U R                  R                  S5        UR                  S5      nUc   S5       eUS   S	-   n[        S	[        U5      5       H!  nX5   X5S	-
     -
  U:w  d  M  [        S
U 35      e   U R                  R!                  U5        O[        SU R                  S    35      eU R                  R#                  U R                  R                  SS5      5        g )N
model_typeqwen2_vl
qwen2_5_vlqwen2_5_omniTfullatt_block_indexesz0fullatt_block_indexes is required for qwen2_5_vlr   r	   zInvalid fullatt_block_indexes: zUnknown QwenVL model type: rms_norm_epsgư>)r   r   rc   global_configr'   add_clip_projector_typer   VisionProjectorTyperP   QWEN25OQWEN25VLadd_vision_use_silur/   rangelen
ValueErroradd_vision_n_wa_pattern"add_vision_attention_layernorm_eps)r   r.   ri   rm   n_wa_patternir   s         r   r   &Qwen2VLVisionModel.set_gguf_parametersM   s   #%""...%%''5
#44T5M5M5U5UV<':+G^+  889Q9Q9Y9YZ  889Q9Q9Z9Z[006$+KK0G$H!(4h6hh403a7L1c"789(+.CE.JJlZ$'FG\F]%^__ : 44\B:4;M;Ml;[:\]^^;;D<N<N<R<RSacg<hir   c                d   > SU;   a  [         R                  R                  $ [        TU ]  XX45      $ )Nz.position_embd.)r   GGMLQuantizationTypeF32r   tensor_force_quantr   rE   new_namebidn_dimsr   s        r   r   %Qwen2VLVisionModel.tensor_force_quanth   s/    (,,000w)$#FFr   c                   > Uu  p#UR                  S5      (       a  UR                  SSS5      nUR                  S5      (       d  g [        TU ]  X#45      $ )Nr;   r<   r	   visual.r>   rB   s       r   rA   !Qwen2VLVisionModel.filter_tensorsm   sQ    	 ??>**<<A6Dy))w%tk22r   c              #  4  >#    SU;   a  UR                   S:X  a  UR                  u  pEOUR                  S   nUS-  S:X  d   eUS-  nUS U nXUS-   nXS-  S  n	[        TU ]  XrR	                  SS5      U5       S h  vN   [        TU ]  XR	                  SS5      U5       S h  vN   [        TU ]  XR	                  SS5      U5       S h  vN   g S	U;   a  UR                  u  ppnA
AAAUS:X  d   S
5       e[
        R                  [
        R                  R                     S-   US S 2S S 2SS4   4v   [
        R                  [
        R                  R                     S-   US S 2S S 2SS4   4v   g [        TU ]  XU5       S h  vN   g  GN N N N7f)Nz.qkv.   r      qkvqkvzpatch_embed.proj.weightz<Current implementation only support temporal_patch_size of 2.weight.z	.weight.1r	   )	ndimshaper   modify_tensorsr@   r   TENSOR_NAMESMODEL_TENSORV_ENC_EMBD_PATCH)r   
data_torchrE   r   c3_cwqwkwvc1c2ktkhkwr   s                  r   r   !Qwen2VLVisionModel.modify_tensorsz   s    d?!#"((A%%a(6Q;;aABQBq1u%BEF#Bw-b,,uc2JCPPPw-b,,uc2JCPPPw-b,,uc2JCPPP&$.!+!1!1BBBBB7ZZZ7$$T%6%6%G%GH9TXbcdfgijlocoXpqq$$T%6%6%G%GH;VXbcdfgijlocoXpqqw-jDDD QPP EsI   A;F>F?'F&F''FFB:F	F
FFFFrH   rI   r   r   rE   strr   z
int | NonerJ   Iterable[tuple[str, Tensor]])rK   rL   rM   rN   rb   r   r   rR   rA   r   rS   rT   rU   s   @r   rW   rW   @   s5    	Vj6G
 
3 
3E Er   rW   c                  X   ^  \ rS rSrSrU 4S jrU 4S jrS	S jrU 4S jrS
S jr	Sr
U =r$ )Qwen25AudioModel   Tc                   > [         TU ]  " U0 UD6  U R                  c   eU R                  S   U R                  S'   U R                  S   U R                  S'   U R                  S   U R                  S'   g )Nd_modelr`   encoder_ffn_dimra   encoder_attention_headsr\   )r   rb   hparams_audiord   s      r   rb   Qwen25AudioModel.__init__   s|    $)&)!!---,0,>,>y,I=)262D2DEV2W./484F4FG`4a01r   c                   > [         TU ]  5         U R                  c   eU R                  R	                  U R                  S   5        U R                  R                  U R                  R                  SS5      5        g )Nnum_mel_binslayer_norm_epsgh㈵>)r   r   r   r'   add_audio_num_mel_bins!add_audio_attention_layernorm_epsr/   r   s    r   r   $Qwen25AudioModel.set_gguf_parameters   sh    #%!!---//0B0B>0RS::4;M;M;Q;QRbdh;ijr   c              #  R  #    U R                   c   eSnSnU R                   S   n[        R                  " U5      US-  S-
  -  n[        R                  " U* [        R
                  " US-  5      R                  5       -  5      n[        R
                  " U5      S S 2[        R                  4   U[        R                  S S 24   -  n[        R                  " [        R                  " U5      [        R                  " U5      /SS9R                  [        R                  S9nSU4v   g 7f)	Ni'  i  r`   r   r	   )dim)dtypez"audio_tower.embed_positions.weight)r   nplogtorchexparangefloatnewaxiscatsincostofloat32)r   max_timescalelengthchannelslog_timescale_incrementinv_timescalesscaled_timepos_embds           r   generate_extra_tensors'Qwen25AudioModel.generate_extra_tensors   s     !!---%%m4"$&&"78q=1;L"M$;#;ell8WX=>Y>_>_>a#abll6*1bjj=9N2::WX=<YY99eii4eii6LMSTUXX_d_l_lXm3X>>s   D%D'c                p   > SU;   a   SU;   a  [         R                  R                  $ [        TU ]  XX45      $ )Nz.convr   )r   r~   F16r   r   r   s        r   r   #Qwen25AudioModel.tensor_force_quant   s5    d?yD0,,000w)$#FFr   c              #     #    SU;   d  SU;   a  UR                  S5      n[        R                  " XX#5       S h  vN   g  N7f)Nz
conv1.biasz
conv2.bias)	unsqueezer
   r   r   r   rE   r   s       r   r   Qwen25AudioModel.modify_tensors   s:     4<4#7#--b1J--dJJJs   8AA ArH   )rJ   r   r   )rK   rL   rM   rN   has_audio_encoderrb   r   r   r   r   rS   rT   rU   s   @r   r   r      s,    bk
?G
K Kr   r   c                  b   ^  \ rS rSrSrSrS	S jrS	S jrU 4S jr\	S
S j5       r
SS jrSrU =r$ )Qwen25OmniModel   Tc                >    U R                   S   R                  S5      $ )Nthinker_configvision_configro   r/   r   s    r   get_vision_config!Qwen25OmniModel.get_vision_config   s    !!"2377HHr   c                >    U R                   S   R                  S5      $ )Nr   audio_configr   r   s    r   get_audio_config Qwen25OmniModel.get_audio_config   s    !!"2377GGr   c                   > [         TU ]  5         U R                  R                  [        R
                  R                  5        g r   )r   r   r'   rp   r   rq   rr   r   s    r   r   #Qwen25OmniModel.set_gguf_parameters   s.    #%001I1I1Q1QRr   c                    Uu  p#UR                  S5      (       d  UR                  S5      (       d  g UR                  S5      (       a  UR                  SS5      nSU;   a  g [        R                  " X#45      $ )Nr   audio_tower.r=   r<   audio_bos_eos_token)r?   r@   r
   rA   )rC   rD   rE   rF   s       r   rA   Qwen25OmniModel.filter_tensors   sh    	y))$//.2Q2Q??:&&<<
B/D D( ))4+66r   c              #     #    SU;   a  [         R                  XX#5       S h  vN   g SU;   a  [        R                  XX#5       S h  vN   g  N* N7f)Nr   r   )rW   r   r   r   s       r   r   Qwen25OmniModel.modify_tensors   sT     )884UUU 	 t#'66tSSS VSs!    AA$AAAArH   )rJ   zdict[str, Any] | NonerI   r   )rK   rL   rM   rN   r   has_vision_encoderr   r   r   rR   rA   r   rS   rT   rU   s   @r   r   r      s?    IHS 7 7  r   r   )
__future__r   typingr   r   r   r   numpyr   r   r   baser
   r   r   r   registerr   rW   r   r   rH   r   r   <module>r      s    " 9 9   9 9 %(	)39 )3)3X N$EGklPE PE mPEf&K{ &KR &'$(*: $ ($r   