+
    UV-j€,  ã                   ó  € ^ RI t ^ RIt^ RIt^ RIt^ RIHtHt ^ RIHt	 ^ RI
Ht ^ RIHt ^ RIHt ^ RIHt ^ RIHt ^ R	IHtHtHtHtHtHtHtHtHtHt ^ R
IH t  ^ RI!H"t" ^ RI#H$t$  ! R R4      t%R t&]'R8X  d
   ]&! 4        R# R# )é    N)ÚDictÚList)Úprint)ÚConsole)ÚMarkdown)ÚPanel)ÚPrompt)Úload)
ÚDEFAULT_KV_GROUP_SIZEÚDEFAULT_KV_QUANT_SCHEMEÚDEFAULT_MAX_TOKENSÚDEFAULT_PREFILL_STEP_SIZEÚDEFAULT_QUANTIZED_KV_STARTÚDEFAULT_TEMPERATUREÚDEFAULT_THINKING_END_TOKENÚDEFAULT_THINKING_START_TOKENÚPromptCacheStateÚstream_generate)Úapply_chat_template)Ú
load_image)ÚVisionFeatureCachec                   ó¢   a € ] tR t^t o RV 3R lR lltV 3R lR ltV 3R lR ltV 3R lR ltV 3R	 lR
 ltV 3R lR lt	V 3R lR lt
RtV tR# )ÚMLXVisionChatc                ó2   <€ V ^8„  d   QhRS[ RS[RS[RS[/# )é   Ú
model_pathÚtemperatureÚ
max_tokensÚverbose)ÚstrÚfloatÚintÚbool)ÚformatÚ__classdict__s   "€Ú]/Volumes/fast/ai/experiments/ui-tars-smoke/.venv/lib/python3.14/site-packages/mlx_vlm/chat.pyÚ__annotate__ÚMLXVisionChat.__annotate__    s3   ø€ ÷ ñ áðñ ðñ ð	ñ
 ñó    c                ó¬  € \        4       V n        W@n        W n        W0n        . V n        R V n        R V n        . V n        \        4       V n
        \        4       V n        WPn        V P                  P                  R4      ;_uu_ 4        \        V4      w  V n        V n        R R R 4       \%        R4       V P'                  4        R #   + '       g   i     L-; i)Nz[bold green]Loading model...z3[bold green]Model loaded successfully![/bold green])r   Úconsoler   r   r   ÚhistoryÚcurrent_imageÚcurrent_image_pathÚimage_pathsr   Úvision_cacher   Úprompt_cache_stateÚstream_kwargsÚstatusr
   ÚmodelÚ	processorÚrprintÚ
print_help)Úselfr   r   r   r   Úkwargss   &&&&&,r&   Ú__init__ÚMLXVisionChat.__init__    s¦   € ô “yˆŒØŒØ&ÔØ$ŒØ#%ˆŒØ!ˆÔØ"&ˆÔØ&(ˆÔÜ.Ó0ˆÔÜ"2Ó"4ˆÔØ#Ôà�\‰\× Ñ Ð!?×@Õ@Ü)-¨jÓ)9Ñ&ˆDŒJ˜œ÷ Aô 	ÐDÔEØ�‰Ö÷	 A×@ús   ÂCÃC	c                ó   <€ V ^8„  d   QhRR/# ©r   ÚreturnN© )r$   r%   s   "€r&   r'   r(   :   s   ø€ ÷ 
Dñ 
D˜Dñ 
Dr)   c                ó8   € Rp\        \        VRRR7      4       R# )zPrint available commands.u&  
[bold yellow]Available Commands:[/bold yellow]
â€¢ /image <path> - Load a new image for discussion
â€¢ /clear - Clear conversation history
â€¢ /help - Show this help message
â€¢ /exit - Exit the chat
â€¢ Any other input will be treated as a question or comment about the current image
        ÚHelpÚblue)ÚtitleÚborder_styleN)r6   r   )r8   Ú	help_texts   & r&   r7   ÚMLXVisionChat.print_help:   s   € ðˆ	ô 	Œu�Y f¸6ÔBÖCr)   c                ó&   <€ V ^8„  d   QhRS[ RS[/# )r   Ú
image_pathr>   ©r    r#   )r$   r%   s   "€r&   r'   r(   F   s   ø€ ÷ ñ ©ð ±ñ r)   c                óp  €  \         P                  P                  V4      '       g   \        RV 24       R# \	        V4      V n        Wn        WP                  9  d   V P                  P                  V4       \        RV 24       R#   \         d#   p\        R\        T4       24        Rp?R# Rp?ii ; i)zJProcess an image and prepare it for the model. Returns True if successful.z2[bold red]Error:[/bold red] Image file not found: Fz%[bold blue]Loaded image:[/bold blue] Tz*[bold red]Error loading image:[/bold red] N)ÚosÚpathÚexistsr6   r   r-   r.   r/   ÚappendÚ	Exceptionr    )r8   rH   Úes   && r&   Úprocess_imageÚMLXVisionChat.process_imageF   sž   € ð	Ü—7‘7—>‘> *×-Ò-ÜØHÈÈÐUôñ ä!+¨JÓ!7ˆDÔØ&0Ô#Ø×!1Ñ!1Ô1Ø× Ñ ×'Ñ'¨
Ô3ÜÐ:¸:¸,ÐGÔHÙøÜô 	ÜÐ?ÄÀAÃ¸xÐHÔIÝûð	ús"   ‚$B §B ·AB ÂB5ÂB0Â0B5c                ó*   <€ V ^8„  d   QhRS[ RS[ RR/# )r   ÚroleÚtextr>   N©r    )r$   r%   s   "€r&   r'   r(   Y   s'   ø€ ÷ @ñ @¡3ð @©cð @°dñ @r)   c                óR   € RRRV/.pV P                   P                  RVRV/4       R# )z*Add a message to the conversation history.ÚtyperU   rT   ÚcontentN)r,   rN   )r8   rT   rU   rY   s   &&& r&   Úadd_to_historyÚMLXVisionChat.add_to_historyY   s.   € à˜F F¨DÐ1Ð2ˆØ�‰×Ñ˜V T¨9°gÐ>Ö?r)   c                ó    <€ V ^8„  d   QhRS[ /# )r   r>   rV   )r$   r%   s   "€r&   r'   r(   ^   s   ø€ ÷ #ñ #¡3ñ #r)   c                ó†  € RV P                   P                  RR4      /pV P                  '       d   ^M^ pV P                  '       d   V P                  .MRp\        V P                  V P
                  P                  V P                  3RV/VB p\        RRRR7       R	p\        V P
                  V P                  V3R
VRV P                  RV P                  RV P                  RV P                  /V P                   B  FA  pWVP                  ,          pV P                  '       g   K)  \        VP                  R	RR7       KC  	  V# )zEGenerate a response from the model based on the conversation history.Úenable_thinkingFNÚ
num_imagesz#[bold green]Assistant:[/bold green]Ú T)ÚendÚflushÚ Úimager   r   r0   r1   )r2   Úgetr.   r   r5   r4   Úconfigr,   r6   r   r   r   r0   r1   rU   r   )r8   Úchat_template_kwargsr_   rd   ÚpromptrU   Úchunks   &      r&   Úgenerate_responseÚMLXVisionChat.generate_response^   sG  € ð ˜t×1Ñ1×5Ñ5Ð6GÈÓOð 
Ðð ×1×1Ð1‘Q°qˆ
Ø-1×-D×-DÐ-D�×(Ñ(Ñ)È$ˆä$Ø�N‰NØ�J‰J×ÑØ�L‰Lñ
ð "ð	
ð
 #ñ
ˆô 	Ð4¸#ÀTÕJàˆÜ$Ø�J‰JØ�N‰NØñ

ð ð	

ð
 —‘ð

ð ×(Ñ(ð

ð ×*Ñ*ð

ð  $×6Ñ6ð

ð × Ñ ô

ˆEð —J‘JÕˆDØ�|�|Š|Ü�u—z‘z r°×6ñ

ð ˆr)   c                ó,   <€ V ^8„  d   QhRS[ RS[ RS[/# )r   ÚcommandÚargsr>   rI   )r$   r%   s   "€r&   r'   r(   ƒ   s"   ø€ ÷ ñ ¡cð ±ð ¹ñ r)   c                ó¤  € VR8X  d   \        R4       R# VR8X  d   V P                  4        R	# VR8X  dQ   V P                  P                  4        V P                  P                  4        \        4       V n        \        R4       R	# VR8X  d7   V'       g   \        R4       R	# V P                  VP                  4       4       R	# \        R
V 24       R	# )zTHandle special commands. Returns True if should continue chat, False if should exit.z/exitz#[bold yellow]Goodbye![/bold yellow]Fz/helpz/clearz4[bold blue]Conversation history cleared.[/bold blue]z/imagez8[bold red]Error:[/bold red] Please provide an image pathTz&[bold red]Unknown command:[/bold red] )	r6   r7   r,   Úclearr/   r   r1   rQ   Ústrip)r8   rm   rn   s   &&&r&   Úhandle_commandÚMLXVisionChat.handle_commandƒ   s¹   € à�gÔÜÐ8Ô9ÙØ˜ÔØ�O‰OÔñ ð ˜Ô Ø�L‰L×ÑÔ Ø×Ñ×"Ñ"Ô$Ü&6Ó&8ˆDÔ#ÜÐIÔJñ ð ˜Ô ßÜÐQÔRÙØ×Ñ˜tŸz™z›|Ô,ñ ô Ð;¸G¸9ÐEÔFÙr)   c                ó   <€ V ^8„  d   QhRR/# r=   r?   )r$   r%   s   "€r&   r'   r(   ˜   s   ø€ ÷  ñ  ˜4ñ  r)   c                óÂ  €   \         P                  ! R4      P                  4       pVP                  R4      '       d`   VP	                  ^R7      pV^ ,          P                  4       p\        V4      ^8”  d
   V^,          MRpV P                  W44      '       g   R# Kœ  V P                  RV4       V P                  4       pV P                  '       g    \        \        \        V4      RR7      4       VP                  RR4      pV P                  R	V4       EK    \         d    \        R
4        EK0  \          d$   p\        R\#        T4       24        Rp?EKX  Rp?ii ; i)zMain chat loop for interaction.z
[bold cyan]You[/bold cyan]Ú/)Úmaxsplitrc   ÚuserÚgreen)rD   z<end_of_utterance>Ú	assistantzD
[bold yellow]Interrupted by user. Type /exit to quit.[/bold yellow]z[bold red]Error:[/bold red] N)r	   Úaskrq   Ú
startswithÚsplitÚlowerÚlenrr   rZ   rj   r   r6   r   r   ÚreplaceÚKeyboardInterruptrO   r    )r8   Ú
user_inputÚpartsrm   rn   ÚresponserP   s   &      r&   Ú	chat_loopÚMLXVisionChat.chat_loop˜   s#  € àðÜ#ŸZšZÐ(FÓG×MÑMÓO�
ð ×(Ñ(¨×-Ò-Ø&×,Ñ,°aÐ,Ó8�EØ# A�hŸn™nÓ.�GÜ'*¨5£z°A¤~˜5 ž8¸2�DØ×.Ñ.¨w×=Ò=ÙÙØ×#Ñ# F¨JÔ7Ø×1Ñ1Ó3�à—|—|�|Üœ5¤¨(Ó!3À'ÔJÔKð $×+Ñ+Ð,@À"ÓE�à×#Ñ# K°×:øä$ô ÜØ[ôó Üô ÜÐ5´c¸!³f°XÐ>Ô?Þûðús6   ƒ:D ¾AD Â3D ÃAD ÄEÄ1EÄ:EÄ;EÅE)r+   r-   r.   r,   r/   r   r4   r5   r1   r2   r   r   r0   N)ú%mlx-community/idefics2-8b-chatty-4bitgffffffæ?iè  F)Ú__name__Ú
__module__Ú__qualname__Ú__firstlineno__r:   r7   rQ   rZ   rj   rr   r…   Ú__static_attributes__Ú__classdictcell__)r%   s   @r&   r   r      sP   ø‡ € ÷ò ÷4
Dð 
D÷ð ÷&@ð @÷
#ð #÷Jð ÷* ö  r)   r   c                  ó2  € \         P                  ! R R7      p V P                  RRRR7       V P                  R\         P                  RRR	7       V P                  R
\        \
        RR7       V P                  R\        \        RR7       V P                  R\        RRRR7       V P                  R\        \        RR7       V P                  R\        RRR7       V P                  R\        RRR7       V P                  R\        \        RR7       V P                  R\        RB\        RR7       V P                  R\        \        R R7       V P                  R!\        RRR"R7       V P                  R#R$R%R&7       V P                  R'R$R(R&7       V P                  R)\        RR*R7       V P                  R+\        \        R,R7       V P                  R-\        \        R.R7       V P                  4       p/ pVP                   eC   . pVP                    F,  p \"        P$                  ! VR/4      pVP'                  V4       K.  	  W2R0&   VP,                  '       d   VP,                  VR1&   VP.                  VR2&   VP0                  e<   VP0                  VR3&   VP2                  VR4&   VP4                  e   VP4                  VR5&   VP6                  e   VP6                  VR6&   VP8                  e=   VP8                  VR7&   VP:                  VR8&   VP<                  VR9&   VP>                  VR:&   VP@                  e   VP@                  VR;&   VPB                  e   VPB                  VR<&    \E        RCR=VPF                  R>VPH                  R?VPJ                  R@VPL                  /VB pVPO                  4        R#   \(        \*        3 d    TP'                  T4        EKÂ  i ; i  \P         d9   p\S        RA\        T4       24       \T        PV                  ! ^4        Rp?R# Rp?ii ; i)DzMLX Vision Chat CLI)Údescriptionz--modelr‡   z%Path to the model or model identifier)ÚdefaultÚhelpz	--verboseTzBStream tokens as they are generated (use --no-verbose to disable).)Úactionr�   r‘   z--temperaturezTemperature for sampling.)rX   r�   r‘   z--max-tokensz%Maximum number of tokens to generate.z--resize-shapeÚ+NzResize shape for the image.)rX   Únargsr�   r‘   z--prefill-step-sizez-Number of tokens to process per prefill step.z--max-kv-sizez%Maximum KV size for the prompt cache.z	--kv-bitsz+Number of bits to quantize the KV cache to.z--kv-group-sizez-Group size for uniform KV cache quantization.z--kv-quant-schemezKV cache quantization backend.)rX   Úchoicesr�   r‘   z--quantized-kv-startz'Start index for the quantized KV cache.z--eos-tokensz#EOS tokens to add to the tokenizer.z--skip-special-tokensÚ
store_truez'Skip special tokens in the detokenizer.)r’   r‘   z--enable-thinkingz*Enable thinking mode in the chat template.z--thinking-budgetzAMaximum number of thinking tokens before forcing end-of-thinking.z--thinking-start-tokenz/Token that marks the start of a thinking block.z--thinking-end-tokenz-Token that marks the end of a thinking block.Úunicode_escapeÚ
eos_tokensÚskip_special_tokensr^   Úthinking_budgetÚthinking_end_tokenÚthinking_start_tokenÚmax_kv_sizeÚkv_bitsÚkv_group_sizeÚkv_quant_schemeÚquantized_kv_startÚresize_shapeÚprefill_step_sizer   r   r   r   z"[bold red]Fatal error:[/bold red] )ÚuniformÚ
turboquantr?   ),ÚargparseÚArgumentParserÚadd_argumentÚBooleanOptionalActionr!   r   r"   r   r   r   r    r   r   r   r   Ú
parse_argsr˜   ÚcodecsÚdecoderN   ÚUnicodeDecodeErrorÚUnicodeErrorr™   r^   rš   r›   rœ   r�   rž   rŸ   r    r¡   r¢   r£   r   r4   r   r   r   r…   rO   r6   ÚsysÚexit)Úparserrn   r9   r˜   ÚtokenÚdecoded_tokenÚchatrP   s           r&   Úmainrµ   »   s|  € Ü×$Ò$Ð1FÔG€FØ
×ÑØØ7Ø4ð ô ð
 ×ÑØÜ×-Ñ-ØØQð	 ô ð ×ÑØÜÜ#Ø(ð	 ô ð ×ÑØÜÜ"Ø4ð	 ô ð ×ÑØÜØØØ*ð ô ð ×ÑØÜÜ)Ø<ð	 ô ð ×ÑØÜØØ4ð	 ô ð ×ÑØÜØØ:ð	 ô ð ×ÑØÜÜ%Ø<ð	 ô ð ×ÑØÜØ)Ü'Ø-ð ô ð ×ÑØÜÜ*Ø6ð	 ô ð ×ÑØÜØØØ2ð ô ð ×ÑØØØ6ð ô ð
 ×ÑØØØ9ð ô ð
 ×ÑØÜØØPð	 ô ð ×ÑØ ÜÜ,Ø>ð	 ô ð ×ÑØÜÜ*Ø<ð	 ô ð ×ÑÓ€Dð €Fà‡�Ò"Øˆ
Ø—_”_ˆEð)Ü &§¢¨eÐ5EÓ F�Ø×!Ñ! -Ö0ñ %ð  *ˆ|Ñà××ÐØ(,×(@Ñ(@ˆÐ$Ñ%ð !%× 4Ñ 4€FÐÑØ×ÑÒ'Ø$(×$8Ñ$8ˆÐ Ñ!Ø'+×'>Ñ'>ˆÐ#Ñ$Ø×$Ñ$Ò0Ø-1×-FÑ-FˆFÐ)Ñ*ð ×ÑÒ#Ø $× 0Ñ 0ˆˆ}ÑØ‡|�|ÒØ ŸL™LˆˆyÑØ"&×"4Ñ"4ˆˆÑØ$(×$8Ñ$8ˆÐ Ñ!Ø'+×'>Ñ'>ˆÐ#Ñ$à×ÑÒ$Ø!%×!2Ñ!2ˆˆ~ÑØ×ÑÒ)Ø&*×&<Ñ&<ˆÐ"Ñ#ðÜñ 
Ø—z‘zð
à×(Ñ(ð
ð —‘ð
ð —L‘Lð	
ð
 ñ
ˆð 	�‰ÖøôK '¬Ð5ô )Ø×!Ñ! %×(Ð(ð)ûôL ô ÜÐ3´C¸³F°8Ð<Ô=Ü�Š��Šûðús+   È(N)ÍAO Î)"OÏOÏPÏ-PÐPÚ__main__)(r¦   r«   rK   r¯   Útypingr   r   Úrichr   r6   Úrich.consoler   Úrich.markdownr   Ú
rich.panelr   Úrich.promptr	   Úmlx_vlmr
   Úmlx_vlm.generater   r   r   r   r   r   r   r   r   r   Úmlx_vlm.prompt_utilsr   Úmlx_vlm.utilsr   Úmlx_vlm.vision_cacher   r   rµ   rˆ   r?   r)   r&   Ú<module>rÂ      sh   ðÛ Û Û 	Û 
ß å  Ý  Ý "Ý Ý å ÷÷ ÷ õ 5Ý $Ý 3÷Yñ Yòx\ð~ ˆzÔÙ†Fñ r)   