Possible Optimizations

[azrahello]

Hi, would it be possible to implement an upscaler? Either using ControlNet like this: https://huggingface.co/jasperai/Flux.1-dev-Controlnet-Upscaler or something similar to https://github.com/ssitu/ComfyUI_UltimateSDUpscale? I find the ability to upscale images very useful. I was wondering if there’s a chance to implement it.
P.S. If you need a tester for your projects, I’m always available! 😉
Edit.

Yesterday, I tested this paper and this repository in ComfyUI using FLUX1 dev in 8-bit GGUF. By following this process, I was able to almost halve the number of iterations and the time required, effectively making it faster than MLX. I was wondering if this could be useful and if it could be implemented. It would be amazing to generate images with FLUX in almost half the traditional time!

[filipstrand]

Hi @azrahello! Upscaling is interesting for sure and should be something we should support! Thanks for linking the ControlNet one, I will look into this to get a sense of how much work it would take to implement... Depending on a few personal factors, I might be able to spend more time than typical on MFLUX over the coming time so I am really happy with your feature suggestions here :)

TeaCache looks awesome, definitely worth investigating! On their site, they have an example to get a sense of the quality degradation, and that looks really promising.

Would you say with your own testing that you got similary good results, or do they vary a lot? Anyway, really cool!

[azrahello]

this is the output of comfy without teacache,
Requested to load Flux
100%|███████████████████████████████████████████| 25/25 [01:17<00:00, 3.09s/it]
Prompt executed in 77.91 seconds
got prompt
this with teacache
Requested to load Flux
100%|███████████████████████████████████████████| 25/25 [00:40<00:00, 1.62s/it]
Prompt executed in 40.91 seconds this with teacache

this without teacache.

with a value of 0.40.
This is made with mflux.

time mflux-generate
--prompt "Luxury food photograph"
--model dev
--steps 25
--seed 1010
--height 512
--width 768
Fetching 8 files: 100%|████████████████████████| 8/8 [00:00<00:00, 75573.05it/s]
Fetching 7 files: 100%|████████████████████████| 7/7 [00:00<00:00, 19277.83it/s]
100%|███████████████████████████████████████████| 25/25 [01:02<00:00, 2.51s/it]
mflux-generate --prompt "Luxury food photograph" --model dev --steps 25 --see 3,65s user 39,06s system 64% cpu 1:06,27 total

[azrahello]

Hi! Still in the same field, I also found this node: Comfy-WaveSpeed, which in the README refers to this: ParaAttention. I’m not sure if they do the same thing, but while teacache works for me with GGUF, these others don’t.

Also, I saw that MLX Explore released this: “Support loading F8_E4M3 from safetensors”, which is a request I often see in nodes for loading the Flux model in ComfyUI (but it’s not usable because it requires too many resources). Maybe this could be useful for you? For me, it’s totally out of reach! :P

[filipstrand]

@azrahello Hi! Thanks for the pointers, I have not seen these before and your great suggestions are always welcome :)

I remember I tried the tea-cache technique a while back but did not give it enough time to properly evaluate it and put it on hold for now. It also felt a bit "hackish" with a few magic numbers in the implementation that made me a bit suspicious..., but can be worth revisiting.

For First Block Cache, it looks like they claim it is even simpler than TeaCache which sounds promising. I'll put this on the reading list and revisit after finishing up the 0.6.0 release!

[azrahello]

Hi, today I found this on Reddit: “https://www.reddit.com/r/StableDiffusion/s/VZHQWZWsc4”. It introduces a lossless compression method that should allow execution on machines with less RAM.
The links included in the article are:
• https://huggingface.co/DFloat11/FLUX.1-dev-DF11
• https://arxiv.org/abs/2504.11651
• https://github.com/LeanModels/DFloat11/tree/master/examples/flux.1
• https://huggingface.co/DFloat11/FLUX.1-schnell-DF11
• https://huggingface.co/DFloat11/FLUX.1-dev-DF11

I’m not sure if this could be applied to MLX or Apple hardware

[IonDen]

I built mlx-teacache since it's what this thread is about: it adds TeaCache to mflux as a standalone wrapper, not a fork. You register it like any other callback and it patches the predict step at runtime, caching a denoising step's transformer output and reusing it on the next step when the predicted change is small enough.

Measured on an M1 Max:

• FLUX.1-dev, 25 steps: 6 of 25 steps skipped, 1.46× wall-clock, SSIM ≥ 0.90 on the gate prompt.
• Non-distilled FLUX.2 Klein base, 50 steps + CFG: 1.23–1.36× depending on the variant.

Distilled 4–8 step Klein schedules don't get algorithmic skipping (adjacent steps differ too much for the gate to fire), so it's targeted at the longer schedules above.

Repo and per-variant notes: https://github.com/IonDen/mlx-teacache

[azrahello]

I built mlx-teacache since it's what this thread is about: it adds TeaCache to mflux as a standalone wrapper, not a fork. You register it like any other callback and it patches the predict step at runtime, caching a denoising step's transformer output and reusing it on the next step when the predicted change is small enough.

Measured on an M1 Max:

• FLUX.1-dev, 25 steps: 6 of 25 steps skipped, 1.46× wall-clock, SSIM ≥ 0.90 on the gate prompt.
• Non-distilled FLUX.2 Klein base, 50 steps + CFG: 1.23–1.36× depending on the variant.

Distilled 4–8 step Klein schedules don't get algorithmic skipping (adjacent steps differ too much for the gate to fire), so it's targeted at the longer schedules above.

Repo and per-variant notes: https://github.com/IonDen/mlx-teacache

Hello! I've seen your work and it's really piqued my interest—what are you working on right now? :D

I have to say, I've been using your tools, especially mlx-taef (in PR). Currently, I'm experimenting with ComfyUI's interface and trying to make omlx, mflux, and mlx-video coexist in order to build complex workflows that I wouldn't know how to handle via CLI.

It's funny, and I think something equivalent and versatile is still missing on Apple.

[IonDen]

Hi, thanks @azrahello I also think Apple has great potential on AI landscape and can be a strong competitor to Nvidia, but so many tools are missing now.

[ianscrivener]

@IonDen , what tools do you feel are missing?

Personally, when I feel I am missing tools... it often turns out that I can actually do it on a Mac (I use MFlux, DrawThings & ComfyUI)... and the actual issue a block/challenge on the creativity side of my process. Granted a H100 will do stuff faster than my Mac Studio Max.

[ianscrivener]

Is Teacache actually worthwhile?
I'm being genuine here. I did extensive tests with Teacache on DrawThings with Flux.1 Base - comparing it to Flux.1 Dev. I was very disappointed - Flux.1 Base + teacahe was still much slower than Flux.1 Dev - and the same (Flux.1 meh) image quality. Broadly the DrawThings Discord community seemed to share my assessment of the usefulness of Teacache for Flux.1.

Is Teacache for Flux.2 Klein Base better + faster than Flux.2 Klein Dev? Test images and benchmark times would be greatly appreciated!!

[IonDen]

@ianscrivener I have a benchmark and test results. Check it here: https://github.com/IonDen/mlx-teacache/blob/main/COMPARISON.md

[IonDen]

Of course you have all major ones on Mac. But still Metal ecosystem is much less mature than CUDA ecosystem.

[ianscrivener]

@ianscrivener I have a benchmark and test results. Check it here: https://github.com/IonDen/mlx-teacache/blob/main/COMPARISON.md

thanks @IonDen,
The comparison that I'm more interested in would be;

• Flux.2 Klein versus Flux.2 Klein Base teacache

ie is there an compelling reason to use base + teacache... or just keep using Dev as the working model for generations?
Sure, base is required for finetunes...

[IonDen]

@ianscrivener, better late than never. No, Base + TeaCache does not beat the distilled Klein on speed, and it is not close.

The arithmetic: distilled Klein runs 4 transformer passes in total. Klein Base at 50 steps with CFG runs about 100, two per step. TeaCache skips 8 of the 48 gated steps on the 4B and 13 on the 9B, which leaves Base with around 20 times more transformer work than distilled.

Measured on an M1 Max 32 GB at q4: Klein Base 4B, 768x1024, 50 steps, guidance 4: 573 s plain, 486 s with TeaCache (1.18x). Klein Base 9B, 512x512: 521 s plain, 379 s with TeaCache (1.37x, SSIM 0.986 against plain). Images and method: https://github.com/IonDen/mlx-teacache/blob/main/COMPARISON.md

So the reason to run Base is not speed. It is guidance, negative prompts and fine-tunes, and TeaCache makes that path 15 to 27 percent cheaper. For everyday generation, keep the distilled model.

I am writing a small script that runs all three (distilled, Base, Base + TeaCache) on one prompt and seed and prints a table. I can only run it at q4 here; if anyone with a 64 GB Mac wants to run it at q8 or bf16, I will post the link.