Skip to content

Diffusion pipeline latent preview - #236

Merged
stikves merged 4 commits into
apple:mainfrom
stikves:sukru/latent-preview
Sep 11, 2026
Merged

stikves merged 4 commits into
apple:mainfrom
stikves:sukru/latent-preview

Conversation

@stikves

@stikves stikves commented Sep 5, 2026 •

Copy link
Copy Markdown
Contributor

Diffusion pipeline latent preview

Surfaces intermediate denoising latents so a UI can render a lightweight live preview without a full VAE decode at every step.

  • PipelineProgress.currentLatent — the current latent is exposed through an optional progress handler; when no handler is set, no preview work is performed.
  • NDArray.asRGB(coefficients:) — projects a latent to an RGB image via a single cblas_sgemm. The projection is a small linear map specific to a model's VAE, supplied by the caller.
  • LatentPreviewTuner and diffusion-runner --tune-preview / --tune-fit — collect latent/image pairs for a model, fit a coefficient set, and export per-step preview images.

Coefficients are fit per model rather than bundled, since the latent-to-RGB mapping depends on the model's VAE. See the LatentRGBCoefficients documentation for the tuning workflow.

Closes: #231

Follow-up (out of scope)

Preview resolution equals the latent resolution. A higher-fidelity option is a small distilled decoder that reconstructs full-resolution output at a fraction of the full VAE cost; tracked separately.

Test plan

  • CI green
  • --tune-preview collects pairs; --tune-fit fits coefficients and exports previews
  • Live preview renders with a fitted coefficient set

@stikves
stikves force-pushed the sukru/latent-preview branch 5 times, most recently from 184015d to f2672a0 Compare September 9, 2026 21:12
Surface intermediate denoising latents so a UI can render a cheap live
preview without a full VAE decode per step:

- PipelineProgress.currentLatent, plumbed through an optional progress
  handler (zero cost when unset) across the SD, SD3, and FLUX.2 pipelines.
- NDArray.asRGB(coefficients:) — a single cblas_sgemm latent-to-RGB
  projection. Coefficients are supplied by the caller; the mapping is
  specific to each model's VAE.
- LatentPreviewTuner and diffusion-runner --tune-preview / --tune-fit to
  collect latent/image pairs, fit a coefficient set, and export per-step
  draft previews.
@stikves
stikves force-pushed the sukru/latent-preview branch from f2672a0 to e65cee3 Compare September 9, 2026 21:16
@stikves
stikves marked this pull request as ready for review September 9, 2026 21:54
@stikves stikves self-assigned this Sep 9, 2026
Comment thread .gitignore Outdated
Comment thread swift/Sources/CoreAIDiffusionPipeline/Pipelines/Flux2Pipeline.swift Outdated
Comment thread swift/Sources/CoreAIDiffusionPipeline/Pipelines/Flux2Pipeline.swift Outdated
Comment thread swift/Sources/CoreAIDiffusionPipeline/Pipelines/Pipeline.swift
Comment thread swift/Sources/CoreAIDiffusionPipeline/Preview/LatentPreview.swift Outdated
Comment thread swift/Sources/CoreAIDiffusionPipeline/Preview/LatentPreview.swift Outdated
Comment thread swift/Sources/CoreAIDiffusionPipeline/Preview/LatentPreviewTuner.swift Outdated
Comment thread swift/Sources/CoreAIDiffusionPipeline/Preview/LatentPreviewTuner.swift Outdated
Comment thread swift/Sources/CoreAIDiffusionPipeline/Preview/LatentPreviewTuner.swift Outdated
Comment thread swift/Sources/CoreAIDiffusionPipeline/Preview/LatentPreviewTuner.swift Outdated
- Flux2Pipeline: fix preview comment typo (16 -> 32 channels after unpatchify)
- LatentPreview: correct vDSP -> cblas_sgemm/BLAS docs; add rgbChannels constant
- LatentPreviewTuner: extract shared solveLeastSquares helper used by both
  fitCoefficients and fitFromDirectory, removing duplicated normal-equations math
- LatentPreviewTuner: report full shape on fit-shape guard; drop thin savePNG
  wrapper; log skipped latents in exportStepPreviewsFromDirectory

@kevchengcodes kevchengcodes left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

looks good, left one more comment

@stikves
stikves merged commit 7304c47 into apple:main Sep 11, 2026
3 checks passed
@stikves
stikves deleted the sukru/latent-preview branch September 11, 2026 00:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

PipelineProgress.currentLatent is never populated, so consumers cannot show a live preview

3 participants