74 lines
2.9 KiB
Swift
74 lines
2.9 KiB
Swift
import Foundation
|
|
|
|
/// A speech model the user can pick in Settings.
|
|
public struct ModelOption: Identifiable, Equatable, Sendable {
|
|
public let id: String // WhisperKit variant name
|
|
public let title: String
|
|
public let detail: String
|
|
public let approxSizeMB: Int
|
|
|
|
public init(id: String, title: String, detail: String, approxSizeMB: Int) {
|
|
self.id = id
|
|
self.title = title
|
|
self.detail = detail
|
|
self.approxSizeMB = approxSizeMB
|
|
}
|
|
}
|
|
|
|
public enum ModelCatalog {
|
|
public static let defaultVariant = "large-v3-v20240930_turbo"
|
|
|
|
/// Sizes measured from the Hugging Face repository `argmaxinc/whisperkit-coreml`.
|
|
public static let all: [ModelOption] = [
|
|
ModelOption(id: defaultVariant, title: "Whisper Large v3 Turbo", detail: "Aanbevolen · nauwkeurigst", approxSizeMB: 1638),
|
|
ModelOption(id: "large-v3-v20240930_turbo_632MB", title: "Whisper Large v3 Turbo (compact)", detail: "Kleiner, iets minder nauwkeurig", approxSizeMB: 646),
|
|
ModelOption(id: "small", title: "Whisper Small", detail: "Snel en licht, minder nauwkeurig", approxSizeMB: 486),
|
|
]
|
|
}
|
|
|
|
/// Where models live on disk: `<base>/models/argmaxinc/whisperkit-coreml/openai_whisper-<variant>`.
|
|
public struct ModelStorage: Sendable {
|
|
public static let repo = "argmaxinc/whisperkit-coreml"
|
|
public static let requiredModels = ["MelSpectrogram", "AudioEncoder", "TextDecoder"]
|
|
|
|
public let base: URL
|
|
|
|
public init(base: URL = ModelStorage.defaultBase) {
|
|
self.base = base
|
|
}
|
|
|
|
/// Application Support, so models are neither in ~/Documents (privacy prompt, iCloud) nor re-fetched.
|
|
public static var defaultBase: URL {
|
|
FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0]
|
|
.appending(path: "Sst/models", directoryHint: .isDirectory)
|
|
}
|
|
|
|
var repoDirectory: URL {
|
|
base.appending(path: "models/\(Self.repo)", directoryHint: .isDirectory)
|
|
}
|
|
|
|
private func folders(for variant: String) -> [URL] {
|
|
let names = (try? FileManager.default.contentsOfDirectory(atPath: repoDirectory.path)) ?? []
|
|
return names
|
|
.filter { $0 == variant || $0 == "openai_whisper-\(variant)" }
|
|
.map { repoDirectory.appending(path: $0, directoryHint: .isDirectory) }
|
|
}
|
|
|
|
/// The folder of a completely downloaded variant, or nil.
|
|
public func folder(for variant: String) -> URL? {
|
|
folders(for: variant).first { folder in
|
|
Self.requiredModels.allSatisfy {
|
|
FileManager.default.fileExists(atPath: folder.appending(path: "\($0).mlmodelc/coremldata.bin").path)
|
|
}
|
|
}
|
|
}
|
|
|
|
public func isInstalled(_ variant: String) -> Bool {
|
|
folder(for: variant) != nil
|
|
}
|
|
|
|
public func delete(_ variant: String) throws {
|
|
for folder in folders(for: variant) { try FileManager.default.removeItem(at: folder) }
|
|
}
|
|
}
|