Skip to content

Cleanup: remove deprecated imports, add LocalizedError, use async URLSession #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mzbac merged 1 commit into from
Feb 11, 2026
Merged

Cleanup: remove deprecated imports, add LocalizedError, use async URLSession #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion Sources/Flux2/Models/TextEncoder/Flux2PixtralVisionTower.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import Foundation
import MLX
import MLXFast
import MLXNN

final class Flux2PixtralVisionTower: Module {
Expand Down
1 change: 0 additions & 1 deletion Sources/Flux2/Models/VAE/Flux2VAEBlocks.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import Foundation
import MLX
import MLXFast
import MLXNN

final class Flux2VAESelfAttention: Module {
Expand Down
17 changes: 16 additions & 1 deletion Sources/Flux2/Pipeline/Flux2DevPipeline.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
25 changes: 24 additions & 1 deletion Sources/Flux2/Pipeline/Flux2KleinPipeline.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 {
Expand Down
14 changes: 12 additions & 2 deletions Sources/Flux2/Pipeline/Flux2LatentPreparation.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import Foundation
import MLX
import MLXRandom

public struct Flux2PreparedLatents {
public let latents: MLXArray
Expand All @@ -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 {
Expand Down
18 changes: 16 additions & 2 deletions Sources/Flux2/Schedulers/FlowMatchEulerDiscreteScheduler.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 {
Expand Down
96 changes: 52 additions & 44 deletions Sources/Flux2CLI/CLI+Generate.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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))
}
Comment on lines +291 to 295

@chatgpt-codex-connector chatgpt-codex-connector Bot Feb 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Keep conditioning MLXArrays on GPU to avoid copies

Here loadConditioningImage is awaited before entering Device.withDefaultDevice(.gpu). That helper builds an MLXArray from the decoded image (see CLI+Image.loadConditioningImage), so this change creates the conditioning tensors on the default device (typically CPU) instead of GPU. When the pipeline later consumes conditioningImages inside the GPU context, MLX has to transfer those arrays, which is an unintended host↔device copy and a measurable performance regression for image conditioning. Consider moving the loading (or at least the MLXArray creation) back inside the GPU block or explicitly placing the array on the GPU.

Useful? React with 👍 / 👎.

@ronaldmannak ronaldmannak Feb 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This suggestion is incorrect for MLX on Apple Silicon.

Apple Silicon uses unified memory (UMA) — the GPU and CPU share the same physical memory space. There is no separate "host" and "device" memory, so there is no "host↔device copy" like you'd see with CUDA on discrete NVIDIA GPUs. An MLXArray created outside Device.withDefaultDevice(.gpu) is already accessible to the GPU without any data transfer.

Device.withDefaultDevice(.gpu) sets the compute context (where operations execute), not where data lives. The conditioning image arrays are read directly by GPU kernels inside the pipeline with zero copy overhead.

Additionally, loadConditioningImage is async (it uses URLSession.data), so it cannot be called inside Device.withDefaultDevice(.gpu), which takes a synchronous closure. The current code correctly hoists the async work before the device block — this is the intended pattern.

@ronaldmannak ronaldmannak Feb 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

human Ronald here. Can I trigger your @codex as well?

@chatgpt-codex-connector chatgpt-codex-connector Bot Feb 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To use Codex here, create a Codex account and connect to github.

@ronaldmannak ronaldmannak Feb 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see. I need to bring my own Codex? That's going to be fun. We'll have three bots talking to each other

@mzbac mzbac Feb 10, 2026

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will pull down the PR on my local machine and testing it late tonight, currently the Mac is occupied by porting glm ocr:)

@ronaldmannak ronaldmannak Feb 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@codex look at the previous comments and give me your take

@ronaldmannak ronaldmannak Feb 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I just connected Claude to Github too. Let's see if I can summon @claude to turn this discussion into a Furby Frenzy

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,
Expand Down
Loading