nexaai 1.0.29__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.
Files changed (580) hide show
  1. nexaai/__init__.py +99 -0
  2. nexaai/_stub.cpython-310-darwin.so +0 -0
  3. nexaai/_version.py +4 -0
  4. nexaai/asr.py +68 -0
  5. nexaai/asr_impl/__init__.py +0 -0
  6. nexaai/asr_impl/mlx_asr_impl.py +93 -0
  7. nexaai/asr_impl/pybind_asr_impl.py +127 -0
  8. nexaai/base.py +39 -0
  9. nexaai/binds/__init__.py +7 -0
  10. nexaai/binds/asr_bind.cpython-310-darwin.so +0 -0
  11. nexaai/binds/common_bind.cpython-310-darwin.so +0 -0
  12. nexaai/binds/cpu_gpu/libggml-base.dylib +0 -0
  13. nexaai/binds/cpu_gpu/libggml-cpu.so +0 -0
  14. nexaai/binds/cpu_gpu/libggml-metal.so +0 -0
  15. nexaai/binds/cpu_gpu/libggml.dylib +0 -0
  16. nexaai/binds/cpu_gpu/libmtmd.dylib +0 -0
  17. nexaai/binds/cpu_gpu/libnexa_cpu_gpu.dylib +0 -0
  18. nexaai/binds/cpu_gpu/libnexa_plugin.dylib +0 -0
  19. nexaai/binds/cv_bind.cpython-310-darwin.so +0 -0
  20. nexaai/binds/diarize_bind.cpython-310-darwin.so +0 -0
  21. nexaai/binds/embedder_bind.cpython-310-darwin.so +0 -0
  22. nexaai/binds/libnexa_bridge.dylib +0 -0
  23. nexaai/binds/llm_bind.cpython-310-darwin.so +0 -0
  24. nexaai/binds/metal/libnexa_plugin.dylib +0 -0
  25. nexaai/binds/metal/py-lib/ml.py +888 -0
  26. nexaai/binds/metal/py-lib/mlx_audio/__init__.py +0 -0
  27. nexaai/binds/metal/py-lib/mlx_audio/codec/__init__.py +1 -0
  28. nexaai/binds/metal/py-lib/mlx_audio/codec/models/__init__.py +5 -0
  29. nexaai/binds/metal/py-lib/mlx_audio/codec/models/bigvgan/__init__.py +1 -0
  30. nexaai/binds/metal/py-lib/mlx_audio/codec/models/bigvgan/activation.py +51 -0
  31. nexaai/binds/metal/py-lib/mlx_audio/codec/models/bigvgan/amp.py +96 -0
  32. nexaai/binds/metal/py-lib/mlx_audio/codec/models/bigvgan/bigvgan.py +149 -0
  33. nexaai/binds/metal/py-lib/mlx_audio/codec/models/bigvgan/conv.py +114 -0
  34. nexaai/binds/metal/py-lib/mlx_audio/codec/models/bigvgan/resample.py +177 -0
  35. nexaai/binds/metal/py-lib/mlx_audio/codec/models/descript/__init__.py +1 -0
  36. nexaai/binds/metal/py-lib/mlx_audio/codec/models/descript/base.py +228 -0
  37. nexaai/binds/metal/py-lib/mlx_audio/codec/models/descript/dac.py +285 -0
  38. nexaai/binds/metal/py-lib/mlx_audio/codec/models/descript/nn/__init__.py +1 -0
  39. nexaai/binds/metal/py-lib/mlx_audio/codec/models/descript/nn/layers.py +129 -0
  40. nexaai/binds/metal/py-lib/mlx_audio/codec/models/descript/nn/quantize.py +149 -0
  41. nexaai/binds/metal/py-lib/mlx_audio/codec/models/encodec/__init__.py +1 -0
  42. nexaai/binds/metal/py-lib/mlx_audio/codec/models/encodec/encodec.py +777 -0
  43. nexaai/binds/metal/py-lib/mlx_audio/codec/models/mimi/__init__.py +1 -0
  44. nexaai/binds/metal/py-lib/mlx_audio/codec/models/mimi/mimi.py +286 -0
  45. nexaai/binds/metal/py-lib/mlx_audio/codec/models/mimi/modules/__init__.py +20 -0
  46. nexaai/binds/metal/py-lib/mlx_audio/codec/models/mimi/modules/conv.py +398 -0
  47. nexaai/binds/metal/py-lib/mlx_audio/codec/models/mimi/modules/kv_cache.py +199 -0
  48. nexaai/binds/metal/py-lib/mlx_audio/codec/models/mimi/modules/quantization.py +179 -0
  49. nexaai/binds/metal/py-lib/mlx_audio/codec/models/mimi/modules/seanet.py +314 -0
  50. nexaai/binds/metal/py-lib/mlx_audio/codec/models/mimi/modules/transformer.py +256 -0
  51. nexaai/binds/metal/py-lib/mlx_audio/codec/models/s3/__init__.py +1 -0
  52. nexaai/binds/metal/py-lib/mlx_audio/codec/models/s3/model.py +260 -0
  53. nexaai/binds/metal/py-lib/mlx_audio/codec/models/s3/model_v2.py +383 -0
  54. nexaai/binds/metal/py-lib/mlx_audio/codec/models/s3/utils.py +122 -0
  55. nexaai/binds/metal/py-lib/mlx_audio/codec/models/snac/__init__.py +1 -0
  56. nexaai/binds/metal/py-lib/mlx_audio/codec/models/snac/attention.py +97 -0
  57. nexaai/binds/metal/py-lib/mlx_audio/codec/models/snac/layers.py +306 -0
  58. nexaai/binds/metal/py-lib/mlx_audio/codec/models/snac/snac.py +154 -0
  59. nexaai/binds/metal/py-lib/mlx_audio/codec/models/snac/vq.py +135 -0
  60. nexaai/binds/metal/py-lib/mlx_audio/codec/models/vocos/__init__.py +1 -0
  61. nexaai/binds/metal/py-lib/mlx_audio/codec/models/vocos/mel.py +33 -0
  62. nexaai/binds/metal/py-lib/mlx_audio/codec/models/vocos/vocos.py +359 -0
  63. nexaai/binds/metal/py-lib/mlx_audio/codec/tests/__init__.py +0 -0
  64. nexaai/binds/metal/py-lib/mlx_audio/codec/tests/test_bigvgan.py +54 -0
  65. nexaai/binds/metal/py-lib/mlx_audio/codec/tests/test_descript.py +109 -0
  66. nexaai/binds/metal/py-lib/mlx_audio/codec/tests/test_encodec.py +58 -0
  67. nexaai/binds/metal/py-lib/mlx_audio/codec/tests/test_mimi.py +22 -0
  68. nexaai/binds/metal/py-lib/mlx_audio/codec/tests/test_s3.py +25 -0
  69. nexaai/binds/metal/py-lib/mlx_audio/codec/tests/test_snac.py +40 -0
  70. nexaai/binds/metal/py-lib/mlx_audio/codec/tests/test_vocos.py +93 -0
  71. nexaai/binds/metal/py-lib/mlx_audio/server.py +525 -0
  72. nexaai/binds/metal/py-lib/mlx_audio/sts/__init__.py +0 -0
  73. nexaai/binds/metal/py-lib/mlx_audio/sts/tests/test_voice_pipeline.py +156 -0
  74. nexaai/binds/metal/py-lib/mlx_audio/sts/voice_pipeline.py +327 -0
  75. nexaai/binds/metal/py-lib/mlx_audio/stt/__init__.py +0 -0
  76. nexaai/binds/metal/py-lib/mlx_audio/stt/generate.py +174 -0
  77. nexaai/binds/metal/py-lib/mlx_audio/stt/models/__init__.py +0 -0
  78. nexaai/binds/metal/py-lib/mlx_audio/stt/models/parakeet/__init__.py +1 -0
  79. nexaai/binds/metal/py-lib/mlx_audio/stt/models/parakeet/alignment.py +248 -0
  80. nexaai/binds/metal/py-lib/mlx_audio/stt/models/parakeet/attention.py +187 -0
  81. nexaai/binds/metal/py-lib/mlx_audio/stt/models/parakeet/audio.py +76 -0
  82. nexaai/binds/metal/py-lib/mlx_audio/stt/models/parakeet/conformer.py +331 -0
  83. nexaai/binds/metal/py-lib/mlx_audio/stt/models/parakeet/ctc.py +34 -0
  84. nexaai/binds/metal/py-lib/mlx_audio/stt/models/parakeet/parakeet.py +604 -0
  85. nexaai/binds/metal/py-lib/mlx_audio/stt/models/parakeet/rnnt.py +157 -0
  86. nexaai/binds/metal/py-lib/mlx_audio/stt/models/parakeet/tokenizer.py +2 -0
  87. nexaai/binds/metal/py-lib/mlx_audio/stt/models/wav2vec/feature_extractor.py +757 -0
  88. nexaai/binds/metal/py-lib/mlx_audio/stt/models/wav2vec/wav2vec.py +738 -0
  89. nexaai/binds/metal/py-lib/mlx_audio/stt/models/whisper/__init__.py +1 -0
  90. nexaai/binds/metal/py-lib/mlx_audio/stt/models/whisper/audio.py +82 -0
  91. nexaai/binds/metal/py-lib/mlx_audio/stt/models/whisper/decoding.py +742 -0
  92. nexaai/binds/metal/py-lib/mlx_audio/stt/models/whisper/timing.py +329 -0
  93. nexaai/binds/metal/py-lib/mlx_audio/stt/models/whisper/tokenizer.py +398 -0
  94. nexaai/binds/metal/py-lib/mlx_audio/stt/models/whisper/whisper.py +862 -0
  95. nexaai/binds/metal/py-lib/mlx_audio/stt/models/whisper/writers.py +268 -0
  96. nexaai/binds/metal/py-lib/mlx_audio/stt/tests/test_models.py +381 -0
  97. nexaai/binds/metal/py-lib/mlx_audio/stt/utils.py +195 -0
  98. nexaai/binds/metal/py-lib/mlx_audio/tts/__init__.py +1 -0
  99. nexaai/binds/metal/py-lib/mlx_audio/tts/audio_player.py +120 -0
  100. nexaai/binds/metal/py-lib/mlx_audio/tts/convert.py +71 -0
  101. nexaai/binds/metal/py-lib/mlx_audio/tts/generate.py +449 -0
  102. nexaai/binds/metal/py-lib/mlx_audio/tts/models/__init__.py +0 -0
  103. nexaai/binds/metal/py-lib/mlx_audio/tts/models/bark/__init__.py +4 -0
  104. nexaai/binds/metal/py-lib/mlx_audio/tts/models/bark/bark.py +528 -0
  105. nexaai/binds/metal/py-lib/mlx_audio/tts/models/bark/isftnet.py +12 -0
  106. nexaai/binds/metal/py-lib/mlx_audio/tts/models/bark/pipeline.py +442 -0
  107. nexaai/binds/metal/py-lib/mlx_audio/tts/models/base.py +84 -0
  108. nexaai/binds/metal/py-lib/mlx_audio/tts/models/dia/__init__.py +1 -0
  109. nexaai/binds/metal/py-lib/mlx_audio/tts/models/dia/audio.py +287 -0
  110. nexaai/binds/metal/py-lib/mlx_audio/tts/models/dia/config.py +256 -0
  111. nexaai/binds/metal/py-lib/mlx_audio/tts/models/dia/dia.py +592 -0
  112. nexaai/binds/metal/py-lib/mlx_audio/tts/models/dia/layers.py +870 -0
  113. nexaai/binds/metal/py-lib/mlx_audio/tts/models/indextts/__init__.py +3 -0
  114. nexaai/binds/metal/py-lib/mlx_audio/tts/models/indextts/attention.py +180 -0
  115. nexaai/binds/metal/py-lib/mlx_audio/tts/models/indextts/bigvgan.py +124 -0
  116. nexaai/binds/metal/py-lib/mlx_audio/tts/models/indextts/conformer.py +247 -0
  117. nexaai/binds/metal/py-lib/mlx_audio/tts/models/indextts/ecapa_tdnn/__init__.py +0 -0
  118. nexaai/binds/metal/py-lib/mlx_audio/tts/models/indextts/ecapa_tdnn/asp.py +59 -0
  119. nexaai/binds/metal/py-lib/mlx_audio/tts/models/indextts/ecapa_tdnn/ecapa_tdnn.py +91 -0
  120. nexaai/binds/metal/py-lib/mlx_audio/tts/models/indextts/ecapa_tdnn/se_res2net.py +132 -0
  121. nexaai/binds/metal/py-lib/mlx_audio/tts/models/indextts/ecapa_tdnn/tdnn.py +42 -0
  122. nexaai/binds/metal/py-lib/mlx_audio/tts/models/indextts/gpt2.py +38 -0
  123. nexaai/binds/metal/py-lib/mlx_audio/tts/models/indextts/indextts.py +412 -0
  124. nexaai/binds/metal/py-lib/mlx_audio/tts/models/indextts/mel.py +37 -0
  125. nexaai/binds/metal/py-lib/mlx_audio/tts/models/indextts/normalize.py +294 -0
  126. nexaai/binds/metal/py-lib/mlx_audio/tts/models/indextts/perceiver.py +62 -0
  127. nexaai/binds/metal/py-lib/mlx_audio/tts/models/interpolate.py +108 -0
  128. nexaai/binds/metal/py-lib/mlx_audio/tts/models/kokoro/__init__.py +4 -0
  129. nexaai/binds/metal/py-lib/mlx_audio/tts/models/kokoro/istftnet.py +979 -0
  130. nexaai/binds/metal/py-lib/mlx_audio/tts/models/kokoro/kokoro.py +331 -0
  131. nexaai/binds/metal/py-lib/mlx_audio/tts/models/kokoro/modules.py +659 -0
  132. nexaai/binds/metal/py-lib/mlx_audio/tts/models/kokoro/pipeline.py +453 -0
  133. nexaai/binds/metal/py-lib/mlx_audio/tts/models/kokoro/voice.py +113 -0
  134. nexaai/binds/metal/py-lib/mlx_audio/tts/models/llama/__init__.py +3 -0
  135. nexaai/binds/metal/py-lib/mlx_audio/tts/models/llama/llama.py +324 -0
  136. nexaai/binds/metal/py-lib/mlx_audio/tts/models/outetts/__init__.py +1 -0
  137. nexaai/binds/metal/py-lib/mlx_audio/tts/models/outetts/audio_processor.py +351 -0
  138. nexaai/binds/metal/py-lib/mlx_audio/tts/models/outetts/dac_interface.py +162 -0
  139. nexaai/binds/metal/py-lib/mlx_audio/tts/models/outetts/outetts.py +255 -0
  140. nexaai/binds/metal/py-lib/mlx_audio/tts/models/outetts/prompt_processor.py +181 -0
  141. nexaai/binds/metal/py-lib/mlx_audio/tts/models/outetts/tokens.py +36 -0
  142. nexaai/binds/metal/py-lib/mlx_audio/tts/models/sesame/__init__.py +3 -0
  143. nexaai/binds/metal/py-lib/mlx_audio/tts/models/sesame/attention.py +195 -0
  144. nexaai/binds/metal/py-lib/mlx_audio/tts/models/sesame/sesame.py +633 -0
  145. nexaai/binds/metal/py-lib/mlx_audio/tts/models/sesame/watermarking.py +105 -0
  146. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/__init__.py +1 -0
  147. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/audio_tokenizer.py +138 -0
  148. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/bicodec.py +269 -0
  149. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/__init__.py +0 -0
  150. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/blocks/__init__.py +0 -0
  151. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/blocks/sampler.py +111 -0
  152. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/encoder_decoder/__init__.py +0 -0
  153. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/encoder_decoder/feat_decoder.py +120 -0
  154. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/encoder_decoder/feat_encoder.py +136 -0
  155. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/encoder_decoder/wave_generator.py +113 -0
  156. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/finite_scalar_quantization.py +238 -0
  157. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/residual.py +209 -0
  158. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/residual_fsq.py +309 -0
  159. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/speaker/__init__.py +1 -0
  160. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/speaker/ecapa_tdnn.py +283 -0
  161. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/speaker/perceiver_encoder.py +326 -0
  162. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/speaker/pooling_layers.py +297 -0
  163. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/speaker/speaker_encoder.py +155 -0
  164. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/spark.py +382 -0
  165. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/utils/audio.py +220 -0
  166. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/utils/file.py +221 -0
  167. nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/utils/token_parser.py +181 -0
  168. nexaai/binds/metal/py-lib/mlx_audio/tts/tests/__init__.py +0 -0
  169. nexaai/binds/metal/py-lib/mlx_audio/tts/tests/test_base.py +66 -0
  170. nexaai/binds/metal/py-lib/mlx_audio/tts/tests/test_convert.py +173 -0
  171. nexaai/binds/metal/py-lib/mlx_audio/tts/tests/test_interpolate.py +88 -0
  172. nexaai/binds/metal/py-lib/mlx_audio/tts/tests/test_models.py +974 -0
  173. nexaai/binds/metal/py-lib/mlx_audio/tts/utils.py +337 -0
  174. nexaai/binds/metal/py-lib/mlx_audio/utils.py +237 -0
  175. nexaai/binds/metal/py-lib/mlx_audio/version.py +1 -0
  176. nexaai/binds/metal/py-lib/profiling.py +239 -0
  177. nexaai/binds/nexaml/libfftw3.3.dylib +0 -0
  178. nexaai/binds/nexaml/libfftw3f.3.dylib +0 -0
  179. nexaai/binds/nexaml/libggml-base.dylib +0 -0
  180. nexaai/binds/nexaml/libggml-cpu.so +0 -0
  181. nexaai/binds/nexaml/libggml-metal.so +0 -0
  182. nexaai/binds/nexaml/libggml.dylib +0 -0
  183. nexaai/binds/nexaml/libmp3lame.0.dylib +0 -0
  184. nexaai/binds/nexaml/libmpg123.0.dylib +0 -0
  185. nexaai/binds/nexaml/libnexa-mm-process.dylib +0 -0
  186. nexaai/binds/nexaml/libnexa-sampling.dylib +0 -0
  187. nexaai/binds/nexaml/libnexa_plugin.dylib +0 -0
  188. nexaai/binds/nexaml/libnexaproc.dylib +0 -0
  189. nexaai/binds/nexaml/libomp.dylib +0 -0
  190. nexaai/binds/nexaml/libqwen3-vl.dylib +0 -0
  191. nexaai/binds/nexaml/libqwen3vl-vision.dylib +0 -0
  192. nexaai/binds/rerank_bind.cpython-310-darwin.so +0 -0
  193. nexaai/binds/vlm_bind.cpython-310-darwin.so +0 -0
  194. nexaai/common.py +106 -0
  195. nexaai/cv.py +95 -0
  196. nexaai/cv_impl/__init__.py +0 -0
  197. nexaai/cv_impl/mlx_cv_impl.py +91 -0
  198. nexaai/cv_impl/pybind_cv_impl.py +124 -0
  199. nexaai/diarize.py +80 -0
  200. nexaai/diarize_impl/__init__.py +1 -0
  201. nexaai/diarize_impl/pybind_diarize_impl.py +125 -0
  202. nexaai/embedder.py +73 -0
  203. nexaai/embedder_impl/__init__.py +0 -0
  204. nexaai/embedder_impl/mlx_embedder_impl.py +118 -0
  205. nexaai/embedder_impl/pybind_embedder_impl.py +96 -0
  206. nexaai/image_gen.py +141 -0
  207. nexaai/image_gen_impl/__init__.py +0 -0
  208. nexaai/image_gen_impl/mlx_image_gen_impl.py +292 -0
  209. nexaai/image_gen_impl/pybind_image_gen_impl.py +85 -0
  210. nexaai/llm.py +98 -0
  211. nexaai/llm_impl/__init__.py +0 -0
  212. nexaai/llm_impl/mlx_llm_impl.py +271 -0
  213. nexaai/llm_impl/pybind_llm_impl.py +238 -0
  214. nexaai/log.py +92 -0
  215. nexaai/mlx_backend/asr/__init__.py +12 -0
  216. nexaai/mlx_backend/asr/interface.py +122 -0
  217. nexaai/mlx_backend/common/__init__.py +0 -0
  218. nexaai/mlx_backend/common/utils.py +25 -0
  219. nexaai/mlx_backend/cv/__init__.py +0 -0
  220. nexaai/mlx_backend/cv/generate.py +195 -0
  221. nexaai/mlx_backend/cv/interface.py +162 -0
  222. nexaai/mlx_backend/cv/main.py +81 -0
  223. nexaai/mlx_backend/cv/modeling/pp_ocr_v4.py +1736 -0
  224. nexaai/mlx_backend/embedding/__init__.py +0 -0
  225. nexaai/mlx_backend/embedding/generate.py +333 -0
  226. nexaai/mlx_backend/embedding/interface.py +617 -0
  227. nexaai/mlx_backend/embedding/main.py +173 -0
  228. nexaai/mlx_backend/embedding/modeling/__init__.py +0 -0
  229. nexaai/mlx_backend/embedding/modeling/nexa_jina_v2.py +399 -0
  230. nexaai/mlx_backend/image_gen/__init__.py +1 -0
  231. nexaai/mlx_backend/image_gen/generate_sd.py +244 -0
  232. nexaai/mlx_backend/image_gen/interface.py +82 -0
  233. nexaai/mlx_backend/image_gen/main.py +281 -0
  234. nexaai/mlx_backend/image_gen/stable_diffusion/__init__.py +306 -0
  235. nexaai/mlx_backend/image_gen/stable_diffusion/clip.py +116 -0
  236. nexaai/mlx_backend/image_gen/stable_diffusion/config.py +65 -0
  237. nexaai/mlx_backend/image_gen/stable_diffusion/model_io.py +386 -0
  238. nexaai/mlx_backend/image_gen/stable_diffusion/sampler.py +105 -0
  239. nexaai/mlx_backend/image_gen/stable_diffusion/tokenizer.py +100 -0
  240. nexaai/mlx_backend/image_gen/stable_diffusion/unet.py +460 -0
  241. nexaai/mlx_backend/image_gen/stable_diffusion/vae.py +274 -0
  242. nexaai/mlx_backend/llm/__init__.py +0 -0
  243. nexaai/mlx_backend/llm/generate.py +149 -0
  244. nexaai/mlx_backend/llm/interface.py +764 -0
  245. nexaai/mlx_backend/llm/main.py +68 -0
  246. nexaai/mlx_backend/ml.py +888 -0
  247. nexaai/mlx_backend/mlx_audio/__init__.py +0 -0
  248. nexaai/mlx_backend/mlx_audio/codec/__init__.py +1 -0
  249. nexaai/mlx_backend/mlx_audio/codec/models/__init__.py +5 -0
  250. nexaai/mlx_backend/mlx_audio/codec/models/bigvgan/__init__.py +1 -0
  251. nexaai/mlx_backend/mlx_audio/codec/models/bigvgan/activation.py +51 -0
  252. nexaai/mlx_backend/mlx_audio/codec/models/bigvgan/amp.py +96 -0
  253. nexaai/mlx_backend/mlx_audio/codec/models/bigvgan/bigvgan.py +149 -0
  254. nexaai/mlx_backend/mlx_audio/codec/models/bigvgan/conv.py +114 -0
  255. nexaai/mlx_backend/mlx_audio/codec/models/bigvgan/resample.py +177 -0
  256. nexaai/mlx_backend/mlx_audio/codec/models/descript/__init__.py +1 -0
  257. nexaai/mlx_backend/mlx_audio/codec/models/descript/base.py +228 -0
  258. nexaai/mlx_backend/mlx_audio/codec/models/descript/dac.py +285 -0
  259. nexaai/mlx_backend/mlx_audio/codec/models/descript/nn/__init__.py +1 -0
  260. nexaai/mlx_backend/mlx_audio/codec/models/descript/nn/layers.py +129 -0
  261. nexaai/mlx_backend/mlx_audio/codec/models/descript/nn/quantize.py +149 -0
  262. nexaai/mlx_backend/mlx_audio/codec/models/encodec/__init__.py +1 -0
  263. nexaai/mlx_backend/mlx_audio/codec/models/encodec/encodec.py +777 -0
  264. nexaai/mlx_backend/mlx_audio/codec/models/mimi/__init__.py +1 -0
  265. nexaai/mlx_backend/mlx_audio/codec/models/mimi/mimi.py +286 -0
  266. nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/__init__.py +20 -0
  267. nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/conv.py +398 -0
  268. nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/kv_cache.py +199 -0
  269. nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/quantization.py +179 -0
  270. nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/seanet.py +314 -0
  271. nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/transformer.py +256 -0
  272. nexaai/mlx_backend/mlx_audio/codec/models/s3/__init__.py +1 -0
  273. nexaai/mlx_backend/mlx_audio/codec/models/s3/model.py +260 -0
  274. nexaai/mlx_backend/mlx_audio/codec/models/s3/model_v2.py +383 -0
  275. nexaai/mlx_backend/mlx_audio/codec/models/s3/utils.py +122 -0
  276. nexaai/mlx_backend/mlx_audio/codec/models/snac/__init__.py +1 -0
  277. nexaai/mlx_backend/mlx_audio/codec/models/snac/attention.py +97 -0
  278. nexaai/mlx_backend/mlx_audio/codec/models/snac/layers.py +306 -0
  279. nexaai/mlx_backend/mlx_audio/codec/models/snac/snac.py +154 -0
  280. nexaai/mlx_backend/mlx_audio/codec/models/snac/vq.py +135 -0
  281. nexaai/mlx_backend/mlx_audio/codec/models/vocos/__init__.py +1 -0
  282. nexaai/mlx_backend/mlx_audio/codec/models/vocos/mel.py +33 -0
  283. nexaai/mlx_backend/mlx_audio/codec/models/vocos/vocos.py +359 -0
  284. nexaai/mlx_backend/mlx_audio/codec/tests/__init__.py +0 -0
  285. nexaai/mlx_backend/mlx_audio/codec/tests/test_bigvgan.py +54 -0
  286. nexaai/mlx_backend/mlx_audio/codec/tests/test_descript.py +109 -0
  287. nexaai/mlx_backend/mlx_audio/codec/tests/test_encodec.py +58 -0
  288. nexaai/mlx_backend/mlx_audio/codec/tests/test_mimi.py +22 -0
  289. nexaai/mlx_backend/mlx_audio/codec/tests/test_s3.py +25 -0
  290. nexaai/mlx_backend/mlx_audio/codec/tests/test_snac.py +40 -0
  291. nexaai/mlx_backend/mlx_audio/codec/tests/test_vocos.py +93 -0
  292. nexaai/mlx_backend/mlx_audio/server.py +525 -0
  293. nexaai/mlx_backend/mlx_audio/sts/__init__.py +0 -0
  294. nexaai/mlx_backend/mlx_audio/sts/tests/test_voice_pipeline.py +156 -0
  295. nexaai/mlx_backend/mlx_audio/sts/voice_pipeline.py +327 -0
  296. nexaai/mlx_backend/mlx_audio/stt/__init__.py +0 -0
  297. nexaai/mlx_backend/mlx_audio/stt/generate.py +174 -0
  298. nexaai/mlx_backend/mlx_audio/stt/models/__init__.py +0 -0
  299. nexaai/mlx_backend/mlx_audio/stt/models/parakeet/__init__.py +1 -0
  300. nexaai/mlx_backend/mlx_audio/stt/models/parakeet/alignment.py +248 -0
  301. nexaai/mlx_backend/mlx_audio/stt/models/parakeet/attention.py +187 -0
  302. nexaai/mlx_backend/mlx_audio/stt/models/parakeet/audio.py +76 -0
  303. nexaai/mlx_backend/mlx_audio/stt/models/parakeet/conformer.py +331 -0
  304. nexaai/mlx_backend/mlx_audio/stt/models/parakeet/ctc.py +34 -0
  305. nexaai/mlx_backend/mlx_audio/stt/models/parakeet/parakeet.py +604 -0
  306. nexaai/mlx_backend/mlx_audio/stt/models/parakeet/rnnt.py +157 -0
  307. nexaai/mlx_backend/mlx_audio/stt/models/parakeet/tokenizer.py +2 -0
  308. nexaai/mlx_backend/mlx_audio/stt/models/wav2vec/feature_extractor.py +757 -0
  309. nexaai/mlx_backend/mlx_audio/stt/models/wav2vec/wav2vec.py +738 -0
  310. nexaai/mlx_backend/mlx_audio/stt/models/whisper/__init__.py +1 -0
  311. nexaai/mlx_backend/mlx_audio/stt/models/whisper/audio.py +82 -0
  312. nexaai/mlx_backend/mlx_audio/stt/models/whisper/decoding.py +742 -0
  313. nexaai/mlx_backend/mlx_audio/stt/models/whisper/timing.py +329 -0
  314. nexaai/mlx_backend/mlx_audio/stt/models/whisper/tokenizer.py +398 -0
  315. nexaai/mlx_backend/mlx_audio/stt/models/whisper/whisper.py +862 -0
  316. nexaai/mlx_backend/mlx_audio/stt/models/whisper/writers.py +268 -0
  317. nexaai/mlx_backend/mlx_audio/stt/tests/test_models.py +381 -0
  318. nexaai/mlx_backend/mlx_audio/stt/utils.py +195 -0
  319. nexaai/mlx_backend/mlx_audio/tts/__init__.py +1 -0
  320. nexaai/mlx_backend/mlx_audio/tts/audio_player.py +120 -0
  321. nexaai/mlx_backend/mlx_audio/tts/convert.py +71 -0
  322. nexaai/mlx_backend/mlx_audio/tts/generate.py +449 -0
  323. nexaai/mlx_backend/mlx_audio/tts/models/__init__.py +0 -0
  324. nexaai/mlx_backend/mlx_audio/tts/models/bark/__init__.py +4 -0
  325. nexaai/mlx_backend/mlx_audio/tts/models/bark/bark.py +528 -0
  326. nexaai/mlx_backend/mlx_audio/tts/models/bark/isftnet.py +12 -0
  327. nexaai/mlx_backend/mlx_audio/tts/models/bark/pipeline.py +442 -0
  328. nexaai/mlx_backend/mlx_audio/tts/models/base.py +84 -0
  329. nexaai/mlx_backend/mlx_audio/tts/models/dia/__init__.py +1 -0
  330. nexaai/mlx_backend/mlx_audio/tts/models/dia/audio.py +287 -0
  331. nexaai/mlx_backend/mlx_audio/tts/models/dia/config.py +256 -0
  332. nexaai/mlx_backend/mlx_audio/tts/models/dia/dia.py +592 -0
  333. nexaai/mlx_backend/mlx_audio/tts/models/dia/layers.py +870 -0
  334. nexaai/mlx_backend/mlx_audio/tts/models/indextts/__init__.py +3 -0
  335. nexaai/mlx_backend/mlx_audio/tts/models/indextts/attention.py +180 -0
  336. nexaai/mlx_backend/mlx_audio/tts/models/indextts/bigvgan.py +124 -0
  337. nexaai/mlx_backend/mlx_audio/tts/models/indextts/conformer.py +247 -0
  338. nexaai/mlx_backend/mlx_audio/tts/models/indextts/ecapa_tdnn/__init__.py +0 -0
  339. nexaai/mlx_backend/mlx_audio/tts/models/indextts/ecapa_tdnn/asp.py +59 -0
  340. nexaai/mlx_backend/mlx_audio/tts/models/indextts/ecapa_tdnn/ecapa_tdnn.py +91 -0
  341. nexaai/mlx_backend/mlx_audio/tts/models/indextts/ecapa_tdnn/se_res2net.py +132 -0
  342. nexaai/mlx_backend/mlx_audio/tts/models/indextts/ecapa_tdnn/tdnn.py +42 -0
  343. nexaai/mlx_backend/mlx_audio/tts/models/indextts/gpt2.py +38 -0
  344. nexaai/mlx_backend/mlx_audio/tts/models/indextts/indextts.py +412 -0
  345. nexaai/mlx_backend/mlx_audio/tts/models/indextts/mel.py +37 -0
  346. nexaai/mlx_backend/mlx_audio/tts/models/indextts/normalize.py +294 -0
  347. nexaai/mlx_backend/mlx_audio/tts/models/indextts/perceiver.py +62 -0
  348. nexaai/mlx_backend/mlx_audio/tts/models/interpolate.py +108 -0
  349. nexaai/mlx_backend/mlx_audio/tts/models/kokoro/__init__.py +4 -0
  350. nexaai/mlx_backend/mlx_audio/tts/models/kokoro/istftnet.py +979 -0
  351. nexaai/mlx_backend/mlx_audio/tts/models/kokoro/kokoro.py +331 -0
  352. nexaai/mlx_backend/mlx_audio/tts/models/kokoro/modules.py +659 -0
  353. nexaai/mlx_backend/mlx_audio/tts/models/kokoro/pipeline.py +453 -0
  354. nexaai/mlx_backend/mlx_audio/tts/models/kokoro/voice.py +113 -0
  355. nexaai/mlx_backend/mlx_audio/tts/models/llama/__init__.py +3 -0
  356. nexaai/mlx_backend/mlx_audio/tts/models/llama/llama.py +324 -0
  357. nexaai/mlx_backend/mlx_audio/tts/models/outetts/__init__.py +1 -0
  358. nexaai/mlx_backend/mlx_audio/tts/models/outetts/audio_processor.py +351 -0
  359. nexaai/mlx_backend/mlx_audio/tts/models/outetts/dac_interface.py +162 -0
  360. nexaai/mlx_backend/mlx_audio/tts/models/outetts/default_speaker.json +461 -0
  361. nexaai/mlx_backend/mlx_audio/tts/models/outetts/outetts.py +255 -0
  362. nexaai/mlx_backend/mlx_audio/tts/models/outetts/prompt_processor.py +181 -0
  363. nexaai/mlx_backend/mlx_audio/tts/models/outetts/tokens.py +36 -0
  364. nexaai/mlx_backend/mlx_audio/tts/models/sesame/__init__.py +3 -0
  365. nexaai/mlx_backend/mlx_audio/tts/models/sesame/attention.py +195 -0
  366. nexaai/mlx_backend/mlx_audio/tts/models/sesame/sesame.py +633 -0
  367. nexaai/mlx_backend/mlx_audio/tts/models/sesame/watermarking.py +105 -0
  368. nexaai/mlx_backend/mlx_audio/tts/models/spark/__init__.py +1 -0
  369. nexaai/mlx_backend/mlx_audio/tts/models/spark/audio_tokenizer.py +138 -0
  370. nexaai/mlx_backend/mlx_audio/tts/models/spark/bicodec.py +269 -0
  371. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/__init__.py +0 -0
  372. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/blocks/__init__.py +0 -0
  373. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/blocks/sampler.py +111 -0
  374. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/encoder_decoder/__init__.py +0 -0
  375. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/encoder_decoder/feat_decoder.py +120 -0
  376. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/encoder_decoder/feat_encoder.py +136 -0
  377. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/encoder_decoder/wave_generator.py +113 -0
  378. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/finite_scalar_quantization.py +238 -0
  379. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/residual.py +209 -0
  380. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/residual_fsq.py +309 -0
  381. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/speaker/__init__.py +1 -0
  382. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/speaker/ecapa_tdnn.py +283 -0
  383. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/speaker/perceiver_encoder.py +326 -0
  384. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/speaker/pooling_layers.py +297 -0
  385. nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/speaker/speaker_encoder.py +155 -0
  386. nexaai/mlx_backend/mlx_audio/tts/models/spark/spark.py +382 -0
  387. nexaai/mlx_backend/mlx_audio/tts/models/spark/utils/audio.py +220 -0
  388. nexaai/mlx_backend/mlx_audio/tts/models/spark/utils/file.py +221 -0
  389. nexaai/mlx_backend/mlx_audio/tts/models/spark/utils/token_parser.py +181 -0
  390. nexaai/mlx_backend/mlx_audio/tts/tests/__init__.py +0 -0
  391. nexaai/mlx_backend/mlx_audio/tts/tests/test_base.py +66 -0
  392. nexaai/mlx_backend/mlx_audio/tts/tests/test_convert.py +173 -0
  393. nexaai/mlx_backend/mlx_audio/tts/tests/test_interpolate.py +88 -0
  394. nexaai/mlx_backend/mlx_audio/tts/tests/test_models.py +974 -0
  395. nexaai/mlx_backend/mlx_audio/tts/utils.py +337 -0
  396. nexaai/mlx_backend/mlx_audio/utils.py +237 -0
  397. nexaai/mlx_backend/mlx_audio/version.py +1 -0
  398. nexaai/mlx_backend/profiling.py +239 -0
  399. nexaai/mlx_backend/rerank/__init__.py +0 -0
  400. nexaai/mlx_backend/rerank/generate.py +174 -0
  401. nexaai/mlx_backend/rerank/interface.py +287 -0
  402. nexaai/mlx_backend/rerank/main.py +127 -0
  403. nexaai/mlx_backend/rerank/modeling/__init__.py +0 -0
  404. nexaai/mlx_backend/rerank/modeling/nexa_jina_rerank.py +330 -0
  405. nexaai/mlx_backend/sd/__init__.py +1 -0
  406. nexaai/mlx_backend/sd/interface.py +362 -0
  407. nexaai/mlx_backend/sd/main.py +286 -0
  408. nexaai/mlx_backend/sd/modeling/__init__.py +306 -0
  409. nexaai/mlx_backend/sd/modeling/clip.py +116 -0
  410. nexaai/mlx_backend/sd/modeling/config.py +65 -0
  411. nexaai/mlx_backend/sd/modeling/model_io.py +385 -0
  412. nexaai/mlx_backend/sd/modeling/sampler.py +105 -0
  413. nexaai/mlx_backend/sd/modeling/tokenizer.py +100 -0
  414. nexaai/mlx_backend/sd/modeling/unet.py +460 -0
  415. nexaai/mlx_backend/sd/modeling/vae.py +274 -0
  416. nexaai/mlx_backend/tts/__init__.py +12 -0
  417. nexaai/mlx_backend/tts/interface.py +276 -0
  418. nexaai/mlx_backend/vlm/__init__.py +3 -0
  419. nexaai/mlx_backend/vlm/generate.py +572 -0
  420. nexaai/mlx_backend/vlm/generate_qwen3_vl.py +374 -0
  421. nexaai/mlx_backend/vlm/generate_qwen3_vl_moe.py +259 -0
  422. nexaai/mlx_backend/vlm/interface.py +559 -0
  423. nexaai/mlx_backend/vlm/main.py +365 -0
  424. nexaai/mlx_backend/vlm/modeling/__init__.py +0 -0
  425. nexaai/mlx_backend/vlm/modeling/convert.py +68 -0
  426. nexaai/mlx_backend/vlm/modeling/models/__init__.py +0 -0
  427. nexaai/mlx_backend/vlm/modeling/models/aya_vision/__init__.py +8 -0
  428. nexaai/mlx_backend/vlm/modeling/models/aya_vision/aya_vision.py +193 -0
  429. nexaai/mlx_backend/vlm/modeling/models/aya_vision/interpolate.py +186 -0
  430. nexaai/mlx_backend/vlm/modeling/models/aya_vision/language.py +233 -0
  431. nexaai/mlx_backend/vlm/modeling/models/aya_vision/vision.py +503 -0
  432. nexaai/mlx_backend/vlm/modeling/models/base.py +202 -0
  433. nexaai/mlx_backend/vlm/modeling/models/cache.py +230 -0
  434. nexaai/mlx_backend/vlm/modeling/models/deepseek_vl_v2/__init__.py +10 -0
  435. nexaai/mlx_backend/vlm/modeling/models/deepseek_vl_v2/conversation.py +264 -0
  436. nexaai/mlx_backend/vlm/modeling/models/deepseek_vl_v2/deepseek_vl_v2.py +472 -0
  437. nexaai/mlx_backend/vlm/modeling/models/deepseek_vl_v2/language.py +591 -0
  438. nexaai/mlx_backend/vlm/modeling/models/deepseek_vl_v2/processing_deepsek_vl_v2.py +526 -0
  439. nexaai/mlx_backend/vlm/modeling/models/deepseek_vl_v2/vision.py +356 -0
  440. nexaai/mlx_backend/vlm/modeling/models/florence2/__init__.py +8 -0
  441. nexaai/mlx_backend/vlm/modeling/models/florence2/florence2.py +366 -0
  442. nexaai/mlx_backend/vlm/modeling/models/florence2/language.py +488 -0
  443. nexaai/mlx_backend/vlm/modeling/models/florence2/vision.py +591 -0
  444. nexaai/mlx_backend/vlm/modeling/models/gemma3/__init__.py +8 -0
  445. nexaai/mlx_backend/vlm/modeling/models/gemma3/gemma3.py +213 -0
  446. nexaai/mlx_backend/vlm/modeling/models/gemma3/language.py +315 -0
  447. nexaai/mlx_backend/vlm/modeling/models/gemma3/vision.py +238 -0
  448. nexaai/mlx_backend/vlm/modeling/models/gemma3n/__init__.py +2 -0
  449. nexaai/mlx_backend/vlm/modeling/models/gemma3n/audio.py +1038 -0
  450. nexaai/mlx_backend/vlm/modeling/models/gemma3n/config.py +139 -0
  451. nexaai/mlx_backend/vlm/modeling/models/gemma3n/gemma3n.py +322 -0
  452. nexaai/mlx_backend/vlm/modeling/models/gemma3n/language.py +629 -0
  453. nexaai/mlx_backend/vlm/modeling/models/gemma3n/vision.py +1022 -0
  454. nexaai/mlx_backend/vlm/modeling/models/idefics2/__init__.py +9 -0
  455. nexaai/mlx_backend/vlm/modeling/models/idefics2/idefics2.py +294 -0
  456. nexaai/mlx_backend/vlm/modeling/models/idefics2/language.py +191 -0
  457. nexaai/mlx_backend/vlm/modeling/models/idefics2/vision.py +267 -0
  458. nexaai/mlx_backend/vlm/modeling/models/idefics3/__init__.py +8 -0
  459. nexaai/mlx_backend/vlm/modeling/models/idefics3/idefics3.py +175 -0
  460. nexaai/mlx_backend/vlm/modeling/models/idefics3/language.py +192 -0
  461. nexaai/mlx_backend/vlm/modeling/models/idefics3/vision.py +233 -0
  462. nexaai/mlx_backend/vlm/modeling/models/internvl_chat/__init__.py +9 -0
  463. nexaai/mlx_backend/vlm/modeling/models/internvl_chat/internvl_chat.py +140 -0
  464. nexaai/mlx_backend/vlm/modeling/models/internvl_chat/language.py +220 -0
  465. nexaai/mlx_backend/vlm/modeling/models/internvl_chat/processor.py +393 -0
  466. nexaai/mlx_backend/vlm/modeling/models/internvl_chat/vision.py +293 -0
  467. nexaai/mlx_backend/vlm/modeling/models/kernels.py +307 -0
  468. nexaai/mlx_backend/vlm/modeling/models/kimi_vl/__init__.py +8 -0
  469. nexaai/mlx_backend/vlm/modeling/models/kimi_vl/kimi_vl.py +143 -0
  470. nexaai/mlx_backend/vlm/modeling/models/kimi_vl/language.py +509 -0
  471. nexaai/mlx_backend/vlm/modeling/models/kimi_vl/vision.py +522 -0
  472. nexaai/mlx_backend/vlm/modeling/models/llama4/__init__.py +8 -0
  473. nexaai/mlx_backend/vlm/modeling/models/llama4/language.py +386 -0
  474. nexaai/mlx_backend/vlm/modeling/models/llama4/llama4.py +138 -0
  475. nexaai/mlx_backend/vlm/modeling/models/llama4/vision.py +560 -0
  476. nexaai/mlx_backend/vlm/modeling/models/llava/__init__.py +8 -0
  477. nexaai/mlx_backend/vlm/modeling/models/llava/language.py +240 -0
  478. nexaai/mlx_backend/vlm/modeling/models/llava/llava.py +153 -0
  479. nexaai/mlx_backend/vlm/modeling/models/llava/vision.py +259 -0
  480. nexaai/mlx_backend/vlm/modeling/models/llava_bunny/__init__.py +9 -0
  481. nexaai/mlx_backend/vlm/modeling/models/llava_bunny/language.py +236 -0
  482. nexaai/mlx_backend/vlm/modeling/models/llava_bunny/llava_bunny.py +256 -0
  483. nexaai/mlx_backend/vlm/modeling/models/llava_bunny/vision.py +303 -0
  484. nexaai/mlx_backend/vlm/modeling/models/llava_next/__init__.py +8 -0
  485. nexaai/mlx_backend/vlm/modeling/models/llava_next/language.py +230 -0
  486. nexaai/mlx_backend/vlm/modeling/models/llava_next/llava_next.py +160 -0
  487. nexaai/mlx_backend/vlm/modeling/models/llava_next/vision.py +243 -0
  488. nexaai/mlx_backend/vlm/modeling/models/mistral3/__init__.py +8 -0
  489. nexaai/mlx_backend/vlm/modeling/models/mistral3/mistral3.py +283 -0
  490. nexaai/mlx_backend/vlm/modeling/models/mllama/__init__.py +8 -0
  491. nexaai/mlx_backend/vlm/modeling/models/mllama/language.py +416 -0
  492. nexaai/mlx_backend/vlm/modeling/models/mllama/mllama.py +172 -0
  493. nexaai/mlx_backend/vlm/modeling/models/mllama/vision.py +499 -0
  494. nexaai/mlx_backend/vlm/modeling/models/molmo/__init__.py +8 -0
  495. nexaai/mlx_backend/vlm/modeling/models/molmo/language.py +243 -0
  496. nexaai/mlx_backend/vlm/modeling/models/molmo/molmo.py +133 -0
  497. nexaai/mlx_backend/vlm/modeling/models/molmo/vision.py +465 -0
  498. nexaai/mlx_backend/vlm/modeling/models/multi_modality/__init__.py +10 -0
  499. nexaai/mlx_backend/vlm/modeling/models/multi_modality/language.py +230 -0
  500. nexaai/mlx_backend/vlm/modeling/models/multi_modality/multi_modality.py +385 -0
  501. nexaai/mlx_backend/vlm/modeling/models/multi_modality/sam.py +557 -0
  502. nexaai/mlx_backend/vlm/modeling/models/multi_modality/vision.py +526 -0
  503. nexaai/mlx_backend/vlm/modeling/models/paligemma/__init__.py +8 -0
  504. nexaai/mlx_backend/vlm/modeling/models/paligemma/language.py +282 -0
  505. nexaai/mlx_backend/vlm/modeling/models/paligemma/paligemma.py +160 -0
  506. nexaai/mlx_backend/vlm/modeling/models/paligemma/vision.py +242 -0
  507. nexaai/mlx_backend/vlm/modeling/models/phi3_v/__init__.py +8 -0
  508. nexaai/mlx_backend/vlm/modeling/models/phi3_v/language.py +21 -0
  509. nexaai/mlx_backend/vlm/modeling/models/phi3_v/phi3_v.py +243 -0
  510. nexaai/mlx_backend/vlm/modeling/models/phi3_v/su_rope.py +71 -0
  511. nexaai/mlx_backend/vlm/modeling/models/phi3_v/vision.py +324 -0
  512. nexaai/mlx_backend/vlm/modeling/models/pixtral/__init__.py +8 -0
  513. nexaai/mlx_backend/vlm/modeling/models/pixtral/language.py +229 -0
  514. nexaai/mlx_backend/vlm/modeling/models/pixtral/pixtral.py +161 -0
  515. nexaai/mlx_backend/vlm/modeling/models/pixtral/vision.py +320 -0
  516. nexaai/mlx_backend/vlm/modeling/models/qwen2_5_vl/__init__.py +2 -0
  517. nexaai/mlx_backend/vlm/modeling/models/qwen2_5_vl/config.py +108 -0
  518. nexaai/mlx_backend/vlm/modeling/models/qwen2_5_vl/language.py +490 -0
  519. nexaai/mlx_backend/vlm/modeling/models/qwen2_5_vl/qwen2_5_vl.py +168 -0
  520. nexaai/mlx_backend/vlm/modeling/models/qwen2_5_vl/vision.py +414 -0
  521. nexaai/mlx_backend/vlm/modeling/models/qwen2_vl/__init__.py +2 -0
  522. nexaai/mlx_backend/vlm/modeling/models/qwen2_vl/config.py +104 -0
  523. nexaai/mlx_backend/vlm/modeling/models/qwen2_vl/language.py +490 -0
  524. nexaai/mlx_backend/vlm/modeling/models/qwen2_vl/qwen2_vl.py +167 -0
  525. nexaai/mlx_backend/vlm/modeling/models/qwen2_vl/vision.py +312 -0
  526. nexaai/mlx_backend/vlm/modeling/models/qwen3_vl/llm_common/__init__.py +0 -0
  527. nexaai/mlx_backend/vlm/modeling/models/qwen3_vl/llm_common/base.py +117 -0
  528. nexaai/mlx_backend/vlm/modeling/models/qwen3_vl/llm_common/cache.py +531 -0
  529. nexaai/mlx_backend/vlm/modeling/models/qwen3_vl/llm_common/generate.py +701 -0
  530. nexaai/mlx_backend/vlm/modeling/models/qwen3_vl/llm_common/rope_utils.py +255 -0
  531. nexaai/mlx_backend/vlm/modeling/models/qwen3_vl/llm_common/sample_utils.py +303 -0
  532. nexaai/mlx_backend/vlm/modeling/models/qwen3_vl/llm_common/tokenizer_utils.py +407 -0
  533. nexaai/mlx_backend/vlm/modeling/models/qwen3_vl/processor.py +476 -0
  534. nexaai/mlx_backend/vlm/modeling/models/qwen3_vl/qwen3vl.py +1262 -0
  535. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/llm_common/__init__.py +0 -0
  536. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/llm_common/base.py +117 -0
  537. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/llm_common/cache.py +531 -0
  538. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/llm_common/generate.py +701 -0
  539. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/llm_common/rope_utils.py +255 -0
  540. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/llm_common/sample_utils.py +303 -0
  541. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/llm_common/tokenizer_utils.py +407 -0
  542. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/processor.py +476 -0
  543. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/qwen3vl_moe.py +1308 -0
  544. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/switch_layers.py +210 -0
  545. nexaai/mlx_backend/vlm/modeling/models/smolvlm/__init__.py +8 -0
  546. nexaai/mlx_backend/vlm/modeling/models/smolvlm/smolvlm.py +62 -0
  547. nexaai/mlx_backend/vlm/modeling/processing_qwen2_5_vl.py +209 -0
  548. nexaai/mlx_backend/vlm/modeling/processing_qwen2_vl.py +215 -0
  549. nexaai/mlx_backend/vlm/modeling/prompt_utils.py +474 -0
  550. nexaai/mlx_backend/vlm/modeling/sample_utils.py +39 -0
  551. nexaai/mlx_backend/vlm/modeling/tokenizer_utils.py +344 -0
  552. nexaai/mlx_backend/vlm/modeling/trainer/__init__.py +9 -0
  553. nexaai/mlx_backend/vlm/modeling/trainer/lora.py +70 -0
  554. nexaai/mlx_backend/vlm/modeling/trainer/trainer.py +296 -0
  555. nexaai/mlx_backend/vlm/modeling/trainer/utils.py +160 -0
  556. nexaai/mlx_backend/vlm/modeling/utils.py +928 -0
  557. nexaai/rerank.py +57 -0
  558. nexaai/rerank_impl/__init__.py +0 -0
  559. nexaai/rerank_impl/mlx_rerank_impl.py +94 -0
  560. nexaai/rerank_impl/pybind_rerank_impl.py +136 -0
  561. nexaai/runtime.py +68 -0
  562. nexaai/runtime_error.py +24 -0
  563. nexaai/tts.py +75 -0
  564. nexaai/tts_impl/__init__.py +0 -0
  565. nexaai/tts_impl/mlx_tts_impl.py +94 -0
  566. nexaai/tts_impl/pybind_tts_impl.py +43 -0
  567. nexaai/utils/decode.py +18 -0
  568. nexaai/utils/manifest_utils.py +531 -0
  569. nexaai/utils/model_manager.py +1745 -0
  570. nexaai/utils/model_types.py +49 -0
  571. nexaai/utils/progress_tracker.py +389 -0
  572. nexaai/utils/quantization_utils.py +245 -0
  573. nexaai/vlm.py +130 -0
  574. nexaai/vlm_impl/__init__.py +0 -0
  575. nexaai/vlm_impl/mlx_vlm_impl.py +259 -0
  576. nexaai/vlm_impl/pybind_vlm_impl.py +275 -0
  577. nexaai-1.0.29.dist-info/METADATA +35 -0
  578. nexaai-1.0.29.dist-info/RECORD +580 -0
  579. nexaai-1.0.29.dist-info/WHEEL +5 -0
  580. nexaai-1.0.29.dist-info/top_level.txt +1 -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
+ )