diff --git a/swift/Sources/CoreAIDiffusionPipeline/Pipelines/Flux2Pipeline.swift b/swift/Sources/CoreAIDiffusionPipeline/Pipelines/Flux2Pipeline.swift index d292bbd8..246bf6f6 100644 --- a/swift/Sources/CoreAIDiffusionPipeline/Pipelines/Flux2Pipeline.swift +++ b/swift/Sources/CoreAIDiffusionPipeline/Pipelines/Flux2Pipeline.swift @@ -420,14 +420,32 @@ public struct Flux2Pipeline: DiffusionPipeline { } } + // Capture the denoising state BEFORE the scheduler advances so the + // preview below can form the x0 estimate. + let previewSigma = scheduler.currentSigma + let sampleBeforeStep = packedLatents packedLatents = scheduler.step(output: output, timeStep: t, sample: packedLatents) try checkLatentsAreFinite(packedLatents, step: step) if let progressHandler { + // Preview the DENOISED estimate, not the raw post-step sample. The + // sample after the Euler step is still mostly noise until the last + // step or two, so on a few-step model (e.g. FLUX.2 Klein at 4 steps) + // the early previews look like static. Flow-matching gives the + // estimate for one multiply-add: with x_t = (1-σ)·x0 + σ·ε and the + // model predicting v = ε - x0, x0 = x_t - σ·v. Blurry on step one, + // but it shows the composition and converges to the final image. + var previewPacked = sampleBeforeStep + if previewSigma > 0 { + var negSigma = -previewSigma + vDSP_vsma( + output, 1, &negSigma, sampleBeforeStep, 1, &previewPacked, 1, + vDSP_Length(output.count)) + } // Unpack → denorm → unpatchify: [1, 128, 64, 64] → [1, 32, 128, 128] // These are array copies, no model call. let spatial = unpackLatentsSpatialFlatten( - packedLatents, channels: inChannels, height: spatialSide, width: spatialSide) + previewPacked, channels: inChannels, height: spatialSide, width: spatialSide) let denormed = applyBatchNormDenorm( spatial, channels: inChannels, height: spatialSide, width: spatialSide) let unpatchified = Self.unpatchifyLatents( diff --git a/swift/Sources/CoreAIDiffusionPipeline/Schedulers/DiscreteFlowScheduler.swift b/swift/Sources/CoreAIDiffusionPipeline/Schedulers/DiscreteFlowScheduler.swift index feb96cb4..57b7257c 100644 --- a/swift/Sources/CoreAIDiffusionPipeline/Schedulers/DiscreteFlowScheduler.swift +++ b/swift/Sources/CoreAIDiffusionPipeline/Schedulers/DiscreteFlowScheduler.swift @@ -15,6 +15,11 @@ public final class DiscreteFlowScheduler { /// The first scheduled sigma after all shifts are applied — use this for img2img noise addition. public var startSigma: Float { sigmas.first ?? 1.0 } + /// Sigma for the step the NEXT `step(...)` call will consume — i.e. the current + /// point on the noise schedule, before advancing. Used to form the denoised x0 + /// estimate for live previews: `x0 = sample − σ·v`. + public var currentSigma: Float { counter < sigmas.count ? sigmas[counter] : 0 } + let trainSteps: Float let shift: Float let mu: Float?