nexaai 1.0.19rc5__cp310-cp310-macosx_14_0_universal2.whl → 1.0.19rc7__cp310-cp310-macosx_14_0_universal2.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of nexaai might be problematic. Click here for more details.

Files changed (221) hide show
  1. nexaai/_stub.cpython-310-darwin.so +0 -0
  2. nexaai/_version.py +1 -1
  3. nexaai/binds/libnexa_bridge.dylib +0 -0
  4. nexaai/binds/nexa_llama_cpp/libggml-base.dylib +0 -0
  5. nexaai/binds/nexa_llama_cpp/libggml-cpu.so +0 -0
  6. nexaai/binds/nexa_llama_cpp/libggml-metal.so +0 -0
  7. nexaai/binds/nexa_llama_cpp/libggml.dylib +0 -0
  8. nexaai/binds/nexa_llama_cpp/libllama.dylib +0 -0
  9. nexaai/binds/nexa_llama_cpp/libmtmd.dylib +0 -0
  10. nexaai/binds/nexa_llama_cpp/libnexa_plugin.dylib +0 -0
  11. nexaai/binds/nexa_mlx/libnexa_plugin.dylib +0 -0
  12. nexaai/binds/nexa_mlx/py-lib/asr/__init__.py +12 -0
  13. nexaai/binds/nexa_mlx/py-lib/asr/interface.py +122 -0
  14. nexaai/binds/nexa_mlx/py-lib/common/__init__.py +0 -0
  15. nexaai/binds/nexa_mlx/py-lib/common/utils.py +25 -0
  16. nexaai/binds/nexa_mlx/py-lib/cv/__init__.py +0 -0
  17. nexaai/binds/nexa_mlx/py-lib/cv/generate.py +195 -0
  18. nexaai/binds/nexa_mlx/py-lib/cv/interface.py +151 -0
  19. nexaai/binds/nexa_mlx/py-lib/cv/main.py +81 -0
  20. nexaai/binds/nexa_mlx/py-lib/cv/modeling/pp_ocr_v4.py +1736 -0
  21. nexaai/binds/nexa_mlx/py-lib/embedding/__init__.py +0 -0
  22. nexaai/binds/nexa_mlx/py-lib/embedding/generate.py +333 -0
  23. nexaai/binds/nexa_mlx/py-lib/embedding/interface.py +617 -0
  24. nexaai/binds/nexa_mlx/py-lib/embedding/main.py +173 -0
  25. nexaai/binds/nexa_mlx/py-lib/embedding/modeling/__init__.py +0 -0
  26. nexaai/binds/nexa_mlx/py-lib/embedding/modeling/nexa_jina_v2.py +399 -0
  27. nexaai/binds/nexa_mlx/py-lib/image_gen/__init__.py +1 -0
  28. nexaai/binds/nexa_mlx/py-lib/image_gen/generate_sd.py +244 -0
  29. nexaai/binds/nexa_mlx/py-lib/image_gen/interface.py +82 -0
  30. nexaai/binds/nexa_mlx/py-lib/image_gen/main.py +281 -0
  31. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/__init__.py +306 -0
  32. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/clip.py +116 -0
  33. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/config.py +65 -0
  34. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/model_io.py +386 -0
  35. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/sampler.py +105 -0
  36. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/tokenizer.py +100 -0
  37. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/unet.py +460 -0
  38. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/vae.py +274 -0
  39. nexaai/binds/nexa_mlx/py-lib/llm/__init__.py +0 -0
  40. nexaai/binds/nexa_mlx/py-lib/llm/generate.py +149 -0
  41. nexaai/binds/nexa_mlx/py-lib/llm/interface.py +764 -0
  42. nexaai/binds/nexa_mlx/py-lib/llm/main.py +68 -0
  43. nexaai/binds/nexa_mlx/py-lib/rerank/__init__.py +0 -0
  44. nexaai/binds/nexa_mlx/py-lib/rerank/generate.py +174 -0
  45. nexaai/binds/nexa_mlx/py-lib/rerank/interface.py +287 -0
  46. nexaai/binds/nexa_mlx/py-lib/rerank/main.py +127 -0
  47. nexaai/binds/nexa_mlx/py-lib/rerank/modeling/__init__.py +0 -0
  48. nexaai/binds/nexa_mlx/py-lib/rerank/modeling/nexa_jina_rerank.py +330 -0
  49. nexaai/binds/nexa_mlx/py-lib/sd/__init__.py +1 -0
  50. nexaai/binds/nexa_mlx/py-lib/sd/interface.py +362 -0
  51. nexaai/binds/nexa_mlx/py-lib/sd/main.py +286 -0
  52. nexaai/binds/nexa_mlx/py-lib/sd/modeling/__init__.py +306 -0
  53. nexaai/binds/nexa_mlx/py-lib/sd/modeling/clip.py +116 -0
  54. nexaai/binds/nexa_mlx/py-lib/sd/modeling/config.py +65 -0
  55. nexaai/binds/nexa_mlx/py-lib/sd/modeling/model_io.py +385 -0
  56. nexaai/binds/nexa_mlx/py-lib/sd/modeling/sampler.py +105 -0
  57. nexaai/binds/nexa_mlx/py-lib/sd/modeling/tokenizer.py +100 -0
  58. nexaai/binds/nexa_mlx/py-lib/sd/modeling/unet.py +460 -0
  59. nexaai/binds/nexa_mlx/py-lib/sd/modeling/vae.py +274 -0
  60. nexaai/binds/nexa_mlx/py-lib/tts/__init__.py +12 -0
  61. nexaai/binds/nexa_mlx/py-lib/tts/interface.py +276 -0
  62. nexaai/binds/nexa_mlx/py-lib/vlm/__init__.py +3 -0
  63. nexaai/binds/nexa_mlx/py-lib/vlm/generate.py +572 -0
  64. nexaai/binds/nexa_mlx/py-lib/vlm/generate_qwen3_vl.py +294 -0
  65. nexaai/binds/nexa_mlx/py-lib/vlm/generate_qwen3_vl_moe.py +276 -0
  66. nexaai/binds/nexa_mlx/py-lib/vlm/interface.py +504 -0
  67. nexaai/binds/nexa_mlx/py-lib/vlm/main.py +320 -0
  68. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/__init__.py +0 -0
  69. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/convert.py +68 -0
  70. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/__init__.py +0 -0
  71. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/__init__.py +8 -0
  72. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/aya_vision.py +193 -0
  73. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/interpolate.py +186 -0
  74. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/language.py +233 -0
  75. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/vision.py +503 -0
  76. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/base.py +202 -0
  77. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/cache.py +230 -0
  78. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/__init__.py +10 -0
  79. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/conversation.py +264 -0
  80. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/deepseek_vl_v2.py +472 -0
  81. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/language.py +591 -0
  82. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/processing_deepsek_vl_v2.py +526 -0
  83. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/vision.py +356 -0
  84. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/florence2/__init__.py +8 -0
  85. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/florence2/florence2.py +366 -0
  86. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/florence2/language.py +488 -0
  87. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/florence2/vision.py +591 -0
  88. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3/__init__.py +8 -0
  89. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3/gemma3.py +213 -0
  90. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3/language.py +315 -0
  91. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3/vision.py +238 -0
  92. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/__init__.py +2 -0
  93. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/audio.py +1038 -0
  94. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/config.py +139 -0
  95. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/gemma3n.py +322 -0
  96. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/language.py +629 -0
  97. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/vision.py +1022 -0
  98. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics2/__init__.py +9 -0
  99. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics2/idefics2.py +294 -0
  100. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics2/language.py +191 -0
  101. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics2/vision.py +267 -0
  102. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics3/__init__.py +8 -0
  103. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics3/idefics3.py +175 -0
  104. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics3/language.py +192 -0
  105. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics3/vision.py +233 -0
  106. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/__init__.py +9 -0
  107. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/internvl_chat.py +140 -0
  108. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/language.py +220 -0
  109. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/processor.py +393 -0
  110. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/vision.py +293 -0
  111. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kernels.py +307 -0
  112. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kimi_vl/__init__.py +8 -0
  113. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kimi_vl/kimi_vl.py +143 -0
  114. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kimi_vl/language.py +509 -0
  115. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kimi_vl/vision.py +522 -0
  116. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llama4/__init__.py +8 -0
  117. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llama4/language.py +386 -0
  118. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llama4/llama4.py +138 -0
  119. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llama4/vision.py +560 -0
  120. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava/__init__.py +8 -0
  121. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava/language.py +240 -0
  122. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava/llava.py +153 -0
  123. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava/vision.py +259 -0
  124. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_bunny/__init__.py +9 -0
  125. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_bunny/language.py +236 -0
  126. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_bunny/llava_bunny.py +256 -0
  127. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_bunny/vision.py +303 -0
  128. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_next/__init__.py +8 -0
  129. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_next/language.py +230 -0
  130. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_next/llava_next.py +160 -0
  131. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_next/vision.py +243 -0
  132. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mistral3/__init__.py +8 -0
  133. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mistral3/mistral3.py +283 -0
  134. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mllama/__init__.py +8 -0
  135. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mllama/language.py +416 -0
  136. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mllama/mllama.py +172 -0
  137. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mllama/vision.py +499 -0
  138. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/molmo/__init__.py +8 -0
  139. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/molmo/language.py +243 -0
  140. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/molmo/molmo.py +133 -0
  141. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/molmo/vision.py +465 -0
  142. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/__init__.py +10 -0
  143. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/language.py +230 -0
  144. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/multi_modality.py +385 -0
  145. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/sam.py +557 -0
  146. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/vision.py +526 -0
  147. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/paligemma/__init__.py +8 -0
  148. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/paligemma/language.py +282 -0
  149. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/paligemma/paligemma.py +160 -0
  150. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/paligemma/vision.py +242 -0
  151. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/__init__.py +8 -0
  152. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/language.py +21 -0
  153. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/phi3_v.py +243 -0
  154. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/su_rope.py +71 -0
  155. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/vision.py +324 -0
  156. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/pixtral/__init__.py +8 -0
  157. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/pixtral/language.py +229 -0
  158. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/pixtral/pixtral.py +161 -0
  159. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/pixtral/vision.py +320 -0
  160. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/__init__.py +2 -0
  161. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/config.py +108 -0
  162. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/language.py +490 -0
  163. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/qwen2_5_vl.py +168 -0
  164. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/vision.py +414 -0
  165. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/__init__.py +2 -0
  166. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/config.py +104 -0
  167. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/language.py +490 -0
  168. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/qwen2_vl.py +167 -0
  169. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/vision.py +312 -0
  170. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/__init__.py +0 -0
  171. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/base.py +117 -0
  172. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/cache.py +531 -0
  173. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/generate.py +701 -0
  174. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/rope_utils.py +255 -0
  175. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/sample_utils.py +303 -0
  176. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/tokenizer_utils.py +407 -0
  177. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/processor.py +476 -0
  178. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/qwen3vl.py +1223 -0
  179. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/__init__.py +0 -0
  180. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/base.py +117 -0
  181. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/cache.py +531 -0
  182. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/generate.py +701 -0
  183. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/rope_utils.py +255 -0
  184. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/sample_utils.py +303 -0
  185. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/tokenizer_utils.py +407 -0
  186. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/processor.py +476 -0
  187. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/qwen3vl_moe.py +1309 -0
  188. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/switch_layers.py +210 -0
  189. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/smolvlm/__init__.py +8 -0
  190. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/smolvlm/smolvlm.py +62 -0
  191. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/processing_qwen2_5_vl.py +209 -0
  192. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/processing_qwen2_vl.py +215 -0
  193. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/prompt_utils.py +474 -0
  194. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/sample_utils.py +39 -0
  195. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/tokenizer_utils.py +344 -0
  196. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/trainer/__init__.py +9 -0
  197. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/trainer/lora.py +70 -0
  198. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/trainer/trainer.py +296 -0
  199. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/trainer/utils.py +160 -0
  200. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/utils.py +928 -0
  201. nexaai/binds/nexa_nexaml/libggml-base.dylib +0 -0
  202. nexaai/binds/nexa_nexaml/libggml-cpu.so +0 -0
  203. nexaai/binds/nexa_nexaml/libggml-metal.so +0 -0
  204. nexaai/binds/nexa_nexaml/libggml.dylib +0 -0
  205. nexaai/mlx_backend/vlm/generate_qwen3_vl_moe.py +276 -0
  206. nexaai/mlx_backend/vlm/interface.py +21 -4
  207. nexaai/mlx_backend/vlm/main.py +6 -2
  208. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/llm_common/__init__.py +0 -0
  209. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/llm_common/base.py +117 -0
  210. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/llm_common/cache.py +531 -0
  211. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/llm_common/generate.py +701 -0
  212. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/llm_common/rope_utils.py +255 -0
  213. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/llm_common/sample_utils.py +303 -0
  214. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/llm_common/tokenizer_utils.py +407 -0
  215. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/processor.py +476 -0
  216. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/qwen3vl_moe.py +1309 -0
  217. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/switch_layers.py +210 -0
  218. {nexaai-1.0.19rc5.dist-info → nexaai-1.0.19rc7.dist-info}/METADATA +1 -1
  219. {nexaai-1.0.19rc5.dist-info → nexaai-1.0.19rc7.dist-info}/RECORD +221 -21
  220. {nexaai-1.0.19rc5.dist-info → nexaai-1.0.19rc7.dist-info}/WHEEL +0 -0
  221. {nexaai-1.0.19rc5.dist-info → nexaai-1.0.19rc7.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,168 @@
1
+ import glob
2
+ import json
3
+ from pathlib import Path
4
+ from typing import Optional
5
+
6
+ import mlx.core as mx
7
+ import mlx.nn as nn
8
+ from huggingface_hub import snapshot_download
9
+
10
+ from .config import ModelConfig, TextConfig, VisionConfig
11
+ from .language import LanguageModel
12
+ from .vision import VisionModel
13
+
14
+
15
+ class Model(nn.Module):
16
+ def __init__(self, config: ModelConfig):
17
+ super().__init__()
18
+ self.config = config
19
+ self.vision_tower = VisionModel(config.vision_config)
20
+ self.language_model = LanguageModel(config.text_config, config)
21
+
22
+ def get_input_embeddings(
23
+ self,
24
+ input_ids: Optional[mx.array] = None,
25
+ pixel_values: Optional[mx.array] = None,
26
+ image_grid_thw: Optional[mx.array] = None,
27
+ ):
28
+ if pixel_values is None:
29
+ return self.language_model.model.embed_tokens(input_ids)
30
+
31
+ dtype = self.vision_tower.patch_embed.proj.weight.dtype
32
+ pixel_values = pixel_values.astype(dtype)
33
+
34
+ # Get the input embeddings from the language model
35
+ inputs_embeds = self.language_model.model.embed_tokens(input_ids)
36
+
37
+ # Get the ouptut hidden states from the vision model
38
+ hidden_states = self.vision_tower(
39
+ pixel_values, image_grid_thw, output_hidden_states=False
40
+ )
41
+
42
+ # Insert special image tokens in the input_ids
43
+ final_inputs_embeds = self.merge_input_ids_with_image_features(
44
+ self.config.image_token_id,
45
+ self.config.video_token_id,
46
+ hidden_states,
47
+ inputs_embeds,
48
+ input_ids,
49
+ )
50
+ return final_inputs_embeds
51
+
52
+ @staticmethod
53
+ def merge_input_ids_with_image_features(
54
+ image_token_id,
55
+ video_token_id,
56
+ image_features,
57
+ inputs_embeds,
58
+ input_ids,
59
+ ):
60
+ """Merge image features into input embeddings at image token positions.
61
+
62
+ Args:
63
+ image_token_id: The token ID for image placeholders
64
+ video_token_id: The token ID for video placeholders (fallback)
65
+ image_features: Vision features from the vision tower [num_features, hidden_dim]
66
+ inputs_embeds: Input embeddings [batch_size, seq_len, hidden_dim]
67
+ input_ids: Input token IDs [batch_size, seq_len]
68
+ grid_thw: Grid dimensions for each image (optional, not used in simple case)
69
+
70
+ Returns:
71
+ Updated input embeddings with image features inserted
72
+ """
73
+ # Find positions of image tokens
74
+ image_positions = input_ids == image_token_id
75
+ if mx.sum(image_positions) == 0:
76
+ image_positions = input_ids == video_token_id
77
+
78
+ # Get dimensions
79
+ batch_size, seq_len = input_ids.shape
80
+
81
+ # Process each batch item
82
+ batch_outputs = []
83
+ feature_start_idx = 0
84
+
85
+ for batch_idx in range(batch_size):
86
+ # Get mask for this batch
87
+ image_mask = image_positions[batch_idx]
88
+ num_positions = mx.sum(image_mask).item()
89
+
90
+ if num_positions > 0:
91
+ # Extract features for this batch
92
+ batch_features = image_features[
93
+ feature_start_idx : feature_start_idx + num_positions
94
+ ]
95
+
96
+ # Validate we have the right number of features
97
+ if batch_features.shape[0] != num_positions:
98
+ raise ValueError(
99
+ f"Number of image token positions ({num_positions}) does not match "
100
+ f"number of image features ({batch_features.shape[0]}) for batch {batch_idx}"
101
+ )
102
+
103
+ # Create indices for gathering
104
+ cumsum = mx.cumsum(image_mask.astype(mx.int32))
105
+ feature_indices = mx.where(image_mask, cumsum - 1, 0)
106
+
107
+ # Gather features
108
+ gathered_features = batch_features[feature_indices]
109
+
110
+ # Combine with original embeddings
111
+ image_mask_expanded = mx.expand_dims(image_mask, axis=-1)
112
+ batch_output = mx.where(
113
+ image_mask_expanded, gathered_features, inputs_embeds[batch_idx]
114
+ )
115
+
116
+ feature_start_idx += num_positions
117
+ else:
118
+ # No image tokens in this batch item
119
+ batch_output = inputs_embeds[batch_idx]
120
+
121
+ batch_outputs.append(batch_output)
122
+
123
+ # Stack all batch outputs
124
+ return mx.stack(batch_outputs, axis=0)
125
+
126
+ @property
127
+ def layers(self):
128
+ return self.language_model.model.layers
129
+
130
+ def __call__(
131
+ self,
132
+ input_ids: mx.array,
133
+ pixel_values: Optional[mx.array] = None,
134
+ mask: Optional[mx.array] = None,
135
+ cache=None,
136
+ **kwargs,
137
+ ):
138
+ image_grid_thw = kwargs.pop("image_grid_thw", None)
139
+ video_grid_thw = kwargs.pop("video_grid_thw", None)
140
+ grid_thw = image_grid_thw if image_grid_thw is not None else video_grid_thw
141
+
142
+ inputs_embeds = self.get_input_embeddings(input_ids, pixel_values, grid_thw)
143
+
144
+ kwargs = {
145
+ "pixel_values": pixel_values,
146
+ "image_grid_thw": image_grid_thw,
147
+ "video_grid_thw": video_grid_thw,
148
+ **kwargs,
149
+ }
150
+
151
+ logits = self.language_model(
152
+ input_ids, inputs_embeds, mask=mask, cache=cache, **kwargs
153
+ )
154
+
155
+ return logits
156
+
157
+ def sanitize(self, weights):
158
+ def transform_key(key):
159
+ if "vision_tower" not in key:
160
+ key = key.replace("visual", "vision_tower")
161
+ if "language_model" not in key:
162
+ if "model" in key:
163
+ key = key.replace("model", "language_model.model")
164
+ elif "lm_head" in key:
165
+ key = key.replace("lm_head", "language_model.lm_head")
166
+ return key
167
+
168
+ return {transform_key(k): v for k, v in weights.items()}
@@ -0,0 +1,414 @@
1
+ from typing import Optional
2
+
3
+ import mlx.core as mx
4
+ import mlx.nn as nn
5
+ import numpy as np
6
+
7
+ from .config import VisionConfig
8
+
9
+
10
+ def check_array_shape(arr):
11
+ shape = arr.shape
12
+
13
+ # Check if the shape has 4 dimensions
14
+ if len(shape) not in [4, 5]:
15
+ return False
16
+
17
+ B, out_channels, kH, KW, t = shape
18
+
19
+ if t == 3:
20
+ return True
21
+
22
+ # Check if out_channels is the largest, and kH and KW are the same
23
+ if (out_channels >= kH) and (out_channels >= KW) and (kH == KW):
24
+ return True
25
+ else:
26
+ return False
27
+
28
+
29
+ def rotate_half(x):
30
+ """Rotates half the hidden dims of the input."""
31
+ x1 = x[..., : x.shape[-1] // 2]
32
+ x2 = x[..., x.shape[-1] // 2 :]
33
+ return mx.concatenate([-x2, x1], axis=-1)
34
+
35
+
36
+ def apply_rotary_pos_emb_vision(tensor, freqs) -> mx.array:
37
+ orig_dtype = tensor.dtype
38
+
39
+ cos = mx.cos(freqs)
40
+ sin = mx.sin(freqs)
41
+
42
+ cos = mx.expand_dims(cos, axis=1) # Equivalent to unsqueeze(1)
43
+ cos = mx.tile(cos, (1, 1, 2)) # Equivalent to repeat(1, 1, 2)
44
+ cos = mx.expand_dims(cos, axis=0) # Equivalent to [None, ...]
45
+
46
+ sin = mx.expand_dims(sin, axis=1) # Equivalent to unsqueeze(1)
47
+ sin = mx.tile(sin, (1, 1, 2)) # Equivalent to repeat(1, 1, 2)
48
+ sin = mx.expand_dims(sin, axis=0) # Equivalent to [None, ...]
49
+
50
+ output = (tensor * cos) + (rotate_half(tensor) * sin)
51
+ return output.astype(orig_dtype)
52
+
53
+
54
+ class VisionRotaryEmbedding(nn.Module):
55
+ def __init__(self, dim: int, theta: float = 10000.0) -> None:
56
+ super().__init__()
57
+ self.dim = dim
58
+ self.theta = theta
59
+
60
+ def __call__(self, seqlen: int) -> mx.array:
61
+ inv_freq = 1.0 / (
62
+ self.theta ** (mx.arange(0, self.dim, 2, dtype=mx.float32) / self.dim)
63
+ )
64
+ seq = mx.arange(seqlen.item(), dtype=inv_freq.dtype)
65
+ freqs = mx.outer(seq, inv_freq)
66
+ return freqs
67
+
68
+
69
+ class PatchEmbed(nn.Module):
70
+ def __init__(
71
+ self,
72
+ patch_size: int = 14,
73
+ temporal_patch_size: int = 2,
74
+ in_channels: int = 3,
75
+ hidden_size: int = 1152,
76
+ ) -> None:
77
+ super().__init__()
78
+ self.patch_size = patch_size
79
+ self.temporal_patch_size = temporal_patch_size
80
+ self.in_channels = in_channels
81
+ self.hidden_size = hidden_size
82
+
83
+ kernel_size = [temporal_patch_size, patch_size, patch_size]
84
+ self.proj = nn.Conv3d(
85
+ in_channels,
86
+ hidden_size,
87
+ kernel_size=kernel_size,
88
+ stride=kernel_size,
89
+ bias=False,
90
+ )
91
+
92
+ def __call__(self, hidden_states: mx.array) -> mx.array:
93
+ hidden_states = hidden_states.reshape(
94
+ -1,
95
+ self.in_channels,
96
+ self.temporal_patch_size,
97
+ self.patch_size,
98
+ self.patch_size,
99
+ ).moveaxis(1, 4)
100
+
101
+ hidden_states = self.proj(hidden_states)
102
+ hidden_states = hidden_states.reshape(-1, self.hidden_size)
103
+ return hidden_states
104
+
105
+
106
+ class PatchMerger(nn.Module):
107
+ def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None:
108
+ super().__init__()
109
+ self.hidden_size = context_dim * (spatial_merge_size**2)
110
+ self.ln_q = nn.RMSNorm(context_dim, eps=1e-6)
111
+ self.mlp = [
112
+ nn.Linear(self.hidden_size, self.hidden_size),
113
+ nn.GELU(),
114
+ nn.Linear(self.hidden_size, dim),
115
+ ]
116
+
117
+ def __call__(self, x: mx.array) -> mx.array:
118
+ x = self.ln_q(x).reshape(-1, self.hidden_size)
119
+ for layer in self.mlp:
120
+ x = layer(x)
121
+ return x
122
+
123
+
124
+ class Attention(nn.Module):
125
+ def __init__(self, dim: int, num_heads: int = 16) -> None:
126
+ super().__init__()
127
+ self.num_heads = num_heads
128
+ self.head_dim = head_dim = dim // num_heads
129
+ self.scale = head_dim**-0.5
130
+ self.qkv = nn.Linear(dim, dim * 3, bias=True)
131
+ self.proj = nn.Linear(dim, dim)
132
+
133
+ def __call__(
134
+ self, x: mx.array, cu_seqlens: mx.array, rotary_pos_emb: mx.array = None
135
+ ) -> mx.array:
136
+ seq_length = x.shape[0]
137
+ qkv = (
138
+ self.qkv(x).reshape(seq_length, 3, self.num_heads, -1).transpose(1, 0, 2, 3)
139
+ )
140
+ q, k, v = mx.split(qkv, 3)
141
+
142
+ q = apply_rotary_pos_emb_vision(mx.expand_dims(q, 0), rotary_pos_emb)[0]
143
+ k = apply_rotary_pos_emb_vision(mx.expand_dims(k, 0), rotary_pos_emb)[0]
144
+ attention_mask = mx.full(
145
+ (1, seq_length, seq_length), mx.finfo(q.dtype).min, dtype=q.dtype
146
+ )
147
+
148
+ for i in range(1, len(cu_seqlens)):
149
+ start = int(cu_seqlens[i - 1])
150
+ end = int(cu_seqlens[i])
151
+ attention_mask[..., start:end, start:end] = 0
152
+
153
+ q = q.transpose(0, 2, 1, 3)
154
+ k = k.transpose(0, 2, 1, 3)
155
+ v = v.transpose(0, 2, 1, 3)
156
+
157
+ output = mx.fast.scaled_dot_product_attention(
158
+ q, k, v, scale=self.scale, mask=attention_mask
159
+ )
160
+ output = output.transpose(0, 2, 1, 3)
161
+ output = output.reshape(seq_length, -1)
162
+ return self.proj(output)
163
+
164
+
165
+ class MLP(nn.Module):
166
+ def __init__(self, dim, hidden_dim):
167
+ super().__init__()
168
+ self.gate_proj = nn.Linear(dim, hidden_dim)
169
+ self.up_proj = nn.Linear(dim, hidden_dim)
170
+ self.down_proj = nn.Linear(hidden_dim, dim)
171
+
172
+ def __call__(self, x: mx.array) -> mx.array:
173
+ return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
174
+
175
+
176
+ class Qwen2VLVisionBlock(nn.Module):
177
+ def __init__(self, config: VisionConfig) -> None:
178
+ super().__init__()
179
+ self.norm1 = nn.RMSNorm(config.hidden_size, eps=1e-6)
180
+ self.norm2 = nn.RMSNorm(config.hidden_size, eps=1e-6)
181
+
182
+ self.attn = Attention(dim=config.hidden_size, num_heads=config.num_heads)
183
+ self.mlp = MLP(dim=config.hidden_size, hidden_dim=config.intermediate_size)
184
+
185
+ def __call__(self, hidden_states, cu_seqlens, rotary_pos_emb) -> mx.array:
186
+ hidden_states = hidden_states + self.attn(
187
+ self.norm1(hidden_states),
188
+ cu_seqlens=cu_seqlens,
189
+ rotary_pos_emb=rotary_pos_emb,
190
+ )
191
+ hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
192
+ return hidden_states
193
+
194
+
195
+ class VisionModel(nn.Module):
196
+
197
+ def __init__(self, config: VisionConfig) -> None:
198
+ super().__init__()
199
+ self.config = config
200
+ self.model_type = config.model_type
201
+ if self.model_type != "qwen2_5_vl":
202
+ raise ValueError(f"Unsupported model type: {self.model_type}")
203
+ self.spatial_merge_size = config.spatial_merge_size
204
+
205
+ self.patch_embed = PatchEmbed(
206
+ patch_size=config.patch_size,
207
+ temporal_patch_size=config.temporal_patch_size,
208
+ in_channels=config.in_channels,
209
+ hidden_size=config.hidden_size,
210
+ )
211
+
212
+ self.window_size = config.window_size
213
+ self.patch_size = config.patch_size
214
+ self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size
215
+ self.fullatt_block_indexes = config.fullatt_block_indexes
216
+ head_dim = config.hidden_size // config.num_heads
217
+ self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)
218
+
219
+ self.blocks = [Qwen2VLVisionBlock(config) for _ in range(config.depth)]
220
+ self.merger = PatchMerger(
221
+ dim=config.out_hidden_size, context_dim=config.hidden_size
222
+ )
223
+
224
+ def rot_pos_emb(self, grid_thw):
225
+ pos_ids = []
226
+
227
+ for t, h, w in grid_thw.tolist():
228
+ hpos_ids = mx.expand_dims(mx.arange(h), 1)
229
+ hpos_ids = mx.repeat(hpos_ids, w, axis=1)
230
+ hpos_ids = hpos_ids.reshape(
231
+ h // self.spatial_merge_size,
232
+ self.spatial_merge_size,
233
+ w // self.spatial_merge_size,
234
+ self.spatial_merge_size,
235
+ )
236
+ hpos_ids = mx.transpose(hpos_ids, (0, 2, 1, 3))
237
+ hpos_ids = hpos_ids.flatten()
238
+
239
+ wpos_ids = mx.expand_dims(mx.arange(w), 0)
240
+ wpos_ids = mx.repeat(wpos_ids, h, axis=0)
241
+ wpos_ids = wpos_ids.reshape(
242
+ h // self.spatial_merge_size,
243
+ self.spatial_merge_size,
244
+ w // self.spatial_merge_size,
245
+ self.spatial_merge_size,
246
+ )
247
+ wpos_ids = mx.transpose(wpos_ids, (0, 2, 1, 3))
248
+ wpos_ids = wpos_ids.flatten()
249
+
250
+ stacked_pos_ids = mx.stack([hpos_ids, wpos_ids], axis=-1)
251
+ pos_ids.append(mx.tile(stacked_pos_ids, (t, 1)))
252
+
253
+ pos_ids = mx.concatenate(pos_ids, axis=0)
254
+ max_grid_size = mx.max(grid_thw[:, 1:])
255
+ rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
256
+ rotary_pos_emb = rotary_pos_emb_full[pos_ids]
257
+
258
+ return rotary_pos_emb.reshape(pos_ids.shape[0], -1)
259
+
260
+ def get_window_index(self, grid_thw):
261
+ window_index = []
262
+ cu_window_seqlens = [0]
263
+ window_index_id = 0
264
+ vit_merger_window_size = (
265
+ self.window_size // self.spatial_merge_size // self.patch_size
266
+ )
267
+
268
+ for grid_t, grid_h, grid_w in grid_thw.tolist():
269
+ llm_grid_h = grid_h // self.spatial_merge_size
270
+ llm_grid_w = grid_w // self.spatial_merge_size
271
+
272
+ index = mx.arange(grid_t * llm_grid_h * llm_grid_w).reshape(
273
+ grid_t, llm_grid_h, llm_grid_w
274
+ )
275
+
276
+ pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size
277
+ pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size
278
+ num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size
279
+ num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size
280
+
281
+ # Replace F.pad with np.pad
282
+ index_padded = mx.pad(
283
+ index,
284
+ ((0, 0), (0, pad_h), (0, pad_w)),
285
+ mode="constant",
286
+ constant_values=-100,
287
+ )
288
+
289
+ index_padded = index_padded.reshape(
290
+ grid_t,
291
+ num_windows_h,
292
+ vit_merger_window_size,
293
+ num_windows_w,
294
+ vit_merger_window_size,
295
+ )
296
+
297
+ # Replace permute with np.transpose
298
+ index_padded = mx.transpose(index_padded, (0, 1, 3, 2, 4)).reshape(
299
+ grid_t,
300
+ num_windows_h * num_windows_w,
301
+ vit_merger_window_size,
302
+ vit_merger_window_size,
303
+ )
304
+
305
+ # Replace torch operations with numpy
306
+ seqlens = mx.sum(index_padded != -100, axis=(2, 3)).reshape(-1)
307
+ index_padded = index_padded.reshape(-1)
308
+ index = np.where(index_padded != -100)[
309
+ 0
310
+ ].tolist() # [i for i, x in enumerate(index_padded) if x != -100]
311
+ index_new = index_padded[index]
312
+
313
+ window_index.append(index_new + window_index_id)
314
+ cu_seqlens_tmp = (
315
+ mx.cumsum(seqlens, axis=0) * self.spatial_merge_unit
316
+ + cu_window_seqlens[-1]
317
+ )
318
+ cu_window_seqlens.extend(cu_seqlens_tmp.tolist())
319
+ window_index_id += int(grid_t * llm_grid_h * llm_grid_w)
320
+
321
+ # Replace torch.cat with np.concatenate
322
+ window_index = mx.concatenate(window_index, axis=0)
323
+ cu_window_seqlens = mx.array(cu_window_seqlens)
324
+
325
+ return window_index, cu_window_seqlens
326
+
327
+ def __call__(
328
+ self,
329
+ hidden_states: mx.array,
330
+ grid_thw: mx.array,
331
+ output_hidden_states: Optional[bool] = None,
332
+ ) -> mx.array:
333
+
334
+ hidden_states = self.patch_embed(hidden_states)
335
+ rotary_pos_emb = self.rot_pos_emb(grid_thw)
336
+ window_index, cu_window_seqlens = self.get_window_index(grid_thw)
337
+
338
+ # Get indices of first occurrence of each unique value
339
+ seen = set()
340
+ idx = []
341
+ for i, x in enumerate(cu_window_seqlens):
342
+ if x not in seen:
343
+ seen.add(x)
344
+ idx.append(i)
345
+
346
+ idx = mx.array(idx, dtype=mx.int32)
347
+ cu_window_seqlens = cu_window_seqlens[idx]
348
+
349
+ seq_len, _ = hidden_states.shape
350
+ hidden_states = hidden_states.reshape(
351
+ seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1
352
+ )
353
+ hidden_states = hidden_states[window_index, :, :]
354
+ hidden_states = hidden_states.reshape(seq_len, -1)
355
+ rotary_pos_emb = rotary_pos_emb.reshape(
356
+ seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1
357
+ )
358
+ rotary_pos_emb = rotary_pos_emb[window_index, :, :]
359
+ rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
360
+
361
+ # Assuming grid_thw has shape (batch_size, 3)
362
+ batch_size = grid_thw.shape[0]
363
+
364
+ # Calculate cu_seqlens for each item in the batch
365
+ cu_seqlens = []
366
+ for i in range(batch_size):
367
+ seq_len = grid_thw[i, 1] * grid_thw[i, 2]
368
+ cu_seqlens.append(mx.repeat(seq_len, grid_thw[i, 0]))
369
+
370
+ # Concatenate the cu_seqlens for all items in the batch
371
+ cu_seqlens = mx.concatenate(cu_seqlens)
372
+
373
+ cu_seqlens = mx.cumsum(cu_seqlens.astype(mx.int32), axis=0)
374
+ cu_seqlens = mx.pad(cu_seqlens, (1, 0), mode="constant", constant_values=0)
375
+
376
+ encoder_states = (hidden_states,) if output_hidden_states else None
377
+
378
+ for layer_num, blk in enumerate(self.blocks):
379
+ if layer_num in self.fullatt_block_indexes:
380
+ cu_seqlens_now = cu_seqlens
381
+ else:
382
+ cu_seqlens_now = cu_window_seqlens
383
+
384
+ hidden_states = blk(
385
+ hidden_states, cu_seqlens=cu_seqlens_now, rotary_pos_emb=rotary_pos_emb
386
+ )
387
+
388
+ if output_hidden_states:
389
+ encoder_states = encoder_states + (hidden_states,)
390
+
391
+ hidden_states = self.merger(hidden_states)
392
+ reverse_indices = mx.argsort(window_index, axis=0)
393
+ hidden_states = hidden_states[reverse_indices, :]
394
+ return hidden_states
395
+
396
+ def sanitize(self, weights):
397
+ sanitized_weights = {}
398
+ for k, v in weights.items():
399
+ if "position_ids" in k:
400
+ # Remove unused position_ids
401
+ continue
402
+ elif "patch_embed.proj.weight" in k:
403
+ # PyTorch conv2d weight tensors have shape:
404
+ # [out_channels, in_channels, kH, KW]
405
+ # MLX conv2d expects the weight be of shape:
406
+ # [out_channels, kH, KW, in_channels]
407
+ if check_array_shape(v):
408
+ sanitized_weights[k] = v
409
+ else:
410
+ sanitized_weights[k] = v.transpose(0, 2, 3, 4, 1)
411
+ else:
412
+ sanitized_weights[k] = v
413
+
414
+ return sanitized_weights
@@ -0,0 +1,2 @@
1
+ from .config import ModelConfig, TextConfig, VisionConfig
2
+ from .qwen2_vl import LanguageModel, Model, VisionModel
@@ -0,0 +1,104 @@
1
+ import inspect
2
+ from dataclasses import dataclass
3
+ from typing import Dict, List, Optional, Union
4
+
5
+
6
+ @dataclass
7
+ class VisionConfig:
8
+ model_type: str = "qwen2_vl"
9
+ depth: int = 32
10
+ embed_dim: int = 1280
11
+ hidden_size: int = 1536
12
+ num_heads: int = 16
13
+ image_size: int = 384
14
+ patch_size: int = 14
15
+ vocab_size: int = 32000
16
+ mlp_ratio: float = 4.0
17
+ in_channels: int = 3
18
+ layer_norm_eps: float = 1e-6
19
+ spatial_patch_size: int = 14
20
+ spatial_merge_size: int = 2
21
+ temporal_patch_size: int = 2
22
+
23
+ @classmethod
24
+ def from_dict(cls, params):
25
+ return cls(
26
+ **{
27
+ k: v
28
+ for k, v in params.items()
29
+ if k in inspect.signature(cls).parameters
30
+ }
31
+ )
32
+
33
+
34
+ @dataclass
35
+ class TextConfig:
36
+ model_type: str
37
+ hidden_size: int
38
+ num_hidden_layers: int
39
+ intermediate_size: int
40
+ num_attention_heads: int
41
+ rms_norm_eps: float
42
+ vocab_size: int
43
+ num_key_value_heads: Optional[int] = 8
44
+ max_position_embeddings: Optional[int] = 40960
45
+ rope_theta: float = 1000000.0
46
+ rope_traditional: bool = False
47
+ rope_scaling: Optional[Dict[str, Union[float, str]]] = None
48
+ tie_word_embeddings: bool = False
49
+ sliding_window: int = 32768
50
+ use_sliding_window: bool = False
51
+ use_cache: bool = True
52
+
53
+ def __post_init__(self):
54
+ if self.num_key_value_heads is None:
55
+ self.num_key_value_heads = self.num_attention_heads
56
+
57
+ if self.rope_scaling:
58
+ required_keys = {"mrope_section", "type"}
59
+ if not all(key in self.rope_scaling for key in required_keys):
60
+ raise ValueError(f"rope_scaling must contain keys {required_keys}")
61
+
62
+ if not self.rope_scaling["type"] in ["mrope", "default"]:
63
+ raise ValueError(f"rope_scaling type must be 'mrope' or 'default'")
64
+
65
+ @classmethod
66
+ def from_dict(cls, params):
67
+ return cls(
68
+ **{
69
+ k: v
70
+ for k, v in params.items()
71
+ if k in inspect.signature(cls).parameters
72
+ }
73
+ )
74
+
75
+
76
+ @dataclass
77
+ class ModelConfig:
78
+ text_config: TextConfig
79
+ vision_config: VisionConfig
80
+ model_type: str
81
+ ignore_index: int = -100
82
+ image_token_id: int = 151655
83
+ video_token_id: int = 151656
84
+ vision_start_token_id: int = 151652
85
+ vision_feature_select_strategy: str = "default"
86
+ vision_feature_layer: int = -2
87
+ vocab_size: int = 32000
88
+ eos_token_id: Optional[List[int]] = None
89
+
90
+ @classmethod
91
+ def from_dict(cls, params):
92
+ # Copy text config parameters from root level
93
+ excluded_keys = {"vision_config"}
94
+ params["text_config"] = dict(
95
+ filter(lambda x: x[0] not in excluded_keys, params.items())
96
+ )
97
+
98
+ return cls(
99
+ **{
100
+ k: v
101
+ for k, v in params.items()
102
+ if k in inspect.signature(cls).parameters
103
+ }
104
+ )