Stepwise images often don’t provide a useful preview of intermediate states

[fortinmike]

I'm working on a nice image gen app for Mac that uses mflux as its main (and only, for now) backend. To be published as open source soon-ish! Works great so far, but I need good mid-generation previews so that the user can have an idea of how their image will turn out and may cancel if a character has one or two extra hands (ha!)

Trying out the experimental stepwise image previews I notice that many models don't produce a useful preview of the image-to-be until very late in the steps, if at all. Compared to something like ComfyUI Latent2RGB this is sorely lacking. Otherwise, mflux is surprisingly capable as a generation backend for an app.

I assume this is a known limitation, but I wanted to create an issue to track this, and if possible get some feedback from @filipstrand WRT if enhancing this is in the short-to-medium term pipeline, and what amount of work it represents.

[ianscrivener]

TAEx is absolutely what you want - see https://github.com/IonDen/mlx-taef. Denis' work is very solid.

Inherently, a 4 or 8 step inference is going to come with a big reveal at the end... you could 'fake' it with a multistep TAEx & show incremental images - that may be better UX than a spinner. TAEx runs on CPU, so doesn't slow the real inference process.

[fortinmike]

Interesting! Thanks for the tip.

At the moment for my MVP I'm simply delegating to the mflux CLI, but I had a hunch I'd need to switch to the Python API sooner than later. Looks like using TAESx would require that.

Makes me wonder:

• How much of a workaround / how fragile is integrating this vs mflux implementing a similar / official integration of TAESx right into the base tool (including for CLI good stepwise preview images by default)?
• Does it support all model types?
• Would it be prone to breaking when mflux itself changes, or is the extension point stable at this point?

[IonDen]

@fortinmike, author of mlx-taef here, late to the thread. Your three questions:

How fragile is it? It is a plain mflux callback, the same in-loop hook mflux's own stepwise image saver uses. Nothing is patched. It broke once: mflux 0.19.0 added a keyword to the before-loop hook and the callback did not accept it. Fixed in mlx-taef 0.8.1, and there is now a test that runs mflux's real dispatcher, so the next new keyword fails my CI instead of your app.

Which models? FLUX.1, FLUX.2 Klein, Z-Image and Z-Image Turbo, Qwen-Image and Qwen-Image-Edit, Krea 2. Lens shares Klein's latent layout and should work, but I have not validated it.

Is the extension point stable? The callback protocol has kept the same shape since mflux 0.17, and mflux relies on it internally. One caveat: each new conditioning type adds a keyword, so accept **kwargs in your own callbacks.

An official integration would be small: register the callback behind a --live-preview flag. I can write that up as a Discussion if the maintainers want it. Decode cost is about 30 ms per step; per-model preview quality is in the README gallery.

[ianscrivener]

• MFlux doesn't have an 'official' TAEx - though this one is listed on the main README.
• Like I said, mlx-taef is solid - certainly the best TAEx option for MLX.
• You could use a pytorch version of TAEx, though this would also require python.
• My guess is you'll hit many blockers just wrapping CLIs - depends how rich a GUI you are wanting.

[fortinmike]

Yep, wrapping the CLI was not the long-term plan, especially because I plan to integrate mlx-teacache at some point. I would have preferred to keep the CLI approach a bit longer for the MVP though.

[ianscrivener]

correction: TAEx - ie Tiny AutoEncoder

[anthonywu]

I built the original stepwise image series generator, the original intent was a debug tool to make sure img2img actually did the prefills of the initial latents. The stepwise generation added significant wall time to get to the image output so I assumed it was only useful as a debugger to fix the algo or to find a optimal steps value.