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.
- package/dist/cli.cjs +23 -4
- package/dist/cli.d.cts +2 -0
- package/dist/cli.d.ts +2 -0
- package/dist/cli.js +23 -4
- package/dist/{ui-Ce-pnUUA.d.ts → ui-Bg11tvlc.d.ts} +7 -7
- package/dist/ui.d.ts +1 -1
- package/dist/window.d.ts +1 -1
- package/package.json +4 -1
- package/templates/image-generation/app/index.tsx +1 -1
- package/templates/native/.gitattributes +1 -0
- package/templates/native/App.tsx +26 -0
- package/templates/native/app.json +12 -0
- package/templates/native/index.ts +8 -0
- package/templates/native/ios/.xcode.env +11 -0
- package/templates/native/ios/Bleam/AI/Flux2/Configuration/Flux2Config.swift +288 -0
- package/templates/native/ios/Bleam/AI/Flux2/Configuration/MemoryConfig.swift +233 -0
- package/templates/native/ios/Bleam/AI/Flux2/Configuration/MemoryOptimizationConfig.swift +157 -0
- package/templates/native/ios/Bleam/AI/Flux2/Configuration/ModelRegistry.swift +425 -0
- package/templates/native/ios/Bleam/AI/Flux2/Configuration/QuantizationConfig.swift +138 -0
- package/templates/native/ios/Bleam/AI/Flux2/Configuration/TransformerRepoProfile.swift +59 -0
- package/templates/native/ios/Bleam/AI/Flux2/Configuration/VAEConfig.swift +134 -0
- package/templates/native/ios/Bleam/AI/Flux2/Flux2Core.swift +55 -0
- package/templates/native/ios/Bleam/AI/Flux2/Loading/WeightLoader.swift +1192 -0
- package/templates/native/ios/Bleam/AI/Flux2/Pipeline/FluxImageRunner.swift +1282 -0
- package/templates/native/ios/Bleam/AI/Flux2/Pipeline/LatentUtils.swift +191 -0
- package/templates/native/ios/Bleam/AI/Flux2/Scheduler/FlowMatchEulerScheduler.swift +279 -0
- package/templates/native/ios/Bleam/AI/Flux2/Text/KleinTextEncoder.swift +380 -0
- package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2Attention.swift +379 -0
- package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2Embeddings.swift +159 -0
- package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2FeedForward.swift +154 -0
- package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2FusedKernels.swift +198 -0
- package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2KVCache.swift +36 -0
- package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2Modulation.swift +193 -0
- package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2ParallelAttention.swift +389 -0
- package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2RoPE.swift +304 -0
- package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2SingleBlock.swift +290 -0
- package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2Transformer.swift +530 -0
- package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2TransformerBlock.swift +362 -0
- package/templates/native/ios/Bleam/AI/Flux2/Utils/Flux2Debug.swift +101 -0
- package/templates/native/ios/Bleam/AI/Flux2/Utils/MLXCheckpoint.swift +118 -0
- package/templates/native/ios/Bleam/AI/Flux2/Utils/MemoryManager.swift +201 -0
- package/templates/native/ios/Bleam/AI/Flux2/VAE/AutoencoderKL.swift +476 -0
- package/templates/native/ios/Bleam/AI/Flux2/VAE/ResnetBlock.swift +316 -0
- package/templates/native/ios/Bleam/AI/Flux2/VAE/VAEDecoder.swift +120 -0
- package/templates/native/ios/Bleam/AI/Flux2/VAE/VAEEncoder.swift +136 -0
- package/templates/native/ios/Bleam/AppDelegate.swift +373 -0
- package/templates/native/ios/Bleam/Appearance.swift +62 -0
- package/templates/native/ios/Bleam/Bridging-Header.h +3 -0
- package/templates/native/ios/Bleam/Images.xcassets/AppIcon.appiconset/Contents.json +13 -0
- package/templates/native/ios/Bleam/Images.xcassets/Contents.json +6 -0
- package/templates/native/ios/Bleam/Info.plist +76 -0
- package/templates/native/ios/Bleam/PrivacyInfo.xcprivacy +48 -0
- package/templates/native/ios/Bleam/SceneDelegate.swift +58 -0
- package/templates/native/ios/Bleam/SplashScreen.storyboard +47 -0
- package/templates/native/ios/Bleam/Supporting/Expo.plist +6 -0
- package/templates/native/ios/Bleam/bleam.entitlements +10 -0
- package/templates/native/ios/Bleam.xcodeproj/project.pbxproj +1146 -0
- package/templates/native/ios/Bleam.xcodeproj/xcshareddata/xcschemes/Bleam.xcscheme +88 -0
- package/templates/native/ios/GenerationService/GenerationService.entitlements +10 -0
- package/templates/native/ios/GenerationService/ImageGenerationRunner.swift +591 -0
- package/templates/native/ios/GenerationService/Info.plist +31 -0
- package/templates/native/ios/GenerationService/main.swift +165 -0
- package/templates/native/ios/PlatformHelper/Info.plist +29 -0
- package/templates/native/ios/PlatformHelper/main.swift +335 -0
- package/templates/native/ios/Podfile +198 -0
- package/templates/native/ios/Podfile.lock +2284 -0
- package/templates/native/ios/Podfile.properties.json +5 -0
- package/templates/native/ios/Shared/Generation/GenerationServiceProtocol.swift +12 -0
- package/templates/native/ios/Shared/Generation/GenerationWorkerProtocol.swift +124 -0
- package/templates/native/metro.config.js +6 -0
- package/templates/native/modules/bleam-runtime/BleamRuntime.podspec +14 -0
- package/templates/native/modules/bleam-runtime/expo-module.config.json +8 -0
- package/templates/native/modules/bleam-runtime/ios/AIModule.swift +1258 -0
- package/templates/native/modules/bleam-runtime/ios/GenerationContracts.swift +135 -0
- package/templates/native/modules/bleam-runtime/ios/PlatformModule.swift +313 -0
- package/templates/native/modules/bleam-runtime/package.json +5 -0
- package/templates/native/package.json +39 -0
- package/templates/native/scripts/build/mlx-frameworks.ts +421 -0
- package/templates/native/scripts/start/index.ts +474 -0
- package/templates/native/scripts/start/known-failures.ts +20 -0
- package/templates/native/scripts/start/loader.ts +76 -0
- package/templates/native/scripts/start/native-state.ts +242 -0
- package/templates/native/scripts/start/package.json +4 -0
- package/templates/native/scripts/start/swift-packages.ts +414 -0
- package/templates/native/scripts/start/xcode-formatter.ts +106 -0
- package/templates/native/tsconfig.json +14 -0
- 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
|
+
}
|