Stepwise images often don’t provide a useful preview of intermediate states #538
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.
All reactions
Replies: 6 comments 1 reply
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.
All reactions
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?
All reactions
@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.
All reactions
- 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.
All reactions
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.
All reactions
-
👍 1
correction: TAEx - ie Tiny AutoEncoder
All reactions
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.