Motivation

Currently FLUX.1-dev has two correctness issues that deteriorate the output quality. This can be seen as a blocky output image and the inability to generate higher resolution (2048 * 2048) images that it should be able to. This PR adds fixes to the two issues, improving the output quality significantly.

Modifications

1. Image seq len calculation properly takes into account the patchification. Without this fix the mu values are too big, pushing sigma values to 1.0 incorrectly.
   2. Change the self-attention to not replicate the un-sharded text tokens. This now follows how it is done with FLUX.2, for example. Already fixed in [diffusion] fix: fix accuracy for some image models #20679

Accuracy Tests

Without fixes:

1 GPU:

Run command:

sglang generate --model-path black-forest-labs/FLUX.1-dev --seed 42 --prompt "A small cat" --height 1024 --width 1024 --num-inference-steps 25 --ulysses-degree 8 --guidance-scale 0.0 --num-gpus 8

8 GPUs:

Run command:

sglang generate --model-path black-forest-labs/FLUX.1-dev --seed 42 --prompt "A small cat" --height 1024 --width 1024 --num-inference-steps 25 --ulysses-degree 8 --guidance-scale 0.0 --num-gpus 8

8 GPUs (2048x2048):

Run command:

sglang generate --model-path black-forest-labs/FLUX.1-dev --seed 42 --prompt "A small cat" --height 2048 --width 2048 --num-inference-steps 25 --ulysses-degree 8 --guidance-scale 0.0 --num-gpus 8

With fixes:

1 GPU:

Run command:

sglang generate --model-path black-forest-labs/FLUX.1-dev --seed 42 --prompt "A small cat" --height 1024 --width 1024 --num-inference-steps 25 --ulysses-degree 1 --guidance-scale 0.0 --num-gpus 1

8 GPUs:

Run command:

sglang generate --model-path black-forest-labs/FLUX.1-dev --seed 42 --prompt "A small cat" --height 1024 --width 1024 --num-inference-steps 25 --ulysses-degree 8 --guidance-scale 0.0 --num-gpus 8

8 GPUs (2048x2048):

Run command:

sglang generate --model-path black-forest-labs/FLUX.1-dev --seed 42 --prompt "A small cat" --height 2048 --width 2048 --num-inference-steps 25 --ulysses-degree 8 --guidance-scale 0.0 --num-gpus 8

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.