bleam 0.0.7 → 0.0.9

package/templates/native/ios/Bleam/AI/Flux2/Pipeline/FluxImageRunner.swift

@@ -0,0 +1,1282 @@
+import CoreGraphics
+import Foundation
+import MLX
+import MLXNN
+
+public struct FluxImageRunnerPaths: Sendable {
+    public let transformer: URL
+    public let vae: URL
+
+    public init(transformer: URL, vae: URL) {
+        self.transformer = transformer
+        self.vae = vae
+    }
+}
+
+public enum FluxImageGenerationMode: Sendable {
+    case textToImage
+    case imageToImage(sourceImage: CGImage?, maskImage: CGImage?, maskMode: FluxImageMaskMode, referenceImages: [CGImage], referenceInfluences: [Float]?, strength: Float)
+}
+
+public enum FluxImageMaskMode: Sendable {
+    case paint
+    case preserve
+}
+
+public struct FluxLoRAConfig: Sendable {
+    public let id: String?
+    public let path: String
+    public let scale: Float
+
+    public init(id: String? = nil, path: String, scale: Float) {
+        self.id = id
+        self.path = path
+        self.scale = scale
+    }
+}
+
+private enum LoRARuntimeWarnings {
+    nonisolated(unsafe) static var skippedShapeMismatchCount = 0
+    static let maxWarnings = 12
+
+    static func warnShapeMismatch(layerName: String?, inputShape: [Int], downShape: [Int], upShape: [Int]) {
+        skippedShapeMismatchCount += 1
+        guard skippedShapeMismatchCount <= maxWarnings else { return }
+        Flux2Debug.warning(
+            "Skipping runtime LoRA shape mismatch: layer=\(layerName ?? "unknown") input=\(inputShape), down=\(downShape), up=\(upShape)"
+        )
+    }
+}
+
+private struct LoRATensorDiagnostic {
+    let key: String
+    let shape: [Int]
+    let min: Float
+    let max: Float
+    let mean: Float
+    let nonFiniteCount: Int
+}
+
+private struct LoRADiagnostics {
+    let fileName: String
+    let tensorCount: Int
+    let layerCount: Int
+    let rank: Int
+    let targetModel: String
+    let sampleTensors: [LoRATensorDiagnostic]
+    let nonFiniteTensorCount: Int
+    let nonFiniteValueCount: Int
+}
+
+private struct LoRADimensionMismatch {
+    let transformerLayer: String
+    let loraLayer: String
+    let loraName: String
+    let expectedInFeatures: Int
+    let expectedOutFeatures: Int
+    let downShape: [Int]
+    let upShape: [Int]
+
+    var logDescription: String {
+        "\(loraName) layer=\(loraLayer) transformer=\(transformerLayer) expected=[out:\(expectedOutFeatures), in:\(expectedInFeatures)] down=\(downShape) up=\(upShape)"
+    }
+}
+
+private struct LoRAFileApplyReport {
+    let name: String
+    let loadedLayerCount: Int
+    let matchedLayerCount: Int
+    let appliedLayerCount: Int
+    let dimensionMismatchCount: Int
+    let sampleMatchedLayers: [String]
+    let sampleUnmatchedLayers: [String]
+    let sampleDimensionMismatches: [LoRADimensionMismatch]
+}
+
+private struct LoRAApplyReport {
+    let loadedLayerCount: Int
+    let transformerLinearLayerCount: Int
+    let matchedLayerCount: Int
+    let replacedLayerCount: Int
+    let updatedLayerCount: Int
+    let dimensionMismatchCount: Int
+    let sampleDimensionMismatches: [LoRADimensionMismatch]
+    let sampleUnmatchedLoRALayers: [String]
+    let sampleTransformerLayers: [String]
+    let loraReports: [LoRAFileApplyReport]
+
+    var appliedLayerCount: Int { replacedLayerCount + updatedLayerCount }
+}
+
+private struct TransformerLinearTarget {
+    let module: Module
+    let inputFeatures: Int
+    let outputFeatures: Int
+}
+
+private struct LoRAContribution {
+    let down: MLXArray
+    let up: MLXArray
+    let scale: Float
+}
+
+private final class LoRALinear: Linear, @unchecked Sendable {
+    private let down: MLXArray
+    private let up: MLXArray
+    private let scale: Float
+    private let debugName: String?
+
+    init(base: Linear, down: MLXArray, up: MLXArray, scale: Float, debugName: String?) {
+        self.down = down
+        self.up = up
+        self.scale = scale
+        self.debugName = debugName
+        super.init(weight: base.weight, bias: base.bias)
+    }
+
+    override func callAsFunction(_ x: MLXArray) -> MLXArray {
+        let result = super.callAsFunction(x)
+        guard down.shape.count == 2,
+              up.shape.count == 2,
+              x.shape.last == down.shape[1],
+              down.shape[0] == up.shape[1] else {
+            LoRARuntimeWarnings.warnShapeMismatch(
+                layerName: debugName,
+                inputShape: x.shape,
+                downShape: down.shape,
+                upShape: up.shape
+            )
+            return result
+        }
+        let delta = matmul(matmul(x, down.transposed(1, 0)), up.transposed(1, 0)) * scale
+        return result + delta.asType(result.dtype)
+    }
+}
+
+private final class LoRAQuantizedLinear: QuantizedLinear, @unchecked Sendable {
+    private let down: MLXArray
+    private let up: MLXArray
+    private let scale: Float
+    private let debugName: String?
+
+    init(base: QuantizedLinear, down: MLXArray, up: MLXArray, scale: Float, debugName: String?) {
+        self.down = down
+        self.up = up
+        self.scale = scale
+        self.debugName = debugName
+        super.init(
+            weight: base.weight,
+            bias: base.bias,
+            scales: base.scales,
+            biases: base.biases,
+            groupSize: base.groupSize,
+            bits: base.bits,
+            mode: base.mode
+        )
+    }
+
+    override func callAsFunction(_ x: MLXArray) -> MLXArray {
+        let result = super.callAsFunction(x)
+        guard down.shape.count == 2,
+              up.shape.count == 2,
+              x.shape.last == down.shape[1],
+              down.shape[0] == up.shape[1] else {
+            LoRARuntimeWarnings.warnShapeMismatch(
+                layerName: debugName,
+                inputShape: x.shape,
+                downShape: down.shape,
+                upShape: up.shape
+            )
+            return result
+        }
+        let delta = matmul(matmul(x, down.transposed(1, 0)), up.transposed(1, 0)) * scale
+        return result + delta.asType(result.dtype)
+    }
+}
+
+public final class FluxImageRunner: @unchecked Sendable {
+    private static let referenceImageMaxLongSide = 768
+
+    private let paths: FluxImageRunnerPaths
+    private let scheduler = FlowMatchEulerScheduler()
+    private var transformer: Flux2Transformer2DModel?
+    private var vae: AutoencoderKLFlux2?
+
+    public init(paths: FluxImageRunnerPaths) {
+        self.paths = paths
+    }
+
+    public func generate(
+        textEmbeddings: MLXArray,
+        height requestedHeight: Int,
+        width requestedWidth: Int,
+        steps: Int,
+        guidance: Float,
+        seed: UInt64?,
+        mode: FluxImageGenerationMode = .textToImage,
+        loras: [FluxLoRAConfig] = [],
+        sigmas: [Float]? = nil,
+        onStage: ((String) -> Void)? = nil,
+        onProgress: (@Sendable (Int, Int) -> Void)? = nil
+    ) throws -> CGImage {
+        let (height, width) = LatentUtils.validateDimensions(height: requestedHeight, width: requestedWidth)
+        try validateTextEmbeddings(textEmbeddings)
+
+        if let seed {
+            MLXRandom.seed(seed)
+        }
+
+        defer {
+            // Failed generations can otherwise leave partially loaded weights
+            // behind. Keep this path conservative until we add explicit model
+            // lifecycle controls and memory telemetry.
+            self.transformer = nil
+            self.vae = nil
+            MLX.Memory.clearCache()
+        }
+
+        let textLength = textEmbeddings.shape[1]
+        let (textIds, outputImageIds, _) = LatentUtils.combinePositionIDs(
+            textLength: textLength,
+            height: height,
+            width: width
+        )
+
+        let conditioning = try encodeImageConditioning(
+            mode: mode,
+            height: height,
+            width: width,
+            steps: steps,
+            sigmas: sigmas
+        )
+        var packedLatents = conditioning.initialLatents
+        let imageIds = conditioning.imageIds ?? outputImageIds
+        let outputSequenceLength = conditioning.outputSequenceLength
+        eval(packedLatents)
+
+        do {
+            onStage?("denoise")
+            // Keep the Flux transformer alive only for denoising. The VAE is
+            // loaded after this scope so both large model families are not
+            // resident together on memory-constrained machines.
+            let transformer = try loadTransformer()
+            try applyLoRAs(loras, to: transformer)
+
+            let effectiveSteps = scheduler.sigmas.count - 1
+            for stepIndex in 0..<effectiveSteps {
+                if Task.isCancelled {
+                    throw Flux2Error.generationCancelled
+                }
+
+                let sigma = scheduler.sigmas[stepIndex]
+                let timestep = MLXArray([sigma])
+                let hiddenStates = conditioning.referenceLatents.map {
+                    concatenated([packedLatents, $0], axis: 1)
+                } ?? packedLatents
+                eval(hiddenStates)
+                let noisePrediction = transformer(
+                    hiddenStates: hiddenStates,
+                    encoderHiddenStates: textEmbeddings,
+                    timestep: timestep,
+                    guidance: nil,
+                    imgIds: imageIds,
+                    txtIds: textIds
+                )
+                let outputNoisePrediction = conditioning.referenceLatents == nil
+                    ? noisePrediction
+                    : noisePrediction[0..., 0..<outputSequenceLength, 0...]
+
+                packedLatents = scheduler.step(
+                    modelOutput: outputNoisePrediction,
+                    timestep: sigma,
+                    sample: packedLatents
+                )
+                eval(packedLatents)
+
+                if let maskedBlend = conditioning.maskedBlend {
+                    let sigmaNext = stepIndex + 1 < scheduler.sigmas.count ? scheduler.sigmas[stepIndex + 1] : 0
+                    let freshNoise = MLXRandom.normal(packedLatents.shape)
+                    let nextSigma = MLXArray([sigmaNext])
+                    let originalNoised = (1 - nextSigma) * maskedBlend.sourceLatents + nextSigma * freshNoise
+                    packedLatents = (1 - maskedBlend.maskLatents) * originalNoised + maskedBlend.maskLatents * packedLatents
+                    eval(packedLatents)
+                }
+
+                onProgress?(stepIndex + 1, effectiveSteps)
+
+                if (stepIndex + 1) % 5 == 0 {
+                    MLX.Memory.clearCache()
+                }
+            }
+        }
+        self.transformer = nil
+        MLX.Memory.clearCache()
+
+        onStage?("decode")
+        let vae = try loadVAE()
+
+        var patchified = LatentUtils.unpackSequenceToPatchified(packedLatents, height: height, width: width)
+        patchified = LatentUtils.denormalizeLatentsWithBatchNorm(
+            patchified,
+            runningMean: vae.batchNormRunningMean,
+            runningVar: vae.batchNormRunningVar
+        )
+        eval(patchified)
+
+        let latents = LatentUtils.unpatchifyLatents(patchified)
+        eval(latents)
+
+        let decoded = vae.decode(latents)
+        eval(decoded)
+
+        guard var image = postprocessVAEOutput(decoded) else {
+            throw Flux2Error.imageProcessingFailed("Failed to convert VAE output to CGImage")
+        }
+        if case let .imageToImage(sourceImage?, maskImage?, maskMode, _, _, _) = mode {
+            image = Self.compositePreservedSource(
+                generated: image,
+                source: sourceImage,
+                mask: maskImage,
+                maskMode: maskMode,
+                targetHeight: height,
+                targetWidth: width
+            ) ?? image
+        }
+
+        // Do not cache the decoder for now. This keeps repeated generations
+        // slower, but avoids carrying VAE weights after the image is written.
+        self.vae = nil
+        MLX.Memory.clearCache()
+        return image
+    }
+
+    private struct ImageConditioning {
+        let initialLatents: MLXArray
+        let referenceLatents: MLXArray?
+        let imageIds: MLXArray?
+        let outputSequenceLength: Int
+        let maskedBlend: MaskedBlend?
+    }
+
+    private struct MaskedBlend {
+        let sourceLatents: MLXArray
+        let maskLatents: MLXArray
+    }
+
+    private func encodeImageConditioning(
+        mode: FluxImageGenerationMode,
+        height: Int,
+        width: Int,
+        steps: Int,
+        sigmas: [Float]?
+    ) throws -> ImageConditioning {
+        let randomNoise = LatentUtils.packPatchifiedToSequence(
+            LatentUtils.generatePatchifiedLatents(height: height, width: width)
+        )
+        eval(randomNoise)
+
+        guard case let .imageToImage(sourceImage, maskImage, maskMode, referenceImages, referenceInfluences, strength) = mode else {
+            scheduler.setTimesteps(numInferenceSteps: steps, imageSeqLen: randomNoise.shape[1], strength: 1.0)
+            if let sigmas { scheduler.setCustomSigmas(sigmas) }
+            return ImageConditioning(
+                initialLatents: randomNoise,
+                referenceLatents: nil,
+                imageIds: nil,
+                outputSequenceLength: randomNoise.shape[1],
+                maskedBlend: nil
+            )
+        }
+
+        guard sourceImage != nil || !referenceImages.isEmpty else {
+            throw Flux2Error.invalidConfiguration("Image-to-image mode requires a source image or references")
+        }
+        guard referenceImages.count <= 3 else {
+            throw Flux2Error.invalidConfiguration("Provide at most 3 reference images")
+        }
+        if let referenceInfluences, referenceInfluences.count != referenceImages.count {
+            throw Flux2Error.invalidConfiguration("referenceInfluences count must match reference image count")
+        }
+        if maskImage != nil, sourceImage == nil {
+            throw Flux2Error.invalidConfiguration("Masked image generation requires a source image")
+        }
+
+        scheduler.setTimesteps(numInferenceSteps: steps, imageSeqLen: randomNoise.shape[1], strength: strength)
+        if let sigmas { scheduler.setCustomSigmas(sigmas) }
+        let vae = try loadVAE()
+        let sourcePacked: MLXArray?
+        if let sourceImage, strength < 1.0 || maskImage != nil {
+            let sourceProcessed = preprocessImageForVAE(sourceImage, targetHeight: height, targetWidth: width)
+            let sourceRaw = vae.encode(sourceProcessed, samplePosterior: false)
+            var sourcePatchified = LatentUtils.packLatentsToPatchified(sourceRaw)
+            sourcePatchified = LatentUtils.normalizeLatentsWithBatchNorm(
+                sourcePatchified,
+                runningMean: vae.batchNormRunningMean,
+                runningVar: vae.batchNormRunningVar
+            )
+            let packed = LatentUtils.packPatchifiedToSequence(sourcePatchified)
+            eval(packed)
+            sourcePacked = packed
+        } else {
+            sourcePacked = nil
+        }
+
+        let maskedBlend: MaskedBlend?
+        let maskLatents: MLXArray?
+        if let sourcePacked, let maskImage {
+            var packedMask = Self.packMaskForLatentBlending(maskImage, targetHeight: height, targetWidth: width)
+            if maskMode == .preserve {
+                packedMask = 1 - packedMask
+            }
+            eval(packedMask)
+            maskLatents = packedMask
+            maskedBlend = MaskedBlend(sourceLatents: sourcePacked, maskLatents: packedMask)
+        } else {
+            maskLatents = nil
+            maskedBlend = nil
+        }
+
+        let initialLatents: MLXArray
+        if let sourcePacked, let maskLatents {
+            let sourceNoised = scheduler.scaleNoise(
+                sample: sourcePacked,
+                sigma: scheduler.initialSigma,
+                noise: randomNoise
+            )
+            initialLatents = (1 - maskLatents) * sourceNoised + maskLatents * randomNoise
+        } else if let sourcePacked, strength < 1.0 {
+            initialLatents = scheduler.scaleNoise(
+                sample: sourcePacked,
+                sigma: scheduler.initialSigma,
+                noise: randomNoise
+            )
+        } else {
+            initialLatents = randomNoise
+        }
+        eval(initialLatents)
+
+        let encodedReferences = try encodeReferenceImages(referenceImages, influences: referenceInfluences, vae: vae)
+        self.vae = nil
+        MLX.Memory.clearCache()
+
+        let imageIds: MLXArray?
+        if let referencePositionIds = encodedReferences.positionIds {
+            let outputImageIds = LatentUtils.generateImagePositionIDs(height: height, width: width)
+            imageIds = concatenated([outputImageIds, referencePositionIds], axis: 0)
+        } else {
+            imageIds = nil
+        }
+
+        return ImageConditioning(
+            initialLatents: initialLatents,
+            referenceLatents: encodedReferences.latents,
+            imageIds: imageIds,
+            outputSequenceLength: randomNoise.shape[1],
+            maskedBlend: maskedBlend
+        )
+    }
+
+    private static func packMaskForLatentBlending(
+        _ mask: CGImage,
+        targetHeight: Int,
+        targetWidth: Int
+    ) -> MLXArray {
+        let latentHeight = targetHeight / 16
+        let latentWidth = targetWidth / 16
+        var pixels = [UInt8](repeating: 0, count: latentHeight * latentWidth)
+
+        guard let context = CGContext(
+            data: &pixels,
+            width: latentWidth,
+            height: latentHeight,
+            bitsPerComponent: 8,
+            bytesPerRow: latentWidth,
+            space: CGColorSpaceCreateDeviceGray(),
+            bitmapInfo: CGImageAlphaInfo.none.rawValue
+        ) else {
+            return MLXArray.zeros([1, latentHeight * latentWidth, 1])
+        }
+
+        context.interpolationQuality = .high
+        context.draw(mask, in: CGRect(x: 0, y: 0, width: latentWidth, height: latentHeight))
+        let values = pixels.map { Float($0) / 255 }
+        return MLXArray(values).reshaped([1, latentHeight * latentWidth, 1])
+    }
+
+    private static func compositePreservedSource(
+        generated: CGImage,
+        source: CGImage,
+        mask: CGImage,
+        maskMode: FluxImageMaskMode,
+        targetHeight: Int,
+        targetWidth: Int
+    ) -> CGImage? {
+        let bytesPerPixel = 4
+        let bytesPerRow = bytesPerPixel * targetWidth
+        var generatedPixels = [UInt8](repeating: 0, count: targetHeight * bytesPerRow)
+        var sourcePixels = [UInt8](repeating: 0, count: targetHeight * bytesPerRow)
+        var maskPixels = [UInt8](repeating: 0, count: targetHeight * targetWidth)
+
+        guard let generatedContext = CGContext(
+            data: &generatedPixels,
+            width: targetWidth,
+            height: targetHeight,
+            bitsPerComponent: 8,
+            bytesPerRow: bytesPerRow,
+            space: CGColorSpaceCreateDeviceRGB(),
+            bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue
+        ), let sourceContext = CGContext(
+            data: &sourcePixels,
+            width: targetWidth,
+            height: targetHeight,
+            bitsPerComponent: 8,
+            bytesPerRow: bytesPerRow,
+            space: CGColorSpaceCreateDeviceRGB(),
+            bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue
+        ), let maskContext = CGContext(
+            data: &maskPixels,
+            width: targetWidth,
+            height: targetHeight,
+            bitsPerComponent: 8,
+            bytesPerRow: targetWidth,
+            space: CGColorSpaceCreateDeviceGray(),
+            bitmapInfo: CGImageAlphaInfo.none.rawValue
+        ) else {
+            return nil
+        }
+
+        let targetRect = CGRect(x: 0, y: 0, width: targetWidth, height: targetHeight)
+        generatedContext.interpolationQuality = .high
+        generatedContext.draw(generated, in: targetRect)
+        sourceContext.interpolationQuality = .high
+        sourceContext.draw(source, in: targetRect)
+        maskContext.interpolationQuality = .high
+        maskContext.draw(mask, in: targetRect)
+
+        for index in 0..<(targetHeight * targetWidth) {
+            let maskValue = maskPixels[index]
+            let shouldPreserve = maskMode == .paint ? maskValue < 128 : maskValue >= 128
+            guard shouldPreserve else { continue }
+            let byteIndex = index * bytesPerPixel
+            generatedPixels[byteIndex] = sourcePixels[byteIndex]
+            generatedPixels[byteIndex + 1] = sourcePixels[byteIndex + 1]
+            generatedPixels[byteIndex + 2] = sourcePixels[byteIndex + 2]
+        }
+
+        return generatedContext.makeImage()
+    }
+
+    private func encodeReferenceImages(
+        _ images: [CGImage],
+        influences: [Float]?,
+        vae: AutoencoderKLFlux2
+    ) throws -> (latents: MLXArray?, positionIds: MLXArray?) {
+        guard !images.isEmpty else {
+            return (nil, nil)
+        }
+
+        var allLatents: [MLXArray] = []
+        var latentHeights: [Int] = []
+        var latentWidths: [Int] = []
+
+        for (index, image) in images.enumerated() {
+            let size = Self.resolveReferenceImageDimensions(image)
+            let processed = preprocessImageForVAE(image, targetHeight: size.height, targetWidth: size.width)
+            let raw = vae.encode(processed, samplePosterior: false)
+            var patchified = LatentUtils.packLatentsToPatchified(raw)
+            patchified = LatentUtils.normalizeLatentsWithBatchNorm(
+                patchified,
+                runningMean: vae.batchNormRunningMean,
+                runningVar: vae.batchNormRunningVar
+            )
+            eval(patchified)
+
+            var packed = LatentUtils.packPatchifiedToSequence(patchified).squeezed(axis: 0)
+            if let influence = influences?[index] {
+                packed = packed * influence
+                eval(packed)
+            }
+            allLatents.append(packed)
+            latentHeights.append(size.height / 16)
+            latentWidths.append(size.width / 16)
+        }
+
+        let latents = concatenated(allLatents, axis: 0).expandedDimensions(axis: 0)
+        let positionIds = LatentUtils.generateReferenceImagePositionIDs(
+            latentHeights: latentHeights,
+            latentWidths: latentWidths,
+            scale: 10
+        )
+        return (latents, positionIds)
+    }
+
+    private static func resolveReferenceImageDimensions(_ image: CGImage) -> (height: Int, width: Int) {
+        let factor = 16
+        let longSide = max(image.height, image.width)
+        let scale = min(1.0, CGFloat(referenceImageMaxLongSide) / CGFloat(longSide))
+
+        func alignedDown(_ value: CGFloat) -> Int {
+            max(factor, Int(value) / factor * factor)
+        }
+
+        return LatentUtils.validateDimensions(
+            height: alignedDown(CGFloat(image.height) * scale),
+            width: alignedDown(CGFloat(image.width) * scale)
+        )
+    }
+
+    private func preprocessImageForVAE(_ image: CGImage, targetHeight: Int, targetWidth: Int) -> MLXArray {
+        let bytesPerPixel = 4
+        let bytesPerRow = bytesPerPixel * targetWidth
+        var pixelData = [UInt8](repeating: 0, count: targetHeight * bytesPerRow)
+
+        guard let context = CGContext(
+            data: &pixelData,
+            width: targetWidth,
+            height: targetHeight,
+            bitsPerComponent: 8,
+            bytesPerRow: bytesPerRow,
+            space: CGColorSpaceCreateDeviceRGB(),
+            bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue
+        ) else {
+            return MLXRandom.normal([1, 3, targetHeight, targetWidth])
+        }
+
+        context.interpolationQuality = .high
+        context.draw(image, in: CGRect(x: 0, y: 0, width: targetWidth, height: targetHeight))
+
+        var floatData = [Float](repeating: 0, count: targetHeight * targetWidth * 3)
+        for y in 0..<targetHeight {
+            for x in 0..<targetWidth {
+                let pixelIndex = y * targetWidth + x
+                let byteIndex = y * bytesPerRow + x * bytesPerPixel
+                floatData[pixelIndex] = Float(pixelData[byteIndex]) / 127.5 - 1.0
+                floatData[targetHeight * targetWidth + pixelIndex] = Float(pixelData[byteIndex + 1]) / 127.5 - 1.0
+                floatData[2 * targetHeight * targetWidth + pixelIndex] = Float(pixelData[byteIndex + 2]) / 127.5 - 1.0
+            }
+        }
+
+        return MLXArray(floatData).reshaped([1, 3, targetHeight, targetWidth])
+    }
+
+    private func validateTextEmbeddings(_ textEmbeddings: MLXArray) throws {
+        guard textEmbeddings.shape.count == 3 else {
+            throw Flux2Error.invalidConfiguration("Expected text embeddings rank 3, got shape \(textEmbeddings.shape)")
+        }
+
+        guard textEmbeddings.shape[0] == 1 else {
+            throw Flux2Error.invalidConfiguration("Only batch size 1 is supported for image generation")
+        }
+
+        guard textEmbeddings.shape[2] == Flux2Model.klein4B.jointAttentionDim else {
+            throw Flux2Error.invalidConfiguration(
+                "Expected Klein 4B text embedding width \(Flux2Model.klein4B.jointAttentionDim), got \(textEmbeddings.shape[2])"
+            )
+        }
+    }
+
+    private func applyLoRAs(_ loras: [FluxLoRAConfig], to model: Flux2Transformer2DModel) throws {
+        guard !loras.isEmpty else { return }
+        var linearModules: [String: TransformerLinearTarget] = [:]
+        for (path, module) in model.leafModules().flattened() {
+            if let quantizedLinear = module as? QuantizedLinear {
+                guard let dims = Self.quantizedLinearDims(quantizedLinear) else { continue }
+                linearModules[path] = TransformerLinearTarget(module: quantizedLinear, inputFeatures: dims.in, outputFeatures: dims.out)
+            } else if let linear = module as? Linear {
+                guard let dims = Self.linearDims(weightShape: linear.weight.shape) else { continue }
+                linearModules[path] = TransformerLinearTarget(module: linear, inputFeatures: dims.in, outputFeatures: dims.out)
+            }
+        }
+
+        var contributionsByTarget: [String: [LoRAContribution]] = [:]
+        var loadedLayerCount = 0
+        var matchedLayers = Set<String>()
+        var dimensionMismatchCount = 0
+        var sampleDimensionMismatches: [LoRADimensionMismatch] = []
+        var sampleUnmatchedLoRALayers: [String] = []
+        var fileReports: [LoRAFileApplyReport] = []
+        for lora in loras {
+            let name = lora.id ?? URL(fileURLWithPath: lora.path).deletingPathExtension().lastPathComponent
+            let tensors = try loadArrays(url: URL(fileURLWithPath: lora.path))
+            let groups = Self.groupLoRATensors(tensors)
+            let diagnostics = Self.loraDiagnostics(filePath: lora.path, tensors: tensors, groups: groups)
+            Self.logLoRADiagnostics(diagnostics)
+            loadedLayerCount += groups.count
+            var fileMatchedLayers = Set<String>()
+            var fileAppliedLayers = Set<String>()
+            var fileMismatchLayers = Set<String>()
+            var fileMismatchSamples: [LoRADimensionMismatch] = []
+            var fileUnmatchedLayers: [String] = []
+            var applied = 0
+            for (target, group) in groups {
+                guard let down = group.down, let up = group.up else { continue }
+                guard let linear = linearModules[target] else {
+                    if fileUnmatchedLayers.count < 8 { fileUnmatchedLayers.append(target) }
+                    if sampleUnmatchedLoRALayers.count < 8 { sampleUnmatchedLoRALayers.append(target) }
+                    continue
+                }
+                fileMatchedLayers.insert(target)
+                matchedLayers.insert(target)
+                guard let normalized = Self.normalizedLoRAPair(
+                    down: down,
+                    up: up,
+                    inputFeatures: linear.inputFeatures,
+                    outputFeatures: linear.outputFeatures
+                ) else {
+                    let mismatch = LoRADimensionMismatch(
+                        transformerLayer: target,
+                        loraLayer: target,
+                        loraName: name,
+                        expectedInFeatures: linear.inputFeatures,
+                        expectedOutFeatures: linear.outputFeatures,
+                        downShape: down.shape,
+                        upShape: up.shape
+                    )
+                    dimensionMismatchCount += 1
+                    fileMismatchLayers.insert(target)
+                    if sampleDimensionMismatches.count < 8 { sampleDimensionMismatches.append(mismatch) }
+                    if fileMismatchSamples.count < 4 { fileMismatchSamples.append(mismatch) }
+                    continue
+                }
+                let rank = max(1, down.shape.first ?? 1)
+                let alphaScale = group.alpha.map { Float($0.item(Float.self)) / Float(rank) } ?? 1
+                contributionsByTarget[target, default: []].append(LoRAContribution(
+                    down: normalized.down,
+                    up: normalized.up,
+                    scale: lora.scale * alphaScale
+                ))
+                fileAppliedLayers.insert(target)
+                applied += 1
+            }
+            Flux2Debug.log("Loaded LoRA \"\(name)\": layers=\(diagnostics.layerCount), rank=\(diagnostics.rank), target=\(diagnostics.targetModel)")
+            fileReports.append(LoRAFileApplyReport(
+                name: name,
+                loadedLayerCount: groups.count,
+                matchedLayerCount: fileMatchedLayers.count,
+                appliedLayerCount: applied,
+                dimensionMismatchCount: fileMismatchLayers.count,
+                sampleMatchedLayers: Array(fileMatchedLayers.sorted().prefix(8)),
+                sampleUnmatchedLayers: fileUnmatchedLayers,
+                sampleDimensionMismatches: fileMismatchSamples
+            ))
+        }
+
+        var moduleUpdates: [String: Module] = [:]
+        for target in contributionsByTarget.keys.sorted() {
+            guard let contributions = contributionsByTarget[target] else { continue }
+            guard let linear = linearModules[target], !contributions.isEmpty else { continue }
+            var combinedDown: MLXArray?
+            var combinedUp: MLXArray?
+            for contribution in contributions {
+                let scaledDown = contribution.down * contribution.scale
+                if let existingDown = combinedDown, let existingUp = combinedUp {
+                    combinedDown = concatenated([existingDown, scaledDown], axis: 0)
+                    combinedUp = concatenated([existingUp, contribution.up], axis: 1)
+                } else {
+                    combinedDown = scaledDown
+                    combinedUp = contribution.up
+                }
+            }
+            guard let combinedDown, let combinedUp else { continue }
+            let materializedDown = combinedDown.asType(.float16)
+            let materializedUp = combinedUp.asType(.float16)
+            eval(materializedDown, materializedUp)
+            if let quantized = linear.module as? QuantizedLinear {
+                moduleUpdates[target] = LoRAQuantizedLinear(
+                    base: quantized,
+                    down: materializedDown,
+                    up: materializedUp,
+                    scale: 1,
+                    debugName: target
+                )
+            } else if let dense = linear.module as? Linear {
+                moduleUpdates[target] = LoRALinear(
+                    base: dense,
+                    down: materializedDown,
+                    up: materializedUp,
+                    scale: 1,
+                    debugName: target
+                )
+            }
+        }
+        contributionsByTarget.removeAll()
+        MLX.Memory.clearCache()
+
+        let report = LoRAApplyReport(
+            loadedLayerCount: loadedLayerCount,
+            transformerLinearLayerCount: linearModules.count,
+            matchedLayerCount: matchedLayers.count,
+            replacedLayerCount: moduleUpdates.count,
+            updatedLayerCount: 0,
+            dimensionMismatchCount: dimensionMismatchCount,
+            sampleDimensionMismatches: sampleDimensionMismatches,
+            sampleUnmatchedLoRALayers: sampleUnmatchedLoRALayers,
+            sampleTransformerLayers: Array(linearModules.keys.sorted().prefix(8)),
+            loraReports: fileReports
+        )
+        Self.logLoRAApplyReport(report, queuedCount: loras.count)
+
+        guard !moduleUpdates.isEmpty else {
+            let sample = report.sampleDimensionMismatches.first?.logDescription
+                ?? report.sampleUnmatchedLoRALayers.first
+                ?? "none"
+            throw Flux2Error.invalidConfiguration(
+                "No compatible LoRA targets found (loaded=\(report.loadedLayerCount), matched=\(report.matchedLayerCount), dimensionMismatches=\(report.dimensionMismatchCount), sample=\(sample))"
+            )
+        }
+        model.update(modules: ModuleChildren.unflattened(moduleUpdates))
+        eval(model.parameters())
+        MLX.Memory.clearCache()
+    }
+
+    private struct LoRATensorGroup {
+        var down: MLXArray?
+        var up: MLXArray?
+        var alpha: MLXArray?
+    }
+
+    private static func linearDims(weightShape: [Int]) -> (out: Int, in: Int)? {
+        guard weightShape.count >= 2 else { return nil }
+        return (out: weightShape[weightShape.count - 2], in: weightShape[weightShape.count - 1])
+    }
+
+    private static func quantizedLinearDims(_ linear: QuantizedLinear) -> (out: Int, in: Int)? {
+        let weightShape = linear.weight.shape
+        guard weightShape.count >= 2, linear.bits > 0, 32 % linear.bits == 0 else { return nil }
+        let out = weightShape[weightShape.count - 2]
+        let packedIn = weightShape[weightShape.count - 1]
+        return (out: out, in: packedIn * (32 / linear.bits))
+    }
+
+    private static func normalizedLoRAPair(
+        down: MLXArray,
+        up: MLXArray,
+        inputFeatures: Int,
+        outputFeatures: Int
+    ) -> (down: MLXArray, up: MLXArray)? {
+        guard down.shape.count == 2, up.shape.count == 2 else { return nil }
+        if down.shape[1] == inputFeatures,
+           up.shape[0] == outputFeatures,
+           down.shape[0] == up.shape[1] {
+            return (down, up)
+        }
+        if down.shape[0] == inputFeatures,
+           up.shape[1] == outputFeatures,
+           down.shape[1] == up.shape[0] {
+            return (down.transposed(1, 0), up.transposed(1, 0))
+        }
+        return nil
+    }
+
+    private static func loraDiagnostics(
+        filePath: String,
+        tensors: [String: MLXArray],
+        groups: [String: LoRATensorGroup]
+    ) -> LoRADiagnostics {
+        var sampleTensors: [LoRATensorDiagnostic] = []
+        var nonFiniteTensorCount = 0
+        var nonFiniteValueCount = 0
+        let visibleKeys = tensors.keys.filter { !$0.hasPrefix("__") }.sorted()
+        for key in visibleKeys.prefix(3) {
+            guard let value = tensors[key] else { continue }
+            let diagnostic = tensorDiagnostic(key: key, value: value)
+            if diagnostic.nonFiniteCount > 0 {
+                nonFiniteTensorCount += 1
+                nonFiniteValueCount += diagnostic.nonFiniteCount
+            }
+            sampleTensors.append(diagnostic)
+        }
+
+        return LoRADiagnostics(
+            fileName: URL(fileURLWithPath: filePath).lastPathComponent,
+            tensorCount: visibleKeys.count,
+            layerCount: groups.count,
+            rank: groups.values.compactMap { $0.down?.shape.first }.first ?? 0,
+            targetModel: detectLoRATargetModel(Array(groups.keys)),
+            sampleTensors: sampleTensors,
+            nonFiniteTensorCount: nonFiniteTensorCount,
+            nonFiniteValueCount: nonFiniteValueCount
+        )
+    }
+
+    private static func tensorDiagnostic(key: String, value: MLXArray) -> LoRATensorDiagnostic {
+        let checked = value.asType(.float32)
+        eval(checked)
+        let values = checked.asArray(Float.self)
+        var minValue = Float.greatestFiniteMagnitude
+        var maxValue = -Float.greatestFiniteMagnitude
+        var sum: Float = 0
+        var finiteCount = 0
+        var nonFiniteCount = 0
+
+        for value in values {
+            guard value.isFinite else {
+                nonFiniteCount += 1
+                continue
+            }
+            minValue = min(minValue, value)
+            maxValue = max(maxValue, value)
+            sum += value
+            finiteCount += 1
+        }
+
+        if finiteCount == 0 {
+            minValue = 0
+            maxValue = 0
+        }
+
+        return LoRATensorDiagnostic(
+            key: key,
+            shape: value.shape,
+            min: minValue,
+            max: maxValue,
+            mean: finiteCount > 0 ? sum / Float(finiteCount) : 0,
+            nonFiniteCount: nonFiniteCount
+        )
+    }
+
+    private static func detectLoRATargetModel(_ layers: [String]) -> String {
+        var doubleBlocks = Set<Int>()
+        var singleBlocks = Set<Int>()
+        for path in layers {
+            if path.hasPrefix("transformerBlocks.") {
+                doubleBlocks.insert(extractBlockIndex(from: path, prefix: "transformerBlocks."))
+            } else if path.hasPrefix("singleTransformerBlocks.") {
+                singleBlocks.insert(extractBlockIndex(from: path, prefix: "singleTransformerBlocks."))
+            }
+        }
+
+        let maxDouble = doubleBlocks.max() ?? 0
+        let maxSingle = singleBlocks.max() ?? 0
+        if maxDouble == 4 && maxSingle == 19 { return "klein-4b" }
+        if maxDouble == 7 && maxSingle == 23 { return "klein-9b" }
+        if maxDouble == 7 && maxSingle == 47 { return "dev" }
+        return "unknown"
+    }
+
+    private static func extractBlockIndex(from path: String, prefix: String) -> Int {
+        guard let range = path.range(of: prefix) else { return 0 }
+        let afterPrefix = path[range.upperBound...]
+        let index = afterPrefix.prefix(while: { $0.isNumber })
+        return Int(index) ?? 0
+    }
+
+    private static func logLoRADiagnostics(_ diagnostics: LoRADiagnostics) {
+        Flux2Debug.log(
+            "LoRA diagnostics file=\(diagnostics.fileName) tensors=\(diagnostics.tensorCount) layers=\(diagnostics.layerCount) rank=\(diagnostics.rank) target=\(diagnostics.targetModel) nonFiniteTensors=\(diagnostics.nonFiniteTensorCount) nonFiniteValues=\(diagnostics.nonFiniteValueCount)"
+        )
+        for tensor in diagnostics.sampleTensors {
+            Flux2Debug.log(
+                "LoRA tensor sample key=\(tensor.key) shape=\(tensor.shape) min=\(tensor.min) max=\(tensor.max) mean=\(tensor.mean) nonFinite=\(tensor.nonFiniteCount)"
+            )
+        }
+        if diagnostics.nonFiniteValueCount > 0 {
+            Flux2Debug.warning("LoRA contains non-finite values: file=\(diagnostics.fileName), values=\(diagnostics.nonFiniteValueCount)")
+        }
+    }
+
+    private static func logLoRAApplyReport(_ report: LoRAApplyReport, queuedCount: Int) {
+        Flux2Debug.log(
+            "Applied \(queuedCount) queued LoRA files dynamically to transformer: loadedLayers=\(report.loadedLayerCount), transformerLinearLayers=\(report.transformerLinearLayerCount), matchedLayers=\(report.matchedLayerCount), appliedLayers=\(report.appliedLayerCount), dimensionMismatches=\(report.dimensionMismatchCount)"
+        )
+        for loraReport in report.loraReports {
+            if loraReport.appliedLayerCount > 0 {
+                Flux2Debug.log(
+                    "LoRA \"\(loraReport.name)\" used: loaded=\(loraReport.loadedLayerCount), matched=\(loraReport.matchedLayerCount), applied=\(loraReport.appliedLayerCount), mismatches=\(loraReport.dimensionMismatchCount)"
+                )
+            } else {
+                Flux2Debug.warning(
+                    "LoRA \"\(loraReport.name)\" loaded but not used: loaded=\(loraReport.loadedLayerCount), matched=\(loraReport.matchedLayerCount), applied=0, mismatches=\(loraReport.dimensionMismatchCount)"
+                )
+            }
+
+            if loraReport.matchedLayerCount < loraReport.loadedLayerCount, !loraReport.sampleUnmatchedLayers.isEmpty {
+                Flux2Debug.warning("LoRA \"\(loraReport.name)\" sample unmatched layers: \(loraReport.sampleUnmatchedLayers.joined(separator: ", "))")
+            }
+            for mismatch in loraReport.sampleDimensionMismatches {
+                Flux2Debug.warning("LoRA \"\(loraReport.name)\" dimension mismatch: \(mismatch.logDescription)")
+            }
+        }
+
+        if report.appliedLayerCount == 0 {
+            for mismatch in report.sampleDimensionMismatches {
+                Flux2Debug.warning("Flux LoRA dimension mismatch: \(mismatch.logDescription)")
+            }
+            Flux2Debug.warning("No Flux LoRA layers were applied. Sample unmatched LoRA layers: \(report.sampleUnmatchedLoRALayers.joined(separator: ", "))")
+            Flux2Debug.warning("Sample transformer linear layers: \(report.sampleTransformerLayers.joined(separator: ", "))")
+        }
+    }
+
+    private static func groupLoRATensors(_ tensors: [String: MLXArray]) -> [String: LoRATensorGroup] {
+        var rawGroups: [String: LoRATensorGroup] = [:]
+        for (rawKey, value) in tensors {
+            guard let parsed = parseLoRAKey(rawKey) else { continue }
+            var group = rawGroups[parsed.target] ?? LoRATensorGroup()
+            switch parsed.kind {
+            case "down": group.down = value
+            case "up": group.up = value
+            case "alpha": group.alpha = value
+            default: break
+            }
+            rawGroups[parsed.target] = group
+        }
+
+        var groups: [String: LoRATensorGroup] = [:]
+        for (rawTarget, group) in rawGroups {
+            if isCombinedQKVLayer(rawTarget),
+               let up = group.up,
+               up.shape.count == 2,
+               up.shape[0].isMultiple(of: 3) {
+                let outputSize = up.shape[0] / 3
+                let targets = qkvTargets(for: rawTarget)
+                guard targets.count == 3 else { continue }
+                for (index, target) in targets.enumerated() {
+                    groups[target] = LoRATensorGroup(
+                        down: group.down,
+                        up: up[(index * outputSize)..<((index + 1) * outputSize), 0...],
+                        alpha: group.alpha
+                    )
+                }
+            } else if let target = mapLoRATarget(rawTarget) {
+                groups[target] = group
+            }
+        }
+        return groups
+    }
+
+    private static func parseLoRAKey(_ rawKey: String) -> (target: String, kind: String)? {
+        var key = rawKey
+        for prefix in ["transformer.", "diffusion_model.", "model.diffusion_model.", "base_model.model."] {
+            if key.hasPrefix(prefix) { key.removeFirst(prefix.count) }
+        }
+
+        let patterns: [(String, String)] = [
+            (".lora_down.weight", "down"),
+            (".lora_A.weight", "down"),
+            (".lora.up.weight", "up"),
+            (".lora_up.weight", "up"),
+            (".lora_B.weight", "up"),
+            (".alpha", "alpha"),
+        ]
+        for (suffix, kind) in patterns where key.hasSuffix(suffix) {
+            let base = normalizeLoRATarget(String(key.dropLast(suffix.count)))
+            return (base, kind)
+        }
+        return nil
+    }
+
+    private static func normalizeLoRATarget(_ rawTarget: String) -> String {
+        var target = rawTarget
+        for prefix in ["model.diffusion_model.", "base_model.model.", "diffusion_model.", "transformer."] {
+            if target.hasPrefix(prefix) {
+                target.removeFirst(prefix.count)
+                break
+            }
+        }
+
+        let doublePrefix = "lora_unet_double_blocks_"
+        let singlePrefix = "lora_unet_single_blocks_"
+        if target.hasPrefix(singlePrefix) {
+            let remainder = String(target.dropFirst(singlePrefix.count))
+            let parts = remainder.split(separator: "_", maxSplits: 1).map(String.init)
+            guard parts.count == 2, Int(parts[0]) != nil else { return target }
+            switch parts[1] {
+            case "linear1": return "single_blocks.\(parts[0]).linear1"
+            case "linear2": return "single_blocks.\(parts[0]).linear2"
+            default: return target
+            }
+        }
+        if target.hasPrefix(doublePrefix) {
+            let remainder = String(target.dropFirst(doublePrefix.count))
+            let parts = remainder.split(separator: "_", maxSplits: 1).map(String.init)
+            guard parts.count == 2, Int(parts[0]) != nil else { return target }
+            let suffix: String
+            switch parts[1] {
+            case "img_attn_qkv": suffix = "img_attn.qkv"
+            case "txt_attn_qkv": suffix = "txt_attn.qkv"
+            case "img_attn_proj": suffix = "img_attn.proj"
+            case "txt_attn_proj": suffix = "txt_attn.proj"
+            case "img_mlp_0": suffix = "img_mlp.0"
+            case "img_mlp_2": suffix = "img_mlp.2"
+            case "txt_mlp_0": suffix = "txt_mlp.0"
+            case "txt_mlp_2": suffix = "txt_mlp.2"
+            default: return target
+            }
+            return "double_blocks.\(parts[0]).\(suffix)"
+        }
+        return target
+    }
+
+    private static func isCombinedQKVLayer(_ target: String) -> Bool {
+        target.contains(".img_attn.qkv") || target.contains(".txt_attn.qkv")
+    }
+
+    private static func qkvTargets(for target: String) -> [String] {
+        let index = extractBlockIndex(from: target, prefix: "double_blocks.")
+        if target.contains(".img_attn.qkv") {
+            return [
+                "transformerBlocks.\(index).attn.toQ",
+                "transformerBlocks.\(index).attn.toK",
+                "transformerBlocks.\(index).attn.toV",
+            ]
+        }
+        if target.contains(".txt_attn.qkv") {
+            return [
+                "transformerBlocks.\(index).attn.addQProj",
+                "transformerBlocks.\(index).attn.addKProj",
+                "transformerBlocks.\(index).attn.addVProj",
+            ]
+        }
+        return []
+    }
+
+    private static func mapLoRATarget(_ key: String) -> String? {
+        let direct = key
+            .replacingOccurrences(of: ".processor", with: "")
+            .replacingOccurrences(of: ".to_q", with: ".toQ")
+            .replacingOccurrences(of: ".to_k", with: ".toK")
+            .replacingOccurrences(of: ".to_v", with: ".toV")
+            .replacingOccurrences(of: ".to_out.0", with: ".toOut")
+            .replacingOccurrences(of: ".add_q_proj", with: ".addQProj")
+            .replacingOccurrences(of: ".add_k_proj", with: ".addKProj")
+            .replacingOccurrences(of: ".add_v_proj", with: ".addVProj")
+            .replacingOccurrences(of: ".to_add_out", with: ".toAddOut")
+        if direct.hasPrefix("transformerBlocks.") || direct.hasPrefix("singleTransformerBlocks.") || direct == "xEmbedder" || direct == "contextEmbedder" || direct == "projOut" {
+            return direct
+        }
+
+        if key.hasPrefix("double_blocks.") {
+            let index = extractBlockIndex(from: key, prefix: "double_blocks.")
+            if key.contains(".img_attn.proj") { return "transformerBlocks.\(index).attn.toOut" }
+            if key.contains(".txt_attn.proj") { return "transformerBlocks.\(index).attn.toAddOut" }
+            if key.contains(".img_mlp.0") { return "transformerBlocks.\(index).ff.activation.proj" }
+            if key.contains(".img_mlp.2") { return "transformerBlocks.\(index).ff.linearOut" }
+            if key.contains(".txt_mlp.0") { return "transformerBlocks.\(index).ffContext.activation.proj" }
+            if key.contains(".txt_mlp.2") { return "transformerBlocks.\(index).ffContext.linearOut" }
+        }
+
+        if key.hasPrefix("single_blocks.") {
+            let index = extractBlockIndex(from: key, prefix: "single_blocks.")
+            if key.contains(".linear1") { return "singleTransformerBlocks.\(index).attn.toQkvMlp" }
+            if key.contains(".linear2") { return "singleTransformerBlocks.\(index).attn.toOut" }
+        }
+
+        if key.hasPrefix("double_blocks.") {
+            let parts = key.split(separator: ".")
+            guard parts.count >= 3, let index = Int(parts[1]) else { return nil }
+            let suffix = parts.dropFirst(2).joined(separator: ".")
+            switch suffix {
+            case "img_to_q": return "transformerBlocks.\(index).attn.toQ"
+            case "img_to_k": return "transformerBlocks.\(index).attn.toK"
+            case "img_to_v": return "transformerBlocks.\(index).attn.toV"
+            case "img_to_out": return "transformerBlocks.\(index).attn.toOut"
+            case "txt_to_q": return "transformerBlocks.\(index).attn.addQProj"
+            case "txt_to_k": return "transformerBlocks.\(index).attn.addKProj"
+            case "txt_to_v": return "transformerBlocks.\(index).attn.addVProj"
+            case "txt_to_out": return "transformerBlocks.\(index).attn.toAddOut"
+            case "img_mlp_in": return "transformerBlocks.\(index).ff.activation.proj"
+            case "img_mlp_out": return "transformerBlocks.\(index).ff.linearOut"
+            case "txt_mlp_in": return "transformerBlocks.\(index).ffContext.activation.proj"
+            case "txt_mlp_out": return "transformerBlocks.\(index).ffContext.linearOut"
+            default: return nil
+            }
+        }
+
+        if key.hasPrefix("single_blocks.") {
+            let parts = key.split(separator: ".")
+            guard parts.count >= 3, let index = Int(parts[1]) else { return nil }
+            let suffix = parts.dropFirst(2).joined(separator: ".")
+            switch suffix {
+            case "to_qkv_mlp", "linear1": return "singleTransformerBlocks.\(index).attn.toQkvMlp"
+            case "to_out", "linear2": return "singleTransformerBlocks.\(index).attn.toOut"
+            default: return nil
+            }
+        }
+
+        if key.hasPrefix("transformer_blocks.") {
+            let parts = key.split(separator: ".")
+            guard parts.count >= 4, let index = Int(parts[1]) else { return nil }
+            let suffix = parts.dropFirst(2).joined(separator: ".")
+            switch suffix {
+            case "attn.to_q": return "transformerBlocks.\(index).attn.toQ"
+            case "attn.to_k": return "transformerBlocks.\(index).attn.toK"
+            case "attn.to_v": return "transformerBlocks.\(index).attn.toV"
+            case "attn.to_out.0", "attn.to_out": return "transformerBlocks.\(index).attn.toOut"
+            case "attn.add_q_proj": return "transformerBlocks.\(index).attn.addQProj"
+            case "attn.add_k_proj": return "transformerBlocks.\(index).attn.addKProj"
+            case "attn.add_v_proj": return "transformerBlocks.\(index).attn.addVProj"
+            case "attn.to_add_out": return "transformerBlocks.\(index).attn.toAddOut"
+            case "ff.net.0.proj", "ff.proj": return "transformerBlocks.\(index).ff.activation.proj"
+            case "ff.net.2", "ff.linear": return "transformerBlocks.\(index).ff.linearOut"
+            case "ff_context.net.0.proj", "ff_context.proj": return "transformerBlocks.\(index).ffContext.activation.proj"
+            case "ff_context.net.2", "ff_context.linear": return "transformerBlocks.\(index).ffContext.linearOut"
+            default: return nil
+            }
+        }
+
+        if key.hasPrefix("single_transformer_blocks.") {
+            let parts = key.split(separator: ".")
+            guard parts.count >= 4, let index = Int(parts[1]) else { return nil }
+            let suffix = parts.dropFirst(2).joined(separator: ".")
+            switch suffix {
+            case "proj_mlp", "attn.to_qkv_mlp", "linear1": return "singleTransformerBlocks.\(index).attn.toQkvMlp"
+            case "proj_out", "attn.to_out", "linear2": return "singleTransformerBlocks.\(index).attn.toOut"
+            default: return nil
+            }
+        }
+
+        switch key {
+        case "img_in", "x_embedder": return "xEmbedder"
+        case "txt_in", "context_embedder": return "contextEmbedder"
+        case "final_layer.linear", "proj_out": return "projOut"
+        default: return nil
+        }
+    }
+
+    private func loadTransformer() throws -> Flux2Transformer2DModel {
+        if let transformer {
+            return transformer
+        }
+
+        let profile = ModelRegistry.TransformerRepo.moxinOrgKlein4B8bit.profile
+        let configURL = paths.transformer.appendingPathComponent("config.json")
+        let config = FileManager.default.fileExists(atPath: configURL.path)
+            ? (try? Flux2TransformerConfig.load(from: configURL)) ?? .klein4B
+            : .klein4B
+
+        let loaded = Flux2Transformer2DModel(
+            config: config,
+            memoryOptimization: .moderate,
+            usesNormOutBias: profile.usesNormOutBias
+        )
+        transformer = loaded
+
+        var weights = try Flux2WeightLoader.loadWeights(from: paths.transformer)
+        try Flux2WeightLoader.applyTransformerWeights(&weights, to: loaded, profile: profile)
+        weights.removeAll()
+        eval(loaded.parameters())
+        MLX.Memory.clearCache()
+
+        return loaded
+    }
+
+    private func loadVAE() throws -> AutoencoderKLFlux2 {
+        if let vae {
+            return vae
+        }
+
+        let configURL = paths.vae.appendingPathComponent("config.json")
+        let config = FileManager.default.fileExists(atPath: configURL.path)
+            ? try VAEConfig.load(from: configURL)
+            : .flux2Dev
+
+        let loaded = AutoencoderKLFlux2(config: config)
+        vae = loaded
+
+        let weights = try Flux2WeightLoader.loadWeights(from: paths.vae)
+        try Flux2WeightLoader.applyVAEWeights(weights, to: loaded)
+        eval(loaded.parameters())
+        MLX.Memory.clearCache()
+
+        return loaded
+    }
+}