Support for Different Samplers and Schdulers?

[andyw-0612]

First of all, thanks for an amazing project! This runs on average 30% faster than flux on ComfyUI. I was wondering if there's any planned support for different schedulers and samplers like how you can choose them in the sample and scheduler nodes in ComfyUI?

It would be great to test out different combination of schedulers as people have mentioned and shared that using different schedulers can improve/change generation style and quality.

[fblissjr]

Also love this project and would love to help contribute and collaborate where possible. Have been using flux on both comfy and the excellent https://github.com/aredden/flux-fp8-api project on my CUDA machine, and excited to see this in native MLX.

Generation speeds are a bit slower for me (at FP8 runtime quant, which is what I typically do in comfy).

Was looking into mlx compilation and trying a few quick things this morning on the code base, but seems like adding compilation will require some more time reading the code and documenting the data flow. Sounds like a job for Claude later on unless someone's already done something like this (I'm a data guy and it helps to see how it all flows).

Also highly interested in assisting with LoRA training if you get there. There's no real standards on this, but I like what kijai is doing in comfy (https://github.com/kijai/ComfyUI-FluxTrainer). Most are using ai-toolkit (github.com/ostris/ai-toolkit/), which I believe uses diffusers and kohya scripts (have tried to trace all the code here, and ostris did an awesome job opening the door on all this, but haven't been able to trace the flux-specific code end-to-end yet.

I've trained a few LoRAs just for tinkering purposes on single layers / full layers / several layers using a few different methods. If trained with runtime quant at FP8, the LoRA wont work unless it's also inferenced at runtime-quant FP8 (pre-quant doesn't work on this project or any others, which is good to see - consistency).

Looking forward to seeing this project grow, and using it as a reference educational tool for learning MLX deeper. Appreciate you open sourcing this, @filipstrand .

[filipstrand]

Thank you both for your kind comments, and very sorry that I have overlooked this issue for so long.... It would definitely be possible to support other schedulers. I have personally not looked into this much yet and I like the simplicity of the current Euler one, but it would of course be nice to support different ones if they have a meaningful impact on the image quality. Due to time constraints, this won't be prioritised from my end for a while, but I would be happy to help out if anyone contributes a solution.

Also: It has been a while since September, but now we finally do have support for LoRA training, so please check it out @fblissjr.

On the topic of compilation, this has also been on my mind a bit lately, along with other optimisation techniques such as gradient checkpointing for training etc.