Flux inference implementation

[entrpn]

This PR implements Flux dev and schnell for inference.

Features:

• Loads encoders, vae and transformer to HBM with encoder offload option. All settings available in the flux configs.
• Optimized flash block sizes for v6e.
• Flash attention
• Sharding - DDP and FSDP, allows loading all models on HBM and keeping them loaded using FSDP sharding, while reducing performance slightly.
• Removed maxdiffusion.transformers and move back to HF transformers library due to support of new models. Originally I ported some of the models from transformers into maxdiffusion.transformers but saw no improvements in sharding them. This will reduce maintenance of those models inside of maxdiffusion.

[jcaraban]

@entrpn I will close #141 since this PR is in a mergable state now
@ksikiric rebased our training on top of flux_impl and opened a continuation PR #147
After that we still plan to submit another PR with flash_attn support for GPUs, alike to MaxText

[entrpn]

@entrpn I will close #141 since this PR is in a mergable state now @ksikiric rebased our training on top of flux_impl and opened a continuation PR #147 After that we still plan to submit another PR with flash_attn support for GPUs, alike to MaxText

thanks @jcaraban and @ksikiric appreciate it very much. I plan to merge this PR by eow. On top of that I've been working on LoRA loading for Flux, which is almost complete and I will also be pushing a few fixes for some issues I found. I'll make sure to keep you informed on these items as they move along.

I've also been working with my team internally to ensure you and your employer is attributed for your work. Thanks again!