You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

107 lines
3.7 KiB

//
// ModelInfo.swift
// Diffusion
//
// Created by Pedro Cuenca on 29/12/22.
// See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE
//
import CoreML
struct ModelInfo {
/// Hugging Face model Id that contains .zip archives with compiled Core ML models
let modelId: String
/// Arbitrary string for presentation purposes. Something like "2.1-base"
let modelVersion: String
/// Suffix of the archive containing the ORIGINAL attention variant. Usually something like "original_compiled"
let originalAttentionSuffix: String
/// Suffix of the archive containing the SPLIT_EINSUM attention variant. Usually something like "split_einsum_compiled"
let splitAttentionName: String
/// Whether the archive contains the VAE Encoder (for image to image tasks). Not yet in use.
let supportsEncoder: Bool
init(modelId: String, modelVersion: String, originalAttentionSuffix: String = "original_compiled", splitAttentionName: String = "split_einsum_compiled", supportsEncoder: Bool = false) {
self.modelId = modelId
self.modelVersion = modelVersion
self.originalAttentionSuffix = originalAttentionSuffix
self.splitAttentionName = splitAttentionName
self.supportsEncoder = supportsEncoder
}
}
extension ModelInfo {
/// Best variant for the current platform.
/// Currently using `split_einsum` for iOS and `original` for macOS, but could vary depending on model.
var bestURL: URL {
// Pattern: https://huggingface.co/pcuenq/coreml-stable-diffusion/resolve/main/coreml-stable-diffusion-v1-5_original_compiled.zip
let suffix = runningOnMac ? originalAttentionSuffix : splitAttentionName
let repo = modelId.split(separator: "/").last!
return URL(string: "https://huggingface.co/\(modelId)/resolve/main/\(repo)_\(suffix).zip")!
}
/// Best units for current platform.
/// Currently using `cpuAndNeuralEngine` for iOS and `cpuAndGPU` for macOS, but could vary depending on model.
/// .all works for v1.4, but not for v1.5.
// TODO: measure performance on different devices.
var bestComputeUnits: MLComputeUnits {
return runningOnMac ? .cpuAndGPU : .cpuAndNeuralEngine
}
var reduceMemory: Bool {
return !runningOnMac
}
}
extension ModelInfo {
// TODO: repo does not exist yet
static let v14Base = ModelInfo(
modelId: "pcuenq/coreml-stable-diffusion-1-4",
modelVersion: "CompVis/stable-diffusion-v1-4"
)
static let v15Base = ModelInfo(
modelId: "pcuenq/coreml-stable-diffusion-v1-5",
modelVersion: "runwayml/stable-diffusion-v1-5"
)
static let v2Base = ModelInfo(
modelId: "pcuenq/coreml-stable-diffusion-2-base",
modelVersion: "stabilityai/stable-diffusion-2-base"
)
static let v21Base = ModelInfo(
modelId: "pcuenq/coreml-stable-diffusion-2-1-base",
modelVersion: "stabilityai/stable-diffusion-2-1-base",
supportsEncoder: true
)
static let ofaSmall = ModelInfo(
modelId: "pcuenq/coreml-small-stable-diffusion-v0",
modelVersion: "OFA-Sys/small-stable-diffusion-v0"
)
static let MODELS = [
ModelInfo.v14Base,
ModelInfo.v15Base,
ModelInfo.v2Base,
ModelInfo.v21Base,
ModelInfo.ofaSmall
]
static func from(modelVersion: String) -> ModelInfo? {
ModelInfo.MODELS.first(where: {$0.modelVersion == modelVersion})
}
static func from(modelId: String) -> ModelInfo? {
ModelInfo.MODELS.first(where: {$0.modelId == modelId})
}
}
extension ModelInfo : Equatable {
static func ==(lhs: ModelInfo, rhs: ModelInfo) -> Bool { lhs.modelId == rhs.modelId }
}