nexaai 1.0.19rc7__cp310-cp310-macosx_14_0_universal2.whl → 1.0.19rc9__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 (200) 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/mlx_backend/vlm/generate_qwen3_vl.py +14 -31
  5. nexaai/mlx_backend/vlm/generate_qwen3_vl_moe.py +15 -32
  6. nexaai/mlx_backend/vlm/modeling/models/qwen3_vl/qwen3vl.py +7 -23
  7. nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/qwen3vl_moe.py +8 -24
  8. {nexaai-1.0.19rc7.dist-info → nexaai-1.0.19rc9.dist-info}/METADATA +1 -1
  9. {nexaai-1.0.19rc7.dist-info → nexaai-1.0.19rc9.dist-info}/RECORD +11 -200
  10. nexaai/binds/nexa_mlx/py-lib/asr/__init__.py +0 -12
  11. nexaai/binds/nexa_mlx/py-lib/asr/interface.py +0 -122
  12. nexaai/binds/nexa_mlx/py-lib/common/__init__.py +0 -0
  13. nexaai/binds/nexa_mlx/py-lib/common/utils.py +0 -25
  14. nexaai/binds/nexa_mlx/py-lib/cv/__init__.py +0 -0
  15. nexaai/binds/nexa_mlx/py-lib/cv/generate.py +0 -195
  16. nexaai/binds/nexa_mlx/py-lib/cv/interface.py +0 -151
  17. nexaai/binds/nexa_mlx/py-lib/cv/main.py +0 -81
  18. nexaai/binds/nexa_mlx/py-lib/cv/modeling/pp_ocr_v4.py +0 -1736
  19. nexaai/binds/nexa_mlx/py-lib/embedding/__init__.py +0 -0
  20. nexaai/binds/nexa_mlx/py-lib/embedding/generate.py +0 -333
  21. nexaai/binds/nexa_mlx/py-lib/embedding/interface.py +0 -617
  22. nexaai/binds/nexa_mlx/py-lib/embedding/main.py +0 -173
  23. nexaai/binds/nexa_mlx/py-lib/embedding/modeling/__init__.py +0 -0
  24. nexaai/binds/nexa_mlx/py-lib/embedding/modeling/nexa_jina_v2.py +0 -399
  25. nexaai/binds/nexa_mlx/py-lib/image_gen/__init__.py +0 -1
  26. nexaai/binds/nexa_mlx/py-lib/image_gen/generate_sd.py +0 -244
  27. nexaai/binds/nexa_mlx/py-lib/image_gen/interface.py +0 -82
  28. nexaai/binds/nexa_mlx/py-lib/image_gen/main.py +0 -281
  29. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/__init__.py +0 -306
  30. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/clip.py +0 -116
  31. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/config.py +0 -65
  32. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/model_io.py +0 -386
  33. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/sampler.py +0 -105
  34. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/tokenizer.py +0 -100
  35. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/unet.py +0 -460
  36. nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/vae.py +0 -274
  37. nexaai/binds/nexa_mlx/py-lib/llm/__init__.py +0 -0
  38. nexaai/binds/nexa_mlx/py-lib/llm/generate.py +0 -149
  39. nexaai/binds/nexa_mlx/py-lib/llm/interface.py +0 -764
  40. nexaai/binds/nexa_mlx/py-lib/llm/main.py +0 -68
  41. nexaai/binds/nexa_mlx/py-lib/rerank/__init__.py +0 -0
  42. nexaai/binds/nexa_mlx/py-lib/rerank/generate.py +0 -174
  43. nexaai/binds/nexa_mlx/py-lib/rerank/interface.py +0 -287
  44. nexaai/binds/nexa_mlx/py-lib/rerank/main.py +0 -127
  45. nexaai/binds/nexa_mlx/py-lib/rerank/modeling/__init__.py +0 -0
  46. nexaai/binds/nexa_mlx/py-lib/rerank/modeling/nexa_jina_rerank.py +0 -330
  47. nexaai/binds/nexa_mlx/py-lib/sd/__init__.py +0 -1
  48. nexaai/binds/nexa_mlx/py-lib/sd/interface.py +0 -362
  49. nexaai/binds/nexa_mlx/py-lib/sd/main.py +0 -286
  50. nexaai/binds/nexa_mlx/py-lib/sd/modeling/__init__.py +0 -306
  51. nexaai/binds/nexa_mlx/py-lib/sd/modeling/clip.py +0 -116
  52. nexaai/binds/nexa_mlx/py-lib/sd/modeling/config.py +0 -65
  53. nexaai/binds/nexa_mlx/py-lib/sd/modeling/model_io.py +0 -385
  54. nexaai/binds/nexa_mlx/py-lib/sd/modeling/sampler.py +0 -105
  55. nexaai/binds/nexa_mlx/py-lib/sd/modeling/tokenizer.py +0 -100
  56. nexaai/binds/nexa_mlx/py-lib/sd/modeling/unet.py +0 -460
  57. nexaai/binds/nexa_mlx/py-lib/sd/modeling/vae.py +0 -274
  58. nexaai/binds/nexa_mlx/py-lib/tts/__init__.py +0 -12
  59. nexaai/binds/nexa_mlx/py-lib/tts/interface.py +0 -276
  60. nexaai/binds/nexa_mlx/py-lib/vlm/__init__.py +0 -3
  61. nexaai/binds/nexa_mlx/py-lib/vlm/generate.py +0 -572
  62. nexaai/binds/nexa_mlx/py-lib/vlm/generate_qwen3_vl.py +0 -294
  63. nexaai/binds/nexa_mlx/py-lib/vlm/generate_qwen3_vl_moe.py +0 -276
  64. nexaai/binds/nexa_mlx/py-lib/vlm/interface.py +0 -504
  65. nexaai/binds/nexa_mlx/py-lib/vlm/main.py +0 -320
  66. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/__init__.py +0 -0
  67. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/convert.py +0 -68
  68. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/__init__.py +0 -0
  69. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/__init__.py +0 -8
  70. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/aya_vision.py +0 -193
  71. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/interpolate.py +0 -186
  72. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/language.py +0 -233
  73. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/vision.py +0 -503
  74. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/base.py +0 -202
  75. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/cache.py +0 -230
  76. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/__init__.py +0 -10
  77. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/conversation.py +0 -264
  78. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/deepseek_vl_v2.py +0 -472
  79. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/language.py +0 -591
  80. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/processing_deepsek_vl_v2.py +0 -526
  81. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/vision.py +0 -356
  82. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/florence2/__init__.py +0 -8
  83. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/florence2/florence2.py +0 -366
  84. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/florence2/language.py +0 -488
  85. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/florence2/vision.py +0 -591
  86. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3/__init__.py +0 -8
  87. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3/gemma3.py +0 -213
  88. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3/language.py +0 -315
  89. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3/vision.py +0 -238
  90. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/__init__.py +0 -2
  91. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/audio.py +0 -1038
  92. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/config.py +0 -139
  93. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/gemma3n.py +0 -322
  94. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/language.py +0 -629
  95. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/vision.py +0 -1022
  96. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics2/__init__.py +0 -9
  97. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics2/idefics2.py +0 -294
  98. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics2/language.py +0 -191
  99. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics2/vision.py +0 -267
  100. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics3/__init__.py +0 -8
  101. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics3/idefics3.py +0 -175
  102. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics3/language.py +0 -192
  103. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics3/vision.py +0 -233
  104. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/__init__.py +0 -9
  105. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/internvl_chat.py +0 -140
  106. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/language.py +0 -220
  107. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/processor.py +0 -393
  108. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/vision.py +0 -293
  109. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kernels.py +0 -307
  110. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kimi_vl/__init__.py +0 -8
  111. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kimi_vl/kimi_vl.py +0 -143
  112. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kimi_vl/language.py +0 -509
  113. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kimi_vl/vision.py +0 -522
  114. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llama4/__init__.py +0 -8
  115. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llama4/language.py +0 -386
  116. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llama4/llama4.py +0 -138
  117. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llama4/vision.py +0 -560
  118. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava/__init__.py +0 -8
  119. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava/language.py +0 -240
  120. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava/llava.py +0 -153
  121. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava/vision.py +0 -259
  122. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_bunny/__init__.py +0 -9
  123. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_bunny/language.py +0 -236
  124. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_bunny/llava_bunny.py +0 -256
  125. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_bunny/vision.py +0 -303
  126. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_next/__init__.py +0 -8
  127. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_next/language.py +0 -230
  128. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_next/llava_next.py +0 -160
  129. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_next/vision.py +0 -243
  130. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mistral3/__init__.py +0 -8
  131. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mistral3/mistral3.py +0 -283
  132. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mllama/__init__.py +0 -8
  133. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mllama/language.py +0 -416
  134. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mllama/mllama.py +0 -172
  135. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mllama/vision.py +0 -499
  136. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/molmo/__init__.py +0 -8
  137. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/molmo/language.py +0 -243
  138. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/molmo/molmo.py +0 -133
  139. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/molmo/vision.py +0 -465
  140. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/__init__.py +0 -10
  141. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/language.py +0 -230
  142. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/multi_modality.py +0 -385
  143. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/sam.py +0 -557
  144. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/vision.py +0 -526
  145. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/paligemma/__init__.py +0 -8
  146. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/paligemma/language.py +0 -282
  147. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/paligemma/paligemma.py +0 -160
  148. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/paligemma/vision.py +0 -242
  149. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/__init__.py +0 -8
  150. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/language.py +0 -21
  151. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/phi3_v.py +0 -243
  152. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/su_rope.py +0 -71
  153. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/vision.py +0 -324
  154. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/pixtral/__init__.py +0 -8
  155. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/pixtral/language.py +0 -229
  156. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/pixtral/pixtral.py +0 -161
  157. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/pixtral/vision.py +0 -320
  158. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/__init__.py +0 -2
  159. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/config.py +0 -108
  160. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/language.py +0 -490
  161. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/qwen2_5_vl.py +0 -168
  162. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/vision.py +0 -414
  163. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/__init__.py +0 -2
  164. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/config.py +0 -104
  165. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/language.py +0 -490
  166. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/qwen2_vl.py +0 -167
  167. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/vision.py +0 -312
  168. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/__init__.py +0 -0
  169. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/base.py +0 -117
  170. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/cache.py +0 -531
  171. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/generate.py +0 -701
  172. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/rope_utils.py +0 -255
  173. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/sample_utils.py +0 -303
  174. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/tokenizer_utils.py +0 -407
  175. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/processor.py +0 -476
  176. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/qwen3vl.py +0 -1223
  177. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/__init__.py +0 -0
  178. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/base.py +0 -117
  179. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/cache.py +0 -531
  180. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/generate.py +0 -701
  181. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/rope_utils.py +0 -255
  182. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/sample_utils.py +0 -303
  183. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/tokenizer_utils.py +0 -407
  184. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/processor.py +0 -476
  185. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/qwen3vl_moe.py +0 -1309
  186. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/switch_layers.py +0 -210
  187. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/smolvlm/__init__.py +0 -8
  188. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/smolvlm/smolvlm.py +0 -62
  189. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/processing_qwen2_5_vl.py +0 -209
  190. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/processing_qwen2_vl.py +0 -215
  191. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/prompt_utils.py +0 -474
  192. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/sample_utils.py +0 -39
  193. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/tokenizer_utils.py +0 -344
  194. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/trainer/__init__.py +0 -9
  195. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/trainer/lora.py +0 -70
  196. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/trainer/trainer.py +0 -296
  197. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/trainer/utils.py +0 -160
  198. nexaai/binds/nexa_mlx/py-lib/vlm/modeling/utils.py +0 -928
  199. {nexaai-1.0.19rc7.dist-info → nexaai-1.0.19rc9.dist-info}/WHEEL +0 -0
  200. {nexaai-1.0.19rc7.dist-info → nexaai-1.0.19rc9.dist-info}/top_level.txt +0 -0
@@ -1,274 +0,0 @@
1
- # Copyright © 2023 Apple Inc.
2
-
3
- import math
4
- from typing import List
5
-
6
- import mlx.core as mx
7
- import mlx.nn as nn
8
-
9
- from .config import AutoencoderConfig
10
- from .unet import ResnetBlock2D, upsample_nearest
11
-
12
-
13
- class Attention(nn.Module):
14
- """A single head unmasked attention for use with the VAE."""
15
-
16
- def __init__(self, dims: int, norm_groups: int = 32):
17
- super().__init__()
18
-
19
- self.group_norm = nn.GroupNorm(norm_groups, dims, pytorch_compatible=True)
20
- self.query_proj = nn.Linear(dims, dims)
21
- self.key_proj = nn.Linear(dims, dims)
22
- self.value_proj = nn.Linear(dims, dims)
23
- self.out_proj = nn.Linear(dims, dims)
24
-
25
- def __call__(self, x):
26
- B, H, W, C = x.shape
27
-
28
- y = self.group_norm(x)
29
-
30
- queries = self.query_proj(y).reshape(B, H * W, C)
31
- keys = self.key_proj(y).reshape(B, H * W, C)
32
- values = self.value_proj(y).reshape(B, H * W, C)
33
-
34
- scale = 1 / math.sqrt(queries.shape[-1])
35
- scores = (queries * scale) @ keys.transpose(0, 2, 1)
36
- attn = mx.softmax(scores, axis=-1)
37
- y = (attn @ values).reshape(B, H, W, C)
38
-
39
- y = self.out_proj(y)
40
- x = x + y
41
-
42
- return x
43
-
44
-
45
- class EncoderDecoderBlock2D(nn.Module):
46
- def __init__(
47
- self,
48
- in_channels: int,
49
- out_channels: int,
50
- num_layers: int = 1,
51
- resnet_groups: int = 32,
52
- add_downsample=True,
53
- add_upsample=True,
54
- ):
55
- super().__init__()
56
-
57
- # Add the resnet blocks
58
- self.resnets = [
59
- ResnetBlock2D(
60
- in_channels=in_channels if i == 0 else out_channels,
61
- out_channels=out_channels,
62
- groups=resnet_groups,
63
- )
64
- for i in range(num_layers)
65
- ]
66
-
67
- # Add an optional downsampling layer
68
- if add_downsample:
69
- self.downsample = nn.Conv2d(
70
- out_channels, out_channels, kernel_size=3, stride=2, padding=0
71
- )
72
-
73
- # or upsampling layer
74
- if add_upsample:
75
- self.upsample = nn.Conv2d(
76
- out_channels, out_channels, kernel_size=3, stride=1, padding=1
77
- )
78
-
79
- def __call__(self, x):
80
- for resnet in self.resnets:
81
- x = resnet(x)
82
-
83
- if "downsample" in self:
84
- x = mx.pad(x, [(0, 0), (0, 1), (0, 1), (0, 0)])
85
- x = self.downsample(x)
86
-
87
- if "upsample" in self:
88
- x = self.upsample(upsample_nearest(x))
89
-
90
- return x
91
-
92
-
93
- class Encoder(nn.Module):
94
- """Implements the encoder side of the Autoencoder."""
95
-
96
- def __init__(
97
- self,
98
- in_channels: int,
99
- out_channels: int,
100
- block_out_channels: List[int] = [64],
101
- layers_per_block: int = 2,
102
- resnet_groups: int = 32,
103
- ):
104
- super().__init__()
105
-
106
- self.conv_in = nn.Conv2d(
107
- in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1
108
- )
109
-
110
- channels = [block_out_channels[0]] + list(block_out_channels)
111
- self.down_blocks = [
112
- EncoderDecoderBlock2D(
113
- in_channels,
114
- out_channels,
115
- num_layers=layers_per_block,
116
- resnet_groups=resnet_groups,
117
- add_downsample=i < len(block_out_channels) - 1,
118
- add_upsample=False,
119
- )
120
- for i, (in_channels, out_channels) in enumerate(zip(channels, channels[1:]))
121
- ]
122
-
123
- self.mid_blocks = [
124
- ResnetBlock2D(
125
- in_channels=block_out_channels[-1],
126
- out_channels=block_out_channels[-1],
127
- groups=resnet_groups,
128
- ),
129
- Attention(block_out_channels[-1], resnet_groups),
130
- ResnetBlock2D(
131
- in_channels=block_out_channels[-1],
132
- out_channels=block_out_channels[-1],
133
- groups=resnet_groups,
134
- ),
135
- ]
136
-
137
- self.conv_norm_out = nn.GroupNorm(
138
- resnet_groups, block_out_channels[-1], pytorch_compatible=True
139
- )
140
- self.conv_out = nn.Conv2d(block_out_channels[-1], out_channels, 3, padding=1)
141
-
142
- def __call__(self, x):
143
- x = self.conv_in(x)
144
-
145
- for l in self.down_blocks:
146
- x = l(x)
147
-
148
- x = self.mid_blocks[0](x)
149
- x = self.mid_blocks[1](x)
150
- x = self.mid_blocks[2](x)
151
-
152
- x = self.conv_norm_out(x)
153
- x = nn.silu(x)
154
- x = self.conv_out(x)
155
-
156
- return x
157
-
158
-
159
- class Decoder(nn.Module):
160
- """Implements the decoder side of the Autoencoder."""
161
-
162
- def __init__(
163
- self,
164
- in_channels: int,
165
- out_channels: int,
166
- block_out_channels: List[int] = [64],
167
- layers_per_block: int = 2,
168
- resnet_groups: int = 32,
169
- ):
170
- super().__init__()
171
-
172
- self.conv_in = nn.Conv2d(
173
- in_channels, block_out_channels[-1], kernel_size=3, stride=1, padding=1
174
- )
175
-
176
- self.mid_blocks = [
177
- ResnetBlock2D(
178
- in_channels=block_out_channels[-1],
179
- out_channels=block_out_channels[-1],
180
- groups=resnet_groups,
181
- ),
182
- Attention(block_out_channels[-1], resnet_groups),
183
- ResnetBlock2D(
184
- in_channels=block_out_channels[-1],
185
- out_channels=block_out_channels[-1],
186
- groups=resnet_groups,
187
- ),
188
- ]
189
-
190
- channels = list(reversed(block_out_channels))
191
- channels = [channels[0]] + channels
192
- self.up_blocks = [
193
- EncoderDecoderBlock2D(
194
- in_channels,
195
- out_channels,
196
- num_layers=layers_per_block,
197
- resnet_groups=resnet_groups,
198
- add_downsample=False,
199
- add_upsample=i < len(block_out_channels) - 1,
200
- )
201
- for i, (in_channels, out_channels) in enumerate(zip(channels, channels[1:]))
202
- ]
203
-
204
- self.conv_norm_out = nn.GroupNorm(
205
- resnet_groups, block_out_channels[0], pytorch_compatible=True
206
- )
207
- self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1)
208
-
209
- def __call__(self, x):
210
- x = self.conv_in(x)
211
-
212
- x = self.mid_blocks[0](x)
213
- x = self.mid_blocks[1](x)
214
- x = self.mid_blocks[2](x)
215
-
216
- for l in self.up_blocks:
217
- x = l(x)
218
-
219
- x = self.conv_norm_out(x)
220
- x = nn.silu(x)
221
- x = self.conv_out(x)
222
-
223
- return x
224
-
225
-
226
- class Autoencoder(nn.Module):
227
- """The autoencoder that allows us to perform diffusion in the latent space."""
228
-
229
- def __init__(self, config: AutoencoderConfig):
230
- super().__init__()
231
-
232
- self.latent_channels = config.latent_channels_in
233
- self.scaling_factor = config.scaling_factor
234
- self.encoder = Encoder(
235
- config.in_channels,
236
- config.latent_channels_out,
237
- config.block_out_channels,
238
- config.layers_per_block,
239
- resnet_groups=config.norm_num_groups,
240
- )
241
- self.decoder = Decoder(
242
- config.latent_channels_in,
243
- config.out_channels,
244
- config.block_out_channels,
245
- config.layers_per_block + 1,
246
- resnet_groups=config.norm_num_groups,
247
- )
248
-
249
- self.quant_proj = nn.Linear(
250
- config.latent_channels_out, config.latent_channels_out
251
- )
252
- self.post_quant_proj = nn.Linear(
253
- config.latent_channels_in, config.latent_channels_in
254
- )
255
-
256
- def decode(self, z):
257
- z = z / self.scaling_factor
258
- return self.decoder(self.post_quant_proj(z))
259
-
260
- def encode(self, x):
261
- x = self.encoder(x)
262
- x = self.quant_proj(x)
263
- mean, logvar = x.split(2, axis=-1)
264
- mean = mean * self.scaling_factor
265
- logvar = logvar + 2 * math.log(self.scaling_factor)
266
-
267
- return mean, logvar
268
-
269
- def __call__(self, x, key=None):
270
- mean, logvar = self.encode(x)
271
- z = mx.random.normal(mean.shape, key=key) * mx.exp(0.5 * logvar) + mean
272
- x_hat = self.decode(z)
273
-
274
- return dict(x_hat=x_hat, z=z, mean=mean, logvar=logvar)
@@ -1,12 +0,0 @@
1
- # patching the _resume method in phonemizer because logger.setLevel(logging.ERROR) doesn't work - the logger instance is created and stored in the package.
2
- try:
3
- from phonemizer.backend.espeak.words_mismatch import BaseWordsMismatch
4
-
5
- def silent_resume(self, nmismatch, nlines):
6
- """Silent version of _resume that suppresses warnings"""
7
- pass
8
-
9
- BaseWordsMismatch._resume = silent_resume
10
-
11
- except ImportError:
12
- pass
@@ -1,276 +0,0 @@
1
- from typing import Any, List, Optional, Sequence
2
- import argparse
3
- import sys
4
- import os
5
- import glob
6
- import tempfile
7
- import time
8
- import soundfile as sf
9
- import mlx.core as mx
10
- import numpy as np
11
-
12
- from ml import TTS, TTSConfig, TTSResult, TTSSamplerConfig, Path as MLPath
13
- from mlx_audio.tts.utils import load_model
14
-
15
- from profiling import ProfilingMixin, StopReason
16
-
17
- class MlxTts(TTS, ProfilingMixin):
18
- """MLX Audio implementation of TTS interface."""
19
-
20
- def __init__(
21
- self,
22
- model_path: MLPath,
23
- vocoder_path: MLPath,
24
- device: Optional[str] = None,
25
- ) -> None:
26
- ProfilingMixin.__init__(self)
27
-
28
- if os.path.isfile(model_path):
29
- model_path = os.path.dirname(model_path)
30
-
31
- # vocoder_path is not used in MLX TTS since the vocoder is integrated
32
- super().__init__(model_path, vocoder_path, device)
33
- self._sampler_config = TTSSamplerConfig()
34
- self.model = None
35
- self._model_loaded = False
36
-
37
- # Load model during initialization (matching C API behavior)
38
- self._load_model()
39
-
40
- def _load_model(self) -> bool:
41
- """Load the TTS model."""
42
- try:
43
- self.model = load_model(self.model_path)
44
- self._model_loaded = True
45
- return True
46
- except Exception as e:
47
- print(f"Failed to load TTS model: {e}")
48
- return False
49
-
50
- def destroy(self) -> None:
51
- """Destroy the model and free resources."""
52
- if self.model is not None:
53
- del self.model
54
- self.model = None
55
- mx.clear_cache()
56
- self._model_loaded = False
57
-
58
- def synthesize(
59
- self,
60
- text: str,
61
- config: Optional[TTSConfig] = None,
62
- output_path: Optional[MLPath] = None,
63
- clear_cache: bool = True,
64
- ) -> TTSResult:
65
- """Synthesize speech from text and save to filesystem."""
66
- # Ensure model is loaded
67
- if not self._model_loaded or self.model is None:
68
- raise RuntimeError("TTS model not loaded")
69
-
70
- # Start profiling
71
- self._start_profiling()
72
- self._prompt_start()
73
-
74
- try:
75
- # Use default config if not provided
76
- if config is None:
77
- config = TTSConfig()
78
-
79
- # Generate output path if not provided
80
- if output_path is None:
81
- timestamp = int(time.time() * 1000)
82
- output_path = os.path.join(tempfile.gettempdir(), f"tts_output_{timestamp}.wav")
83
-
84
- # Resolve voice path for Kokoro models
85
- voice = config.voice
86
- if voice and not voice.endswith(".pt") and not os.path.isabs(voice):
87
- # For relative voice names like "af_heart", construct full path
88
- voice_path = os.path.join(self.model_path, "voices", f"{voice}.pt")
89
- if os.path.exists(voice_path):
90
- voice = voice_path
91
-
92
- # End prompt processing, start decode
93
- self._prompt_end()
94
- self._decode_start()
95
-
96
- results = self.model.generate(
97
- text=text,
98
- voice=voice,
99
- speed=config.speed,
100
- temperature=self._sampler_config.temperature,
101
- seed=config.seed if config.seed != -1 else None,
102
- verbose=False,
103
- stream=False,
104
- join_audio=True,
105
- )
106
-
107
- # Get the results (should be a generator)
108
- audio_list = []
109
- sample_rate = None
110
- for result in results:
111
- audio_list.append(result.audio)
112
- sample_rate = result.sample_rate
113
-
114
- if not audio_list:
115
- raise RuntimeError("No audio generated")
116
-
117
- # Concatenate audio if multiple chunks
118
- if len(audio_list) > 1:
119
- audio = mx.concatenate(audio_list, axis=0)
120
- else:
121
- audio = audio_list[0]
122
-
123
- # Convert MLX array to numpy for saving
124
- if isinstance(audio, mx.array):
125
- audio_np = np.array(audio)
126
-
127
- else:
128
- audio_np = audio
129
-
130
- # Save audio to file
131
- sf.write(output_path, audio_np, sample_rate)
132
-
133
- if clear_cache:
134
- mx.clear_cache()
135
-
136
- # Calculate metadata
137
- channels = 1 if len(audio_np.shape) == 1 else audio_np.shape[1]
138
- num_samples = len(audio_np)
139
- duration_seconds = num_samples / sample_rate
140
-
141
- # End decode and profiling
142
- self._decode_end()
143
- self._set_stop_reason(StopReason.ML_STOP_REASON_COMPLETED)
144
- self._end_profiling()
145
-
146
- return TTSResult(
147
- audio_path=output_path,
148
- duration_seconds=duration_seconds,
149
- sample_rate=sample_rate,
150
- channels=channels,
151
- num_samples=num_samples
152
- )
153
- except Exception as e:
154
- # End profiling on error
155
- self._end_profiling()
156
- raise e
157
-
158
-
159
-
160
- def list_available_voices(self) -> List[str]:
161
- """List available voices."""
162
- # Common MLX TTS voice names - this could be enhanced to discover voices dynamically
163
- default_voices = [
164
- "af_heart", "af_bella", "af_nicole", "af_sarah", "af_sky", "af_sunshine",
165
- "am_adam", "am_michael", "am_mead", "an_nova", "an_michael",
166
- "bf_emma", "bf_isabella", "bm_george", "bm_lewis"
167
- ]
168
-
169
- # Try to discover voices from model directory if available
170
- if self.model_path and os.path.exists(self.model_path):
171
- discovered_voices = []
172
- voice_patterns = [
173
- "*.pt", # Voice files in model root
174
- "voices/*.pt", # Voice files in voices subdirectory
175
- ]
176
-
177
- for pattern in voice_patterns:
178
- voice_files = glob.glob(os.path.join(self.model_path, pattern))
179
- for voice_file in voice_files:
180
- voice_name = os.path.splitext(os.path.basename(voice_file))[0]
181
- discovered_voices.append(voice_name)
182
-
183
- if discovered_voices:
184
- return discovered_voices
185
-
186
- return default_voices
187
-
188
-
189
-
190
-
191
-
192
- def main():
193
- """Main function for command line text-to-speech synthesis."""
194
- parser = argparse.ArgumentParser(description="Synthesize speech using MLX TTS")
195
- parser.add_argument("model_path", help="Path to the TTS model")
196
- parser.add_argument("text", help="Text to synthesize")
197
- parser.add_argument("--voice", "-v", default="af_heart", help="Voice to use (default: af_heart)")
198
- parser.add_argument("--speed", "-s", type=float, default=1.0, help="Speech speed (default: 1.0)")
199
- parser.add_argument("--output", "-o", default="output.wav", help="Output audio file (default: output.wav)")
200
- parser.add_argument("--sample-rate", "-sr", type=int, default=24000, help="Sample rate (default: 24000)")
201
- parser.add_argument("--temperature", "-t", type=float, default=0.7, help="Temperature for sampling (default: 0.7)")
202
- parser.add_argument("--seed", type=int, default=-1, help="Random seed (-1 for random)")
203
- parser.add_argument("--list-voices", action="store_true", help="List available voices")
204
-
205
- args = parser.parse_args()
206
-
207
- # Check if model path exists
208
- if not os.path.exists(args.model_path):
209
- print(f"Error: Model path does not exist: {args.model_path}")
210
- sys.exit(1)
211
-
212
- # Initialize TTS adapter
213
- print(f"Initializing TTS with model: {args.model_path}")
214
- try:
215
- tts = MlxTts(
216
- model_path=args.model_path,
217
- vocoder_path="", # Not used in MLX TTS
218
- device=None
219
- )
220
-
221
- print("TTS model loaded successfully")
222
-
223
- # List voices if requested
224
- if args.list_voices:
225
- voices = tts.list_available_voices()
226
- print(f"Available voices: {', '.join(voices)}")
227
- return
228
-
229
- except Exception as e:
230
- print(f"Error initializing TTS: {e}")
231
- sys.exit(1)
232
-
233
- # Set up synthesis config
234
- sampler_config = TTSSamplerConfig(
235
- temperature=args.temperature,
236
- noise_scale=0.667,
237
- length_scale=1.0
238
- )
239
- tts._sampler_config = sampler_config
240
-
241
- config = TTSConfig(
242
- voice=args.voice,
243
- speed=args.speed,
244
- seed=args.seed,
245
- sample_rate=args.sample_rate
246
- )
247
-
248
- # Synthesize speech
249
- print(f"Synthesizing text: '{args.text}'")
250
- print(f"Using voice: {args.voice}")
251
- print(f"Speed: {args.speed}x")
252
- print("-" * 50)
253
-
254
- try:
255
- result = tts.synthesize(args.text, config, args.output)
256
-
257
- # Print results
258
- print("Synthesis Results:")
259
- print("=" * 50)
260
- print(f"Audio generated:")
261
- print(f" Duration: {result.duration_seconds:.2f} seconds")
262
- print(f" Sample rate: {result.sample_rate} Hz")
263
- print(f" Channels: {result.channels}")
264
- print(f" Samples: {result.num_samples}")
265
- print(f"✅ Audio saved to: {result.audio_path}")
266
-
267
- except Exception as e:
268
- print(f"Error during synthesis: {e}")
269
- sys.exit(1)
270
- finally:
271
- # Clean up
272
- tts.destroy()
273
-
274
-
275
- if __name__ == "__main__":
276
- main()
@@ -1,3 +0,0 @@
1
- import logging
2
-
3
- logging.getLogger("transformers").setLevel(logging.ERROR)