MLX model support

.github/workflows/development-tests.yml

@@ -2,7 +2,6 @@ name: Development Tests
 
 on:
   pull_request:
-    branches: ["main"]
   pull_request_review:
     types: [submitted]
   workflow_dispatch:

.github/workflows/unit-tests.yml

@@ -20,26 +20,34 @@ jobs:
           - {
               name: "macOS",
               condition: true,
-              clean-destination: "generic/platform=macOS",
               test-destination: "platform=macOS,arch=arm64",
+              test-cases: "-only-testing WhisperKitTests/UnitTests -only-testing WhisperKitMLXTests/MLXUnitTests",
+              mlx-disabled: "0",
+              scheme: "whisperkit-Package",
             }
           - {
               name: "iOS",
               condition: true,
-              clean-destination: "generic/platform=iOS",
               test-destination: "platform=iOS Simulator,OS=${{ inputs.ios-version }},name=iPhone 15",
+              test-cases: "-only-testing WhisperKitTests/UnitTests",
+              mlx-disabled: "1",
+              scheme: "whisperkit",
             }
           - {
               name: "watchOS",
               condition: "${{ inputs.macos-runner == 'macos-14' }}",
-              clean-destination: "generic/platform=watchOS",
-              test-destination: "platform=watchOS Simulator,OS=10.2,name=Apple Watch Ultra 2 (49mm)",
+              test-destination: "platform=watchOS Simulator,OS=10.5,name=Apple Watch Ultra 2 (49mm)",
+              test-cases: "-only-testing WhisperKitTests/UnitTests",
+              mlx-disabled: "1",
+              scheme: "whisperkit",
             }
           - {
               name: "visionOS",
               condition: "${{ inputs.macos-runner == 'macos-14' }}",
-              clean-destination: "generic/platform=visionOS",
               test-destination: "platform=visionOS Simulator,name=Apple Vision Pro",
+              test-cases: "-only-testing WhisperKitTests/UnitTests",
+              mlx-disabled: "1",
+              scheme: "whisperkit",
             }
     timeout-minutes: 30
     steps:
@@ -57,12 +65,18 @@ jobs:
           key: ${{ runner.os }}-models
       - name: Download Models
         if: steps.model-cache.outputs.cache-hit != 'true'
-        run: make download-model MODEL=tiny
+        run: |
+          make download-model MODEL=tiny 
+          make download-mlx-model MODEL=tiny
       - name: Install and discover destinations
+        env:
+          MLX_DISABLED: ${{ matrix.run-config['mlx-disabled'] }}
         run: |
+          echo "Available schemes:"
+          xcodebuild -list
           xcodebuild -downloadAllPlatforms
           echo "Destinations for testing:"
-          xcodebuild test-without-building -only-testing WhisperKitTests/UnitTests -scheme whisperkit-Package -showdestinations
+          export ${{ matrix.run-config['compiler-flags'] }} && xcodebuild test-without-building -only-testing WhisperKitTests/UnitTests -scheme ${{ matrix.run-config['scheme'] }} -showdestinations -skipPackagePluginValidation
       - name: Boot Simulator and Wait
         if: ${{ matrix.run-config['name'] != 'macOS' }} && ${{ inputs.macos-runner == 'macos-14' }}
         # Slower runners require some time to fully boot the simulator
@@ -73,8 +87,13 @@ jobs:
           sleep 15
           xcrun simctl list devices
       - name: Build and Test - ${{ matrix.run-config['name'] }}
+        env:
+          MLX_DISABLED: ${{ matrix.run-config['mlx-disabled'] }}
         if: ${{ matrix.run-config['condition'] == true }}
         run: |
           set -o pipefail
-          xcodebuild clean build-for-testing -scheme whisperkit-Package -destination '${{ matrix.run-config['clean-destination'] }}' | xcpretty
-          xcodebuild test -only-testing WhisperKitTests/UnitTests -scheme whisperkit-Package -destination '${{ matrix.run-config['test-destination'] }}'
+          xcodebuild clean build-for-testing test \
+            ${{ matrix.run-config['test-cases'] }} \
+            -scheme ${{ matrix.run-config['scheme'] }} \
+            -destination '${{ matrix.run-config['test-destination'] }}' \
+            -skipPackagePluginValidation

.gitignore

@@ -9,6 +9,7 @@ DerivedData/
 **/*.xcscheme
 .netrc
 .env
+/.vscode
 
 # Core ML Model Files
 Models
@@ -23,6 +24,7 @@ Models
 **/*.mp3
 **/*.m4a
 **/*.flac
+!Sources/WhisperKitTestsUtils/Resources/*.*
 
 ## Xcode
 # Build generated
@@ -62,4 +64,4 @@ fastlane/test_output
 !*.xcodeproj/project.pbxproj
 !*.xcodeproj/xcshareddata/
 !*.xcworkspace/contents.xcworkspacedata
-/*.gcno
+/*.gcno

.swiftpm/configuration/Package.resolved

@@ -0,0 +1,41 @@
+{
+  "pins" : [
+    {
+      "identity" : "mlx-swift",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/ml-explore/mlx-swift",
+      "state" : {
+        "revision" : "597aaa5f465b4b9a17c8646b751053f84e37925b",
+        "version" : "0.16.0"
+      }
+    },
+    {
+      "identity" : "swift-argument-parser",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-argument-parser.git",
+      "state" : {
+        "revision" : "c8ed701b513cf5177118a175d85fbbbcd707ab41",
+        "version" : "1.3.0"
+      }
+    },
+    {
+      "identity" : "swift-numerics",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-numerics",
+      "state" : {
+        "revision" : "0a5bc04095a675662cf24757cc0640aa2204253b",
+        "version" : "1.0.2"
+      }
+    },
+    {
+      "identity" : "swift-transformers",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/huggingface/swift-transformers.git",
+      "state" : {
+        "revision" : "74b94211bdc741694ed7e700a1104c72e5ba68fe",
+        "version" : "0.1.7"
+      }
+    }
+  ],
+  "version" : 2
+}

Examples/WhisperAX/WhisperAX.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved

@@ -1,6 +1,15 @@
 {
-  "originHash" : "cd17206b47bb810af9459722192530e3838d8e6629a970988e32a432aaa05f6e",
+  "originHash" : "829222b514832cb61fe0002e0eebda98f23a75169c63f7d6ed7a320d57d5318f",
   "pins" : [
+    {
+      "identity" : "mlx-swift",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/ml-explore/mlx-swift",
+      "state" : {
+        "revision" : "597aaa5f465b4b9a17c8646b751053f84e37925b",
+        "version" : "0.16.0"
+      }
+    },
     {
       "identity" : "networkimage",
       "kind" : "remoteSourceControl",
@@ -24,8 +33,17 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/gonzalezreal/swift-markdown-ui.git",
       "state" : {
-        "revision" : "ae799d015a5374708f7b4c85f3294c05f2a564e2",
-        "version" : "2.3.0"
+        "revision" : "9a8119b37e09a770367eeb26e05267c75d854053",
+        "version" : "2.3.1"
+      }
+    },
+    {
+      "identity" : "swift-numerics",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-numerics",
+      "state" : {
+        "revision" : "0a5bc04095a675662cf24757cc0640aa2204253b",
+        "version" : "1.0.2"
       }
     },
     {

Makefile

@@ -1,12 +1,15 @@
-.PHONY: setup setup-huggingface-cli setup-model-repo download-models download-model build build-cli test clean-package-caches
+.PHONY: setup setup-huggingface-cli setup-model-repo download-models download-model download-mlx-models download-mlx-model build build-cli test clean-package-caches
 
 PIP_COMMAND := pip3
 PYTHON_COMMAND := python3
 
 # Define model repository and directories
 MODEL_REPO := argmaxinc/whisperkit-coreml
+MLX_MODEL_REPO := argmaxinc/whisperkit-mlx
+
 MODEL_REPO_DIR := ./Models/whisperkit-coreml
-BASE_COMPILED_DIR := ./Models
+MLX_MODEL_REPO_DIR := ./Models/whisperkit-mlx
+BASE_MODEL_DIR := ./Models
 
 
 setup:
@@ -29,9 +32,7 @@ setup:
 
 
 setup-huggingface-cli:
-	@if huggingface-cli whoami; then \
-		echo "Already logged in to Hugging Face."; \
-	else \
+	@if huggingface-cli whoami 2>&1 | grep -q "Not logged in"; then \
 		echo "Not logged in to Hugging Face."; \
 		if [ -z "$$HF_TOKEN" ]; then \
 			echo "Environment variable HF_TOKEN is not set. Running normal login."; \
@@ -40,12 +41,15 @@ setup-huggingface-cli:
 			echo "Using HF_TOKEN from environment variable."; \
 			huggingface-cli login --token $$HF_TOKEN; \
 		fi; \
+	else \
+		echo "Already logged in to Hugging Face."; \
+		huggingface-cli whoami; \
 	fi
 
 
 setup-model-repo:
 	@echo "Setting up repository..."
-	@mkdir -p $(BASE_COMPILED_DIR)
+	@mkdir -p $(BASE_MODEL_DIR)
 	@if [ -d "$(MODEL_REPO_DIR)/.git" ]; then \
 		echo "Repository exists, resetting..."; \
 		export GIT_LFS_SKIP_SMUDGE=1; \
@@ -56,37 +60,90 @@ setup-model-repo:
 		git clone https://huggingface.co/$(MODEL_REPO) $(MODEL_REPO_DIR); \
 	fi
 
+
+setup-mlx-model-repo:
+	@echo "Setting up mlx repository..."
+	@mkdir -p $(BASE_MODEL_DIR)
+	@if [ -d "$(MLX_MODEL_REPO_DIR)/.git" ]; then \
+		echo "Repository exists, resetting..."; \
+		export GIT_LFS_SKIP_SMUDGE=1; \
+		cd $(MLX_MODEL_REPO_DIR) && git fetch --all && git reset --hard origin/main && git clean -fdx; \
+	else \
+		echo "Repository not found, initializing..."; \
+		export GIT_LFS_SKIP_SMUDGE=1; \
+		git clone https://huggingface.co/$(MLX_MODEL_REPO) $(MLX_MODEL_REPO_DIR); \
+	fi
+
+
 # Download all models
 download-models: setup-model-repo
 	@echo "Downloading all models..."
 	@cd $(MODEL_REPO_DIR) && \
 	git lfs pull
+	@echo "CoreML models downloaded to $(MODEL_REPO_DIR)"
+
 
 # Download a specific model
-download-model:
+download-model: setup-model-repo
 	@if [ -z "$(MODEL)" ]; then \
-		echo "Error: MODEL is not set. Usage: make download-model MODEL=base"; \
+		echo "Error: MODEL is not set. Usage: make download-model MODEL=tiny"; \
 		exit 1; \
 	fi
 	@echo "Downloading model $(MODEL)..."
-	@$(MAKE) setup-model-repo
-	@echo "Fetching model $(MODEL)..."
 	@cd $(MODEL_REPO_DIR) && \
 	git lfs pull --include="openai_whisper-$(MODEL)/*"
+	@echo "CoreML model $(MODEL) downloaded to $(MODEL_REPO_DIR)/openai_whisper-$(MODEL)"
+
+
+download-mlx-models: setup-mlx-model-repo
+	@echo "Downloading all mlx models..."
+	@cd $(MLX_MODEL_REPO_DIR) && \
+	git lfs pull
+	@echo "MLX models downloaded to $(MLX_MODEL_REPO_DIR)"
+
+
+download-mlx-model: setup-mlx-model-repo
+	@if [ -z "$(MODEL)" ]; then \
+		echo "Error: MODEL is not set. Usage: make download-mlx-model MODEL=tiny"; \
+		exit 1; \
+	fi
+	@echo "Downloading mlx model $(MODEL)..."
+	@cd $(MLX_MODEL_REPO_DIR) && \
+	git lfs pull --include="openai_whisper-$(MODEL)/*"
+	@echo "MLX model $(MODEL) downloaded to $(MLX_MODEL_REPO_DIR)/openai_whisper-$(MODEL)"
+
 
 build:
 	@echo "Building WhisperKit..."
-	@swift build -v
+	@xcodebuild CLANG_ENABLE_CODE_COVERAGE=NO VALID_ARCHS=arm64 clean build \
+		-configuration Release \
+		-scheme whisperkit-Package \
+		-destination generic/platform=macOS \
+		-derivedDataPath .build/.xcodebuild/ \
+		-clonedSourcePackagesDirPath .build/ \
+		-skipPackagePluginValidation
 
 
 build-cli:
 	@echo "Building WhisperKit CLI..."
-	@swift build -c release --product whisperkit-cli
+	@xcodebuild CLANG_ENABLE_CODE_COVERAGE=NO VALID_ARCHS=arm64 clean build \
+		-configuration Release \
+		-scheme whisperkit-cli \
+		-destination generic/platform=macOS \
+		-derivedDataPath .build/.xcodebuild/ \
+		-clonedSourcePackagesDirPath .build/ \
+		-skipPackagePluginValidation
 
 
 test:
 	@echo "Running tests..."
-	@swift test -v
+	@xcodebuild clean build-for-testing test \
+		-scheme whisperkit-Package \
+		-only-testing WhisperKitMLXTests/MLXUnitTests \
+		-only-testing WhisperKitTests/UnitTests \
+		-destination 'platform=macOS,arch=arm64' \
+		-skipPackagePluginValidation
+
 
 clean-package-caches:
 	@trash ~/Library/Caches/org.swift.swiftpm/repositories

Package.resolved

@@ -1,5 +1,14 @@
 {
   "pins" : [
+    {
+      "identity" : "mlx-swift",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/ml-explore/mlx-swift",
+      "state" : {
+        "revision" : "597aaa5f465b4b9a17c8646b751053f84e37925b",
+        "version" : "0.16.0"
+      }
+    },
     {
       "identity" : "swift-argument-parser",
       "kind" : "remoteSourceControl",
@@ -9,6 +18,15 @@
         "version" : "1.3.0"
       }
     },
+    {
+      "identity" : "swift-numerics",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-numerics",
+      "state" : {
+        "revision" : "0a5bc04095a675662cf24757cc0640aa2204253b",
+        "version" : "1.0.2"
+      }
+    },
     {
       "identity" : "swift-transformers",
       "kind" : "remoteSourceControl",