From 010ca03c262c20a7bce0cf3cc30491c3dd86d876 Mon Sep 17 00:00:00 2001 From: Ronald Mannak Date: Sat, 7 Feb 2026 20:27:27 -0800 Subject: [PATCH] Cleanup: remove deprecated imports, add LocalizedError, use async URLSession - Remove standalone `import MLXRandom` and `import MLXFast` where namespaced calls work via `import MLX` - Add LocalizedError conformance with errorDescription to all public error enums for better consumer diagnostics - Replace DispatchSemaphore + callback-based URLSession.dataTask with modern async URLSession.data(for:) in fetchRemoteData() - Hoist async image loading out of Device.withDefaultDevice closure Co-Authored-By: Claude Opus 4.6 --- .../TextEncoder/Flux2PixtralVisionTower.swift | 1 - Sources/Flux2/Models/VAE/Flux2VAEBlocks.swift | 1 - Sources/Flux2/Pipeline/Flux2DevPipeline.swift | 17 +++- .../Flux2/Pipeline/Flux2KleinPipeline.swift | 25 ++++- .../Pipeline/Flux2LatentPreparation.swift | 14 ++- .../FlowMatchEulerDiscreteScheduler.swift | 18 +++- Sources/Flux2CLI/CLI+Generate.swift | 96 ++++++++++--------- Sources/Flux2CLI/CLI+Image.swift | 80 +++++----------- 8 files changed, 143 insertions(+), 109 deletions(-) diff --git a/Sources/Flux2/Models/TextEncoder/Flux2PixtralVisionTower.swift b/Sources/Flux2/Models/TextEncoder/Flux2PixtralVisionTower.swift index 25b8e33..0fa8640 100644 --- a/Sources/Flux2/Models/TextEncoder/Flux2PixtralVisionTower.swift +++ b/Sources/Flux2/Models/TextEncoder/Flux2PixtralVisionTower.swift @@ -1,6 +1,5 @@ import Foundation import MLX -import MLXFast import MLXNN final class Flux2PixtralVisionTower: Module { diff --git a/Sources/Flux2/Models/VAE/Flux2VAEBlocks.swift b/Sources/Flux2/Models/VAE/Flux2VAEBlocks.swift index a36b4ea..904c33e 100644 --- a/Sources/Flux2/Models/VAE/Flux2VAEBlocks.swift +++ b/Sources/Flux2/Models/VAE/Flux2VAEBlocks.swift @@ -1,6 +1,5 @@ import Foundation import MLX -import MLXFast import MLXNN final class Flux2VAESelfAttention: Module { diff --git a/Sources/Flux2/Pipeline/Flux2DevPipeline.swift b/Sources/Flux2/Pipeline/Flux2DevPipeline.swift index 42bcf6b..b08651b 100644 --- a/Sources/Flux2/Pipeline/Flux2DevPipeline.swift +++ b/Sources/Flux2/Pipeline/Flux2DevPipeline.swift @@ -12,12 +12,27 @@ public struct Flux2DevPipelineOutput { public let imageLatentIds: MLXArray? } -public enum Flux2DevPipelineError: Error { +public enum Flux2DevPipelineError: Error, LocalizedError { case promptEncoderReleased case missingProcessor case invalidLatentChannels(Int, Int) case invalidNumInferenceSteps(Int) case invalidImageCount(Int) + + public var errorDescription: String? { + switch self { + case .promptEncoderReleased: + return "Prompt encoder has been released and is no longer available." + case .missingProcessor: + return "Processor is required but was not loaded." + case .invalidLatentChannels(let inChannels, let patchArea): + return "Latent channels (\(inChannels)) must be divisible by patch area (\(patchArea))." + case .invalidNumInferenceSteps(let steps): + return "Number of inference steps must be positive, got \(steps)." + case .invalidImageCount(let count): + return "Image list must not be empty, got \(count) images." + } + } } public final class Flux2DevPipeline { diff --git a/Sources/Flux2/Pipeline/Flux2KleinPipeline.swift b/Sources/Flux2/Pipeline/Flux2KleinPipeline.swift index c6c9bbb..1681920 100644 --- a/Sources/Flux2/Pipeline/Flux2KleinPipeline.swift +++ b/Sources/Flux2/Pipeline/Flux2KleinPipeline.swift @@ -12,7 +12,7 @@ public struct Flux2KleinPipelineOutput { public let imageLatentIds: MLXArray? } -public enum Flux2KleinPipelineError: Error { +public enum Flux2KleinPipelineError: Error, LocalizedError { case promptEncoderReleased case missingTokenizer case missingNegativeTokens @@ -22,6 +22,29 @@ public enum Flux2KleinPipelineError: Error { case invalidLatentChannels(Int, Int) case invalidNumInferenceSteps(Int) case invalidImageCount(Int) + + public var errorDescription: String? { + switch self { + case .promptEncoderReleased: + return "Prompt encoder has been released and is no longer available." + case .missingTokenizer: + return "Tokenizer is required but was not loaded." + case .missingNegativeTokens: + return "Negative tokens are required for classifier-free guidance but were not provided." + case .invalidNegativeTokens: + return "Both negative input IDs and attention mask must be provided together." + case .negativeInputIdsShapeMismatch(let expected, let got): + return "Negative input IDs shape \(got) must match positive input IDs shape \(expected)." + case .negativeAttentionMaskShapeMismatch(let expected, let got): + return "Negative attention mask shape \(got) must match positive attention mask shape \(expected)." + case .invalidLatentChannels(let inChannels, let patchArea): + return "Latent channels (\(inChannels)) must be divisible by patch area (\(patchArea))." + case .invalidNumInferenceSteps(let steps): + return "Number of inference steps must be positive, got \(steps)." + case .invalidImageCount(let count): + return "Image list must not be empty, got \(count) images." + } + } } public final class Flux2KleinPipeline { diff --git a/Sources/Flux2/Pipeline/Flux2LatentPreparation.swift b/Sources/Flux2/Pipeline/Flux2LatentPreparation.swift index 174f9fe..52b5859 100644 --- a/Sources/Flux2/Pipeline/Flux2LatentPreparation.swift +++ b/Sources/Flux2/Pipeline/Flux2LatentPreparation.swift @@ -1,6 +1,5 @@ import Foundation import MLX -import MLXRandom public struct Flux2PreparedLatents { public let latents: MLXArray @@ -14,10 +13,21 @@ public struct Flux2PreparedImageLatents { public let ids: MLXArray } -public enum Flux2LatentPreparationError: Error { +public enum Flux2LatentPreparationError: Error, LocalizedError { case failedToCreateLatents case emptyImages case invalidImageShape(index: Int, shape: [Int]) + + public var errorDescription: String? { + switch self { + case .failedToCreateLatents: + return "Failed to create latent tensor." + case .emptyImages: + return "Image list must not be empty." + case .invalidImageShape(let index, let shape): + return "Image at index \(index) has invalid shape \(shape); expected 4 dimensions." + } + } } public enum Flux2LatentPreparation { diff --git a/Sources/Flux2/Schedulers/FlowMatchEulerDiscreteScheduler.swift b/Sources/Flux2/Schedulers/FlowMatchEulerDiscreteScheduler.swift index 9c471c6..f9c0e7a 100644 --- a/Sources/Flux2/Schedulers/FlowMatchEulerDiscreteScheduler.swift +++ b/Sources/Flux2/Schedulers/FlowMatchEulerDiscreteScheduler.swift @@ -1,17 +1,31 @@ import Foundation import MLX -import MLXRandom public struct FlowMatchEulerDiscreteSchedulerOutput { public let prevSample: MLXArray } -public enum FlowMatchEulerDiscreteSchedulerError: Error { +public enum FlowMatchEulerDiscreteSchedulerError: Error, LocalizedError { case configNotFound(URL) case invalidTimeShiftType(String) case betaSigmasUnsupported case missingDynamicShiftMu case invalidSchedule(String) + + public var errorDescription: String? { + switch self { + case .configNotFound(let url): + return "Scheduler configuration not found at \(url.path)." + case .invalidTimeShiftType(let type): + return "Invalid time shift type '\(type)'; expected 'exponential' or 'linear'." + case .betaSigmasUnsupported: + return "Beta sigmas are not supported." + case .missingDynamicShiftMu: + return "Dynamic shifting requires a mu value." + case .invalidSchedule(let reason): + return "Invalid schedule: \(reason)." + } + } } public final class FlowMatchEulerDiscreteScheduler { diff --git a/Sources/Flux2CLI/CLI+Generate.swift b/Sources/Flux2CLI/CLI+Generate.swift index e1d74f2..bc066a9 100644 --- a/Sources/Flux2CLI/CLI+Generate.swift +++ b/Sources/Flux2CLI/CLI+Generate.swift @@ -196,32 +196,36 @@ extension CLI { } let dtype = resolvedDType + let clock = ContinuousClock() - try Device.withDefaultDevice(.gpu) { - let clock = ContinuousClock() + guard let prompt = options.prompt else { + throw CLIError.missingArgument("--prompt") + } - guard let prompt = options.prompt else { - throw CLIError.missingArgument("--prompt") + var height = options.height + var width = options.width + let imageSpecs = options.imageSpecs + let conditioning: [ConditioningImage] + if imageSpecs.isEmpty { + conditioning = [] + } else { + let loadStart = clock.now + var loaded: [ConditioningImage] = [] + for spec in imageSpecs { + try await loaded.append(loadConditioningImage(spec: spec)) } + conditioning = loaded + stageTimes["conditioning_load_s"] = seconds(clock.now - loadStart) + height = height ?? conditioning.first?.height + width = width ?? conditioning.first?.width + } + let conditioningImages = conditioning.isEmpty ? nil : conditioning.map(\.array) - var height = options.height - var width = options.width - let steps = options.steps ?? 50 - let guidanceScale = options.guidanceScale - let imageIdScale = options.imageIdScale - let imageSpecs = options.imageSpecs - let conditioning: [ConditioningImage] - if imageSpecs.isEmpty { - conditioning = [] - } else { - let loadStart = clock.now - conditioning = try imageSpecs.map { try loadConditioningImage(spec: $0) } - stageTimes["conditioning_load_s"] = seconds(clock.now - loadStart) - height = height ?? conditioning.first?.height - width = width ?? conditioning.first?.width - } - let conditioningImages = conditioning.isEmpty ? nil : conditioning.map(\.array) + let steps = options.steps ?? 50 + let guidanceScale = options.guidanceScale + let imageIdScale = options.imageIdScale + try Device.withDefaultDevice(.gpu) { let initStart = clock.now let pipeline = try Flux2KleinPipeline( snapshot: snapshotURL, @@ -271,34 +275,38 @@ extension CLI { } let dtype = resolvedDType + let clock = ContinuousClock() - try Device.withDefaultDevice(.gpu) { - let clock = ContinuousClock() + guard let prompt = options.prompt else { + throw CLIError.missingArgument("--prompt") + } - guard let prompt = options.prompt else { - throw CLIError.missingArgument("--prompt") + var height = options.height + var width = options.width + let imageSpecs = options.imageSpecs + let conditioning: [ConditioningImage] + if imageSpecs.isEmpty { + conditioning = [] + } else { + let loadStart = clock.now + var loaded: [ConditioningImage] = [] + for spec in imageSpecs { + try await loaded.append(loadConditioningImage(spec: spec)) } + conditioning = loaded + stageTimes["conditioning_load_s"] = seconds(clock.now - loadStart) + height = height ?? conditioning.first?.height + width = width ?? conditioning.first?.width + } + let conditioningImages = conditioning.isEmpty ? nil : conditioning.map(\.array) + let upsampleImages = conditioning.isEmpty ? nil : conditioning.map(\.original) - var height = options.height - var width = options.width - let steps = options.steps ?? 50 - let guidanceScale = options.guidanceScale - let imageIdScale = options.imageIdScale - let maxLength = options.maxLength - let imageSpecs = options.imageSpecs - let conditioning: [ConditioningImage] - if imageSpecs.isEmpty { - conditioning = [] - } else { - let loadStart = clock.now - conditioning = try imageSpecs.map { try loadConditioningImage(spec: $0) } - stageTimes["conditioning_load_s"] = seconds(clock.now - loadStart) - height = height ?? conditioning.first?.height - width = width ?? conditioning.first?.width - } - let conditioningImages = conditioning.isEmpty ? nil : conditioning.map(\.array) - let upsampleImages = conditioning.isEmpty ? nil : conditioning.map(\.original) + let steps = options.steps ?? 50 + let guidanceScale = options.guidanceScale + let imageIdScale = options.imageIdScale + let maxLength = options.maxLength + try Device.withDefaultDevice(.gpu) { let initStart = clock.now let pipeline = try Flux2DevPipeline( snapshot: snapshotURL, diff --git a/Sources/Flux2CLI/CLI+Image.swift b/Sources/Flux2CLI/CLI+Image.swift index eb5aab4..5eb6de0 100644 --- a/Sources/Flux2CLI/CLI+Image.swift +++ b/Sources/Flux2CLI/CLI+Image.swift @@ -21,7 +21,7 @@ extension CLI { _ spec: String, remoteTimeout: TimeInterval = 60, remoteMaximumBytes: Int = 50 * 1024 * 1024 - ) throws -> Data { + ) async throws -> Data { let trimmed = spec.trimmingCharacters(in: .whitespacesAndNewlines) guard !trimmed.isEmpty else { throw CLIError.invalidOption("Empty image path") @@ -32,7 +32,7 @@ extension CLI { guard let url = URL(string: trimmed) else { throw CLIError.invalidOption("Invalid image URL: \(trimmed)") } - return try fetchRemoteData(url: url, timeout: remoteTimeout, maxBytes: remoteMaximumBytes) + return try await fetchRemoteData(url: url, timeout: remoteTimeout, maxBytes: remoteMaximumBytes) } if lower.hasPrefix("file://") { guard let url = URL(string: trimmed) else { @@ -48,7 +48,7 @@ extension CLI { return try Data(contentsOf: url) } - private static func fetchRemoteData(url: URL, timeout: TimeInterval, maxBytes: Int) throws -> Data { + private static func fetchRemoteData(url: URL, timeout: TimeInterval, maxBytes: Int) async throws -> Data { let config = URLSessionConfiguration.ephemeral config.timeoutIntervalForRequest = timeout config.timeoutIntervalForResource = timeout @@ -58,72 +58,38 @@ extension CLI { var request = URLRequest(url: url) request.timeoutInterval = timeout - let semaphore = DispatchSemaphore(value: 0) - var result: Result? + let (data, response) = try await session.data(for: request) - let task = session.dataTask(with: request) { data, response, error in - defer { semaphore.signal() } - - if let error { - result = .failure(CLIError.invalidOption("Failed to download image: \(url.absoluteString) (\(error.localizedDescription))")) - return - } - - guard let http = response as? HTTPURLResponse else { - result = .failure(CLIError.invalidOption("Invalid response while downloading image: \(url.absoluteString)")) - return - } - - guard (200..<300).contains(http.statusCode) else { - result = .failure(CLIError.invalidOption("Failed to download image: \(url.absoluteString) (HTTP \(http.statusCode))")) - return - } - - if http.expectedContentLength > 0, http.expectedContentLength > Int64(maxBytes) { - result = .failure( - CLIError.invalidOption( - "Remote image too large: \(url.absoluteString) (expected \(http.expectedContentLength) bytes > max \(maxBytes))" - ) - ) - return - } - - let resolved = data ?? Data() - if resolved.count > maxBytes { - result = .failure( - CLIError.invalidOption( - "Remote image too large: \(url.absoluteString) (downloaded \(resolved.count) bytes > max \(maxBytes))" - ) - ) - return - } + guard let http = response as? HTTPURLResponse else { + throw CLIError.invalidOption("Invalid response while downloading image: \(url.absoluteString)") + } - result = .success(resolved) + guard (200..<300).contains(http.statusCode) else { + throw CLIError.invalidOption("Failed to download image: \(url.absoluteString) (HTTP \(http.statusCode))") } - task.resume() - if semaphore.wait(timeout: .now() + timeout + 5) == .timedOut { - task.cancel() - throw CLIError.invalidOption("Timed out downloading image: \(url.absoluteString)") + if http.expectedContentLength > 0, http.expectedContentLength > Int64(maxBytes) { + throw CLIError.invalidOption( + "Remote image too large: \(url.absoluteString) (expected \(http.expectedContentLength) bytes > max \(maxBytes))" + ) } - switch result { - case .success(let data): - return data - case .failure(let error): - throw error - case .none: - throw CLIError.invalidOption("Invalid response while downloading image: \(url.absoluteString)") + if data.count > maxBytes { + throw CLIError.invalidOption( + "Remote image too large: \(url.absoluteString) (downloaded \(data.count) bytes > max \(maxBytes))" + ) } + + return data } - static func loadImage(spec: String, height: Int, width: Int) throws -> MLXArray { + static func loadImage(spec: String, height: Int, width: Int) async throws -> MLXArray { guard height > 0, width > 0 else { throw CLIError.invalidOption("Invalid image size \(width)x\(height)") } let trimmed = spec.trimmingCharacters(in: .whitespacesAndNewlines) - let data = try loadDataFromSpec(trimmed) + let data = try await loadDataFromSpec(trimmed) let cgImage = try decodeCGImage(data: data, spec: trimmed) let rgba = try renderRGBA(image: cgImage, width: width, height: height, resizeMode: .crop) @@ -151,7 +117,7 @@ extension CLI { spec: String, maxArea: Int = 1024 * 1024, multipleOf: Int = 16 - ) throws -> ConditioningImage { + ) async throws -> ConditioningImage { guard maxArea > 0 else { throw CLIError.invalidOption("Invalid maxArea: \(maxArea)") } @@ -160,7 +126,7 @@ extension CLI { } let trimmed = spec.trimmingCharacters(in: .whitespacesAndNewlines) - let data = try loadDataFromSpec(trimmed) + let data = try await loadDataFromSpec(trimmed) var cgImage = try decodeCGImage(data: data, spec: trimmed) let originalImage = cgImage