bleam 0.0.8 → 0.0.9

package/templates/native/ios/Bleam/AI/Flux2/Loading/WeightLoader.swift

@@ -0,0 +1,1192 @@
+// WeightLoader.swift - Load model weights from safetensors
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+import MLX
+import MLXNN
+
+/// Utilities for loading Flux.2 model weights from safetensors files
+public class Flux2WeightLoader {
+
+    /// Load all weights from a model directory
+    /// - Parameter modelPath: Path to directory containing safetensors files
+    /// - Returns: Dictionary of weight name to MLXArray
+    public static func loadWeights(from modelPath: String) throws -> [String: MLXArray] {
+        let fm = FileManager.default
+        let contents = try fm.contentsOfDirectory(atPath: modelPath)
+        let safetensorFiles = contents.filter { $0.hasSuffix(".safetensors") }.sorted()
+
+        if safetensorFiles.isEmpty {
+            throw Flux2WeightLoaderError.noWeightsFound(modelPath)
+        }
+
+        Flux2Debug.log("Found \(safetensorFiles.count) safetensor files in \(modelPath)")
+
+        var allWeights: [String: MLXArray] = [:]
+
+        for filename in safetensorFiles {
+            let filePath = "\(modelPath)/\(filename)"
+            let weights = try loadArrays(url: URL(fileURLWithPath: filePath))
+
+            for (key, value) in weights {
+                let normalizedKey = normalizeUREVTransformerKey(key, filename: filename, modelPath: modelPath)
+                allWeights[normalizedKey] = value
+            }
+
+            Flux2Debug.log("Loaded \(weights.count) tensors from \(filename)")
+        }
+
+        return allWeights
+    }
+
+    private static func normalizeUREVTransformerKey(_ key: String, filename: String, modelPath: String) -> String {
+        let isUREVCheckpoint = filename == "urevF2K4BFp16.safetensors"
+            && modelPath.split(separator: "/").contains("urev")
+        let prefix = "model.diffusion_model."
+        guard isUREVCheckpoint, key.hasPrefix(prefix) else { return key }
+        return String(key.dropFirst(prefix.count))
+    }
+
+    /// Load weights from URL
+    public static func loadWeights(from url: URL) throws -> [String: MLXArray] {
+        try loadWeights(from: url.path)
+    }
+
+    // MARK: - Transformer Weight Mapping
+
+    /// Detect if weights are in BFL (Black Forest Labs) native format
+    private static func isBFLFormat(_ weights: [String: MLXArray]) -> Bool {
+        // BFL format uses "double_blocks" and "single_blocks" directly
+        // Diffusers format uses "transformer_blocks" and "single_transformer_blocks"
+        return weights.keys.contains { $0.hasPrefix("double_blocks.") || $0.hasPrefix("single_blocks.") }
+    }
+
+    /// Convert HuggingFace transformer weight keys to Swift module paths
+    /// Also handles dequantization of quanto qint8 weights
+    /// Supports both Diffusers format and BFL native format
+    /// Note: For Diffusers quanto format, entries are removed from `weights` as they are mapped
+    /// to avoid holding both raw quanto (~32GB) and dequantified float16 (~60GB) simultaneously.
+    static func mapTransformerWeights(
+        _ weights: inout [String: MLXArray],
+        profile: Flux2TransformerRepoProfile
+    ) -> [String: MLXArray] {
+        if profile.checkpointFormat == .bfl || isBFLFormat(weights) {
+            Flux2Debug.log("Detected BFL native format weights")
+            return mapBFLTransformerWeights(&weights)
+        }
+
+        Flux2Debug.log("Detected Diffusers format weights for \(profile.repoId)")
+        return mapDiffusersTransformerWeights(&weights, profile: profile)
+    }
+
+    /// Detect whether the raw transformer checkpoint is already stored in a packed
+    /// quantized format that should be dequantized during load rather than requantized
+    /// again at runtime.
+    static func isNativePrequantizedTransformerCheckpoint(_ weights: [String: MLXArray]) -> Bool {
+        let hasPackedWeights = weights.contains { key, value in
+            key.hasSuffix(".weight") && value.dtype == .uint32
+        }
+        let hasQuantizationSidecars = weights.keys.contains { $0.hasSuffix(".scales") || $0.hasSuffix(".biases") }
+
+        return hasPackedWeights && hasQuantizationSidecars
+    }
+
+    static func isNativePrequantizedTransformerCheckpoint(
+        _ weights: [String: MLXArray],
+        profile: Flux2TransformerRepoProfile
+    ) -> Bool {
+        switch profile.checkpointFormat {
+        case .packedDiffusers, .bfl:
+            return isNativePrequantizedTransformerCheckpoint(weights)
+        case .diffusers, .diffusersQuanto:
+            return false
+        }
+    }
+
+    /// Map weights in BFL (Black Forest Labs) native format
+    /// BFL format uses fused QKV and different naming conventions
+    /// Entries are removed from the input dict as they are processed to minimize peak memory,
+    /// avoiding holding both raw bf16 weights and mapped/converted weights simultaneously.
+    private static func mapBFLTransformerWeights(_ weights: inout [String: MLXArray]) -> [String: MLXArray] {
+        var mapped: [String: MLXArray] = [:]
+
+        // Debug: count double_blocks keys
+        let doubleBlockKeys = weights.keys.filter { $0.hasPrefix("double_blocks.") }
+        Flux2Debug.log("BFL format: found \(doubleBlockKeys.count) double_blocks keys")
+        for key in doubleBlockKeys.prefix(5) {
+            Flux2Debug.verbose("  double_block key: \(key)")
+        }
+
+        // Iterate over snapshot of keys; remove each entry after processing to free memory
+        let allKeys = Array(weights.keys)
+        for key in allKeys {
+            guard let value = weights.removeValue(forKey: key) else { continue }
+            // BFL format uses fused QKV weights that need to be split
+            // double_blocks.{i}.img_attn.qkv.weight -> split into toQ, toK, toV
+            // double_blocks.{i}.txt_attn.qkv.weight -> split into addQProj, addKProj, addVProj
+
+            if let mappedKey = mapCurrentBFLTransformerKey(key) {
+                mapped[mappedKey] = value
+                continue
+            }
+
+            if key.contains(".img_attn.qkv.weight") {
+                Flux2Debug.verbose("Splitting img_attn QKV: \(key) shape=\(value.shape)")
+                // Extract block index
+                let blockIdx = extractBlockIndex(from: key, prefix: "double_blocks.")
+                // Split fused QKV into Q, K, V (shape: [3*dim, dim] -> 3x [dim, dim])
+                let dim = value.shape[0] / 3
+                let q = value[0..<dim, 0...]
+                let k = value[dim..<(2*dim), 0...]
+                let v = value[(2*dim)..., 0...]
+                mapped["transformerBlocks.\(blockIdx).attn.toQ.weight"] = q
+                mapped["transformerBlocks.\(blockIdx).attn.toK.weight"] = k
+                mapped["transformerBlocks.\(blockIdx).attn.toV.weight"] = v
+            } else if key.contains(".txt_attn.qkv.weight") {
+                let blockIdx = extractBlockIndex(from: key, prefix: "double_blocks.")
+                let dim = value.shape[0] / 3
+                let q = value[0..<dim, 0...]
+                let k = value[dim..<(2*dim), 0...]
+                let v = value[(2*dim)..., 0...]
+                mapped["transformerBlocks.\(blockIdx).attn.addQProj.weight"] = q
+                mapped["transformerBlocks.\(blockIdx).attn.addKProj.weight"] = k
+                mapped["transformerBlocks.\(blockIdx).attn.addVProj.weight"] = v
+            } else if key.contains(".img_attn.proj.weight") {
+                let blockIdx = extractBlockIndex(from: key, prefix: "double_blocks.")
+                mapped["transformerBlocks.\(blockIdx).attn.toOut.weight"] = value
+            } else if key.contains(".txt_attn.proj.weight") {
+                let blockIdx = extractBlockIndex(from: key, prefix: "double_blocks.")
+                mapped["transformerBlocks.\(blockIdx).attn.toAddOut.weight"] = value
+            } else if key.contains(".img_attn.norm.query_norm.scale") {
+                let blockIdx = extractBlockIndex(from: key, prefix: "double_blocks.")
+                mapped["transformerBlocks.\(blockIdx).attn.normQ.weight"] = value
+            } else if key.contains(".img_attn.norm.key_norm.scale") {
+                let blockIdx = extractBlockIndex(from: key, prefix: "double_blocks.")
+                mapped["transformerBlocks.\(blockIdx).attn.normK.weight"] = value
+            } else if key.contains(".txt_attn.norm.query_norm.scale") {
+                let blockIdx = extractBlockIndex(from: key, prefix: "double_blocks.")
+                mapped["transformerBlocks.\(blockIdx).attn.normAddedQ.weight"] = value
+            } else if key.contains(".txt_attn.norm.key_norm.scale") {
+                let blockIdx = extractBlockIndex(from: key, prefix: "double_blocks.")
+                mapped["transformerBlocks.\(blockIdx).attn.normAddedK.weight"] = value
+            } else if key.contains(".img_mlp.0.weight") {
+                // BFL: img_mlp.0 is the gated linear (produces 2x for SwiGLU)
+                let blockIdx = extractBlockIndex(from: key, prefix: "double_blocks.")
+                mapped["transformerBlocks.\(blockIdx).ff.activation.proj.weight"] = value
+            } else if key.contains(".img_mlp.2.weight") {
+                let blockIdx = extractBlockIndex(from: key, prefix: "double_blocks.")
+                mapped["transformerBlocks.\(blockIdx).ff.linearOut.weight"] = value
+            } else if key.contains(".txt_mlp.0.weight") {
+                let blockIdx = extractBlockIndex(from: key, prefix: "double_blocks.")
+                mapped["transformerBlocks.\(blockIdx).ffContext.activation.proj.weight"] = value
+            } else if key.contains(".txt_mlp.2.weight") {
+                let blockIdx = extractBlockIndex(from: key, prefix: "double_blocks.")
+                mapped["transformerBlocks.\(blockIdx).ffContext.linearOut.weight"] = value
+            } else if key.hasPrefix("single_blocks.") && key.contains(".linear1.weight") {
+                // Single blocks have fused QKV+MLP
+                let blockIdx = extractBlockIndex(from: key, prefix: "single_blocks.")
+                mapped["singleTransformerBlocks.\(blockIdx).attn.toQkvMlp.weight"] = value
+            } else if key.hasPrefix("single_blocks.") && key.contains(".linear2.weight") {
+                let blockIdx = extractBlockIndex(from: key, prefix: "single_blocks.")
+                mapped["singleTransformerBlocks.\(blockIdx).attn.toOut.weight"] = value
+            } else if key.hasPrefix("single_blocks.") && key.contains(".norm.query_norm.scale") {
+                let blockIdx = extractBlockIndex(from: key, prefix: "single_blocks.")
+                mapped["singleTransformerBlocks.\(blockIdx).attn.normQ.weight"] = value
+            } else if key.hasPrefix("single_blocks.") && key.contains(".norm.key_norm.scale") {
+                let blockIdx = extractBlockIndex(from: key, prefix: "single_blocks.")
+                mapped["singleTransformerBlocks.\(blockIdx).attn.normK.weight"] = value
+            } else if key == "img_in.weight" {
+                Flux2Debug.log("BFL: mapping img_in -> xEmbedder: \(value.shape)")
+                mapped["xEmbedder.weight"] = value
+            } else if key == "txt_in.weight" {
+                Flux2Debug.log("BFL: mapping txt_in -> contextEmbedder: \(value.shape)")
+                mapped["contextEmbedder.weight"] = value
+            } else if key == "time_in.in_layer.weight" {
+                Flux2Debug.log("BFL: mapping time_in.in_layer -> timestepEmbedder.linear1: \(value.shape)")
+                mapped["timeGuidanceEmbed.timestepEmbedder.linear1.weight"] = value
+            } else if key == "time_in.out_layer.weight" {
+                Flux2Debug.log("BFL: mapping time_in.out_layer -> timestepEmbedder.linear2: \(value.shape)")
+                mapped["timeGuidanceEmbed.timestepEmbedder.linear2.weight"] = value
+            } else if key == "double_stream_modulation_img.lin.weight" {
+                mapped["doubleStreamModulationImg.linear.weight"] = value
+            } else if key == "double_stream_modulation_txt.lin.weight" {
+                mapped["doubleStreamModulationTxt.linear.weight"] = value
+            } else if key == "single_stream_modulation.lin.weight" {
+                mapped["singleStreamModulation.linear.weight"] = value
+            } else if key == "final_layer.adaLN_modulation.1.weight" {
+                // IMPORTANT: BFL format vs Diffusers format weight layout difference
+                //
+                // The normOut layer (AdaLayerNormContinuous) projects conditioning to scale+shift:
+                //   params = linear(silu(conditioning))  // [B, dim * 2]
+                //   scale = params[0:dim]
+                //   shift = params[dim:]
+                //
+                // BFL format stores the linear weight as [shift_weights | scale_weights]
+                // Diffusers format stores it as [scale_weights | shift_weights]
+                //
+                // Without this swap, bf16 BFL models produce inverted scale/shift values,
+                // causing ~10x higher output magnitude and posterized/noisy images.
+                //
+                // Fix: Swap the two halves of the weight matrix to match diffusers layout.
+                let dim = value.shape[0] / 2  // 3072 for Klein 4B, 6144 for Dev
+                let shiftRows = value[0..<dim, 0...]  // First half (shift in BFL)
+                let scaleRows = value[dim..., 0...]   // Second half (scale in BFL)
+                let swapped = concatenated([scaleRows, shiftRows], axis: 0)
+                Flux2Debug.log("BFL: Swapped normOut.linear.weight halves (dim=\(dim)): \(value.shape) -> \(swapped.shape)")
+                mapped["normOut.linear.weight"] = swapped
+            } else if key == "final_layer.linear.weight" {
+                mapped["projOut.weight"] = value
+            } else {
+                Flux2Debug.verbose("BFL format: unmapped key \(key)")
+            }
+        }
+
+        // Input dict is now empty — raw BFL weight references freed
+        Flux2Debug.log("Input dict cleared — freed raw BFL weight entries")
+        Flux2Debug.log("Mapped \(mapped.count) BFL format weights")
+
+        // Debug: print mapped keys by category
+        let transformerBlockKeys2 = mapped.keys.filter { $0.hasPrefix("transformerBlocks.") }
+        let singleBlockKeys = mapped.keys.filter { $0.hasPrefix("singleTransformerBlocks.") }
+        Flux2Debug.log("  - transformerBlocks keys: \(transformerBlockKeys2.count)")
+        Flux2Debug.log("  - singleTransformerBlocks keys: \(singleBlockKeys.count)")
+
+        Flux2Debug.verbose("transformerBlocks keys:")
+        for key in transformerBlockKeys2.sorted().prefix(10) {
+            Flux2Debug.verbose("  - \(key): \(mapped[key]!.shape)")
+        }
+
+        // DEBUG: Print raw BFL weight statistics BEFORE conversion
+        Flux2Debug.log("=== Raw BFL Weight Statistics (BEFORE conversion) ===")
+        for key in ["xEmbedder.weight", "contextEmbedder.weight",
+                    "transformerBlocks.0.attn.toQ.weight",
+                    "transformerBlocks.0.ff.activation.proj.weight",
+                    "singleTransformerBlocks.0.attn.toQkvMlp.weight"] {
+            if let w = mapped[key] {
+                eval(w)
+                // Convert to float32 for accurate stats
+                let wf32 = w.asType(.float32)
+                let absVal = MLX.abs(wf32)
+                let meanVal = mean(absVal).item(Float.self)
+                let maxVal = MLX.max(absVal).item(Float.self)
+                let minVal = MLX.min(absVal).item(Float.self)
+                let hasNaN = any(isNaN(wf32)).item(Bool.self)
+                let numInf = sum(MLX.abs(wf32) .> 1e30).item(Int.self)  // Count very large values as proxy for inf
+                Flux2Debug.log("  \(key): dtype=\(w.dtype), shape=\(w.shape), mean=\(meanVal), max=\(maxVal), min=\(minVal), hasNaN=\(hasNaN), numInf=\(numInf)")
+            }
+        }
+
+        // Convert bfloat16 → float16 directly (no intermediate f32)
+        // bf16 has wider exponent range (8-bit) than f16 (5-bit), so values outside
+        // f16 range [−65504, 65504] would overflow to inf. MLX handles this correctly
+        // by saturating to ±inf, which is the same behavior as the f32 intermediate path.
+        // Going direct saves ~24GB of temporary f32 allocations for Dev-size models.
+        // Progressive removal from `mapped` avoids holding both dicts simultaneously.
+        var converted: [String: MLXArray] = [:]
+        var convertedCount = 0
+        let mappedKeys = Array(mapped.keys)
+        for key in mappedKeys {
+            guard let value = mapped.removeValue(forKey: key) else { continue }
+            if value.dtype == .bfloat16 {
+                converted[key] = value.asType(.float16)
+                convertedCount += 1
+            } else {
+                converted[key] = value
+            }
+        }
+        Flux2Debug.log("Converted \(convertedCount) bfloat16 weights to float16 (direct)")
+
+        // DEBUG: Print weight statistics for key layers (AFTER processing)
+        Flux2Debug.log("=== BFL Weight Statistics (AFTER processing) ===")
+        for key in ["xEmbedder.weight", "contextEmbedder.weight",
+                    "transformerBlocks.0.attn.toQ.weight",
+                    "transformerBlocks.0.ff.activation.proj.weight",
+                    "singleTransformerBlocks.0.attn.toQkvMlp.weight"] {
+            if let w = converted[key] {
+                eval(w)
+                let wf32 = w.asType(.float32)
+                let absVal = MLX.abs(wf32)
+                let meanVal = mean(absVal).item(Float.self)
+                let maxVal = MLX.max(absVal).item(Float.self)
+                let minVal = MLX.min(absVal).item(Float.self)
+                let hasNaN = any(isNaN(wf32)).item(Bool.self)
+                let numInf = sum(MLX.abs(wf32) .> 1e30).item(Int.self)  // Count very large values as proxy for inf
+                Flux2Debug.log("  \(key): dtype=\(w.dtype), shape=\(w.shape), mean=\(meanVal), max=\(maxVal), min=\(minVal), hasNaN=\(hasNaN), numInf=\(numInf)")
+            } else {
+                Flux2Debug.log("  \(key): NOT FOUND")
+            }
+        }
+
+        return converted
+    }
+
+    private static func mapCurrentBFLTransformerKey(_ key: String) -> String? {
+        if key.hasPrefix("double_blocks.") {
+            let blockIdx = extractBlockIndex(from: key, prefix: "double_blocks.")
+
+            if let suffix = key.split(separator: ".", maxSplits: 2, omittingEmptySubsequences: false).last {
+                switch String(suffix) {
+                case "img_to_q.weight": return "transformerBlocks.\(blockIdx).attn.toQ.weight"
+                case "img_to_q.scales": return "transformerBlocks.\(blockIdx).attn.toQ.scales"
+                case "img_to_q.biases": return "transformerBlocks.\(blockIdx).attn.toQ.biases"
+                case "img_to_k.weight": return "transformerBlocks.\(blockIdx).attn.toK.weight"
+                case "img_to_k.scales": return "transformerBlocks.\(blockIdx).attn.toK.scales"
+                case "img_to_k.biases": return "transformerBlocks.\(blockIdx).attn.toK.biases"
+                case "img_to_v.weight": return "transformerBlocks.\(blockIdx).attn.toV.weight"
+                case "img_to_v.scales": return "transformerBlocks.\(blockIdx).attn.toV.scales"
+                case "img_to_v.biases": return "transformerBlocks.\(blockIdx).attn.toV.biases"
+                case "img_to_out.weight": return "transformerBlocks.\(blockIdx).attn.toOut.weight"
+                case "img_to_out.scales": return "transformerBlocks.\(blockIdx).attn.toOut.scales"
+                case "img_to_out.biases": return "transformerBlocks.\(blockIdx).attn.toOut.biases"
+                case "txt_to_q.weight": return "transformerBlocks.\(blockIdx).attn.addQProj.weight"
+                case "txt_to_q.scales": return "transformerBlocks.\(blockIdx).attn.addQProj.scales"
+                case "txt_to_q.biases": return "transformerBlocks.\(blockIdx).attn.addQProj.biases"
+                case "txt_to_k.weight": return "transformerBlocks.\(blockIdx).attn.addKProj.weight"
+                case "txt_to_k.scales": return "transformerBlocks.\(blockIdx).attn.addKProj.scales"
+                case "txt_to_k.biases": return "transformerBlocks.\(blockIdx).attn.addKProj.biases"
+                case "txt_to_v.weight": return "transformerBlocks.\(blockIdx).attn.addVProj.weight"
+                case "txt_to_v.scales": return "transformerBlocks.\(blockIdx).attn.addVProj.scales"
+                case "txt_to_v.biases": return "transformerBlocks.\(blockIdx).attn.addVProj.biases"
+                case "txt_to_out.weight": return "transformerBlocks.\(blockIdx).attn.toAddOut.weight"
+                case "txt_to_out.scales": return "transformerBlocks.\(blockIdx).attn.toAddOut.scales"
+                case "txt_to_out.biases": return "transformerBlocks.\(blockIdx).attn.toAddOut.biases"
+                case "img_norm_q.weight": return "transformerBlocks.\(blockIdx).attn.normQ.weight"
+                case "img_norm_k.weight": return "transformerBlocks.\(blockIdx).attn.normK.weight"
+                case "txt_norm_q.weight": return "transformerBlocks.\(blockIdx).attn.normAddedQ.weight"
+                case "txt_norm_k.weight": return "transformerBlocks.\(blockIdx).attn.normAddedK.weight"
+                case "img_mlp_in.weight": return "transformerBlocks.\(blockIdx).ff.activation.proj.weight"
+                case "img_mlp_in.scales": return "transformerBlocks.\(blockIdx).ff.activation.proj.scales"
+                case "img_mlp_in.biases": return "transformerBlocks.\(blockIdx).ff.activation.proj.biases"
+                case "img_mlp_out.weight": return "transformerBlocks.\(blockIdx).ff.linearOut.weight"
+                case "img_mlp_out.scales": return "transformerBlocks.\(blockIdx).ff.linearOut.scales"
+                case "img_mlp_out.biases": return "transformerBlocks.\(blockIdx).ff.linearOut.biases"
+                case "txt_mlp_in.weight": return "transformerBlocks.\(blockIdx).ffContext.activation.proj.weight"
+                case "txt_mlp_in.scales": return "transformerBlocks.\(blockIdx).ffContext.activation.proj.scales"
+                case "txt_mlp_in.biases": return "transformerBlocks.\(blockIdx).ffContext.activation.proj.biases"
+                case "txt_mlp_out.weight": return "transformerBlocks.\(blockIdx).ffContext.linearOut.weight"
+                case "txt_mlp_out.scales": return "transformerBlocks.\(blockIdx).ffContext.linearOut.scales"
+                case "txt_mlp_out.biases": return "transformerBlocks.\(blockIdx).ffContext.linearOut.biases"
+                default: return nil
+                }
+            }
+        }
+
+        if key.hasPrefix("single_blocks.") {
+            let blockIdx = extractBlockIndex(from: key, prefix: "single_blocks.")
+
+            if let suffix = key.split(separator: ".", maxSplits: 2, omittingEmptySubsequences: false).last {
+                switch String(suffix) {
+                case "to_qkv_mlp.weight": return "singleTransformerBlocks.\(blockIdx).attn.toQkvMlp.weight"
+                case "to_qkv_mlp.scales": return "singleTransformerBlocks.\(blockIdx).attn.toQkvMlp.scales"
+                case "to_qkv_mlp.biases": return "singleTransformerBlocks.\(blockIdx).attn.toQkvMlp.biases"
+                case "to_out.weight": return "singleTransformerBlocks.\(blockIdx).attn.toOut.weight"
+                case "to_out.scales": return "singleTransformerBlocks.\(blockIdx).attn.toOut.scales"
+                case "to_out.biases": return "singleTransformerBlocks.\(blockIdx).attn.toOut.biases"
+                case "norm_q.weight": return "singleTransformerBlocks.\(blockIdx).attn.normQ.weight"
+                case "norm_k.weight": return "singleTransformerBlocks.\(blockIdx).attn.normK.weight"
+                default: return nil
+                }
+            }
+        }
+
+        switch key {
+        case "x_embedder.weight":
+            return "xEmbedder.weight"
+        case "x_embedder.scales":
+            return "xEmbedder.scales"
+        case "x_embedder.biases":
+            return "xEmbedder.biases"
+        case "context_embedder.weight":
+            return "contextEmbedder.weight"
+        case "context_embedder.scales":
+            return "contextEmbedder.scales"
+        case "context_embedder.biases":
+            return "contextEmbedder.biases"
+        case "time_embed_1.weight":
+            return "timeGuidanceEmbed.timestepEmbedder.linear1.weight"
+        case "time_embed_1.scales":
+            return "timeGuidanceEmbed.timestepEmbedder.linear1.scales"
+        case "time_embed_1.biases":
+            return "timeGuidanceEmbed.timestepEmbedder.linear1.biases"
+        case "time_embed_2.weight":
+            return "timeGuidanceEmbed.timestepEmbedder.linear2.weight"
+        case "time_embed_2.scales":
+            return "timeGuidanceEmbed.timestepEmbedder.linear2.scales"
+        case "time_embed_2.biases":
+            return "timeGuidanceEmbed.timestepEmbedder.linear2.biases"
+        case "double_mod_img.linear.weight":
+            return "doubleStreamModulationImg.linear.weight"
+        case "double_mod_img.linear.scales":
+            return "doubleStreamModulationImg.linear.scales"
+        case "double_mod_img.linear.biases":
+            return "doubleStreamModulationImg.linear.biases"
+        case "double_mod_txt.linear.weight":
+            return "doubleStreamModulationTxt.linear.weight"
+        case "double_mod_txt.linear.scales":
+            return "doubleStreamModulationTxt.linear.scales"
+        case "double_mod_txt.linear.biases":
+            return "doubleStreamModulationTxt.linear.biases"
+        case "single_mod.linear.weight":
+            return "singleStreamModulation.linear.weight"
+        case "single_mod.linear.scales":
+            return "singleStreamModulation.linear.scales"
+        case "single_mod.linear.biases":
+            return "singleStreamModulation.linear.biases"
+        case "norm_out.weight":
+            return "normOut.linear.weight"
+        case "norm_out.scales":
+            return "normOut.linear.scales"
+        case "norm_out.biases":
+            return "normOut.linear.biases"
+        case "proj_out.weight":
+            return "projOut.weight"
+        case "proj_out.scales":
+            return "projOut.scales"
+        case "proj_out.biases":
+            return "projOut.biases"
+        default:
+            return nil
+        }
+    }
+
+    private static func quantizedBaseKey(for mappedKey: String) -> String? {
+        if mappedKey.hasSuffix(".scales") {
+            return String(mappedKey.dropLast(7)) + ".weight"
+        }
+        if mappedKey.hasSuffix(".biases") {
+            return String(mappedKey.dropLast(7)) + ".weight"
+        }
+        return nil
+    }
+
+    private static func inferQuantizationBits(packedShape: [Int], expectedShape: [Int]) -> Int? {
+        guard packedShape.count == expectedShape.count,
+              let packedCols = packedShape.last,
+              let expectedCols = expectedShape.last,
+              packedCols > 0,
+              expectedCols > 0 else {
+            return nil
+        }
+
+        let numerator = packedCols * 32
+        guard numerator % expectedCols == 0 else {
+            return nil
+        }
+
+        let bits = numerator / expectedCols
+        return bits > 0 ? bits : nil
+    }
+
+    private static func inferGroupSize(expectedShape: [Int], scalesShape: [Int]) -> Int? {
+        guard expectedShape.count == scalesShape.count,
+              let expectedCols = expectedShape.last,
+              let scaleGroups = scalesShape.last,
+              scaleGroups > 0,
+              expectedCols > 0,
+              expectedCols % scaleGroups == 0 else {
+            return nil
+        }
+
+        return expectedCols / scaleGroups
+    }
+
+    private static func maybeSwapNormOutRows(
+        _ weight: MLXArray,
+        mappedKey: String,
+        profile: Flux2TransformerRepoProfile
+    ) -> MLXArray {
+        guard profile.needsNormOutSwap,
+              mappedKey == "normOut.linear.weight",
+              weight.shape.count == 2 else {
+            return weight
+        }
+
+        let dim = weight.shape[0] / 2
+        guard dim > 0 else { return weight }
+        let shiftRows = weight[0..<dim, 0...]
+        let scaleRows = weight[dim..., 0...]
+        return concatenated([scaleRows, shiftRows], axis: 0)
+    }
+
+    private static func requiresStrictPackedValidation(for profile: Flux2TransformerRepoProfile) -> Bool {
+        profile.checkpointFormat == .packedDiffusers
+    }
+
+    private static func isBonsaiTernaryProfile(_ profile: Flux2TransformerRepoProfile) -> Bool {
+        profile.repoId == ModelRegistry.TransformerRepo.prismMLBonsaiKlein4B2bit.rawValue
+    }
+
+    private static func validateCriticalTransformerCoverage(
+        updates: [String: MLXArray],
+        packedModulePaths: Set<String> = [],
+        flatParameters: [String: MLXArray],
+        profile: Flux2TransformerRepoProfile
+    ) throws {
+        guard requiresStrictPackedValidation(for: profile) else {
+            return
+        }
+
+        let missingParameters = flatParameters.keys.sorted().filter { key in
+            updates[key] == nil && !packedModulePaths.contains(parameterModulePath(for: key))
+        }
+        guard missingParameters.isEmpty else {
+            let sample = missingParameters.prefix(12).joined(separator: ", ")
+            let missingByTopLevel = Dictionary(grouping: missingParameters) { key in
+                key.split(separator: ".").first.map(String.init) ?? key
+            }
+            let summary = missingByTopLevel
+                .map { "\($0.key): \($0.value.count)" }
+                .sorted()
+                .joined(separator: ", ")
+            throw Flux2WeightLoaderError.weightMismatch(
+                "Missing \(missingParameters.count) transformer parameters for \(profile.repoId): \(sample). Summary: \(summary)"
+            )
+        }
+    }
+
+    private static func parameterModulePath(for key: String) -> String {
+        if key.hasSuffix(".weight") {
+            return String(key.dropLast(7))
+        }
+        if key.hasSuffix(".bias") {
+            return String(key.dropLast(5))
+        }
+        if key.hasSuffix(".scales") {
+            return String(key.dropLast(7))
+        }
+        if key.hasSuffix(".biases") {
+            return String(key.dropLast(7))
+        }
+        return key
+    }
+
+    private static func supportsPackedAffineRuntime(bits: Int) -> Bool {
+        switch bits {
+        case 2, 3, 4, 5, 6, 8:
+            return true
+        default:
+            return false
+        }
+    }
+
+    /// Extract block index from key like "double_blocks.3.img_attn.qkv.weight"
+    private static func extractBlockIndex(from key: String, prefix: String) -> Int {
+        guard key.hasPrefix(prefix) else { return 0 }
+        let afterPrefix = String(key.dropFirst(prefix.count))
+        guard let dotIndex = afterPrefix.firstIndex(of: ".") else { return 0 }
+        let indexStr = String(afterPrefix[..<dotIndex])
+        return Int(indexStr) ?? 0
+    }
+
+    /// Map weights in Diffusers format (HuggingFace/quanto quantized)
+    /// Entries are removed from the input dict as they are processed to minimize peak memory.
+    /// This avoids holding both raw quanto data (~32GB) and dequantified float16 (~60GB) simultaneously.
+    private static func mapDiffusersTransformerWeights(
+        _ weights: inout [String: MLXArray],
+        profile: Flux2TransformerRepoProfile
+    ) -> [String: MLXArray] {
+        var mapped: [String: MLXArray] = [:]
+        var quantizedData: [String: MLXArray] = [:]  // base_key -> ._data
+        var quantizedScales: [String: MLXArray] = [:]  // base_key -> ._scale
+        let hasExplicitTimeTextEmbed = weights.keys.contains { key in
+            let normalizedKey = key.hasPrefix("transformer.") ? String(key.dropFirst(12)) : key
+            return normalizedKey.hasPrefix("time_text_embed.")
+        }
+
+        // First pass: extract entries from input dict and remove them to free memory
+        let allKeys = Array(weights.keys)
+        for key in allKeys {
+            guard let value = weights.removeValue(forKey: key) else { continue }
+
+            var processedKey = key
+
+            // Remove common prefixes
+            if processedKey.hasPrefix("transformer.") {
+                processedKey = String(processedKey.dropFirst(12))
+            }
+
+            // Check for quanto quantization suffixes
+            if processedKey.hasSuffix("._data") {
+                let baseKey = String(processedKey.dropLast(6))  // Remove "._data"
+                let mappedKey = mapTransformerKeySimple(baseKey, hasExplicitTimeTextEmbed: hasExplicitTimeTextEmbed)
+                quantizedData[mappedKey] = value
+            } else if processedKey.hasSuffix("._scale") {
+                let baseKey = String(processedKey.dropLast(7))  // Remove "._scale"
+                let mappedKey = mapTransformerKeySimple(baseKey, hasExplicitTimeTextEmbed: hasExplicitTimeTextEmbed)
+                quantizedScales[mappedKey] = value
+            } else if processedKey.hasSuffix(".input_scale") || processedKey.hasSuffix(".output_scale") {
+                // Skip activation scales for quanto (value is released)
+                continue
+            } else {
+                // Non-quantized weight, map directly
+                let mappedKey = mapTransformerKeySimple(processedKey, hasExplicitTimeTextEmbed: hasExplicitTimeTextEmbed)
+                mapped[mappedKey] = value
+            }
+        }
+        // Input dict is now empty — raw safetensors references freed
+        Flux2Debug.log("Input dict cleared — freed raw weight entries")
+
+        // Second pass: dequantize qint8 weights to float16
+        // Release quanto data/scale entries progressively to keep peak ≈ 60GB (float16 only)
+        var dequantCount = 0
+        let quantizedKeys = Array(quantizedData.keys)
+        for key in quantizedKeys {
+            guard let data = quantizedData.removeValue(forKey: key) else { continue }
+
+            if let scale = quantizedScales.removeValue(forKey: key) {
+                // Dequantize: weight = data.float() * scale, then convert to float16
+                let dequantized = (data.asType(.float32) * scale).asType(.float16)
+                mapped[key] = dequantized
+                dequantCount += 1
+            } else {
+                // No scale found, use data as-is
+                mapped[key] = data.asType(.float16)
+                Flux2Debug.warning("No scale found for quantized weight: \(key)")
+            }
+            // Periodically materialize and clear GPU cache
+            if dequantCount % 50 == 0 {
+                eval(mapped.values.map { $0 })
+                MLX.Memory.clearCache()
+            }
+        }
+        Flux2Debug.log("Dequantized \(dequantCount) qint8 weights to float16")
+
+        // DEBUG: Print weight statistics for key layers
+        Flux2Debug.log("=== Diffusers Weight Statistics ===")
+        for key in ["xEmbedder.weight", "contextEmbedder.weight",
+                    "transformerBlocks.0.attn.toQ.weight",
+                    "transformerBlocks.0.ff.activation.proj.weight",
+                    "singleTransformerBlocks.0.attn.toQkvMlp.weight"].prefix(5) {
+            if let w = mapped[key] {
+                eval(w)
+                let absVal = MLX.abs(w)
+                let meanVal = mean(absVal).item(Float.self)
+                let maxVal = MLX.max(absVal).item(Float.self)
+                let minVal = MLX.min(absVal).item(Float.self)
+                Flux2Debug.log("  \(key): shape=\(w.shape), mean=\(meanVal), max=\(maxVal), min=\(minVal)")
+            } else {
+                Flux2Debug.log("  \(key): NOT FOUND")
+            }
+        }
+
+        return mapped
+    }
+
+    /// Map key without handling quantization suffixes (called after suffix processing)
+    private static func mapTransformerKeySimple(_ key: String, hasExplicitTimeTextEmbed: Bool = false) -> String {
+        var mapped = key
+
+        // Map single stream blocks (checkpoint may use various formats)
+        mapped = mapped.replacingOccurrences(of: "single_transformer_blocks.", with: "singleTransformerBlocks.")
+        mapped = mapped.replacingOccurrences(of: "single_transformerBlocks.", with: "singleTransformerBlocks.")
+
+        // Map double stream blocks (both formats)
+        mapped = mapped.replacingOccurrences(of: "transformer_blocks.", with: "transformerBlocks.")
+
+        // Map attention components - both underscore and camelCase from checkpoint
+        mapped = mapped.replacingOccurrences(of: "attn.to_q.", with: "attn.toQ.")
+        mapped = mapped.replacingOccurrences(of: "attn.to_k.", with: "attn.toK.")
+        mapped = mapped.replacingOccurrences(of: "attn.to_v.", with: "attn.toV.")
+        mapped = mapped.replacingOccurrences(of: "attn.toQ.", with: "attn.toQ.")
+        mapped = mapped.replacingOccurrences(of: "attn.toK.", with: "attn.toK.")
+        mapped = mapped.replacingOccurrences(of: "attn.toV.", with: "attn.toV.")
+        mapped = mapped.replacingOccurrences(of: "attn.to_out.0.", with: "attn.toOut.")
+        mapped = mapped.replacingOccurrences(of: "attn.toOut.0.", with: "attn.toOut.")
+        mapped = mapped.replacingOccurrences(of: "attn.add_q_proj.", with: "attn.addQProj.")
+        mapped = mapped.replacingOccurrences(of: "attn.add_k_proj.", with: "attn.addKProj.")
+        mapped = mapped.replacingOccurrences(of: "attn.add_v_proj.", with: "attn.addVProj.")
+        mapped = mapped.replacingOccurrences(of: "attn.addQProj.", with: "attn.addQProj.")
+        mapped = mapped.replacingOccurrences(of: "attn.addKProj.", with: "attn.addKProj.")
+        mapped = mapped.replacingOccurrences(of: "attn.addVProj.", with: "attn.addVProj.")
+        mapped = mapped.replacingOccurrences(of: "attn.to_add_out.", with: "attn.toAddOut.")
+        mapped = mapped.replacingOccurrences(of: "attn.toAddOut.", with: "attn.toAddOut.")
+
+        // Map single block attention (fused QKV+MLP)
+        // Model uses toQkvMlp, checkpoint uses to_qkv_mlp_proj
+        mapped = mapped.replacingOccurrences(of: "attn.to_qkv_mlp_proj.", with: "attn.toQkvMlp.")
+
+        // Map to_out without .0. (single stream blocks)
+        mapped = mapped.replacingOccurrences(of: "attn.to_out.", with: "attn.toOut.")
+
+        // Map norms - both formats
+        mapped = mapped.replacingOccurrences(of: "attn.norm_q.", with: "attn.normQ.")
+        mapped = mapped.replacingOccurrences(of: "attn.norm_k.", with: "attn.normK.")
+        mapped = mapped.replacingOccurrences(of: "attn.normQ.", with: "attn.normQ.")
+        mapped = mapped.replacingOccurrences(of: "attn.normK.", with: "attn.normK.")
+        mapped = mapped.replacingOccurrences(of: "attn.norm_added_q.", with: "attn.normAddedQ.")
+        mapped = mapped.replacingOccurrences(of: "attn.norm_added_k.", with: "attn.normAddedK.")
+        mapped = mapped.replacingOccurrences(of: "attn.normAddedQ.", with: "attn.normAddedQ.")
+        mapped = mapped.replacingOccurrences(of: "attn.normAddedK.", with: "attn.normAddedK.")
+        mapped = mapped.replacingOccurrences(of: "norm1_context.", with: "norm1Context.")
+        mapped = mapped.replacingOccurrences(of: "norm2_context.", with: "norm2Context.")
+
+        // Map feedforward
+        // Model uses SwiGLU with activation.proj for input, checkpoint uses linear_in
+        mapped = mapped.replacingOccurrences(of: "ff_context.linear_in.", with: "ffContext.activation.proj.")
+        mapped = mapped.replacingOccurrences(of: "ff_context.linear_out.", with: "ffContext.linearOut.")
+        mapped = mapped.replacingOccurrences(of: ".ff.linear_in.", with: ".ff.activation.proj.")
+        mapped = mapped.replacingOccurrences(of: ".ff.linear_out.", with: ".ff.linearOut.")
+        // Generic fallbacks
+        mapped = mapped.replacingOccurrences(of: "ff_context.", with: "ffContext.")
+        mapped = mapped.replacingOccurrences(of: "linear_in.", with: "activation.proj.")
+        mapped = mapped.replacingOccurrences(of: "linear_out.", with: "linearOut.")
+
+        // Map embeddings
+        mapped = mapped.replacingOccurrences(of: "x_embedder.", with: "xEmbedder.")
+        mapped = mapped.replacingOccurrences(of: "context_embedder.", with: "contextEmbedder.")
+        mapped = mapped.replacingOccurrences(of: "contextEmbedder.", with: "contextEmbedder.")
+        mapped = mapped.replacingOccurrences(of: "time_text_embed.linear_1.", with: "timeGuidanceEmbed.timestepEmbedder.linear1.")
+        mapped = mapped.replacingOccurrences(of: "time_text_embed.linear_2.", with: "timeGuidanceEmbed.timestepEmbedder.linear2.")
+        if hasExplicitTimeTextEmbed {
+            mapped = mapped.replacingOccurrences(of: "time_guidance_embed.linear_1.", with: "timeGuidanceEmbed.guidanceEmbedder.linear1.")
+            mapped = mapped.replacingOccurrences(of: "time_guidance_embed.linear_2.", with: "timeGuidanceEmbed.guidanceEmbedder.linear2.")
+        } else {
+            mapped = mapped.replacingOccurrences(of: "time_guidance_embed.linear_1.", with: "timeGuidanceEmbed.timestepEmbedder.linear1.")
+            mapped = mapped.replacingOccurrences(of: "time_guidance_embed.linear_2.", with: "timeGuidanceEmbed.timestepEmbedder.linear2.")
+        }
+        mapped = mapped.replacingOccurrences(of: "time_text_embed.", with: "timeGuidanceEmbed.")
+        mapped = mapped.replacingOccurrences(of: "time_guidance_embed.", with: "timeGuidanceEmbed.")
+
+        // Map timestep embedder
+        mapped = mapped.replacingOccurrences(of: "timestep_embedder.", with: "timestepEmbedder.")
+        mapped = mapped.replacingOccurrences(of: "guidance_embedder.", with: "guidanceEmbedder.")
+        mapped = mapped.replacingOccurrences(of: "linear_1.", with: "linear1.")
+        mapped = mapped.replacingOccurrences(of: "linear_2.", with: "linear2.")
+
+        // Map modulation layers
+        mapped = mapped.replacingOccurrences(of: "double_stream_modulation_img.", with: "doubleStreamModulationImg.")
+        mapped = mapped.replacingOccurrences(of: "double_stream_modulation_txt.", with: "doubleStreamModulationTxt.")
+        mapped = mapped.replacingOccurrences(of: "single_stream_modulation.", with: "singleStreamModulation.")
+
+        // Map output
+        mapped = mapped.replacingOccurrences(of: "norm_out.", with: "normOut.")
+        mapped = mapped.replacingOccurrences(of: "proj_out.", with: "projOut.")
+        mapped = mapped.replacingOccurrences(of: "normOut.", with: "normOut.")
+        mapped = mapped.replacingOccurrences(of: "projOut.", with: "projOut.")
+
+        return mapped
+    }
+
+    // MARK: - VAE Weight Mapping
+
+    /// Convert HuggingFace VAE weight keys to Swift module paths
+    /// Also transposes Conv2d weights from PyTorch OIHW format to MLX OHWI format
+    public static func mapVAEWeights(_ weights: [String: MLXArray]) -> [String: MLXArray] {
+        var mapped: [String: MLXArray] = [:]
+
+        for (key, var value) in weights {
+            var newKey = key
+
+            // Transpose Conv2d weights from PyTorch OIHW to MLX OHWI format
+            // Conv2d weights have 4 dimensions and key ends with ".weight"
+            // PyTorch: [out_channels, in_channels, kernel_h, kernel_w]
+            // MLX: [out_channels, kernel_h, kernel_w, in_channels]
+            if key.hasSuffix(".weight") && value.ndim == 4 {
+                value = value.transposed(0, 2, 3, 1)  // OIHW -> OHWI
+            }
+
+            // Map encoder/decoder
+            newKey = newKey.replacingOccurrences(of: "encoder.conv_in.", with: "encoder.convIn.")
+            newKey = newKey.replacingOccurrences(of: "encoder.conv_out.", with: "encoder.convOut.")
+            newKey = newKey.replacingOccurrences(of: "encoder.conv_norm_out.", with: "encoder.convNormOut.")
+            newKey = newKey.replacingOccurrences(of: "decoder.conv_in.", with: "decoder.convIn.")
+            newKey = newKey.replacingOccurrences(of: "decoder.conv_out.", with: "decoder.convOut.")
+            newKey = newKey.replacingOccurrences(of: "decoder.conv_norm_out.", with: "decoder.convNormOut.")
+
+            // Map down/up blocks
+            newKey = newKey.replacingOccurrences(of: "down_blocks.", with: "downBlocks.")
+            newKey = newKey.replacingOccurrences(of: "up_blocks.", with: "upBlocks.")
+
+            // Map mid_block - special handling for tuple structure
+            // mid_block.resnets.0 -> midBlock.0 (resnet1)
+            // mid_block.attentions.0 -> midBlock.1 (attention)
+            // mid_block.resnets.1 -> midBlock.2 (resnet2)
+            // Apply to both encoder and decoder
+            newKey = newKey.replacingOccurrences(of: "encoder.mid_block.resnets.0.", with: "encoder.midBlock.0.")
+            newKey = newKey.replacingOccurrences(of: "encoder.mid_block.attentions.0.", with: "encoder.midBlock.1.")
+            newKey = newKey.replacingOccurrences(of: "encoder.mid_block.resnets.1.", with: "encoder.midBlock.2.")
+            newKey = newKey.replacingOccurrences(of: "decoder.mid_block.resnets.0.", with: "decoder.midBlock.0.")
+            newKey = newKey.replacingOccurrences(of: "decoder.mid_block.attentions.0.", with: "decoder.midBlock.1.")
+            newKey = newKey.replacingOccurrences(of: "decoder.mid_block.resnets.1.", with: "decoder.midBlock.2.")
+
+            // Map ResNet blocks in down/up blocks
+            // The model uses tuples: (blocks: [ResnetBlock2D], upsample/downsample: Upsample2D?)
+            // This flattens as: upBlocks.{idx}.0.{resnet_idx} for blocks, upBlocks.{idx}.1 for upsample
+            // Checkpoint uses: up_blocks.{idx}.resnets.{resnet_idx}
+            newKey = newKey.replacingOccurrences(of: "resnets.", with: "0.")
+            newKey = newKey.replacingOccurrences(of: "conv_shortcut.", with: "convShortcut.")
+
+            // Map upsamplers and downsamplers
+            // upsamplers.0.conv -> 1.conv (position 1 in tuple)
+            // downsamplers.0.conv -> 1.conv (position 1 in tuple)
+            newKey = newKey.replacingOccurrences(of: "upsamplers.0.conv.", with: "1.conv.")
+            newKey = newKey.replacingOccurrences(of: "downsamplers.0.conv.", with: "1.conv.")
+
+            // Map attention
+            newKey = newKey.replacingOccurrences(of: "attentions.", with: "attentions.")
+            newKey = newKey.replacingOccurrences(of: "group_norm.", with: "groupNorm.")
+            newKey = newKey.replacingOccurrences(of: "to_q.", with: "toQ.")
+            newKey = newKey.replacingOccurrences(of: "to_k.", with: "toK.")
+            newKey = newKey.replacingOccurrences(of: "to_v.", with: "toV.")
+            newKey = newKey.replacingOccurrences(of: "to_out.0.", with: "toOut.")  // VAE attention has indexed to_out
+            newKey = newKey.replacingOccurrences(of: "to_out.", with: "toOut.")
+
+            // Map quant conv - order matters! More specific patterns first
+            newKey = newKey.replacingOccurrences(of: "post_quant_conv.", with: "postQuantConv.")
+            newKey = newKey.replacingOccurrences(of: "quant_conv.", with: "quantConv.")
+
+            // Map batch norm (bn -> latentBatchNorm)
+            newKey = newKey.replacingOccurrences(of: "bn.", with: "latentBatchNorm.")
+            newKey = newKey.replacingOccurrences(of: "running_mean", with: "runningMean")
+            newKey = newKey.replacingOccurrences(of: "running_var", with: "runningVar")
+            newKey = newKey.replacingOccurrences(of: "num_batches_tracked", with: "numBatchesTracked")
+
+            mapped[newKey] = value
+        }
+
+        return mapped
+    }
+
+    // MARK: - Weight Application
+
+    /// Apply weights to a transformer model
+    /// Note: `weights` is passed inout — for Diffusers quanto format, entries are freed during mapping
+    /// to reduce peak memory from ~98GB to ~71GB.
+    public static func applyTransformerWeights(
+        _ weights: inout [String: MLXArray],
+        to model: Flux2Transformer2DModel,
+        profile: Flux2TransformerRepoProfile
+    ) throws {
+        let mapped = mapTransformerWeights(&weights, profile: profile)
+
+        // Use MLX's built-in weight loading - flatten to get full paths
+        let flattenedArray = model.parameters().flattened()
+        var flatParameters: [String: MLXArray] = [:]
+        for (key, value) in flattenedArray {
+            flatParameters[key] = value
+        }
+
+        var updates: [String: MLXArray] = [:]
+        var moduleUpdates: [String: Module] = [:]
+        var packedModulePaths = Set<String>()
+        var notFound = 0
+
+        // Debug: print flattened model parameter keys by category
+        let modelTransformerBlocks = flatParameters.keys.filter { $0.hasPrefix("transformerBlocks.") }
+        let modelSingleBlocks = flatParameters.keys.filter { $0.hasPrefix("singleTransformerBlocks.") }
+        let modelTimeEmbed = flatParameters.keys.filter { $0.contains("timeGuidanceEmbed") || $0.contains("xEmbedder") }
+        Flux2Debug.verbose("Model parameter keys:")
+        Flux2Debug.verbose("  - transformerBlocks params: \(modelTransformerBlocks.count)")
+        Flux2Debug.verbose("  - singleTransformerBlocks params: \(modelSingleBlocks.count)")
+        Flux2Debug.verbose("  - time/embedder params: \(modelTimeEmbed.count)")
+        Flux2Debug.verbose("  - total params: \(flatParameters.count)")
+        for key in modelTimeEmbed.sorted() {
+            Flux2Debug.verbose("    \(key)")
+        }
+
+        var linearModules: [String: Linear] = [:]
+        for (path, module) in model.leafModules().flattened() {
+            if let linear = module as? Linear {
+                linearModules[path] = linear
+            }
+        }
+
+        var dequantizedWeights: [String: MLXArray] = [:]
+        var packedRuntimeCount = 0
+        var denseFallbackCount = 0
+        var unsupportedFallbackBits = Set<Int>()
+        let requiresBonsaiTernaryPackedRuntime = isBonsaiTernaryProfile(profile)
+
+        for (key, value) in mapped where value.dtype == .uint32 {
+            guard let scales = mapped["\(key.dropLast(6))scales"],
+                  let biases = mapped["\(key.dropLast(6))biases"] else {
+                throw Flux2WeightLoaderError.weightMismatch(
+                    "Packed transformer weight \(key) is missing matching scales/biases tensors for \(profile.repoId)"
+                )
+            }
+
+            guard let expectedShape = flatParameters[key]?.shape else {
+                if requiresBonsaiTernaryPackedRuntime {
+                    throw Flux2WeightLoaderError.weightMismatch(
+                        "Bonsai ternary packed tensor \(key) is not present in the instantiated Klein transformer for \(profile.repoId)"
+                    )
+                }
+
+                Flux2Debug.warning(
+                    "Skipping extra packed transformer tensor \(key) not present in instantiated model for \(profile.repoId)"
+                )
+                continue
+            }
+
+            guard let bits = inferQuantizationBits(packedShape: value.shape, expectedShape: expectedShape),
+                  let groupSize = inferGroupSize(expectedShape: expectedShape, scalesShape: scales.shape) else {
+                throw Flux2WeightLoaderError.weightMismatch(
+                    "Unable to infer quantization metadata for \(key): packed=\(value.shape) scales=\(scales.shape) expected=\(expectedShape)"
+                )
+            }
+
+            if requiresBonsaiTernaryPackedRuntime && (bits != 2 || groupSize != 128) {
+                throw Flux2WeightLoaderError.weightMismatch(
+                    "Bonsai ternary requires packed 2-bit/groupSize=128 runtime weights, got bits=\(bits), groupSize=\(groupSize) for \(key) in \(profile.repoId)"
+                )
+            }
+
+            let modulePath = parameterModulePath(for: key)
+            if supportsPackedAffineRuntime(bits: bits), let linear = linearModules[modulePath] {
+                let quantizedLinear = QuantizedLinear(
+                    weight: value,
+                    bias: linear.bias,
+                    scales: scales,
+                    biases: biases,
+                    groupSize: groupSize,
+                    bits: bits,
+                    mode: .affine
+                )
+                moduleUpdates[modulePath] = quantizedLinear
+                packedModulePaths.insert(modulePath)
+                packedRuntimeCount += 1
+                continue
+            }
+
+            if supportsPackedAffineRuntime(bits: bits) {
+                throw Flux2WeightLoaderError.weightMismatch(
+                    "Packed transformer weight \(key) inferred as \(bits)-bit/groupSize=\(groupSize), but no matching Linear module exists at \(modulePath) for \(profile.repoId)"
+                )
+            }
+
+            if !supportsPackedAffineRuntime(bits: bits) {
+                if requiresBonsaiTernaryPackedRuntime {
+                    throw Flux2WeightLoaderError.weightMismatch(
+                        "Bonsai ternary requires packed 2-bit MLX affine runtime; \(key) inferred unsupported bits=\(bits) for \(profile.repoId)"
+                    )
+                }
+
+                unsupportedFallbackBits.insert(bits)
+            }
+
+            let dense = maybeSwapNormOutRows(
+                dequantized(
+                    value,
+                    scales: scales,
+                    biases: biases,
+                    groupSize: groupSize,
+                    bits: bits
+                ).asType(.float16),
+                mappedKey: key,
+                profile: profile
+            )
+            dequantizedWeights[key] = dense
+            denseFallbackCount += 1
+        }
+
+        if !moduleUpdates.isEmpty {
+            model.update(modules: ModuleChildren.unflattened(moduleUpdates))
+            Flux2Debug.log("Loaded \(packedRuntimeCount) packed transformer weights as QuantizedLinear modules")
+        }
+
+        if !unsupportedFallbackBits.isEmpty {
+            let bits = unsupportedFallbackBits.sorted().map(String.init).joined(separator: ", ")
+            Flux2Debug.warning(
+                "Packed transformer uses unsupported affine runtime bits [\(bits)]; falling back to dense float16 for \(profile.repoId)"
+            )
+        }
+
+        if !dequantizedWeights.isEmpty {
+            Flux2Debug.log("Dequantized \(denseFallbackCount) packed transformer weights to float16 fallback")
+        }
+
+        if requiresBonsaiTernaryPackedRuntime {
+            guard denseFallbackCount == 0 else {
+                throw Flux2WeightLoaderError.weightMismatch(
+                    "Bonsai ternary must not dequantize packed transformer weights; \(denseFallbackCount) weights fell back to dense for \(profile.repoId)"
+                )
+            }
+            Flux2Debug.log("Bonsai ternary packed runtime verified: \(packedRuntimeCount) 2-bit/groupSize=128 QuantizedLinear modules")
+        }
+
+        for (key, value) in mapped {
+            if key.hasSuffix(".scales") || key.hasSuffix(".biases") {
+                continue
+            }
+
+            if packedModulePaths.contains(parameterModulePath(for: key)) {
+                continue
+            }
+
+            let effectiveValue: MLXArray
+            if let dequantized = dequantizedWeights[key] {
+                effectiveValue = dequantized
+            } else {
+                effectiveValue = maybeSwapNormOutRows(value, mappedKey: key, profile: profile)
+            }
+
+            if flatParameters.keys.contains(key) {
+                updates[key] = effectiveValue
+            } else {
+                notFound += 1
+                if notFound <= 10 {
+                    Flux2Debug.log("Warning: No parameter found for key: \(key)")
+                }
+            }
+        }
+
+        if notFound > 10 {
+            Flux2Debug.log("... and \(notFound - 10) more missing parameters")
+        }
+
+        // Update model with new weights using the flattened format
+        try validateCriticalTransformerCoverage(
+            updates: updates,
+            packedModulePaths: packedModulePaths,
+            flatParameters: flatParameters,
+            profile: profile
+        )
+        _ = model.update(parameters: ModuleParameters.unflattened(updates))
+
+        // Debug: print some weight shapes
+        for (key, value) in updates.sorted(by: { $0.key < $1.key }).prefix(10) {
+            Flux2Debug.log("  Weight \(key): shape=\(value.shape), dtype=\(value.dtype)")
+        }
+
+        Flux2Debug.log("Applied \(updates.count) weights to transformer (\(notFound) not found)")
+    }
+
+    /// Apply weights to a VAE model
+    public static func applyVAEWeights(
+        _ weights: [String: MLXArray],
+        to model: AutoencoderKLFlux2
+    ) throws {
+        let mapped = mapVAEWeights(weights)
+
+        // Check for BatchNorm running stats
+        if let runningMean = mapped["latentBatchNorm.runningMean"],
+           let runningVar = mapped["latentBatchNorm.runningVar"] {
+            model.loadBatchNormStats(runningMean: runningMean, runningVar: runningVar)
+        }
+
+        // Flatten model parameters for key matching
+        let flattenedArray = model.parameters().flattened()
+        var flatParameters: [String: MLXArray] = [:]
+        for (key, value) in flattenedArray {
+            flatParameters[key] = value
+        }
+
+        // Debug: print some VAE parameter keys
+        Flux2Debug.verbose("VAE model parameter keys (first 20 of \(flatParameters.count) flattened):")
+        for key in flatParameters.keys.sorted().prefix(20) {
+            Flux2Debug.verbose("  - \(key)")
+        }
+
+        var updates: [String: MLXArray] = [:]
+        var notFound = 0
+
+        for (key, value) in mapped {
+            if key.contains("running") { continue }  // Skip running stats, handled above
+
+            if flatParameters.keys.contains(key) {
+                updates[key] = value
+            } else {
+                notFound += 1
+                if notFound <= 10 {
+                    Flux2Debug.log("VAE Warning: No parameter found for key: \(key)")
+                }
+            }
+        }
+
+        if notFound > 10 {
+            Flux2Debug.log("... and \(notFound - 10) more missing VAE parameters")
+        }
+
+        _ = model.update(parameters: ModuleParameters.unflattened(updates))
+
+        Flux2Debug.log("Applied \(updates.count) weights to VAE (\(notFound) not found)")
+    }
+
+    // MARK: - Quantized Weight Loading
+
+    /// Load quantized transformer weights
+    public static func loadQuantizedTransformer(
+        from modelPath: String,
+        quantization: TransformerQuantization
+    ) throws -> [String: MLXArray] {
+        let weights = try loadWeights(from: modelPath)
+
+        // For qint8, weights may already be quantized in safetensors
+        // If not, we need to quantize them here
+        if quantization != .bf16 {
+            Flux2Debug.log("Loading \(quantization.rawValue) quantized weights")
+        }
+
+        return weights
+    }
+
+    // MARK: - Verification
+
+    /// Verify that all required weights are present
+    public static func verifyTransformerWeights(_ weights: [String: MLXArray]) -> (complete: Bool, missing: [String]) {
+        let requiredPrefixes = [
+            "xEmbedder",
+            "contextEmbedder",
+            "timeGuidanceEmbed",
+            "transformerBlocks.0",
+            "singleTransformerBlocks.0",
+            "normOut",
+            "projOut"
+        ]
+
+        var missing: [String] = []
+
+        for prefix in requiredPrefixes {
+            let hasPrefix = weights.keys.contains { $0.hasPrefix(prefix) }
+            if !hasPrefix {
+                missing.append(prefix)
+            }
+        }
+
+        return (missing.isEmpty, missing)
+    }
+
+    /// Get summary of loaded weights
+    public static func summarizeWeights(_ weights: [String: MLXArray]) {
+        var totalParams: Int64 = 0
+        var byPrefix: [String: Int64] = [:]
+
+        for (key, array) in weights {
+            let params = Int64(array.shape.reduce(1, *))
+            totalParams += params
+
+            // Group by first component
+            let prefix = String(key.split(separator: ".").first ?? Substring(key))
+            byPrefix[prefix, default: 0] += params
+        }
+
+        Flux2Debug.log("Weight Summary:")
+        for (prefix, params) in byPrefix.sorted(by: { $0.value > $1.value }) {
+            let gb = Float(params * 2) / 1_000_000_000  // bf16 = 2 bytes
+            Flux2Debug.log("  \(prefix): \(params) params (~\(String(format: "%.2f", gb))GB)")
+        }
+        Flux2Debug.log("Total: \(totalParams) parameters")
+    }
+}
+
+// MARK: - Errors
+
+public enum Flux2WeightLoaderError: LocalizedError {
+    case noWeightsFound(String)
+    case weightMismatch(String)
+    case fileNotFound(String)
+    case incompatibleShape(expected: [Int], got: [Int], key: String)
+
+    public var errorDescription: String? {
+        switch self {
+        case .noWeightsFound(let path):
+            return "No safetensors files found in: \(path)"
+        case .weightMismatch(let message):
+            return "Weight mismatch: \(message)"
+        case .fileNotFound(let path):
+            return "File not found: \(path)"
+        case .incompatibleShape(let expected, let got, let key):
+            return "Incompatible shape for \(key): expected \(expected), got \(got)"
+        }
+    }
+}