From df1c5ff32c7ad5b0af14ffb9a9e176c139dc7689 Mon Sep 17 00:00:00 2001 From: Rocky <72559939+hkgood@users.noreply.github.com> Date: Tue, 23 Jun 2026 00:46:58 +0800 Subject: [PATCH] feat: migrate on-device Qwen3 ASR to CoreML for background Flow dictation Replace MLX GPU inference with CoreML bundles so transcription continues while the host app is backgrounded. Adds model download and warm-up, vendored Qwen3Speech, and updates onboarding, settings, and copy for the ~1.6 GB CoreML package (iOS 18+). --- .github/workflows/ci.yml | 4 +- CHANGELOG.md | 26 +- CONTRIBUTING.md | 2 +- OSGKeyboard/OSGKeyboardApp.swift | 22 +- OSGKeyboard/Resources/PrivacyPolicy.html | 98 ++ OSGKeyboard/Services/AppPermissions.swift | 7 +- OSGKeyboard/Services/FlowDiagnostics.swift | 22 + OSGKeyboard/Services/FlowSessionManager.swift | 272 ++++- OSGKeyboard/Services/LegalLinks.swift | 4 +- .../Services/ModelDownloadSourcePicker.swift | 126 +++ OSGKeyboard/Services/ModelManager.swift | 492 +++++++++ .../Services/OnDeviceModelWarmup.swift | 197 ++++ .../Services/OpenSourceLicenseCatalog.swift | 115 +++ OSGKeyboard/Services/Qwen3ASRService.swift | 257 +++++ .../ThirdParty/Qwen3Speech/Package.swift | 99 ++ .../Sources/AudioCommon/AudioFileLoader.swift | 382 +++++++ .../Sources/AudioCommon/AudioIO.swift | 174 ++++ .../Sources/AudioCommon/AudioModelError.swift | 34 + .../Sources/AudioCommon/AudioRingBuffer.swift | 75 ++ .../AudioCommon/CoreMLComputeUnits.swift | 49 + .../Sources/AudioCommon/CoreMLLoader.swift | 125 +++ .../AudioCommon/HuggingFaceDownloader.swift | 428 ++++++++ .../Sources/AudioCommon/Logging.swift | 13 + .../Sources/AudioCommon/ModelLoader.swift | 175 ++++ .../Sources/AudioCommon/ModelRegistry.swift | 9 + .../AudioCommon/ModelScopeDownloader.swift | 335 +++++++ .../Sources/AudioCommon/PipelineLLM.swift | 53 + .../Sources/AudioCommon/Protocols.swift | 282 ++++++ .../AudioCommon/SentencePieceModel.swift | 182 ++++ .../AudioCommon/StreamingAudioPlayer.swift | 511 ++++++++++ .../Sources/AudioCommon/Tokenizer.swift | 335 +++++++ .../Sources/AudioCommon/WAVWriter.swift | 105 ++ .../Sources/MLXCommon/EncodecLSTM.swift | 78 ++ .../Sources/MLXCommon/MetalBudget.swift | 58 ++ .../Sources/MLXCommon/ModuleMemory.swift | 26 + .../MLXCommon/PreQuantizedEmbedding.swift | 50 + .../Sources/MLXCommon/QuantizedMLP.swift | 56 ++ .../Qwen3Speech/Sources/MLXCommon/SDPA.swift | 102 ++ .../Sources/MLXCommon/WeightLoading.swift | 292 ++++++ .../Sources/Qwen3ASR/AudioEncoder.swift | 512 ++++++++++ .../Sources/Qwen3ASR/AudioPreprocessing.swift | 490 +++++++++ .../Sources/Qwen3ASR/Configuration.swift | 158 +++ .../Sources/Qwen3ASR/CoreMLASRModel.swift | 389 ++++++++ .../Sources/Qwen3ASR/CoreMLEncoder.swift | 231 +++++ .../Sources/Qwen3ASR/CoreMLTextDecoder.swift | 566 +++++++++++ .../Sources/Qwen3ASR/ExportedImports.swift | 3 + .../Sources/Qwen3ASR/FloatTextDecoder.swift | 240 +++++ .../Qwen3ASR/ForcedAligner+Protocols.swift | 9 + .../Sources/Qwen3ASR/ForcedAligner.swift | 482 +++++++++ .../Qwen3ASR/QuantizedTextDecoder.swift | 252 +++++ .../Sources/Qwen3ASR/Qwen3ASR+Memory.swift | 27 + .../Sources/Qwen3ASR/Qwen3ASR+Protocols.swift | 11 + .../Sources/Qwen3ASR/Qwen3ASR.swift | 930 ++++++++++++++++++ .../Sources/Qwen3ASR/StreamingASR.swift | 277 ++++++ .../Sources/Qwen3ASR/TextPreprocessing.swift | 308 ++++++ .../Qwen3ASR/TimestampCorrection.swift | 145 +++ .../Sources/Qwen3ASR/WeightLoading.swift | 321 ++++++ .../Sources/Qwen3Chat/ChatSampler.swift | 107 ++ .../Sources/Qwen3Chat/ChatTemplate.swift | 104 ++ .../Sources/Qwen3Chat/ChatTokenizer.swift | 272 +++++ .../Sources/Qwen3Chat/MLXGenerator.swift | 425 ++++++++ .../Sources/Qwen3Chat/Qwen35CoreMLChat.swift | 381 +++++++ .../Sources/Qwen3Chat/Qwen35Model.swift | 710 +++++++++++++ .../Sources/Qwen3Chat/Qwen35PipelineLLM.swift | 79 ++ .../Qwen3Chat/Qwen35WeightLoading.swift | 226 +++++ .../Sources/Qwen3Chat/Qwen3ChatConfig.swift | 146 +++ .../Sources/Qwen3Chat/Qwen3ChatError.swift | 25 + .../Sources/SpeechVAD/BiLSTM.swift | 100 ++ .../Sources/SpeechVAD/Configuration.swift | 92 ++ .../SpeechVAD/CoreMLSileroInference.swift | 64 ++ .../SpeechVAD/CoreMLWeSpeakerInference.swift | 76 ++ .../Sources/SpeechVAD/DERScoring.swift | 408 ++++++++ .../SpeechVAD/DiarizationHelpers.swift | 183 ++++ .../SpeechVAD/DiarizationPipeline.swift | 570 +++++++++++ .../Sources/SpeechVAD/FireRedVAD.swift | 505 ++++++++++ .../SpeechVAD/MelFeatureExtractor.swift | 238 +++++ .../Sources/SpeechVAD/PowersetDecoder.swift | 73 ++ .../SpeechVAD/PyannoteVAD+Memory.swift | 16 + .../Sources/SpeechVAD/Segmentation.swift | 97 ++ .../Sources/SpeechVAD/SileroModel.swift | 186 ++++ .../Sources/SpeechVAD/SileroVAD+Memory.swift | 19 + .../Sources/SpeechVAD/SileroVAD.swift | 321 ++++++ .../SpeechVAD/SileroWeightLoading.swift | 36 + .../Sources/SpeechVAD/SincNet.swift | 129 +++ .../Sources/SpeechVAD/SortformerConfig.swift | 114 +++ .../SpeechVAD/SortformerDiarizer.swift | 432 ++++++++ .../SpeechVAD/SortformerMelExtractor.swift | 205 ++++ .../Sources/SpeechVAD/SortformerModel.swift | 161 +++ .../SpeechVAD/SpeechVAD+Protocols.swift | 29 + .../Sources/SpeechVAD/SpeechVAD.swift | 142 +++ .../SpeechVAD/StreamingVADProcessor.swift | 210 ++++ .../Sources/SpeechVAD/VADPipeline.swift | 181 ++++ .../Sources/SpeechVAD/WeSpeaker+Memory.swift | 19 + .../Sources/SpeechVAD/WeSpeaker.swift | 231 +++++ .../Sources/SpeechVAD/WeSpeakerModel.swift | 167 ++++ .../SpeechVAD/WeSpeakerWeightLoading.swift | 33 + .../Sources/SpeechVAD/WeightLoading.swift | 37 + OSGKeyboard/Utilities/ASRLocaleLabels.swift | 32 + OSGKeyboard/Utilities/AppL10n.swift | 40 + OSGKeyboard/Views/APISettingsCard.swift | 30 +- .../Views/Components/LegalWebView.swift | 50 + .../Views/Components/RemoteWebView.swift | 52 + .../Views/Components/SafariSheet.swift | 5 +- OSGKeyboard/Views/DictationCaptureView.swift | 51 +- OSGKeyboard/Views/DownloadConfirmSheet.swift | 85 ++ OSGKeyboard/Views/EnginePickerSection.swift | 32 +- OSGKeyboard/Views/HelpFeedbackView.swift | 35 + OSGKeyboard/Views/HomeView.swift | 160 ++- OSGKeyboard/Views/KeyboardPreviewSheet.swift | 33 +- .../Views/LocalEngineSettingsRows.swift | 91 ++ OSGKeyboard/Views/OnDeviceModelsView.swift | 184 ++++ OSGKeyboard/Views/OnboardingView.swift | 118 ++- .../Views/OpenSourceLicensesView.swift | 115 +++ OSGKeyboard/Views/PrivacyPolicyView.swift | 33 + OSGKeyboard/Views/SettingsView.swift | 255 ++--- .../Views/SystemPromptSettingsView.swift | 53 + OSGKeyboard/en.lproj/Localizable.strings | 95 +- OSGKeyboard/zh-Hans.lproj/Localizable.strings | 95 +- OSGKeyboardExt/KeyboardViewController.swift | 87 +- .../Services/AppGroupPersistor.swift | 40 +- OSGKeyboardExt/Utilities/ExtL10n.swift | 13 +- OSGKeyboardExt/Views/KeyboardRootView.swift | 108 +- OSGKeyboardExt/en.lproj/Keyboard.strings | 10 +- OSGKeyboardExt/zh-Hans.lproj/Keyboard.strings | 8 +- .../Localization/SharedL10n.swift | 41 + OSGKeyboardShared/Models/AppUILanguage.swift | 87 ++ .../Models/EngineServiceLabel.swift | 37 +- .../Models/FlowUtteranceChunkConfig.swift | 69 ++ .../Models/LocalASRBackend.swift | 56 ++ OSGKeyboardShared/Models/OnDeviceModel.swift | 54 + OSGKeyboardShared/Models/ProviderConfig.swift | 34 +- .../Models/TranscriptionDelivery.swift | 17 + OSGKeyboardShared/Services/ASRService.swift | 253 ++++- .../Services/AppGroupStore.swift | 25 +- .../Services/ChunkedUtterancePipeline.swift | 160 +++ .../Services/DictationBridge.swift | 27 +- .../Services/FlowAppLifecycle.swift | 39 + .../Services/FlowContinuousCapture.swift | 18 +- .../Services/FlowSessionBridge.swift | 35 +- .../Services/FlowSessionKeys.swift | 28 +- .../Services/KeyboardState.swift | 12 +- OSGKeyboardShared/Services/LLMClient.swift | 21 +- .../Services/LiveDictationController.swift | 73 +- .../Services/OnDeviceModelStatus.swift | 104 ++ .../Services/PolishingService.swift | 31 +- ...essiveDictationTranscriptAccumulator.swift | 64 ++ .../Utilities/ProviderDisplayName.swift | 7 +- .../Utilities/UtteranceStreamChunker.swift | 101 ++ .../UtteranceTranscriptStitcher.swift | 109 ++ OSGKeyboardShared/en.lproj/Shared.strings | 30 + .../zh-Hans.lproj/Shared.strings | 30 + .../ChunkedUtterancePipelineTests.swift | 114 +++ OSGKeyboardTests/FlowSessionBridgeTests.swift | 26 +- OSGKeyboardTests/LLMClientTests.swift | 65 +- ...eDictationTranscriptAccumulatorTests.swift | 44 + .../UtteranceStreamChunkerTests.swift | 42 + .../UtteranceTranscriptStitcherTests.swift | 23 + Scripts/generate-xcodeproj.sh | 10 + Scripts/patch-icon-composer.sh | 138 +++ project.yml | 66 +- 160 files changed, 22080 insertions(+), 492 deletions(-) create mode 100644 OSGKeyboard/Resources/PrivacyPolicy.html create mode 100644 OSGKeyboard/Services/FlowDiagnostics.swift create mode 100644 OSGKeyboard/Services/ModelDownloadSourcePicker.swift create mode 100644 OSGKeyboard/Services/ModelManager.swift create mode 100644 OSGKeyboard/Services/OnDeviceModelWarmup.swift create mode 100644 OSGKeyboard/Services/OpenSourceLicenseCatalog.swift create mode 100644 OSGKeyboard/Services/Qwen3ASRService.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Package.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioFileLoader.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioIO.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioModelError.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioRingBuffer.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/CoreMLComputeUnits.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/CoreMLLoader.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/HuggingFaceDownloader.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Logging.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelLoader.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelRegistry.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelScopeDownloader.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/PipelineLLM.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Protocols.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/SentencePieceModel.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/StreamingAudioPlayer.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Tokenizer.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/WAVWriter.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/EncodecLSTM.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/MetalBudget.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/ModuleMemory.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/PreQuantizedEmbedding.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/QuantizedMLP.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/SDPA.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/MLXCommon/WeightLoading.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/AudioEncoder.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/AudioPreprocessing.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Configuration.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/CoreMLASRModel.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/CoreMLEncoder.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/CoreMLTextDecoder.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/ExportedImports.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/FloatTextDecoder.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/ForcedAligner+Protocols.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/ForcedAligner.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/QuantizedTextDecoder.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Qwen3ASR+Memory.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Qwen3ASR+Protocols.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/Qwen3ASR.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/StreamingASR.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/TextPreprocessing.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/TimestampCorrection.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3ASR/WeightLoading.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/ChatSampler.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/ChatTemplate.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/ChatTokenizer.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/MLXGenerator.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen35CoreMLChat.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen35Model.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen35PipelineLLM.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen35WeightLoading.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen3ChatConfig.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/Qwen3Chat/Qwen3ChatError.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/BiLSTM.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/Configuration.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/CoreMLSileroInference.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/CoreMLWeSpeakerInference.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/DERScoring.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/DiarizationHelpers.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/DiarizationPipeline.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/FireRedVAD.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/MelFeatureExtractor.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/PowersetDecoder.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/PyannoteVAD+Memory.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/Segmentation.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroModel.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroVAD+Memory.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroVAD.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SileroWeightLoading.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SincNet.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SortformerConfig.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SortformerDiarizer.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SortformerMelExtractor.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SortformerModel.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SpeechVAD+Protocols.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/SpeechVAD.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/StreamingVADProcessor.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/VADPipeline.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeaker+Memory.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeaker.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeakerModel.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeSpeakerWeightLoading.swift create mode 100644 OSGKeyboard/ThirdParty/Qwen3Speech/Sources/SpeechVAD/WeightLoading.swift create mode 100644 OSGKeyboard/Utilities/ASRLocaleLabels.swift create mode 100644 OSGKeyboard/Utilities/AppL10n.swift create mode 100644 OSGKeyboard/Views/Components/LegalWebView.swift create mode 100644 OSGKeyboard/Views/Components/RemoteWebView.swift create mode 100644 OSGKeyboard/Views/DownloadConfirmSheet.swift create mode 100644 OSGKeyboard/Views/HelpFeedbackView.swift create mode 100644 OSGKeyboard/Views/LocalEngineSettingsRows.swift create mode 100644 OSGKeyboard/Views/OnDeviceModelsView.swift create mode 100644 OSGKeyboard/Views/OpenSourceLicensesView.swift create mode 100644 OSGKeyboard/Views/PrivacyPolicyView.swift create mode 100644 OSGKeyboard/Views/SystemPromptSettingsView.swift create mode 100644 OSGKeyboardShared/Localization/SharedL10n.swift create mode 100644 OSGKeyboardShared/Models/AppUILanguage.swift create mode 100644 OSGKeyboardShared/Models/FlowUtteranceChunkConfig.swift create mode 100644 OSGKeyboardShared/Models/LocalASRBackend.swift create mode 100644 OSGKeyboardShared/Models/OnDeviceModel.swift create mode 100644 OSGKeyboardShared/Models/TranscriptionDelivery.swift create mode 100644 OSGKeyboardShared/Services/ChunkedUtterancePipeline.swift create mode 100644 OSGKeyboardShared/Services/FlowAppLifecycle.swift create mode 100644 OSGKeyboardShared/Services/OnDeviceModelStatus.swift create mode 100644 OSGKeyboardShared/Utilities/ProgressiveDictationTranscriptAccumulator.swift create mode 100644 OSGKeyboardShared/Utilities/UtteranceStreamChunker.swift create mode 100644 OSGKeyboardShared/Utilities/UtteranceTranscriptStitcher.swift create mode 100644 OSGKeyboardShared/en.lproj/Shared.strings create mode 100644 OSGKeyboardShared/zh-Hans.lproj/Shared.strings create mode 100644 OSGKeyboardTests/ChunkedUtterancePipelineTests.swift create mode 100644 OSGKeyboardTests/ProgressiveDictationTranscriptAccumulatorTests.swift create mode 100644 OSGKeyboardTests/UtteranceStreamChunkerTests.swift create mode 100644 OSGKeyboardTests/UtteranceTranscriptStitcherTests.swift create mode 100755 Scripts/generate-xcodeproj.sh create mode 100755 Scripts/patch-icon-composer.sh diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f6141dc..73b87c8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -50,7 +50,7 @@ jobs: - name: Install XcodeGen run: brew install xcodegen - name: Generate project - run: xcodegen generate + run: ./Scripts/generate-xcodeproj.sh - name: Build run: | set -o pipefail @@ -76,7 +76,7 @@ jobs: - name: Install XcodeGen run: brew install xcodegen - name: Generate project - run: xcodegen generate + run: ./Scripts/generate-xcodeproj.sh - name: Run tests run: | set -o pipefail diff --git a/CHANGELOG.md b/CHANGELOG.md index dc4a73e..36a0ef4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,13 +7,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -### Changed -- **iPhone only**: Set `TARGETED_DEVICE_FAMILY` to `"1"` for both `OSGKeyboard` and `OSGKeyboardExt` targets. Removed `UIRequiresFullScreen` and trimmed `UISupportedInterfaceOrientations` to Portrait only (iPad is no longer a supported device). -- **Remove top divider line**: Deleted the 0.5 pt `palette.divider` overlay from `KeyboardRootView` — the subtle highlight gradient is retained; the hard separator line is gone. -- **Keyboard preview always dark**: `KeyboardPreviewSheet` now injects `.environment(\.themePalette, Palette.dark)` alongside `.environment(\.colorScheme, .dark)` on `KeyboardPreviewStub`, so the preview palette is always the dark variant regardless of the app's active theme. -- **Docs consistency**: README/README.zh now consistently describe the currently implemented capability set (`iOS 26+`, on-device `SpeechAnalyzer` + `DictationTranscriber`) with no deferred-ASR wording. +## [0.2.0] - 2026-06-22 -## [0.1.2] - In Progress +### Added +- **Local engine with on-device Qwen models**: Optional Qwen3-ASR 0.6B speech recognition and Qwen3.5-0.8B text polish, fully offline after download. +- **On-device model management**: Download, progress, delete, and readiness status in Settings; mirror auto-selection between ModelScope and Hugging Face with fallback. +- **Engine picker**: Choose between local (on-device ASR + polish) and cloud (ASR + user-configured LLM polish). +- **Flow session dictation**: TypeWhisper-style continuous capture in the host app with keyboard handoff via App Group. +- **Open-source licenses** screen for bundled third-party components. + +### Changed +- **Settings simplified**: Merged language and model sections; cloud mode always enables polish (removed off/transcribe mode picker). +- **Keyboard UI**: Local/cloud engine badges replace the mode menu; shows model-not-downloaded guidance when the local stack is incomplete. +- **iPhone only**: Set `TARGETED_DEVICE_FAMILY` to `"1"` for both targets; portrait-only orientations. +- **Keyboard preview always dark**: Preview injects the dark palette regardless of app theme. +- **Docs consistency**: README/README.zh aligned to the iOS 26+ capability set. + +### Fixed +- **ModelScope download progress**: Progress now tracks byte counts instead of jumping to 50% after the first file. +- **Light/Dark mode consistency** for shared button/card modifiers via `@Environment(\.themePalette)`. + +## [0.1.2] - 2026-06-20 ### Fixed - **Light/Dark mode consistency**: `cardSurface()`, `primaryButton()`, `secondaryButton()`, and `pillChip()` view modifiers in `Theme.swift` now use `ViewModifier` structs that read from `@Environment(\.themePalette)`. Previously they used hardcoded dark `Palette` constants, causing cards and buttons to always render in dark mode even when the main App was in light mode. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d2a7a9f..2f6bb68 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -29,7 +29,7 @@ Open an issue using the **Feature request** template. Briefly describe: 2. **Generate the project locally:** ```bash brew install xcodegen swiftlint - xcodegen generate + xcodegen generate # or: ./Scripts/generate-xcodeproj.sh ``` 3. **Code style.** SwiftLint config lives in `.swiftlint.yml` — keep it green. We use Swift 6 strict concurrency, no `Sendable` shims where avoidable. 4. **Tests.** Add XCTest coverage in `OSGKeyboardTests/` for any non-trivial logic. diff --git a/OSGKeyboard/OSGKeyboardApp.swift b/OSGKeyboard/OSGKeyboardApp.swift index 181df28..6e610f3 100644 --- a/OSGKeyboard/OSGKeyboardApp.swift +++ b/OSGKeyboard/OSGKeyboardApp.swift @@ -14,6 +14,14 @@ struct OSGKeyboardApp: App { init() { MaterialIconsFont.registerIfNeeded() + // Register backend-specific ASR providers. The shared + // framework ships a built-in SpeechAnalyzer provider; we + // install the Qwen3-ASR provider here because linking + // `Qwen3ASR` pulls in mlx-swift, which the keyboard + // extension's `APPLICATION_EXTENSION_API_ONLY` build would + // refuse. Doing it in the host app's `init` keeps the heavy + // dependency localised. + ASRServiceFactory.providers[.qwen3ASR] = Qwen3ASRServiceProvider() } var body: some Scene { @@ -29,8 +37,10 @@ struct OSGKeyboardApp: App { AppGroupErrorView() } } + .environment(\.locale, config.uiLanguage.swiftUILocale) .environmentObject(flowManager) .onAppear { + FlowAppLifecycle.shared.setForeground(scenePhase == .active) flowManager.setAppForeground(scenePhase == .active) } .onOpenURL { url in @@ -51,11 +61,19 @@ struct OSGKeyboardApp: App { ) } .onChange(of: config.hasCompletedOnboarding) { _, done in - if done { flowManager.autoStartIfNeeded() } + if done { + flowManager.autoStartIfNeeded() + if config.isLocalEngine { + OnDeviceModelWarmup.shared.warmUpIfNeeded() + } + } } .onChange(of: scenePhase) { _, phase in - flowManager.setAppForeground(phase == .active) + flowManager.handleScenePhase(phase) guard phase == .active, AppGroup.isAvailable, config.hasCompletedOnboarding else { return } + if config.isLocalEngine { + OnDeviceModelWarmup.shared.ensureReadyAfterBackground() + } if flowManager.isActive { flowManager.extendSession() } else { diff --git a/OSGKeyboard/Resources/PrivacyPolicy.html b/OSGKeyboard/Resources/PrivacyPolicy.html new file mode 100644 index 0000000..d5629d3 --- /dev/null +++ b/OSGKeyboard/Resources/PrivacyPolicy.html @@ -0,0 +1,98 @@ + + + + + + OSGKeyboard Privacy Policy + + + +

中文 · English

+
+

OSGKeyboard Privacy Policy

+

Last updated: June 19, 2026

+

OSGKeyboard is a custom iOS keyboard that turns your voice into text. This policy explains what data the app processes and how it is used.

+ +

What we collect

+ + +

What we do not collect

+ + +

Permissions

+ + +

Third parties

+

When you choose Cloud polish mode, transcribed text is sent to the API endpoint you configure. That provider’s privacy policy applies to those requests.

+ +

Data retention

+

Settings and API keys remain on your device until you delete the app or reset settings. Transcription results are passed to the host app you are typing in and are not stored long-term by OSGKeyboard.

+

Voice history — successful transcripts may be saved locally in the main app’s History tab for your convenience. This history stays on your device only, is never uploaded, and can be cleared at any time from History or by resetting settings.

+ +

Contact

+

Questions: open an issue at github.com/hkgood/OSGKeyboard.

+
+ +
+

OSGKeyboard 隐私政策

+

更新日期：2026 年 6 月 19 日

+

OSGKeyboard 是一款 iOS 自定义键盘，可将语音转为文字。本政策说明应用处理哪些数据及用途。

+ +

我们处理的数据

+ + +

我们不收集的内容

+ + +

权限说明

+ + +

第三方

+

选择云端润色时，转写文字会发往你配置的 API，该服务商的隐私政策适用于相关请求。

+ +

数据保留

+

设置与 API Key 保留在设备上，直至卸载或重置。识别结果写入你正在使用的宿主 App，OSGKeyboard 不会长期存储。

+

语音历史 — 成功的转写可能保存在主 App「历史」页，仅供本机查看，不会上传，可随时在历史页清空或通过重置设置清除。

+ +

联系

+

问题反馈：github.com/hkgood/OSGKeyboard

+ + diff --git a/OSGKeyboard/Services/AppPermissions.swift b/OSGKeyboard/Services/AppPermissions.swift index cf46743..8aa5f2d 100644 --- a/OSGKeyboard/Services/AppPermissions.swift +++ b/OSGKeyboard/Services/AppPermissions.swift @@ -72,6 +72,7 @@ enum AppPermissions { } } + @MainActor static func openSystemSettings() { guard let url = URL(string: UIApplication.openSettingsURLString) else { return } UIApplication.shared.open(url) @@ -82,12 +83,12 @@ enum AppPermissions { let micMissing = micStatus != .granted let speechMissing = speechStatus != .granted if micMissing && speechMissing { - return NSLocalizedString("home.setup.permission.both", comment: "") + return AppL10n.string("home.setup.permission.both") } if micMissing { - return NSLocalizedString("home.setup.permission.mic", comment: "") + return AppL10n.string("home.setup.permission.mic") } - return NSLocalizedString("home.setup.permission.speech", comment: "") + return AppL10n.string("home.setup.permission.speech") } /// True when at least one permission can still be requested in-app. diff --git a/OSGKeyboard/Services/FlowDiagnostics.swift b/OSGKeyboard/Services/FlowDiagnostics.swift new file mode 100644 index 0000000..3861257 --- /dev/null +++ b/OSGKeyboard/Services/FlowDiagnostics.swift @@ -0,0 +1,22 @@ +// FlowDiagnostics.swift +// OSGKeyboard · Main App +// +// Structured logging for the Flow dictation pipeline. Visible in Xcode +// console (DEBUG) and Console.app via `subsystem: com.osgkeyboard.ios`. + +import Foundation +import os + +enum FlowDiagnostics { + private static let logger = Logger( + subsystem: "com.osgkeyboard.ios", + category: "Flow" + ) + + static func log(_ message: String) { + logger.info("\(message, privacy: .public)") + #if DEBUG + print("🌊[OSGFlow] \(message)") + #endif + } +} diff --git a/OSGKeyboard/Services/FlowSessionManager.swift b/OSGKeyboard/Services/FlowSessionManager.swift index 6a278a2..3880147 100644 --- a/OSGKeyboard/Services/FlowSessionManager.swift +++ b/OSGKeyboard/Services/FlowSessionManager.swift @@ -3,12 +3,14 @@ // // Session Owner for TypeWhisper-style Flow dictation: continuous // `.playAndRecord` capture for the whole session, utterance gating for -// ASR, optional LLM polish, and App Group result delivery. +// ASR and cloud LLM polish, with App Group result delivery. import Foundation import AVFoundation import Speech import OSGKeyboardShared +import UIKit +import SwiftUI @MainActor final class FlowSessionManager: ObservableObject { @@ -19,9 +21,22 @@ final class FlowSessionManager: ObservableObject { @Published private(set) var sessionWarning: String? private let capture = FlowContinuousCapture() - private let asr: ASRService = ASRServiceFactory.make() - private let polisher = PolishingService() private let store = AppGroupStore() + /// Cloud-engine polish only; local engine delivers raw ASR text. + private var polisher: PolishingService { + PolishingService() + } + /// Cached ASR instance shared with `OnDeviceModelWarmup`. + private var sessionASR: ASRService? + private var asr: ASRService { + if let sessionASR { return sessionASR } + let service = OnDeviceModelWarmup.shared.asrService( + engineMode: store.engineMode, + localBackend: store.localASRBackend + ) + sessionASR = service + return service + } private var pollingTask: Task? private var heartbeatTask: Task? @@ -33,10 +48,13 @@ final class FlowSessionManager: ObservableObject { private var isUtteranceProcessing = false private var finalizeTask: Task? private var asrTask: Task? + private var chunkedPipeline: ChunkedUtterancePipeline? private var currentPartial = "" private var lastFinal = "" + private var chunkWarnings: [String] = [] /// True while the host app scene is `.active` — drives foreground renewal. private var isAppForeground = false + private var backgroundTaskID: UIBackgroundTaskIdentifier = .invalid init() { Task { @MainActor [weak self] in @@ -126,6 +144,9 @@ final class FlowSessionManager: ObservableObject { startLevelPublishing() scheduleExpiry(after: remaining) + OnDeviceModelWarmup.shared.warmUpIfNeeded() + bindSessionASR() + debug("Flow session restored (\(Int(remaining))s remaining)") } @@ -149,13 +170,17 @@ final class FlowSessionManager: ObservableObject { if isUtteranceRecording || isUtteranceProcessing { capture.cancelUtterance() asrTask?.cancel() + Task { await chunkedPipeline?.cancel() } asr.cancel() } asrTask = nil + chunkedPipeline = nil isUtteranceRecording = false isUtteranceProcessing = false capture.stop() + endBackgroundKeepAlive() + sessionASR = nil FlowSessionBridge.markSessionInactive() FlowSessionDarwin.postSessionChanged() isActive = false @@ -179,6 +204,81 @@ final class FlowSessionManager: ObservableObject { } } + /// Full scene lifecycle — keeps Flow + ASR alive across app switches. + func handleScenePhase(_ phase: ScenePhase) { + switch phase { + case .active: + FlowAppLifecycle.shared.setForeground(true) + setAppForeground(true) + resumeAfterForeground() + case .inactive: + writeHeartbeatIfActive() + case .background: + FlowAppLifecycle.shared.setForeground(false) + setAppForeground(false) + beginBackgroundKeepAlive() + @unknown default: + break + } + } + + private func writeHeartbeatIfActive() { + guard isActive else { return } + FlowSessionBridge.writeHeartbeat() + } + + private func beginBackgroundKeepAlive() { + guard isActive else { return } + FlowSessionBridge.writeHeartbeat() + + guard backgroundTaskID == .invalid else { return } + backgroundTaskID = UIApplication.shared.beginBackgroundTask { [weak self] in + self?.endBackgroundKeepAlive() + } + debug("background keep-alive started") + } + + private func endBackgroundKeepAlive() { + guard backgroundTaskID != .invalid else { return } + UIApplication.shared.endBackgroundTask(backgroundTaskID) + backgroundTaskID = .invalid + debug("background keep-alive ended") + } + + private func resumeAfterForeground() { + guard isActive else { + endBackgroundKeepAlive() + return + } + + FlowSessionBridge.writeHeartbeat() + endBackgroundKeepAlive() + + Task { @MainActor [weak self] in + await self?.reactivateCaptureIfNeeded() + OnDeviceModelWarmup.shared.ensureReadyAfterBackground() + self?.bindSessionASR() + } + } + + private func reactivateCaptureIfNeeded() async { + guard isActive else { return } + + if capture.running { + capture.reassertIfRunning() + return + } + + do { + try capture.start() + debug("capture restarted after foreground") + } catch { + let message = (error as? LocalizedError)?.errorDescription ?? error.localizedDescription + sessionWarning = message + debug("capture restart failed: \(message)") + } + } + /// Extend the session before it expires while the host app stays in foreground. private func renewSessionIfNeededWhileForeground() { guard isActive, isAppForeground else { return } @@ -220,14 +320,24 @@ final class FlowSessionManager: ObservableObject { startLevelPublishing() scheduleExpiry(after: duration) + OnDeviceModelWarmup.shared.warmUpIfNeeded() + bindSessionASR() + debug("Flow session started (\(Int(duration))s), continuous capture running") } + private func bindSessionASR() { + sessionASR = OnDeviceModelWarmup.shared.asrService( + engineMode: store.engineMode, + localBackend: store.localASRBackend + ) + } + private func permissionWarningMessage() -> String { if AppPermissions.micStatus != .granted { - return NSLocalizedString("flow.error.micRequired", comment: "") + return AppL10n.string("flow.error.micRequired") } - return NSLocalizedString("flow.error.speechRequired", comment: "") + return AppL10n.string("flow.error.speechRequired") } // MARK: - Polling @@ -259,19 +369,20 @@ final class FlowSessionManager: ObservableObject { private func beginUtterance() { guard capture.running else { - failUtterance(message: NSLocalizedString("flow.error.audioUnavailable", comment: "")) + failUtterance(message: AppL10n.string("flow.error.audioUnavailable")) return } - // Mirror `LiveDictationController.start`: only begin when the previous - // utterance fully finished. Never cancel an in-flight analyzer here — - // that was the source of intermittent CancellationError / noSpeech. guard !isUtteranceProcessing else { debug("beginUtterance ignored — previous utterance still processing") return } + // Honor engine / ASR backend changes without restarting the session. + bindSessionASR() + currentPartial = "" lastFinal = "" + chunkWarnings = [] let localeId = store.localeId FlowSessionBridge.setTranscriptionLanguage(localeId) @@ -279,28 +390,42 @@ final class FlowSessionManager: ObservableObject { let locale = SpeechLocaleResolver.resolve(localeId) let stream = capture.beginUtterance() - let events = asr.transcribe(stream: stream, locale: locale) + let pipeline = ChunkedUtterancePipeline(asr: asr, locale: locale) + chunkedPipeline = pipeline isUtteranceRecording = true + FlowDiagnostics.log( + "beginUtterance engine=\(store.engineMode) asr=\(store.localASRBackend.rawValue) " + + "modelsInMemory=\(OnDeviceModelStatus.modelsLoadedInMemory()) " + + "asrType=\(type(of: asr)) pipelined=true max=\(Int(FlowSessionKeys.maxUtteranceDuration))s" + ) - asrTask = Task { @MainActor [weak self] in - guard let self else { return } - for await event in events { - switch event { - case .capability: - break - case .partial(let text): - self.currentPartial = text - case .final(let text): - self.lastFinal = text.trimmingCharacters(in: .whitespacesAndNewlines) - self.currentPartial = "" - case .error(let message): - self.debug("asr error: \(message)") - if self.isUtteranceRecording { - self.failUtterance(message: message) - } else if self.isUtteranceProcessing { - self.finishProcessing(withError: message) + asrTask = Task.detached(priority: .userInitiated) { [weak manager = self] in + let outcome = await pipeline.transcribe(stream: stream) { partial in + Task { @MainActor in + manager?.currentPartial = partial + } + } + await MainActor.run { + guard let manager else { return } + FlowDiagnostics.log( + "chunkedASR finished partialLen=\(manager.currentPartial.count) " + + "finalPending=\(manager.lastFinal.isEmpty)" + ) + switch outcome { + case .success(let success): + manager.lastFinal = success.text + manager.chunkWarnings = success.chunkWarnings + manager.currentPartial = "" + case .failure(let message): + manager.debug("asr error: \(message)") + if manager.isUtteranceRecording { + manager.failUtterance(message: message) + } else if manager.isUtteranceProcessing { + manager.finishProcessing(withError: message) } + case .cancelled: + break } } } @@ -335,10 +460,13 @@ final class FlowSessionManager: ObservableObject { finalizeTask?.cancel() finalizeTask = nil asrTask?.cancel() + Task { await chunkedPipeline?.cancel() } + chunkedPipeline = nil asr.cancel() capture.cancelUtterance() currentPartial = "" lastFinal = "" + chunkWarnings = [] FlowSessionBridge.setRecordingState(.idle) debug("utterance aborted") } @@ -349,10 +477,13 @@ final class FlowSessionManager: ObservableObject { finalizeTask?.cancel() finalizeTask = nil asrTask?.cancel() + Task { await chunkedPipeline?.cancel() } + chunkedPipeline = nil asr.cancel() capture.cancelUtterance() currentPartial = "" lastFinal = "" + chunkWarnings = [] FlowSessionBridge.storeTranscriptionError(message) FlowSessionBridge.setRecordingState(.idle) debug("utterance failed: \(message)") @@ -362,26 +493,43 @@ final class FlowSessionManager: ObservableObject { isUtteranceProcessing = false finalizeTask?.cancel() finalizeTask = nil + chunkedPipeline = nil currentPartial = "" lastFinal = "" + chunkWarnings = [] FlowSessionBridge.storeTranscriptionError(message) FlowSessionBridge.setRecordingState(.idle) debug("utterance processing failed: \(message)") } private func finalizeUtterance() async { + let pipelineStarted = Date() defer { isUtteranceProcessing = false FlowSessionBridge.setRecordingState(.idle) } - let deadline = Date().addingTimeInterval(30) - while Date() < deadline { + let asrWait = asrWaitTimeout() + FlowDiagnostics.log( + "finalize start asrWait=\(Int(asrWait))s engine=\(store.engineMode) " + + "backend=\(store.localASRBackend.rawValue)" + ) + + let asrDeadline = Date().addingTimeInterval(asrWait) + while Date() < asrDeadline { if !lastFinal.isEmpty { break } if asrTask?.isCancelled == true { break } try? await Task.sleep(nanoseconds: 100_000_000) } + if lastFinal.isEmpty, let asrTask { + FlowDiagnostics.log("ASR wait elapsed — awaiting asrTask completion") + _ = await asrTask.value + } + + let asrElapsed = Date().timeIntervalSince(pipelineStarted) + FlowDiagnostics.log("ASR phase done in \(String(format: "%.1f", asrElapsed))s finalLen=\(lastFinal.count)") + var text = lastFinal.trimmingCharacters(in: .whitespacesAndNewlines) if text.isEmpty { text = currentPartial.trimmingCharacters(in: .whitespacesAndNewlines) @@ -390,36 +538,72 @@ final class FlowSessionManager: ObservableObject { let key = (asrTask?.isCancelled == true) ? "flow.error.recognitionInterrupted" : "flow.error.noSpeech" + FlowDiagnostics.log("finalize failed: empty transcript after \(String(format: "%.1f", asrElapsed))s") FlowSessionBridge.storeTranscriptionError( - NSLocalizedString(key, comment: "") + AppL10n.string(key) ) return } let engineMode = store.engineMode - let modeId = store.modeId - let shouldPolish = engineMode != "local" && modeId == "polish" + + if engineMode == "local" { + let warning = Self.chunkWarningMessage(chunkWarnings) + FlowSessionBridge.storeTranscriptionResult(text, polishWarning: warning) + FlowDiagnostics.log( + "finalize ASR-only total=\(String(format: "%.1f", Date().timeIntervalSince(pipelineStarted)))s " + + "len=\(text.count)" + ) + SpeechHistoryStore.shared.append(text: text, engineMode: engineMode) + currentPartial = "" + lastFinal = "" + chunkWarnings = [] + debug("utterance finalized length=\(text.count)") + return + } var delivered = text - if shouldPolish { - do { - let polished = try await polisher.polish(text) - delivered = polished - FlowSessionBridge.storeTranscriptionResult(polished) - } catch { - FlowSessionBridge.storeTranscriptionResult(text) - } - } else { - FlowSessionBridge.storeTranscriptionResult(text) + let chunkNote = Self.chunkWarningMessage(chunkWarnings) + let polishStarted = Date() + do { + let polished = try await polisher.polish(text) + delivered = polished + FlowSessionBridge.storeTranscriptionResult(polished, polishWarning: chunkNote) + FlowDiagnostics.log( + "polish done in \(String(format: "%.1f", Date().timeIntervalSince(polishStarted)))s " + + "total=\(String(format: "%.1f", Date().timeIntervalSince(pipelineStarted)))s" + ) + } catch { + FlowDiagnostics.log( + "polish failed after \(String(format: "%.1f", Date().timeIntervalSince(polishStarted)))s: " + + "\(error.localizedDescription)" + ) + FlowSessionBridge.storeTranscriptionResult(text, polishWarning: chunkNote) } SpeechHistoryStore.shared.append(text: delivered, engineMode: engineMode) currentPartial = "" lastFinal = "" + chunkWarnings = [] + chunkedPipeline = nil debug("utterance finalized length=\(text.count)") } + private static func chunkWarningMessage(_ warnings: [String]) -> String? { + guard !warnings.isEmpty else { return nil } + return warnings.joined(separator: "\n") + } + + private func asrWaitTimeout() -> TimeInterval { + if store.engineMode == "local" { + return store.localASRBackend == .qwen3ASR + ? FlowSessionKeys.localQwen3ASRWaitTimeout + : FlowSessionKeys.localASRWaitTimeout + } + return FlowSessionKeys.cloudASRWaitTimeout + } + // MARK: - Level publishing (main thread only) private func startLevelPublishing() { @@ -461,8 +645,6 @@ final class FlowSessionManager: ObservableObject { } private func debug(_ message: String) { - #if DEBUG - print("🌊[FlowSession] \(message)") - #endif + FlowDiagnostics.log(message) } } diff --git a/OSGKeyboard/Services/LegalLinks.swift b/OSGKeyboard/Services/LegalLinks.swift index 0c9136e..2f751cb 100644 --- a/OSGKeyboard/Services/LegalLinks.swift +++ b/OSGKeyboard/Services/LegalLinks.swift @@ -4,7 +4,9 @@ import Foundation enum LegalLinks { - /// Public privacy policy (GitHub Pages). + static let repositoryURL = URL(string: "https://github.com/hkgood/OSGKeyboard")! + + /// Public privacy policy (GitHub Pages). Also bundled in-app as PrivacyPolicy.html. static var privacyPolicyURL: URL? { URL(string: "https://hkgood.github.io/OSGKeyboard/privacy/") } diff --git a/OSGKeyboard/Services/ModelDownloadSourcePicker.swift b/OSGKeyboard/Services/ModelDownloadSourcePicker.swift new file mode 100644 index 0000000..9cecae5 --- /dev/null +++ b/OSGKeyboard/Services/ModelDownloadSourcePicker.swift @@ -0,0 +1,126 @@ +// ModelDownloadSourcePicker.swift +// OSGKeyboard · Main App +// +// Picks ModelScope vs Hugging Face by probing both mirrors on the +// user's current network. Result is cached briefly so consecutive +// downloads in one session don't re-probe. + +import Foundation +import os + +enum ModelDownloadSourcePicker { + + private struct CacheState { + var source: ModelDownloadSource? + var expiresAt: Date? + } + + private static let lock = OSAllocatedUnfairLock(initialState: CacheState()) + private static let cacheTTL: TimeInterval = 300 + + /// Resolves the fastest reachable mirror for the current network. + static func resolve() async -> ModelDownloadSource { + if let cached = cachedValue() { return cached } + + let winner = await probeFastest() ?? defaultHeuristic() + storeCache(winner) + return winner + } + + /// Alternate mirror — used when the first download attempt fails. + static func alternate(to source: ModelDownloadSource) -> ModelDownloadSource { + switch source { + case .modelScope: return .huggingface + case .huggingface: return .modelScope + } + } + + // MARK: - Probe + + private static func probeFastest() async -> ModelDownloadSource? { + await withTaskGroup(of: (ModelDownloadSource, TimeInterval)?.self) { group in + for source in ModelDownloadSource.allCases { + group.addTask { + guard let latency = await probeLatency(for: source) else { return nil } + return (source, latency) + } + } + + var best: (ModelDownloadSource, TimeInterval)? + for await candidate in group { + guard let candidate else { continue } + if best == nil || candidate.1 < best!.1 { + best = candidate + } + } + return best?.0 + } + } + + private static func probeLatency(for source: ModelDownloadSource) async -> TimeInterval? { + guard let url = probeURL(for: source) else { return nil } + + var request = URLRequest(url: url) + request.httpMethod = "HEAD" + request.timeoutInterval = 4 + request.cachePolicy = .reloadIgnoringLocalCacheData + + let started = CFAbsoluteTimeGetCurrent() + do { + let (_, response) = try await URLSession.shared.data(for: request) + guard let http = response as? HTTPURLResponse else { return nil } + guard (200...399).contains(http.statusCode) else { return nil } + return CFAbsoluteTimeGetCurrent() - started + } catch { + // Some hosts reject HEAD — retry with a tiny GET. + var get = URLRequest(url: url) + get.httpMethod = "GET" + get.timeoutInterval = 4 + get.cachePolicy = .reloadIgnoringLocalCacheData + do { + let (_, response) = try await URLSession.shared.data(for: get) + guard let http = response as? HTTPURLResponse else { return nil } + guard (200...399).contains(http.statusCode) else { return nil } + return CFAbsoluteTimeGetCurrent() - started + } catch { + return nil + } + } + } + + private static func probeURL(for source: ModelDownloadSource) -> URL? { + switch source { + case .modelScope: + return URL(string: "https://modelscope.cn") + case .huggingface: + return URL(string: "https://huggingface.co") + } + } + + /// When both probes fail (offline, captive portal, etc.). + private static func defaultHeuristic() -> ModelDownloadSource { + if Locale.current.region?.identifier == "CN" { return .modelScope } + if TimeZone.current.identifier.hasPrefix("Asia/Shanghai") { return .modelScope } + return .huggingface + } + + // MARK: - Cache + + private static func cachedValue() -> ModelDownloadSource? { + lock.withLock { state in + guard let source = state.source, + let expiresAt = state.expiresAt, + expiresAt > Date() else { + return nil + } + return source + } + } + + private static func storeCache(_ source: ModelDownloadSource) { + lock.withLock { state in + state.source = source + state.expiresAt = Date().addingTimeInterval(cacheTTL) + } + } +} diff --git a/OSGKeyboard/Services/ModelManager.swift b/OSGKeyboard/Services/ModelManager.swift new file mode 100644 index 0000000..aa891c0 --- /dev/null +++ b/OSGKeyboard/Services/ModelManager.swift @@ -0,0 +1,492 @@ +// ModelManager.swift +// OSGKeyboard · Main App +// +// Owns the lifecycle of on-device ML models that back the local +// ASR backend (Qwen3-ASR-0.6B CoreML, ~1.6 GB). +// +// `runDownload` fetches CoreML bundles + tokenizer files — it does +// not load models into memory (warm-up happens in `OnDeviceModelWarmup`). +// Weights land under `~/Library/Caches/qwen3-speech/` using the Hub +// layout from `HuggingFaceDownloader`. +// +// Why this lives in the host app: the ASR model is loaded via +// soniqo/speech-swift, which is only linked into the main App +// target (Qwen3Speech pulls mlx-swift as a transitive dependency). +// +// Mirror selection: resolved automatically at download time via +// `ModelDownloadSourcePicker` (latency probe + locale fallback). + +import Foundation +import SwiftUI +import OSGKeyboardShared +import Qwen3ASR + +private enum Qwen3CoreMLDownloadArtifacts { + static let coreMLBundleGlobs = [ + "encoder.mlmodelc/**", + "embedding.mlmodelc/**", + "decoder_part1.mlmodelc/**", + "decoder_part2.mlmodelc/**", + "config.json", + ] + + static let tokenizerFiles = [ + "vocab.json", + "merges.txt", + "tokenizer_config.json", + ] +} + +/// Where on-device model weights are downloaded from. +enum ModelDownloadSource: String, CaseIterable, Identifiable, Sendable { + case huggingface + case modelScope + + var id: String { rawValue } + + var registry: ModelRegistry { + switch self { + case .huggingface: return .huggingFace() + case .modelScope: return .modelScope() + } + } + + /// Host shown in error messages. + var hostLabel: String { + switch self { + case .huggingface: return "huggingface.co" + case .modelScope: return "modelscope.cn" + } + } +} + +enum ModelDownloadState: Equatable, Sendable { + case notDownloaded + case downloading(progress: Double) + case downloaded + case failed(String) + + var isTerminal: Bool { + switch self { + case .downloaded, .failed: return true + case .notDownloaded, .downloading: return false + } + } + + var downloadProgress: Double? { + if case .downloading(let progress) = self { return progress } + return nil + } +} + +/// Per-model state tracked by `ModelManager`. The manager keeps a +/// dictionary of these and re-emits it on the main actor whenever +/// any field changes. +struct ModelState: Equatable, Sendable { + var download: ModelDownloadState + var lastError: String? +} + +/// Observable holder that the Settings UI binds to. All mutating +/// methods dispatch onto the main actor so SwiftUI views can +/// observe without ceremony. +@MainActor +final class ModelManager: ObservableObject { + + static let shared = ModelManager() + + @Published private(set) var states: [OnDeviceModel: ModelState] = [:] + @Published private(set) var activeDownloads: Set = [] + + private var downloadTasks: [OnDeviceModel: Task] = [:] + + init() { + for model in OnDeviceModel.allCases { + states[model] = ModelState(download: .notDownloaded, lastError: nil) + } + refreshAll() + } + + // MARK: - Queries + + /// Synchronous check on whether the model is already on disk. + /// Used by the UI to decide whether to show "Download" or + /// "Delete". Doesn't touch the network. + func isDownloaded(_ model: OnDeviceModel) -> Bool { + Self.weightsOnDisk(for: model) + } + + /// Disk-only probe safe to call from background ASR tasks. + /// Returns `false` until the user downloads via Settings. + nonisolated static func weightsOnDisk(for model: OnDeviceModel) -> Bool { + existingCacheDirectory(for: model) != nil + } + + /// Approximate on-disk bytes used by the model directory. Used + /// by the Settings "Storage" badge. + func onDiskBytes(_ model: OnDeviceModel) -> Int64 { + guard let dir = Self.existingCacheDirectory(for: model) else { return 0 } + guard let enumerator = FileManager.default.enumerator( + at: dir, + includingPropertiesForKeys: [.totalFileAllocatedSizeKey, .isRegularFileKey] + ) else { return 0 } + var total: Int64 = 0 + for case let url as URL in enumerator { + let values = try? url.resourceValues(forKeys: [.totalFileAllocatedSizeKey, .isRegularFileKey]) + if values?.isRegularFile == true { + total += Int64(values?.totalFileAllocatedSize ?? 0) + } + } + return total + } + + // MARK: - Mutations + + /// Triggers a background download of the model. The call returns + /// immediately; observe `states[model].download` for progress. + /// Calling this while a download is in progress is a no-op. + func startDownload(_ model: OnDeviceModel) { + if activeDownloads.contains(model) { return } + if isDownloaded(model) { + states[model]?.download = .downloaded + return + } + activeDownloads.insert(model) + states[model] = ModelState(download: .downloading(progress: 0), lastError: nil) + publishStatusToAppGroup() + + let task = Task.detached(priority: .userInitiated) { [weak self] in + guard let self else { return } + let primary = await ModelDownloadSourcePicker.resolve() + do { + try await self.runDownload(model, registry: primary.registry) + } catch is CancellationError { + await self.finishDownloadCancelled(model) + } catch { + let fallback = ModelDownloadSourcePicker.alternate(to: primary) + await self.reportDownloadProgress(model, fraction: 0, monotonic: false) + do { + try await self.runDownload(model, registry: fallback.registry) + } catch is CancellationError { + await self.finishDownloadCancelled(model) + } catch { + await self.finishDownloadFailed(model, error: error) + } + } + } + downloadTasks[model] = task + } + + func cancelDownload(_ model: OnDeviceModel) { + downloadTasks[model]?.cancel() + downloadTasks[model] = nil + activeDownloads.remove(model) + states[model] = ModelState(download: .notDownloaded, lastError: nil) + publishStatusToAppGroup() + } + + func deleteModel(_ model: OnDeviceModel) { + for dir in Self.candidateCacheDirectories(for: model) { + try? FileManager.default.removeItem(at: dir) + } + states[model] = ModelState(download: .notDownloaded, lastError: nil) + publishStatusToAppGroup() + OnDeviceModelWarmup.shared.invalidate() + } + + /// Cheap refresh that re-reads the on-disk state for every + /// tracked model. Called from `init` and after a successful + /// download so the Settings row updates from "Downloading…" to + /// "Downloaded · 1.4 GB" without needing a separate notifier. + func refreshAll() { + for model in OnDeviceModel.allCases { + if activeDownloads.contains(model) { continue } + if isDownloaded(model) { + states[model] = ModelState(download: .downloaded, lastError: nil) + } else if case .failed = states[model]?.download { + // Preserve any existing failure message so the UI + // can show "Download failed: " instead of + // resetting it back to "Not downloaded" every + // refresh. + continue + } else { + states[model] = ModelState(download: .notDownloaded, lastError: states[model]?.lastError) + } + } + publishStatusToAppGroup() + } + + /// Mirror disk/download state into the App Group for the keyboard + /// extension, which cannot probe the host app's Caches folder. + private func publishStatusToAppGroup() { + for model in OnDeviceModel.allCases { + let downloaded: Bool + let progress: Double? + switch states[model]?.download { + case .downloaded: + downloaded = true + progress = nil + case .downloading(let fraction): + downloaded = false + progress = fraction + case .failed, .notDownloaded, .none: + downloaded = isDownloaded(model) + progress = nil + } + OnDeviceModelStatus.setDownloaded(downloaded, for: model) + OnDeviceModelStatus.setProgress(progress, for: model) + } + scheduleWarmupIfNeeded() + } + + private func scheduleWarmupIfNeeded() { + let config = ProviderConfig.shared + guard config.isLocalEngine else { + OnDeviceModelWarmup.shared.invalidate() + return + } + if OnDeviceModelStatus.isLocalStackReady(asrBackend: config.localASRBackend) { + // Do not force-restart an in-flight warm-up — `publishStatusToAppGroup` + // runs on download progress ticks and would otherwise cancel load + // mid-flight, leaving the UI stuck on "warming". + OnDeviceModelWarmup.shared.warmUpIfNeeded() + } else { + OnDeviceModelWarmup.shared.invalidate() + } + } + + // MARK: - Internals + + /// Runs off the main actor; updates `@Published` state via `MainActor.run`. + /// Throws on failure so `startDownload` can fall back to the alternate mirror. + nonisolated private func runDownload(_ model: OnDeviceModel, registry: ModelRegistry) async throws { + switch model { + case .qwen3ASR: + try await Self.downloadQwen3CoreMLWeights( + model: model, + registry: registry, + progressHandler: { @Sendable [weak self] fraction, _ in + Task { @MainActor [weak self] in + guard let self else { return } + self.reportDownloadProgress(model, fraction: fraction) + } + } + ) + } + + await MainActor.run { [weak self] in + guard let self else { return } + self.activeDownloads.remove(model) + self.downloadTasks[model] = nil + self.states[model] = ModelState(download: .downloaded, lastError: nil) + self.publishStatusToAppGroup() + let config = ProviderConfig.shared + if config.isLocalEngine, + OnDeviceModelStatus.isLocalStackReady(asrBackend: config.localASRBackend) { + // Retry warm-up after a prior load failure once weights land on disk. + OnDeviceModelWarmup.shared.warmUpIfNeeded(force: true) + } + } + } + + nonisolated private func finishDownloadCancelled(_ model: OnDeviceModel) async { + await MainActor.run { [weak self] in + guard let self else { return } + self.activeDownloads.remove(model) + self.downloadTasks[model] = nil + self.states[model] = ModelState(download: .notDownloaded, lastError: nil) + self.publishStatusToAppGroup() + } + } + + nonisolated private func finishDownloadFailed(_ model: OnDeviceModel, error: Error) async { + let message = Self.userFacingDownloadError(error) + await MainActor.run { [weak self] in + guard let self else { return } + self.activeDownloads.remove(model) + self.downloadTasks[model] = nil + self.states[model] = ModelState(download: .failed(message), lastError: message) + self.publishStatusToAppGroup() + } + } + + /// Updates UI progress. By default keeps the bar monotonic so brief + /// per-file jumps inside the downloader never move backwards. + private func reportDownloadProgress( + _ model: OnDeviceModel, + fraction: Double, + monotonic: Bool = true + ) { + let clamped = min(max(fraction, 0), 1) + let previous = states[model]?.download.downloadProgress ?? 0 + let value = monotonic ? max(previous, clamped) : clamped + states[model] = ModelState(download: .downloading(progress: value), lastError: nil) + publishStatusToAppGroup() + } + + /// Short, user-readable download failure text for Settings UI. + nonisolated private static func userFacingDownloadError(_ error: Error) -> String { + let raw = (error as? LocalizedError)?.errorDescription ?? error.localizedDescription + if raw.localizedCaseInsensitiveContains("metadata") { + return AppL10n.string("settings.models.error.metadata") + } + if raw.localizedCaseInsensitiveContains("offline mode") { + return AppL10n.string("settings.models.error.offline") + } + if raw.count > 280 { + return String(raw.prefix(277)) + "…" + } + return raw + } + + /// Resolve on-disk cache directories for a model. Matches the layout + /// `HuggingFaceDownloader.getCacheDirectory(for:)` uses in Qwen3Speech. + nonisolated static func candidateCacheDirectories(for model: OnDeviceModel) -> [URL] { + let base = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first! + .appendingPathComponent("qwen3-speech", isDirectory: true) + let repoId = model.repoId + var candidates: [URL] = [] + + // Hub-style path (current default). + let parts = repoId.split(separator: "/", omittingEmptySubsequences: true) + if parts.count == 2 { + candidates.append( + base + .appendingPathComponent("models/\(parts[0])/\(parts[1])", isDirectory: true) + ) + } + + // Legacy flat path kept by HuggingFaceDownloader for backward compat. + let sanitized = repoId.replacingOccurrences(of: "/", with: "_") + candidates.append(base.appendingPathComponent(sanitized, isDirectory: true)) + + // Older OSGKeyboard probe paths (pre-alignment); delete still sweeps these. + switch model { + case .qwen3ASR: + candidates.append(base.appendingPathComponent("Qwen3ASR", isDirectory: true)) + candidates.append( + base.appendingPathComponent("models/aufklarer/Qwen3-ASR-0.6B-MLX-4bit", isDirectory: true) + ) + candidates.append(base.appendingPathComponent("aufklarer_Qwen3-ASR-0.6B-MLX-4bit", isDirectory: true)) + } + + return candidates + } + + /// First candidate directory that already contains downloaded weights. + nonisolated static func existingCacheDirectory(for model: OnDeviceModel) -> URL? { + candidateCacheDirectories(for: model).first { dir in + weightsExist(in: dir, model: model) + } + } + + nonisolated private static func weightsExist(in directory: URL, model: OnDeviceModel) -> Bool { + let fm = FileManager.default + switch model { + case .qwen3ASR: + let encoder = directory.appendingPathComponent("encoder.mlmodelc", isDirectory: true) + let decoder = directory.appendingPathComponent("decoder_part1.mlmodelc", isDirectory: true) + let vocab = directory.appendingPathComponent("vocab.json") + return fm.fileExists(atPath: encoder.path) + && fm.fileExists(atPath: decoder.path) + && fm.fileExists(atPath: vocab.path) + } + } + + // MARK: - CoreML download + + /// Downloads CoreML encoder/decoder bundles and tokenizer files into one cache dir. + nonisolated static func downloadQwen3CoreMLWeights( + model: OnDeviceModel, + registry: ModelRegistry, + progressHandler: @escaping @Sendable (Double, String) -> Void + ) async throws { + let coreMLId = model.repoId + let tokenizerId = model.tokenizerRepoId + let dir = try HuggingFaceDownloader.getCacheDirectory(for: coreMLId) + + switch registry { + case .huggingFace(let hubEndpoint): + try await HuggingFaceDownloader.downloadWeights( + modelId: coreMLId, + to: dir, + additionalFiles: Qwen3CoreMLDownloadArtifacts.coreMLBundleGlobs, + hubEndpoint: hubEndpoint, + progressHandler: { progressHandler($0 * 0.85, "CoreML") } + ) + try await HuggingFaceDownloader.downloadWeights( + modelId: tokenizerId, + to: dir, + additionalFiles: Qwen3CoreMLDownloadArtifacts.tokenizerFiles, + hubEndpoint: hubEndpoint, + progressHandler: { progressHandler(0.85 + $0 * 0.15, "Tokenizer") } + ) + case .modelScope(let baseURL, let revision): + try await downloadQwen3CoreMLViaModelScope( + coreMLId: coreMLId, + tokenizerId: tokenizerId, + to: dir, + baseURL: baseURL, + revision: revision, + progressHandler: progressHandler + ) + } + progressHandler(1.0, "Ready") + } + + nonisolated private static func downloadQwen3CoreMLViaModelScope( + coreMLId: String, + tokenizerId: String, + to directory: URL, + baseURL: String, + revision: String, + progressHandler: @escaping @Sendable (Double, String) -> Void + ) async throws { + let coreListed = try await ModelScopeDownloader.listAllFiles( + modelId: coreMLId, + baseURL: baseURL, + revision: revision + ) + let corePaths = coreListed.map(\.path).filter { path in + path.contains(".mlmodelc/") || path == "config.json" + } + guard !corePaths.isEmpty else { + throw DownloadError.failedToDownload("\(coreMLId): no CoreML files on ModelScope") + } + let coreSizes = Dictionary(uniqueKeysWithValues: coreListed.map { ($0.path, $0.size) }) + try await ModelScopeDownloader.downloadFiles( + modelId: coreMLId, + to: directory, + files: corePaths, + fileSizes: coreSizes, + baseURL: baseURL, + revision: revision, + progressHandler: { progressHandler($0 * 0.85, "CoreML") } + ) + + let tokListed = try await ModelScopeDownloader.listAllFiles( + modelId: tokenizerId, + baseURL: baseURL, + revision: revision + ) + let tokPaths = Qwen3CoreMLDownloadArtifacts.tokenizerFiles.filter { name in + tokListed.contains { $0.path == name } + } + let tokSizes = Dictionary(uniqueKeysWithValues: tokListed.map { ($0.path, $0.size) }) + try await ModelScopeDownloader.downloadFiles( + modelId: tokenizerId, + to: directory, + files: tokPaths, + fileSizes: tokSizes, + baseURL: baseURL, + revision: revision, + progressHandler: { progressHandler(0.85 + $0 * 0.15, "Tokenizer") } + ) + } + + /// Preferred cache directory for display / storage badges. + nonisolated static func cacheDirectory(for model: OnDeviceModel) -> URL { + existingCacheDirectory(for: model) + ?? candidateCacheDirectories(for: model).first! + } +} diff --git a/OSGKeyboard/Services/OnDeviceModelWarmup.swift b/OSGKeyboard/Services/OnDeviceModelWarmup.swift new file mode 100644 index 0000000..48b93b9 --- /dev/null +++ b/OSGKeyboard/Services/OnDeviceModelWarmup.swift @@ -0,0 +1,197 @@ +// OnDeviceModelWarmup.swift +// OSGKeyboard · Main App +// +// Preloads on-device ASR weights for Flow sessions. + +import Foundation +import OSGKeyboardShared + +@MainActor +final class OnDeviceModelWarmup: ObservableObject { + + static let shared = OnDeviceModelWarmup() + + enum Phase: Equatable { + case idle + case warming + case ready + case failed(String) + case notNeeded + + var isFailed: Bool { + if case .failed = self { return true } + return false + } + } + + @Published private(set) var phase: Phase = .idle + + /// Bumped on `invalidate()` and each new warm-up so cancelled tasks + /// cannot leave `phase` stuck on `.warming`. + private var warmupGeneration = 0 + private var warmupTask: Task? + private var qwenASRService: Qwen3ASRService? + private var speechAnalyzerService: ASRService? + private var cloudASRService: ASRService? + + private init() {} + + /// Loads ASR into memory when the local stack is ready on disk. + func warmUpIfNeeded(force: Bool = false) { + let store = AppGroupStore() + guard store.engineMode == "local" else { + resetInstances() + phase = .notNeeded + publishMemoryReady(false) + return + } + + guard store.localASRBackend != .qwen3ASR || OnDeviceMLRuntime.supportsOnDeviceQwen3 else { + resetInstances() + phase = .notNeeded + publishMemoryReady(false) + return + } + + guard OnDeviceModelStatus.isLocalStackReady(asrBackend: store.localASRBackend) else { + resetInstances() + phase = .idle + publishMemoryReady(false) + return + } + + var shouldForce = force + if phase == .ready, !shouldForce { + if needsModelReload() { + shouldForce = true + } else { + publishMemoryReady(true) + return + } + } + if phase == .warming { return } + if case .failed = phase, !shouldForce { return } + + warmupTask?.cancel() + warmupGeneration += 1 + let generation = warmupGeneration + phase = .warming + publishMemoryReady(false) + + let asrBackend = store.localASRBackend + warmupTask = Task { @MainActor [weak self] in + guard let self else { return } + do { + try await self.performWarmup(asrBackend: asrBackend) + guard generation == self.warmupGeneration, !Task.isCancelled else { return } + self.phase = .ready + self.publishMemoryReady(true) + } catch { + guard generation == self.warmupGeneration, !Task.isCancelled else { return } + let message = (error as? LocalizedError)?.errorDescription + ?? error.localizedDescription + self.phase = .failed(message) + self.publishMemoryReady(false) + } + } + } + + func invalidate() { + warmupGeneration += 1 + warmupTask?.cancel() + warmupTask = nil + resetInstances() + phase = .idle + publishMemoryReady(false) + } + + /// Called when returning from background — re-verify CoreML weights and + /// unstick a warmup that was frozen while the app was suspended. + func ensureReadyAfterBackground() { + let store = AppGroupStore() + guard store.engineMode == "local" else { + phase = .notNeeded + publishMemoryReady(false) + return + } + + guard OnDeviceModelStatus.isLocalStackReady(asrBackend: store.localASRBackend) else { + phase = .idle + publishMemoryReady(false) + return + } + + switch phase { + case .warming, .ready: + if needsModelReload() { + warmUpIfNeeded(force: true) + } + case .failed, .idle: + warmUpIfNeeded(force: true) + case .notNeeded: + break + } + } + + func asrService(engineMode: String, localBackend: LocalASRBackend) -> ASRService { + if engineMode != "local" { + if cloudASRService == nil { + cloudASRService = ASRServiceFactory.make( + engineMode: engineMode, + localBackend: localBackend + ) + } + return cloudASRService! + } + + switch localBackend { + case .qwen3ASR: + if qwenASRService == nil { + qwenASRService = Qwen3ASRService() + } + return qwenASRService! + case .speechAnalyzer: + if speechAnalyzerService == nil { + speechAnalyzerService = ASRServiceFactory.make( + engineMode: "local", + localBackend: .speechAnalyzer + ) + } + return speechAnalyzerService! + } + } + + // MARK: - Internals + + private func performWarmup(asrBackend: LocalASRBackend) async throws { + switch asrBackend { + case .qwen3ASR: + if qwenASRService == nil { + qwenASRService = Qwen3ASRService() + } + FlowDiagnostics.log("warmup ASR start backend=qwen3ASR") + try await qwenASRService!.warmUp() + FlowDiagnostics.log("warmup ASR done") + case .speechAnalyzer: + FlowDiagnostics.log("warmup skipped — speechAnalyzer backend") + } + } + + private func resetInstances() { + qwenASRService = nil + speechAnalyzerService = nil + cloudASRService = nil + } + + private func publishMemoryReady(_ ready: Bool) { + OnDeviceModelStatus.setModelsLoadedInMemory(ready) + } + + private func needsModelReload() -> Bool { + let store = AppGroupStore() + guard store.engineMode == "local", store.localASRBackend == .qwen3ASR else { + return false + } + return qwenASRService?.isModelInMemory != true + } +} diff --git a/OSGKeyboard/Services/OpenSourceLicenseCatalog.swift b/OSGKeyboard/Services/OpenSourceLicenseCatalog.swift new file mode 100644 index 0000000..33be0ac --- /dev/null +++ b/OSGKeyboard/Services/OpenSourceLicenseCatalog.swift @@ -0,0 +1,115 @@ +// OpenSourceLicenseCatalog.swift +// OSGKeyboard · Main App +// +// Single source of truth for third-party open-source components shipped +// with or downloaded by OSGKeyboard. Consumed by Settings → About → +// Third-Party Licenses. +// +// Keep this list aligned with `project.yml` package dependencies and +// the default model ID in `Qwen3ASRService`. + +import Foundation + +enum OpenSourceLicenseCatalog { + + struct Entry: Identifiable, Hashable { + let id: String + let name: String + let licenseName: String + /// One-line explanation shown in the popup and above the full text. + let purpose: String + let url: URL? + /// Verbatim license body for the long-scroll disclosure page. + let licenseText: String + } + + /// Bundled libraries and runtime model artefacts referenced by the app. + static let entries: [Entry] = [ + .init( + id: "speech-swift", + name: "soniqo/speech-swift", + licenseName: "Apache-2.0", + purpose: "On-device ASR runtime. Vendored locally as the Qwen3Speech SPM package (Qwen3ASR CoreML path, AudioCommon, SpeechVAD).", + url: URL(string: "https://github.com/soniqo/speech-swift"), + licenseText: apache2Text + ), + .init( + id: "swift-transformers", + name: "huggingface/swift-transformers", + licenseName: "Apache-2.0", + purpose: "Hugging Face Hub client and tokenizer bindings. Used to resolve and download model snapshots at runtime.", + url: URL(string: "https://github.com/huggingface/swift-transformers"), + licenseText: apache2Text + ), + .init( + id: "qwen3-asr-coreml", + name: "aufklarer/Qwen3-ASR-CoreML", + licenseName: "Apache-2.0", + purpose: "CoreML INT8 weights for Qwen3-ASR-0.6B (derived from Alibaba Qwen team). Downloaded on first use (~1.6 GB); not bundled in the app binary.", + url: URL(string: "https://huggingface.co/aufklarer/Qwen3-ASR-CoreML"), + licenseText: apache2Text + ), + .init( + id: "qwen3-asr-upstream", + name: "Qwen/Qwen3-ASR-0.6B", + licenseName: "Apache-2.0", + purpose: "Original ASR model by Alibaba's Qwen team. CoreML bundle and tokenizer files are derived from these weights.", + url: URL(string: "https://huggingface.co/Qwen/Qwen3-ASR-0.6B"), + licenseText: apache2Text + ), + .init( + id: "material-icons", + name: "Google Material Icons", + licenseName: "Apache-2.0", + purpose: "MaterialIcons-Regular.ttf bundled for Settings and navigation iconography.", + url: URL(string: "https://github.com/google/material-design-icons"), + licenseText: apache2Text + ), + ] + + // MARK: - License bodies + + static let apache2Text = """ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied. See the License for the specific language governing + permissions and limitations under the License. + """ + + static let mitText = """ + MIT License + + Copyright (c) 2023 Apple Inc. + + Permission is hereby granted, free of charge, to any person + obtaining a copy of this software and associated documentation + files (the "Software"), to deal in the Software without + restriction, including without limitation the rights to use, copy, + modify, merge, publish, distribute, sublicense, and/or sell copies + of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be + included in all copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, + EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES + OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND + NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT + HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, + WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR + OTHER DEALINGS IN THE SOFTWARE. + """ +} diff --git a/OSGKeyboard/Services/Qwen3ASRService.swift b/OSGKeyboard/Services/Qwen3ASRService.swift new file mode 100644 index 0000000..a3f5cfb --- /dev/null +++ b/OSGKeyboard/Services/Qwen3ASRService.swift @@ -0,0 +1,257 @@ +// Qwen3ASRService.swift +// OSGKeyboard · Main App +// +// On-device ASR via Qwen3-ASR-0.6B CoreML (Neural Engine + CPU). Uses the +// MLX-free `transcribeBackgroundSafe` path so Flow dictation works while the +// host app is backgrounded (no Metal GPU). + +import Foundation +import AVFoundation +import os +import OSGKeyboardShared +@preconcurrency import Qwen3ASR + +struct Qwen3ASRServiceProvider: ASRServiceProvider { + let backend: LocalASRBackend = .qwen3ASR + func make() -> ASRService { Qwen3ASRService() } +} + +final class Qwen3ASRService: ASRService, @unchecked Sendable { + + private enum TranscribeConstants { + static let sampleRate = 16_000 + /// CoreML encoder is exported for 30 s windows — align with Flow chunking. + static let chunkDurationSeconds = Int(FlowUtteranceChunkConfig.flowDefault.maxChunkDurationSeconds) + } + + private let lock = OSAllocatedUnfairLock() + private var currentTask: Task? + private var cancelled = false + + private var model: CoreMLASRModel? + private var loadError: Error? + private var loadingTask: Task? + + private func resolveModel() async throws -> CoreMLASRModel { + if let model = lock.withLock({ self.model }) { return model } + if let err = lock.withLock({ self.loadError }) { throw err } + + guard ModelManager.weightsOnDisk(for: .qwen3ASR) else { + throw ASRServiceError.modelNotDownloaded + } + + guard OnDeviceMLRuntime.supportsOnDeviceQwen3 else { + throw ASRServiceError.unsupportedOS + } + + let task: Task = lock.withLock { + if let existing = loadingTask { return existing } + let new = Task { [weak self] in + guard let self else { throw ASRServiceError.notReady } + let cacheDir = ModelManager.cacheDirectory(for: .qwen3ASR) + let loaded = try await CoreMLASRModel.fromPretrained( + tokenizerModelId: OnDeviceModel.qwen3ASR.tokenizerRepoId, + cacheDir: cacheDir, + offlineMode: true, + progressHandler: { @Sendable _, _ in } + ) + try loaded.warmUp() + self.lock.withLock { self.model = loaded } + return loaded + } + loadingTask = new + return new + } + do { + let model = try await task.value + return model + } catch { + lock.withLock { self.loadError = error } + throw error + } + } + + func warmUp() async throws { + FlowDiagnostics.log("Qwen3ASR CoreML warmUp start") + resetForNewUtterance() + _ = try await resolveModel() + FlowDiagnostics.log("Qwen3ASR CoreML warmUp done") + } + + func resetForNewUtterance() { + lock.withLock { cancelled = false } + } + + var isModelInMemory: Bool { + lock.withLock { model != nil } + } + + func transcribe( + stream: AsyncStream, + locale: Locale + ) -> AsyncStream { + AsyncStream { continuation in + continuation.yield(.capability(onDeviceSupported: true)) + + let task = Task { [weak self] in + guard let self else { return } + defer { self.lock.withLock { self.currentTask = nil } } + + var samples: [Float] = [] + samples.reserveCapacity( + Int(Double(TranscribeConstants.sampleRate) * FlowSessionKeys.maxUtteranceDuration) + 16_000 + ) + for await snap in stream { + if Task.isCancelled || self.cancelledNow() { break } + samples.append(contentsOf: snap.samples) + } + guard !Task.isCancelled, !self.cancelledNow() else { + continuation.finish() + return + } + if samples.isEmpty { + continuation.yield(.error(SharedL10n.string("error.asr.noSpeech"))) + continuation.finish() + return + } + + do { + let model = try await self.resolveModel() + let language = Self.languageHint(from: locale) + let durationSec = Double(samples.count) / 16_000.0 + FlowDiagnostics.log( + "Qwen3ASR CoreML transcribe start samples=\(samples.count) " + + "duration=\(String(format: "%.1f", durationSec))s" + ) + let text = self.transcribeInChunks( + model: model, + samples: samples, + language: language + ) + FlowDiagnostics.log("Qwen3ASR CoreML transcribe done chars=\(text.count)") + let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) + if trimmed.isEmpty { + continuation.yield(.error(SharedL10n.string("error.asr.noSpeech"))) + } else { + continuation.yield(.final(trimmed)) + } + continuation.finish() + } catch { + Self.debug("Qwen3ASR.transcribe failed: \(error.localizedDescription)") + continuation.yield(.error(error.localizedDescription)) + continuation.finish() + } + } + self.lock.withLock { self.currentTask = task } + + continuation.onTermination = { @Sendable [weak self] _ in + self?.cancel() + } + } + } + + func cancel() { + lock.withLock { + cancelled = true + currentTask?.cancel() + currentTask = nil + } + } + + func transcribeChunk(samples: [Float], locale: Locale) async -> ASRChunkResult { + if cancelledNow() || Task.isCancelled { return .cancelled } + guard !samples.isEmpty else { return .success("") } + + do { + let model = try await resolveModel() + let language = Self.languageHint(from: locale) + let text = model.transcribeBackgroundSafe( + audio: samples, + sampleRate: TranscribeConstants.sampleRate, + language: language + ) + .trimmingCharacters(in: .whitespacesAndNewlines) + if text.hasPrefix("[CoreML error:") { + return .failure(text) + } + return .success(text) + } catch { + let message = (error as? LocalizedError)?.errorDescription ?? error.localizedDescription + return .failure(message) + } + } + + private func cancelledNow() -> Bool { + lock.withLock { cancelled } + } + + private func transcribeInChunks( + model: CoreMLASRModel, + samples: [Float], + language: String? + ) -> String { + let chunkSize = TranscribeConstants.sampleRate * TranscribeConstants.chunkDurationSeconds + guard samples.count > chunkSize else { + return model.transcribeBackgroundSafe( + audio: samples, + sampleRate: TranscribeConstants.sampleRate, + language: language + ) + } + + var parts: [String] = [] + parts.reserveCapacity((samples.count + chunkSize - 1) / chunkSize) + var offset = 0 + var chunkIndex = 0 + while offset < samples.count { + let end = min(offset + chunkSize, samples.count) + let chunk = Array(samples[offset.. String? { + let id = locale.identifier.lowercased() + if id.hasPrefix("zh") { return "zh" } + if id.hasPrefix("en") { return "en" } + return locale.language.languageCode?.identifier + } + + private static func debug(_ message: String) { + #if DEBUG + print("🎙️[Qwen3ASR] \(message)") + #endif + } +} + +private enum ASRServiceError: Error, LocalizedError { + case notReady + case modelNotDownloaded + case unsupportedOS + + var errorDescription: String? { + switch self { + case .notReady: + return nil + case .modelNotDownloaded: + return AppL10n.string("asr.error.modelNotDownloaded") + case .unsupportedOS: + return AppL10n.string("asr.error.unsupportedOS") + } + } +} diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Package.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Package.swift new file mode 100644 index 0000000..74a28c9 --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Package.swift @@ -0,0 +1,99 @@ +// swift-tools-version: 5.10 +import PackageDescription + +// Local fork of soniqo/speech-swift that ships ONLY what OSGKeyboard +// consumes: Qwen3ASR + Qwen3Chat. The original repo's `Package.swift` +// references a `CSpeechCore` binary target whose URL doesn't match +// its declared filename (`SpeechCore.xcframework.zip` vs target +// name `CSpeechCore`), which breaks SwiftPM resolve on a clean +// checkout. The only thing we need from speech-swift for the +// OSGKeyboard on-device path is the two Qwen3 modules and the +// AudioCommon / MLXCommon / SpeechVAD slices they depend on — the +// AudioServer / AudioCLI / AudioCLILib targets that pulled in +// SpeechCore aren't part of our build graph. +// +// Source provenance: every `.swift` file in `Sources//` is +// copied from https://github.com/soniqo/speech-swift (commit pinned +// to v0.0.21 of the upstream tag tree). Original copyright +// headers are preserved in each file. Apache-2.0 license. +// +// Track upstream: when soniqo fixes the binary-target mismatch in +// their main `Package.swift`, delete this local package and +// re-enable the upstream dependency in the host project. + +let package = Package( + name: "Qwen3Speech", + platforms: [ + .iOS("18.0"), + .macOS("15.0") + ], + products: [ + .library(name: "Qwen3ASR", targets: ["Qwen3ASR"]), + .library(name: "Qwen3Chat", targets: ["Qwen3Chat"]), + ], + dependencies: [ + // mlx-swift is the Apple MLX array framework bindings; Qwen3 + // runtime depends on the GPU side, the chat runtime depends + // on the linear-attention kernels exposed by MLXNN / MLXFast. + // + // We pin to a local flattened copy at `~/.local/mlx-swift` + // (an exported snapshot of mlx-swift 0.31.4 with its Cmlx / + // mlx-c submodules baked in as plain directories) because + // SwiftPM can't reliably fetch the upstream's git submodules + // on this network — the Cmlx/mlx submodule is ~700 MB of + // history and the clone drops mid-fetch. The snapshot is + // generated once on a healthy network, kept outside the + // project, and re-used on every resolve. + .package(path: "/Users/rocky/.local/mlx-swift"), + // swift-transformers exposes Hugging Face Hub and tokenizers + // — AudioCommon uses Hub to resolve repo → snapshot path. + .package(url: "https://github.com/huggingface/swift-transformers", from: "1.1.6"), + ], + targets: [ + .target( + name: "AudioCommon", + dependencies: [ + .product(name: "Hub", package: "swift-transformers"), + ] + ), + .target( + name: "MLXCommon", + dependencies: [ + "AudioCommon", + .product(name: "MLX", package: "mlx-swift"), + .product(name: "MLXNN", package: "mlx-swift"), + .product(name: "MLXFast", package: "mlx-swift"), + ] + ), + .target( + name: "SpeechVAD", + dependencies: [ + "AudioCommon", + "MLXCommon", + .product(name: "MLX", package: "mlx-swift"), + .product(name: "MLXNN", package: "mlx-swift"), + ] + ), + .target( + name: "Qwen3ASR", + dependencies: [ + "AudioCommon", + "MLXCommon", + "SpeechVAD", + .product(name: "MLX", package: "mlx-swift"), + .product(name: "MLXNN", package: "mlx-swift"), + .product(name: "MLXFast", package: "mlx-swift"), + ] + ), + .target( + name: "Qwen3Chat", + dependencies: [ + "AudioCommon", + "MLXCommon", + .product(name: "MLX", package: "mlx-swift"), + .product(name: "MLXNN", package: "mlx-swift"), + .product(name: "MLXFast", package: "mlx-swift"), + ] + ), + ] +) diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioFileLoader.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioFileLoader.swift new file mode 100644 index 0000000..42812d6 --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioFileLoader.swift @@ -0,0 +1,382 @@ +import Foundation +import AVFoundation + +/// Sample-rate-conversion quality. Both options fully drain the converter and +/// produce exact-length output; they differ only in the SRC filter. +public enum ResampleQuality { + /// Framework-default band-limited SRC (`Normal` algorithm). Anti-aliases + /// steep downsamples and retains high frequencies well below Nyquist; + /// rolls off slightly more near Nyquist than `.mastering`. The right + /// default for speech/voice, which is band-limited and usually + /// downsampled (e.g. 44.1k→16k for ASR), where mastering-grade filtering + /// is wasted cost. + case standard + /// Mastering algorithm at maximum quality — fullest high-frequency + /// retention right up to Nyquist, at higher cost. Use for music (source + /// separation) and upsampling/super-resolution, where full-band fidelity + /// matters. + case mastering +} + +/// Loads audio files and converts to float samples +public enum AudioFileLoader { + /// Load audio file and return samples at target sample rate. + /// `quality` selects the SRC filter when resampling (default `.standard`; + /// pass `.mastering` for music/upsampling). + public static func load(url: URL, targetSampleRate: Int = 24000, quality: ResampleQuality = .standard) throws -> [Float] { + let audioFile = try AVAudioFile(forReading: url) + let format = audioFile.processingFormat + let frameCount = AVAudioFrameCount(audioFile.length) + + guard let buffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: frameCount) else { + throw AudioLoadError.bufferCreationFailed + } + + try audioFile.read(into: buffer) + + guard let floatData = buffer.floatChannelData else { + throw AudioLoadError.noFloatData + } + + // Get mono samples (use first channel) + let samples = Array(UnsafeBufferPointer(start: floatData[0], count: Int(buffer.frameLength))) + + // Resample if needed + let inputSampleRate = Int(format.sampleRate) + if inputSampleRate != targetSampleRate { + return resample(samples, from: inputSampleRate, to: targetSampleRate, quality: quality) + } + + return samples + } + + /// Load audio file and return stereo channels at target sample rate. + /// Returns `[left, right]` — mono files are duplicated to stereo. + /// `quality` selects the SRC filter when resampling (default `.standard`; + /// pass `.mastering` for music). + public static func loadStereo(url: URL, targetSampleRate: Int = 44100, quality: ResampleQuality = .standard) throws -> [[Float]] { + let audioFile = try AVAudioFile(forReading: url) + let format = audioFile.processingFormat + let frameCount = AVAudioFrameCount(audioFile.length) + + guard let buffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: frameCount) else { + throw AudioLoadError.bufferCreationFailed + } + + try audioFile.read(into: buffer) + + guard let floatData = buffer.floatChannelData else { + throw AudioLoadError.noFloatData + } + + let count = Int(buffer.frameLength) + let left = Array(UnsafeBufferPointer(start: floatData[0], count: count)) + let right: [Float] + if format.channelCount >= 2 { + right = Array(UnsafeBufferPointer(start: floatData[1], count: count)) + } else { + right = left // Mono → duplicate + } + + let inputSampleRate = Int(format.sampleRate) + if inputSampleRate != targetSampleRate { + // Resample both channels in one converter pass so L/R stay + // phase-aligned (two independent converters can drift). + return resampleStereo([left, right], from: inputSampleRate, to: targetSampleRate, quality: quality) + } + + return [left, right] + } + + /// Load WAV file directly (for 16-bit PCM) + public static func loadWAV(url: URL) throws -> (samples: [Float], sampleRate: Int) { + let data = try Data(contentsOf: url) + + // Parse WAV header + guard data.count > 44 else { + throw AudioLoadError.invalidWAVFile + } + + // Check RIFF header + let riff = String(data: data[0..<4], encoding: .ascii) + guard riff == "RIFF" else { + throw AudioLoadError.invalidWAVFile + } + + // Check WAVE format + let wave = String(data: data[8..<12], encoding: .ascii) + guard wave == "WAVE" else { + throw AudioLoadError.invalidWAVFile + } + + // Parse format chunk (handle unaligned reads) + let audioFormat = data[20..<22].withUnsafeBytes { $0.loadUnaligned(as: UInt16.self) } + let numChannels = data[22..<24].withUnsafeBytes { $0.loadUnaligned(as: UInt16.self) } + let sampleRate = data[24..<28].withUnsafeBytes { $0.loadUnaligned(as: UInt32.self) } + let bitsPerSample = data[34..<36].withUnsafeBytes { $0.loadUnaligned(as: UInt16.self) } + + guard audioFormat == 1 else { // PCM + throw AudioLoadError.unsupportedFormat("Not PCM format") + } + + guard numChannels > 0 else { + throw AudioLoadError.invalidWAVFile + } + + guard bitsPerSample == 16 else { + throw AudioLoadError.unsupportedFormat("Not 16-bit") + } + + // Find data chunk + var dataOffset = 36 + var dataChunkSize: UInt32? = nil + while dataOffset < data.count - 8 { + let chunkId = String(data: data[dataOffset..<(dataOffset+4)], encoding: .ascii) + let chunkSize = data[(dataOffset+4)..<(dataOffset+8)].withUnsafeBytes { $0.loadUnaligned(as: UInt32.self) } + + if chunkId == "data" { + dataOffset += 8 + dataChunkSize = chunkSize + break + } + + // Validate chunk advance to avoid out-of-bounds. + let nextOffset = dataOffset + 8 + Int(chunkSize) + guard nextOffset >= dataOffset, nextOffset <= data.count else { + throw AudioLoadError.invalidWAVFile + } + dataOffset = nextOffset + } + + // Read samples + guard let chunkSize = dataChunkSize else { + throw AudioLoadError.invalidWAVFile + } + let chunkSizeInt = Int(chunkSize) + guard dataOffset >= 0, dataOffset <= data.count, dataOffset + chunkSizeInt <= data.count else { + throw AudioLoadError.invalidWAVFile + } + + let sampleData = data[dataOffset..<(dataOffset + chunkSizeInt)] + let channels = Int(numChannels) + let bytesPerSample = 2 + let frameSize = bytesPerSample * channels + let sampleCount = sampleData.count / frameSize + + var samples = [Float](repeating: 0, count: sampleCount) + sampleData.withUnsafeBytes { ptr in + let int16Ptr = ptr.bindMemory(to: Int16.self) + for i in 0.. [Float] { + guard inputRate != outputRate, !samples.isEmpty else { return samples } + + guard let sourceFormat = AVAudioFormat( + commonFormat: .pcmFormatFloat32, sampleRate: Double(inputRate), + channels: 1, interleaved: false), + let targetFormat = AVAudioFormat( + commonFormat: .pcmFormatFloat32, sampleRate: Double(outputRate), + channels: 1, interleaved: false), + let converter = AVAudioConverter(from: sourceFormat, to: targetFormat), + let sourceBuffer = AVAudioPCMBuffer( + pcmFormat: sourceFormat, frameCapacity: AVAudioFrameCount(samples.count)) + else { + return samples + } + + configureSRC(converter, quality: quality) + sourceBuffer.frameLength = AVAudioFrameCount(samples.count) + samples.withUnsafeBufferPointer { src in + sourceBuffer.floatChannelData![0].update(from: src.baseAddress!, count: samples.count) + } + + let ratio = Double(outputRate) / Double(inputRate) + guard let out = convertDrained( + converter: converter, source: sourceBuffer, targetFormat: targetFormat, + inputFrames: samples.count, ratio: ratio, channels: 1) + else { + return samples + } + return out[0] + } + + /// Resample a stereo signal in a single converter pass so the two channels + /// stay phase-aligned. `channels[0]` = left, `channels[1]` = right; both + /// must have equal length. `quality` selects the SRC filter (default + /// `.standard`; pass `.mastering` for music). Falls back to per-channel + /// mono resampling for non-stereo input or on converter-setup failure. + public static func resampleStereo(_ channels: [[Float]], from inputRate: Int, to outputRate: Int, quality: ResampleQuality = .standard) -> [[Float]] { + guard channels.count == 2 else { + return channels.map { resample($0, from: inputRate, to: outputRate, quality: quality) } + } + let n = channels[0].count + guard inputRate != outputRate, n > 0, channels[1].count == n else { + return channels + } + + guard let sourceFormat = AVAudioFormat( + commonFormat: .pcmFormatFloat32, sampleRate: Double(inputRate), + channels: 2, interleaved: false), + let targetFormat = AVAudioFormat( + commonFormat: .pcmFormatFloat32, sampleRate: Double(outputRate), + channels: 2, interleaved: false), + let converter = AVAudioConverter(from: sourceFormat, to: targetFormat), + let sourceBuffer = AVAudioPCMBuffer( + pcmFormat: sourceFormat, frameCapacity: AVAudioFrameCount(n)) + else { + return channels.map { resample($0, from: inputRate, to: outputRate, quality: quality) } + } + + configureSRC(converter, quality: quality) + sourceBuffer.frameLength = AVAudioFrameCount(n) + channels[0].withUnsafeBufferPointer { + sourceBuffer.floatChannelData![0].update(from: $0.baseAddress!, count: n) + } + channels[1].withUnsafeBufferPointer { + sourceBuffer.floatChannelData![1].update(from: $0.baseAddress!, count: n) + } + + let ratio = Double(outputRate) / Double(inputRate) + guard let out = convertDrained( + converter: converter, source: sourceBuffer, targetFormat: targetFormat, + inputFrames: n, ratio: ratio, channels: 2) + else { + return channels.map { resample($0, from: inputRate, to: outputRate, quality: quality) } + } + return out + } + + /// Configure the converter's SRC filter. Must be set before the first + /// `convert`. `.standard` leaves the framework default (`Normal`); only + /// `.mastering` opts into the slower, full-band Mastering algorithm. + private static func configureSRC(_ converter: AVAudioConverter, quality: ResampleQuality) { + switch quality { + case .standard: + break // framework default Normal SRC — already drains + exact length + case .mastering: + converter.sampleRateConverterAlgorithm = AVSampleRateConverterAlgorithm_Mastering + converter.sampleRateConverterQuality = .max + } + } + + /// Run the converter to completion, draining its internal tail via + /// `.endOfStream`, and return one Float array per channel normalized to the + /// exact expected frame count. + /// + /// Returns `nil` unless the converter reaches `.endOfStream` cleanly. Only + /// `.endOfStream` is success: `.error` (or a thrown `NSError`) is a hard + /// failure, and a no-progress step that isn't end-of-stream means the + /// converter is stuck. In every non-success case the partial output is + /// discarded rather than returned, so callers can fall back instead of + /// silently propagating a truncated buffer (which would desync downstream + /// audio/video). + private static func convertDrained( + converter: AVAudioConverter, + source: AVAudioPCMBuffer, + targetFormat: AVAudioFormat, + inputFrames: Int, + ratio: Double, + channels: Int + ) -> [[Float]]? { + // ceil + headroom for the sinc filter's priming/tail latency. + let capacity = AVAudioFrameCount(ceil(Double(inputFrames) * ratio)) + 4096 + guard let target = AVAudioPCMBuffer(pcmFormat: targetFormat, frameCapacity: capacity) else { + return nil + } + + var out = [[Float]](repeating: [], count: channels) + for c in 0.. 0, let chans = target.floatChannelData { + for c in 0.. 0, !out[0].isEmpty else { return nil } + for c in 0.. expected { + out[c].removeLast(out[c].count - expected) + } else if out[c].count < expected { + out[c].append(contentsOf: repeatElement(0, count: expected - out[c].count)) + } + } + return out + } +} + +public enum AudioLoadError: Error, LocalizedError { + case bufferCreationFailed + case noFloatData + case invalidWAVFile + case unsupportedFormat(String) + + public var errorDescription: String? { + switch self { + case .bufferCreationFailed: + return "Failed to create audio buffer" + case .noFloatData: + return "No float channel data available" + case .invalidWAVFile: + return "Invalid WAV file format" + case .unsupportedFormat(let reason): + return "Unsupported audio format: \(reason)" + } + } +} diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioIO.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioIO.swift new file mode 100644 index 0000000..ad7fd3e --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioIO.swift @@ -0,0 +1,174 @@ +#if canImport(AVFoundation) +import AVFoundation +import os + +/// Reusable audio I/O manager — handles mic capture, resampling, and playback. +/// +/// Eliminates AVAudioEngine boilerplate that every demo app reimplements. +/// +/// ```swift +/// let audio = AudioIO() +/// try audio.startMicrophone(targetSampleRate: 16000) { samples in +/// pipeline.pushAudio(samples) +/// } +/// audio.player.scheduleChunk(ttsOutput) +/// audio.stopMicrophone() +/// ``` +public final class AudioIO { + /// Microphone state. + public enum MicrophoneState: Sendable { + case stopped, running, error(String) + } + + /// Audio player for TTS output. Attached to the engine when mic starts. + public let player = StreamingAudioPlayer() + + /// Current microphone state. + public private(set) var microphoneState: MicrophoneState = .stopped + + /// RMS audio level (0.0–1.0) for UI meters. Updated on each mic buffer. + public private(set) var audioLevel: Float = 0 + + /// Whether to enable Voice Processing I/O for echo cancellation. + public let enableAEC: Bool + + /// Playback sample rate (for TTS output). + public let playbackSampleRate: Double + + private var engine: AVAudioEngine? + private static let log = Logger(subsystem: "audio.soniqo", category: "AudioIO") + + public init(enableAEC: Bool = false, playbackSampleRate: Double = 24000) { + self.enableAEC = enableAEC + self.playbackSampleRate = playbackSampleRate + } + + /// Start microphone capture, resampled to targetSampleRate. + /// + /// Also attaches the player to the engine for simultaneous playback. + /// Call `player.scheduleChunk()` to play audio while recording. + /// + /// - Parameters: + /// - targetSampleRate: Output sample rate for onSamples (default 16kHz for VAD/ASR) + /// - onSamples: Callback with resampled mono Float32 samples (called on audio thread) + public func startMicrophone( + targetSampleRate: Int = 16000, + onSamples: @escaping ([Float]) -> Void + ) throws { + stopMicrophone() + + #if os(iOS) + let session = AVAudioSession.sharedInstance() + if enableAEC { + try session.setCategory(.playAndRecord, mode: .voiceChat, options: [.defaultToSpeaker, .allowBluetoothHFP]) + } else { + try session.setCategory(.playAndRecord, mode: .default, options: [.defaultToSpeaker, .allowBluetoothHFP]) + } + try session.setActive(true) + #endif + + let engine = AVAudioEngine() + let inputNode = engine.inputNode + let hwFormat = inputNode.outputFormat(forBus: 0) + + // Mono intermediate at hardware rate + guard let monoFormat = AVAudioFormat( + commonFormat: .pcmFormatFloat32, + sampleRate: hwFormat.sampleRate, + channels: 1, + interleaved: false + ) else { + microphoneState = .error("Cannot create mono format") + return + } + + // Target format for VAD/ASR + guard let targetFormat = AVAudioFormat( + commonFormat: .pcmFormatFloat32, + sampleRate: Double(targetSampleRate), + channels: 1, + interleaved: false + ) else { + microphoneState = .error("Cannot create target format") + return + } + + guard let resampler = AVAudioConverter(from: monoFormat, to: targetFormat) else { + microphoneState = .error("Cannot create resampler") + return + } + + inputNode.installTap(onBus: 0, bufferSize: 1024, format: hwFormat) { [weak self] buffer, _ in + guard let self else { return } + guard let srcData = buffer.floatChannelData else { return } + let frameLen = Int(buffer.frameLength) + guard frameLen > 0 else { return } + + // Extract channel 0 into mono buffer + guard let monoBuffer = AVAudioPCMBuffer(pcmFormat: monoFormat, frameCapacity: buffer.frameCapacity) else { return } + monoBuffer.frameLength = buffer.frameLength + memcpy(monoBuffer.floatChannelData![0], srcData[0], frameLen * MemoryLayout.size) + + // Resample + let outFrameCount = AVAudioFrameCount(Double(frameLen) * Double(targetSampleRate) / hwFormat.sampleRate) + guard outFrameCount > 0, + let outBuffer = AVAudioPCMBuffer(pcmFormat: targetFormat, frameCapacity: outFrameCount) else { return } + + var error: NSError? + resampler.convert(to: outBuffer, error: &error) { _, outStatus in + outStatus.pointee = .haveData + return monoBuffer + } + if error != nil { return } + + guard let outData = outBuffer.floatChannelData else { return } + let count = Int(outBuffer.frameLength) + guard count > 0 else { return } + let samples = Array(UnsafeBufferPointer(start: outData[0], count: count)) + + // RMS for audio level + var sum: Float = 0 + for s in samples { sum += s * s } + self.audioLevel = sqrt(sum / max(Float(count), 1)) + + onSamples(samples) + } + + // Attach player for TTS output + guard let playerFormat = AVAudioFormat( + commonFormat: .pcmFormatFloat32, + sampleRate: playbackSampleRate, + channels: 1, + interleaved: false + ) else { return } + player.attach(to: engine, format: playerFormat) + + do { + try engine.start() + player.startPlayback() + self.engine = engine + microphoneState = .running + Self.log.info("Microphone started at \(targetSampleRate)Hz, player at \(self.playbackSampleRate)Hz") + } catch { + microphoneState = .error(error.localizedDescription) + throw error + } + } + + /// Stop microphone capture and detach player. + public func stopMicrophone() { + if let engine { + engine.inputNode.removeTap(onBus: 0) + player.detach(from: engine) + engine.stop() + } + engine = nil + audioLevel = 0 + microphoneState = .stopped + } + + deinit { + stopMicrophone() + } +} +#endif diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioModelError.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioModelError.swift new file mode 100644 index 0000000..b07fb28 --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioModelError.swift @@ -0,0 +1,34 @@ +import Foundation + +/// Unified error type for audio model operations. +public enum AudioModelError: Error, LocalizedError { + /// Model failed to load from disk or network. + case modelLoadFailed(modelId: String, reason: String, underlying: Error? = nil) + /// Weight file could not be read or parsed. + case weightLoadingFailed(path: String, underlying: Error? = nil) + /// Inference or generation step failed. + case inferenceFailed(operation: String, reason: String) + /// Model configuration is invalid or incompatible. + case invalidConfiguration(model: String, reason: String) + /// Voice preset file not found. + case voiceNotFound(voice: String, searchPath: String) + + public var errorDescription: String? { + switch self { + case .modelLoadFailed(let modelId, let reason, let underlying): + var msg = "Failed to load model '\(modelId)': \(reason)" + if let underlying { msg += " (\(underlying.localizedDescription))" } + return msg + case .weightLoadingFailed(let path, let underlying): + var msg = "Failed to load weights from '\(path)'" + if let underlying { msg += ": \(underlying.localizedDescription)" } + return msg + case .inferenceFailed(let operation, let reason): + return "Inference failed during \(operation): \(reason)" + case .invalidConfiguration(let model, let reason): + return "Invalid configuration for '\(model)': \(reason)" + case .voiceNotFound(let voice, let searchPath): + return "Voice preset '\(voice)' not found at '\(searchPath)'" + } + } +} diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioRingBuffer.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioRingBuffer.swift new file mode 100644 index 0000000..9aad18a --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/AudioRingBuffer.swift @@ -0,0 +1,75 @@ +import Foundation +import os + +/// Thread-safe ring buffer for passing audio between the audio capture thread and the MLX +/// inference thread. Writes drop oldest data when full; reads return zeros on underrun. +/// +/// Uses `os_unfair_lock` for priority inheritance — safe to call `write` from a real-time +/// Core Audio I/O thread without risking priority inversion. +public final class AudioRingBuffer: @unchecked Sendable { + private var buffer: [Float] + private var readPos = 0 + private var writePos = 0 + private var count = 0 + private var _lock = os_unfair_lock() + private let capacity: Int + + public init(capacity: Int) { + self.capacity = capacity + self.buffer = [Float](repeating: 0, count: capacity) + } + + /// Called from audio capture thread — non-blocking; drops oldest data if full. + public func write(_ samples: [Float]) { + os_unfair_lock_lock(&_lock) + defer { os_unfair_lock_unlock(&_lock) } + for sample in samples { + if count == capacity { + // Drop oldest sample + readPos = (readPos + 1) % capacity + count -= 1 + } + buffer[writePos] = sample + writePos = (writePos + 1) % capacity + count += 1 + } + } + + /// Zero-copy write from a raw pointer — preferred on real-time audio threads + /// to avoid heap allocation from `Array(UnsafeBufferPointer(...))`. + public func write(from pointer: UnsafePointer, count sampleCount: Int) { + os_unfair_lock_lock(&_lock) + defer { os_unfair_lock_unlock(&_lock) } + for i in 0.. [Float] { + os_unfair_lock_lock(&_lock) + defer { os_unfair_lock_unlock(&_lock) } + var result = [Float](repeating: 0, count: n) + let available = min(n, count) + for i in 0.. MLComputeUnits { + guard let raw = ProcessInfo.processInfo.environment[envKey]? + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased(), !raw.isEmpty + else { + return fallback + } + switch raw { + case "ane", "cpuandneuralengine", "neuralengine": + return .cpuAndNeuralEngine + case "gpu", "cpuandgpu": + return .cpuAndGPU + case "cpu", "cpuonly": + return .cpuOnly + case "all": + return .all + default: + return fallback + } + } +} +#endif diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/CoreMLLoader.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/CoreMLLoader.swift new file mode 100644 index 0000000..35e0933 --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/CoreMLLoader.swift @@ -0,0 +1,125 @@ +import CoreML +import Foundation +#if canImport(os) +import os +#endif + +/// CoreML model loader that surfaces Neural Engine fallback. +/// +/// `MLModel(contentsOf:configuration:)` silently succeeds when +/// ``MILCompilerForANE`` fails — the model just runs on CPU instead of +/// ANE. Users only see the performance cliff: RTF jumps from ~0.04 to +/// ~1.8 on wake-word, ASR slows 5–20×, etc. They have no way to +/// correlate this with the CoreML runtime's ``E5RT encountered an STL +/// exception. msg = MILCompilerForANE error`` stderr message. +/// +/// This helper: +/// 1. Times the load. +/// 2. Logs a single structured line per model with name + compute +/// units + elapsed ms. +/// 3. When the requested compute units include `.cpuAndNeuralEngine` +/// (or `.all`) and the load completes faster than a typical ANE +/// compile, emits a one-time warning pointing users at the +/// fallback diagnostic. +/// +/// Usage: +/// ```swift +/// let encoder = try CoreMLLoader.load( +/// url: cacheDir.appendingPathComponent("encoder.mlmodelc"), +/// computeUnits: .cpuAndNeuralEngine, +/// name: "parakeet-eou-encoder" +/// ) +/// ``` +public enum CoreMLLoader { + + /// Seconds under which an ANE-eligible load is considered suspicious + /// (likely CPU fallback). Calibrated against observed behaviour: + /// - Successful ANE compile: ~200–800 ms on cold cache, ~20–50 ms + /// cached. + /// - CPU fallback after ANE compile failure: <10 ms regardless of + /// cache state. + /// + /// Picking 15 ms keeps false positives low on warm caches while + /// still catching the silent-fallback case on cold systems. + private static let aneCompileFloorSeconds: Double = 0.015 + + /// Track which model names we've already warned about so we don't + /// spam the log. Protected by ``warnedQueue``. + private static var warnedNames = Set() + private static let warnedQueue = DispatchQueue( + label: "com.qwen3speech.coreml-loader.warned" + ) + + /// Load a compiled CoreML model with instrumentation. + public static func load( + url: URL, + computeUnits: MLComputeUnits, + name: String? = nil + ) throws -> MLModel { + let config = MLModelConfiguration() + config.computeUnits = computeUnits + return try load(url: url, configuration: config, name: name) + } + + /// Load with an explicit ``MLModelConfiguration``. + public static func load( + url: URL, + configuration: MLModelConfiguration, + name: String? = nil + ) throws -> MLModel { + // Honor the SPEECH_COREML_COMPUTE_UNITS override (CI forces cpuOnly to + // skip the runner's hanging ANE/GPU first-load compile). No-op on device. + configuration.computeUnits = CoreMLComputeUnitsResolver.resolved( + default: configuration.computeUnits) + let label = name ?? url.deletingPathExtension().lastPathComponent + let unitsLabel = describe(units: configuration.computeUnits) + let start = Date() + let model = try MLModel(contentsOf: url, configuration: configuration) + let elapsed = Date().timeIntervalSince(start) + let ms = Int((elapsed * 1000).rounded()) + AudioLog.modelLoading.info("CoreML loaded \(label) in \(ms)ms (units=\(unitsLabel))") + + let aneEligible = + configuration.computeUnits == .cpuAndNeuralEngine || + configuration.computeUnits == .all + if aneEligible && elapsed < aneCompileFloorSeconds { + maybeWarn( + name: label, + message: """ + CoreML model '\(label)' loaded in \(ms)ms with compute units \ + \(unitsLabel). This is faster than a typical Neural Engine \ + compile (~200–800 ms cold, ~20–50 ms cached). If console logs \ + show 'MILCompilerForANE error', the model has fallen back to \ + CPU and inference may be 5–20× slower than expected. + """ + ) + } + return model + } + + // MARK: - Private + + private static func maybeWarn(name: String, message: String) { + warnedQueue.sync { + guard !warnedNames.contains(name) else { return } + warnedNames.insert(name) + AudioLog.modelLoading.warning("\(message)") + } + } + + private static func describe(units: MLComputeUnits) -> String { + switch units { + case .cpuOnly: return "cpuOnly" + case .cpuAndGPU: return "cpuAndGPU" + case .all: return "all" + case .cpuAndNeuralEngine: return "cpuAndNeuralEngine" + @unknown default: return "unknown(\(units.rawValue))" + } + } + + /// Reset the per-process warning set. Exposed for tests so a fresh + /// run of the helper can emit a warning again. + public static func resetWarningState() { + warnedQueue.sync { warnedNames.removeAll() } + } +} diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/HuggingFaceDownloader.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/HuggingFaceDownloader.swift new file mode 100644 index 0000000..cc8d95f --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/HuggingFaceDownloader.swift @@ -0,0 +1,428 @@ +import Foundation +import Hub +import os + +/// Download errors +public enum DownloadError: Error, LocalizedError { + case failedToDownload(String) + case invalidRemoteFileName(String) + /// A download attempt made no progress for `seconds` and was aborted + /// so the caller's retry loop can fire instead of hanging. + case stalled(modelId: String, seconds: Int) + + public var errorDescription: String? { + switch self { + case .failedToDownload(let file): + return "Failed to download: \(file)" + case .invalidRemoteFileName(let file): + return "Refusing to write unsafe remote file name: \(file)" + case .stalled(let modelId, let seconds): + return "Download stalled for \(modelId): no progress in \(seconds)s" + } + } +} + +/// HuggingFace model downloader — shared between ASR, TTS, VAD, etc. +/// +/// Uses `HubApi` from the swift-transformers `Hub` module for downloads, +/// which provides HF token auth and metadata tracking. Files that finished +/// downloading are skipped on retry (etag/commit-hash check), but a file +/// interrupted mid-transfer restarts from byte 0 — there is no usable +/// mid-file resume in the current Hub stack, which is why the stall guard +/// and retry ladder below favor patience over fast abort. +public enum HuggingFaceDownloader { + + // MARK: - Cache Directory + + /// Get cache directory for a model. + /// + /// Returns the old flat cache path if it already contains model files (preserving + /// ~10 GB of existing cached models), otherwise returns the new Hub-style path. + public static func getCacheDirectory(for modelId: String, basePath: URL? = nil, cacheDirName: String = "qwen3-speech") throws -> URL { + let base = basePath ?? resolveBaseCacheDir(cacheDirName: cacheDirName) + let fm = FileManager.default + + // Check old (flat) cache path for backward compat: + // ~/Library/Caches/qwen3-speech/aufklarer_Qwen3-ASR-0.6B-MLX-4bit/ + let oldDir = base.appendingPathComponent(sanitizedCacheKey(for: modelId), isDirectory: true) + if weightsExist(in: oldDir) { + return oldDir + } + + // New Hub-style path: + // ~/Library/Caches/qwen3-speech/models/aufklarer/Qwen3-ASR-0.6B-MLX-4bit/ + let hub = HubApi(downloadBase: base) + let repo = Hub.Repo(id: modelId) + let dir = hub.localRepoLocation(repo) + try fm.createDirectory(at: dir, withIntermediateDirectories: true) + return dir + } + + // MARK: - Weight Existence Check + + /// Extensions recognised as cached model weights: the canonical + /// HF `.safetensors` layout plus Apple CoreML bundle directories + /// (`.mlmodelc`, `.mlpackage`) shipped by CoreML-only repos. + public static let weightFileExtensions: Set = [ + "safetensors", "mlmodelc", "mlpackage" + ] + + /// Returns `true` when `directory` contains at least one entry + /// whose extension matches `weightFileExtensions`. Used by + /// `downloadWeights` to short-circuit network requests when + /// `offlineMode: true` is set on caches that contain only CoreML + /// bundles and no `.safetensors` files. + public static func weightsExist(in directory: URL) -> Bool { + let fm = FileManager.default + guard fm.fileExists(atPath: directory.path) else { return false } + let contents: [URL] + do { + contents = try fm.contentsOfDirectory(at: directory, includingPropertiesForKeys: nil) + } catch { + AudioLog.download.debug("Could not list directory \(directory.path): \(error)") + contents = [] + } + return contents.contains { weightFileExtensions.contains($0.pathExtension) } + } + + // MARK: - Download + + /// Download model files from HuggingFace using `HubApi.snapshot()`. + /// + /// Builds glob patterns from the file list: + /// - Always includes `config.json` + /// - If `additionalFiles` doesn't contain `.safetensors` files, adds `*.safetensors` + /// and `model.safetensors.index.json` to discover sharded weights automatically + /// - All entries in `additionalFiles` are added as-is (they work as glob patterns) + public static func downloadWeights( + modelId: String, + to directory: URL, + additionalFiles: [String] = [], + offlineMode: Bool = false, + hubEndpoint: String? = nil, + retryDelaysSeconds: [Int]? = nil, + progressHandler: ((Double) -> Void)? = nil + ) async throws { + // Skip network requests when weights are already cached + if offlineMode && weightsExist(in: directory) { + progressHandler?(1.0) + return + } + + prepareRepoDirectoryForDownload(at: directory) + + var globs: [String] = ["config.json"] + + let hasExplicitWeights = additionalFiles.contains { $0.hasSuffix(".safetensors") } + if !hasExplicitWeights { + globs.append("*.safetensors") + globs.append("model.safetensors.index.json") + } + for file in additionalFiles where !globs.contains(file) { + globs.append(file) + } + + // Derive the download base from the directory. + // getCacheDirectory returns either: + // old: base/cacheKey (flat, already has weights — won't reach here) + // new: base/models/org/model (Hub-style) + // For Hub API we need `base` as downloadBase. + // + // Forward `offlineMode` explicitly so HubApi doesn't fall through to + // its internal NWPathMonitor auto-detect, which on macOS can briefly + // report `.unsatisfied` and then refuse to download (manifesting as + // "Offline mode error: No files available locally for this repository" + // for a freshly-requested model). + let hub = makeHubApi(for: modelId, repoDir: directory, offlineMode: offlineMode, hubEndpoint: hubEndpoint) + let repo = Hub.Repo(id: modelId) + + // Retry with capped backoff — HuggingFace can timeout on slow + // connections or rate-limit, and flaky networks (hotspots, captive + // portals) drop out for minutes at a time. Each attempt is wrapped + // in a progress-stall guard so a wedged mid-transfer (which + // `hub.snapshot` won't surface on its own) aborts and retries + // instead of hanging until the CI job is killed. + // + // No retries in offline mode: the failure is a deterministic local + // cache miss, and 110 s of backoff can't change what's on disk. + let delays = offlineMode ? [] : (retryDelaysSeconds ?? downloadRetryDelaysSeconds) + let maxAttempts = delays.count + 1 + var lastError: Error? + for attempt in 1...maxAttempts { + do { + try await withDownloadStallGuard(modelId: modelId) { reportProgress in + try await hub.snapshot(from: repo, matching: globs) { progress in + reportProgress(progress.fractionCompleted) + progressHandler?(progress.fractionCompleted) + } + } + return // Success + } catch { + lastError = error + if isRecoverableHubCacheError(error) { + prepareRepoDirectoryForDownload(at: directory, force: true) + } + if attempt < maxAttempts { + try await Task.sleep(for: .seconds(delays[attempt - 1])) + } + } + } + throw DownloadError.failedToDownload( + "\(modelId) after \(maxAttempts) attempt\(maxAttempts == 1 ? "" : "s") " + + "(target: \(directory.path)): " + + (lastError?.localizedDescription ?? "unknown")) + } + + /// Download an explicit list of files from HuggingFace without adding any + /// implicit weight globs. This is useful for overlaying tokenizer or config + /// assets from a second repository on top of an existing cache. + public static func downloadFiles( + modelId: String, + to directory: URL, + files: [String], + offlineMode: Bool = false, + hubEndpoint: String? = nil, + retryDelaysSeconds: [Int]? = nil, + progressHandler: ((Double) -> Void)? = nil + ) async throws { + if files.isEmpty { + progressHandler?(1.0) + return + } + + prepareRepoDirectoryForDownload(at: directory) + + let hub = makeHubApi(for: modelId, repoDir: directory, offlineMode: offlineMode, hubEndpoint: hubEndpoint) + let repo = Hub.Repo(id: modelId) + + let globs = files.map { $0 } + // Same retry semantics as downloadWeights, including the offline + // no-retry rule — keep the two loops in lockstep. + let delays = offlineMode ? [] : (retryDelaysSeconds ?? downloadRetryDelaysSeconds) + let maxAttempts = delays.count + 1 + var lastError: Error? + for attempt in 1...maxAttempts { + do { + try await withDownloadStallGuard(modelId: modelId) { reportProgress in + try await hub.snapshot(from: repo, matching: globs) { progress in + reportProgress(progress.fractionCompleted) + progressHandler?(progress.fractionCompleted) + } + } + return + } catch { + lastError = error + if isRecoverableHubCacheError(error) { + prepareRepoDirectoryForDownload(at: directory, force: true) + } + if attempt < maxAttempts { + try await Task.sleep(for: .seconds(delays[attempt - 1])) + } + } + } + throw DownloadError.failedToDownload( + "\(modelId) after \(maxAttempts) attempt\(maxAttempts == 1 ? "" : "s") " + + "(target: \(directory.path)): " + + (lastError?.localizedDescription ?? "unknown")) + } + + // MARK: - Retry ladder + + /// Delays between download attempts. One more attempt than entries: + /// 5 attempts with 5/15/30/60 s pauses (~110 s of backoff on top of the + /// per-attempt stall patience). Generous on purpose — abandoned attempts + /// restart files from byte 0 with the current Hub stack, so the cheap + /// resource here is wall-clock, not bytes. A network that's down for a + /// couple of minutes (AP roam, hotspot sleep, captive-portal re-auth) + /// should not kill a 2.75 GB first-run download. + static let downloadRetryDelaysSeconds = [5, 15, 30, 60] + + /// Total attempts per download (retries + the initial try). + static var downloadMaxAttempts: Int { downloadRetryDelaysSeconds.count + 1 } + + // MARK: - Download stall guard + + /// Seconds of zero download progress after which an attempt is + /// considered wedged and aborted. `hub.snapshot` reports + /// `fractionCompleted` continuously while bytes flow, so a healthy + /// (even slow) transfer keeps resetting the clock; only a genuinely + /// stalled connection trips this. + /// + /// The default is tuned for end users, not CI: aborted attempts restart + /// each file from byte 0 (the Hub stack's mid-file resume never engages + /// on a fresh download), so firing the guard on a connection that would + /// have recovered throws away every byte of that attempt. Flaky networks + /// — AP roams, captive-portal re-auth, hotspot sleep — routinely stall + /// for 1–3 minutes and then recover, hence 300 s. CI pins + /// `HF_DOWNLOAD_STALL_TIMEOUT=90` to keep failing fast (app users can't + /// set env vars; CI can). + static var downloadStallTimeoutSeconds: Int { + if let raw = ProcessInfo.processInfo.environment["HF_DOWNLOAD_STALL_TIMEOUT"], + let v = Int(raw), v > 0 { + return v + } + return 300 + } + + /// Thread-safe last-progress timestamp. `hub.snapshot`'s progress + /// callback may fire from a background queue, so guard with a lock. + private final class ProgressClock: @unchecked Sendable { + private let lock = NSLock() + private var last = Date() + func tick() { lock.lock(); last = Date(); lock.unlock() } + func idleSeconds() -> Double { + lock.lock(); defer { lock.unlock() } + return Date().timeIntervalSince(last) + } + } + + /// Run a download `operation` that reports fractional progress, and + /// abort it if progress stalls for `downloadStallTimeoutSeconds`. + /// On stall the in-flight `hub.snapshot` task is cancelled (URLSession + /// honors cancellation) and `DownloadError.stalled` is thrown so the + /// caller's retry loop fires instead of hanging indefinitely. + static func withDownloadStallGuard( + modelId: String, + stallTimeoutSeconds: Int? = nil, + _ operation: @escaping (@escaping @Sendable (Double) -> Void) async throws -> Void + ) async throws { + let stall = stallTimeoutSeconds ?? downloadStallTimeoutSeconds + let clock = ProgressClock() + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + try await operation { _ in clock.tick() } + } + group.addTask { + // Poll on a fraction of the window so we detect a stall + // within ~stall..stall+pollStep seconds. + let pollStep = max(1, stall / 3) + while true { + try await Task.sleep(for: .seconds(pollStep)) + if clock.idleSeconds() >= Double(stall) { + throw DownloadError.stalled(modelId: modelId, seconds: stall) + } + } + } + // Whichever finishes first wins; cancel the other (the poller + // on success, or the download on stall). + defer { group.cancelAll() } + try await group.next() + } + } + + // MARK: - Security Helpers (kept for backward compat + security tests) + + /// Convert an arbitrary modelId into a single, safe path component for on-disk caching. + public static func sanitizedCacheKey(for modelId: String) -> String { + let replaced = modelId.replacingOccurrences(of: "/", with: "_") + + let allowed = CharacterSet(charactersIn: "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._-") + var scalars: [UnicodeScalar] = [] + scalars.reserveCapacity(replaced.unicodeScalars.count) + for s in replaced.unicodeScalars { + scalars.append(allowed.contains(s) ? s : "_") + } + + var cleaned = String(String.UnicodeScalarView(scalars)) + cleaned = cleaned.trimmingCharacters(in: CharacterSet(charactersIn: "._")) + + if cleaned.isEmpty || cleaned == "." || cleaned == ".." { + cleaned = "model" + } + + return cleaned + } + + /// Validate that a remote file name is safe. + public static func validatedRemoteFileName(_ file: String) throws -> String { + let base = URL(fileURLWithPath: file).lastPathComponent + guard base == file else { + throw DownloadError.invalidRemoteFileName(file) + } + guard !base.isEmpty, !base.hasPrefix("."), !base.contains("..") else { + throw DownloadError.invalidRemoteFileName(file) + } + guard base.range(of: #"^[A-Za-z0-9._-]+$"#, options: .regularExpression) != nil else { + throw DownloadError.invalidRemoteFileName(file) + } + return base + } + + /// Validate that a local path stays within the expected directory. + public static func validatedLocalPath(directory: URL, fileName: String) throws -> URL { + let local = directory.appendingPathComponent(fileName, isDirectory: false) + let dirPath = directory.standardizedFileURL.path + let localPath = local.standardizedFileURL.path + let prefix = dirPath.hasSuffix("/") ? dirPath : (dirPath + "/") + guard localPath.hasPrefix(prefix) else { + throw DownloadError.invalidRemoteFileName(fileName) + } + return local + } + + // MARK: - Private Helpers + + /// Remove a repo folder that has Hub metadata but no complete weights. + /// Stale partial caches trigger "File metadata must have been retrieved from server". + static func prepareRepoDirectoryForDownload(at directory: URL, force: Bool = false) { + let fm = FileManager.default + guard fm.fileExists(atPath: directory.path) else { return } + if !force && weightsExist(in: directory) { return } + try? fm.removeItem(at: directory) + try? fm.createDirectory(at: directory, withIntermediateDirectories: true) + } + + private static func isRecoverableHubCacheError(_ error: Error) -> Bool { + let text = (error as? LocalizedError)?.errorDescription + ?? error.localizedDescription + return text.localizedCaseInsensitiveContains("metadata") + || text.localizedCaseInsensitiveContains("offline mode") + } + + /// Resolve the base cache directory from env vars or system default. + private static func resolveBaseCacheDir(cacheDirName: String) -> URL { + let fm = FileManager.default + let root: URL + if let override = ProcessInfo.processInfo.environment["QWEN3_CACHE_DIR"], + !override.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty { + root = URL(fileURLWithPath: override, isDirectory: true) + } else if let override = ProcessInfo.processInfo.environment["QWEN3_ASR_CACHE_DIR"], + !override.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty { + // Legacy env var support + root = URL(fileURLWithPath: override, isDirectory: true) + } else { + root = fm.urls(for: .cachesDirectory, in: .userDomainMask).first! + } + return root.appendingPathComponent(cacheDirName, isDirectory: true) + } + + /// Create a `HubApi` whose `downloadBase` is derived from the repo directory that + /// `getCacheDirectory` returned (strips the `models//` suffix). + /// + /// `offlineMode` is forwarded as `useOfflineMode` so callers get the mode + /// they asked for instead of relying on `NWPathMonitor` auto-detection, + /// which can spuriously report `.unsatisfied` on macOS. + private static func makeHubApi( + for modelId: String, + repoDir: URL, + offlineMode: Bool, + hubEndpoint: String? + ) -> HubApi { + // repoDir is base/models/org/model + // We need base + let repo = Hub.Repo(id: modelId) + let suffix = "/\(repo.type.rawValue)/\(repo.id)" + let repoDirPath = repoDir.path + let downloadBase: URL + if repoDirPath.hasSuffix(suffix) { + let basePath = String(repoDirPath.dropLast(suffix.count)) + downloadBase = URL(fileURLWithPath: basePath, isDirectory: true) + } else { + // Fallback: old-style flat dir — use its parent as downloadBase. + // Hub won't match this path, so we derive base from env/defaults. + downloadBase = resolveBaseCacheDir(cacheDirName: repoDir.deletingLastPathComponent().lastPathComponent) + } + return HubApi(downloadBase: downloadBase, endpoint: hubEndpoint, useOfflineMode: offlineMode) + } +} diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Logging.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Logging.swift new file mode 100644 index 0000000..204debb --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Logging.swift @@ -0,0 +1,13 @@ +import os + +/// Centralized loggers for audio model subsystems. +public enum AudioLog { + /// Logger for model weight loading and initialization. + public static let modelLoading = Logger(subsystem: "com.qwen3speech", category: "ModelLoading") + /// Logger for inference and generation. + public static let inference = Logger(subsystem: "com.qwen3speech", category: "Inference") + /// Logger for HuggingFace downloads and caching. + public static let download = Logger(subsystem: "com.qwen3speech", category: "Download") + /// Logger for voice pipeline events. + public static let pipeline = Logger(subsystem: "com.qwen3speech", category: "Pipeline") +} diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelLoader.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelLoader.swift new file mode 100644 index 0000000..67583dd --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelLoader.swift @@ -0,0 +1,175 @@ +import Foundation +import os + +/// Loaded model set — holds references to all loaded models. +public struct ModelSet { + public let vad: (any StreamingVADProvider)? + public let stt: (any SpeechRecognitionModel)? + public let tts: (any SpeechGenerationModel)? + + public init( + vad: (any StreamingVADProvider)? = nil, + stt: (any SpeechRecognitionModel)? = nil, + tts: (any SpeechGenerationModel)? = nil + ) { + self.vad = vad + self.stt = stt + self.tts = tts + } +} + +/// A model to load, with its factory closure and progress weight. +public struct ModelSpec: Sendable { + let name: String + let weight: Double + let group: Int // 0 = parallel group 1, 1 = sequential group 2 + let loader: @Sendable (_ progress: @escaping (Double, String) -> Void) async throws -> any Sendable + + /// VAD model spec. + public static func vad( + _ factory: @escaping @Sendable (_ progress: @escaping (Double, String) -> Void) async throws -> any StreamingVADProvider + ) -> ModelSpec { + ModelSpec(name: "VAD", weight: 1, group: 0, loader: { progress in + try await factory(progress) as any Sendable + }) + } + + /// Speech-to-text model spec. + public static func stt( + _ factory: @escaping @Sendable (_ progress: @escaping (Double, String) -> Void) async throws -> any SpeechRecognitionModel + ) -> ModelSpec { + ModelSpec(name: "ASR", weight: 15, group: 0, loader: { progress in + try await factory(progress) as any Sendable + }) + } + + /// Text-to-speech model spec. + public static func tts( + _ factory: @escaping @Sendable (_ progress: @escaping (Double, String) -> Void) async throws -> any SpeechGenerationModel + ) -> ModelSpec { + ModelSpec(name: "TTS", weight: 20, group: 1, loader: { progress in + try await factory(progress) as any Sendable + }) + } +} + +/// Unified model loading orchestrator with aggregated progress. +/// +/// Loads multiple speech models with coordinated progress reporting. +/// Group 0 models (VAD, ASR) load in parallel; Group 1 (TTS) loads after +/// to reduce peak memory. +/// +/// ```swift +/// let models = try await ModelLoader.load([ +/// .vad { p in try await SileroVADModel.fromPretrained(engine: .coreml, progressHandler: p) }, +/// .stt { p in try await ParakeetASRModel.fromPretrained(progressHandler: p) }, +/// .tts { p in try await KokoroTTSModel.fromPretrained(progressHandler: p) }, +/// ], onProgress: { progress, stage in +/// self.loadProgress = progress +/// self.loadingStatus = stage +/// }) +/// // models.vad, models.stt, models.tts are ready +/// ``` +public enum ModelLoader { + + private static let log = Logger(subsystem: "audio.soniqo", category: "ModelLoader") + + /// Load the requested models with aggregated progress reporting. + public static func load( + _ specs: [ModelSpec], + onProgress: @escaping @Sendable (_ progress: Double, _ stage: String) -> Void = { _, _ in } + ) async throws -> ModelSet { + let totalWeight = specs.reduce(0.0) { $0 + $1.weight } + guard totalWeight > 0 else { return ModelSet() } + + let state = LoadState(totalWeight: totalWeight) + + // Group 0: parallel (VAD + ASR) + let group0 = specs.filter { $0.group == 0 } + // Group 1: sequential after group 0 (TTS — heavy, reduce peak memory) + let group1 = specs.filter { $0.group != 0 } + + var results: [(String, any Sendable)] = [] + + // Load group 0 in parallel + if !group0.isEmpty { + try await withThrowingTaskGroup(of: (String, any Sendable).self) { group in + for spec in group0 { + group.addTask { + let model = try await loadSpec(spec, state: state, onProgress: onProgress) + return (spec.name, model) + } + } + for try await result in group { + results.append(result) + } + } + } + + // Load group 1 sequentially + for spec in group1 { + let model = try await loadSpec(spec, state: state, onProgress: onProgress) + results.append((spec.name, model)) + } + + onProgress(1.0, "Ready") + log.info("All models loaded") + + // Build ModelSet from results + var vad: (any StreamingVADProvider)? + var stt: (any SpeechRecognitionModel)? + var tts: (any SpeechGenerationModel)? + + for (_, model) in results { + if let m = model as? any StreamingVADProvider { vad = m } + if let m = model as? any SpeechRecognitionModel { stt = m } + if let m = model as? any SpeechGenerationModel { tts = m } + } + + return ModelSet(vad: vad, stt: stt, tts: tts) + } + + // MARK: - Internal + + private final class LoadState: @unchecked Sendable { + let totalWeight: Double + private var completed: Double = 0 + private let lock = NSLock() + + init(totalWeight: Double) { self.totalWeight = totalWeight } + + func addCompleted(_ w: Double) { + lock.lock(); completed += w; lock.unlock() + } + + var completedFraction: Double { + lock.lock(); defer { lock.unlock() } + return completed / totalWeight + } + + func overallProgress(specWeight: Double, localFraction: Double) -> Double { + lock.lock(); defer { lock.unlock() } + return (completed + localFraction * specWeight) / totalWeight + } + } + + private static func loadSpec( + _ spec: ModelSpec, + state: LoadState, + onProgress: @escaping @Sendable (Double, String) -> Void + ) async throws -> any Sendable { + log.info("Loading \(spec.name)...") + onProgress(state.completedFraction, "\(spec.name)...") + + let adapter: @Sendable (Double, String) -> Void = { fraction, status in + let overall = state.overallProgress(specWeight: spec.weight, localFraction: fraction) + let stage = status.isEmpty ? spec.name : "\(spec.name): \(status)" + onProgress(overall, stage) + } + + let model = try await spec.loader(adapter) + state.addCompleted(spec.weight) + log.info("\(spec.name) loaded") + return model + } +} diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelRegistry.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelRegistry.swift new file mode 100644 index 0000000..9196d08 --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelRegistry.swift @@ -0,0 +1,9 @@ +import Foundation + +/// Remote registry used when fetching on-device model weights. +public enum ModelRegistry: Sendable, Equatable { + /// Official Hugging Face Hub (`swift-transformers` / `HubApi`). + case huggingFace(hubEndpoint: String? = nil) + /// ModelScope.cn — same `owner/model` ids as Hugging Face for aufklarer MLX repos. + case modelScope(baseURL: String = ModelScopeDownloader.defaultBaseURL, revision: String = "master") +} diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelScopeDownloader.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelScopeDownloader.swift new file mode 100644 index 0000000..5c4a41b --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/ModelScopeDownloader.swift @@ -0,0 +1,335 @@ +import Foundation + +/// Downloads model files from [ModelScope](https://www.modelscope.cn) using the +/// public repo API. Uses the same `owner/model` ids as Hugging Face for repos +/// mirrored on ModelScope (e.g. `aufklarer/Qwen3-ASR-0.6B-MLX-4bit`). +public enum ModelScopeDownloader { + + public static let defaultBaseURL = "https://modelscope.cn" + + private struct FilesPayload: Decodable { + struct Entry: Decodable { + let Path: String + let Size: Int64? + let entryType: String? + + enum CodingKeys: String, CodingKey { + case Path + case Size + case entryType = "Type" + } + } + let Files: [Entry] + } + + private struct APIResponse: Decodable { + let Data: FilesPayload + } + + public struct RemoteFile: Sendable { + public let path: String + public let size: Int64 + } + + // MARK: - Public API + + /// Mirror of `HuggingFaceDownloader.downloadWeights` for ModelScope. + public static func downloadWeights( + modelId: String, + to directory: URL, + additionalFiles: [String] = [], + baseURL: String = defaultBaseURL, + revision: String = "master", + retryDelaysSeconds: [Int]? = nil, + progressHandler: ((Double) -> Void)? = nil + ) async throws { + HuggingFaceDownloader.prepareRepoDirectoryForDownload(at: directory) + + let listed = try await listAllFiles(modelId: modelId, baseURL: baseURL, revision: revision) + var selected = Set(["config.json"]) + for file in additionalFiles { + selected.insert(file) + } + + let hasExplicitWeights = additionalFiles.contains { $0.hasSuffix(".safetensors") } + if !hasExplicitWeights { + for file in listed where file.path.hasSuffix(".safetensors") { + selected.insert(file.path) + } + if listed.contains(where: { $0.path == "model.safetensors.index.json" }) { + selected.insert("model.safetensors.index.json") + } + } + + let files = listed.filter { selected.contains($0.path) }.map(\.path) + guard !files.isEmpty else { + throw DownloadError.failedToDownload("\(modelId): no matching files on ModelScope") + } + + try await downloadFiles( + modelId: modelId, + to: directory, + files: files, + fileSizes: Dictionary(uniqueKeysWithValues: listed.map { ($0.path, $0.size) }), + baseURL: baseURL, + revision: revision, + retryDelaysSeconds: retryDelaysSeconds, + progressHandler: progressHandler + ) + } + + /// Download an explicit list of repo-relative paths into `directory`. + public static func downloadFiles( + modelId: String, + to directory: URL, + files: [String], + fileSizes: [String: Int64] = [:], + baseURL: String = defaultBaseURL, + revision: String = "master", + retryDelaysSeconds: [Int]? = nil, + progressHandler: ((Double) -> Void)? = nil + ) async throws { + if files.isEmpty { + progressHandler?(1.0) + return + } + + HuggingFaceDownloader.prepareRepoDirectoryForDownload(at: directory) + + let ordered = files.sorted() + var sizes = fileSizes + for path in ordered where sizes[path] == nil { + sizes[path] = 0 + } + + // Without byte sizes the old logic fell back to `(index + 1) / count`, + // which jumps to 50% as soon as two small JSON files finish. Resolve + // sizes from the repo listing whenever any entry is missing. + if ordered.contains(where: { (sizes[$0] ?? 0) <= 0 }) { + let listed = try await listAllFiles( + modelId: modelId, + baseURL: baseURL, + revision: revision + ) + let listedMap = Dictionary(uniqueKeysWithValues: listed.map { ($0.path, $0.size) }) + for path in ordered where (sizes[path] ?? 0) <= 0 { + if let remote = listedMap[path], remote > 0 { + sizes[path] = remote + } + } + } + + let totalBytes = max(ordered.reduce(Int64(0)) { $0 + (sizes[$1] ?? 0) }, 1) + var completedBytes: Int64 = 0 + + let delays = retryDelaysSeconds ?? HuggingFaceDownloader.downloadRetryDelaysSeconds + let maxAttempts = delays.count + 1 + + for (index, path) in ordered.enumerated() { + let destination = directory.appendingPathComponent(path, isDirectory: false) + try FileManager.default.createDirectory( + at: destination.deletingLastPathComponent(), + withIntermediateDirectories: true + ) + + var lastError: Error? + for attempt in 1...maxAttempts { + do { + try await HuggingFaceDownloader.withDownloadStallGuard(modelId: modelId) { reportProgress in + try await fetchFile( + modelId: modelId, + filePath: path, + to: destination, + baseURL: baseURL, + revision: revision + ) { fileBytes, fileExpectedBytes in + reportProgress(1.0) + let fileSize = sizes[path] ?? 0 + let expected = fileSize > 0 ? fileSize : fileExpectedBytes + let overall: Double + if expected > 0, totalBytes > 1 { + overall = Double(completedBytes + min(fileBytes, expected)) / Double(totalBytes) + } else { + // Last resort when listing omits sizes: spread each + // file's slice by bytes received vs Content-Length. + let slice = 1.0 / Double(ordered.count) + let base = Double(index) * slice + let inFile = expected > 0 + ? min(Double(fileBytes) / Double(expected), 1.0) * slice + : slice + overall = base + inFile + } + progressHandler?(min(max(overall, 0), 1)) + } + } + lastError = nil + break + } catch { + lastError = error + try? FileManager.default.removeItem(at: destination) + if attempt < maxAttempts { + try await Task.sleep(for: .seconds(delays[attempt - 1])) + } + } + } + + if let lastError { + throw DownloadError.failedToDownload( + "\(modelId)/\(path) on ModelScope: \(lastError.localizedDescription)" + ) + } + + completedBytes += sizes[path] ?? 0 + progressHandler?(min(Double(completedBytes) / Double(totalBytes), 1)) + } + + progressHandler?(1.0) + } + + // MARK: - Listing + + /// Recursively lists every file in a ModelScope repo (used for CoreML bundles). + public static func listAllFiles( + modelId: String, + baseURL: String, + revision: String + ) async throws -> [RemoteFile] { + var collected: [RemoteFile] = [] + try await listFiles( + modelId: modelId, + root: nil, + into: &collected, + baseURL: baseURL, + revision: revision + ) + return collected + } + + private static func listFiles( + modelId: String, + root: String?, + into collected: inout [RemoteFile], + baseURL: String, + revision: String + ) async throws { + guard let url = listingURL(modelId: modelId, baseURL: baseURL, revision: revision, root: root) else { + throw DownloadError.failedToDownload("Invalid ModelScope listing URL for \(modelId)") + } + + let (data, response) = try await URLSession.shared.data(from: url) + guard let http = response as? HTTPURLResponse, (200...299).contains(http.statusCode) else { + throw DownloadError.failedToDownload("ModelScope listing failed for \(modelId)") + } + + let payload = try JSONDecoder().decode(APIResponse.self, from: data) + for entry in payload.Data.Files { + if isDirectoryEntry(entry) { + try await listFiles( + modelId: modelId, + root: entry.Path, + into: &collected, + baseURL: baseURL, + revision: revision + ) + } else { + collected.append(RemoteFile(path: entry.Path, size: entry.Size ?? 0)) + } + } + } + + private static func isDirectoryEntry(_ entry: FilesPayload.Entry) -> Bool { + if entry.entryType?.lowercased() == "tree" { return true } + let size = entry.Size ?? 0 + return size == 0 && !entry.Path.contains(".") + } + + // MARK: - Transfer + + /// Streams a single repo file. `onBytes` receives `(bytesWritten, expectedBytes)`. + private static func fetchFile( + modelId: String, + filePath: String, + to destination: URL, + baseURL: String, + revision: String, + onBytes: @escaping (Int64, Int64) -> Void + ) async throws { + guard let url = fileURL(modelId: modelId, baseURL: baseURL, revision: revision, filePath: filePath) else { + throw DownloadError.invalidRemoteFileName(filePath) + } + + var request = URLRequest(url: url) + request.timeoutInterval = 3600 + + let (asyncBytes, response) = try await URLSession.shared.bytes(for: request) + guard let http = response as? HTTPURLResponse else { + throw DownloadError.failedToDownload(filePath) + } + guard (200...299).contains(http.statusCode) else { + throw DownloadError.failedToDownload("\(filePath) HTTP \(http.statusCode)") + } + + let expectedBytes = http.value(forHTTPHeaderField: "Content-Length") + .flatMap(Int64.init) ?? 0 + + if FileManager.default.fileExists(atPath: destination.path) { + try FileManager.default.removeItem(at: destination) + } + FileManager.default.createFile(atPath: destination.path, contents: nil) + let handle = try FileHandle(forWritingTo: destination) + defer { try? handle.close() } + + var buffer = Data() + buffer.reserveCapacity(1_048_576) + var written: Int64 = 0 + + for try await byte in asyncBytes { + try Task.checkCancellation() + buffer.append(byte) + if buffer.count >= 1_048_576 { + try handle.write(contentsOf: buffer) + written += Int64(buffer.count) + buffer.removeAll(keepingCapacity: true) + onBytes(written, expectedBytes) + } + } + if !buffer.isEmpty { + try handle.write(contentsOf: buffer) + written += Int64(buffer.count) + } + onBytes(written, expectedBytes) + } + + // MARK: - URLs + + private static func listingURL( + modelId: String, + baseURL: String, + revision: String, + root: String? + ) -> URL? { + var components = URLComponents(string: "\(baseURL)/api/v1/models/\(modelId)/repo/files") + var items = [ + URLQueryItem(name: "Revision", value: revision), + ] + if let root, !root.isEmpty { + items.append(URLQueryItem(name: "Root", value: root)) + } + components?.queryItems = items + return components?.url + } + + private static func fileURL( + modelId: String, + baseURL: String, + revision: String, + filePath: String + ) -> URL? { + var components = URLComponents(string: "\(baseURL)/api/v1/models/\(modelId)/repo") + components?.queryItems = [ + URLQueryItem(name: "Revision", value: revision), + URLQueryItem(name: "FilePath", value: filePath), + ] + return components?.url + } +} diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/PipelineLLM.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/PipelineLLM.swift new file mode 100644 index 0000000..17b77fc --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/PipelineLLM.swift @@ -0,0 +1,53 @@ +// MARK: - LLM Protocol + +/// Protocol for language model integration with voice pipelines. +/// +/// Conforming types bridge an LLM (local or remote) to the VoicePipeline's +/// ASR → LLM → TTS flow. The pipeline calls `chat()` on a background thread +/// and expects blocking behavior (return when generation is complete). +public protocol PipelineLLM: AnyObject { + /// Generate a response given conversation messages. + /// + /// Called on the pipeline's worker thread (blocking). Emit tokens via + /// `onToken(text, isFinal)` — the pipeline forwards them to TTS. + func chat(messages: [(role: MessageRole, content: String)], + onToken: @escaping (String, Bool) -> Void) + + /// Cancel in-progress generation. Thread-safe. + func cancel() +} + +/// Message roles for LLM conversation. +public enum MessageRole: Int, Sendable { + case system = 0 + case user = 1 + case assistant = 2 + case tool = 3 +} + +// MARK: - Tool Calling + +/// A tool that can be invoked by the LLM during voice pipeline execution. +public struct PipelineTool { + public let name: String + public let description: String + public let handler: (String) -> String + public let cooldown: Int + + /// - Parameters: + /// - name: Tool name (used by LLM to invoke) + /// - description: What the tool does (included in LLM system prompt) + /// - cooldown: Minimum seconds between invocations (0 = no limit) + /// - handler: Synchronous handler `(arguments) -> result`. Called on pipeline worker thread. + public init( + name: String, + description: String, + cooldown: Int = 0, + handler: @escaping (String) -> String + ) { + self.name = name + self.description = description + self.cooldown = cooldown + self.handler = handler + } +} diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Protocols.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Protocols.swift new file mode 100644 index 0000000..0762968 --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/Protocols.swift @@ -0,0 +1,282 @@ +import Foundation + +// MARK: - Model Memory Management + +/// Memory statistics for a loaded model. +public struct ModelMemoryStats: Sendable { + /// Estimated weight memory in bytes + public let weightMemory: Int + /// Current active GPU memory in bytes (MLX only) + public let activeMemory: Int + + public init(weightMemory: Int, activeMemory: Int = 0) { + self.weightMemory = weightMemory + self.activeMemory = activeMemory + } +} + +/// A model that supports explicit memory management. +/// +/// Call `unload()` to release model weights and free GPU memory. +/// After unloading, the model cannot be used for inference until re-loaded. +public protocol ModelMemoryManageable: AnyObject { + /// Whether the model is currently loaded and ready for inference. + var isLoaded: Bool { get } + + /// Release model weights and free GPU memory. + /// + /// After calling this, `isLoaded` returns false and inference methods will fail. + /// To use the model again, create a new instance via `fromPretrained()`. + func unload() + + /// Estimated memory footprint of the loaded model weights in bytes. + /// Returns 0 if the model is not loaded. + var memoryFootprint: Int { get } +} + +// MARK: - Unified Audio Chunk + +/// A chunk of audio produced during streaming synthesis or generation. +public struct AudioChunk: Sendable { + /// PCM audio samples (Float32) + public let samples: [Float] + /// Sample rate in Hz (e.g. 24000) + public let sampleRate: Int + /// Index of the first frame in this chunk + public let frameIndex: Int + /// True if this is the last chunk + public let isFinal: Bool + /// Wall-clock seconds since generation started (nil if not tracked) + public let elapsedTime: Double? + /// Text tokens generated alongside audio (populated on final chunk if available) + public let textTokens: [Int32] + + public init( + samples: [Float], + sampleRate: Int, + frameIndex: Int, + isFinal: Bool, + elapsedTime: Double? = nil, + textTokens: [Int32] = [] + ) { + self.samples = samples + self.sampleRate = sampleRate + self.frameIndex = frameIndex + self.isFinal = isFinal + self.elapsedTime = elapsedTime + self.textTokens = textTokens + } +} + +// MARK: - Aligned Word + +/// A word with its aligned start and end timestamps (in seconds). +public struct AlignedWord: Sendable { + public let text: String + public let startTime: Float + public let endTime: Float + + public init(text: String, startTime: Float, endTime: Float) { + self.text = text + self.startTime = startTime + self.endTime = endTime + } +} + +// MARK: - Speech Generation (TTS) + +/// A text-to-speech model that generates audio from text. +public protocol SpeechGenerationModel: AnyObject { + /// Output sample rate in Hz + var sampleRate: Int { get } + /// Synthesize audio from text (returns full waveform) + func generate(text: String, language: String?) async throws -> [Float] + /// Synthesize audio from text with streaming output. + /// Default implementation wraps `generate()` as a single chunk. + func generateStream(text: String, language: String?) -> AsyncThrowingStream +} + +extension SpeechGenerationModel { + /// Default: wraps `generate()` as a single-chunk stream. + public func generateStream(text: String, language: String?) -> AsyncThrowingStream { + let rate = sampleRate + return AsyncThrowingStream { continuation in + Task { + do { + let samples = try await self.generate(text: text, language: language) + continuation.yield(AudioChunk(samples: samples, sampleRate: rate, frameIndex: 0, isFinal: true)) + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + } + } +} + +// MARK: - Speech Recognition (STT) + +/// A word with its confidence score. +public struct WordConfidence: Sendable { + public let word: String + /// Confidence score (0.0–1.0) derived from mean token log-probability. + public let confidence: Float + + public init(word: String, confidence: Float) { + self.word = word + self.confidence = confidence + } +} + +/// Result of speech recognition including detected language. +public struct TranscriptionResult: Sendable { + public let text: String + /// Detected language (e.g. "english", "russian"). Nil if model doesn't detect. + public let language: String? + /// Confidence score (0.0–1.0). Higher = more confident transcription. + /// Derived from average token log-probability. 0.0 if model doesn't provide. + public let confidence: Float + /// Per-word confidence scores. Nil if model doesn't provide. + public let words: [WordConfidence]? + + public init(text: String, language: String? = nil, confidence: Float = 0.0, words: [WordConfidence]? = nil) { + self.text = text + self.language = language + self.confidence = confidence + self.words = words + } +} + +/// A speech-to-text model that transcribes audio. +public protocol SpeechRecognitionModel: AnyObject { + /// Expected input sample rate in Hz + var inputSampleRate: Int { get } + /// Transcribe audio to text + func transcribe(audio: [Float], sampleRate: Int, language: String?) -> String + /// Transcribe audio to text with language detection + func transcribeWithLanguage(audio: [Float], sampleRate: Int, language: String?) -> TranscriptionResult +} + +/// Default implementation: delegates to transcribe() with no language detection. +public extension SpeechRecognitionModel { + func transcribeWithLanguage(audio: [Float], sampleRate: Int, language: String?) -> TranscriptionResult { + TranscriptionResult(text: transcribe(audio: audio, sampleRate: sampleRate, language: language)) + } +} + +// MARK: - Forced Alignment + +/// A model that aligns text to audio at the word level. +public protocol ForcedAlignmentModel: AnyObject { + /// Align text to audio, returning word-level timestamps + func align(audio: [Float], text: String, sampleRate: Int, language: String?) -> [AlignedWord] +} + +// MARK: - Speech-to-Speech + +/// A speech-to-speech model that generates a spoken response to spoken input. +public protocol SpeechToSpeechModel: AnyObject { + /// Output sample rate in Hz + var sampleRate: Int { get } + /// Generate response audio from input audio (blocking) + func respond(userAudio: [Float]) -> [Float] + /// Generate response audio from input audio with streaming output + func respondStream(userAudio: [Float]) -> AsyncThrowingStream +} + +// MARK: - Voice Activity Detection + +/// A time segment where speech was detected. +public struct SpeechSegment: Sendable { + /// Start time in seconds + public let startTime: Float + /// End time in seconds + public let endTime: Float + + public init(startTime: Float, endTime: Float) { + self.startTime = startTime + self.endTime = endTime + } + + /// Duration in seconds + public var duration: Float { endTime - startTime } +} + +/// A model that detects speech activity regions in audio. +public protocol VoiceActivityDetectionModel: AnyObject { + /// Expected input sample rate in Hz + var inputSampleRate: Int { get } + /// Detect speech segments in audio + func detectSpeech(audio: [Float], sampleRate: Int) -> [SpeechSegment] +} + +/// A streaming VAD that processes fixed-size audio chunks and returns speech probability. +/// +/// Maps directly to speech-core's `sc_vad_vtable_t` for pipeline integration. +public protocol StreamingVADProvider: AnyObject { + /// Expected input sample rate in Hz + var inputSampleRate: Int { get } + /// Number of samples per chunk + var chunkSize: Int { get } + /// Process a single audio chunk, returns speech probability in [0, 1] + func processChunk(_ samples: [Float]) -> Float + /// Reset internal state (LSTM hidden state, context buffer, etc.) + func resetState() +} + +// MARK: - Speaker Diarization + +/// A speech segment with an assigned speaker identity. +public struct DiarizedSegment: Sendable { + /// Start time in seconds + public let startTime: Float + /// End time in seconds + public let endTime: Float + /// Speaker identifier (0-based) + public let speakerId: Int + + public init(startTime: Float, endTime: Float, speakerId: Int) { + self.startTime = startTime + self.endTime = endTime + self.speakerId = speakerId + } + + /// Duration in seconds + public var duration: Float { endTime - startTime } +} + +/// A model that produces speaker embeddings from audio. +public protocol SpeakerEmbeddingModel: AnyObject { + /// Expected input sample rate in Hz + var inputSampleRate: Int { get } + /// Embedding vector dimension + var embeddingDimension: Int { get } + /// Extract a speaker embedding from audio + func embed(audio: [Float], sampleRate: Int) -> [Float] +} + +// MARK: - Speech Enhancement + +/// A model that enhances speech by removing noise. +public protocol SpeechEnhancementModel: AnyObject { + /// Expected input sample rate in Hz + var inputSampleRate: Int { get } + /// Enhance audio by removing noise + func enhance(audio: [Float], sampleRate: Int) throws -> [Float] +} + +/// A model that assigns speaker identities to speech segments. +public protocol SpeakerDiarizationModel: AnyObject { + /// Expected input sample rate in Hz + var inputSampleRate: Int { get } + /// Diarize audio into speaker-labeled segments + func diarize(audio: [Float], sampleRate: Int) -> [DiarizedSegment] +} + +/// A diarization model that also supports extracting a specific speaker's segments +/// using a reference embedding. Not all engines support this (e.g. Sortformer is +/// end-to-end and does not produce speaker embeddings). +public protocol SpeakerExtractionCapable: SpeakerDiarizationModel { + /// Extract segments belonging to a target speaker identified by a reference embedding. + func extractSpeaker(audio: [Float], sampleRate: Int, targetEmbedding: [Float]) -> [SpeechSegment] +} diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/SentencePieceModel.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/SentencePieceModel.swift new file mode 100644 index 0000000..28b9211 --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/SentencePieceModel.swift @@ -0,0 +1,182 @@ +import Foundation + +/// Minimal SentencePiece `.model` (`sentencepiece_model.proto`) reader. +/// +/// Extracts the vocabulary list — `(text, score, type)` for every piece — +/// without requiring a protobuf runtime dependency. Modules build their own +/// encode/decode logic on top: this struct only owns the wire-format parse +/// and the raw piece array. +/// +/// `sentencepiece_model.proto` excerpt: +/// ``` +/// message ModelProto { +/// repeated SentencePiece pieces = 1; // field 1, length-delimited submsg +/// ... +/// } +/// message SentencePiece { +/// optional string piece = 1; // field 1, length-delimited string +/// optional float score = 2; // field 2, fixed32 (wire type 5) +/// optional Type type = 3; // field 3, varint (wire type 0) +/// } +/// ``` +public struct SentencePieceModel: Sendable { + + /// Piece type constants from `sentencepiece_model.proto`. Values not in + /// this enum are surfaced as `.unknown(rawValue)` so callers can apply + /// their own special-token handling. + public enum PieceType: Int32, Sendable { + case normal = 1 + case unknown = 2 + case control = 3 + case userDefined = 4 + case unused = 5 + case byte = 6 + } + + public struct Piece: Sendable, Equatable { + public let text: String + public let score: Float + public let type: Int32 + + public init(text: String, score: Float, type: Int32) { + self.text = text + self.score = score + self.type = type + } + + public var pieceType: PieceType? { PieceType(rawValue: type) } + + public var isControlOrUnknown: Bool { + type == PieceType.control.rawValue || + type == PieceType.unknown.rawValue || + type == PieceType.unused.rawValue || + type == PieceType.byte.rawValue + } + } + + public let pieces: [Piece] + + public var count: Int { pieces.count } + + public subscript(_ id: Int) -> Piece? { + guard id >= 0, id < pieces.count else { return nil } + return pieces[id] + } + + public init(contentsOf url: URL) throws { + let data = try Data(contentsOf: url) + try self.init(data: data) + } + + public init(modelPath: String) throws { + try self.init(contentsOf: URL(fileURLWithPath: modelPath)) + } + + public init(data: Data) throws { + var parsed: [Piece] = [] + var offset = 0 + + while offset < data.count { + let (fieldNumber, wireType, afterTag) = Self.readTag(data: data, offset: offset) + offset = afterTag + + // Top-level field 1 = repeated SentencePiece, length-delimited (wire 2) + guard fieldNumber == 1, wireType == 2 else { + offset = Self.skipField(data: data, offset: offset, wireType: wireType) + continue + } + + let (length, afterLen) = Self.readVarint(data: data, offset: afterTag) + offset = afterLen + let end = offset + length + + var piece = "" + var score: Float = 0 + var type: Int32 = PieceType.normal.rawValue + + var sub = offset + while sub < end { + let (subField, subWire, afterSubTag) = Self.readTag(data: data, offset: sub) + sub = afterSubTag + switch (subField, subWire) { + case (1, 2): // piece string + let (strLen, afterStrLen) = Self.readVarint(data: data, offset: sub) + sub = afterStrLen + if let s = String(data: data[sub..<(sub + strLen)], encoding: .utf8) { + piece = s + } + sub += strLen + case (2, 5): // score (fixed32 / wire type 5) + score = data[sub..<(sub + 4)].withUnsafeBytes { $0.loadUnaligned(as: Float.self) } + sub += 4 + case (3, 0): // type varint + let (typeValue, afterType) = Self.readVarint(data: data, offset: sub) + sub = afterType + type = Int32(typeValue) + default: + sub = Self.skipField(data: data, offset: sub, wireType: subWire) + } + } + + parsed.append(Piece(text: piece, score: score, type: type)) + offset = end + } + + guard !parsed.isEmpty else { + throw SentencePieceModelError.emptyModel + } + self.pieces = parsed + } + + // MARK: - Protobuf wire helpers + + private static func readVarint(data: Data, offset: Int) -> (value: Int, newOffset: Int) { + var result = 0 + var shift = 0 + var off = offset + while off < data.count { + let byte = Int(data[off]) + off += 1 + result |= (byte & 0x7F) << shift + if byte & 0x80 == 0 { break } + shift += 7 + } + return (result, off) + } + + private static func readTag(data: Data, offset: Int) -> (fieldNumber: Int, wireType: Int, newOffset: Int) { + let (tag, newOffset) = readVarint(data: data, offset: offset) + return (tag >> 3, tag & 0x07, newOffset) + } + + private static func skipField(data: Data, offset: Int, wireType: Int) -> Int { + switch wireType { + case 0: + let (_, newOffset) = readVarint(data: data, offset: offset) + return newOffset + case 1: + return offset + 8 + case 2: + let (length, newOffset) = readVarint(data: data, offset: offset) + return newOffset + length + case 5: + return offset + 4 + default: + return data.count + } + } +} + +public enum SentencePieceModelError: Error, CustomStringConvertible { + case emptyModel + case invalidFile(URL) + + public var description: String { + switch self { + case .emptyModel: + return "SentencePiece model contained no pieces" + case .invalidFile(let url): + return "Could not read SentencePiece model at \(url.path)" + } + } +} diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/StreamingAudioPlayer.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/StreamingAudioPlayer.swift new file mode 100644 index 0000000..5be907e --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/StreamingAudioPlayer.swift @@ -0,0 +1,511 @@ +#if canImport(AVFoundation) +import AVFoundation +import os + +/// Lock-free SPSC ring buffer for audio samples. +/// Producer (TTS thread) writes, consumer (audio render thread) reads. +public final class AudioSampleRingBuffer: @unchecked Sendable { + private let buffer: UnsafeMutableBufferPointer + private let capacity: Int + private var writePos: Int = 0 // only written by producer + private var readPos: Int = 0 // only written by consumer + + public init(capacity: Int) { + self.capacity = capacity + let ptr = UnsafeMutablePointer.allocate(capacity: capacity) + ptr.initialize(repeating: 0, count: capacity) + self.buffer = UnsafeMutableBufferPointer(start: ptr, count: capacity) + } + + deinit { + buffer.baseAddress?.deinitialize(count: capacity) + buffer.baseAddress?.deallocate() + } + + /// Number of samples available to read. + public var availableToRead: Int { + let w = writePos + let r = readPos + return w >= r ? w - r : capacity - r + w + } + + /// Number of free slots for writing. + public var availableToWrite: Int { + return capacity - availableToRead - 1 + } + + /// Write samples into the buffer. Returns number actually written. + @discardableResult + public func write(_ samples: [Float]) -> Int { + let count = min(samples.count, availableToWrite) + guard count > 0 else { return 0 } + + samples.withUnsafeBufferPointer { src in + let w = writePos + let firstChunk = min(count, capacity - w) + buffer.baseAddress!.advanced(by: w).update(from: src.baseAddress!, count: firstChunk) + if firstChunk < count { + buffer.baseAddress!.update(from: src.baseAddress!.advanced(by: firstChunk), count: count - firstChunk) + } + } + writePos = (writePos + count) % capacity + return count + } + + /// Read samples from the buffer into dst. Returns number actually read. + @discardableResult + public func read(into dst: UnsafeMutablePointer, count: Int) -> Int { + let available = min(count, availableToRead) + guard available > 0 else { return 0 } + + let r = readPos + let firstChunk = min(available, capacity - r) + dst.update(from: buffer.baseAddress!.advanced(by: r), count: firstChunk) + if firstChunk < available { + dst.advanced(by: firstChunk).update(from: buffer.baseAddress!, count: available - firstChunk) + } + readPos = (readPos + available) % capacity + return available + } + + /// Reset both pointers (call when not actively reading/writing). + public func reset() { + readPos = 0 + writePos = 0 + } +} + +/// Streams TTS audio via AVAudioEngine using an event-driven render callback. +/// +/// Architecture: +/// ``` +/// TTS (producer) → [Ring Buffer] → AVAudioSourceNode render callback (consumer) +/// pre-fill N sec hardware pulls when it needs data +/// ``` +/// +/// The render thread calls our callback when it needs audio. We read from the +/// ring buffer. If the buffer is empty (underflow), we output silence. +/// +/// `preBufferDuration` controls how much audio must accumulate before playback +/// starts. This is the latency-quality tradeoff: +/// - Higher = more resilient to TTS jitter, but more latency +/// - Lower = less latency, but risk of underflow gaps +/// +/// Typical values: +/// - 0s: single-pass TTS (Kokoro) where all audio arrives at once +/// - 2s: streaming TTS (Qwen3-TTS, RTF ~0.53) +public final class StreamingAudioPlayer: @unchecked Sendable { + private var engine: AVAudioEngine? + private var sourceNode: AVAudioSourceNode? + private var format: AVAudioFormat? + private let lock = NSLock() + + private var ringBuffer: AudioSampleRingBuffer? + private var playbackStarted = false + private var generationComplete = false + private var isFirstChunk = true + private var upsampler: AVAudioConverter? + private var preBufferSamples: Int = 0 + public private(set) var totalWritten: Int = 0 + /// Number of samples written for external diagnostics. + public var totalWrittenSamples: Int { totalWritten } + private var totalRead: Int = 0 + + public private(set) var isPlaying = false + private var playbackFinishedFired = false + + /// Pre-buffer duration in seconds. Playback starts after this much audio accumulates. + /// Default 1.0s — sufficient for streaming TTS at RTF < 0.6. + public var preBufferDuration: Double = 1.0 + + /// Callback when all audio has finished playing. + public var onPlaybackFinished: (() -> Void)? + + /// Ring buffer capacity in seconds. Default 30s — enough for any TTS response. + public var ringBufferDuration: Double = 30 + + public init() {} + + // MARK: - Standalone mode + + /// Start playback engine at the given sample rate. + public func start(sampleRate: Double = 24000) throws { + stop() + let eng = AVAudioEngine() + guard let fmt = AVAudioFormat( + commonFormat: .pcmFormatFloat32, + sampleRate: sampleRate, + channels: 1, + interleaved: false + ) else { return } + + setupSourceNode(engine: eng, format: fmt) + try eng.start() + self.engine = eng + self.format = fmt + } + + /// Create a standalone engine at the hardware's native sample rate. + public func ensureStandaloneEngine() { + guard sourceNode == nil else { return } + let eng = AVAudioEngine() + let mixerFormat = eng.mainMixerNode.outputFormat(forBus: 0) + guard let monoFormat = AVAudioFormat( + commonFormat: .pcmFormatFloat32, + sampleRate: mixerFormat.sampleRate, + channels: 1, + interleaved: false + ) else { return } + setupSourceNode(engine: eng, format: monoFormat) + do { + try eng.start() + self.engine = eng + self.format = monoFormat + } catch {} + } + + // MARK: - Attached mode + + /// Attach to an existing AVAudioEngine. + public func attach(to engine: AVAudioEngine, format: AVAudioFormat) { + setupSourceNode(engine: engine, format: format) + self.format = format + } + + /// Start the source node (for use when attaching before engine.start()). + public func startPlayback() { + // Source node is always running once attached — no-op + } + + /// Detach from an external engine. + public func detach(from engine: AVAudioEngine) { + if let node = sourceNode { + engine.disconnectNodeOutput(node) + engine.detach(node) + } + sourceNode = nil + format = nil + upsampler = nil + ringBuffer?.reset() + } + + // MARK: - Audio Scheduling + + /// Write a chunk of audio samples into the ring buffer. + /// If pre-buffer threshold is reached, playback begins automatically. + public func scheduleChunk(_ samples: [Float]) { + guard !samples.isEmpty else { return } + + var output = samples + + // Drop near-silent warmup chunks at start of generation + if isFirstChunk { + var sumSq: Float = 0 + for s in samples { sumSq += s * s } + let rms = sqrt(sumSq / Float(samples.count)) + if rms < 0.005 { return } // Only drop near-silence (codec init noise) + isFirstChunk = false + // 5ms fade-in to prevent pop + if let fmt = format { + let fadeFrames = min(samples.count, Int(fmt.sampleRate * 0.005)) + for i in 0.. 0 { + if (ringBuffer?.availableToRead ?? 0) >= preBufferSamples { + playbackStarted = true + } + } else if preBufferSamples == 0 { + playbackStarted = true + } + lock.unlock() + } + + /// Write samples with resampling from sourceSampleRate to the player's rate. + public func play(samples: [Float], sampleRate: Int) throws { + guard let fmt = format else { return } + if Double(sampleRate) == fmt.sampleRate { + scheduleChunk(samples) + } else { + guard let srcFmt = AVAudioFormat(commonFormat: .pcmFormatFloat32, sampleRate: Double(sampleRate), channels: 1, interleaved: false) else { return } + if upsampler == nil || upsampler?.inputFormat.sampleRate != Double(sampleRate) { + upsampler = AVAudioConverter(from: srcFmt, to: fmt) + } + guard let converter = upsampler else { return } + guard let inputBuffer = AVAudioPCMBuffer(pcmFormat: srcFmt, frameCapacity: AVAudioFrameCount(samples.count)) else { return } + inputBuffer.frameLength = AVAudioFrameCount(samples.count) + samples.withUnsafeBufferPointer { ptr in + inputBuffer.floatChannelData![0].update(from: ptr.baseAddress!, count: samples.count) + } + let outFrameCount = AVAudioFrameCount(Double(samples.count) * fmt.sampleRate / Double(sampleRate)) + guard let outputBuffer = AVAudioPCMBuffer(pcmFormat: fmt, frameCapacity: outFrameCount) else { return } + var consumed = false + var error: NSError? + converter.convert(to: outputBuffer, error: &error) { _, outStatus in + if consumed { outStatus.pointee = .noDataNow; return nil } + consumed = true + outStatus.pointee = .haveData + return inputBuffer + } + let count = Int(outputBuffer.frameLength) + guard count > 0, let data = outputBuffer.floatChannelData else { return } + let resampled = Array(UnsafeBufferPointer(start: data[0], count: count)) + scheduleChunk(resampled) + } + } + + // MARK: - Completion + + /// Signal that TTS generation is complete — no more chunks will arrive. + /// The render callback will drain remaining samples, then fire onPlaybackFinished. + public func markGenerationComplete() { + lock.lock() + generationComplete = true + playbackStarted = true + let hasEngine = sourceNode != nil + let empty = (ringBuffer?.availableToRead ?? 0) == 0 + let written = totalWritten + lock.unlock() + + // No engine or nothing was written — fire immediately + if !hasEngine || (empty && written == 0) { + guard !playbackFinishedFired else { return } + playbackFinishedFired = true + isPlaying = false + onPlaybackFinished?() + return + } + + // Start polling: the render callback normally fires onPlaybackFinished + // when the buffer drains, but if the render thread isn't running (e.g. + // simulator, or audio route change), we poll the buffer to detect + // completion reliably. Works on both device and simulator. + startCompletionPolling() + } + + private var completionPollTimer: DispatchSourceTimer? + private var lastPolledRead: Int = 0 + private var noProgressPolls: Int = 0 + + private func startCompletionPolling() { + completionPollTimer?.cancel() + lastPolledRead = -1 + noProgressPolls = 0 + let timer = DispatchSource.makeTimerSource(queue: .main) + timer.schedule(deadline: .now() + 0.2, repeating: 0.2) + timer.setEventHandler { [weak self] in + guard let self else { return } + // Already fired by render callback — stop polling + guard !self.playbackFinishedFired else { + self.completionPollTimer?.cancel() + self.completionPollTimer = nil + return + } + self.lock.lock() + let complete = self.generationComplete + let remaining = self.ringBuffer?.availableToRead ?? 0 + let read = self.totalRead + let written = self.totalWritten + self.lock.unlock() + + // All samples consumed (or render thread never started reading) + let drained = remaining == 0 && read >= written && written > 0 + // Render thread never started — audio engine not running + let stalled = complete && read == 0 && written > 0 + + // Render thread stalled mid-stream (partial read, no progress for + // 3 consecutive polls = 600 ms). Seen on virtualized macOS CI runners + // and on real iOS when an audio-session interrupt freezes the + // render thread between buffers. + if complete && read > 0 && read < written { + if read == self.lastPolledRead { + self.noProgressPolls += 1 + } else { + self.noProgressPolls = 0 + self.lastPolledRead = read + } + } + let frozen = complete && self.noProgressPolls >= 3 && read > 0 && read < written + + if complete && (drained || stalled || frozen) { + self.completionPollTimer?.cancel() + self.completionPollTimer = nil + guard !self.playbackFinishedFired else { return } + self.playbackFinishedFired = true + self.isPlaying = false + self.onPlaybackFinished?() + } + } + completionPollTimer = timer + timer.resume() + } + + /// Reset for a new generation cycle. + public func resetGeneration() { + completionPollTimer?.cancel() + completionPollTimer = nil + lastPolledRead = -1 + noProgressPolls = 0 + lock.lock() + generationComplete = false + playbackFinishedFired = false + playbackStarted = false + isFirstChunk = true + totalWritten = 0 + totalRead = 0 + ringBuffer?.reset() + lock.unlock() + } + + /// Wait until all audio has finished playing. + public func waitForCompletion() async { + while isPlaying { + try? await Task.sleep(nanoseconds: 50_000_000) // 50ms poll + } + } + + /// Stop immediately. + public func fadeOutAndStop() { + lock.lock() + generationComplete = false + playbackStarted = false + isFirstChunk = true + totalWritten = 0 + totalRead = 0 + ringBuffer?.reset() + lock.unlock() + isPlaying = false + } + + /// Stop and release resources. + public func stop() { + completionPollTimer?.cancel() + completionPollTimer = nil + if let eng = engine, let node = sourceNode { + eng.disconnectNodeOutput(node) + eng.detach(node) + } + engine?.stop() + engine = nil + sourceNode = nil + format = nil + upsampler = nil + lock.lock() + generationComplete = false + playbackStarted = false + isFirstChunk = true + totalWritten = 0 + totalRead = 0 + ringBuffer?.reset() + lock.unlock() + isPlaying = false + } + + // MARK: - Private + + private func setupSourceNode(engine: AVAudioEngine, format: AVAudioFormat) { + let bufferCapacity = Int(format.sampleRate * ringBufferDuration) + let rb = AudioSampleRingBuffer(capacity: bufferCapacity) + self.ringBuffer = rb + self.preBufferSamples = Int(format.sampleRate * preBufferDuration) + + let node = AVAudioSourceNode(format: format) { [weak self] _, _, frameCount, bufferList -> OSStatus in + guard let self else { return noErr } + + let ablPointer = UnsafeMutableAudioBufferListPointer(bufferList) + guard let dst = ablPointer[0].mData?.assumingMemoryBound(to: Float.self) else { + return noErr + } + let frames = Int(frameCount) + + self.lock.lock() + let started = self.playbackStarted + let complete = self.generationComplete + let available = rb.availableToRead + self.lock.unlock() + + if !started { + // Pre-buffer not full yet — output silence + dst.update(repeating: 0, count: frames) + return noErr + } + + if available > 0 { + let read = rb.read(into: dst, count: min(frames, available)) + // Zero-fill remainder if not enough + if read < frames { + dst.advanced(by: read).update(repeating: 0, count: frames - read) + } + self.lock.lock() + self.totalRead += read + self.lock.unlock() + } else if complete && !self.playbackFinishedFired { + // Buffer empty + generation done = playback finished (fire once) + self.playbackFinishedFired = true + dst.update(repeating: 0, count: frames) + DispatchQueue.main.async { + self.isPlaying = false + self.onPlaybackFinished?() + } + } else { + // Underflow — output silence, keep waiting for more data + dst.update(repeating: 0, count: frames) + } + + return noErr + } + + engine.attach(node) + engine.connect(node, to: engine.mainMixerNode, format: format) + self.sourceNode = node + } + + /// Compress long silent gaps to at most `maxSilence` samples. + /// TTS models produce long pauses between sentences (500ms+). + /// This shortens them while keeping a natural brief pause. + static func compressSilence(_ samples: [Float], maxSilence: Int, threshold: Float) -> [Float] { + guard samples.count > maxSilence else { return samples } + + var result = [Float]() + result.reserveCapacity(samples.count) + var silenceRun = 0 + + // Process in small frames (240 samples = 10ms at 24kHz) + let frameSize = 240 + var offset = 0 + + while offset < samples.count { + let end = min(offset + frameSize, samples.count) + let frame = samples[offset..text) and basic BPE encoding (text->ids) via merges.txt +public class Qwen3Tokenizer { + private var idToToken: [Int: String] = [:] + private var tokenToId: [String: Int] = [:] + private var bpeMerges: [(String, String)] = [] + private var bpeMergeRanks: [String: Int] = [:] + + public var eosTokenId: Int = 151643 + public var padTokenId: Int = 151643 + public var bosTokenId: Int = 151644 + + public init() {} + + /// Test-only initializer with pre-built token mappings + internal init(idToToken: [Int: String]) { + self.idToToken = idToToken + for (id, token) in idToToken { tokenToId[token] = id } + } + + /// Load tokenizer from vocab.json file (direct token->id mapping) + public func load(from url: URL) throws { + let data = try Data(contentsOf: url) + + // vocab.json is a direct {token: id} mapping + guard let vocab = try JSONSerialization.jsonObject(with: data) as? [String: Int] else { + throw TokenizerError.invalidFormat("Expected {token: id} dictionary") + } + + for (token, id) in vocab { + idToToken[id] = token + tokenToId[token] = id + } + + // Also load added tokens from tokenizer_config.json if it exists + let configUrl = url.deletingLastPathComponent().appendingPathComponent("tokenizer_config.json") + if FileManager.default.fileExists(atPath: configUrl.path) { + try loadAddedTokens(from: configUrl) + } + + // Load BPE merges if available + let mergesUrl = url.deletingLastPathComponent().appendingPathComponent("merges.txt") + if FileManager.default.fileExists(atPath: mergesUrl.path) { + try loadMerges(from: mergesUrl) + } + + logTokenizer("Loaded tokenizer with \(idToToken.count) tokens, \(bpeMerges.count) merges") + } + + /// Load added tokens from tokenizer_config.json + private func loadAddedTokens(from url: URL) throws { + let data = try Data(contentsOf: url) + + guard let config = try JSONSerialization.jsonObject(with: data) as? [String: Any] else { + return // Not a valid config, skip + } + + // added_tokens_decoder is a dict with string keys (token IDs) and object values with "content" field + if let addedTokens = config["added_tokens_decoder"] as? [String: [String: Any]] { + var addedCount = 0 + for (idString, tokenInfo) in addedTokens { + guard let id = Int(idString), + let content = tokenInfo["content"] as? String else { + continue + } + + // Add to our mappings (overwrite if exists) + idToToken[id] = content + tokenToId[content] = id + addedCount += 1 + } + logTokenizer("Loaded \(addedCount) added tokens from tokenizer_config.json") + } + } + + /// Load BPE merge rules from merges.txt + private func loadMerges(from url: URL) throws { + let content = try String(contentsOf: url, encoding: .utf8) + let lines = content.components(separatedBy: .newlines) + + for (index, line) in lines.enumerated() { + // Skip header line and empty lines + if line.hasPrefix("#") || line.isEmpty { continue } + + let parts = line.components(separatedBy: " ") + guard parts.count == 2 else { continue } + + bpeMerges.append((parts[0], parts[1])) + bpeMergeRanks["\(parts[0]) \(parts[1])"] = index + } + } + + /// Decode token IDs to text using a unified byte buffer. + /// Collects all bytes before converting to UTF-8, so multi-byte characters + /// split across BPE tokens (e.g. CJK) decode correctly. + public func decode(tokens: [Int]) -> String { + var buffer: [UInt8] = [] + + for tokenId in tokens { + guard let token = idToToken[tokenId] else { continue } + + // Skip <|...|> special tokens + if token.hasPrefix("<|") && token.hasSuffix("|>") { + continue + } + + // Keep and similar markers — append their UTF-8 bytes + if token.hasPrefix("<") && token.hasSuffix(">") && !token.contains("|") { + buffer.append(contentsOf: Array(token.utf8)) + continue + } + + // Convert each char via unicodeToByte (Ġ→0x20 space is handled + // automatically since unicodeToByte maps Ġ (U+0120) → byte 32) + for char in token { + if let byte = Self.unicodeToByte[char] { + buffer.append(byte) + } else { + buffer.append(contentsOf: String(char).utf8) + } + } + } + + let text = String(bytes: buffer, encoding: .utf8) + ?? String(decoding: buffer, as: UTF8.self) + return text.trimmingCharacters(in: .whitespaces) + } + + /// Byte-to-unicode mapping table (GPT-2 style) + /// Built lazily on first use + private static var byteToUnicode: [UInt8: Character] = { + var mapping: [UInt8: Character] = [:] + var n = 0 + + // Printable ASCII and some extended chars map directly + let ranges: [(ClosedRange)] = [ + (UInt8(ascii: "!")...UInt8(ascii: "~")), // 33-126 + (0xA1...0xAC), // 161-172 + (0xAE...0xFF), // 174-255 + ] + + for range in ranges { + for b in range { + mapping[b] = Character(UnicodeScalar(b)) + } + } + + // Remaining bytes (0-32, 127-160, 173) map to U+0100 onwards + for b: UInt8 in 0...255 { + if mapping[b] == nil { + mapping[b] = Character(UnicodeScalar(0x100 + n)!) + n += 1 + } + } + + return mapping + }() + + /// Unicode-to-byte reverse mapping + private static var unicodeToByte: [Character: UInt8] = { + var reverse: [Character: UInt8] = [:] + for (byte, char) in byteToUnicode { + reverse[char] = byte + } + return reverse + }() + + /// Encode a byte-level BPE token string from raw text bytes + private func encodeByteLevelToken(_ text: String) -> String { + var result = "" + for byte in text.utf8 { + if let char = Self.byteToUnicode[byte] { + result.append(char) + } + } + return result + } + + /// BPE encode text to token IDs + public func encode(_ text: String) -> [Int] { + guard !bpeMerges.isEmpty else { + // Fallback: character-level encoding + return characterEncode(text) + } + + // Split text into words (whitespace-aware, GPT-2 style pre-tokenization) + // Simple approach: split on word boundaries, preserving leading spaces as Ġ + let words = preTokenize(text) + + var tokens: [Int] = [] + for word in words { + // Convert word to byte-level BPE representation + let bpeTokens = bpe(word) + for bpeToken in bpeTokens { + if let id = tokenToId[bpeToken] { + tokens.append(id) + } + } + } + + return tokens + } + + /// Pre-tokenize text into words (GPT-2 style) + private func preTokenize(_ text: String) -> [String] { + // Split on whitespace boundaries while preserving leading spaces as part of the next word + var words: [String] = [] + var current = "" + + for char in text { + if char == " " || char == "\n" || char == "\t" { + if !current.isEmpty { + words.append(encodeByteLevelToken(current)) + current = "" + } + current = String(char) + } else { + current.append(char) + } + } + if !current.isEmpty { + words.append(encodeByteLevelToken(current)) + } + + return words + } + + /// Apply BPE merges to a word + private func bpe(_ word: String) -> [String] { + var pieces = word.map { String($0) } + + while pieces.count > 1 { + // Find the pair with lowest merge rank + var bestPair: (String, String)? + var bestRank = Int.max + + for i in 0..<(pieces.count - 1) { + let pair = "\(pieces[i]) \(pieces[i + 1])" + if let rank = bpeMergeRanks[pair], rank < bestRank { + bestRank = rank + bestPair = (pieces[i], pieces[i + 1]) + } + } + + guard let (first, second) = bestPair else { break } + + // Merge the pair + var newPieces: [String] = [] + var i = 0 + while i < pieces.count { + if i < pieces.count - 1 && pieces[i] == first && pieces[i + 1] == second { + newPieces.append(first + second) + i += 2 + } else { + newPieces.append(pieces[i]) + i += 1 + } + } + pieces = newPieces + } + + return pieces + } + + /// Simple character-level encoding fallback + private func characterEncode(_ text: String) -> [Int] { + var tokens: [Int] = [] + for char in text { + if let id = tokenToId[String(char)] { + tokens.append(id) + } + } + return tokens + } + + /// Get token ID for a specific token string + public func getTokenId(for token: String) -> Int? { + return tokenToId[token] + } + + /// Get token string for a specific ID + public func getToken(for id: Int) -> String? { + return idToToken[id] + } + + /// Debug: print token mappings for common words + public func debugTokenMappings() { + let commonTokens = [ + "<|im_start|>", "<|im_end|>", "<|audio_start|>", "<|audio_end|>", + "<|audio_pad|>", "", "<|endoftext|>", + "system", "user", "assistant", "language", "English", + "Ġsystem", "Ġuser", "Ġassistant", "Ġlanguage", "ĠEnglish", + "\n", "Ċ" // newline representations + ] + + print("Token ID mappings:") + for token in commonTokens { + if let id = tokenToId[token] { + print(" '\(token)' -> \(id)") + } else { + print(" '\(token)' -> NOT FOUND") + } + } + } +} + +/// Protocol for tokenizer to allow different implementations +public protocol TokenizerProtocol { + func decode(tokens: [Int]) -> String + func encode(_ text: String) -> [Int] +} + +extension Qwen3Tokenizer: TokenizerProtocol {} diff --git a/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/WAVWriter.swift b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/WAVWriter.swift new file mode 100644 index 0000000..d190329 --- /dev/null +++ b/OSGKeyboard/ThirdParty/Qwen3Speech/Sources/AudioCommon/WAVWriter.swift @@ -0,0 +1,105 @@ +import Foundation + +/// Write float audio samples to WAV file +public enum WAVWriter { + + /// Write mono float samples to a 16-bit PCM WAV file + /// - Parameters: + /// - samples: Float audio samples in [-1.0, 1.0] range + /// - sampleRate: Sample rate in Hz (default 24000) + /// - url: Output file URL + public static func write(samples: [Float], sampleRate: Int = 24000, to url: URL) throws { + let numChannels: UInt16 = 1 + let bitsPerSample: UInt16 = 16 + let bytesPerSample = Int(bitsPerSample) / 8 + let dataSize = samples.count * bytesPerSample + let fileSize = 36 + dataSize + + var data = Data(capacity: fileSize + 8) + + // RIFF header + data.append(contentsOf: "RIFF".utf8) + appendUInt32(&data, UInt32(fileSize)) + data.append(contentsOf: "WAVE".utf8) + + // fmt chunk + data.append(contentsOf: "fmt ".utf8) + appendUInt32(&data, 16) // chunk size + appendUInt16(&data, 1) // PCM format + appendUInt16(&data, numChannels) + appendUInt32(&data, UInt32(sampleRate)) + appendUInt32(&data, UInt32(sampleRate * Int(numChannels) * bytesPerSample)) // byte rate + appendUInt16(&data, numChannels * UInt16(bytesPerSample)) // block align + appendUInt16(&data, bitsPerSample) + + // data chunk + data.append(contentsOf: "data".utf8) + appendUInt32(&data, UInt32(dataSize)) + + // Convert float samples to 16-bit PCM + for sample in samples { + let clamped = max(-1.0, min(1.0, sample)) + let int16Value = Int16(clamped * 32767.0) + appendInt16(&data, int16Value) + } + + try data.write(to: url) + } + + /// Write stereo float samples to a 16-bit PCM WAV file. + /// - Parameters: + /// - left: Left channel float samples in [-1.0, 1.0] + /// - right: Right channel float samples in [-1.0, 1.0] + /// - sampleRate: Sample rate in Hz + /// - url: Output file URL + public static func writeStereo(left: [Float], right: [Float], sampleRate: Int = 44100, to url: URL) throws { + let numChannels: UInt16 = 2 + let bitsPerSample: UInt16 = 16 + let bytesPerSample = Int(bitsPerSample) / 8 + let frameCount = min(left.count, right.count) + let dataSize = frameCount * Int(numChannels) * bytesPerSample + let fileSize = 36 + dataSize + + var data = Data(capacity: fileSize + 8) + + data.append(contentsOf: "RIFF".utf8) + appendUInt32(&data, UInt32(fileSize)) + data.append(contentsOf: "WAVE".utf8) + + data.append(contentsOf: "fmt ".utf8) + appendUInt32(&data, 16) + appendUInt16(&data, 1) // PCM + appendUInt16(&data, numChannels) + appendUInt32(&data, UInt32(sampleRate)) + appendUInt32(&data, UInt32(sampleRate * Int(numChannels) * bytesPerSample)) + appendUInt16(&data, numChannels * UInt16(bytesPerSample)) + appendUInt16(&data, bitsPerSample) + + data.append(contentsOf: "data".utf8) + appendUInt32(&data, UInt32(dataSize)) + + for i in 0..