Upgrade Flux to new diffusers format

[avjves]

What?

Updated Flux to use the latest diffusers format.

Why?

Current flux implementation depended on old diffusers version. This made it impossible to upgrade diffusers without breaking support for Flux. Upgrading diffusers is necessary for supporting future models. It also allows us to take advantage of new diffusers features in the current models.

How?

I took the new version of Flux attention processor from diffusers and modified it to support both sequence parallelism (Ulysses/Ring) and pipeline parallelism. The transformer required a couple of changes as well, but nothing major. Flux pipeline is left as-is to support pipeline parallelism as well. Also added gating mechanisms, so certain code paths are not loaded at all if the version isn't correct. This allows different models to require different diffusers versions.

The attention processor now also lives in the transformer_flux.py file, not in attention_processor.py. This follows the style diffusers uses and allows us to do a clean version gating without having bunch of if-else statements.

Tests

Tested both with Ulysses/Ring and with pipefusion:

Tested running:

torchrun --nproc_per_node=8 examples/flux_example.py     --model black-forest-labs/FLUX.1-dev   --seed 42     --prompt "A small cat"     --height 1024     --width 1024     --num_inference_steps 25     --max_sequence_length 256        --no_use_resolution_binning      --ulysses_degree 8
torchrun --nproc_per_node=8 examples/flux_usp_example.py     --model black-forest-labs/FLUX.1-dev   --seed 42     --prompt "A small cat"     --height 1024     --width 1024     --num_inference_steps 25     --max_sequence_length 256        --no_use_resolution_binning      --ulysses_degree 8
torchrun --nproc_per_node=8 examples/flux_example.py     --model black-forest-labs/FLUX.1-dev   --seed 42     --prompt "A small cat"     --height 1024     --width 1024     --num_inference_steps 25     --max_sequence_length 256        --no_use_resolution_binning      --ulysses_degree 2 --pipefusion_parallel_degree 4

Other

• Updated the setup.py
  • Diffusers is clear requirements nowadays - xDiT can't be ran without it due to all the imports. I set the minimum version to 0.33.0, as this is the lowest one one can run xDiT with due to pre-existing diffusers imports in the code.
• Updated the README.md
  • It now reflects the diffusers requirement.
  • Added a table for supported diffusers version for models. Currently only has Flux, but more will come.

[jcaraban]

LGTM