nexaai 1.0.4rc10__py3-none-macosx_11_0_arm64.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.
- nexaai/__init__.py +71 -0
- nexaai/_version.py +4 -0
- nexaai/asr.py +60 -0
- nexaai/asr_impl/__init__.py +0 -0
- nexaai/asr_impl/mlx_asr_impl.py +91 -0
- nexaai/asr_impl/pybind_asr_impl.py +43 -0
- nexaai/base.py +39 -0
- nexaai/binds/__init__.py +3 -0
- nexaai/binds/common_bind.cpython-310-darwin.so +0 -0
- nexaai/binds/embedder_bind.cpython-310-darwin.so +0 -0
- nexaai/binds/libnexa_bridge.dylib +0 -0
- nexaai/binds/llm_bind.cpython-310-darwin.so +0 -0
- nexaai/binds/nexa_llama_cpp/libggml-base.dylib +0 -0
- nexaai/binds/nexa_llama_cpp/libggml-cpu.so +0 -0
- nexaai/binds/nexa_llama_cpp/libggml-metal.so +0 -0
- nexaai/binds/nexa_llama_cpp/libggml.dylib +0 -0
- nexaai/binds/nexa_llama_cpp/libllama.dylib +0 -0
- nexaai/binds/nexa_llama_cpp/libmtmd.dylib +0 -0
- nexaai/binds/nexa_llama_cpp/libnexa_plugin.dylib +0 -0
- nexaai/binds/nexa_mlx/libnexa_plugin.dylib +0 -0
- nexaai/binds/nexa_mlx/py-lib/ml.py +842 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/__init__.py +5 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/bigvgan/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/bigvgan/activation.py +51 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/bigvgan/amp.py +96 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/bigvgan/bigvgan.py +149 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/bigvgan/conv.py +114 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/bigvgan/resample.py +177 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/descript/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/descript/base.py +228 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/descript/dac.py +285 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/descript/nn/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/descript/nn/layers.py +129 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/descript/nn/quantize.py +149 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/encodec/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/encodec/encodec.py +777 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/mimi/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/mimi/mimi.py +286 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/mimi/modules/__init__.py +20 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/mimi/modules/conv.py +398 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/mimi/modules/kv_cache.py +199 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/mimi/modules/quantization.py +179 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/mimi/modules/seanet.py +314 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/mimi/modules/transformer.py +256 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/s3/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/s3/model.py +260 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/s3/model_v2.py +383 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/s3/utils.py +122 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/snac/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/snac/attention.py +97 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/snac/layers.py +306 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/snac/snac.py +154 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/snac/vq.py +135 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/vocos/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/vocos/mel.py +33 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/models/vocos/vocos.py +359 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/tests/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/tests/test_bigvgan.py +54 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/tests/test_descript.py +109 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/tests/test_encodec.py +58 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/tests/test_mimi.py +22 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/tests/test_s3.py +25 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/tests/test_snac.py +40 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/codec/tests/test_vocos.py +93 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/server.py +525 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/sts/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/sts/tests/test_voice_pipeline.py +156 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/sts/voice_pipeline.py +327 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/generate.py +174 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/parakeet/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/parakeet/alignment.py +248 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/parakeet/attention.py +187 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/parakeet/audio.py +76 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/parakeet/conformer.py +331 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/parakeet/ctc.py +34 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/parakeet/parakeet.py +604 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/parakeet/rnnt.py +157 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/parakeet/tokenizer.py +2 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/wav2vec/feature_extractor.py +757 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/wav2vec/wav2vec.py +738 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/whisper/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/whisper/audio.py +82 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/whisper/decoding.py +742 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/whisper/timing.py +329 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/whisper/tokenizer.py +398 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/whisper/whisper.py +862 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/models/whisper/writers.py +268 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/tests/test_models.py +381 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/stt/utils.py +195 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/audio_player.py +120 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/convert.py +71 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/generate.py +449 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/bark/__init__.py +4 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/bark/bark.py +528 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/bark/isftnet.py +12 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/bark/pipeline.py +442 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/base.py +84 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/dia/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/dia/audio.py +287 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/dia/config.py +256 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/dia/dia.py +592 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/dia/layers.py +870 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/indextts/__init__.py +3 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/indextts/attention.py +180 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/indextts/bigvgan.py +124 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/indextts/conformer.py +247 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/indextts/ecapa_tdnn/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/indextts/ecapa_tdnn/asp.py +59 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/indextts/ecapa_tdnn/ecapa_tdnn.py +91 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/indextts/ecapa_tdnn/se_res2net.py +132 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/indextts/ecapa_tdnn/tdnn.py +42 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/indextts/gpt2.py +38 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/indextts/indextts.py +412 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/indextts/mel.py +37 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/indextts/normalize.py +294 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/indextts/perceiver.py +62 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/interpolate.py +108 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/kokoro/__init__.py +4 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/kokoro/istftnet.py +979 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/kokoro/kokoro.py +331 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/kokoro/modules.py +659 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/kokoro/pipeline.py +453 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/kokoro/voice.py +113 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/llama/__init__.py +3 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/llama/llama.py +324 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/outetts/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/outetts/audio_processor.py +351 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/outetts/dac_interface.py +162 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/outetts/outetts.py +255 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/outetts/prompt_processor.py +181 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/outetts/tokens.py +36 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/sesame/__init__.py +3 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/sesame/attention.py +195 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/sesame/sesame.py +633 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/sesame/watermarking.py +105 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/audio_tokenizer.py +138 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/bicodec.py +269 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/blocks/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/blocks/sampler.py +111 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/encoder_decoder/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/encoder_decoder/feat_decoder.py +120 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/encoder_decoder/feat_encoder.py +136 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/encoder_decoder/wave_generator.py +113 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/finite_scalar_quantization.py +238 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/residual.py +209 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/residual_fsq.py +309 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/speaker/__init__.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/speaker/ecapa_tdnn.py +283 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/speaker/perceiver_encoder.py +326 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/speaker/pooling_layers.py +297 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/modules/speaker/speaker_encoder.py +155 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/spark.py +382 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/utils/audio.py +220 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/utils/file.py +221 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/models/spark/utils/token_parser.py +181 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/tests/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/tests/test_base.py +66 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/tests/test_convert.py +173 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/tests/test_interpolate.py +88 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/tests/test_models.py +974 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/tts/utils.py +337 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/utils.py +237 -0
- nexaai/binds/nexa_mlx/py-lib/mlx_audio/version.py +1 -0
- nexaai/binds/nexa_mlx/py-lib/profiling.py +239 -0
- nexaai/common.py +61 -0
- nexaai/cv.py +87 -0
- nexaai/cv_impl/__init__.py +0 -0
- nexaai/cv_impl/mlx_cv_impl.py +88 -0
- nexaai/cv_impl/pybind_cv_impl.py +31 -0
- nexaai/embedder.py +68 -0
- nexaai/embedder_impl/__init__.py +0 -0
- nexaai/embedder_impl/mlx_embedder_impl.py +114 -0
- nexaai/embedder_impl/pybind_embedder_impl.py +91 -0
- nexaai/image_gen.py +136 -0
- nexaai/image_gen_impl/__init__.py +0 -0
- nexaai/image_gen_impl/mlx_image_gen_impl.py +291 -0
- nexaai/image_gen_impl/pybind_image_gen_impl.py +84 -0
- nexaai/llm.py +89 -0
- nexaai/llm_impl/__init__.py +0 -0
- nexaai/llm_impl/mlx_llm_impl.py +249 -0
- nexaai/llm_impl/pybind_llm_impl.py +207 -0
- nexaai/mlx_backend/asr/__init__.py +12 -0
- nexaai/mlx_backend/asr/interface.py +122 -0
- nexaai/mlx_backend/common/__init__.py +0 -0
- nexaai/mlx_backend/common/utils.py +25 -0
- nexaai/mlx_backend/cv/__init__.py +0 -0
- nexaai/mlx_backend/cv/generate.py +195 -0
- nexaai/mlx_backend/cv/interface.py +151 -0
- nexaai/mlx_backend/cv/main.py +81 -0
- nexaai/mlx_backend/cv/modeling/pp_ocr_v4.py +1736 -0
- nexaai/mlx_backend/embedding/__init__.py +0 -0
- nexaai/mlx_backend/embedding/generate.py +130 -0
- nexaai/mlx_backend/embedding/interface.py +312 -0
- nexaai/mlx_backend/embedding/main.py +82 -0
- nexaai/mlx_backend/embedding/modeling/__init__.py +0 -0
- nexaai/mlx_backend/embedding/modeling/nexa_jina_v2.py +399 -0
- nexaai/mlx_backend/llm/__init__.py +0 -0
- nexaai/mlx_backend/llm/generate.py +149 -0
- nexaai/mlx_backend/llm/interface.py +764 -0
- nexaai/mlx_backend/llm/main.py +68 -0
- nexaai/mlx_backend/ml.py +842 -0
- nexaai/mlx_backend/mlx_audio/__init__.py +0 -0
- nexaai/mlx_backend/mlx_audio/codec/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/codec/models/__init__.py +5 -0
- nexaai/mlx_backend/mlx_audio/codec/models/bigvgan/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/codec/models/bigvgan/activation.py +51 -0
- nexaai/mlx_backend/mlx_audio/codec/models/bigvgan/amp.py +96 -0
- nexaai/mlx_backend/mlx_audio/codec/models/bigvgan/bigvgan.py +149 -0
- nexaai/mlx_backend/mlx_audio/codec/models/bigvgan/conv.py +114 -0
- nexaai/mlx_backend/mlx_audio/codec/models/bigvgan/resample.py +177 -0
- nexaai/mlx_backend/mlx_audio/codec/models/descript/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/codec/models/descript/base.py +228 -0
- nexaai/mlx_backend/mlx_audio/codec/models/descript/dac.py +285 -0
- nexaai/mlx_backend/mlx_audio/codec/models/descript/nn/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/codec/models/descript/nn/layers.py +129 -0
- nexaai/mlx_backend/mlx_audio/codec/models/descript/nn/quantize.py +149 -0
- nexaai/mlx_backend/mlx_audio/codec/models/encodec/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/codec/models/encodec/encodec.py +777 -0
- nexaai/mlx_backend/mlx_audio/codec/models/mimi/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/codec/models/mimi/mimi.py +286 -0
- nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/__init__.py +20 -0
- nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/conv.py +398 -0
- nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/kv_cache.py +199 -0
- nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/quantization.py +179 -0
- nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/seanet.py +314 -0
- nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/transformer.py +256 -0
- nexaai/mlx_backend/mlx_audio/codec/models/s3/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/codec/models/s3/model.py +260 -0
- nexaai/mlx_backend/mlx_audio/codec/models/s3/model_v2.py +383 -0
- nexaai/mlx_backend/mlx_audio/codec/models/s3/utils.py +122 -0
- nexaai/mlx_backend/mlx_audio/codec/models/snac/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/codec/models/snac/attention.py +97 -0
- nexaai/mlx_backend/mlx_audio/codec/models/snac/layers.py +306 -0
- nexaai/mlx_backend/mlx_audio/codec/models/snac/snac.py +154 -0
- nexaai/mlx_backend/mlx_audio/codec/models/snac/vq.py +135 -0
- nexaai/mlx_backend/mlx_audio/codec/models/vocos/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/codec/models/vocos/mel.py +33 -0
- nexaai/mlx_backend/mlx_audio/codec/models/vocos/vocos.py +359 -0
- nexaai/mlx_backend/mlx_audio/codec/tests/__init__.py +0 -0
- nexaai/mlx_backend/mlx_audio/codec/tests/test_bigvgan.py +54 -0
- nexaai/mlx_backend/mlx_audio/codec/tests/test_descript.py +109 -0
- nexaai/mlx_backend/mlx_audio/codec/tests/test_encodec.py +58 -0
- nexaai/mlx_backend/mlx_audio/codec/tests/test_mimi.py +22 -0
- nexaai/mlx_backend/mlx_audio/codec/tests/test_s3.py +25 -0
- nexaai/mlx_backend/mlx_audio/codec/tests/test_snac.py +40 -0
- nexaai/mlx_backend/mlx_audio/codec/tests/test_vocos.py +93 -0
- nexaai/mlx_backend/mlx_audio/server.py +525 -0
- nexaai/mlx_backend/mlx_audio/sts/__init__.py +0 -0
- nexaai/mlx_backend/mlx_audio/sts/tests/test_voice_pipeline.py +156 -0
- nexaai/mlx_backend/mlx_audio/sts/voice_pipeline.py +327 -0
- nexaai/mlx_backend/mlx_audio/stt/__init__.py +0 -0
- nexaai/mlx_backend/mlx_audio/stt/generate.py +174 -0
- nexaai/mlx_backend/mlx_audio/stt/models/__init__.py +0 -0
- nexaai/mlx_backend/mlx_audio/stt/models/parakeet/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/stt/models/parakeet/alignment.py +248 -0
- nexaai/mlx_backend/mlx_audio/stt/models/parakeet/attention.py +187 -0
- nexaai/mlx_backend/mlx_audio/stt/models/parakeet/audio.py +76 -0
- nexaai/mlx_backend/mlx_audio/stt/models/parakeet/conformer.py +331 -0
- nexaai/mlx_backend/mlx_audio/stt/models/parakeet/ctc.py +34 -0
- nexaai/mlx_backend/mlx_audio/stt/models/parakeet/parakeet.py +604 -0
- nexaai/mlx_backend/mlx_audio/stt/models/parakeet/rnnt.py +157 -0
- nexaai/mlx_backend/mlx_audio/stt/models/parakeet/tokenizer.py +2 -0
- nexaai/mlx_backend/mlx_audio/stt/models/wav2vec/feature_extractor.py +757 -0
- nexaai/mlx_backend/mlx_audio/stt/models/wav2vec/wav2vec.py +738 -0
- nexaai/mlx_backend/mlx_audio/stt/models/whisper/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/stt/models/whisper/audio.py +82 -0
- nexaai/mlx_backend/mlx_audio/stt/models/whisper/decoding.py +742 -0
- nexaai/mlx_backend/mlx_audio/stt/models/whisper/timing.py +329 -0
- nexaai/mlx_backend/mlx_audio/stt/models/whisper/tokenizer.py +398 -0
- nexaai/mlx_backend/mlx_audio/stt/models/whisper/whisper.py +862 -0
- nexaai/mlx_backend/mlx_audio/stt/models/whisper/writers.py +268 -0
- nexaai/mlx_backend/mlx_audio/stt/tests/test_models.py +381 -0
- nexaai/mlx_backend/mlx_audio/stt/utils.py +195 -0
- nexaai/mlx_backend/mlx_audio/tts/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/tts/audio_player.py +120 -0
- nexaai/mlx_backend/mlx_audio/tts/convert.py +71 -0
- nexaai/mlx_backend/mlx_audio/tts/generate.py +449 -0
- nexaai/mlx_backend/mlx_audio/tts/models/__init__.py +0 -0
- nexaai/mlx_backend/mlx_audio/tts/models/bark/__init__.py +4 -0
- nexaai/mlx_backend/mlx_audio/tts/models/bark/bark.py +528 -0
- nexaai/mlx_backend/mlx_audio/tts/models/bark/isftnet.py +12 -0
- nexaai/mlx_backend/mlx_audio/tts/models/bark/pipeline.py +442 -0
- nexaai/mlx_backend/mlx_audio/tts/models/base.py +84 -0
- nexaai/mlx_backend/mlx_audio/tts/models/dia/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/tts/models/dia/audio.py +287 -0
- nexaai/mlx_backend/mlx_audio/tts/models/dia/config.py +256 -0
- nexaai/mlx_backend/mlx_audio/tts/models/dia/dia.py +592 -0
- nexaai/mlx_backend/mlx_audio/tts/models/dia/layers.py +870 -0
- nexaai/mlx_backend/mlx_audio/tts/models/indextts/__init__.py +3 -0
- nexaai/mlx_backend/mlx_audio/tts/models/indextts/attention.py +180 -0
- nexaai/mlx_backend/mlx_audio/tts/models/indextts/bigvgan.py +124 -0
- nexaai/mlx_backend/mlx_audio/tts/models/indextts/conformer.py +247 -0
- nexaai/mlx_backend/mlx_audio/tts/models/indextts/ecapa_tdnn/__init__.py +0 -0
- nexaai/mlx_backend/mlx_audio/tts/models/indextts/ecapa_tdnn/asp.py +59 -0
- nexaai/mlx_backend/mlx_audio/tts/models/indextts/ecapa_tdnn/ecapa_tdnn.py +91 -0
- nexaai/mlx_backend/mlx_audio/tts/models/indextts/ecapa_tdnn/se_res2net.py +132 -0
- nexaai/mlx_backend/mlx_audio/tts/models/indextts/ecapa_tdnn/tdnn.py +42 -0
- nexaai/mlx_backend/mlx_audio/tts/models/indextts/gpt2.py +38 -0
- nexaai/mlx_backend/mlx_audio/tts/models/indextts/indextts.py +412 -0
- nexaai/mlx_backend/mlx_audio/tts/models/indextts/mel.py +37 -0
- nexaai/mlx_backend/mlx_audio/tts/models/indextts/normalize.py +294 -0
- nexaai/mlx_backend/mlx_audio/tts/models/indextts/perceiver.py +62 -0
- nexaai/mlx_backend/mlx_audio/tts/models/interpolate.py +108 -0
- nexaai/mlx_backend/mlx_audio/tts/models/kokoro/__init__.py +4 -0
- nexaai/mlx_backend/mlx_audio/tts/models/kokoro/istftnet.py +979 -0
- nexaai/mlx_backend/mlx_audio/tts/models/kokoro/kokoro.py +331 -0
- nexaai/mlx_backend/mlx_audio/tts/models/kokoro/modules.py +659 -0
- nexaai/mlx_backend/mlx_audio/tts/models/kokoro/pipeline.py +453 -0
- nexaai/mlx_backend/mlx_audio/tts/models/kokoro/voice.py +113 -0
- nexaai/mlx_backend/mlx_audio/tts/models/llama/__init__.py +3 -0
- nexaai/mlx_backend/mlx_audio/tts/models/llama/llama.py +324 -0
- nexaai/mlx_backend/mlx_audio/tts/models/outetts/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/tts/models/outetts/audio_processor.py +351 -0
- nexaai/mlx_backend/mlx_audio/tts/models/outetts/dac_interface.py +162 -0
- nexaai/mlx_backend/mlx_audio/tts/models/outetts/default_speaker.json +461 -0
- nexaai/mlx_backend/mlx_audio/tts/models/outetts/outetts.py +255 -0
- nexaai/mlx_backend/mlx_audio/tts/models/outetts/prompt_processor.py +181 -0
- nexaai/mlx_backend/mlx_audio/tts/models/outetts/tokens.py +36 -0
- nexaai/mlx_backend/mlx_audio/tts/models/sesame/__init__.py +3 -0
- nexaai/mlx_backend/mlx_audio/tts/models/sesame/attention.py +195 -0
- nexaai/mlx_backend/mlx_audio/tts/models/sesame/sesame.py +633 -0
- nexaai/mlx_backend/mlx_audio/tts/models/sesame/watermarking.py +105 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/audio_tokenizer.py +138 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/bicodec.py +269 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/__init__.py +0 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/blocks/__init__.py +0 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/blocks/sampler.py +111 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/encoder_decoder/__init__.py +0 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/encoder_decoder/feat_decoder.py +120 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/encoder_decoder/feat_encoder.py +136 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/encoder_decoder/wave_generator.py +113 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/finite_scalar_quantization.py +238 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/residual.py +209 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/residual_fsq.py +309 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/speaker/__init__.py +1 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/speaker/ecapa_tdnn.py +283 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/speaker/perceiver_encoder.py +326 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/speaker/pooling_layers.py +297 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/speaker/speaker_encoder.py +155 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/spark.py +382 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/utils/audio.py +220 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/utils/file.py +221 -0
- nexaai/mlx_backend/mlx_audio/tts/models/spark/utils/token_parser.py +181 -0
- nexaai/mlx_backend/mlx_audio/tts/tests/__init__.py +0 -0
- nexaai/mlx_backend/mlx_audio/tts/tests/test_base.py +66 -0
- nexaai/mlx_backend/mlx_audio/tts/tests/test_convert.py +173 -0
- nexaai/mlx_backend/mlx_audio/tts/tests/test_interpolate.py +88 -0
- nexaai/mlx_backend/mlx_audio/tts/tests/test_models.py +974 -0
- nexaai/mlx_backend/mlx_audio/tts/utils.py +337 -0
- nexaai/mlx_backend/mlx_audio/utils.py +237 -0
- nexaai/mlx_backend/mlx_audio/version.py +1 -0
- nexaai/mlx_backend/profiling.py +239 -0
- nexaai/mlx_backend/rerank/__init__.py +0 -0
- nexaai/mlx_backend/rerank/generate.py +174 -0
- nexaai/mlx_backend/rerank/interface.py +287 -0
- nexaai/mlx_backend/rerank/main.py +127 -0
- nexaai/mlx_backend/rerank/modeling/__init__.py +0 -0
- nexaai/mlx_backend/rerank/modeling/nexa_jina_rerank.py +330 -0
- nexaai/mlx_backend/sd/__init__.py +1 -0
- nexaai/mlx_backend/sd/interface.py +362 -0
- nexaai/mlx_backend/sd/main.py +286 -0
- nexaai/mlx_backend/sd/modeling/__init__.py +306 -0
- nexaai/mlx_backend/sd/modeling/clip.py +116 -0
- nexaai/mlx_backend/sd/modeling/config.py +65 -0
- nexaai/mlx_backend/sd/modeling/model_io.py +330 -0
- nexaai/mlx_backend/sd/modeling/sampler.py +105 -0
- nexaai/mlx_backend/sd/modeling/tokenizer.py +100 -0
- nexaai/mlx_backend/sd/modeling/unet.py +460 -0
- nexaai/mlx_backend/sd/modeling/vae.py +274 -0
- nexaai/mlx_backend/tts/__init__.py +12 -0
- nexaai/mlx_backend/tts/interface.py +276 -0
- nexaai/mlx_backend/vlm/__init__.py +3 -0
- nexaai/mlx_backend/vlm/generate.py +572 -0
- nexaai/mlx_backend/vlm/interface.py +406 -0
- nexaai/mlx_backend/vlm/main.py +157 -0
- nexaai/mlx_backend/vlm/modeling/__init__.py +0 -0
- nexaai/mlx_backend/vlm/modeling/convert.py +68 -0
- nexaai/mlx_backend/vlm/modeling/models/__init__.py +0 -0
- nexaai/mlx_backend/vlm/modeling/models/aya_vision/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/aya_vision/aya_vision.py +193 -0
- nexaai/mlx_backend/vlm/modeling/models/aya_vision/interpolate.py +186 -0
- nexaai/mlx_backend/vlm/modeling/models/aya_vision/language.py +233 -0
- nexaai/mlx_backend/vlm/modeling/models/aya_vision/vision.py +503 -0
- nexaai/mlx_backend/vlm/modeling/models/base.py +202 -0
- nexaai/mlx_backend/vlm/modeling/models/cache.py +230 -0
- nexaai/mlx_backend/vlm/modeling/models/deepseek_vl_v2/__init__.py +10 -0
- nexaai/mlx_backend/vlm/modeling/models/deepseek_vl_v2/conversation.py +264 -0
- nexaai/mlx_backend/vlm/modeling/models/deepseek_vl_v2/deepseek_vl_v2.py +472 -0
- nexaai/mlx_backend/vlm/modeling/models/deepseek_vl_v2/language.py +591 -0
- nexaai/mlx_backend/vlm/modeling/models/deepseek_vl_v2/processing_deepsek_vl_v2.py +526 -0
- nexaai/mlx_backend/vlm/modeling/models/deepseek_vl_v2/vision.py +356 -0
- nexaai/mlx_backend/vlm/modeling/models/florence2/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/florence2/florence2.py +366 -0
- nexaai/mlx_backend/vlm/modeling/models/florence2/language.py +488 -0
- nexaai/mlx_backend/vlm/modeling/models/florence2/vision.py +591 -0
- nexaai/mlx_backend/vlm/modeling/models/gemma3/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/gemma3/gemma3.py +213 -0
- nexaai/mlx_backend/vlm/modeling/models/gemma3/language.py +315 -0
- nexaai/mlx_backend/vlm/modeling/models/gemma3/vision.py +238 -0
- nexaai/mlx_backend/vlm/modeling/models/gemma3n/__init__.py +2 -0
- nexaai/mlx_backend/vlm/modeling/models/gemma3n/audio.py +1038 -0
- nexaai/mlx_backend/vlm/modeling/models/gemma3n/config.py +139 -0
- nexaai/mlx_backend/vlm/modeling/models/gemma3n/gemma3n.py +322 -0
- nexaai/mlx_backend/vlm/modeling/models/gemma3n/language.py +629 -0
- nexaai/mlx_backend/vlm/modeling/models/gemma3n/vision.py +1022 -0
- nexaai/mlx_backend/vlm/modeling/models/idefics2/__init__.py +9 -0
- nexaai/mlx_backend/vlm/modeling/models/idefics2/idefics2.py +294 -0
- nexaai/mlx_backend/vlm/modeling/models/idefics2/language.py +191 -0
- nexaai/mlx_backend/vlm/modeling/models/idefics2/vision.py +267 -0
- nexaai/mlx_backend/vlm/modeling/models/idefics3/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/idefics3/idefics3.py +175 -0
- nexaai/mlx_backend/vlm/modeling/models/idefics3/language.py +192 -0
- nexaai/mlx_backend/vlm/modeling/models/idefics3/vision.py +233 -0
- nexaai/mlx_backend/vlm/modeling/models/internvl_chat/__init__.py +9 -0
- nexaai/mlx_backend/vlm/modeling/models/internvl_chat/internvl_chat.py +140 -0
- nexaai/mlx_backend/vlm/modeling/models/internvl_chat/language.py +220 -0
- nexaai/mlx_backend/vlm/modeling/models/internvl_chat/processor.py +393 -0
- nexaai/mlx_backend/vlm/modeling/models/internvl_chat/vision.py +293 -0
- nexaai/mlx_backend/vlm/modeling/models/kernels.py +307 -0
- nexaai/mlx_backend/vlm/modeling/models/kimi_vl/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/kimi_vl/kimi_vl.py +143 -0
- nexaai/mlx_backend/vlm/modeling/models/kimi_vl/language.py +509 -0
- nexaai/mlx_backend/vlm/modeling/models/kimi_vl/vision.py +522 -0
- nexaai/mlx_backend/vlm/modeling/models/llama4/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/llama4/language.py +386 -0
- nexaai/mlx_backend/vlm/modeling/models/llama4/llama4.py +138 -0
- nexaai/mlx_backend/vlm/modeling/models/llama4/vision.py +560 -0
- nexaai/mlx_backend/vlm/modeling/models/llava/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/llava/language.py +240 -0
- nexaai/mlx_backend/vlm/modeling/models/llava/llava.py +153 -0
- nexaai/mlx_backend/vlm/modeling/models/llava/vision.py +259 -0
- nexaai/mlx_backend/vlm/modeling/models/llava_bunny/__init__.py +9 -0
- nexaai/mlx_backend/vlm/modeling/models/llava_bunny/language.py +236 -0
- nexaai/mlx_backend/vlm/modeling/models/llava_bunny/llava_bunny.py +256 -0
- nexaai/mlx_backend/vlm/modeling/models/llava_bunny/vision.py +303 -0
- nexaai/mlx_backend/vlm/modeling/models/llava_next/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/llava_next/language.py +230 -0
- nexaai/mlx_backend/vlm/modeling/models/llava_next/llava_next.py +160 -0
- nexaai/mlx_backend/vlm/modeling/models/llava_next/vision.py +243 -0
- nexaai/mlx_backend/vlm/modeling/models/mistral3/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/mistral3/mistral3.py +283 -0
- nexaai/mlx_backend/vlm/modeling/models/mllama/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/mllama/language.py +416 -0
- nexaai/mlx_backend/vlm/modeling/models/mllama/mllama.py +172 -0
- nexaai/mlx_backend/vlm/modeling/models/mllama/vision.py +499 -0
- nexaai/mlx_backend/vlm/modeling/models/molmo/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/molmo/language.py +243 -0
- nexaai/mlx_backend/vlm/modeling/models/molmo/molmo.py +133 -0
- nexaai/mlx_backend/vlm/modeling/models/molmo/vision.py +465 -0
- nexaai/mlx_backend/vlm/modeling/models/multi_modality/__init__.py +10 -0
- nexaai/mlx_backend/vlm/modeling/models/multi_modality/language.py +230 -0
- nexaai/mlx_backend/vlm/modeling/models/multi_modality/multi_modality.py +385 -0
- nexaai/mlx_backend/vlm/modeling/models/multi_modality/sam.py +557 -0
- nexaai/mlx_backend/vlm/modeling/models/multi_modality/vision.py +526 -0
- nexaai/mlx_backend/vlm/modeling/models/paligemma/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/paligemma/language.py +282 -0
- nexaai/mlx_backend/vlm/modeling/models/paligemma/paligemma.py +160 -0
- nexaai/mlx_backend/vlm/modeling/models/paligemma/vision.py +242 -0
- nexaai/mlx_backend/vlm/modeling/models/phi3_v/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/phi3_v/language.py +21 -0
- nexaai/mlx_backend/vlm/modeling/models/phi3_v/phi3_v.py +243 -0
- nexaai/mlx_backend/vlm/modeling/models/phi3_v/su_rope.py +71 -0
- nexaai/mlx_backend/vlm/modeling/models/phi3_v/vision.py +324 -0
- nexaai/mlx_backend/vlm/modeling/models/pixtral/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/pixtral/language.py +229 -0
- nexaai/mlx_backend/vlm/modeling/models/pixtral/pixtral.py +161 -0
- nexaai/mlx_backend/vlm/modeling/models/pixtral/vision.py +320 -0
- nexaai/mlx_backend/vlm/modeling/models/qwen2_5_vl/__init__.py +2 -0
- nexaai/mlx_backend/vlm/modeling/models/qwen2_5_vl/config.py +108 -0
- nexaai/mlx_backend/vlm/modeling/models/qwen2_5_vl/language.py +490 -0
- nexaai/mlx_backend/vlm/modeling/models/qwen2_5_vl/qwen2_5_vl.py +168 -0
- nexaai/mlx_backend/vlm/modeling/models/qwen2_5_vl/vision.py +414 -0
- nexaai/mlx_backend/vlm/modeling/models/qwen2_vl/__init__.py +2 -0
- nexaai/mlx_backend/vlm/modeling/models/qwen2_vl/config.py +104 -0
- nexaai/mlx_backend/vlm/modeling/models/qwen2_vl/language.py +490 -0
- nexaai/mlx_backend/vlm/modeling/models/qwen2_vl/qwen2_vl.py +167 -0
- nexaai/mlx_backend/vlm/modeling/models/qwen2_vl/vision.py +312 -0
- nexaai/mlx_backend/vlm/modeling/models/smolvlm/__init__.py +8 -0
- nexaai/mlx_backend/vlm/modeling/models/smolvlm/smolvlm.py +62 -0
- nexaai/mlx_backend/vlm/modeling/processing_qwen2_5_vl.py +209 -0
- nexaai/mlx_backend/vlm/modeling/processing_qwen2_vl.py +215 -0
- nexaai/mlx_backend/vlm/modeling/prompt_utils.py +474 -0
- nexaai/mlx_backend/vlm/modeling/sample_utils.py +39 -0
- nexaai/mlx_backend/vlm/modeling/tokenizer_utils.py +344 -0
- nexaai/mlx_backend/vlm/modeling/trainer/__init__.py +9 -0
- nexaai/mlx_backend/vlm/modeling/trainer/lora.py +70 -0
- nexaai/mlx_backend/vlm/modeling/trainer/trainer.py +296 -0
- nexaai/mlx_backend/vlm/modeling/trainer/utils.py +160 -0
- nexaai/mlx_backend/vlm/modeling/utils.py +928 -0
- nexaai/rerank.py +51 -0
- nexaai/rerank_impl/__init__.py +0 -0
- nexaai/rerank_impl/mlx_rerank_impl.py +91 -0
- nexaai/rerank_impl/pybind_rerank_impl.py +42 -0
- nexaai/runtime.py +64 -0
- nexaai/tts.py +70 -0
- nexaai/tts_impl/__init__.py +0 -0
- nexaai/tts_impl/mlx_tts_impl.py +93 -0
- nexaai/tts_impl/pybind_tts_impl.py +42 -0
- nexaai/utils/avatar_fetcher.py +104 -0
- nexaai/utils/decode.py +18 -0
- nexaai/utils/model_manager.py +1195 -0
- nexaai/utils/progress_tracker.py +372 -0
- nexaai/vlm.py +120 -0
- nexaai/vlm_impl/__init__.py +0 -0
- nexaai/vlm_impl/mlx_vlm_impl.py +205 -0
- nexaai/vlm_impl/pybind_vlm_impl.py +228 -0
- nexaai-1.0.4rc10.dist-info/METADATA +26 -0
- nexaai-1.0.4rc10.dist-info/RECORD +519 -0
- nexaai-1.0.4rc10.dist-info/WHEEL +5 -0
- nexaai-1.0.4rc10.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,312 @@
|
|
|
1
|
+
import inspect
|
|
2
|
+
from dataclasses import dataclass
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
import mlx.core as mx
|
|
6
|
+
import mlx.nn as nn
|
|
7
|
+
|
|
8
|
+
from .config import VisionConfig
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def check_array_shape(arr):
|
|
12
|
+
shape = arr.shape
|
|
13
|
+
|
|
14
|
+
# Check if the shape has 4 dimensions
|
|
15
|
+
if len(shape) not in [4, 5]:
|
|
16
|
+
return False
|
|
17
|
+
|
|
18
|
+
B, out_channels, kH, KW, t = shape
|
|
19
|
+
|
|
20
|
+
if t == 3:
|
|
21
|
+
return True
|
|
22
|
+
|
|
23
|
+
# Check if out_channels is the largest, and kH and KW are the same
|
|
24
|
+
if (out_channels >= kH) and (out_channels >= KW) and (kH == KW):
|
|
25
|
+
return True
|
|
26
|
+
else:
|
|
27
|
+
return False
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def rotate_half(x):
|
|
31
|
+
"""Rotates half the hidden dims of the input."""
|
|
32
|
+
x1 = x[..., : x.shape[-1] // 2]
|
|
33
|
+
x2 = x[..., x.shape[-1] // 2 :]
|
|
34
|
+
return mx.concatenate([-x2, x1], axis=-1)
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def apply_rotary_pos_emb_vision(tensor, freqs) -> mx.array:
|
|
38
|
+
orig_dtype = tensor.dtype
|
|
39
|
+
|
|
40
|
+
cos = mx.cos(freqs)
|
|
41
|
+
sin = mx.sin(freqs)
|
|
42
|
+
|
|
43
|
+
cos = mx.expand_dims(cos, axis=1) # Equivalent to unsqueeze(1)
|
|
44
|
+
cos = mx.tile(cos, (1, 1, 2)) # Equivalent to repeat(1, 1, 2)
|
|
45
|
+
cos = mx.expand_dims(cos, axis=0) # Equivalent to [None, ...]
|
|
46
|
+
|
|
47
|
+
sin = mx.expand_dims(sin, axis=1) # Equivalent to unsqueeze(1)
|
|
48
|
+
sin = mx.tile(sin, (1, 1, 2)) # Equivalent to repeat(1, 1, 2)
|
|
49
|
+
sin = mx.expand_dims(sin, axis=0) # Equivalent to [None, ...]
|
|
50
|
+
|
|
51
|
+
output = (tensor * cos) + (rotate_half(tensor) * sin)
|
|
52
|
+
return output.astype(orig_dtype)
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
class VisionRotaryEmbedding(nn.Module):
|
|
56
|
+
def __init__(self, dim: int, theta: float = 10000.0) -> None:
|
|
57
|
+
super().__init__()
|
|
58
|
+
self.dim = dim
|
|
59
|
+
self.theta = theta
|
|
60
|
+
|
|
61
|
+
def __call__(self, seqlen: int) -> mx.array:
|
|
62
|
+
inv_freq = 1.0 / (
|
|
63
|
+
self.theta ** (mx.arange(0, self.dim, 2, dtype=mx.float32) / self.dim)
|
|
64
|
+
)
|
|
65
|
+
seq = mx.arange(seqlen.tolist(), dtype=inv_freq.dtype)
|
|
66
|
+
freqs = mx.outer(seq, inv_freq)
|
|
67
|
+
return freqs
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
class PatchEmbed(nn.Module):
|
|
71
|
+
def __init__(
|
|
72
|
+
self,
|
|
73
|
+
patch_size: int = 14,
|
|
74
|
+
temporal_patch_size: int = 2,
|
|
75
|
+
in_channels: int = 3,
|
|
76
|
+
embed_dim: int = 1152,
|
|
77
|
+
) -> None:
|
|
78
|
+
super().__init__()
|
|
79
|
+
self.patch_size = patch_size
|
|
80
|
+
self.temporal_patch_size = temporal_patch_size
|
|
81
|
+
self.in_channels = in_channels
|
|
82
|
+
self.embed_dim = embed_dim
|
|
83
|
+
|
|
84
|
+
kernel_size = [temporal_patch_size, patch_size, patch_size]
|
|
85
|
+
self.proj = nn.Conv3d(
|
|
86
|
+
in_channels,
|
|
87
|
+
embed_dim,
|
|
88
|
+
kernel_size=kernel_size,
|
|
89
|
+
stride=kernel_size,
|
|
90
|
+
bias=False,
|
|
91
|
+
)
|
|
92
|
+
|
|
93
|
+
def __call__(self, hidden_states: mx.array) -> mx.array:
|
|
94
|
+
hidden_states = hidden_states.reshape(
|
|
95
|
+
-1,
|
|
96
|
+
self.in_channels,
|
|
97
|
+
self.temporal_patch_size,
|
|
98
|
+
self.patch_size,
|
|
99
|
+
self.patch_size,
|
|
100
|
+
).moveaxis(1, 4)
|
|
101
|
+
|
|
102
|
+
hidden_states = self.proj(hidden_states)
|
|
103
|
+
hidden_states = hidden_states.reshape(-1, self.embed_dim)
|
|
104
|
+
return hidden_states
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
class PatchMerger(nn.Module):
|
|
108
|
+
def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None:
|
|
109
|
+
super().__init__()
|
|
110
|
+
self.hidden_size = context_dim * (spatial_merge_size**2)
|
|
111
|
+
self.ln_q = nn.LayerNorm(context_dim, eps=1e-6)
|
|
112
|
+
self.mlp = [
|
|
113
|
+
nn.Linear(self.hidden_size, self.hidden_size),
|
|
114
|
+
nn.GELU(),
|
|
115
|
+
nn.Linear(self.hidden_size, dim),
|
|
116
|
+
]
|
|
117
|
+
|
|
118
|
+
def __call__(self, x: mx.array) -> mx.array:
|
|
119
|
+
x = self.ln_q(x).reshape(-1, self.hidden_size)
|
|
120
|
+
for layer in self.mlp:
|
|
121
|
+
x = layer(x)
|
|
122
|
+
return x
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
class Attention(nn.Module):
|
|
126
|
+
def __init__(self, dim: int, num_heads: int = 16) -> None:
|
|
127
|
+
super().__init__()
|
|
128
|
+
self.num_heads = num_heads
|
|
129
|
+
self.head_dim = head_dim = dim // num_heads
|
|
130
|
+
self.scale = head_dim**-0.5
|
|
131
|
+
self.qkv = nn.Linear(dim, dim * 3, bias=True)
|
|
132
|
+
self.proj = nn.Linear(dim, dim)
|
|
133
|
+
|
|
134
|
+
def __call__(
|
|
135
|
+
self, x: mx.array, cu_seqlens: mx.array, rotary_pos_emb: mx.array = None
|
|
136
|
+
) -> mx.array:
|
|
137
|
+
seq_length = x.shape[0]
|
|
138
|
+
qkv = (
|
|
139
|
+
self.qkv(x).reshape(seq_length, 3, self.num_heads, -1).transpose(1, 0, 2, 3)
|
|
140
|
+
)
|
|
141
|
+
q, k, v = mx.split(qkv, 3)
|
|
142
|
+
|
|
143
|
+
q = apply_rotary_pos_emb_vision(mx.expand_dims(q, 0), rotary_pos_emb)[0]
|
|
144
|
+
k = apply_rotary_pos_emb_vision(mx.expand_dims(k, 0), rotary_pos_emb)[0]
|
|
145
|
+
attention_mask = mx.ones((1, seq_length, seq_length), dtype=x.dtype)
|
|
146
|
+
|
|
147
|
+
for i in range(1, len(cu_seqlens)):
|
|
148
|
+
start = int(cu_seqlens[i - 1])
|
|
149
|
+
end = int(cu_seqlens[i])
|
|
150
|
+
attention_mask[start:end, start:end] = 0
|
|
151
|
+
|
|
152
|
+
q = q.transpose(0, 2, 1, 3)
|
|
153
|
+
k = k.transpose(0, 2, 1, 3)
|
|
154
|
+
v = v.transpose(0, 2, 1, 3)
|
|
155
|
+
|
|
156
|
+
output = mx.fast.scaled_dot_product_attention(
|
|
157
|
+
q, k, v, scale=self.scale, mask=attention_mask
|
|
158
|
+
)
|
|
159
|
+
output = output.transpose(0, 2, 1, 3)
|
|
160
|
+
output = output.reshape(seq_length, -1)
|
|
161
|
+
return self.proj(output)
|
|
162
|
+
|
|
163
|
+
|
|
164
|
+
class MLP(nn.Module):
|
|
165
|
+
def __init__(self, dim, hidden_dim):
|
|
166
|
+
super().__init__()
|
|
167
|
+
self.activation_fn = nn.GELU(approx="fast")
|
|
168
|
+
self.fc1 = nn.Linear(dim, hidden_dim)
|
|
169
|
+
self.fc2 = nn.Linear(hidden_dim, dim)
|
|
170
|
+
|
|
171
|
+
def __call__(self, x: mx.array) -> mx.array:
|
|
172
|
+
x = self.activation_fn(self.fc1(x))
|
|
173
|
+
x = self.fc2(x)
|
|
174
|
+
return x
|
|
175
|
+
|
|
176
|
+
|
|
177
|
+
class Qwen2VLVisionBlock(nn.Module):
|
|
178
|
+
def __init__(self, config: VisionConfig) -> None:
|
|
179
|
+
super().__init__()
|
|
180
|
+
self.norm1 = nn.LayerNorm(config.embed_dim, eps=1e-6)
|
|
181
|
+
self.norm2 = nn.LayerNorm(config.embed_dim, eps=1e-6)
|
|
182
|
+
mlp_hidden_dim = int(config.embed_dim * config.mlp_ratio)
|
|
183
|
+
|
|
184
|
+
self.attn = Attention(dim=config.embed_dim, num_heads=config.num_heads)
|
|
185
|
+
self.mlp = MLP(dim=config.embed_dim, hidden_dim=mlp_hidden_dim)
|
|
186
|
+
|
|
187
|
+
def __call__(self, hidden_states, cu_seqlens, rotary_pos_emb) -> mx.array:
|
|
188
|
+
hidden_states = hidden_states + self.attn(
|
|
189
|
+
self.norm1(hidden_states),
|
|
190
|
+
cu_seqlens=cu_seqlens,
|
|
191
|
+
rotary_pos_emb=rotary_pos_emb,
|
|
192
|
+
)
|
|
193
|
+
hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
|
|
194
|
+
return hidden_states
|
|
195
|
+
|
|
196
|
+
|
|
197
|
+
class VisionModel(nn.Module):
|
|
198
|
+
|
|
199
|
+
def __init__(self, config: VisionConfig) -> None:
|
|
200
|
+
super().__init__()
|
|
201
|
+
self.config = config
|
|
202
|
+
self.model_type = config.model_type
|
|
203
|
+
if self.model_type != "qwen2_vl":
|
|
204
|
+
raise ValueError(f"Unsupported model type: {self.model_type}")
|
|
205
|
+
self.spatial_merge_size = config.spatial_merge_size
|
|
206
|
+
|
|
207
|
+
self.patch_embed = PatchEmbed(
|
|
208
|
+
patch_size=config.patch_size,
|
|
209
|
+
temporal_patch_size=config.temporal_patch_size,
|
|
210
|
+
in_channels=config.in_channels,
|
|
211
|
+
embed_dim=config.embed_dim,
|
|
212
|
+
)
|
|
213
|
+
|
|
214
|
+
head_dim = config.embed_dim // config.num_heads
|
|
215
|
+
self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)
|
|
216
|
+
|
|
217
|
+
self.blocks = [Qwen2VLVisionBlock(config) for _ in range(config.depth)]
|
|
218
|
+
self.merger = PatchMerger(dim=config.hidden_size, context_dim=config.embed_dim)
|
|
219
|
+
|
|
220
|
+
def rot_pos_emb(self, grid_thw):
|
|
221
|
+
pos_ids = []
|
|
222
|
+
|
|
223
|
+
for t, h, w in grid_thw:
|
|
224
|
+
h, w = int(h), int(w) # Ensure h and w are integers
|
|
225
|
+
hpos_ids = mx.expand_dims(mx.arange(h), 1)
|
|
226
|
+
hpos_ids = mx.repeat(hpos_ids, w, axis=1)
|
|
227
|
+
hpos_ids = hpos_ids.reshape(
|
|
228
|
+
h // self.spatial_merge_size,
|
|
229
|
+
self.spatial_merge_size,
|
|
230
|
+
w // self.spatial_merge_size,
|
|
231
|
+
self.spatial_merge_size,
|
|
232
|
+
)
|
|
233
|
+
hpos_ids = mx.transpose(hpos_ids, (0, 2, 1, 3))
|
|
234
|
+
hpos_ids = hpos_ids.flatten()
|
|
235
|
+
|
|
236
|
+
wpos_ids = mx.expand_dims(mx.arange(w), 0)
|
|
237
|
+
wpos_ids = mx.repeat(wpos_ids, h, axis=0)
|
|
238
|
+
wpos_ids = wpos_ids.reshape(
|
|
239
|
+
h // self.spatial_merge_size,
|
|
240
|
+
self.spatial_merge_size,
|
|
241
|
+
w // self.spatial_merge_size,
|
|
242
|
+
self.spatial_merge_size,
|
|
243
|
+
)
|
|
244
|
+
wpos_ids = mx.transpose(wpos_ids, (0, 2, 1, 3))
|
|
245
|
+
wpos_ids = wpos_ids.flatten()
|
|
246
|
+
|
|
247
|
+
stacked_pos_ids = mx.stack([hpos_ids, wpos_ids], axis=-1)
|
|
248
|
+
pos_ids.append(mx.tile(stacked_pos_ids, (t, 1)))
|
|
249
|
+
|
|
250
|
+
pos_ids = mx.concatenate(pos_ids, axis=0)
|
|
251
|
+
max_grid_size = mx.max(grid_thw[:, 1:])
|
|
252
|
+
rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
|
|
253
|
+
|
|
254
|
+
rotary_pos_emb_full = rotary_pos_emb_full[pos_ids]
|
|
255
|
+
|
|
256
|
+
return rotary_pos_emb_full.reshape(pos_ids.shape[0], -1)
|
|
257
|
+
|
|
258
|
+
def __call__(
|
|
259
|
+
self,
|
|
260
|
+
hidden_states: mx.array,
|
|
261
|
+
grid_thw: mx.array,
|
|
262
|
+
output_hidden_states: Optional[bool] = None,
|
|
263
|
+
) -> mx.array:
|
|
264
|
+
|
|
265
|
+
hidden_states = self.patch_embed(hidden_states)
|
|
266
|
+
rotary_pos_emb = self.rot_pos_emb(grid_thw)
|
|
267
|
+
|
|
268
|
+
# Assuming grid_thw has shape (batch_size, 3)
|
|
269
|
+
batch_size = grid_thw.shape[0]
|
|
270
|
+
|
|
271
|
+
# Calculate cu_seqlens for each item in the batch
|
|
272
|
+
cu_seqlens = []
|
|
273
|
+
for i in range(batch_size):
|
|
274
|
+
seq_len = grid_thw[i, 1] * grid_thw[i, 2]
|
|
275
|
+
cu_seqlens.append(mx.repeat(seq_len, grid_thw[i, 0]))
|
|
276
|
+
|
|
277
|
+
# Concatenate the cu_seqlens for all items in the batch
|
|
278
|
+
cu_seqlens = mx.concatenate(cu_seqlens)
|
|
279
|
+
|
|
280
|
+
cu_seqlens = mx.cumsum(cu_seqlens.astype(mx.int32), axis=0)
|
|
281
|
+
cu_seqlens = mx.pad(cu_seqlens, (1, 0), mode="constant", constant_values=0)
|
|
282
|
+
|
|
283
|
+
encoder_states = (hidden_states,) if output_hidden_states else None
|
|
284
|
+
|
|
285
|
+
for blk in self.blocks:
|
|
286
|
+
hidden_states = blk(
|
|
287
|
+
hidden_states, cu_seqlens=cu_seqlens, rotary_pos_emb=rotary_pos_emb
|
|
288
|
+
)
|
|
289
|
+
if output_hidden_states:
|
|
290
|
+
encoder_states = encoder_states + (hidden_states,)
|
|
291
|
+
|
|
292
|
+
return self.merger(hidden_states)
|
|
293
|
+
|
|
294
|
+
def sanitize(self, weights):
|
|
295
|
+
sanitized_weights = {}
|
|
296
|
+
for k, v in weights.items():
|
|
297
|
+
if "position_ids" in k:
|
|
298
|
+
# Remove unused position_ids
|
|
299
|
+
continue
|
|
300
|
+
elif "patch_embed.proj.weight" in k:
|
|
301
|
+
# PyTorch conv2d weight tensors have shape:
|
|
302
|
+
# [out_channels, in_channels, kH, KW]
|
|
303
|
+
# MLX conv2d expects the weight be of shape:
|
|
304
|
+
# [out_channels, kH, KW, in_channels]
|
|
305
|
+
if check_array_shape(v):
|
|
306
|
+
sanitized_weights[k] = v
|
|
307
|
+
else:
|
|
308
|
+
sanitized_weights[k] = v.transpose(0, 2, 3, 4, 1)
|
|
309
|
+
else:
|
|
310
|
+
sanitized_weights[k] = v
|
|
311
|
+
|
|
312
|
+
return sanitized_weights
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
import mlx.core as mx
|
|
2
|
+
import numpy as np
|
|
3
|
+
|
|
4
|
+
from ..idefics3 import LanguageModel
|
|
5
|
+
from ..idefics3 import Model as Idefics3Model
|
|
6
|
+
from ..idefics3 import ModelConfig, TextConfig, VisionConfig, VisionModel
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class Model(Idefics3Model):
|
|
10
|
+
def _prepare_inputs_for_multimodal(self, image_features, inputs_embeds, input_ids):
|
|
11
|
+
# Assumes bs == 1
|
|
12
|
+
|
|
13
|
+
B, T, D_text = inputs_embeds.shape
|
|
14
|
+
N, S, D_img = image_features.shape
|
|
15
|
+
|
|
16
|
+
image_offset = 0
|
|
17
|
+
cur_embeds = inputs_embeds[0]
|
|
18
|
+
|
|
19
|
+
# Find positions of <image> tokens in the text
|
|
20
|
+
image_token_index = self.config.image_token_index
|
|
21
|
+
image_positions = np.where(input_ids == image_token_index)[1].tolist()
|
|
22
|
+
num_image_tokens = len(image_positions)
|
|
23
|
+
|
|
24
|
+
# If no <image> => text-only
|
|
25
|
+
if num_image_tokens == 0:
|
|
26
|
+
empty_slice = image_features[0][:0, :] # shape (0, D)
|
|
27
|
+
return mx.concatenate([cur_embeds, empty_slice], axis=0)
|
|
28
|
+
|
|
29
|
+
# Typically, if each image is S embeddings, we expect the total # of <image> tokens
|
|
30
|
+
# in this sample to be multiple of S => each group of S tokens = 1 image
|
|
31
|
+
if num_image_tokens % S != 0:
|
|
32
|
+
raise ValueError(
|
|
33
|
+
f"Input has {num_image_tokens} <image> tokens, not a multiple of S={S}. "
|
|
34
|
+
"Cannot map them to blocks of shape (S, D)."
|
|
35
|
+
)
|
|
36
|
+
|
|
37
|
+
chunks = [image_positions[i : i + S] for i in range(0, num_image_tokens, S)]
|
|
38
|
+
|
|
39
|
+
segments = []
|
|
40
|
+
text_start = 0
|
|
41
|
+
|
|
42
|
+
# For each chunk (each chunk => 1 image)
|
|
43
|
+
for chunk in chunks:
|
|
44
|
+
cur_block = image_features[image_offset]
|
|
45
|
+
image_offset += 1
|
|
46
|
+
|
|
47
|
+
# We'll iterate over the S positions in ascending order
|
|
48
|
+
for i_s, pos in enumerate(chunk):
|
|
49
|
+
if pos > text_start:
|
|
50
|
+
segments.append(cur_embeds[text_start:pos])
|
|
51
|
+
# Then add one row from cur_block => shape (1, D)
|
|
52
|
+
row_of_block = cur_block[i_s : i_s + 1, :]
|
|
53
|
+
segments.append(row_of_block)
|
|
54
|
+
text_start = pos + 1
|
|
55
|
+
|
|
56
|
+
# leftover text after the final <image> token
|
|
57
|
+
if text_start < T:
|
|
58
|
+
segments.append(cur_embeds[text_start:])
|
|
59
|
+
|
|
60
|
+
# cat them into a single (T_b, D) tensor
|
|
61
|
+
merged_sample = mx.concatenate(segments, axis=0)
|
|
62
|
+
return mx.expand_dims(merged_sample, axis=0)
|
|
@@ -0,0 +1,209 @@
|
|
|
1
|
+
# Copied from transformers. Removed video-related code.
|
|
2
|
+
|
|
3
|
+
from typing import Optional, Union
|
|
4
|
+
|
|
5
|
+
import numpy as np
|
|
6
|
+
|
|
7
|
+
from transformers.feature_extraction_utils import BatchFeature
|
|
8
|
+
from transformers.image_utils import ImageInput
|
|
9
|
+
from transformers.processing_utils import ImagesKwargs, MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack, VideosKwargs
|
|
10
|
+
from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class Qwen2_5_VLImagesKwargs(ImagesKwargs):
|
|
14
|
+
min_pixels: Optional[int]
|
|
15
|
+
max_pixels: Optional[int]
|
|
16
|
+
patch_size: Optional[int]
|
|
17
|
+
temporal_patch_size: Optional[int]
|
|
18
|
+
merge_size: Optional[int]
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class Qwen2_5_VLProcessorKwargs(ProcessingKwargs, total=False):
|
|
22
|
+
images_kwargs: Qwen2_5_VLImagesKwargs
|
|
23
|
+
_defaults = {
|
|
24
|
+
"text_kwargs": {
|
|
25
|
+
"padding": False,
|
|
26
|
+
"return_mm_token_type_ids": False,
|
|
27
|
+
},
|
|
28
|
+
}
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
class Qwen2_5_VLProcessor(ProcessorMixin):
|
|
32
|
+
r"""
|
|
33
|
+
Constructs a Qwen2.5-VL processor which wraps a Qwen2.5-VL image processor and a Qwen2 tokenizer into a single processor.
|
|
34
|
+
[`Qwen2_5_VLProcessor`] offers all the functionalities of [`Qwen2VLImageProcessor`] and [`Qwen2TokenizerFast`]. See the
|
|
35
|
+
[`~Qwen2_5_VLProcessor.__call__`] and [`~Qwen2_5_VLProcessor.decode`] for more information.
|
|
36
|
+
Args:
|
|
37
|
+
image_processor ([`Qwen2VLImageProcessor`], *optional*):
|
|
38
|
+
The image processor is a required input.
|
|
39
|
+
tokenizer ([`Qwen2TokenizerFast`], *optional*):
|
|
40
|
+
The tokenizer is a required input.
|
|
41
|
+
chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
|
|
42
|
+
in a chat into a tokenizable string.
|
|
43
|
+
"""
|
|
44
|
+
|
|
45
|
+
attributes = ["image_processor", "tokenizer"]
|
|
46
|
+
|
|
47
|
+
image_processor_class = "AutoImageProcessor"
|
|
48
|
+
tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
|
|
49
|
+
|
|
50
|
+
def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs):
|
|
51
|
+
self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
|
|
52
|
+
self.image_token_id = (
|
|
53
|
+
tokenizer.image_token_id
|
|
54
|
+
if getattr(tokenizer, "image_token_id", None)
|
|
55
|
+
else tokenizer.convert_tokens_to_ids(self.image_token)
|
|
56
|
+
)
|
|
57
|
+
super().__init__(image_processor, tokenizer, chat_template=chat_template)
|
|
58
|
+
|
|
59
|
+
def __call__(
|
|
60
|
+
self,
|
|
61
|
+
images: ImageInput = None,
|
|
62
|
+
text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,
|
|
63
|
+
**kwargs: Unpack[Qwen2_5_VLProcessorKwargs],
|
|
64
|
+
) -> BatchFeature:
|
|
65
|
+
"""
|
|
66
|
+
Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
|
|
67
|
+
and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode
|
|
68
|
+
the text. To prepare the vision inputs, this method forwards the `vision_infos` and `kwrags` arguments to
|
|
69
|
+
Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`] if `vision_infos` is not `None`.
|
|
70
|
+
|
|
71
|
+
Args:
|
|
72
|
+
images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):
|
|
73
|
+
The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
|
|
74
|
+
tensor. Both channels-first and channels-last formats are supported.
|
|
75
|
+
text (`str`, `list[str]`, `list[list[str]]`):
|
|
76
|
+
The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
|
|
77
|
+
(pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
|
|
78
|
+
`is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
|
|
79
|
+
return_tensors (`str` or [`~utils.TensorType`], *optional*):
|
|
80
|
+
If set, will return tensors of a particular framework. Acceptable values are:
|
|
81
|
+
- `'tf'`: Return TensorFlow `tf.constant` objects.
|
|
82
|
+
- `'pt'`: Return PyTorch `torch.Tensor` objects.
|
|
83
|
+
- `'np'`: Return NumPy `np.ndarray` objects.
|
|
84
|
+
- `'jax'`: Return JAX `jnp.ndarray` objects.
|
|
85
|
+
|
|
86
|
+
Returns:
|
|
87
|
+
[`BatchFeature`]: A [`BatchFeature`] with the following fields:
|
|
88
|
+
|
|
89
|
+
- **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
|
|
90
|
+
- **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
|
|
91
|
+
`return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
|
|
92
|
+
`None`).
|
|
93
|
+
- **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
|
|
94
|
+
- **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`.
|
|
95
|
+
"""
|
|
96
|
+
output_kwargs = self._merge_kwargs(
|
|
97
|
+
Qwen2_5_VLProcessorKwargs,
|
|
98
|
+
tokenizer_init_kwargs=self.tokenizer.init_kwargs,
|
|
99
|
+
**kwargs,
|
|
100
|
+
)
|
|
101
|
+
|
|
102
|
+
image_inputs = {}
|
|
103
|
+
if images is not None:
|
|
104
|
+
image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
|
|
105
|
+
image_grid_thw = image_inputs["image_grid_thw"]
|
|
106
|
+
|
|
107
|
+
if not isinstance(text, list):
|
|
108
|
+
text = [text]
|
|
109
|
+
|
|
110
|
+
text = text.copy() # below lines change text in-place
|
|
111
|
+
if images is not None:
|
|
112
|
+
merge_length = self.image_processor.merge_size**2
|
|
113
|
+
index = 0
|
|
114
|
+
for i in range(len(text)):
|
|
115
|
+
while self.image_token in text[i]:
|
|
116
|
+
num_image_tokens = image_grid_thw[index].prod() // merge_length
|
|
117
|
+
text[i] = text[i].replace(self.image_token, "<|placeholder|>" * num_image_tokens, 1)
|
|
118
|
+
index += 1
|
|
119
|
+
text[i] = text[i].replace("<|placeholder|>", self.image_token)
|
|
120
|
+
|
|
121
|
+
return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
|
|
122
|
+
return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", None)
|
|
123
|
+
text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
|
|
124
|
+
self._check_special_mm_tokens(text, text_inputs, modalities=["image"])
|
|
125
|
+
|
|
126
|
+
if return_mm_token_type_ids:
|
|
127
|
+
array_ids = np.array(text_inputs["input_ids"])
|
|
128
|
+
mm_token_type_ids = np.zeros_like(text_inputs["input_ids"])
|
|
129
|
+
mm_token_type_ids[array_ids == self.image_token_id] = 1
|
|
130
|
+
text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist()
|
|
131
|
+
|
|
132
|
+
return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors)
|
|
133
|
+
|
|
134
|
+
def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs):
|
|
135
|
+
"""
|
|
136
|
+
Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.
|
|
137
|
+
Args:
|
|
138
|
+
image_sizes (`list[list[int]]`, *optional*):
|
|
139
|
+
The input sizes formatted as (height, width) per each image.
|
|
140
|
+
Returns:
|
|
141
|
+
`MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided
|
|
142
|
+
input modalities, along with other useful data.
|
|
143
|
+
"""
|
|
144
|
+
|
|
145
|
+
vision_data = {}
|
|
146
|
+
if image_sizes is not None:
|
|
147
|
+
images_kwargs = Qwen2_5_VLProcessorKwargs._defaults.get("images_kwargs", {})
|
|
148
|
+
images_kwargs.update(kwargs)
|
|
149
|
+
merge_size = images_kwargs.get("merge_size", None) or self.image_processor.merge_size
|
|
150
|
+
|
|
151
|
+
num_image_patches = [
|
|
152
|
+
self.image_processor.get_number_of_image_patches(*image_size, images_kwargs)
|
|
153
|
+
for image_size in image_sizes
|
|
154
|
+
]
|
|
155
|
+
num_image_tokens = [(num_patches // merge_size**2) for num_patches in num_image_patches]
|
|
156
|
+
vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})
|
|
157
|
+
|
|
158
|
+
return MultiModalData(**vision_data)
|
|
159
|
+
|
|
160
|
+
def batch_decode(self, *args, **kwargs):
|
|
161
|
+
"""
|
|
162
|
+
This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
|
|
163
|
+
refer to the docstring of this method for more information.
|
|
164
|
+
"""
|
|
165
|
+
return self.tokenizer.batch_decode(*args, **kwargs)
|
|
166
|
+
|
|
167
|
+
def decode(self, *args, **kwargs):
|
|
168
|
+
"""
|
|
169
|
+
This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
|
|
170
|
+
the docstring of this method for more information.
|
|
171
|
+
"""
|
|
172
|
+
return self.tokenizer.decode(*args, **kwargs)
|
|
173
|
+
|
|
174
|
+
def post_process_image_text_to_text(
|
|
175
|
+
self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs
|
|
176
|
+
):
|
|
177
|
+
"""
|
|
178
|
+
Post-process the output of the model to decode the text.
|
|
179
|
+
|
|
180
|
+
Args:
|
|
181
|
+
generated_outputs (`torch.Tensor` or `np.ndarray`):
|
|
182
|
+
The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
|
|
183
|
+
or `(sequence_length,)`.
|
|
184
|
+
skip_special_tokens (`bool`, *optional*, defaults to `True`):
|
|
185
|
+
Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.
|
|
186
|
+
clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
|
|
187
|
+
Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method.
|
|
188
|
+
**kwargs:
|
|
189
|
+
Additional arguments to be passed to the tokenizer's `batch_decode method`.
|
|
190
|
+
|
|
191
|
+
Returns:
|
|
192
|
+
`list[str]`: The decoded text.
|
|
193
|
+
"""
|
|
194
|
+
return self.tokenizer.batch_decode(
|
|
195
|
+
generated_outputs,
|
|
196
|
+
skip_special_tokens=skip_special_tokens,
|
|
197
|
+
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
|
|
198
|
+
**kwargs,
|
|
199
|
+
)
|
|
200
|
+
|
|
201
|
+
@property
|
|
202
|
+
def model_input_names(self):
|
|
203
|
+
tokenizer_input_names = self.tokenizer.model_input_names
|
|
204
|
+
image_processor_input_names = self.image_processor.model_input_names
|
|
205
|
+
names_from_processor = list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
|
|
206
|
+
return names_from_processor
|
|
207
|
+
|
|
208
|
+
|
|
209
|
+
__all__ = ["Qwen2_5_VLProcessor"]
|