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

[james-333i]

What are you trying to build?

PipelineProgress has carried a currentLatent field since it was introduced, but Flux2Pipeline, SD3Pipeline and StableDiffusionPipeline all pass nil. Apps that want to show the image forming during generation have nothing to render.

Proposed change for FLUX.2, on my fork: main...james-333i:coreai-models:flux-preview-latents

After each Euler step it computes the flow-matching x0 prediction (x0 = x_t minus sigma times v), converts it to VAE latent space (unpack, batch-norm denorm, unpatchify) and hands it to the progress handler as a [1, 32, H/8, W/8] array. The final step is skipped since the decoded image follows immediately. The cost is a few array copies per step, negligible next to a transformer pass.

On the consumer side a cheap channel-to-RGB projection is enough for a recognizable preview. I fitted a 32x3 linear map against the FLUX.2 VAE by encoding a set of images and regressing against their downscaled RGB, and the previews show subject, color and composition clearly at every step. I can include those factors in a follow-up if a preview helper is wanted in the package itself. Glad to send this as a PR if PR creation opens up.

Where are the current docs or utilities unclear?

N/A

Expected improvement

Allows consumer to see progress on image generation steps.

Additional context

No response

[stikves]

Thanks for the report.

The progress is an expensive operation. However we can present the current latents and the RGB image on demand.

[james-333i]

@stikves Agreed, it should be designed to cost nothing unless asked for. Two options that fit and happy to rework the branch to either:

1. Opt-in on the configuration, for example PipelineConfiguration.emitsPreviewLatents, defaulting to off. When set, the loop computes the x0 prediction per step (the cost is only the unpack, denorm and unpatchify copies; no decode) and passes it in currentLatent.

2. On demand, with PipelineProgress carrying closures such as previewLatent() and previewImage() instead of a value, so nothing is computed unless a consumer calls them from the progress handler.

For the RGB image I would suggest a cheap linear channel-to-RGB projection rather than a VAE decode per step; I fitted a 32x3 map against the FLUX.2 VAE that gives recognizable previews at every step, and can contribute those factors alongside. If you have a preference between the two options I can update the branch to match.

[james-333i]

@stikves Thanks for adding this. One suggestion on what the handler receives.

currentLatent is the sample after the scheduler step, so it is still mostly noise until the last step or two. On a four step model like FLUX.2 Klein, the first three previews look like static. The denoised estimate is available at the same point for one multiply-add. With x_t = (1 - σ)·x0 + σ·ε and the model predicting v = ε - x0, the estimate is x0 = x_t - σ·v. Both arrays are already in hand after scheduler.step, so the extra cost is a vDSP_vsma over the packed latent before the same unpack, denorm and unpatchify you do now. The estimate is blurry on step one but shows the composition, and it converges to the same result on the final step.

I ran previews this way against the FLUX.2 VAE with a fitted [32 x 3] projection and got a recognizable image from the first step. Happy to send a small PR that emits the estimate, or emits both if you would rather keep the raw sample available.

[stikves]

Thanks for the follow up.

We have yet to open for public PRs, but would this work for you?

#254