Request for new models

[karthik4579]

Will support for new models be added like Auraflow and Flux be added ?

https://huggingface.co/fal/AuraFlow
https://huggingface.co/black-forest-labs/FLUX.1-dev

[entrpn]

@karthik4579 there are plans to add Flux in the future

[gabeweisz]

Can you comment on the expected timeline for adding flux? Would you be open to a PR that adds it (we haven't built it yet, but could work on it if you are open to such a PR)

[entrpn]

@gabeweisz PRs are always welcomed :)

I am planning to start working on flux in beginning of 2025. There is already a JAX Flux repo that runs on both GPUs and TPUs, but it needs to be optimized. repo

[jcaraban]

@entrpn together with @gabeweisz we have a working implementation of Flux.1 in a private fork of MaxDiffusion. We plan to submit a PR soon, but I'm reaching now as I see @jfacevedo-google has resumed work in his flux_impl branch this week. We originally started from jflux using Flax.nnx, but once that worked we migrated to linen.nn on top of Juan's flux stale branch. At this point our implementation is probably not too different from what's brewing in flux_impl. Our FSDP sharding is not optimal but works. Performance is on par with HF.diffusers.FluxTransformer2DModel or slightly better for MI300x. What do you think?

[entrpn]

@jcaraban and @gabeweisz thanks for working on this. It sounds like you took the same approach I'm taking. Please open a PR and we can work together on merging it to main. Excited to test it out!

[gabeweisz]

We plan on submitting an initial PR on Monday. We do expect you to have feedback on how we have structured the code and done some of the implementation, and are looking forward to working with you on this.