diff --git a/Sources/App/AppState.swift b/Sources/App/AppState.swift index 31a99d0..5d335dc 100644 --- a/Sources/App/AppState.swift +++ b/Sources/App/AppState.swift @@ -4,6 +4,7 @@ import Combine enum AppPhase: Equatable { case idle case downloading + case loadingModel case recording case transcribing case processing diff --git a/Sources/App/VoicePipeline+Models.swift b/Sources/App/VoicePipeline+Models.swift index 0d416ce..b0ca06b 100644 --- a/Sources/App/VoicePipeline+Models.swift +++ b/Sources/App/VoicePipeline+Models.swift @@ -46,50 +46,25 @@ extension VoicePipeline { let catalog = ModelCatalog.shared catalog.refreshStatus() let modelStatus = catalog.llmModels.first(where: { $0.id == model })?.status - let shouldShowDownload = !(modelStatus == .downloaded || modelStatus == .ready) - if shouldShowDownload { - appState.phase = .downloading - appState.statusMessage = L("pipeline.downloading") - appState.resetDownloadProgress() - catalog.updateLLMStatus(model, status: .downloading) - } else { - appState.statusMessage = L("pipeline.loading_llm") - catalog.updateLLMStatus(model, status: .loading, detail: L("model.loading")) + guard modelStatus == .downloaded || modelStatus == .ready else { + let message = L("model.download_required") + catalog.updateLLMStatus(model, status: .error(message)) + appState.statusMessage = showFailureInStatus ? message : L("status.ready") + return } - let estimatedDownloadBytes = catalog.estimatedLLMDownloadBytes(model) - let loaded = await textProcessor.warmUpLLM( - model: model, - estimatedDownloadBytes: estimatedDownloadBytes - ) { [weak self] info in - guard shouldShowDownload else { return } - Task { @MainActor in - guard let self, self.appState.isDownloading else { return } - self.appState.updateDownloadProgress(info) - self.appState.statusMessage = L("pipeline.downloading") + " \(info.percentText)" - if let index = catalog.llmModels.firstIndex(where: { $0.id == model }) { - catalog.llmModels[index].status = .downloading - catalog.llmModels[index].downloadProgress = info.fraction - catalog.llmModels[index].downloadDetail = info.detailText - } - } - } + appState.statusMessage = L("pipeline.loading_llm") + catalog.updateLLMStatus(model, status: .loading, detail: L("model.loading")) + + let loaded = await textProcessor.warmUpLLM(model: model) let ready = await textProcessor.isLLMReady appState.llmModelReady = loaded && ready if appState.llmModelReady { - if shouldShowDownload { - appState.phase = .idle - appState.resetDownloadProgress() - } catalog.updateLLMStatus(model, status: .ready) Log.info("[VoicePipeline] LLM model loaded into memory, ready for instant inference") appState.statusMessage = L("status.ready") } else { - if shouldShowDownload { - appState.phase = .idle - appState.resetDownloadProgress() - } catalog.updateLLMStatus(model, status: .error(L("pipeline.model_load_failed"))) Log.info("[VoicePipeline] LLM warmup failed, will retry on demand") appState.statusMessage = showFailureInStatus ? L("pipeline.model_load_failed") : L("status.ready") @@ -119,17 +94,6 @@ extension VoicePipeline { case .qwen3: let settings = appState.settings let catalog = ModelCatalog.shared - if !LocalASRRuntime.isReady(for: .qwen3), - !catalog.asrModelPath(for: settings.qwenASRModel).isEmpty { - do { - _ = try await LocalASRRuntime.ensurePythonPath( - for: .qwen3, - preferredPath: settings.localASRPythonPath - ) - } catch { - Log.error("[VoicePipeline] Qwen ASR runtime migration failed: \(error.localizedDescription)") - } - } guard localASRIsAvailable(settings.qwenASRModel) else { qwenSpeechEngine = nil markSpeechModelDownloadRequired(showInStatus: requestPermission) @@ -191,10 +155,10 @@ extension VoicePipeline { let engine = WhisperEngine(modelName: modelID) whisperEngine = engine - appState.phase = .downloading + appState.phase = .loadingModel appState.statusMessage = L("pipeline.preparing_model") appState.resetDownloadProgress() - catalog.updateWhisperStatus(modelID, status: .downloading) + catalog.updateWhisperStatus(modelID, status: .loading, detail: L("model.loading")) do { try await engine.loadModel { [weak self] progress in @@ -204,21 +168,10 @@ extension VoicePipeline { switch progress.stage { case .downloading: - if alreadyDownloaded { - self.appState.statusMessage = L("pipeline.loading_model") - self.appState.resetDownloadProgress() - self.appState.downloadProgress = progress.fraction - catalog.updateWhisperStatus(modelID, status: .loading, detail: L("model.loading")) - } else { - self.appState.updateDownloadProgress(progress.info) - self.appState.downloadProgress = progress.fraction - self.appState.statusMessage = L("pipeline.downloading") + " \(progress.info.percentText)" - catalog.updateWhisperStatus( - modelID, - status: .downloading, - detail: progress.detailText - ) - } + self.appState.statusMessage = L("pipeline.loading_model") + self.appState.resetDownloadProgress() + self.appState.downloadProgress = progress.fraction + catalog.updateWhisperStatus(modelID, status: .loading, detail: L("model.loading")) case .compiling: self.appState.statusMessage = L("pipeline.loading_model") self.appState.resetDownloadProgress() diff --git a/Sources/App/VoicePipeline.swift b/Sources/App/VoicePipeline.swift index 7feb546..ad63493 100644 --- a/Sources/App/VoicePipeline.swift +++ b/Sources/App/VoicePipeline.swift @@ -36,14 +36,19 @@ final class VoicePipeline { func warmUp() async { let settings = appState.settings + let catalog = ModelCatalog.shared + catalog.refreshStatus(recheckingErrors: true) + let llmStatus = catalog.llmModels.first(where: { $0.id == settings.llmModel })?.status let shouldLoadSpeech = StartupModelPreloadPolicy.shouldPreloadSpeechModel( enabled: settings.preloadSpeechModelOnLaunch, - speechEngine: settings.speechEngine + speechEngine: settings.speechEngine, + modelDownloaded: catalog.isWhisperDownloaded(settings.whisperModel) ) let shouldLoadFormatting = StartupModelPreloadPolicy.shouldPreloadFormattingModel( enabled: settings.preloadFormattingModelOnLaunch, useRemoteLLM: settings.useRemoteLLM, - modelID: settings.llmModel + modelID: settings.llmModel, + modelDownloaded: llmStatus == .downloaded || llmStatus == .ready ) if shouldLoadSpeech { @@ -184,16 +189,21 @@ final class VoicePipeline { enum StartupModelPreloadPolicy { static func shouldPreloadSpeechModel( enabled: Bool, - speechEngine: SpeechEngineType + speechEngine: SpeechEngineType, + modelDownloaded: Bool ) -> Bool { - enabled && speechEngine == .whisper + enabled && speechEngine == .whisper && modelDownloaded } static func shouldPreloadFormattingModel( enabled: Bool, useRemoteLLM: Bool, - modelID: String + modelID: String, + modelDownloaded: Bool ) -> Bool { - enabled && !useRemoteLLM && !modelID.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + enabled && + !useRemoteLLM && + modelDownloaded && + !modelID.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty } } diff --git a/Sources/Config/ModelCatalog.swift b/Sources/Config/ModelCatalog.swift index 9bf7cdf..dae8973 100644 --- a/Sources/Config/ModelCatalog.swift +++ b/Sources/Config/ModelCatalog.swift @@ -1,7 +1,5 @@ import Foundation import WhisperKit -import MLXLMCommon -import MLXLLM @MainActor final class ModelCatalog: ObservableObject { @@ -11,7 +9,8 @@ final class ModelCatalog: ObservableObject { @Published var llmModels: [ModelEntry] = [] @Published var asrModels: [ModelEntry] = [] - private let settings = AppSettings.shared + let settings = AppSettings.shared + let downloadTasks = ModelDownloadTasks() /// LLM model family categories enum ModelFamily: String, CaseIterable { @@ -65,7 +64,9 @@ final class ModelCatalog: ObservableObject { } enum ModelStatus: Equatable { - case notDownloaded, downloading, compiling, loading, downloaded, ready, error(String) + case notDownloaded, downloading, compiling, loading, downloaded, ready + case unavailable(String) + case error(String) var isDownloading: Bool { if case .downloading = self { return true }; return false } var isError: Bool { if case .error = self { return true }; return false } @@ -73,7 +74,10 @@ final class ModelCatalog: ObservableObject { switch self { case .downloading, .compiling, .loading: return true; default: return false } } var canDelete: Bool { - switch self { case .downloaded, .ready, .error: return true; default: return false } + switch self { + case .downloaded, .ready, .unavailable, .error: return true + default: return false + } } } @@ -178,10 +182,13 @@ final class ModelCatalog: ObservableObject { let size = whisperVariantSize(id) whisperModels[i].cacheSize = size if recheckingErrors || (whisperModels[i].status != .ready && !whisperModels[i].status.isError) { - if ModelStorage.localWhisperURL(id) != nil, size == 0 { + let isComplete = isWhisperDownloaded(id) + if ModelStorage.localWhisperURL(id) != nil, !isComplete { whisperModels[i].status = .error(L("model.local_missing")) } else { - whisperModels[i].status = size > 0 ? .downloaded : .notDownloaded + whisperModels[i].status = isComplete + ? .downloaded + : (size > 0 ? .error(L("model.download_incomplete")) : .notDownloaded) } } } @@ -190,163 +197,19 @@ final class ModelCatalog: ObservableObject { let size = llmRepoSize(id) llmModels[i].cacheSize = size if recheckingErrors || (llmModels[i].status != .ready && !llmModels[i].status.isError) { - if ModelStorage.localLLMURL(id) != nil, !llmRepoHasConfig(id) { + let isComplete = llmRepoIsComplete(id) + if ModelStorage.localLLMURL(id) != nil, !isComplete { llmModels[i].status = .error(L("model.local_missing")) } else { - llmModels[i].status = size > 0 ? .downloaded : .notDownloaded + llmModels[i].status = isComplete + ? .downloaded + : (size > 0 ? .error(L("model.download_incomplete")) : .notDownloaded) } } } refreshASRStatus(recheckingErrors: recheckingErrors) } - // MARK: - Whisper Operations - - func downloadWhisper(_ id: String) async { - guard let idx = whisperModels.firstIndex(where: { $0.id == id }), - !whisperModels[idx].status.isDownloading else { return } - - whisperModels[idx].status = .downloading - whisperModels[idx].downloadProgress = 0 - - do { - let modelDir = ModelStorage.whisperVariantDir(id) - let tracker = DownloadProgressTracker(initialBytes: ModelStorage.directorySize(at: modelDir)) - - _ = try await WhisperKit.download( - variant: id, - downloadBase: Self.whisperDownloadBase, - progressCallback: { [weak self] p in - Task { @MainActor in - guard let self, let i = self.whisperModels.firstIndex(where: { $0.id == id }) else { return } - let completedBytes = ModelStorage.directorySize(at: modelDir) - let info = tracker.update( - completedBytes: completedBytes > 0 ? completedBytes : p.completedUnitCount, - totalBytes: p.totalUnitCount, - fraction: p.fractionCompleted - ) - self.whisperModels[i].downloadProgress = info.fraction - self.whisperModels[i].downloadDetail = info.detailText - } - } - ) - whisperModels[idx].status = .downloaded - whisperModels[idx].cacheSize = whisperVariantSize(id) - whisperModels[idx].downloadDetail = "" - } catch is CancellationError { - if let i = whisperModels.firstIndex(where: { $0.id == id }) { - whisperModels[i].status = .notDownloaded - whisperModels[i].downloadDetail = "" - } - } catch { - whisperModels[idx].status = .error(error.localizedDescription) - whisperModels[idx].downloadDetail = "" - } - } - - func deleteWhisper(_ id: String) { - guard let idx = whisperModels.firstIndex(where: { $0.id == id }) else { return } - if ModelStorage.localWhisperURL(id) != nil { - var paths = settings.localWhisperModelPaths - paths.removeValue(forKey: id) - settings.localWhisperModelPaths = paths - whisperModels.remove(at: idx) - } else { - deleteWhisperVariant(id) - whisperModels[idx].status = .notDownloaded - whisperModels[idx].cacheSize = 0 - } - - if settings.whisperModel == id { - settings.whisperModel = nextAvailableWhisper(excluding: id) ?? whisperModels.first?.id ?? "" - } - } - - func nextAvailableWhisper(excluding id: String) -> String? { - whisperModels.first { $0.id != id && ($0.status == .downloaded || $0.status == .ready) }?.id - } - - // MARK: - LLM Operations - - func downloadLLM(_ id: String) async { - guard let idx = llmModels.firstIndex(where: { $0.id == id }), - !llmModels[idx].status.isDownloading else { return } - if ModelStorage.localLLMURL(id) != nil { - llmModels[idx].status = llmRepoHasConfig(id) ? .downloaded : .error(L("model.local_missing")) - llmModels[idx].cacheSize = llmRepoSize(id) - return - } - if let dir = llmRepoDir(id), !llmRepoHasConfig(id) { - Log.info("[ModelCatalog] removing incomplete LLM cache before redownload: \(dir.path)") - try? FileManager.default.removeItem(at: dir) - } - - llmModels[idx].status = .downloading - llmModels[idx].downloadProgress = 0 - - do { - let estimatedTotalBytes = estimatedLLMDownloadBytes(id) ?? 0 - let repoDir = ModelStorage.hubModelRepoDir(id) - let tracker = DownloadProgressTracker(initialBytes: ModelStorage.directorySize(at: repoDir)) - let config = ModelConfiguration(id: id) - _ = try await LLMModelFactory.shared.loadContainer( - from: MLXModelLoading.downloader, - using: MLXModelLoading.tokenizerLoader, - configuration: config - ) { [weak self] p in - Task { @MainActor in - guard let self, let i = self.llmModels.firstIndex(where: { $0.id == id }) else { return } - let completedBytes = ModelStorage.directorySize(at: repoDir) - let info = tracker.update( - completedBytes: completedBytes, - totalBytes: estimatedTotalBytes, - fraction: p.fractionCompleted - ) - self.llmModels[i].downloadProgress = info.fraction - self.llmModels[i].downloadDetail = info.detailText - } - } - if let i = llmModels.firstIndex(where: { $0.id == id }) { - llmModels[i].status = .downloaded - llmModels[i].cacheSize = llmRepoSize(id) - llmModels[i].downloadDetail = "" - } - } catch is CancellationError { - if let i = llmModels.firstIndex(where: { $0.id == id }) { - llmModels[i].status = .notDownloaded - llmModels[i].downloadDetail = "" - } - } catch { - if let i = llmModels.firstIndex(where: { $0.id == id }) { - llmModels[i].status = .error(error.localizedDescription) - llmModels[i].cacheSize = llmRepoSize(id) - llmModels[i].downloadDetail = "" - } - } - } - - func deleteLLM(_ id: String) { - guard let idx = llmModels.firstIndex(where: { $0.id == id }) else { return } - if ModelStorage.localLLMURL(id) != nil { - var paths = settings.localLLMModelPaths - paths.removeValue(forKey: id) - settings.localLLMModelPaths = paths - llmModels.remove(at: idx) - } else { - if let dir = llmRepoDir(id) { try? FileManager.default.removeItem(at: dir) } - llmModels[idx].status = .notDownloaded - llmModels[idx].cacheSize = 0 - } - - if settings.llmModel == id { - settings.llmModel = nextAvailableLLM(excluding: id) ?? llmModels.first?.id ?? "" - } - } - - func nextAvailableLLM(excluding id: String) -> String? { - llmModels.first { $0.id != id && ($0.status == .downloaded || $0.status == .ready) }?.id - } - func addCustomLLM(_ modelID: String) { guard !modelID.isEmpty, !llmModels.contains(where: { $0.id == modelID }) else { return } let name = modelID.components(separatedBy: "/").last ?? modelID @@ -405,46 +268,6 @@ final class ModelCatalog: ObservableObject { return "\(bytes) B" } - static var whisperDownloadBase: URL { ModelStorage.huggingFaceBase } - - // MARK: Whisper cache - - private func whisperVariantDir(_ variant: String) -> URL { - ModelStorage.localWhisperURL(variant) ?? ModelStorage.whisperVariantDir(variant) - } - - private func whisperVariantSize(_ variant: String) -> Int64 { - let dir = whisperVariantDir(variant) - guard FileManager.default.fileExists(atPath: dir.path) else { return 0 } - return ModelStorage.directorySize(at: dir) - } - - /// Returns true if the Whisper variant is already downloaded (skip download progress UI). - func isWhisperDownloaded(_ variant: String) -> Bool { - whisperVariantSize(variant) > 0 - } - - private func deleteWhisperVariant(_ variant: String) { - let dir = whisperVariantDir(variant) - try? FileManager.default.removeItem(at: dir) - } - - // MARK: LLM cache - - private func llmRepoDir(_ modelID: String) -> URL? { - ModelStorage.llmRepoDir(modelID) - } - - private func llmRepoHasConfig(_ modelID: String) -> Bool { - guard let dir = llmRepoDir(modelID) else { return false } - return FileManager.default.fileExists(atPath: dir.appendingPathComponent("config.json").path) - } - - private func llmRepoSize(_ modelID: String) -> Int64 { - guard let dir = llmRepoDir(modelID), llmRepoHasConfig(modelID) else { return 0 } - return ModelStorage.directorySize(at: dir) - } - private func appendLocalWhisperModels() { for (id, path) in settings.localWhisperModelPaths.sorted(by: { $0.key < $1.key }) { guard !whisperModels.contains(where: { $0.id == id }) else { continue } diff --git a/Sources/Config/ModelCatalogASR.swift b/Sources/Config/ModelCatalogASR.swift index e68b548..0138939 100644 --- a/Sources/Config/ModelCatalogASR.swift +++ b/Sources/Config/ModelCatalogASR.swift @@ -31,6 +31,11 @@ extension ModelCatalog { Self.defaultASRModels.first { $0.id == id }?.provider } + func asrRuntimeAvailability(for id: String) -> LocalASRRuntimeAvailability { + guard let provider = asrProvider(for: id) else { return .supported } + return LocalASRRuntime.availability(for: provider) + } + func asrModelPath(for id: String) -> String { asrSingleRepoIsComplete(id) ? ModelStorage.asrRepoDir(id)?.path ?? "" : "" } @@ -50,6 +55,10 @@ extension ModelCatalog { let id = asrModels[i].id let size = asrRepoSize(id) asrModels[i].cacheSize = size + if case .unavailable(let message) = asrRuntimeAvailability(for: id) { + asrModels[i].status = .unavailable(message) + continue + } if recheckingErrors || (asrModels[i].status != .ready && !asrModels[i].status.isError) { asrModels[i].status = asrRepoIsComplete(id) ? .downloaded @@ -59,7 +68,20 @@ extension ModelCatalog { } func downloadASR(_ id: String, onProgress: ((DownloadProgressInfo) -> Void)? = nil) async { + await downloadTasks.run(key: ModelDownloadKey(kind: .asr, modelID: id)) { [weak self] in + await self?.performASRDownload(id, onProgress: onProgress) + } + } + + private func performASRDownload( + _ id: String, + onProgress: ((DownloadProgressInfo) -> Void)? + ) async { guard let idx = asrModels.firstIndex(where: { $0.id == id }), !asrModels[idx].status.isDownloading else { return } + if case .unavailable(let message) = asrRuntimeAvailability(for: id) { + asrModels[idx].status = .unavailable(message) + return + } if asrRepoIsComplete(id) { asrModels[idx].status = .downloaded @@ -77,17 +99,15 @@ extension ModelCatalog { let api = HubApi(downloadBase: Self.asrDownloadBase) let startedAt = Date() let estimatedTotalBytes = estimatedASRDownloadBytes(id) ?? 0 - if asrProvider(for: id) == .mimo { - asrModels[idx].downloadDetail = L("model.asr_preparing_runtime") - try await ensureMimoRepository() - } if !asrModelFilesAreComplete(id) { let tracker = DownloadProgressTracker(startDate: startedAt, initialBytes: asrRepoSize(id)) for (repoIndex, repoID) in repos.enumerated() { _ = try await api.snapshot(from: ModelStorage.hubModelRepo(repoID)) { [weak self] progress in Task { @MainActor in guard let self, let i = self.asrModels.firstIndex(where: { $0.id == id }) else { return } - let fraction = (Double(repoIndex) + progress.fractionCompleted) / Double(repos.count) + let repositoryFraction = + (Double(repoIndex) + progress.fractionCompleted) / Double(repos.count) + let fraction = repositoryFraction * 0.9 let completedBytes = self.asrRepoSize(id) let info = tracker.update( completedBytes: completedBytes, @@ -102,12 +122,14 @@ extension ModelCatalog { } } if asrProvider(for: id) == .qwen3 { + asrModels[idx].downloadProgress = max(asrModels[idx].downloadProgress, 0.9) asrModels[idx].downloadDetail = L("model.asr_installing_runtime") _ = try await LocalASRRuntime.ensurePythonPath( for: .qwen3, preferredPath: AppSettings.shared.localASRPythonPath ) } + try Task.checkCancellation() if let i = asrModels.firstIndex(where: { $0.id == id }) { asrModels[i].status = asrRepoIsComplete(id) ? .downloaded : .error(L("model.asr_incomplete")) asrModels[i].cacheSize = asrRepoSize(id) @@ -115,14 +137,15 @@ extension ModelCatalog { } } catch is CancellationError { if let i = asrModels.firstIndex(where: { $0.id == id }) { - asrModels[i].status = asrRepoIsComplete(id) - ? .downloaded - : asrMissingStatus(for: id, size: asrRepoSize(id)) + asrModels[i].status = .error(L("model.download_paused")) + asrModels[i].cacheSize = asrRepoSize(id) + asrModels[i].downloadProgress = 0 asrModels[i].downloadDetail = "" } } catch { if let i = asrModels.firstIndex(where: { $0.id == id }) { - asrModels[i].status = .error(error.localizedDescription) + Log.error("[ModelCatalog] ASR download failed: \(error.localizedDescription)") + asrModels[i].status = .error(ModelDownloadFailureMessage.userFacing(error)) asrModels[i].cacheSize = asrRepoSize(id) asrModels[i].downloadDetail = "" } @@ -138,8 +161,11 @@ extension ModelCatalog { if provider == .mimo { try? FileManager.default.removeItem(at: ModelStorage.mimoASRRepositoryDir()) } - asrModels[idx].status = .notDownloaded + if let provider { + try? FileManager.default.removeItem(at: ModelStorage.localASRRuntimeDir(for: provider)) + } asrModels[idx].cacheSize = 0 + asrModels[idx].status = asrMissingStatus(for: id, size: 0) asrModels[idx].downloadDetail = "" let settings = AppSettings.shared @@ -172,7 +198,9 @@ extension ModelCatalog { case .qwen3: return hasFiles && LocalASRRuntime.isReady(for: .qwen3) case .mimo: - return hasFiles && Self.mimoRepositoryIsReady(at: ModelStorage.mimoASRRepositoryDir()) + return hasFiles && + Self.mimoRepositoryIsReady(at: ModelStorage.mimoASRRepositoryDir()) && + LocalASRRuntime.isReady(for: .mimo) case nil: return hasFiles } @@ -183,6 +211,9 @@ extension ModelCatalog { } private func asrMissingStatus(for id: String, size: Int64) -> ModelStatus { + if case .unavailable(let message) = asrRuntimeAvailability(for: id) { + return .unavailable(message) + } guard size > 0 else { return .notDownloaded } if asrModelFilesAreComplete(id), asrProvider(for: id) == .qwen3 { return .error(L("model.asr_runtime_missing")) @@ -192,8 +223,12 @@ extension ModelCatalog { private func asrRepoSize(_ id: String) -> Int64 { let modelSize = asrRequiredRepoIDs(for: id).reduce(Int64(0)) { $0 + asrSingleRepoSize($1) } - guard asrProvider(for: id) == .mimo else { return modelSize } - return modelSize + ModelStorage.directorySize(at: ModelStorage.mimoASRRepositoryDir()) + guard let provider = asrProvider(for: id) else { return modelSize } + let runtimeSize = ModelStorage.directorySize(at: ModelStorage.localASRRuntimeDir(for: provider)) + let repositorySize = provider == .mimo + ? ModelStorage.directorySize(at: ModelStorage.mimoASRRepositoryDir()) + : 0 + return modelSize + runtimeSize + repositorySize } private func asrSingleRepoIsComplete(_ id: String) -> Bool { @@ -248,57 +283,4 @@ extension ModelCatalog { } } - static func mimoRepositoryIsReady(at dir: URL) -> Bool { - FileManager.default.fileExists( - atPath: dir.appendingPathComponent("src/mimo_audio/mimo_audio.py").path - ) - } - - private func ensureMimoRepository() async throws { - let dir = ModelStorage.mimoASRRepositoryDir() - if Self.mimoRepositoryIsReady(at: dir) { return } - - let parent = dir.deletingLastPathComponent() - let temp = parent.appendingPathComponent(".MiMo-V2.5-ASR.download") - try? FileManager.default.removeItem(at: temp) - try FileManager.default.createDirectory(at: parent, withIntermediateDirectories: true) - - try await runProcess( - executable: "/usr/bin/env", - arguments: [ - "git", "clone", "--quiet", "--depth", "1", - LocalASRConfiguration.mimoRepositoryURL, - temp.path, - ] - ) - - try? FileManager.default.removeItem(at: dir) - try FileManager.default.moveItem(at: temp, to: dir) - guard Self.mimoRepositoryIsReady(at: dir) else { throw ASRDownloadError.incompleteRuntime } - } - - private func runProcess(executable: String, arguments: [String]) async throws { - try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in - let process = Process() - let stderr = Pipe() - process.executableURL = URL(fileURLWithPath: executable) - process.arguments = arguments - process.standardError = stderr - process.terminationHandler = { process in - let data = stderr.fileHandleForReading.readDataToEndOfFile() - let message = String(data: data, encoding: .utf8)? - .trimmingCharacters(in: .whitespacesAndNewlines) - guard process.terminationStatus == 0 else { - continuation.resume(throwing: ASRDownloadError.processFailed(message ?? "")) - return - } - continuation.resume(returning: ()) - } - do { - try process.run() - } catch { - continuation.resume(throwing: error) - } - } - } } diff --git a/Sources/Config/ModelCatalogDownloads.swift b/Sources/Config/ModelCatalogDownloads.swift new file mode 100644 index 0000000..a547954 --- /dev/null +++ b/Sources/Config/ModelCatalogDownloads.swift @@ -0,0 +1,228 @@ +import Foundation +import MLXLLM +import MLXLMCommon +import WhisperKit + +@MainActor +extension ModelCatalog { + static var whisperDownloadBase: URL { ModelStorage.huggingFaceBase } + + func downloadWhisper(_ id: String) async { + await downloadTasks.run(key: ModelDownloadKey(kind: .whisper, modelID: id)) { [weak self] in + await self?.performWhisperDownload(id) + } + } + + private func performWhisperDownload(_ id: String) async { + guard let idx = whisperModels.firstIndex(where: { $0.id == id }), + !whisperModels[idx].status.isDownloading else { return } + + if isWhisperDownloaded(id) { + whisperModels[idx].status = .downloaded + whisperModels[idx].cacheSize = whisperVariantSize(id) + whisperModels[idx].downloadDetail = "" + return + } + + whisperModels[idx].status = .downloading + whisperModels[idx].downloadProgress = 0 + + do { + let modelDir = ModelStorage.whisperVariantDir(id) + let tracker = DownloadProgressTracker(initialBytes: ModelStorage.directorySize(at: modelDir)) + _ = try await WhisperKit.download( + variant: id, + downloadBase: Self.whisperDownloadBase, + progressCallback: { [weak self] progress in + Task { @MainActor in + guard let self, + let i = self.whisperModels.firstIndex(where: { $0.id == id }) else { return } + let downloadedBytes = ModelStorage.directorySize(at: modelDir) + let info = tracker.update( + completedBytes: downloadedBytes > 0 + ? downloadedBytes + : progress.completedUnitCount, + totalBytes: progress.totalUnitCount, + fraction: progress.fractionCompleted + ) + self.whisperModels[i].downloadProgress = info.fraction + self.whisperModels[i].downloadDetail = info.detailText + } + } + ) + try Task.checkCancellation() + if let i = whisperModels.firstIndex(where: { $0.id == id }) { + whisperModels[i].status = isWhisperDownloaded(id) + ? .downloaded + : .error(L("model.download_incomplete")) + whisperModels[i].cacheSize = whisperVariantSize(id) + whisperModels[i].downloadDetail = "" + } + } catch is CancellationError { + if let i = whisperModels.firstIndex(where: { $0.id == id }) { + whisperModels[i].status = .error(L("model.download_paused")) + whisperModels[i].cacheSize = whisperVariantSize(id) + whisperModels[i].downloadProgress = 0 + whisperModels[i].downloadDetail = "" + } + } catch { + Log.error("[ModelCatalog] Whisper download failed: \(error.localizedDescription)") + whisperModels[idx].status = .error(ModelDownloadFailureMessage.userFacing(error)) + whisperModels[idx].cacheSize = whisperVariantSize(id) + whisperModels[idx].downloadDetail = "" + } + } + + func deleteWhisper(_ id: String) { + guard let idx = whisperModels.firstIndex(where: { $0.id == id }) else { return } + if ModelStorage.localWhisperURL(id) != nil { + var paths = settings.localWhisperModelPaths + paths.removeValue(forKey: id) + settings.localWhisperModelPaths = paths + whisperModels.remove(at: idx) + } else { + try? FileManager.default.removeItem(at: whisperVariantDir(id)) + whisperModels[idx].status = .notDownloaded + whisperModels[idx].cacheSize = 0 + } + + if settings.whisperModel == id { + settings.whisperModel = + nextAvailableWhisper(excluding: id) ?? whisperModels.first?.id ?? "" + } + } + + func nextAvailableWhisper(excluding id: String) -> String? { + whisperModels.first { + $0.id != id && ($0.status == .downloaded || $0.status == .ready) + }?.id + } + + func downloadLLM(_ id: String) async { + await downloadTasks.run(key: ModelDownloadKey(kind: .llm, modelID: id)) { [weak self] in + await self?.performLLMDownload(id) + } + } + + private func performLLMDownload(_ id: String) async { + guard let idx = llmModels.firstIndex(where: { $0.id == id }), + !llmModels[idx].status.isDownloading else { return } + if ModelStorage.localLLMURL(id) != nil { + llmModels[idx].status = llmRepoIsComplete(id) + ? .downloaded + : .error(L("model.local_missing")) + llmModels[idx].cacheSize = llmRepoSize(id) + return + } + if llmRepoIsComplete(id) { + llmModels[idx].status = .downloaded + llmModels[idx].cacheSize = llmRepoSize(id) + llmModels[idx].downloadDetail = "" + return + } + + llmModels[idx].status = .downloading + llmModels[idx].downloadProgress = 0 + + do { + let estimatedTotalBytes = estimatedLLMDownloadBytes(id) ?? 0 + let repoDir = ModelStorage.hubModelRepoDir(id) + let tracker = DownloadProgressTracker( + initialBytes: ModelStorage.directorySize(at: repoDir) + ) + _ = try await LLMModelFactory.shared.loadContainer( + from: MLXModelLoading.downloader, + using: MLXModelLoading.tokenizerLoader, + configuration: ModelConfiguration(id: id) + ) { [weak self] progress in + Task { @MainActor in + guard let self, + let i = self.llmModels.firstIndex(where: { $0.id == id }) else { return } + let info = tracker.update( + completedBytes: ModelStorage.directorySize(at: repoDir), + totalBytes: estimatedTotalBytes, + fraction: progress.fractionCompleted + ) + self.llmModels[i].downloadProgress = info.fraction + self.llmModels[i].downloadDetail = info.detailText + } + } + try Task.checkCancellation() + if let i = llmModels.firstIndex(where: { $0.id == id }) { + llmModels[i].status = llmRepoIsComplete(id) + ? .downloaded + : .error(L("model.download_incomplete")) + llmModels[i].cacheSize = llmRepoSize(id) + llmModels[i].downloadDetail = "" + } + } catch is CancellationError { + if let i = llmModels.firstIndex(where: { $0.id == id }) { + llmModels[i].status = .error(L("model.download_paused")) + llmModels[i].cacheSize = llmRepoSize(id) + llmModels[i].downloadProgress = 0 + llmModels[i].downloadDetail = "" + } + } catch { + if let i = llmModels.firstIndex(where: { $0.id == id }) { + Log.error("[ModelCatalog] LLM download failed: \(error.localizedDescription)") + llmModels[i].status = .error(ModelDownloadFailureMessage.userFacing(error)) + llmModels[i].cacheSize = llmRepoSize(id) + llmModels[i].downloadDetail = "" + } + } + } + + func deleteLLM(_ id: String) { + guard let idx = llmModels.firstIndex(where: { $0.id == id }) else { return } + if ModelStorage.localLLMURL(id) != nil { + var paths = settings.localLLMModelPaths + paths.removeValue(forKey: id) + settings.localLLMModelPaths = paths + llmModels.remove(at: idx) + } else { + if let dir = llmRepoDir(id) { + try? FileManager.default.removeItem(at: dir) + } + llmModels[idx].status = .notDownloaded + llmModels[idx].cacheSize = 0 + } + + if settings.llmModel == id { + settings.llmModel = nextAvailableLLM(excluding: id) ?? llmModels.first?.id ?? "" + } + } + + func nextAvailableLLM(excluding id: String) -> String? { + llmModels.first { + $0.id != id && ($0.status == .downloaded || $0.status == .ready) + }?.id + } + + func whisperVariantDir(_ variant: String) -> URL { + ModelStorage.localWhisperURL(variant) ?? ModelStorage.whisperVariantDir(variant) + } + + func whisperVariantSize(_ variant: String) -> Int64 { + let dir = whisperVariantDir(variant) + guard FileManager.default.fileExists(atPath: dir.path) else { return 0 } + return ModelStorage.directorySize(at: dir) + } + + func isWhisperDownloaded(_ variant: String) -> Bool { + ModelStorage.whisperModelIsComplete(at: whisperVariantDir(variant)) + } + + func llmRepoDir(_ modelID: String) -> URL? { + ModelStorage.llmRepoDir(modelID) + } + + func llmRepoIsComplete(_ modelID: String) -> Bool { + guard let dir = llmRepoDir(modelID) else { return false } + return ModelStorage.llmRepoIsComplete(at: dir) + } + + func llmRepoSize(_ modelID: String) -> Int64 { + guard let dir = llmRepoDir(modelID) else { return 0 } + return ModelStorage.directorySize(at: dir) + } +} diff --git a/Sources/Config/ModelCatalogMiMoRuntime.swift b/Sources/Config/ModelCatalogMiMoRuntime.swift new file mode 100644 index 0000000..fdb3d53 --- /dev/null +++ b/Sources/Config/ModelCatalogMiMoRuntime.swift @@ -0,0 +1,10 @@ +import Foundation + +@MainActor +extension ModelCatalog { + static func mimoRepositoryIsReady(at dir: URL) -> Bool { + FileManager.default.fileExists( + atPath: dir.appendingPathComponent("src/mimo_audio/mimo_audio.py").path + ) + } +} diff --git a/Sources/Config/ModelDownloadFailureMessage.swift b/Sources/Config/ModelDownloadFailureMessage.swift new file mode 100644 index 0000000..589cb8d --- /dev/null +++ b/Sources/Config/ModelDownloadFailureMessage.swift @@ -0,0 +1,52 @@ +import Foundation + +enum ModelDownloadFailureMessage { + static func userFacing(_ error: Error) -> String { + let errors = errorChain(from: error) + + if errors.contains(where: { isTimeout($0) }) { + return L("model.download_failed_timeout") + } + if errors.contains(where: { isNetworkFailure($0) }) { + return L("model.download_failed_network") + } + if errors.contains(where: { $0.domain == NSCocoaErrorDomain && $0.code == NSFileWriteOutOfSpaceError }) { + return L("model.download_failed_disk_space") + } + if errors.contains(where: { isFilePermissionFailure($0) }) { + return L("model.download_failed_permission") + } + return L("model.download_failed_retry") + } + + private static func errorChain(from error: Error) -> [NSError] { + var result: [NSError] = [] + var current: NSError? = error as NSError + var seen = Set() + + while let item = current { + let identity = ObjectIdentifier(item) + guard seen.insert(identity).inserted else { break } + result.append(item) + current = item.userInfo[NSUnderlyingErrorKey] as? NSError + } + return result + } + + private static func isTimeout(_ error: NSError) -> Bool { + error.domain == NSURLErrorDomain && error.code == URLError.timedOut.rawValue + } + + private static func isNetworkFailure(_ error: NSError) -> Bool { + error.domain == NSURLErrorDomain || error.domain.hasPrefix("Network.") + } + + private static func isFilePermissionFailure(_ error: NSError) -> Bool { + guard error.domain == NSCocoaErrorDomain else { return false } + return [ + NSFileReadNoPermissionError, + NSFileWriteNoPermissionError, + NSFileWriteVolumeReadOnlyError, + ].contains(error.code) + } +} diff --git a/Sources/Config/ModelDownloadTasks.swift b/Sources/Config/ModelDownloadTasks.swift new file mode 100644 index 0000000..b4eb57f --- /dev/null +++ b/Sources/Config/ModelDownloadTasks.swift @@ -0,0 +1,54 @@ +import Foundation + +enum ModelDownloadKind: Hashable { + case whisper + case llm + case asr +} + +struct ModelDownloadKey: Hashable { + let kind: ModelDownloadKind + let modelID: String +} + +@MainActor +final class ModelDownloadTasks { + private struct Entry { + let token: UUID + let task: Task + } + + private var entries: [ModelDownloadKey: Entry] = [:] + + func run( + key: ModelDownloadKey, + operation: @escaping @MainActor () async -> Void + ) async { + if let existing = entries[key] { + await existing.task.value + return + } + + let token = UUID() + let task = Task { @MainActor in + await operation() + } + entries[key] = Entry(token: token, task: task) + await task.value + + if entries[key]?.token == token { + entries.removeValue(forKey: key) + } + } + + func cancel(_ key: ModelDownloadKey) { + entries[key]?.task.cancel() + } +} + +@MainActor +extension ModelCatalog { + func cancelDownload(_ id: String, kind: ModelDownloadKind) { + downloadTasks.cancel(ModelDownloadKey(kind: kind, modelID: id)) + } +} diff --git a/Sources/Config/ModelStorage.swift b/Sources/Config/ModelStorage.swift index 6e348ca..baf4f0f 100644 --- a/Sources/Config/ModelStorage.swift +++ b/Sources/Config/ModelStorage.swift @@ -34,8 +34,8 @@ enum ModelStorage { root.appendingPathComponent("runtimes") } - static func qwenASRRuntimeDir() -> URL { - asrRuntimeBase.appendingPathComponent("qwen3-asr") + static func localASRRuntimeDir(for provider: LocalASRConfiguration.Provider) -> URL { + asrRuntimeBase.appendingPathComponent("\(provider.rawValue)-asr") } static func whisperVariantDir(_ variant: String) -> URL { @@ -58,6 +58,11 @@ enum ModelStorage { return FileManager.default.fileExists(atPath: dir.path) ? dir : nil } + static func installedLLMURL(_ modelID: String) -> URL? { + guard let dir = llmRepoDir(modelID), llmRepoIsComplete(at: dir) else { return nil } + return dir + } + static func asrRepoDir(_ modelID: String) -> URL? { let dir = hubModelRepoDir(modelID) return FileManager.default.fileExists(atPath: dir.path) ? dir : nil @@ -73,6 +78,46 @@ enum ModelStorage { return URL(fileURLWithPath: NSString(string: path).expandingTildeInPath) } + static func whisperModelIsComplete(at dir: URL) -> Bool { + ["MelSpectrogram", "AudioEncoder", "TextDecoder"].allSatisfy { name in + resourceHasContent(dir.appendingPathComponent("\(name).mlmodelc")) || + resourceHasContent(dir.appendingPathComponent("\(name).mlpackage")) + } + } + + static func llmRepoIsComplete(at dir: URL) -> Bool { + guard fileExists(dir.appendingPathComponent("config.json")) else { return false } + let indexURL = dir.appendingPathComponent("model.safetensors.index.json") + if fileExists(indexURL), + let data = try? Data(contentsOf: indexURL), + let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let weightMap = object["weight_map"] as? [String: String] { + let shards = Set(weightMap.values) + return !shards.isEmpty && shards.allSatisfy { + fileExists(dir.appendingPathComponent($0)) + } + } + + if fileExists(dir.appendingPathComponent("model.safetensors")) || + fileExists(dir.appendingPathComponent("weights.safetensors")) { + return true + } + + guard let enumerator = FileManager.default.enumerator( + at: dir, + includingPropertiesForKeys: [.isRegularFileKey], + options: [.skipsHiddenFiles] + ) else { return false } + + for case let fileURL as URL in enumerator { + if fileURL.pathExtension.lowercased() == "npz", + fileExists(fileURL) { + return true + } + } + return false + } + static func directorySize(at url: URL) -> Int64 { guard let enumerator = FileManager.default.enumerator( at: url, includingPropertiesForKeys: [.fileSizeKey] @@ -86,6 +131,22 @@ enum ModelStorage { return total } + private static func fileExists(_ url: URL) -> Bool { + var isDirectory = ObjCBool(false) + guard FileManager.default.fileExists(atPath: url.path, isDirectory: &isDirectory), + !isDirectory.boolValue else { return false } + let attributes = try? FileManager.default.attributesOfItem(atPath: url.path) + return (attributes?[.size] as? NSNumber)?.int64Value ?? 0 > 0 + } + + private static func resourceHasContent(_ url: URL) -> Bool { + var isDirectory = ObjCBool(false) + guard FileManager.default.fileExists(atPath: url.path, isDirectory: &isDirectory) else { + return false + } + return isDirectory.boolValue ? directorySize(at: url) > 0 : fileExists(url) + } + static func makeLocalID(prefix: String, folderName: String, existing: Set) -> String { let cleanName = folderName.isEmpty ? "model" : folderName let base = "local/\(prefix)-\(cleanName)" diff --git a/Sources/LLM/LLMEngine.swift b/Sources/LLM/LLMEngine.swift index 7139e52..387dfec 100644 --- a/Sources/LLM/LLMEngine.swift +++ b/Sources/LLM/LLMEngine.swift @@ -7,46 +7,19 @@ actor LLMEngine { private var container: ModelContainer? private var currentModelID: String? - func loadModel( - id: String, - estimatedDownloadBytes: Int64? = nil, - progress: (@Sendable (DownloadProgressInfo) -> Void)? = nil - ) async throws { + func loadModel(id: String) async throws { if currentModelID == id, container != nil { return } Log.info("[LLMEngine] loading model: \(id)") let t0 = CFAbsoluteTimeGetCurrent() - if let localURL = ModelStorage.localLLMURL(id) { - container = try await LLMModelFactory.shared.loadContainer( - from: localURL, - using: MLXModelLoading.tokenizerLoader - ) - progress?(DownloadProgressInfo( - fraction: 1, - elapsedSeconds: CFAbsoluteTimeGetCurrent() - t0, - completedBytes: 0, - totalBytes: 0, - speedBytesPerSecond: 0 - )) - } else { - let estimatedTotalBytes = estimatedDownloadBytes ?? 0 - let repoDir = ModelStorage.hubModelRepoDir(id) - let tracker = DownloadProgressTracker(initialBytes: ModelStorage.directorySize(at: repoDir)) - let config = Self.modelConfiguration(for: id) - container = try await LLMModelFactory.shared.loadContainer( - from: MLXModelLoading.downloader, - using: MLXModelLoading.tokenizerLoader, - configuration: config - ) { p in - let completedBytes = ModelStorage.directorySize(at: repoDir) - progress?(tracker.update( - completedBytes: completedBytes, - totalBytes: estimatedTotalBytes, - fraction: p.fractionCompleted - )) - } + guard let localURL = ModelStorage.installedLLMURL(id) else { + throw LLMError.modelNotDownloaded } + container = try await LLMModelFactory.shared.loadContainer( + from: localURL, + using: MLXModelLoading.tokenizerLoader + ) currentModelID = id let elapsed = CFAbsoluteTimeGetCurrent() - t0 @@ -153,10 +126,12 @@ actor LLMEngine { enum LLMError: LocalizedError { case modelNotLoaded + case modelNotDownloaded var errorDescription: String? { switch self { - case .modelNotLoaded: return "LLM 模型未加载" + case .modelNotLoaded: return L("error.llm_not_loaded") + case .modelNotDownloaded: return L("error.llm_not_downloaded") } } } diff --git a/Sources/LLM/VLMEngine.swift b/Sources/LLM/VLMEngine.swift index 04f1054..66773f9 100644 --- a/Sources/LLM/VLMEngine.swift +++ b/Sources/LLM/VLMEngine.swift @@ -14,19 +14,13 @@ actor VLMEngine { Log.info("[VLMEngine] loading model: \(id)") let started = CFAbsoluteTimeGetCurrent() - if let localURL = ModelStorage.localLLMURL(id) { - container = try await VLMModelFactory.shared.loadContainer( - from: localURL, - using: MLXModelLoading.tokenizerLoader - ) - } else { - let config = LLMEngine.modelConfiguration(for: id) - container = try await VLMModelFactory.shared.loadContainer( - from: MLXModelLoading.downloader, - using: MLXModelLoading.tokenizerLoader, - configuration: config - ) + guard let localURL = ModelStorage.installedLLMURL(id) else { + throw LLMError.modelNotDownloaded } + container = try await VLMModelFactory.shared.loadContainer( + from: localURL, + using: MLXModelLoading.tokenizerLoader + ) currentModelID = id let elapsed = CFAbsoluteTimeGetCurrent() - started diff --git a/Sources/Processing/TextProcessor.swift b/Sources/Processing/TextProcessor.swift index d5e7ed6..9ae2ee7 100644 --- a/Sources/Processing/TextProcessor.swift +++ b/Sources/Processing/TextProcessor.swift @@ -21,18 +21,10 @@ final class TextProcessor { } @discardableResult - func warmUpLLM( - model: String, - estimatedDownloadBytes: Int64? = nil, - progress: (@Sendable (DownloadProgressInfo) -> Void)? = nil - ) async -> Bool { + func warmUpLLM(model: String) async -> Bool { if AppSettings.shared.useRemoteLLM { return true } do { - try await llm.loadModel( - id: model, - estimatedDownloadBytes: estimatedDownloadBytes, - progress: progress - ) + try await llm.loadModel(id: model) return true } catch { Log.error("[TextProcessor] LLM warmup failed: \(error.localizedDescription)") diff --git a/Sources/Resources/en.lproj/Localizable.strings b/Sources/Resources/en.lproj/Localizable.strings index 4a21317..4e222f2 100644 --- a/Sources/Resources/en.lproj/Localizable.strings +++ b/Sources/Resources/en.lproj/Localizable.strings @@ -36,8 +36,8 @@ "pipeline.busy" = "Previous task still processing…"; "pipeline.loading_llm" = "Loading LLM…"; "pipeline.model_not_ready" = "Speech model not ready — check your network"; -"pipeline.model_load_failed" = "Model load failed"; -"pipeline.model_load_failed_network" = "Model load failed — check network"; +"pipeline.model_load_failed" = "The model could not be loaded. Try again; if it still fails, redownload or switch models in Models."; +"pipeline.model_load_failed_network" = "The model could not be loaded. Check your network and try again."; "pipeline.recording" = "Recording…"; "pipeline.recording_translation" = "Recording for translation…"; "pipeline.transcribing" = "Transcribing…"; @@ -70,10 +70,10 @@ "pipeline.insert_failed_body" = "The text has been copied to your clipboard. Press ⌘V to paste it manually.\n\nReason: "; "pipeline.mic_failed_permissions" = "Microphone failed — check permissions"; "pipeline.mic_unavailable" = "Microphone unavailable"; -"pipeline.download_failed" = "Model download failed — check your network"; -"pipeline.compile_failed" = "Model compile failed — try again or switch model"; -"pipeline.load_failed" = "Model load failed — try again"; -"pipeline.model_load_failed_prefix" = "Model load failed: "; +"pipeline.download_failed" = "The model download did not finish. Existing files are preserved; check your network and resume."; +"pipeline.compile_failed" = "The model was downloaded but could not be compiled. Retry, then redownload or switch models if it continues."; +"pipeline.load_failed" = "The model was downloaded but could not be loaded. Retry, then redownload or switch models if it continues."; +"pipeline.model_load_failed_prefix" = "The model could not be loaded: "; /* ── Menu Bar ── */ "menubar.listening" = "Listening…"; @@ -211,11 +211,34 @@ "model.local" = "Local"; "model.custom_local" = "Custom / Local"; "model.legacy_group" = "Traditional Models"; -"model.local_missing" = "Local model folder missing"; +"model.local_missing" = "The local model folder cannot be found. Select it again or remove this reference."; "model.storage.title" = "Model Storage"; "model.storage.choose" = "Choose Folder…"; "model.storage.reveal" = "Show in Finder"; "model.storage.reset" = "Use Default"; +"model.storage.usage" = "%@ managed locally"; +"model.storage.change_confirm_title" = "Existing models will not be moved"; +"model.storage.change_confirm_message" = "%@ remains at %@. Switching to %@ can create duplicate downloads unless you move or reuse those files yourself."; +"model.storage.switch_anyway" = "Switch Anyway"; +"model.downloads_active" = "%d Active Download(s)"; +"model.download_confirm_title" = "Start model download?"; +"model.download_confirm_message" = "%@ needs about %@ of network data. Files will be saved to %@. This may incur data charges; the download starts only after you confirm and can be cancelled."; +"model.download_runtime_note" = "A provider-specific runtime may also be downloaded into OpenType's managed storage."; +"model.delete_confirm_title" = "Delete local model files?"; +"model.delete_confirm_message" = "Delete %@ and reclaim %@? The model can be downloaded again later."; +"model.remove_reference_confirm_title" = "Remove imported model?"; +"model.remove_reference_confirm_message" = "Remove %@ from OpenType? The original files will stay where they are."; +"model.download_incomplete" = "Only part of this model was downloaded. Select Resume to continue; existing files are preserved."; +"model.download_paused" = "Download paused. Select Resume to continue from the files already saved."; +"model.download_required" = "The model files have not been downloaded. Open Settings → Models and confirm the data usage first."; +"model.download_stalled_help" = "No progress for a while? Cancel and resume; files already downloaded will be preserved."; +"model.download_failed_timeout" = "The download connection timed out. Existing files are preserved; check your network and select Resume."; +"model.download_failed_network" = "The network connection was interrupted and the download paused. Restore the connection and select Resume."; +"model.download_failed_disk_space" = "The model storage disk is out of space. Free space or change the storage location, then resume."; +"model.download_failed_permission" = "OpenType cannot write to the model folder. Check its permissions or choose another storage location."; +"model.download_failed_retry" = "The download did not finish. Select Resume; if progress still does not change, cancel and try again."; +"model.resume" = "Resume"; +"model.mimo_macos_unavailable" = "Unavailable on macOS: Xiaomi's official runtime currently requires Linux, CUDA, and flash-attn"; "download.elapsed_format" = "Time %@"; "download.progress_format" = "Progress %@"; "download.remaining_format" = "Left %@"; @@ -370,9 +393,13 @@ "mimo_asr.config_hint" = "Choose MiMo-V2.5-ASR, then click Download to fetch the model, audio tokenizer, and runtime files."; "model.qwen3_asr_quality" = "Local ASR through MLX, ~4.1 GB"; "model.mimo_asr_quality" = "Local ASR model plus audio tokenizer ~36 GB"; -"model.asr_incomplete" = "Download incomplete"; +"model.asr_incomplete" = "Only part of the model or runtime was downloaded. Select Resume to finish."; "model.asr_preparing_runtime" = "Preparing runtime files"; "model.asr_installing_runtime" = "Installing local runtime"; -"model.asr_runtime_missing" = "Runtime not installed"; +"model.asr_runtime_missing" = "The model files are present, but the runtime is not installed. Select Resume to finish setup."; "model.asr_runtime_download_failed" = "Runtime files download failed"; "model.asr_runtime_incomplete" = "Runtime files are incomplete"; +"error.llm_not_loaded" = "The model is stored locally but is not loaded into memory. Run the action again and OpenType will retry."; +"error.llm_not_downloaded" = "The model files have not been downloaded. Open Settings → Models and confirm the data usage first."; +"onboarding.download_notice" = "No download starts automatically. Review the size, then confirm if you want this local model."; +"onboarding.download_size" = "Download (%@)"; diff --git a/Sources/Resources/zh-Hans.lproj/Localizable.strings b/Sources/Resources/zh-Hans.lproj/Localizable.strings index 0f677f0..c440d50 100644 --- a/Sources/Resources/zh-Hans.lproj/Localizable.strings +++ b/Sources/Resources/zh-Hans.lproj/Localizable.strings @@ -36,8 +36,8 @@ "pipeline.busy" = "上一条语音仍在处理中…"; "pipeline.loading_llm" = "正在加载 LLM…"; "pipeline.model_not_ready" = "语音模型未就绪,请检查网络后重试"; -"pipeline.model_load_failed" = "模型加载失败"; -"pipeline.model_load_failed_network" = "模型加载失败,请检查网络"; +"pipeline.model_load_failed" = "模型未能加载,请重试;若仍失败,请在模型页重新下载或切换模型"; +"pipeline.model_load_failed_network" = "模型未能加载,请检查网络后重试"; "pipeline.recording" = "录音中…"; "pipeline.recording_translation" = "正在录音并准备翻译…"; "pipeline.transcribing" = "识别中…"; @@ -70,10 +70,10 @@ "pipeline.insert_failed_body" = "文本已复制到剪贴板,请按 ⌘V 手动粘贴。\n\n原因:"; "pipeline.mic_failed_permissions" = "麦克风启动失败,请检查权限"; "pipeline.mic_unavailable" = "麦克风不可用"; -"pipeline.download_failed" = "模型下载失败,请检查网络连接"; -"pipeline.compile_failed" = "模型编译失败,请重试或更换模型"; -"pipeline.load_failed" = "模型加载失败,请重试"; -"pipeline.model_load_failed_prefix" = "模型加载失败: "; +"pipeline.download_failed" = "模型下载未完成。已下载内容会保留,请检查网络后继续下载"; +"pipeline.compile_failed" = "模型文件已下载,但编译失败。请重试;若仍失败,请重新下载或切换模型"; +"pipeline.load_failed" = "模型文件已下载,但加载失败。请重试;若仍失败,请重新下载或切换模型"; +"pipeline.model_load_failed_prefix" = "模型未能加载:"; /* ── Menu Bar ── */ "menubar.listening" = "正在聆听…"; @@ -211,11 +211,34 @@ "model.local" = "本地"; "model.custom_local" = "自定义 / 本地"; "model.legacy_group" = "传统模型"; -"model.local_missing" = "本地模型目录不存在"; +"model.local_missing" = "找不到本地模型目录。请重新选择模型目录,或移除该引用"; "model.storage.title" = "模型存储位置"; "model.storage.choose" = "选择文件夹…"; "model.storage.reveal" = "在 Finder 中显示"; "model.storage.reset" = "使用默认位置"; +"model.storage.usage" = "本地已管理 %@"; +"model.storage.change_confirm_title" = "现有模型不会自动迁移"; +"model.storage.change_confirm_message" = "%@ 模型仍会留在 %@。切换到 %@ 后,如果未自行迁移或复用这些文件,可能造成重复下载。"; +"model.storage.switch_anyway" = "仍然切换"; +"model.downloads_active" = "%d 个下载任务"; +"model.download_confirm_title" = "开始下载模型?"; +"model.download_confirm_message" = "%@ 预计需要约 %@ 网络流量,文件将保存到 %@。下载可能产生流量费用;只有确认后才会开始,并且可以取消。"; +"model.download_runtime_note" = "对应模型的专用运行环境也可能下载到 OpenType 的统一模型目录。"; +"model.delete_confirm_title" = "删除本地模型文件?"; +"model.delete_confirm_message" = "删除 %@ 并释放 %@ 空间?之后仍可重新下载。"; +"model.remove_reference_confirm_title" = "移除导入的模型?"; +"model.remove_reference_confirm_message" = "从 OpenType 中移除 %@?原始文件仍会保留在原位置。"; +"model.download_incomplete" = "模型只下载了一部分。点击“继续下载”即可接着下载,已有文件会保留"; +"model.download_paused" = "下载已暂停。点击“继续下载”可从现有文件接着下载"; +"model.download_required" = "模型文件尚未下载。请前往 设置 → 模型,确认流量后下载"; +"model.download_stalled_help" = "进度长时间不变?可先取消再继续下载,已下载文件会保留。"; +"model.download_failed_timeout" = "下载连接超时。已下载内容会保留,请检查网络后点击“继续下载”"; +"model.download_failed_network" = "网络连接中断,下载已暂停。请恢复网络后点击“继续下载”"; +"model.download_failed_disk_space" = "模型目录磁盘空间不足。请清理空间或更换存储位置后继续下载"; +"model.download_failed_permission" = "无法写入模型目录。请检查目录权限,或在模型页更换存储位置"; +"model.download_failed_retry" = "下载未完成。请点击“继续下载”;若进度仍不变化,可先取消再重试"; +"model.resume" = "继续"; +"model.mimo_macos_unavailable" = "macOS 暂不可用:小米官方运行时目前依赖 Linux、CUDA 和 flash-attn"; "download.elapsed_format" = "已用 %@"; "download.progress_format" = "进度 %@"; "download.remaining_format" = "剩余 %@"; @@ -370,9 +393,13 @@ "mimo_asr.config_hint" = "选择 MiMo-V2.5-ASR 后,请点击下载按钮获取模型、音频 tokenizer 和运行文件。"; "model.qwen3_asr_quality" = "本地 MLX 语音识别,约 4.1 GB"; "model.mimo_asr_quality" = "本地语音识别模型 + 音频 tokenizer,约 36 GB"; -"model.asr_incomplete" = "下载不完整"; +"model.asr_incomplete" = "模型或运行环境只下载了一部分。点击“继续下载”即可接着完成"; "model.asr_preparing_runtime" = "准备运行文件"; "model.asr_installing_runtime" = "安装本地运行环境"; -"model.asr_runtime_missing" = "运行环境未安装"; +"model.asr_runtime_missing" = "模型文件已下载,但运行环境尚未安装。请点击“继续下载”完成安装"; "model.asr_runtime_download_failed" = "运行文件下载失败"; "model.asr_runtime_incomplete" = "运行文件不完整"; +"error.llm_not_loaded" = "模型文件已在本地,但当前尚未加载到内存。请重新执行;OpenType 会再次尝试加载"; +"error.llm_not_downloaded" = "模型文件尚未下载。请前往 设置 → 模型,确认流量后下载"; +"onboarding.download_notice" = "这里不会自动下载。请先确认体积,再决定是否下载这个本地模型。"; +"onboarding.download_size" = "下载(%@)"; diff --git a/Sources/Speech/LocalASREngine.swift b/Sources/Speech/LocalASREngine.swift index bc14bb1..c344396 100644 --- a/Sources/Speech/LocalASREngine.swift +++ b/Sources/Speech/LocalASREngine.swift @@ -132,10 +132,7 @@ final class LocalASREngine: SpeechEngine, @unchecked Sendable { func prepare() async { guard configuration.hasRequiredFiles, let runnerURL = Self.runnerScriptURL(), - let pythonPath = try? await LocalASRRuntime.ensurePythonPath( - for: configuration.provider, - preferredPath: configuration.pythonPath - ) else { + let pythonPath = try? LocalASRRuntime.pythonPath(for: configuration.provider) else { return } await server.warmUp(runnerURL: runnerURL, pythonPath: pythonPath) @@ -143,10 +140,7 @@ final class LocalASREngine: SpeechEngine, @unchecked Sendable { func transcribe(audioURL: URL?, language: String?) async throws -> String { guard configuration.hasRequiredFiles else { throw LocalASRError.notConfigured } - let pythonPath = try await LocalASRRuntime.ensurePythonPath( - for: configuration.provider, - preferredPath: configuration.pythonPath - ) + let pythonPath = try LocalASRRuntime.pythonPath(for: configuration.provider) guard let audioURL else { throw LocalASRError.noAudioFile } guard let runnerURL = Self.runnerScriptURL() else { throw LocalASRError.runnerMissing } diff --git a/Sources/Speech/LocalASRRuntime.swift b/Sources/Speech/LocalASRRuntime.swift index 68d7c8f..2171f49 100644 --- a/Sources/Speech/LocalASRRuntime.swift +++ b/Sources/Speech/LocalASRRuntime.swift @@ -7,6 +7,17 @@ enum LocalASRRuntime { private static let markerName = ".opentype-runtime-ready" private static let nativeMarkerName = ".opentype-native-runtime-ready" + static func availability( + for provider: LocalASRConfiguration.Provider + ) -> LocalASRRuntimeAvailability { + switch provider { + case .qwen3: + return .supported + case .mimo: + return .unavailable(L("model.mimo_macos_unavailable")) + } + } + static func isReady(for provider: LocalASRConfiguration.Provider) -> Bool { switch provider { case .qwen3: @@ -15,10 +26,22 @@ enum LocalASRRuntime { qwenMarkerIsCurrent(at: qwenMarkerURL()) && qwenMarkerIsCurrent(at: qwenNativeMarkerURL()) case .mimo: - return LocalASRConfiguration.resolvePythonPath() != nil + return false } } + static func pythonPath(for provider: LocalASRConfiguration.Provider) throws -> String { + guard case .supported = availability(for: provider) else { + throw LocalASRRuntimeError.unsupported(L("model.mimo_macos_unavailable")) + } + guard isReady(for: provider) else { + throw LocalASRRuntimeError.notInstalled + } + return qwenPythonURL().path + } + + /// Installs a managed runtime. Call only from an explicit, user-confirmed + /// model download or repair action. static func ensurePythonPath( for provider: LocalASRConfiguration.Provider, preferredPath: String @@ -27,15 +50,12 @@ enum LocalASRRuntime { case .qwen3: return try await ensureQwenRuntime(preferredPath: preferredPath) case .mimo: - guard let python = LocalASRConfiguration.resolvePythonPath(preferredPath: preferredPath) else { - throw LocalASRRuntimeError.pythonMissing - } - return python + throw LocalASRRuntimeError.unsupported(L("model.mimo_macos_unavailable")) } } private static func ensureQwenRuntime(preferredPath: String) async throws -> String { - let runtimeDir = ModelStorage.qwenASRRuntimeDir() + let runtimeDir = ModelStorage.localASRRuntimeDir(for: .qwen3) let runtimePython = qwenPythonURL().path if isReady(for: .qwen3) { return runtimePython } @@ -103,15 +123,15 @@ enum LocalASRRuntime { } private static func qwenPythonURL() -> URL { - ModelStorage.qwenASRRuntimeDir().appendingPathComponent("bin/python") + ModelStorage.localASRRuntimeDir(for: .qwen3).appendingPathComponent("bin/python") } private static func qwenMarkerURL() -> URL { - ModelStorage.qwenASRRuntimeDir().appendingPathComponent(markerName) + ModelStorage.localASRRuntimeDir(for: .qwen3).appendingPathComponent(markerName) } private static func qwenNativeMarkerURL() -> URL { - ModelStorage.qwenASRRuntimeDir().appendingPathComponent(nativeMarkerName) + ModelStorage.localASRRuntimeDir(for: .qwen3).appendingPathComponent(nativeMarkerName) } private static func prepareNativeExtensions(in runtimeDir: URL) async throws { @@ -149,40 +169,109 @@ enum LocalASRRuntime { } private static func runProcess(executable: String, arguments: [String]) async throws { - try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in - let process = Process() - let stderr = Pipe() - process.executableURL = URL(fileURLWithPath: executable) - process.arguments = arguments - process.standardOutput = FileHandle(forWritingAtPath: "/dev/null") - process.standardError = stderr - process.terminationHandler = { process in - let data = stderr.fileHandleForReading.readDataToEndOfFile() - let message = String(data: data, encoding: .utf8)? - .trimmingCharacters(in: .whitespacesAndNewlines) - guard process.terminationStatus == 0 else { - continuation.resume(throwing: LocalASRRuntimeError.processFailed(message ?? "")) - return + let cancellation = ProcessCancellation() + try await withTaskCancellationHandler { + try Task.checkCancellation() + try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + let process = Process() + let stderr = Pipe() + process.executableURL = URL(fileURLWithPath: executable) + process.arguments = arguments + process.standardOutput = FileHandle(forWritingAtPath: "/dev/null") + process.standardError = stderr + process.terminationHandler = { process in + let data = stderr.fileHandleForReading.readDataToEndOfFile() + let message = String(data: data, encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines) + cancellation.clear(process) + if cancellation.isCancelled { + continuation.resume(throwing: CancellationError()) + } else if process.terminationStatus != 0 { + continuation.resume(throwing: LocalASRRuntimeError.processFailed(message ?? "")) + } else { + continuation.resume(returning: ()) + } + } + cancellation.register(process) + do { + try process.run() + cancellation.terminateIfCancelled() + } catch { + cancellation.clear(process) + continuation.resume(throwing: error) } - continuation.resume(returning: ()) - } - do { - try process.run() - } catch { - continuation.resume(throwing: error) } + } onCancel: { + cancellation.cancel() + } + } +} + +private final class ProcessCancellation: @unchecked Sendable { + private let lock = NSLock() + private var process: Process? + private var wasCancelled = false + + var isCancelled: Bool { + lock.lock() + defer { lock.unlock() } + return wasCancelled + } + + func register(_ process: Process) { + lock.lock() + self.process = process + lock.unlock() + } + + func clear(_ process: Process) { + lock.lock() + if self.process === process { + self.process = nil } + lock.unlock() } + + func cancel() { + lock.lock() + wasCancelled = true + let process = self.process + lock.unlock() + if process?.isRunning == true { + process?.terminate() + } + } + + func terminateIfCancelled() { + lock.lock() + let shouldTerminate = wasCancelled + let process = self.process + lock.unlock() + if shouldTerminate, process?.isRunning == true { + process?.terminate() + } + } +} + +enum LocalASRRuntimeAvailability: Equatable { + case supported + case unavailable(String) } enum LocalASRRuntimeError: LocalizedError { case pythonMissing + case notInstalled + case unsupported(String) case processFailed(String) var errorDescription: String? { switch self { case .pythonMissing: return L("error.local_asr_python_missing") + case .notInstalled: + return L("model.asr_runtime_missing") + case .unsupported(let message): + return message case .processFailed(let message): return message.isEmpty ? L("error.local_asr_runtime_failed") : message } diff --git a/Sources/Speech/SpeechEngineProvider.swift b/Sources/Speech/SpeechEngineProvider.swift index 13a0bc2..53eb96b 100644 --- a/Sources/Speech/SpeechEngineProvider.swift +++ b/Sources/Speech/SpeechEngineProvider.swift @@ -43,17 +43,6 @@ final class SpeechEngineProvider { resourceId: settings.volcResourceId ) case .qwen3: - if !LocalASRRuntime.isReady(for: .qwen3), - !ModelCatalog.shared.asrModelPath(for: settings.qwenASRModel).isEmpty { - do { - _ = try await LocalASRRuntime.ensurePythonPath( - for: .qwen3, - preferredPath: settings.localASRPythonPath - ) - } catch { - Log.error("[SpeechEngineProvider] Qwen ASR runtime migration failed: \(error.localizedDescription)") - } - } guard localASRIsAvailable(settings.qwenASRModel) else { qwenSpeechEngine = nil Log.info("[SpeechEngineProvider] Qwen ASR model requires manual download: \(settings.qwenASRModel)") diff --git a/Sources/Speech/WhisperEngine.swift b/Sources/Speech/WhisperEngine.swift index ae23fbe..83a97d3 100644 --- a/Sources/Speech/WhisperEngine.swift +++ b/Sources/Speech/WhisperEngine.swift @@ -80,46 +80,12 @@ final class WhisperEngine: SpeechEngine, @unchecked Sendable { } Log.info("[WhisperEngine] using model: \(selectedModel)") - progress(dp(0.02, stage: .downloading)) - - let modelDir = ModelStorage.whisperVariantDir(selectedModel) - let tracker = DownloadProgressTracker(initialBytes: ModelStorage.directorySize(at: modelDir)) - - let folder: URL - if let localFolder { - folder = localFolder - } else { - do { - folder = try await WhisperKit.download( - variant: selectedModel, - downloadBase: ModelCatalog.whisperDownloadBase, - progressCallback: { p in - let downloadedBytes = ModelStorage.directorySize(at: modelDir) - let completed = downloadedBytes > 0 ? downloadedBytes : p.completedUnitCount - let total = p.totalUnitCount - - let frac = 0.02 + p.fractionCompleted * 0.58 - let info = tracker.update( - completedBytes: completed, - totalBytes: total, - fraction: frac - ) - progress(DownloadProgress( - fraction: frac, completedBytes: completed, - totalBytes: total, - speedBytesPerSec: info.speedBytesPerSecond, - elapsedSeconds: info.elapsedSeconds, - downloadFraction: p.fractionCompleted, - stage: .downloading - )) - } - ) - } catch { - isLoading = false - throw WhisperError.downloadFailed(error.localizedDescription) - } + let folder = localFolder ?? ModelStorage.whisperVariantDir(selectedModel) + guard ModelStorage.whisperModelIsComplete(at: folder) else { + isLoading = false + throw WhisperError.modelNotLoaded(L("model.download_required")) } - Log.info("[WhisperEngine] download complete") + Log.info("[WhisperEngine] loading local model assets") progress(dp(0.62, stage: .compiling)) diff --git a/Sources/UI/MenuBarView.swift b/Sources/UI/MenuBarView.swift index 24b7db5..2291911 100644 --- a/Sources/UI/MenuBarView.swift +++ b/Sources/UI/MenuBarView.swift @@ -90,7 +90,7 @@ struct MenuBarView: View { private var showActiveStatus: Bool { switch appState.phase { - case .transcribing, .processing, .inserting, .error: return true + case .loadingModel, .transcribing, .processing, .inserting, .error: return true default: return false } } @@ -243,6 +243,7 @@ struct MenuBarView: View { private var statusColor: Color { switch appState.phase { + case .loadingModel: return .blue case .transcribing, .processing: return .orange case .inserting: return .yellow case .error: return .red diff --git a/Sources/UI/ModelManagementActions.swift b/Sources/UI/ModelManagementActions.swift new file mode 100644 index 0000000..742ae07 --- /dev/null +++ b/Sources/UI/ModelManagementActions.swift @@ -0,0 +1,213 @@ +import SwiftUI + +extension ModelManagementView { + struct PendingModelAction: Identifiable { + enum Kind { + case download + case delete + } + + let id = UUID() + let kind: Kind + let model: ModelCatalog.ModelEntry + let type: ModelType + let isActive: Bool + } + + struct ActiveModelDownload: Identifiable { + let model: ModelCatalog.ModelEntry + let type: ModelType + + var id: String { + "\(type.downloadKind)-\(model.id)" + } + } + + var activeDownloads: [ActiveModelDownload] { + catalog.whisperModels + .filter(\.status.isDownloading) + .map { ActiveModelDownload(model: $0, type: .whisper) } + + catalog.llmModels + .filter(\.status.isDownloading) + .map { ActiveModelDownload(model: $0, type: .llm) } + + catalog.asrModels + .filter(\.status.isDownloading) + .map { ActiveModelDownload(model: $0, type: .asr) } + } + + var hasActiveDownloads: Bool { + !activeDownloads.isEmpty + } + + var knownStorageBytes: Int64 { + catalog.whisperModels + .filter { ModelStorage.localWhisperURL($0.id) == nil } + .reduce(0) { $0 + $1.cacheSize } + + catalog.llmModels + .filter { ModelStorage.localLLMURL($0.id) == nil } + .reduce(0) { $0 + $1.cacheSize } + + catalog.asrModels.reduce(0) { $0 + $1.cacheSize } + } + + var activeDownloadsSection: some View { + VStack(alignment: .leading, spacing: 10) { + Label( + String(format: L("model.downloads_active"), activeDownloads.count), + systemImage: "arrow.down.circle.fill" + ) + .font(.headline) + + ForEach(activeDownloads) { download in + VStack(alignment: .leading, spacing: 6) { + HStack { + Text(download.model.displayName) + .font(.system(size: 12, weight: .semibold)) + Spacer() + Text("\(Int(download.model.downloadProgress * 100))%") + .font(.system(size: 11, design: .monospaced)) + .monospacedDigit() + Button(L("common.cancel")) { + catalog.cancelDownload( + download.model.id, + kind: download.type.downloadKind + ) + } + .controlSize(.mini) + } + ProgressView(value: download.model.downloadProgress) + .progressViewStyle(.linear) + if !download.model.downloadDetail.isEmpty { + Text(download.model.downloadDetail) + .font(.system(size: 10, design: .monospaced)) + .foregroundStyle(.secondary) + } + Text(L("model.download_stalled_help")) + .font(.system(size: 10)) + .foregroundStyle(.tertiary) + } + .padding(10) + .background(.regularMaterial, in: RoundedRectangle(cornerRadius: 8)) + } + } + } + + func requestDownload( + _ model: ModelCatalog.ModelEntry, + type: ModelType, + isActive: Bool + ) { + pendingModelAction = PendingModelAction( + kind: .download, + model: model, + type: type, + isActive: isActive + ) + } + + func requestDelete( + _ model: ModelCatalog.ModelEntry, + type: ModelType, + isActive: Bool + ) { + pendingModelAction = PendingModelAction( + kind: .delete, + model: model, + type: type, + isActive: isActive + ) + } + + func modelActionAlert(_ action: PendingModelAction) -> Alert { + switch action.kind { + case .download: + return Alert( + title: Text(L("model.download_confirm_title")), + message: Text(downloadConfirmationMessage(action)), + primaryButton: .default(Text(L("common.download"))) { + Task { await download(action.model, type: action.type) } + }, + secondaryButton: .cancel() + ) + case .delete: + if isImportedLocal(action.model, type: action.type) { + return Alert( + title: Text(L("model.remove_reference_confirm_title")), + message: Text(String( + format: L("model.remove_reference_confirm_message"), + action.model.displayName + )), + primaryButton: .destructive(Text(L("common.delete"))) { + delete(action.model, isActive: action.isActive, type: action.type) + }, + secondaryButton: .cancel() + ) + } + return Alert( + title: Text(L("model.delete_confirm_title")), + message: Text(String( + format: L("model.delete_confirm_message"), + action.model.displayName, + ModelCatalog.formatBytes(action.model.cacheSize) + )), + primaryButton: .destructive(Text(L("common.delete"))) { + delete(action.model, isActive: action.isActive, type: action.type) + }, + secondaryButton: .cancel() + ) + } + } + + private func downloadConfirmationMessage(_ action: PendingModelAction) -> String { + let estimate = estimatedDownloadBytes(for: action.model, type: action.type) + let remaining = estimate.map { max($0 - action.model.cacheSize, 0) } + let sizeText = remaining.map(ModelCatalog.formatBytes) ?? L("download.unknown") + var message = String( + format: L("model.download_confirm_message"), + action.model.displayName, + sizeText, + ModelStorage.root.path + ) + if action.type == .asr { + message += "\n\n" + L("model.download_runtime_note") + } + return message + } + + private func estimatedDownloadBytes( + for model: ModelCatalog.ModelEntry, + type: ModelType + ) -> Int64? { + switch type { + case .whisper: + return ModelCatalog.estimatedDownloadBytes(from: model.id) + case .llm: + return catalog.estimatedLLMDownloadBytes(model.id) + case .asr: + return catalog.estimatedASRDownloadBytes(model.id) + } + } + + private func isImportedLocal( + _ model: ModelCatalog.ModelEntry, + type: ModelType + ) -> Bool { + switch type { + case .whisper: + return ModelStorage.localWhisperURL(model.id) != nil + case .llm: + return ModelStorage.localLLMURL(model.id) != nil + case .asr: + return false + } + } +} + +extension ModelManagementView.ModelType { + var downloadKind: ModelDownloadKind { + switch self { + case .whisper: return .whisper + case .llm: return .llm + case .asr: return .asr + } + } +} diff --git a/Sources/UI/ModelManagementRows.swift b/Sources/UI/ModelManagementRows.swift index 6717451..0a3f739 100644 --- a/Sources/UI/ModelManagementRows.swift +++ b/Sources/UI/ModelManagementRows.swift @@ -1,7 +1,7 @@ import SwiftUI extension ModelManagementView { - enum ModelType { + enum ModelType: Equatable { case whisper case llm case asr @@ -27,60 +27,65 @@ extension ModelManagementView { func modelRow( _ model: ModelCatalog.ModelEntry, isActive: Bool, type: ModelType ) -> some View { - HStack(spacing: 10) { - statusDot(model.status) + VStack(alignment: .leading, spacing: 7) { + HStack(spacing: 10) { + statusDot(model.status) - VStack(alignment: .leading, spacing: 2) { - HStack(spacing: 6) { - Text(model.displayName) - .font(.system(size: 12, weight: isActive ? .semibold : .regular)) - if model.tier == .recommended { - Text(L("common.recommended")) - .font(.system(size: 9, weight: .medium)) - .padding(.horizontal, 5) - .padding(.vertical, 1) - .background(Color.green.opacity(0.15)) - .foregroundStyle(.green) - .clipShape(Capsule()) + VStack(alignment: .leading, spacing: 2) { + HStack(spacing: 6) { + Text(model.displayName) + .font(.system(size: 12, weight: isActive ? .semibold : .regular)) + if model.tier == .recommended { + Text(L("common.recommended")) + .font(.system(size: 9, weight: .medium)) + .padding(.horizontal, 5) + .padding(.vertical, 1) + .background(Color.green.opacity(0.15)) + .foregroundStyle(.green) + .clipShape(Capsule()) + } + if isActive { + Text(L("model.active")) + .font(.system(size: 9, weight: .medium)) + .padding(.horizontal, 5) + .padding(.vertical, 1) + .background(Color.accentColor.opacity(0.15)) + .foregroundStyle(Color.accentColor) + .clipShape(Capsule()) + } } - if isActive { - Text(L("model.active")) - .font(.system(size: 9, weight: .medium)) - .padding(.horizontal, 5) - .padding(.vertical, 1) - .background(Color.accentColor.opacity(0.15)) - .foregroundStyle(Color.accentColor) - .clipShape(Capsule()) + HStack(spacing: 6) { + Text(secondaryText(for: model)) + .font(.system(size: 10)) + .foregroundStyle(.secondary) + if let tps = model.benchmarkTPS { + Text(String(format: "%.1f tok/s", tps)) + .font(.system(size: 9, weight: .medium, design: .monospaced)) + .padding(.horizontal, 5) + .padding(.vertical, 1) + .background(Color.orange.opacity(0.12)) + .foregroundStyle(.orange) + .clipShape(Capsule()) + } } } - HStack(spacing: 6) { - Text(secondaryText(for: model)) + + Spacer() + + if !model.status.isBusy && model.cacheSize > 0 { + Text(ModelCatalog.formatBytes(model.cacheSize)) .font(.system(size: 10)) .foregroundStyle(.secondary) - if let tps = model.benchmarkTPS { - Text(String(format: "%.1f tok/s", tps)) - .font(.system(size: 9, weight: .medium, design: .monospaced)) - .padding(.horizontal, 5) - .padding(.vertical, 1) - .background(Color.orange.opacity(0.12)) - .foregroundStyle(.orange) - .clipShape(Capsule()) - } + .monospacedDigit() } - } - Spacer() + rowActions(model, isActive: isActive, type: type) + } if model.status.isBusy { busyModelStatus(model) - } else if model.cacheSize > 0 { - Text(ModelCatalog.formatBytes(model.cacheSize)) - .font(.system(size: 10)) - .foregroundStyle(.secondary) - .monospacedDigit() + .padding(.leading, 18) } - - rowActions(model, isActive: isActive, type: type) } .padding(.horizontal, 12) .padding(.vertical, 8) @@ -89,10 +94,10 @@ extension ModelManagementView { @ViewBuilder func busyModelStatus(_ model: ModelCatalog.ModelEntry) -> some View { - VStack(alignment: .trailing, spacing: 2) { + VStack(alignment: .leading, spacing: 4) { if model.status.isDownloading { ProgressView(value: model.downloadProgress) - .frame(width: 86) + .progressViewStyle(.linear) } else { ProgressView() .controlSize(.small) @@ -101,20 +106,25 @@ extension ModelManagementView { Text(model.downloadDetail) .font(.system(size: 9, design: .monospaced)) .foregroundStyle(.secondary) - .lineLimit(2) - .multilineTextAlignment(.trailing) - .frame(maxWidth: 220, alignment: .trailing) + .lineLimit(3) } } + .frame(maxWidth: .infinity, alignment: .leading) } @ViewBuilder func rowActions( _ model: ModelCatalog.ModelEntry, isActive: Bool, type: ModelType ) -> some View { - if model.status == .notDownloaded || model.status.isError { - Button(L("common.download")) { - Task { await download(model, type: type) } + if model.status.isDownloading { + Button(L("common.cancel")) { + catalog.cancelDownload(model.id, kind: type.downloadKind) + } + .controlSize(.mini) + } else if canDownload(model, type: type) && + (model.status == .notDownloaded || model.status.isError) { + Button(model.status.isError ? L("model.resume") : L("common.download")) { + requestDownload(model, type: type, isActive: isActive) } .controlSize(.mini) } @@ -130,9 +140,9 @@ extension ModelManagementView { benchmarkButton(model) } - if model.status.canDelete { - Button(role: .destructive) { - delete(model, isActive: isActive, type: type) + if model.status.canDelete && model.cacheSize > 0 { + Button { + requestDelete(model, type: type, isActive: isActive) } label: { Image(systemName: "trash") .font(.system(size: 10)) @@ -171,6 +181,20 @@ extension ModelManagementView { } } + func canDownload(_ model: ModelCatalog.ModelEntry, type: ModelType) -> Bool { + switch type { + case .whisper: + return ModelStorage.localWhisperURL(model.id) == nil + case .llm: + return ModelStorage.localLLMURL(model.id) == nil + case .asr: + if case .supported = catalog.asrRuntimeAvailability(for: model.id) { + return true + } + return false + } + } + func select(_ model: ModelCatalog.ModelEntry, type: ModelType) { switch type { case .whisper: @@ -239,7 +263,12 @@ extension ModelManagementView { } func secondaryText(for model: ModelCatalog.ModelEntry) -> String { - model.hint + switch model.status { + case .unavailable(let message), .error(let message): + return message + default: + return model.hint + } } func statusDot(_ status: ModelCatalog.ModelStatus) -> some View { @@ -251,6 +280,8 @@ extension ModelManagementView { ProgressView().controlSize(.mini) case .downloaded, .ready: Circle().fill(.green) + case .unavailable: + Circle().fill(.orange) case .error: Circle().fill(.red) } diff --git a/Sources/UI/ModelManagementSections.swift b/Sources/UI/ModelManagementSections.swift index 2282072..46b928e 100644 --- a/Sources/UI/ModelManagementSections.swift +++ b/Sources/UI/ModelManagementSections.swift @@ -13,6 +13,13 @@ extension ModelManagementView { .lineLimit(2) .textSelection(.enabled) + Text(String( + format: L("model.storage.usage"), + ModelCatalog.formatBytes(knownStorageBytes) + )) + .font(.system(size: 10)) + .foregroundStyle(.secondary) + HStack(spacing: 8) { Button(L("model.storage.choose")) { chooseModelStorageLocation() @@ -21,10 +28,11 @@ extension ModelManagementView { NSWorkspace.shared.activateFileViewerSelecting([ModelStorage.root]) } Button(L("model.storage.reset")) { - updateModelStoragePath(ModelStorage.defaultRoot.path) + requestModelStoragePath(ModelStorage.defaultRoot.path) } } .controlSize(.small) + .disabled(hasActiveDownloads) } } diff --git a/Sources/UI/ModelManagementView.swift b/Sources/UI/ModelManagementView.swift index f7f3510..c9c16a3 100644 --- a/Sources/UI/ModelManagementView.swift +++ b/Sources/UI/ModelManagementView.swift @@ -16,11 +16,16 @@ struct ModelManagementView: View { @State var importErrorMessage = "" @State var selectedModelFamily: ModelCatalog.ModelFamily? = .qwen @State var showLegacyModels = false + @State var pendingModelAction: PendingModelAction? let benchmarkEngine = LLMEngine() var body: some View { ScrollView { VStack(alignment: .leading, spacing: 20) { + if hasActiveDownloads { + activeDownloadsSection + Divider() + } storageSection Divider() preloadSection @@ -44,7 +49,7 @@ struct ModelManagementView: View { .padding(20) } .onAppear { - catalog.refreshStatus() + catalog.refreshStatus(recheckingErrors: true) syncSelectedFamilyFromActiveModel() } .onChange(of: settings.llmModel) { _, _ in syncSelectedFamilyFromActiveModel() } @@ -52,6 +57,7 @@ struct ModelManagementView: View { .onChange(of: settings.mimoASRRepoPath) { _, _ in onUnloadLocalASR?() } .onChange(of: settings.qwenASRModel) { _, _ in onUnloadLocalASR?() } .onChange(of: settings.mimoASRModel) { _, _ in onUnloadLocalASR?() } + .alert(item: $pendingModelAction, content: modelActionAlert) } } @@ -65,10 +71,34 @@ extension ModelManagementView { panel.directoryURL = ModelStorage.root panel.message = L("model.storage.choose") if panel.runModal() == .OK, let url = panel.url { - updateModelStoragePath(url.path) + requestModelStoragePath(url.path) } } + func requestModelStoragePath(_ path: String) { + let currentURL = ModelStorage.root.standardizedFileURL + let nextURL = URL(fileURLWithPath: path).standardizedFileURL + guard currentURL != nextURL else { return } + + let currentSize = ModelStorage.directorySize(at: currentURL) + if currentSize > 0 { + let alert = NSAlert() + alert.alertStyle = .warning + alert.messageText = L("model.storage.change_confirm_title") + alert.informativeText = String( + format: L("model.storage.change_confirm_message"), + ModelCatalog.formatBytes(currentSize), + currentURL.path, + nextURL.path + ) + alert.addButton(withTitle: L("common.cancel")) + alert.addButton(withTitle: L("model.storage.switch_anyway")) + guard alert.runModal() == .alertSecondButtonReturn else { return } + } + + updateModelStoragePath(nextURL.path) + } + func updateModelStoragePath(_ path: String) { onUnloadWhisper?() onUnloadLLM?() @@ -103,7 +133,7 @@ extension ModelManagementView { panel.canCreateDirectories = false panel.message = L("model.import_local") if panel.runModal() == .OK, let url = panel.url { - guard FileManager.default.fileExists(atPath: url.appendingPathComponent("config.json").path) else { + guard ModelStorage.llmRepoIsComplete(at: url) else { importErrorMessage = "" showImportError = true return @@ -114,9 +144,6 @@ extension ModelManagementView { } private func isValidWhisperFolder(_ url: URL) -> Bool { - ["MelSpectrogram", "AudioEncoder", "TextDecoder"].allSatisfy { name in - FileManager.default.fileExists(atPath: url.appendingPathComponent("\(name).mlmodelc").path) - || FileManager.default.fileExists(atPath: url.appendingPathComponent("\(name).mlpackage").path) - } + ModelStorage.whisperModelIsComplete(at: url) } } diff --git a/Sources/UI/OnboardingView.swift b/Sources/UI/OnboardingView.swift index c4eb302..1c2b236 100644 --- a/Sources/UI/OnboardingView.swift +++ b/Sources/UI/OnboardingView.swift @@ -8,6 +8,7 @@ struct OnboardingView: View { @ObservedObject private var settings = AppSettings.shared @ObservedObject private var catalog = ModelCatalog.shared @State private var skippedModelDownload = false + @State private var showModelDownloadConfirmation = false var body: some View { VStack(spacing: 0) { @@ -26,6 +27,14 @@ struct OnboardingView: View { navigationBar } .frame(width: 480, height: 420) + .alert(L("model.download_confirm_title"), isPresented: $showModelDownloadConfirmation) { + Button(L("common.cancel"), role: .cancel) { } + Button(L("common.download")) { + Task { await catalog.downloadLLM(settings.llmModel) } + } + } message: { + Text(onboardingDownloadConfirmationMessage) + } } // MARK: - Welcome @@ -248,6 +257,17 @@ struct OnboardingView: View { : L("onboarding.downloading_model")) .font(.system(size: 11)) .foregroundStyle(.secondary) + if model.status.isDownloading { + Button(L("common.cancel")) { + catalog.cancelDownload(model.id, kind: .llm) + } + .controlSize(.small) + } + } + if model.status.isDownloading { + Text(L("model.download_stalled_help")) + .font(.system(size: 10)) + .foregroundStyle(.tertiary) } } else if model.status == .downloaded || model.status == .ready { HStack(spacing: 6) { @@ -272,7 +292,16 @@ struct OnboardingView: View { } } Button(L("common.retry")) { - Task { await catalog.downloadLLM(settings.llmModel) } + showModelDownloadConfirmation = true + } + .controlSize(.small) + } else { + Text(L("onboarding.download_notice")) + .font(.system(size: 11)) + .foregroundStyle(.secondary) + + Button(onboardingDownloadButtonTitle) { + showModelDownloadConfirmation = true } .controlSize(.small) } @@ -291,16 +320,32 @@ struct OnboardingView: View { } .padding(.horizontal, 32) .onAppear { - if !skippedModelDownload, - let model = catalog.llmModels.first(where: { $0.id == settings.llmModel }), - model.status != .downloaded && model.status != .ready && !model.status.isBusy { - Task { - await catalog.downloadLLM(settings.llmModel) - } - } + catalog.refreshStatus(recheckingErrors: true) } } + private var onboardingDownloadButtonTitle: String { + guard let bytes = catalog.estimatedLLMDownloadBytes(settings.llmModel) else { + return L("common.download") + } + return String( + format: L("onboarding.download_size"), + ModelCatalog.formatBytes(bytes) + ) + } + + private var onboardingDownloadConfirmationMessage: String { + let model = catalog.llmModels.first(where: { $0.id == settings.llmModel }) + let estimate = catalog.estimatedLLMDownloadBytes(settings.llmModel) + let remaining = estimate.map { max($0 - (model?.cacheSize ?? 0), 0) } + return String( + format: L("model.download_confirm_message"), + model?.displayName ?? settings.llmModel, + remaining.map(ModelCatalog.formatBytes) ?? L("download.unknown"), + ModelStorage.root.path + ) + } + private var canContinueFromModelPrep: Bool { skippedModelDownload || catalog.llmModels.first(where: { $0.id == settings.llmModel })?.status == .downloaded || diff --git a/Sources/UI/OverlayPanelContent.swift b/Sources/UI/OverlayPanelContent.swift index ffc1105..28303ca 100644 --- a/Sources/UI/OverlayPanelContent.swift +++ b/Sources/UI/OverlayPanelContent.swift @@ -268,6 +268,9 @@ struct OverlayContentView: View { case .processing, .inserting: Image(systemName: "brain") .foregroundStyle(Color(red: 1.0, green: 0.6, blue: 0.25)) + case .loadingModel: + Image(systemName: "shippingbox.fill") + .foregroundStyle(.blue) case .downloading: Image(systemName: "arrow.down.circle.fill") .foregroundStyle(.blue) diff --git a/Tests/OpenTypeTests/ConfigurationTests.swift b/Tests/OpenTypeTests/ConfigurationTests.swift index 0eddd21..d18e391 100644 --- a/Tests/OpenTypeTests/ConfigurationTests.swift +++ b/Tests/OpenTypeTests/ConfigurationTests.swift @@ -270,36 +270,62 @@ final class ConfigurationTests: XCTestCase { } func testStartupPreloadPolicyLoadsOnlyWhisperSpeechModel() { - XCTAssertTrue(StartupModelPreloadPolicy.shouldPreloadSpeechModel(enabled: true, speechEngine: .whisper)) - XCTAssertFalse(StartupModelPreloadPolicy.shouldPreloadSpeechModel(enabled: true, speechEngine: .apple)) - XCTAssertFalse(StartupModelPreloadPolicy.shouldPreloadSpeechModel(enabled: true, speechEngine: .volc)) - XCTAssertFalse(StartupModelPreloadPolicy.shouldPreloadSpeechModel(enabled: true, speechEngine: .qwen3)) - XCTAssertFalse(StartupModelPreloadPolicy.shouldPreloadSpeechModel(enabled: true, speechEngine: .mimo)) - XCTAssertFalse(StartupModelPreloadPolicy.shouldPreloadSpeechModel(enabled: false, speechEngine: .whisper)) + XCTAssertTrue(StartupModelPreloadPolicy.shouldPreloadSpeechModel( + enabled: true, speechEngine: .whisper, modelDownloaded: true + )) + XCTAssertFalse(StartupModelPreloadPolicy.shouldPreloadSpeechModel( + enabled: true, speechEngine: .whisper, modelDownloaded: false + )) + for engine in [SpeechEngineType.apple, .volc, .qwen3, .mimo] { + XCTAssertFalse(StartupModelPreloadPolicy.shouldPreloadSpeechModel( + enabled: true, speechEngine: engine, modelDownloaded: true + )) + } + XCTAssertFalse(StartupModelPreloadPolicy.shouldPreloadSpeechModel( + enabled: false, speechEngine: .whisper, modelDownloaded: true + )) } - func testStartupPreloadPolicyLoadsOnlyLocalFormattingModelWithID() { + func testStartupPreloadPolicyLoadsOnlyDownloadedLocalFormattingModelWithID() { XCTAssertTrue(StartupModelPreloadPolicy.shouldPreloadFormattingModel( enabled: true, useRemoteLLM: false, - modelID: "mlx-community/Qwen3.5-2B-4bit" + modelID: "mlx-community/Qwen3.5-2B-4bit", + modelDownloaded: true + )) + XCTAssertFalse(StartupModelPreloadPolicy.shouldPreloadFormattingModel( + enabled: true, + useRemoteLLM: false, + modelID: "mlx-community/Qwen3.5-2B-4bit", + modelDownloaded: false )) XCTAssertFalse(StartupModelPreloadPolicy.shouldPreloadFormattingModel( enabled: true, useRemoteLLM: true, - modelID: "gpt-4.1-mini" + modelID: "gpt-4.1-mini", + modelDownloaded: true )) XCTAssertFalse(StartupModelPreloadPolicy.shouldPreloadFormattingModel( enabled: true, useRemoteLLM: false, - modelID: " " + modelID: " ", + modelDownloaded: true )) XCTAssertFalse(StartupModelPreloadPolicy.shouldPreloadFormattingModel( enabled: false, useRemoteLLM: false, - modelID: "mlx-community/Qwen3.5-2B-4bit" + modelID: "mlx-community/Qwen3.5-2B-4bit", + modelDownloaded: true )) } + + func testLocalASRRuntimeCapabilityDoesNotTreatSystemPythonAsMiMoRuntime() { + XCTAssertEqual( + LocalASRRuntime.availability(for: .mimo), + .unavailable(L("model.mimo_macos_unavailable")) + ) + XCTAssertFalse(LocalASRRuntime.isReady(for: .mimo)) + } } private func writeTestFiles(_ paths: [String], under dir: URL) throws { diff --git a/Tests/OpenTypeTests/ModelDownloadFailureMessageTests.swift b/Tests/OpenTypeTests/ModelDownloadFailureMessageTests.swift new file mode 100644 index 0000000..fea9a4f --- /dev/null +++ b/Tests/OpenTypeTests/ModelDownloadFailureMessageTests.swift @@ -0,0 +1,31 @@ +import XCTest +@testable import OpenType + +final class ModelDownloadFailureMessageTests: XCTestCase { + func testTimeoutMessageExplainsResumeBehavior() { + XCTAssertEqual( + ModelDownloadFailureMessage.userFacing(URLError(.timedOut)), + L("model.download_failed_timeout") + ) + } + + func testDiskSpaceMessageProvidesStorageAction() { + let error = NSError(domain: NSCocoaErrorDomain, code: NSFileWriteOutOfSpaceError) + XCTAssertEqual( + ModelDownloadFailureMessage.userFacing(error), + L("model.download_failed_disk_space") + ) + } + + func testWrappedNetworkErrorIsRecognized() { + let error = NSError( + domain: "Hub.Download", + code: 1, + userInfo: [NSUnderlyingErrorKey: URLError(.networkConnectionLost)] + ) + XCTAssertEqual( + ModelDownloadFailureMessage.userFacing(error), + L("model.download_failed_network") + ) + } +} diff --git a/Tests/OpenTypeTests/ModelDownloadTasksTests.swift b/Tests/OpenTypeTests/ModelDownloadTasksTests.swift new file mode 100644 index 0000000..c9fcfbc --- /dev/null +++ b/Tests/OpenTypeTests/ModelDownloadTasksTests.swift @@ -0,0 +1,51 @@ +import XCTest +@testable import OpenType + +@MainActor +final class ModelDownloadTasksTests: XCTestCase { + func testDuplicateRequestsShareOneDownloadTask() async { + let downloads = ModelDownloadTasks() + let key = ModelDownloadKey(kind: .llm, modelID: "test/model") + var starts = 0 + + let first = Task { @MainActor in + await downloads.run(key: key) { + starts += 1 + try? await Task.sleep(nanoseconds: 50_000_000) + } + } + await Task.yield() + let second = Task { @MainActor in + await downloads.run(key: key) { + starts += 1 + } + } + + await first.value + await second.value + XCTAssertEqual(starts, 1) + } + + func testCancelPropagatesToActiveDownloadTask() async { + let downloads = ModelDownloadTasks() + let key = ModelDownloadKey(kind: .asr, modelID: "test/model") + var observedCancellation = false + + let task = Task { @MainActor in + await downloads.run(key: key) { + do { + try await Task.sleep(nanoseconds: 5_000_000_000) + } catch is CancellationError { + observedCancellation = true + } catch { + XCTFail("Unexpected cancellation error: \(error)") + } + } + } + await Task.yield() + downloads.cancel(key) + await task.value + + XCTAssertTrue(observedCancellation) + } +} diff --git a/Tests/OpenTypeTests/UtilityTests.swift b/Tests/OpenTypeTests/UtilityTests.swift index f143fc6..ef12721 100644 --- a/Tests/OpenTypeTests/UtilityTests.swift +++ b/Tests/OpenTypeTests/UtilityTests.swift @@ -44,6 +44,63 @@ final class UtilityTests: XCTestCase { XCTAssertTrue(suffix.hasSuffix("/models/XiaomiMiMo/MiMo-V2.5-ASR")) } + func testModelStorageRequiresWeightsBeforeLLMIsComplete() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("OpenTypeTests-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: root, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: root) } + + try Data("{}".utf8).write(to: root.appendingPathComponent("config.json")) + XCTAssertFalse(ModelStorage.llmRepoIsComplete(at: root)) + + try Data("weights".utf8).write(to: root.appendingPathComponent("model.safetensors")) + XCTAssertTrue(ModelStorage.llmRepoIsComplete(at: root)) + } + + func testModelStorageRequiresEveryIndexedLLMShard() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("OpenTypeTests-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: root, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: root) } + + try Data("{}".utf8).write(to: root.appendingPathComponent("config.json")) + let index = """ + {"weight_map":{"first":"model-00001-of-00002.safetensors","second":"model-00002-of-00002.safetensors"}} + """ + try Data(index.utf8).write(to: root.appendingPathComponent("model.safetensors.index.json")) + try Data("one".utf8).write(to: root.appendingPathComponent("model-00001-of-00002.safetensors")) + XCTAssertFalse(ModelStorage.llmRepoIsComplete(at: root)) + + try Data("two".utf8).write(to: root.appendingPathComponent("model-00002-of-00002.safetensors")) + XCTAssertTrue(ModelStorage.llmRepoIsComplete(at: root)) + } + + func testModelStorageRequiresAllWhisperComponents() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("OpenTypeTests-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: root, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: root) } + + try FileManager.default.createDirectory( + at: root.appendingPathComponent("MelSpectrogram.mlmodelc"), + withIntermediateDirectories: true + ) + try Data("model".utf8).write( + to: root.appendingPathComponent("MelSpectrogram.mlmodelc/model.bin") + ) + XCTAssertFalse(ModelStorage.whisperModelIsComplete(at: root)) + + for name in ["AudioEncoder", "TextDecoder"] { + let component = root.appendingPathComponent("\(name).mlmodelc") + try FileManager.default.createDirectory( + at: component, + withIntermediateDirectories: true + ) + try Data("model".utf8).write(to: component.appendingPathComponent("model.bin")) + } + XCTAssertTrue(ModelStorage.whisperModelIsComplete(at: root)) + } + @MainActor func testDownloadEstimateParsesModelHints() { XCTAssertEqual(