bleam 0.0.8 → 0.0.9

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (87) hide show
  1. package/dist/cli.cjs +23 -4
  2. package/dist/cli.d.cts +2 -0
  3. package/dist/cli.d.ts +2 -0
  4. package/dist/cli.js +23 -4
  5. package/dist/{ui-Ce-pnUUA.d.ts → ui-Bg11tvlc.d.ts} +7 -7
  6. package/dist/ui.d.ts +1 -1
  7. package/dist/window.d.ts +1 -1
  8. package/package.json +4 -1
  9. package/templates/image-generation/app/index.tsx +1 -1
  10. package/templates/native/.gitattributes +1 -0
  11. package/templates/native/App.tsx +26 -0
  12. package/templates/native/app.json +12 -0
  13. package/templates/native/index.ts +8 -0
  14. package/templates/native/ios/.xcode.env +11 -0
  15. package/templates/native/ios/Bleam/AI/Flux2/Configuration/Flux2Config.swift +288 -0
  16. package/templates/native/ios/Bleam/AI/Flux2/Configuration/MemoryConfig.swift +233 -0
  17. package/templates/native/ios/Bleam/AI/Flux2/Configuration/MemoryOptimizationConfig.swift +157 -0
  18. package/templates/native/ios/Bleam/AI/Flux2/Configuration/ModelRegistry.swift +425 -0
  19. package/templates/native/ios/Bleam/AI/Flux2/Configuration/QuantizationConfig.swift +138 -0
  20. package/templates/native/ios/Bleam/AI/Flux2/Configuration/TransformerRepoProfile.swift +59 -0
  21. package/templates/native/ios/Bleam/AI/Flux2/Configuration/VAEConfig.swift +134 -0
  22. package/templates/native/ios/Bleam/AI/Flux2/Flux2Core.swift +55 -0
  23. package/templates/native/ios/Bleam/AI/Flux2/Loading/WeightLoader.swift +1192 -0
  24. package/templates/native/ios/Bleam/AI/Flux2/Pipeline/FluxImageRunner.swift +1282 -0
  25. package/templates/native/ios/Bleam/AI/Flux2/Pipeline/LatentUtils.swift +191 -0
  26. package/templates/native/ios/Bleam/AI/Flux2/Scheduler/FlowMatchEulerScheduler.swift +279 -0
  27. package/templates/native/ios/Bleam/AI/Flux2/Text/KleinTextEncoder.swift +380 -0
  28. package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2Attention.swift +379 -0
  29. package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2Embeddings.swift +159 -0
  30. package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2FeedForward.swift +154 -0
  31. package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2FusedKernels.swift +198 -0
  32. package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2KVCache.swift +36 -0
  33. package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2Modulation.swift +193 -0
  34. package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2ParallelAttention.swift +389 -0
  35. package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2RoPE.swift +304 -0
  36. package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2SingleBlock.swift +290 -0
  37. package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2Transformer.swift +530 -0
  38. package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2TransformerBlock.swift +362 -0
  39. package/templates/native/ios/Bleam/AI/Flux2/Utils/Flux2Debug.swift +101 -0
  40. package/templates/native/ios/Bleam/AI/Flux2/Utils/MLXCheckpoint.swift +118 -0
  41. package/templates/native/ios/Bleam/AI/Flux2/Utils/MemoryManager.swift +201 -0
  42. package/templates/native/ios/Bleam/AI/Flux2/VAE/AutoencoderKL.swift +476 -0
  43. package/templates/native/ios/Bleam/AI/Flux2/VAE/ResnetBlock.swift +316 -0
  44. package/templates/native/ios/Bleam/AI/Flux2/VAE/VAEDecoder.swift +120 -0
  45. package/templates/native/ios/Bleam/AI/Flux2/VAE/VAEEncoder.swift +136 -0
  46. package/templates/native/ios/Bleam/AppDelegate.swift +373 -0
  47. package/templates/native/ios/Bleam/Appearance.swift +62 -0
  48. package/templates/native/ios/Bleam/Bridging-Header.h +3 -0
  49. package/templates/native/ios/Bleam/Images.xcassets/AppIcon.appiconset/Contents.json +13 -0
  50. package/templates/native/ios/Bleam/Images.xcassets/Contents.json +6 -0
  51. package/templates/native/ios/Bleam/Info.plist +76 -0
  52. package/templates/native/ios/Bleam/PrivacyInfo.xcprivacy +48 -0
  53. package/templates/native/ios/Bleam/SceneDelegate.swift +58 -0
  54. package/templates/native/ios/Bleam/SplashScreen.storyboard +47 -0
  55. package/templates/native/ios/Bleam/Supporting/Expo.plist +6 -0
  56. package/templates/native/ios/Bleam/bleam.entitlements +10 -0
  57. package/templates/native/ios/Bleam.xcodeproj/project.pbxproj +1146 -0
  58. package/templates/native/ios/Bleam.xcodeproj/xcshareddata/xcschemes/Bleam.xcscheme +88 -0
  59. package/templates/native/ios/GenerationService/GenerationService.entitlements +10 -0
  60. package/templates/native/ios/GenerationService/ImageGenerationRunner.swift +591 -0
  61. package/templates/native/ios/GenerationService/Info.plist +31 -0
  62. package/templates/native/ios/GenerationService/main.swift +165 -0
  63. package/templates/native/ios/PlatformHelper/Info.plist +29 -0
  64. package/templates/native/ios/PlatformHelper/main.swift +335 -0
  65. package/templates/native/ios/Podfile +198 -0
  66. package/templates/native/ios/Podfile.lock +2284 -0
  67. package/templates/native/ios/Podfile.properties.json +5 -0
  68. package/templates/native/ios/Shared/Generation/GenerationServiceProtocol.swift +12 -0
  69. package/templates/native/ios/Shared/Generation/GenerationWorkerProtocol.swift +124 -0
  70. package/templates/native/metro.config.js +6 -0
  71. package/templates/native/modules/bleam-runtime/BleamRuntime.podspec +14 -0
  72. package/templates/native/modules/bleam-runtime/expo-module.config.json +8 -0
  73. package/templates/native/modules/bleam-runtime/ios/AIModule.swift +1258 -0
  74. package/templates/native/modules/bleam-runtime/ios/GenerationContracts.swift +135 -0
  75. package/templates/native/modules/bleam-runtime/ios/PlatformModule.swift +313 -0
  76. package/templates/native/modules/bleam-runtime/package.json +5 -0
  77. package/templates/native/package.json +39 -0
  78. package/templates/native/scripts/build/mlx-frameworks.ts +421 -0
  79. package/templates/native/scripts/start/index.ts +474 -0
  80. package/templates/native/scripts/start/known-failures.ts +20 -0
  81. package/templates/native/scripts/start/loader.ts +76 -0
  82. package/templates/native/scripts/start/native-state.ts +242 -0
  83. package/templates/native/scripts/start/package.json +4 -0
  84. package/templates/native/scripts/start/swift-packages.ts +414 -0
  85. package/templates/native/scripts/start/xcode-formatter.ts +106 -0
  86. package/templates/native/tsconfig.json +14 -0
  87. package/templates/native/yarn.lock +3749 -0
@@ -0,0 +1,201 @@
1
+ // MemoryManager.swift - GPU and system memory management
2
+ // Copyright 2025 Vincent Gourbin
3
+
4
+ import Foundation
5
+ import MLX
6
+
7
+ /// Memory manager for Flux.2 generation
8
+ ///
9
+ /// Monitors GPU memory usage and helps manage the two-phase
10
+ /// pipeline to fit within available RAM.
11
+ public final class Flux2MemoryManager: @unchecked Sendable {
12
+
13
+ /// Shared instance
14
+ public static let shared = Flux2MemoryManager()
15
+
16
+ /// System physical memory in bytes
17
+ public var physicalMemory: UInt64 {
18
+ ProcessInfo.processInfo.physicalMemory
19
+ }
20
+
21
+ /// System physical memory in GB
22
+ public var physicalMemoryGB: Int {
23
+ Int(physicalMemory / 1_073_741_824)
24
+ }
25
+
26
+ /// Estimated available memory (rough heuristic)
27
+ public var estimatedAvailableMemoryGB: Int {
28
+ // Reserve some memory for system
29
+ max(0, physicalMemoryGB - 8)
30
+ }
31
+
32
+ private init() {}
33
+
34
+ // MARK: - Memory Checks
35
+
36
+ /// Check if we have enough memory for a configuration
37
+ public func canRun(config: Flux2QuantizationConfig) -> Bool {
38
+ let required = config.estimatedTotalMemoryGB
39
+ return required <= estimatedAvailableMemoryGB
40
+ }
41
+
42
+ /// Get recommended configuration for current system
43
+ public func recommendedConfig() -> Flux2QuantizationConfig {
44
+ ModelRegistry.recommendedConfig(forRAMGB: physicalMemoryGB)
45
+ }
46
+
47
+ /// Check memory before text encoding phase
48
+ public func checkTextEncodingPhase(config: Flux2QuantizationConfig) -> MemoryCheckResult {
49
+ let required = config.textEncodingPhaseMemoryGB
50
+ let available = estimatedAvailableMemoryGB
51
+
52
+ if required > available {
53
+ return .insufficientMemory(
54
+ required: required,
55
+ available: available,
56
+ suggestion: "Use a lower text encoder quantization (4-bit instead of 8-bit)"
57
+ )
58
+ }
59
+
60
+ return .ok
61
+ }
62
+
63
+ /// Check memory before image generation phase
64
+ public func checkImageGenerationPhase(config: Flux2QuantizationConfig) -> MemoryCheckResult {
65
+ let required = config.imageGenerationPhaseMemoryGB
66
+ let available = estimatedAvailableMemoryGB
67
+
68
+ if required > available {
69
+ return .insufficientMemory(
70
+ required: required,
71
+ available: available,
72
+ suggestion: "Use qint8 transformer quantization or reduce image size"
73
+ )
74
+ }
75
+
76
+ return .ok
77
+ }
78
+
79
+ // MARK: - Memory Cleanup
80
+
81
+ /// Clear GPU cache
82
+ /// Call this between phases or periodically during generation
83
+ public func clearCache() {
84
+ MLX.Memory.clearCache()
85
+ Flux2Debug.log("GPU cache cleared")
86
+ }
87
+
88
+ /// Suggest garbage collection
89
+ /// Note: Swift's ARC handles this automatically, but we can hint
90
+ public func suggestCleanup() {
91
+ // Force autoreleasepool drain on next opportunity
92
+ autoreleasepool { }
93
+ }
94
+
95
+ /// Full memory cleanup (between phases)
96
+ public func fullCleanup() {
97
+ clearCache()
98
+ suggestCleanup()
99
+
100
+ Flux2Debug.log("Full memory cleanup performed")
101
+ }
102
+
103
+ // MARK: - Memory Monitoring
104
+
105
+ /// Get current memory usage summary
106
+ public func memorySummary() -> String {
107
+ var info = mach_task_basic_info()
108
+ var count = mach_msg_type_number_t(MemoryLayout<mach_task_basic_info>.size) / 4
109
+
110
+ let result = withUnsafeMutablePointer(to: &info) {
111
+ $0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) {
112
+ task_info(mach_task_self_, task_flavor_t(MACH_TASK_BASIC_INFO), $0, &count)
113
+ }
114
+ }
115
+
116
+ let metalActiveMB = MLX.Memory.activeMemory / 1_048_576
117
+ let metalPeakMB = MLX.Memory.peakMemory / 1_048_576
118
+ let metalCacheMB = MLX.Memory.cacheMemory / 1_048_576
119
+
120
+ if result == KERN_SUCCESS {
121
+ let usedMB = info.resident_size / 1_048_576
122
+
123
+ return """
124
+ Memory Usage:
125
+ Metal GPU — active: \(metalActiveMB) MB, peak: \(metalPeakMB) MB, cache: \(metalCacheMB) MB
126
+ Process Resident: \(usedMB) MB | System RAM: \(physicalMemoryGB) GB
127
+ """
128
+ }
129
+
130
+ return """
131
+ Memory Usage:
132
+ Metal GPU — active: \(metalActiveMB) MB, peak: \(metalPeakMB) MB, cache: \(metalCacheMB) MB
133
+ Process info unavailable | System RAM: \(physicalMemoryGB) GB
134
+ """
135
+ }
136
+
137
+ /// Log current memory state
138
+ public func logMemoryState() {
139
+ Flux2Debug.log(memorySummary())
140
+ }
141
+ }
142
+
143
+ // MARK: - Memory Check Result
144
+
145
+ public enum MemoryCheckResult {
146
+ case ok
147
+ case insufficientMemory(required: Int, available: Int, suggestion: String)
148
+ case warning(message: String)
149
+
150
+ public var isOk: Bool {
151
+ if case .ok = self { return true }
152
+ return false
153
+ }
154
+
155
+ public var message: String {
156
+ switch self {
157
+ case .ok:
158
+ return "Memory check passed"
159
+ case .insufficientMemory(let required, let available, let suggestion):
160
+ return "Insufficient memory: need ~\(required)GB, have ~\(available)GB. \(suggestion)"
161
+ case .warning(let message):
162
+ return "Warning: \(message)"
163
+ }
164
+ }
165
+ }
166
+
167
+ // MARK: - Memory-aware Generation
168
+
169
+ extension Flux2MemoryManager {
170
+
171
+ /// Determine optimal batch size for given image dimensions
172
+ public func optimalBatchSize(
173
+ width: Int,
174
+ height: Int,
175
+ config: Flux2QuantizationConfig
176
+ ) -> Int {
177
+ // For now, always return 1 (single image generation)
178
+ // Batch generation requires more sophisticated memory planning
179
+ return 1
180
+ }
181
+
182
+ /// Check if image dimensions are feasible
183
+ public func checkImageSize(width: Int, height: Int) -> MemoryCheckResult {
184
+ let pixels = width * height
185
+
186
+ // Very large images need more working memory
187
+ if pixels > 2048 * 2048 {
188
+ return .warning(message: "Large image size may cause memory pressure")
189
+ }
190
+
191
+ if pixels > 4096 * 4096 {
192
+ return .insufficientMemory(
193
+ required: 100,
194
+ available: estimatedAvailableMemoryGB,
195
+ suggestion: "Reduce image size to 2048x2048 or smaller"
196
+ )
197
+ }
198
+
199
+ return .ok
200
+ }
201
+ }
@@ -0,0 +1,476 @@
1
+ // AutoencoderKL.swift - Flux.2 VAE (AutoencoderKL)
2
+ // Copyright 2025 Vincent Gourbin
3
+
4
+ import Foundation
5
+ import CoreGraphics
6
+ import MLX
7
+ import MLXNN
8
+
9
+ /// Configuration for spatial tiling during VAE decoding
10
+ /// Used to decode large images in tiles to reduce peak memory usage
11
+ public struct VAETilingConfig: Sendable {
12
+ /// Tile size in latent space (e.g., 64 = 512px after 8x upscale)
13
+ public var tileSize: Int
14
+
15
+ /// Overlap between tiles in latent space for blending (e.g., 8 = 64px)
16
+ public var tileOverlap: Int
17
+
18
+ /// Minimum image dimension (in latent space) to trigger tiling
19
+ /// Images smaller than this are decoded without tiling
20
+ public var minTileThreshold: Int
21
+
22
+ public init(tileSize: Int = 64, tileOverlap: Int = 8, minTileThreshold: Int = 128) {
23
+ self.tileSize = tileSize
24
+ self.tileOverlap = tileOverlap
25
+ self.minTileThreshold = minTileThreshold
26
+ }
27
+
28
+ /// Default tiling config (tiles of 512px with 64px overlap)
29
+ public static let `default` = VAETilingConfig()
30
+
31
+ /// Aggressive tiling for very large images (tiles of 256px)
32
+ public static let aggressive = VAETilingConfig(tileSize: 32, tileOverlap: 4, minTileThreshold: 64)
33
+
34
+ /// Disabled - no tiling
35
+ public static let disabled = VAETilingConfig(tileSize: 9999, tileOverlap: 0, minTileThreshold: 9999)
36
+ }
37
+
38
+ /// Variational Autoencoder for Flux.2
39
+ ///
40
+ /// Key differences from Flux.1 VAE:
41
+ /// - 32 latent channels (vs 4)
42
+ /// - Uses BatchNorm2d for latent normalization
43
+ /// - Patch size (2, 2) for latent packing
44
+ /// - Supports spatial tiling for large image decoding
45
+ public class AutoencoderKLFlux2: Module, @unchecked Sendable {
46
+ let config: VAEConfig
47
+
48
+ let encoder: VAEEncoder
49
+ let decoder: VAEDecoder
50
+
51
+ // Quantization convolutions (optional)
52
+ let quantConv: Conv2d?
53
+ let postQuantConv: Conv2d?
54
+
55
+ // BatchNorm for latent normalization (Flux.2 specific)
56
+ let latentBatchNorm: BatchNorm2d
57
+
58
+ /// Scaling factor for latent space
59
+ let scalingFactor: Float
60
+
61
+ public init(config: VAEConfig = .flux2Dev) {
62
+ self.config = config
63
+ self.scalingFactor = config.scalingFactor
64
+
65
+ self.encoder = VAEEncoder(config: config)
66
+ self.decoder = VAEDecoder(config: config)
67
+
68
+ // Optional quant convs (some models use them)
69
+ self.quantConv = Conv2d(
70
+ inputChannels: config.latentChannels * 2,
71
+ outputChannels: config.latentChannels * 2,
72
+ kernelSize: 1
73
+ )
74
+ self.postQuantConv = Conv2d(
75
+ inputChannels: config.latentChannels,
76
+ outputChannels: config.latentChannels,
77
+ kernelSize: 1
78
+ )
79
+
80
+ // BatchNorm for latent normalization
81
+ self.latentBatchNorm = BatchNorm2d(numFeatures: config.latentChannels)
82
+ }
83
+
84
+ /// Encode image to latent space
85
+ /// - Parameters:
86
+ /// - x: Input image [B, 3, H, W] normalized to [-1, 1]
87
+ /// - samplePosterior: Whether to sample from posterior or use mean
88
+ /// - Returns: Latent representation [B, 32, H/8, W/8]
89
+ public func encode(_ x: MLXArray, samplePosterior: Bool = true) -> MLXArray {
90
+ var h = encoder(x) // Output is NCHW
91
+
92
+ // Apply quant conv if present - needs NHWC
93
+ if let qc = quantConv {
94
+ let hNHWC = h.transposed(0, 2, 3, 1) // NCHW -> NHWC
95
+ let qcOut = qc(hNHWC)
96
+ h = qcOut.transposed(0, 3, 1, 2) // NHWC -> NCHW
97
+ }
98
+
99
+ // Sample from posterior (NCHW format)
100
+ let mean = h[0..., 0..<config.latentChannels, 0..., 0...]
101
+ let logVar = h[0..., config.latentChannels..., 0..., 0...]
102
+
103
+ var latent: MLXArray
104
+ if samplePosterior {
105
+ let std = exp(0.5 * logVar)
106
+ let noise = MLXRandom.normal(mean.shape)
107
+ latent = mean + std * noise
108
+ } else {
109
+ latent = mean
110
+ }
111
+
112
+ // NOTE: BatchNorm is NOT applied during encode for Flux.2
113
+ // The BatchNorm weights are for patchified format (128 channels)
114
+ // and encode produces unpatchified latents (32 channels).
115
+ // Normalization is handled later in the pipeline after patchifying.
116
+
117
+ // NOTE: Flux.2 VAE does NOT apply scaling factor during encode!
118
+ // (Standard SD VAE uses scalingFactor=0.18215, but Flux.2 does not)
119
+ // See: diffusers AutoencoderKL.encode() for Flux.2
120
+
121
+ return latent
122
+ }
123
+
124
+ /// Decode latent to image
125
+ /// - Parameter z: Latent representation [B, 32, H/8, W/8] (NCHW format)
126
+ /// - Returns: Decoded image [B, 3, H, W] in [-1, 1] (NCHW format)
127
+ /// - Note: Flux.2 VAE does NOT use scaling factor (unlike SD VAE)
128
+ public func decode(_ z: MLXArray) -> MLXArray {
129
+ // NOTE: Flux.2 VAE does NOT apply scaling factor!
130
+ // (Standard SD VAE uses scalingFactor=0.18215, but Flux.2 does not)
131
+ var latent = z
132
+
133
+ // Apply post-quant conv if present - needs NHWC
134
+ if let pqc = postQuantConv {
135
+ let latentNHWC = latent.transposed(0, 2, 3, 1) // NCHW -> NHWC
136
+ let pqcOut = pqc(latentNHWC)
137
+ latent = pqcOut.transposed(0, 3, 1, 2) // NHWC -> NCHW
138
+ }
139
+
140
+ // Decode - decoder handles NCHW->NHWC->NCHW internally
141
+ return decoder(latent)
142
+ }
143
+
144
+ /// Full forward pass (for training/debugging)
145
+ public func callAsFunction(_ x: MLXArray) -> MLXArray {
146
+ let z = encode(x)
147
+ return decode(z)
148
+ }
149
+
150
+ // MARK: - Tiled Decoding
151
+
152
+ /// Decode latent to image with optional spatial tiling
153
+ /// Use this for large images to reduce peak memory usage
154
+ /// - Parameters:
155
+ /// - z: Latent representation [B, 32, H/8, W/8] (NCHW format)
156
+ /// - tiling: Tiling configuration (nil = auto-detect based on size)
157
+ /// - Returns: Decoded image [B, 3, H, W] in [-1, 1] (NCHW format)
158
+ public func decodeWithTiling(_ z: MLXArray, tiling: VAETilingConfig? = nil) -> MLXArray {
159
+ let config = tiling ?? .default
160
+ let H = z.shape[2]
161
+ let W = z.shape[3]
162
+
163
+ // Check if tiling is needed
164
+ if H <= config.minTileThreshold && W <= config.minTileThreshold {
165
+ Flux2Debug.verbose("VAE: Image small enough, decoding without tiling")
166
+ return decode(z)
167
+ }
168
+
169
+ Flux2Debug.log("VAE: Using tiled decoding for \(H*8)x\(W*8) image")
170
+ return decodeTiled(z, tileSize: config.tileSize, overlap: config.tileOverlap)
171
+ }
172
+
173
+ /// Tiled decoding implementation with overlap cropping
174
+ /// Decodes large latents in tiles and concatenates results
175
+ private func decodeTiled(_ z: MLXArray, tileSize: Int, overlap: Int) -> MLXArray {
176
+ let H = z.shape[2]
177
+ let W = z.shape[3]
178
+
179
+ let outOverlap = overlap * 8
180
+ let stride = tileSize - overlap
181
+
182
+ // Calculate tiles
183
+ let numTilesH = max(1, Int(ceil(Float(H - overlap) / Float(stride))))
184
+ let numTilesW = max(1, Int(ceil(Float(W - overlap) / Float(stride))))
185
+
186
+ Flux2Debug.verbose("VAE tiling: \(numTilesH)x\(numTilesW) tiles, size=\(tileSize), overlap=\(overlap)")
187
+
188
+ // Decode all tiles
189
+ var tiles: [[MLXArray]] = []
190
+ var tileHeights: [[Int]] = []
191
+ var tileWidths: [[Int]] = []
192
+
193
+ for tileY in 0..<numTilesH {
194
+ var row: [MLXArray] = []
195
+ var heights: [Int] = []
196
+ var widths: [Int] = []
197
+
198
+ for tileX in 0..<numTilesW {
199
+ let y0 = min(tileY * stride, max(0, H - tileSize))
200
+ let x0 = min(tileX * stride, max(0, W - tileSize))
201
+ let y1 = min(y0 + tileSize, H)
202
+ let x1 = min(x0 + tileSize, W)
203
+
204
+ let tile = z[0..., 0..., y0..<y1, x0..<x1]
205
+ let decoded = decode(tile)
206
+ eval(decoded)
207
+
208
+ row.append(decoded)
209
+ heights.append((y1 - y0) * 8)
210
+ widths.append((x1 - x0) * 8)
211
+
212
+ // Clear cache after each tile to manage memory
213
+ MLX.Memory.clearCache()
214
+
215
+ Flux2Debug.verbose("VAE tile [\(tileY),\(tileX)] decoded")
216
+ }
217
+ tiles.append(row)
218
+ tileHeights.append(heights)
219
+ tileWidths.append(widths)
220
+ }
221
+
222
+ Flux2Debug.log("VAE: Blending \(numTilesH * numTilesW) tiles")
223
+
224
+ // Reconstruct by cropping overlaps and concatenating
225
+ var rows: [MLXArray] = []
226
+ for tileY in 0..<numTilesH {
227
+ var rowTiles: [MLXArray] = []
228
+ for tileX in 0..<numTilesW {
229
+ let tile = tiles[tileY][tileX]
230
+ let h = tileHeights[tileY][tileX]
231
+ let w = tileWidths[tileY][tileX]
232
+
233
+ // Crop overlap regions (take center half of overlap from each side)
234
+ let cropTop = (tileY > 0) ? outOverlap / 2 : 0
235
+ let cropLeft = (tileX > 0) ? outOverlap / 2 : 0
236
+ let cropBottom = (tileY < numTilesH - 1) ? outOverlap / 2 : 0
237
+ let cropRight = (tileX < numTilesW - 1) ? outOverlap / 2 : 0
238
+
239
+ let cropped = tile[0..., 0..., cropTop..<(h - cropBottom), cropLeft..<(w - cropRight)]
240
+ rowTiles.append(cropped)
241
+ }
242
+ // Concatenate horizontally
243
+ let row = concatenated(rowTiles, axis: 3)
244
+ rows.append(row)
245
+ }
246
+ // Concatenate vertically
247
+ let result = concatenated(rows, axis: 2)
248
+
249
+ Flux2Debug.log("VAE: Tiled decoding completed, output shape: \(result.shape)")
250
+ return result
251
+ }
252
+
253
+ /// Create a blending weight mask for tile overlap
254
+ private func createBlendMask(size: Int, overlap: Int) -> MLXArray {
255
+ // Create linear ramp weights for blending
256
+ var weights = [Float](repeating: 1.0, count: size * size)
257
+
258
+ // Apply linear falloff in overlap regions
259
+ for y in 0..<size {
260
+ for x in 0..<size {
261
+ var weight: Float = 1.0
262
+
263
+ // Top edge
264
+ if y < overlap {
265
+ weight *= Float(y) / Float(overlap)
266
+ }
267
+ // Bottom edge
268
+ if y >= size - overlap {
269
+ weight *= Float(size - 1 - y) / Float(overlap)
270
+ }
271
+ // Left edge
272
+ if x < overlap {
273
+ weight *= Float(x) / Float(overlap)
274
+ }
275
+ // Right edge
276
+ if x >= size - overlap {
277
+ weight *= Float(size - 1 - x) / Float(overlap)
278
+ }
279
+
280
+ weights[y * size + x] = weight
281
+ }
282
+ }
283
+
284
+ return MLXArray(weights).reshaped([1, 1, size, size])
285
+ }
286
+ }
287
+
288
+ // MARK: - Latent Packing/Unpacking
289
+
290
+ extension AutoencoderKLFlux2 {
291
+ /// Pack latents for transformer input
292
+ /// Converts [B, C, H, W] to [B, H*W, C] with patch handling
293
+ /// - Parameters:
294
+ /// - latents: Encoded latents [B, 32, H, W]
295
+ /// - patchSize: Patch size (2, 2) for Flux.2
296
+ /// - Returns: Packed latents [B, (H/p)*(W/p), C*p*p]
297
+ public func packLatents(_ latents: MLXArray, patchSize: (Int, Int) = (2, 2)) -> MLXArray {
298
+ let shape = latents.shape
299
+ let B = shape[0]
300
+ let C = shape[1]
301
+ let H = shape[2]
302
+ let W = shape[3]
303
+
304
+ let pH = patchSize.0
305
+ let pW = patchSize.1
306
+
307
+ // Reshape to patches
308
+ // [B, C, H, W] -> [B, C, H/pH, pH, W/pW, pW]
309
+ var packed = latents.reshaped([B, C, H / pH, pH, W / pW, pW])
310
+
311
+ // Permute to [B, H/pH, W/pW, C, pH, pW]
312
+ packed = packed.transposed(0, 2, 4, 1, 3, 5)
313
+
314
+ // Flatten patches and spatial dims
315
+ // [B, H/pH, W/pW, C, pH, pW] -> [B, (H/pH)*(W/pW), C*pH*pW]
316
+ let numPatches = (H / pH) * (W / pW)
317
+ let patchDim = C * pH * pW
318
+
319
+ packed = packed.reshaped([B, numPatches, patchDim])
320
+
321
+ return packed
322
+ }
323
+
324
+ /// Unpack latents from transformer output
325
+ /// Converts [B, H*W, C] back to [B, C, H, W]
326
+ /// - Parameters:
327
+ /// - packed: Packed latents [B, (H/p)*(W/p), C*p*p]
328
+ /// - height: Target height (after unpacking)
329
+ /// - width: Target width (after unpacking)
330
+ /// - patchSize: Patch size (2, 2) for Flux.2
331
+ /// - Returns: Unpacked latents [B, C, H, W]
332
+ public func unpackLatents(
333
+ _ packed: MLXArray,
334
+ height: Int,
335
+ width: Int,
336
+ patchSize: (Int, Int) = (2, 2)
337
+ ) -> MLXArray {
338
+ let shape = packed.shape
339
+ let B = shape[0]
340
+
341
+ let pH = patchSize.0
342
+ let pW = patchSize.1
343
+
344
+ let outH = height / 8 // VAE downsamples by 8
345
+ let outW = width / 8
346
+
347
+ let numPatchesH = outH / pH
348
+ let numPatchesW = outW / pW
349
+ let C = config.latentChannels
350
+
351
+ // [B, numPatches, C*pH*pW] -> [B, numPatchesH, numPatchesW, C, pH, pW]
352
+ var unpacked = packed.reshaped([B, numPatchesH, numPatchesW, C, pH, pW])
353
+
354
+ // Permute to [B, C, numPatchesH, pH, numPatchesW, pW]
355
+ unpacked = unpacked.transposed(0, 3, 1, 4, 2, 5)
356
+
357
+ // Reshape to final [B, C, H, W]
358
+ unpacked = unpacked.reshaped([B, C, outH, outW])
359
+
360
+ return unpacked
361
+ }
362
+ }
363
+
364
+ // MARK: - Weight Loading
365
+
366
+ extension AutoencoderKLFlux2 {
367
+ /// Load weights from safetensors
368
+ public func loadWeights(from url: URL) throws {
369
+ // Implementation in WeightLoader.swift
370
+ fatalError("Weight loading not yet implemented")
371
+ }
372
+
373
+ /// Load BatchNorm running statistics
374
+ public func loadBatchNormStats(runningMean: MLXArray, runningVar: MLXArray) {
375
+ latentBatchNorm.runningMean = runningMean
376
+ latentBatchNorm.runningVar = runningVar
377
+ }
378
+
379
+ /// Get BatchNorm running mean for latent normalization
380
+ public var batchNormRunningMean: MLXArray {
381
+ latentBatchNorm.runningMean
382
+ }
383
+
384
+ /// Get BatchNorm running variance for latent normalization
385
+ public var batchNormRunningVar: MLXArray {
386
+ latentBatchNorm.runningVar
387
+ }
388
+ }
389
+
390
+ // MARK: - Image Processing Utilities
391
+
392
+ /// Process image for VAE input
393
+ /// - Parameter image: CGImage to process
394
+ /// - Returns: MLXArray normalized to [-1, 1], shape [1, 3, H, W]
395
+ public func preprocessImageForVAE(_ image: CGImage) -> MLXArray {
396
+ let width = image.width
397
+ let height = image.height
398
+
399
+ // Get pixel data
400
+ guard let colorSpace = CGColorSpace(name: CGColorSpace.sRGB),
401
+ let context = CGContext(
402
+ data: nil,
403
+ width: width,
404
+ height: height,
405
+ bitsPerComponent: 8,
406
+ bytesPerRow: width * 4,
407
+ space: colorSpace,
408
+ bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue
409
+ ) else {
410
+ fatalError("Failed to create graphics context")
411
+ }
412
+
413
+ context.draw(image, in: CGRect(x: 0, y: 0, width: width, height: height))
414
+
415
+ guard let data = context.data else {
416
+ fatalError("Failed to get pixel data")
417
+ }
418
+
419
+ let pixelData = data.assumingMemoryBound(to: UInt8.self)
420
+ var floatData: [Float] = []
421
+
422
+ // Convert RGBA to RGB and normalize to [-1, 1]
423
+ for y in 0..<height {
424
+ for x in 0..<width {
425
+ let offset = (y * width + x) * 4
426
+ floatData.append(Float(pixelData[offset]) / 127.5 - 1.0) // R
427
+ floatData.append(Float(pixelData[offset + 1]) / 127.5 - 1.0) // G
428
+ floatData.append(Float(pixelData[offset + 2]) / 127.5 - 1.0) // B
429
+ }
430
+ }
431
+
432
+ // Create MLXArray [H, W, 3] and permute to [1, 3, H, W]
433
+ let array = MLXArray(floatData).reshaped([height, width, 3])
434
+ return array.transposed(2, 0, 1).expandedDimensions(axis: 0)
435
+ }
436
+
437
+ /// Convert VAE output to CGImage
438
+ /// - Parameter tensor: MLXArray in [-1, 1], shape [1, 3, H, W]
439
+ /// - Returns: CGImage
440
+ public func postprocessVAEOutput(_ tensor: MLXArray) -> CGImage? {
441
+ // Remove batch dim and permute to [H, W, 3]
442
+ let image = tensor.squeezed(axis: 0).transposed(1, 2, 0)
443
+
444
+ // Clamp and convert to [0, 255]
445
+ let clamped = clip(image, min: -1, max: 1)
446
+ let normalized = ((clamped + 1) * 127.5).asType(.uint8)
447
+
448
+ // Get dimensions
449
+ let shape = normalized.shape
450
+ let height = shape[0]
451
+ let width = shape[1]
452
+
453
+ // Evaluate and get data
454
+ eval(normalized)
455
+ let data = normalized.asArray(UInt8.self)
456
+
457
+ // Create CGImage
458
+ guard let colorSpace = CGColorSpace(name: CGColorSpace.sRGB),
459
+ let provider = CGDataProvider(data: Data(data) as CFData) else {
460
+ return nil
461
+ }
462
+
463
+ return CGImage(
464
+ width: width,
465
+ height: height,
466
+ bitsPerComponent: 8,
467
+ bitsPerPixel: 24,
468
+ bytesPerRow: width * 3,
469
+ space: colorSpace,
470
+ bitmapInfo: CGBitmapInfo(rawValue: CGImageAlphaInfo.none.rawValue),
471
+ provider: provider,
472
+ decode: nil,
473
+ shouldInterpolate: true,
474
+ intent: .defaultIntent
475
+ )
476
+ }