bleam 0.0.8 → 0.0.10

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

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

@@ -0,0 +1,134 @@
+// VAEConfig.swift - Flux.2 VAE Configuration
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+
+/// Configuration for the Flux.2 VAE (AutoencoderKL)
+public struct VAEConfig: Codable, Sendable {
+    /// Number of input/output channels (3 for RGB)
+    public var inChannels: Int
+
+    /// Number of output channels (same as input)
+    public var outChannels: Int
+
+    /// Number of latent channels (32 for Flux.2, vs 4 for Flux.1)
+    public var latentChannels: Int
+
+    /// Channel multipliers for each block
+    public var blockOutChannels: [Int]
+
+    /// Number of ResNet blocks per level
+    public var layersPerBlock: Int
+
+    /// Activation function
+    public var activationFunction: String
+
+    /// Normalization layer epsilon
+    public var normEps: Float
+
+    /// Number of groups for group normalization
+    public var normNumGroups: Int
+
+    /// Scaling factor for latent space
+    public var scalingFactor: Float
+
+    /// Whether to use BatchNorm for latent normalization
+    public var useBatchNorm: Bool
+
+    /// Patch size for latent packing
+    public var patchSize: (Int, Int)
+
+    public init(
+        inChannels: Int = 3,
+        outChannels: Int = 3,
+        latentChannels: Int = 32,
+        blockOutChannels: [Int] = [128, 256, 512, 512],
+        layersPerBlock: Int = 2,
+        activationFunction: String = "silu",
+        normEps: Float = 1e-6,
+        normNumGroups: Int = 32,
+        scalingFactor: Float = 0.18215,
+        useBatchNorm: Bool = true,
+        patchSize: (Int, Int) = (2, 2)
+    ) {
+        self.inChannels = inChannels
+        self.outChannels = outChannels
+        self.latentChannels = latentChannels
+        self.blockOutChannels = blockOutChannels
+        self.layersPerBlock = layersPerBlock
+        self.activationFunction = activationFunction
+        self.normEps = normEps
+        self.normNumGroups = normNumGroups
+        self.scalingFactor = scalingFactor
+        self.useBatchNorm = useBatchNorm
+        self.patchSize = patchSize
+    }
+
+    /// Default Flux.2 VAE configuration
+    public static let flux2Dev = VAEConfig()
+
+    // MARK: - Codable
+
+    enum CodingKeys: String, CodingKey {
+        case inChannels = "in_channels"
+        case outChannels = "out_channels"
+        case latentChannels = "latent_channels"
+        case blockOutChannels = "block_out_channels"
+        case layersPerBlock = "layers_per_block"
+        case activationFunction = "act_fn"
+        case normEps = "norm_eps"
+        case normNumGroups = "norm_num_groups"
+        case scalingFactor = "scaling_factor"
+        case useBatchNorm = "use_batch_norm"
+    }
+
+    public init(from decoder: Decoder) throws {
+        let container = try decoder.container(keyedBy: CodingKeys.self)
+
+        inChannels = try container.decodeIfPresent(Int.self, forKey: .inChannels) ?? 3
+        outChannels = try container.decodeIfPresent(Int.self, forKey: .outChannels) ?? 3
+        latentChannels = try container.decodeIfPresent(Int.self, forKey: .latentChannels) ?? 32
+        blockOutChannels = try container.decodeIfPresent([Int].self, forKey: .blockOutChannels) ?? [128, 256, 512, 512]
+        layersPerBlock = try container.decodeIfPresent(Int.self, forKey: .layersPerBlock) ?? 2
+        activationFunction = try container.decodeIfPresent(String.self, forKey: .activationFunction) ?? "silu"
+        normEps = try container.decodeIfPresent(Float.self, forKey: .normEps) ?? 1e-6
+        normNumGroups = try container.decodeIfPresent(Int.self, forKey: .normNumGroups) ?? 32
+        scalingFactor = try container.decodeIfPresent(Float.self, forKey: .scalingFactor) ?? 0.18215
+        useBatchNorm = try container.decodeIfPresent(Bool.self, forKey: .useBatchNorm) ?? true
+        patchSize = (2, 2) // Not in JSON, hardcoded for Flux.2
+    }
+
+    public func encode(to encoder: Encoder) throws {
+        var container = encoder.container(keyedBy: CodingKeys.self)
+
+        try container.encode(inChannels, forKey: .inChannels)
+        try container.encode(outChannels, forKey: .outChannels)
+        try container.encode(latentChannels, forKey: .latentChannels)
+        try container.encode(blockOutChannels, forKey: .blockOutChannels)
+        try container.encode(layersPerBlock, forKey: .layersPerBlock)
+        try container.encode(activationFunction, forKey: .activationFunction)
+        try container.encode(normEps, forKey: .normEps)
+        try container.encode(normNumGroups, forKey: .normNumGroups)
+        try container.encode(scalingFactor, forKey: .scalingFactor)
+        try container.encode(useBatchNorm, forKey: .useBatchNorm)
+    }
+
+    /// Load configuration from a JSON file
+    public static func load(from url: URL) throws -> VAEConfig {
+        let data = try Data(contentsOf: url)
+        return try JSONDecoder().decode(VAEConfig.self, from: data)
+    }
+}
+
+extension VAEConfig: CustomStringConvertible {
+    public var description: String {
+        """
+        VAEConfig(
+            channels: \(inChannels) -> latent \(latentChannels) -> \(outChannels),
+            blocks: \(blockOutChannels),
+            batchNorm: \(useBatchNorm),
+            patch: \(patchSize)
+        )
+        """
+    }
+}

package/templates/native/ios/Bleam/AI/Flux2/Flux2Core.swift

@@ -0,0 +1,55 @@
+import Foundation
+import MLX
+import MLXNN
+
+public enum Flux2Core {
+    public static let version = "0.1.0-image-studio"
+}
+
+public enum Flux2Error: Error, LocalizedError {
+    case modelNotLoaded(String)
+    case invalidConfiguration(String)
+    case insufficientMemory(required: Int, available: Int)
+    case weightLoadingFailed(String)
+    case imageProcessingFailed(String)
+    case generationFailed(String)
+    case generationCancelled
+
+    public var errorCode: String {
+        switch self {
+        case .modelNotLoaded:
+            return "model_not_loaded"
+        case .invalidConfiguration:
+            return "invalid_configuration"
+        case .insufficientMemory:
+            return "insufficient_memory"
+        case .weightLoadingFailed:
+            return "weight_loading_failed"
+        case .imageProcessingFailed:
+            return "image_processing_failed"
+        case .generationFailed:
+            return "generation_failed"
+        case .generationCancelled:
+            return "generation_cancelled"
+        }
+    }
+
+    public var errorDescription: String? {
+        switch self {
+        case .modelNotLoaded(let component):
+            return "Model component not loaded: \(component)"
+        case .invalidConfiguration(let message):
+            return "Invalid configuration: \(message)"
+        case .insufficientMemory(let required, let available):
+            return "Insufficient memory: required \(required / 1_000_000_000)GB, available \(available / 1_000_000_000)GB"
+        case .weightLoadingFailed(let message):
+            return "Failed to load weights: \(message)"
+        case .imageProcessingFailed(let message):
+            return "Image processing failed: \(message)"
+        case .generationFailed(let message):
+            return "Generation failed: \(message)"
+        case .generationCancelled:
+            return "Generation was cancelled"
+        }
+    }
+}