bleam 0.0.8 → 0.0.9

package/templates/native/ios/Bleam/AI/Flux2/Configuration/ModelRegistry.swift

@@ -0,0 +1,425 @@
+import Foundation
+
+public enum ModelRegistry {
+    public enum TransformerRepo: String, CaseIterable, Sendable {
+        case blackForestLabsKlein4B = "black-forest-labs/FLUX.2-klein-4B"
+        case moxinOrgKlein4B8bit = "moxin-org/FLUX.2-klein-4B-8bit-mlx"
+        case themindstudioKlein4B4bit = "themindstudio/flux2-klein-4b-mlx-4bit"
+        // These can be presented as Bonsai models while still using the Flux2 pipeline/runtime.
+        case prismMLBonsaiKlein4B2bit = "prism-ml/bonsai-image-ternary-4B-mlx-2bit"
+        case themindstudioKlein9B4bit = "themindstudio/flux2-klein-9b-mlx-4bit"
+        case aitraderKlein9B8bit = "AITRADER/FLUX2-klein-9B-mlx-8bit"
+        case mlxCommunityKlein9B4bit = "mlx-community/FLUX.2-klein-9B"
+
+        public var profile: Flux2TransformerRepoProfile {
+            switch self {
+            case .blackForestLabsKlein4B:
+                return Flux2TransformerRepoProfile(
+                    repoId: rawValue,
+                    model: .klein4B,
+                    quantization: .bf16,
+                    relativePathComponents: ["models--black-forest-labs--FLUX.2-klein-4B"],
+                    requiredSnapshotPatterns: [
+                        "config.json",
+                        "model_index.json",
+                        "flux-2-klein-4b*.safetensors",
+                        "flux-2-klein-4b*.safetensors.index.json",
+                    ],
+                    estimatedSizeGB: 8,
+                    checkpointFormat: .diffusers
+                )
+            case .moxinOrgKlein4B8bit:
+                return Flux2TransformerRepoProfile(
+                    repoId: rawValue,
+                    model: .klein4B,
+                    quantization: .qint8,
+                    relativePathComponents: ["models--moxin-org--FLUX.2-klein-4B-8bit-mlx", "transformer"],
+                    requiredSnapshotPatterns: [
+                        "transformer/config.json",
+                        "transformer/quantized_8bit.safetensors",
+                    ],
+                    estimatedSizeGB: 4,
+                    checkpointFormat: .diffusersQuanto
+                )
+            case .themindstudioKlein4B4bit:
+                return Flux2TransformerRepoProfile(
+                    repoId: rawValue,
+                    model: .klein4B,
+                    quantization: .int4,
+                    relativePathComponents: ["models--themindstudio--flux2-klein-4b-mlx-4bit", "transformer"],
+                    requiredSnapshotPatterns: [
+                        "transformer/model.safetensors.index.json",
+                        "transformer/*.safetensors",
+                    ],
+                    estimatedSizeGB: 5,
+                    checkpointFormat: .packedDiffusers,
+                    needsNormOutSwap: false,
+                    usesNormOutBias: true
+                )
+            case .prismMLBonsaiKlein4B2bit:
+                return Flux2TransformerRepoProfile(
+                    repoId: rawValue,
+                    model: .klein4B,
+                    quantization: .int4,
+                    relativePathComponents: ["models--prism-ml--bonsai-image-ternary-4B-mlx-2bit", "transformer-packed-mflux"],
+                    requiredSnapshotPatterns: [
+                        "transformer-packed-mflux/config.json",
+                        "transformer-packed-mflux/diffusion_pytorch_model.safetensors",
+                        "transformer-packed-mflux/quantization_config.json",
+                    ],
+                    estimatedSizeGB: 3,
+                    checkpointFormat: .packedDiffusers,
+                    needsNormOutSwap: false,
+                    usesNormOutBias: false
+                )
+            case .themindstudioKlein9B4bit:
+                return Flux2TransformerRepoProfile(
+                    repoId: rawValue,
+                    model: .klein9B,
+                    quantization: .int4,
+                    relativePathComponents: ["models--themindstudio--flux2-klein-9b-mlx-4bit", "transformer"],
+                    requiredSnapshotPatterns: [
+                        "transformer/model.safetensors.index.json",
+                        "transformer/*.safetensors",
+                    ],
+                    estimatedSizeGB: 20,
+                    checkpointFormat: .packedDiffusers,
+                    needsNormOutSwap: false,
+                    usesNormOutBias: true
+                )
+            case .aitraderKlein9B8bit:
+                return Flux2TransformerRepoProfile(
+                    repoId: rawValue,
+                    model: .klein9B,
+                    quantization: .qint8,
+                    relativePathComponents: ["models--AITRADER--FLUX2-klein-9B-mlx-8bit", "transformer"],
+                    requiredSnapshotPatterns: [
+                        "transformer/model.safetensors.index.json",
+                        "transformer/*.safetensors",
+                    ],
+                    estimatedSizeGB: 10,
+                    checkpointFormat: .packedDiffusers,
+                    needsNormOutSwap: false,
+                    usesNormOutBias: true
+                )
+            case .mlxCommunityKlein9B4bit:
+                return Flux2TransformerRepoProfile(
+                    repoId: rawValue,
+                    model: .klein9B,
+                    quantization: .int4,
+                    relativePathComponents: ["models--mlx-community--FLUX.2-klein-9B", "transformer"],
+                    requiredSnapshotPatterns: [
+                        "transformer/model.safetensors.index.json",
+                        "transformer/*.safetensors",
+                    ],
+                    estimatedSizeGB: 20,
+                    checkpointFormat: .packedDiffusers,
+                    needsNormOutSwap: false,
+                    usesNormOutBias: true
+                )
+            }
+        }
+
+        public static func repo(for model: Flux2Model, quantization: TransformerQuantization) -> TransformerRepo {
+            switch model {
+            case .klein4B:
+                switch quantization {
+                case .bf16:
+                    return .blackForestLabsKlein4B
+                case .int4:
+                    return .themindstudioKlein4B4bit
+                case .qint8:
+                    return .moxinOrgKlein4B8bit
+                }
+            case .klein9B:
+                switch quantization {
+                case .qint8:
+                    return .aitraderKlein9B8bit
+                case .int4:
+                    return .themindstudioKlein9B4bit
+                case .bf16:
+                    return .mlxCommunityKlein9B4bit
+                }
+            }
+        }
+
+        public static func profile(for repoId: String) -> Flux2TransformerRepoProfile? {
+            allCases.first { $0.rawValue == repoId }?.profile
+        }
+    }
+
+    public enum TransformerVariant: String, CaseIterable, Sendable {
+        case klein4B_bf16 = "klein4b-bf16"
+        case klein4B_8bit = "klein4b-8bit"
+        case klein4B_4bit = "klein4b-4bit"
+        case bonsaiKlein4B_2bit = "bonsai-klein4b-2bit"
+        case bonsaiKlein4B_1bit = "bonsai-klein4b-1bit"
+        case klein9B_8bit = "klein9b-8bit"
+        case klein9B_4bit = "klein9b-4bit"
+
+        public var huggingFaceRepo: String {
+            switch self {
+            case .klein4B_bf16:
+                return "black-forest-labs/FLUX.2-klein-4B"
+            case .klein4B_8bit:
+                return "moxin-org/FLUX.2-klein-4B-8bit-mlx"
+            case .klein4B_4bit:
+                return "themindstudio/flux2-klein-4b-mlx-4bit"
+            case .bonsaiKlein4B_2bit:
+                return "prism-ml/bonsai-image-ternary-4B-mlx-2bit"
+            case .bonsaiKlein4B_1bit:
+                return "prism-ml/bonsai-image-binary-4B-mlx-1bit"
+            case .klein9B_8bit:
+                return "AITRADER/FLUX2-klein-9B-mlx-8bit"
+            case .klein9B_4bit:
+                return "themindstudio/flux2-klein-9b-mlx-4bit"
+            }
+        }
+
+        public var huggingFaceSubfolder: String? {
+            switch self {
+            case .klein4B_bf16:
+                return nil
+            case .klein4B_8bit, .klein4B_4bit, .bonsaiKlein4B_2bit, .bonsaiKlein4B_1bit, .klein9B_8bit, .klein9B_4bit:
+                return "transformer"
+            }
+        }
+
+        public var requiredSnapshotPatterns: [String] {
+            switch self {
+            case .klein4B_bf16:
+                return [
+                    "config.json",
+                    "model_index.json",
+                    "flux-2-klein-4b*.safetensors",
+                    "flux-2-klein-4b*.safetensors.index.json",
+                ]
+            case .klein4B_8bit:
+                return [
+                    "transformer/config.json",
+                    "transformer/quantized_8bit.safetensors",
+                ]
+            case .klein4B_4bit:
+                return [
+                    "transformer/model.safetensors.index.json",
+                    "transformer/*.safetensors",
+                ]
+            case .bonsaiKlein4B_2bit:
+                return [
+                    "transformer-packed-mflux/config.json",
+                    "transformer-packed-mflux/diffusion_pytorch_model.safetensors",
+                    "transformer-packed-mflux/quantization_config.json",
+                ]
+            case .bonsaiKlein4B_1bit:
+                return [
+                    "transformer-packed-mflux/config.json",
+                    "transformer-packed-mflux/diffusion_pytorch_model.safetensors",
+                    "transformer-packed-mflux/quantization_config.json",
+                ]
+            case .klein9B_8bit:
+                return [
+                    "transformer/model.safetensors.index.json",
+                    "transformer/*.safetensors",
+                ]
+            case .klein9B_4bit:
+                return [
+                    "transformer/model.safetensors.index.json",
+                    "transformer/*.safetensors",
+                ]
+            }
+        }
+
+        public var estimatedSizeGB: Int {
+            switch self {
+            case .klein4B_bf16: return 8
+            case .klein4B_8bit: return 4
+            case .klein4B_4bit: return 5
+            case .bonsaiKlein4B_2bit: return 3
+            case .bonsaiKlein4B_1bit: return 2
+            case .klein9B_8bit: return 10
+            case .klein9B_4bit: return 20
+            }
+        }
+
+        public var isGated: Bool { false }
+        public var license: String { "Apache 2.0" }
+        public var isCommercialUseAllowed: Bool { true }
+        public var defaultSteps: Int { 4 }
+        public var defaultGuidance: Float { 1.0 }
+        public var maxReferenceImages: Int { 4 }
+
+        public var quantization: TransformerQuantization {
+            switch self {
+            case .klein4B_bf16: return .bf16
+            case .klein4B_8bit: return .qint8
+            case .klein4B_4bit: return .int4
+            case .bonsaiKlein4B_2bit: return .int4
+            case .bonsaiKlein4B_1bit: return .int4
+            case .klein9B_8bit: return .qint8
+            case .klein9B_4bit: return .int4
+            }
+        }
+
+        public var modelType: Flux2Model {
+            switch self {
+            case .klein4B_bf16, .klein4B_8bit, .klein4B_4bit, .bonsaiKlein4B_2bit, .bonsaiKlein4B_1bit: return .klein4B
+            case .klein9B_8bit, .klein9B_4bit: return .klein9B
+            }
+        }
+
+        public var isForInference: Bool { true }
+        public var isForTraining: Bool { false }
+
+        public static func variant(for model: Flux2Model, quantization: TransformerQuantization) -> TransformerVariant {
+            switch model {
+            case .klein4B:
+                switch quantization {
+                case .bf16: return .klein4B_bf16
+                case .qint8: return .klein4B_8bit
+                case .int4: return .klein4B_4bit
+                }
+            case .klein9B:
+                switch quantization {
+                case .qint8: return .klein9B_8bit
+                case .int4: return .klein9B_4bit
+                case .bf16: return .klein9B_4bit
+                }
+            }
+        }
+    }
+
+    public enum VAEVariant: String, CaseIterable, Sendable {
+        case standard = "standard"
+        case smallDecoder = "small-decoder"
+        case bonsaiTernary = "bonsai-ternary"
+
+        public var huggingFaceRepo: String {
+            switch self {
+            case .standard:
+                return "black-forest-labs/FLUX.2-klein-4B"
+            case .smallDecoder:
+                return "black-forest-labs/FLUX.2-dev"
+            case .bonsaiTernary:
+                return "prism-ml/bonsai-image-ternary-4B-mlx-2bit"
+            }
+        }
+
+        public var huggingFaceSubfolder: String {
+            switch self {
+            case .standard, .smallDecoder, .bonsaiTernary:
+                return "vae"
+            }
+        }
+
+        public var requiredSnapshotPatterns: [String] {
+            switch self {
+            case .standard, .bonsaiTernary:
+                return [
+                    "vae/config.json",
+                    "vae/diffusion_pytorch_model.safetensors",
+                    "vae/diffusion_pytorch_model.safetensors.index.json",
+                    "vae/diffusion_pytorch_model-*.safetensors",
+                ]
+            case .smallDecoder:
+                return [
+                    "vae/config.json",
+                    "vae/diffusion_pytorch_model.safetensors",
+                    "vae/diffusion_pytorch_model.safetensors.index.json",
+                    "vae/diffusion_pytorch_model-*.safetensors",
+                ]
+            }
+        }
+
+        public var estimatedSizeGB: Int {
+            switch self {
+            case .standard, .bonsaiTernary: return 3
+            case .smallDecoder: return 1
+            }
+        }
+        public var isGated: Bool { false }
+        public var license: String { "Apache 2.0" }
+        public var isCommercialUseAllowed: Bool { true }
+    }
+
+    public enum ModelComponent: Hashable, Sendable {
+        case transformer(Flux2TransformerRepoProfile)
+        case textEncoder
+        case vae(VAEVariant)
+
+        public var displayName: String {
+            switch self {
+            case .transformer(let profile):
+                return "Flux.2 Transformer (\(profile.repoId))"
+            case .textEncoder:
+                return "Qwen3 Text Encoder"
+            case .vae:
+                return "Flux.2 VAE"
+            }
+        }
+
+        public var estimatedSizeGB: Int {
+            switch self {
+            case .transformer(let profile): return profile.estimatedSizeGB
+            case .textEncoder: return 4
+            case .vae(let variant): return variant.estimatedSizeGB
+            }
+        }
+
+        public var requiredSnapshotPatterns: [String] {
+            switch self {
+            case .transformer(let profile):
+                return profile.requiredSnapshotPatterns
+            case .textEncoder:
+                return []
+            case .vae(let variant):
+                return variant.requiredSnapshotPatterns
+            }
+        }
+    }
+
+    /// The `models/flux/` subdirectory inside the given container root.
+    public static func familyDirectory(containerRoot: URL) -> URL {
+        containerRoot
+            .appendingPathComponent("models", isDirectory: true)
+            .appendingPathComponent("flux", isDirectory: true)
+    }
+
+    /// The `models/flux/LoRAs/` subdirectory inside the given container root.
+    public static func loRAsDirectory(containerRoot: URL) -> URL {
+        familyDirectory(containerRoot: containerRoot)
+            .appendingPathComponent("LoRAs", isDirectory: true)
+    }
+
+    public static func localPath(for component: ModelComponent, containerRoot: URL) -> URL {
+        let family = familyDirectory(containerRoot: containerRoot)
+        switch component {
+        case .transformer(let profile):
+            return profile.relativePathComponents.reduce(family) { partial, component in
+                partial.appendingPathComponent(component, isDirectory: true)
+            }
+        case .textEncoder:
+            return family
+        case .vae(let variant):
+            let repoDir = variant.huggingFaceRepo
+                .replacingOccurrences(of: "/", with: "--")
+            return family
+                .appendingPathComponent("models--\(repoDir)")
+                .appendingPathComponent("vae")
+        }
+    }
+
+    public static func recommendedConfig(forRAMGB ram: Int) -> Flux2QuantizationConfig {
+        switch ram {
+        case 0..<32:
+            return .ultraMinimal
+        case 32..<48:
+            return .minimal
+        case 48..<64:
+            return .memoryEfficient
+        case 64..<96:
+            return .balanced
+        default:
+            return .highQuality
+        }
+    }
+}

package/templates/native/ios/Bleam/AI/Flux2/Configuration/QuantizationConfig.swift

@@ -0,0 +1,138 @@
+// QuantizationConfig.swift - Fine-grained quantization configuration
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+
+/// Quantization level for the Mistral text encoder
+public enum MistralQuantization: String, CaseIterable, Codable, Sendable {
+    case bf16 = "bf16"      // Full precision ~48GB
+    case mlx8bit = "8bit"   // 8-bit ~25GB
+    case mlx6bit = "6bit"   // 6-bit ~19GB
+    case mlx4bit = "4bit"   // 4-bit ~14GB
+
+    public var estimatedMemoryGB: Int {
+        switch self {
+        case .bf16: return 48
+        case .mlx8bit: return 25
+        case .mlx6bit: return 19
+        case .mlx4bit: return 14
+        }
+    }
+
+    public var displayName: String {
+        switch self {
+        case .bf16: return "Full Precision (bf16)"
+        case .mlx8bit: return "8-bit"
+        case .mlx6bit: return "6-bit"
+        case .mlx4bit: return "4-bit"
+        }
+    }
+}
+
+/// Quantization level for the Flux.2 diffusion transformer
+public enum TransformerQuantization: String, CaseIterable, Codable, Sendable {
+    case bf16 = "bf16"      // Full precision ~64GB
+    case qint8 = "qint8"    // 8-bit ~32GB
+    case int4 = "int4"      // 4-bit ~16GB (on-the-fly quantization)
+
+    public var estimatedMemoryGB: Int {
+        switch self {
+        case .bf16: return 64
+        case .qint8: return 32
+        case .int4: return 16
+        }
+    }
+
+    public var displayName: String {
+        switch self {
+        case .bf16: return "Full Precision (bf16)"
+        case .qint8: return "8-bit (qint8)"
+        case .int4: return "4-bit (int4)"
+        }
+    }
+
+    public var bits: Int {
+        switch self {
+        case .bf16: return 16
+        case .qint8: return 8
+        case .int4: return 4
+        }
+    }
+
+    public var groupSize: Int { 64 }
+}
+
+/// Independent quantization configuration for text encoder and transformer
+public struct Flux2QuantizationConfig: Codable, Sendable {
+    /// Quantization for the Mistral text encoder
+    public var textEncoder: MistralQuantization
+
+    /// Quantization for the Flux.2 diffusion transformer
+    public var transformer: TransformerQuantization
+
+    public init(
+        textEncoder: MistralQuantization,
+        transformer: TransformerQuantization
+    ) {
+        self.textEncoder = textEncoder
+        self.transformer = transformer
+    }
+
+    /// Total estimated memory requirement in GB
+    public var estimatedTotalMemoryGB: Int {
+        // Text encoder and transformer are never loaded simultaneously
+        // So we take the max of the two, plus VAE (~3GB) and working memory (~5GB)
+        max(textEncoder.estimatedMemoryGB, transformer.estimatedMemoryGB) + 8
+    }
+
+    /// Peak memory during text encoding phase
+    public var textEncodingPhaseMemoryGB: Int {
+        textEncoder.estimatedMemoryGB + 1  // +1GB for embeddings
+    }
+
+    /// Peak memory during image generation phase
+    public var imageGenerationPhaseMemoryGB: Int {
+        transformer.estimatedMemoryGB + 3 + 5  // +3GB VAE, +5GB working memory
+    }
+
+    // MARK: - Presets
+
+    /// High quality preset - requires ~90GB+ RAM
+    public static let highQuality = Flux2QuantizationConfig(
+        textEncoder: .bf16,
+        transformer: .bf16
+    )
+
+    /// Balanced preset - requires ~57GB RAM (recommended for 64GB Macs)
+    public static let balanced = Flux2QuantizationConfig(
+        textEncoder: .mlx8bit,
+        transformer: .qint8
+    )
+
+    /// Memory efficient preset - requires ~47GB RAM
+    public static let memoryEfficient = Flux2QuantizationConfig(
+        textEncoder: .mlx4bit,
+        transformer: .qint8
+    )
+
+    /// Minimal preset - requires ~47GB RAM
+    public static let minimal = Flux2QuantizationConfig(
+        textEncoder: .mlx4bit,
+        transformer: .qint8
+    )
+
+    /// Ultra-minimal preset - requires ~30GB RAM (4-bit transformer)
+    public static let ultraMinimal = Flux2QuantizationConfig(
+        textEncoder: .mlx4bit,
+        transformer: .int4
+    )
+
+    /// Default preset (balanced)
+    public static let `default` = balanced
+}
+
+extension Flux2QuantizationConfig: CustomStringConvertible {
+    public var description: String {
+        "Flux2QuantizationConfig(text: \(textEncoder.rawValue), transformer: \(transformer.rawValue), ~\(estimatedTotalMemoryGB)GB)"
+    }
+}

package/templates/native/ios/Bleam/AI/Flux2/Configuration/TransformerRepoProfile.swift

@@ -0,0 +1,59 @@
+import Foundation
+
+public enum Flux2TransformerCheckpointFormat: String, Sendable {
+    case diffusers
+    case diffusersQuanto
+    case packedDiffusers
+    case bfl
+}
+
+public struct Flux2TransformerRepoProfile: Hashable, Sendable {
+    public let repoId: String
+    public let model: Flux2Model
+    public let quantization: TransformerQuantization
+    public let relativePathComponents: [String]
+    public let requiredSnapshotPatterns: [String]
+    public let estimatedSizeGB: Int
+    public let checkpointFormat: Flux2TransformerCheckpointFormat
+    public let needsNormOutSwap: Bool
+    public let usesNormOutBias: Bool
+
+    public init(
+        repoId: String,
+        model: Flux2Model,
+        quantization: TransformerQuantization,
+        relativePathComponents: [String],
+        requiredSnapshotPatterns: [String],
+        estimatedSizeGB: Int,
+        checkpointFormat: Flux2TransformerCheckpointFormat,
+        needsNormOutSwap: Bool = false,
+        usesNormOutBias: Bool = false
+    ) {
+        self.repoId = repoId
+        self.model = model
+        self.quantization = quantization
+        self.relativePathComponents = relativePathComponents
+        self.requiredSnapshotPatterns = requiredSnapshotPatterns
+        self.estimatedSizeGB = estimatedSizeGB
+        self.checkpointFormat = checkpointFormat
+        self.needsNormOutSwap = needsNormOutSwap
+        self.usesNormOutBias = usesNormOutBias
+    }
+
+    public func withRuntimeOverrides(
+        needsNormOutSwap: Bool? = nil,
+        usesNormOutBias: Bool? = nil
+    ) -> Flux2TransformerRepoProfile {
+        Flux2TransformerRepoProfile(
+            repoId: repoId,
+            model: model,
+            quantization: quantization,
+            relativePathComponents: relativePathComponents,
+            requiredSnapshotPatterns: requiredSnapshotPatterns,
+            estimatedSizeGB: estimatedSizeGB,
+            checkpointFormat: checkpointFormat,
+            needsNormOutSwap: needsNormOutSwap ?? self.needsNormOutSwap,
+            usesNormOutBias: usesNormOutBias ?? self.usesNormOutBias
+        )
+    }
+}