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.
Screenshot 2025年01月14日 at 19 27 49

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

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
ComfyUI_00135_
this without teacache.
ComfyUI_00134_
with a value of 0.40.
This is made with mflux.
image
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

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

@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!

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

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, ×ばつ wall-clock, SSIM ≥ 0.90 on the gate prompt.
• Non-distilled FLUX.2 Klein base, 50 steps + CFG: ×ばつ 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

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, ×ばつ wall-clock, SSIM ≥ 0.90 on the gate prompt.
• Non-distilled FLUX.2 Klein base, 50 steps + CFG: ×ばつ 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.

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.

@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.

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!!

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

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

@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...