From b6c23c5c1925859f6ac55562d1cd7d8d38fe3e20 Mon Sep 17 00:00:00 2001 From: idevlab Date: Sun, 21 Jun 2026 17:26:49 +0800 Subject: [PATCH] Fix download remaining and speed metrics --- Sources/App/VoicePipeline+Models.swift | 6 +++- Sources/Config/DownloadProgressInfo.swift | 4 +-- Sources/Config/ModelCatalog.swift | 9 +++++- Sources/Config/ModelCatalogASR.swift | 8 ++++- .../ModelCatalogDownloadEstimates.swift | 32 +++++++++++++++++++ Sources/LLM/LLMEngine.swift | 15 +++++++-- Sources/Processing/TextProcessor.swift | 7 +++- Tests/OpenTypeTests/UtilityTests.swift | 26 +++++++++++++++ 8 files changed, 99 insertions(+), 8 deletions(-) create mode 100644 Sources/Config/ModelCatalogDownloadEstimates.swift diff --git a/Sources/App/VoicePipeline+Models.swift b/Sources/App/VoicePipeline+Models.swift index f35acdf0..755ad06c 100644 --- a/Sources/App/VoicePipeline+Models.swift +++ b/Sources/App/VoicePipeline+Models.swift @@ -57,7 +57,11 @@ extension VoicePipeline { catalog.updateLLMStatus(model, status: .loading, detail: L("model.loading")) } - let loaded = await textProcessor.warmUpLLM(model: model) { [weak self] info in + 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 } diff --git a/Sources/Config/DownloadProgressInfo.swift b/Sources/Config/DownloadProgressInfo.swift index e28e64f3..8a80f634 100644 --- a/Sources/Config/DownloadProgressInfo.swift +++ b/Sources/Config/DownloadProgressInfo.swift @@ -27,8 +27,8 @@ struct DownloadProgressInfo: Equatable, Sendable { } var speedText: String { - guard speedBytesPerSecond > 0 else { return L("download.unknown") } - return "\(Self.formatBytes(Int64(speedBytesPerSecond)))/s" + guard speedBytesPerSecond >= 1 else { return L("download.unknown") } + return "\(Self.formatBytes(Int64(speedBytesPerSecond.rounded())))/s" } var detailText: String { diff --git a/Sources/Config/ModelCatalog.swift b/Sources/Config/ModelCatalog.swift index b367df62..b866591c 100644 --- a/Sources/Config/ModelCatalog.swift +++ b/Sources/Config/ModelCatalog.swift @@ -280,6 +280,8 @@ final class ModelCatalog: ObservableObject { do { let tracker = DownloadProgressTracker() + let estimatedTotalBytes = estimatedLLMDownloadBytes(id) ?? 0 + let repoDir = ModelStorage.hubModelRepoDir(id) let config = ModelConfiguration(id: id) _ = try await LLMModelFactory.shared.loadContainer( from: MLXModelLoading.downloader, @@ -288,7 +290,12 @@ final class ModelCatalog: ObservableObject { ) { [weak self] p in Task { @MainActor in guard let self, let i = self.llmModels.firstIndex(where: { $0.id == id }) else { return } - let info = tracker.update(progress: p) + 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 } diff --git a/Sources/Config/ModelCatalogASR.swift b/Sources/Config/ModelCatalogASR.swift index 3f823363..48fdf0e1 100644 --- a/Sources/Config/ModelCatalogASR.swift +++ b/Sources/Config/ModelCatalogASR.swift @@ -76,6 +76,7 @@ extension ModelCatalog { let repos = asrRequiredRepoIDs(for: id) 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() @@ -87,7 +88,12 @@ extension ModelCatalog { 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 info = tracker.update(progress: progress, fraction: fraction) + let completedBytes = self.asrRepoSize(id) + let info = tracker.update( + completedBytes: completedBytes, + totalBytes: estimatedTotalBytes, + fraction: fraction + ) self.asrModels[i].downloadProgress = info.fraction self.asrModels[i].downloadDetail = info.detailText onProgress?(info) diff --git a/Sources/Config/ModelCatalogDownloadEstimates.swift b/Sources/Config/ModelCatalogDownloadEstimates.swift new file mode 100644 index 00000000..0007bd57 --- /dev/null +++ b/Sources/Config/ModelCatalogDownloadEstimates.swift @@ -0,0 +1,32 @@ +import Foundation + +@MainActor +extension ModelCatalog { + func estimatedLLMDownloadBytes(_ id: String) -> Int64? { + guard let model = llmModels.first(where: { $0.id == id }) else { return nil } + return Self.estimatedDownloadBytes(from: model.hint) + } + + func estimatedASRDownloadBytes(_ id: String) -> Int64? { + guard let model = asrModels.first(where: { $0.id == id }) else { return nil } + return Self.estimatedDownloadBytes(from: model.hint) + } + + static func estimatedDownloadBytes(from text: String) -> Int64? { + let pattern = #"([0-9]+(?:\.[0-9]+)?)\s*(GB|MB)"# + guard let regex = try? NSRegularExpression(pattern: pattern, options: [.caseInsensitive]) else { + return nil + } + + let range = NSRange(text.startIndex.. Void)? = nil) async throws { + func loadModel( + id: String, + estimatedDownloadBytes: Int64? = nil, + progress: (@Sendable (DownloadProgressInfo) -> Void)? = nil + ) async throws { if currentModelID == id, container != nil { return } Log.info("[LLMEngine] loading model: \(id)") @@ -27,13 +31,20 @@ actor LLMEngine { )) } else { let tracker = DownloadProgressTracker() + let estimatedTotalBytes = estimatedDownloadBytes ?? 0 + let repoDir = ModelStorage.hubModelRepoDir(id) let config = Self.modelConfiguration(for: id) container = try await LLMModelFactory.shared.loadContainer( from: MLXModelLoading.downloader, using: MLXModelLoading.tokenizerLoader, configuration: config ) { p in - progress?(tracker.update(progress: p)) + let completedBytes = ModelStorage.directorySize(at: repoDir) + progress?(tracker.update( + completedBytes: completedBytes, + totalBytes: estimatedTotalBytes, + fraction: p.fractionCompleted + )) } } diff --git a/Sources/Processing/TextProcessor.swift b/Sources/Processing/TextProcessor.swift index 4ad47daf..c58c3182 100644 --- a/Sources/Processing/TextProcessor.swift +++ b/Sources/Processing/TextProcessor.swift @@ -23,11 +23,16 @@ final class TextProcessor { @discardableResult func warmUpLLM( model: String, + estimatedDownloadBytes: Int64? = nil, progress: (@Sendable (DownloadProgressInfo) -> Void)? = nil ) async -> Bool { if AppSettings.shared.useRemoteLLM { return true } do { - try await llm.loadModel(id: model, progress: progress) + try await llm.loadModel( + id: model, + estimatedDownloadBytes: estimatedDownloadBytes, + progress: progress + ) return true } catch { Log.error("[TextProcessor] LLM warmup failed: \(error.localizedDescription)") diff --git a/Tests/OpenTypeTests/UtilityTests.swift b/Tests/OpenTypeTests/UtilityTests.swift index ec81c13b..06ca94da 100644 --- a/Tests/OpenTypeTests/UtilityTests.swift +++ b/Tests/OpenTypeTests/UtilityTests.swift @@ -44,6 +44,32 @@ final class UtilityTests: XCTestCase { XCTAssertTrue(suffix.hasSuffix("/models/XiaomiMiMo/MiMo-V2.5-ASR")) } + @MainActor + func testDownloadEstimateParsesModelHints() { + XCTAssertEqual( + ModelCatalog.estimatedDownloadBytes(from: "整理质量最佳 5-bit ~5.5 GB"), + 5_500_000_000 + ) + XCTAssertEqual( + ModelCatalog.estimatedDownloadBytes(from: "Qwen3.5 极速 ~620 MB"), + 620_000_000 + ) + XCTAssertNil(ModelCatalog.estimatedDownloadBytes(from: "本地语音识别模型 + audio tokenizer")) + } + + func testDownloadSpeedHidesSubByteNoise() { + let info = DownloadProgressInfo( + fraction: 0.81, + elapsedSeconds: 1701, + completedBytes: 4_500_000_000, + totalBytes: 5_500_000_000, + speedBytesPerSecond: 0.4 + ) + + XCTAssertEqual(info.remainingText, "1.0 GB") + XCTAssertEqual(info.speedText, L("download.unknown")) + } + func testGzipRoundTripForTextAndBinaryData() throws { let text = Data("OpenType voice input. 你好，世界。".utf8) let compressedText = try XCTUnwrap(Gzip.compress(text))