bleam 0.0.8 → 0.0.10

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

@@ -0,0 +1,157 @@
+// MemoryOptimizationConfig.swift - Memory optimization settings
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+
+/// Configuration for memory optimization during transformer inference
+///
+/// These settings control how aggressively the transformer manages memory
+/// by forcing evaluation of the computation graph at intervals.
+///
+/// # Background
+/// MLX uses lazy evaluation and builds a computation graph. For large models,
+/// this graph can grow very large and consume significant memory. Periodically
+/// calling `eval()` materializes intermediate results and allows the graph to
+/// be garbage collected.
+///
+/// # Trade-offs
+/// - Lower `evalFrequency` = more memory savings, but slower inference
+/// - Higher `evalFrequency` = faster inference, but more memory usage
+/// - `clearCacheOnEval` = additional memory savings at cost of recomputation
+///
+/// # Reference
+/// Based on optimizations from LTX-2-MLX: https://github.com/Acelogic/LTX-2-MLX
+public struct MemoryOptimizationConfig: Sendable, Equatable {
+
+    /// How often to evaluate the computation graph during transformer forward pass.
+    ///
+    /// - `0` = disabled (fastest, most memory)
+    /// - `4` = evaluate every 4 blocks (aggressive memory savings)
+    /// - `8` = evaluate every 8 blocks (moderate, recommended)
+    /// - `16` = evaluate every 16 blocks (light memory savings)
+    ///
+    /// For Flux.2 Dev with 8 double + 48 single blocks:
+    /// - `evalFrequency: 8` will eval ~7 times during forward pass
+    /// - `evalFrequency: 4` will eval ~14 times during forward pass
+    public var evalFrequency: Int
+
+    /// Whether to clear GPU cache after each evaluation.
+    ///
+    /// When enabled, calls `MLX.Memory.clearCache()` after each periodic eval.
+    /// This provides additional memory savings but may cause recomputation
+    /// of cached values.
+    ///
+    /// Recommended: `false` for most cases, `true` only if running out of memory
+    public var clearCacheOnEval: Bool
+
+    /// Whether to evaluate after all double-stream blocks complete.
+    ///
+    /// This ensures the transition between double-stream and single-stream
+    /// phases has a clean memory state.
+    public var evalBetweenPhases: Bool
+
+    public init(
+        evalFrequency: Int = 0,
+        clearCacheOnEval: Bool = false,
+        evalBetweenPhases: Bool = true
+    ) {
+        self.evalFrequency = max(0, evalFrequency)
+        self.clearCacheOnEval = clearCacheOnEval
+        self.evalBetweenPhases = evalBetweenPhases
+    }
+
+    // MARK: - Presets
+
+    /// Disabled - no periodic evaluation (fastest, most memory)
+    ///
+    /// Use when you have plenty of RAM and want maximum speed.
+    public static let disabled = MemoryOptimizationConfig(
+        evalFrequency: 0,
+        clearCacheOnEval: false,
+        evalBetweenPhases: false
+    )
+
+    /// Light - evaluate every 16 blocks
+    ///
+    /// Minimal impact on speed with some memory savings.
+    public static let light = MemoryOptimizationConfig(
+        evalFrequency: 16,
+        clearCacheOnEval: false,
+        evalBetweenPhases: true
+    )
+
+    /// Moderate - evaluate every 8 blocks (recommended default)
+    ///
+    /// Good balance between speed and memory usage.
+    public static let moderate = MemoryOptimizationConfig(
+        evalFrequency: 8,
+        clearCacheOnEval: false,
+        evalBetweenPhases: true
+    )
+
+    /// Aggressive - evaluate every 4 blocks with cache clearing
+    ///
+    /// Use when running low on memory. Slower but uses significantly less RAM.
+    public static let aggressive = MemoryOptimizationConfig(
+        evalFrequency: 4,
+        clearCacheOnEval: true,
+        evalBetweenPhases: true
+    )
+
+    /// Ultra low memory - evaluate every 2 blocks with cache clearing
+    ///
+    /// Maximum memory savings at significant speed cost.
+    /// Use only if other presets cause out-of-memory errors.
+    public static let ultraLowMemory = MemoryOptimizationConfig(
+        evalFrequency: 2,
+        clearCacheOnEval: true,
+        evalBetweenPhases: true
+    )
+}
+
+// MARK: - Description
+
+extension MemoryOptimizationConfig: CustomStringConvertible {
+    public var description: String {
+        if evalFrequency == 0 && !evalBetweenPhases {
+            return "MemoryOptimization(disabled)"
+        }
+
+        var parts: [String] = []
+        if evalFrequency > 0 {
+            parts.append("evalEvery: \(evalFrequency)")
+        }
+        if clearCacheOnEval {
+            parts.append("clearCache: true")
+        }
+        if evalBetweenPhases {
+            parts.append("evalBetweenPhases: true")
+        }
+
+        return "MemoryOptimization(\(parts.joined(separator: ", ")))"
+    }
+}
+
+// MARK: - Recommendation
+
+extension MemoryOptimizationConfig {
+
+    /// Get recommended configuration based on available RAM
+    ///
+    /// - Parameter ramGB: System RAM in gigabytes
+    /// - Returns: Recommended memory optimization preset
+    public static func recommended(forRAMGB ramGB: Int) -> MemoryOptimizationConfig {
+        switch ramGB {
+        case 0..<32:
+            return .ultraLowMemory
+        case 32..<64:
+            return .aggressive
+        case 64..<96:
+            return .moderate
+        case 96..<128:
+            return .light
+        default:
+            return .disabled
+        }
+    }
+}

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
+        }
+    }
+}