Support MLX in WhisperAX

[ZachNagengast]

Add support for the WhisperAX example app, as well as various refactors and cleanup.

Note: this is using an unreleased version of MLX, pending merge of https://github.com/ml-explore/mlx-swift/pull/130

There are still some memory issues to address, will look more into this soon.

CoreML:

MLX:

Important note, first time running this will likely throw an error about PrepareMetalShaders, which requires Trust & Enable on this popup when selecting the error.

[iandundas]

Hey all, just a heads-up, the example project gets this SPM error:

[davidkoski]

Hey all, just a heads-up, the example project gets this SPM error:

that fork was merged and is now the 0.16.2 tag on https://github.com/ml-explore/mlx-swift

[ZachNagengast]

Awesome, thanks for the update @davidkoski!

[latenitefilms]

@ZachNagengast - This looks awesome! Apologies - rookie question... It looks like currently the MLX repo only has the base and tiny models. How hard is it for mere mortals to "build" some of the larger models for testing this out with larger models?

[latenitefilms]

Or... can you use existing models with MLX?

[ZachNagengast]

Sorry missed your original message! We will fill in the mlx repo with the remaining models as part of this release, we just made these copies for consistency with our swift package. Any MLX whisper model currently existing with the same naming scheme will work in theory 👍 @jkrukowski may be able to confirm or deny.

[latenitefilms]

Legend, thanks @ZachNagengast! So basically, we do need new models for MLX, we can use the existing WhisperKit models? They need to be optimised or something?

[ZachNagengast]

Yep the existing WhisperKit models are optimized for CoreML, the ones in this repo we will fill out with the equivalent weights that are compatible with this MLX PR

[latenitefilms]

Sorry for all the rookie questions, but when you say "optimised for CoreML" - does this mean they ONLY work on CoreML, or can you use these CoreML models in MLX and they're just not as fast/accurate?

Apologies - this whole Whisper world is very new to me, so I very much appreciate all your wisdom and support!

[maxlund]

Yep the existing WhisperKit models are optimized for CoreML, the ones in this repo we will fill out with the equivalent weights that are compatible with this MLX PR

@ZachNagengast Is there any model conversion script we can run, or any other source we can use, in order to create/obtain more MLX compatible model versions?

[ZachNagengast]

@latenitefilms Yes the .mlmodelc models only work with CoreML at the moment. @maxlund There is a script made by @jkrukowski to do the conversion here https://github.com/argmaxinc/WhisperKit/pull/169, we'll integrate this into https://github.com/argmaxinc/whisperkittools in the future.

[latenitefilms]

Legend, thanks so much @ZachNagengast! Do you have a rough/ballpark ETA of when you're hoping to finish and merge in MLX support? No rush or pressure - just wondering if it's worth trying to convert our own models or not.

Let me know if there's anything I can do to help with MLX testing/release! Would love to see this in action ASAP!

Thanks for EVERYTHING you do! Appreciate it!

[ZachNagengast]

There are just a few optimizations to fix up to make it ready for release, specifically memory usage. Current issues are:

• MLX does not require a prewarm stage, so it should skip that. Currently its loading the model twice without freeing up the memory. Can also be solved by setting a cache limit or clearing the cache after load
• KV cache should use this instead of mlmultiarrays
• Sampling can be compiled for some easy speedups
• Attention should use SDPA instead of current logic

These are paraphrased from @davidkoski and @awni

Will be revisiting this after the upcoming release but feel free to test with this current branch if you see any other potential speedup besides these, all the interfaces should be the same in its final form, just faster and more memory efficient with these changes.

[latenitefilms]

Amazing! Thanks so much! Will test out and let you know if I break anything.