AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-23 09:09:01 +03:00
Code Activity

Flux3 image (#15046)

Browse Source 
* Bump the actions group with 17 updates (#12)

Bumps the actions group with 17 updates:

| Package | From | To |
| --- | --- | --- |
| [actions/checkout](https://github.com/actions/checkout) | `6` | `7` |
| [actions/upload-artifact](https://github.com/actions/upload-artifact) | `4.6.2` | `7.0.1` |
| [docker/setup-buildx-action](https://github.com/docker/setup-buildx-action) | `3.12.0` | `4.2.0` |
| [docker/login-action](https://github.com/docker/login-action) | `3.7.0` | `4.6.0` |
| [docker/build-push-action](https://github.com/docker/build-push-action) | `6.19.2` | `7.3.0` |
| [huggingface/doc-builder/.github/workflows/build_main_documentation.yml](https://github.com/huggingface/doc-builder) | `e60a538eea9817ab312196d0d233604b01697265` | `af49b8dc3859c3d80c7220c8c0ba146431132e49` |
| [actions/setup-python](https://github.com/actions/setup-python) | `5.6.0` | `7.0.0` |
| [huggingface/doc-builder/.github/workflows/build_pr_documentation.yml](https://github.com/huggingface/doc-builder) | `e60a538eea9817ab312196d0d233604b01697265` | `af49b8dc3859c3d80c7220c8c0ba146431132e49` |
| [anthropics/claude-code-action](https://github.com/anthropics/claude-code-action) | `1.0.155` | `1.0.193` |
| [huggingface/security-workflows/.github/workflows/codeql-reusable.yml](https://github.com/huggingface/security-workflows) | `1.1.0` | `1.2.0` |
| [actions/download-artifact](https://github.com/actions/download-artifact) | `4.3.0` | `8.0.1` |
| [actions/labeler](https://github.com/actions/labeler) | `5.0.0` | `7.0.0` |
| [pypa/gh-action-pypi-publish](https://github.com/pypa/gh-action-pypi-publish) | `1.13.0` | `1.14.2` |
| [huggingface/tailscale-action](https://github.com/huggingface/tailscale-action) | `7d53c9737e53934c30290b5524d1c9b4a7c98c8a` | `07d36e4c3b68707bd3fe070dbb96f03f993a6efd` |
| [trufflesecurity/trufflehog](https://github.com/trufflesecurity/trufflehog) | `3.94.2` | `3.97.0` |
| [crate-ci/typos](https://github.com/crate-ci/typos) | `1.42.1` | `1.49.0` |
| [huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml](https://github.com/huggingface/doc-builder) | `9ad2de8582b56c017cb530c1165116d40433f1c6` | `af49b8dc3859c3d80c7220c8c0ba146431132e49` |


Updates `actions/checkout` from 6 to 7
- [Release notes](https://github.com/actions/checkout/releases)
- [Commits](https://github.com/actions/checkout/compare/v6...v7)

Updates `actions/upload-artifact` from 4.6.2 to 7.0.1
- [Release notes](https://github.com/actions/upload-artifact/releases)
- [Commits](https://github.com/actions/upload-artifact/compare/v4.6.2...v7.0.1)

Updates `docker/setup-buildx-action` from 3.12.0 to 4.2.0
- [Release notes](https://github.com/docker/setup-buildx-action/releases)
- [Commits](8d2750c68a...bb05f3f551)

Updates `docker/login-action` from 3.7.0 to 4.6.0
- [Release notes](https://github.com/docker/login-action/releases)
- [Commits](c94ce9fb46...dbcb813823)

Updates `docker/build-push-action` from 6.19.2 to 7.3.0
- [Release notes](https://github.com/docker/build-push-action/releases)
- [Commits](10e90e3645...53b7df96c9)

Updates `huggingface/doc-builder/.github/workflows/build_main_documentation.yml` from e60a538eea9817ab312196d0d233604b01697265 to af49b8dc3859c3d80c7220c8c0ba146431132e49
- [Release notes](https://github.com/huggingface/doc-builder/releases)
- [Commits](e60a538eea...af49b8dc38)

Updates `actions/setup-python` from 5.6.0 to 7.0.0
- [Release notes](https://github.com/actions/setup-python/releases)
- [Commits](https://github.com/actions/setup-python/compare/v5.6.0...v7)

Updates `huggingface/doc-builder/.github/workflows/build_pr_documentation.yml` from e60a538eea9817ab312196d0d233604b01697265 to af49b8dc3859c3d80c7220c8c0ba146431132e49
- [Release notes](https://github.com/huggingface/doc-builder/releases)
- [Commits](e60a538eea...af49b8dc38)

Updates `anthropics/claude-code-action` from 1.0.155 to 1.0.193
- [Release notes](https://github.com/anthropics/claude-code-action/releases)
- [Commits](80b3182633...9d7150bc8a)

Updates `huggingface/security-workflows/.github/workflows/codeql-reusable.yml` from 1.1.0 to 1.2.0
- [Release notes](https://github.com/huggingface/security-workflows/releases)
- [Commits](dc6ca34688...e0d40aaf3c)

Updates `actions/download-artifact` from 4.3.0 to 8.0.1
- [Release notes](https://github.com/actions/download-artifact/releases)
- [Commits](https://github.com/actions/download-artifact/compare/v4.3.0...3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c)

Updates `actions/labeler` from 5.0.0 to 7.0.0
- [Release notes](https://github.com/actions/labeler/releases)
- [Commits](8558fd7429...bf12e9b00b)

Updates `pypa/gh-action-pypi-publish` from 1.13.0 to 1.14.2
- [Release notes](https://github.com/pypa/gh-action-pypi-publish/releases)
- [Commits](ed0c53931b...dc37677b2e)

Updates `huggingface/tailscale-action` from 7d53c9737e53934c30290b5524d1c9b4a7c98c8a to 07d36e4c3b68707bd3fe070dbb96f03f993a6efd
- [Release notes](https://github.com/huggingface/tailscale-action/releases)
- [Commits](7d53c9737e...07d36e4c3b)

Updates `trufflesecurity/trufflehog` from 3.94.2 to 3.97.0
- [Release notes](https://github.com/trufflesecurity/trufflehog/releases)
- [Commits](6bd2d14f7a...bcfcf73aaf)

Updates `crate-ci/typos` from 1.42.1 to 1.49.0
- [Release notes](https://github.com/crate-ci/typos/releases)
- [Changelog](https://github.com/crate-ci/typos/blob/master/CHANGELOG.md)
- [Commits](65120634e7...8a48f81b6c)

Updates `huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml` from 9ad2de8582b56c017cb530c1165116d40433f1c6 to af49b8dc3859c3d80c7220c8c0ba146431132e49
- [Release notes](https://github.com/huggingface/doc-builder/releases)
- [Commits](9ad2de8582...af49b8dc38)

---
updated-dependencies:
- dependency-name: actions/checkout
  dependency-version: '7'
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: actions/upload-artifact
  dependency-version: 7.0.1
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: docker/setup-buildx-action
  dependency-version: 4.2.0
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: docker/login-action
  dependency-version: 4.6.0
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: docker/build-push-action
  dependency-version: 7.3.0
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: huggingface/doc-builder/.github/workflows/build_main_documentation.yml
  dependency-version: af49b8dc3859c3d80c7220c8c0ba146431132e49
  dependency-type: direct:production
  dependency-group: actions
- dependency-name: actions/setup-python
  dependency-version: 7.0.0
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml
  dependency-version: af49b8dc3859c3d80c7220c8c0ba146431132e49
  dependency-type: direct:production
  dependency-group: actions
- dependency-name: anthropics/claude-code-action
  dependency-version: 1.0.193
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: actions
- dependency-name: huggingface/security-workflows/.github/workflows/codeql-reusable.yml
  dependency-version: 1.2.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: actions
- dependency-name: actions/download-artifact
  dependency-version: 8.0.1
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: actions/labeler
  dependency-version: 7.0.0
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: pypa/gh-action-pypi-publish
  dependency-version: 1.14.2
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: actions
- dependency-name: huggingface/tailscale-action
  dependency-version: 07d36e4c3b68707bd3fe070dbb96f03f993a6efd
  dependency-type: direct:production
  dependency-group: actions
- dependency-name: trufflesecurity/trufflehog
  dependency-version: 3.97.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: actions
- dependency-name: crate-ci/typos
  dependency-version: 1.49.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: actions
- dependency-name: huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml
  dependency-version: af49b8dc3859c3d80c7220c8c0ba146431132e49
  dependency-type: direct:production
  dependency-group: actions
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

* Scope GITHUB_TOKEN permissions per job (#16)

A job with no `permissions:` block inherits whatever the repository
hands out. Each block added here grants what that job's own steps
need and nothing more.

Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com>

* Pin GitHub Actions to commit SHAs (#17)

A tag is mutable: whoever controls the action's repository can move
`v4` to different code, and the next run picks it up with no diff to
review. Pinning to the commit SHA freezes the code that runs today.

Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com>

* Bump doc-builder workflow pin (#22)

* Bump doc-builder pin to 9fc8a41 in build_documentation.yml

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Bump doc-builder pin to 9fc8a41 in build_pr_documentation.yml

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Bump doc-builder pin to 9fc8a41 in upload_pr_documentation.yml

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

* Add FLUX.3 image: Flux3Transformer2DModel and modular text-to-image / image-to-image pipeline

Adds the FLUX.3 image transformer, the modular pipeline (text-to-image and image-to-image with one
prompt per call and one or more reference images), Flux3ImageProcessor, the converter for BFL's
model package, tests and docs. Verified against BFL's reference on the real checkpoint
(integration-notes/parity_bf16.py, temporary folder).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* integration-notes: ship the knight example image with the doc examples

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* integration-notes: drop the tokenizer config workaround (not needed with transformers 5.x)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Address review: rename the mode channel dict and the final norm, name the reference workflow
image_conditioned, document the guidance scale

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Add the step-distilled FLUX.3 image model (flow map)

- transformer: `flow_map` config and `target_timestep_embedder`; forward takes
  `target_timestep` (the timestep each step jumps to), added to the modulation
  of every mode and the final layer
- convert script: map `ref_time_in.*` and read `flow_map` from the package config
- modular: dedicated `Flux3Distilled*` blocks (text encoder, text input, text
  RoPE inputs, denoiser) with no guider and no negative prompt, 16 steps by
  default; `Flux3DistilledAutoBlocks` and `Flux3DistilledModularPipeline`
- tests for the flow-map transformer path and the distilled modular pipeline
- docs: step-distilled section and autodoc entries

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Add the step-distilled examples script

Same four cases as flux3_examples.py, run against the distilled Hub repo.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* feat: add a single file support in Flux3 image. (#32)

* Warn about the denoiser input fields that are not passed to the transformer

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Document reproducible generation with user-provided latents

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Run the QK norm and RoPE in fp32 with a single cast back, matching the BFL implementation

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* parity script: run against the Hub repos and print the imported diffusers source

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Wrap a forward docstring line for the doc-builder style check

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* support lora loading (including ostris) in flux3 image pipeline. (#25)

* feat: support loading and performing inference with single-file formatted FP8 and NVFP4 checkpoints (#24)

* feat: support loading siungle file format checkpoints including quants.

* remove updates.patch

* remove fluff.

* feat: add a single file support in Flux3 image.

* bfl prefix

* version guard.

* prpoagate changes.

* add tests

* doc mods and remove integration notes

* remove script for quant

---------

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

* add LoRA training for Flux 3 image small

* tests: clarify FLUX3_TINY_REPO comment

The tiny repo is uploaded to hf-internal-testing and the env var is a
local-dev escape hatch, not a pre-PR TODO.

* examples/flux3: drop unused args and dead class

Remove the text-encoder weight-decay arg from both trainers (neither
trains the text encoder) and, from the img2img trainer, the
prior-preservation leftovers (--class_data_dir, --sample_batch_size,
PromptDataset) that it never uses.

* bump yarn example rank to 32 in flux3 small readme

* Document <ref_image_n> references and layout prompts for FLUX.3

Matches the prompt grammar in BFL's reference inference code (bfl-hf-flux3-confidential#3): references are addressed as <ref_image_n> in the order passed to image, and regions can be placed with a trailing JSON list of id / bbox / desc (from / src_bbox / tgt_bbox when editing).

* examples/flux3: set klein model id, document ref image prompts, clean leftover refs

* point flux3 image docs at the klein release model ids

* fix: TestFlux3DistilledModularPipelineWorkflow::test_workflow_defaults test (#35)

* fix checkpoint ids in the flux3 docs. (#38)

* [docs] update pipe loading with vae fp32 dtype (#39)

* update pipe loading with vae fp32 dtype

* updating `load_componenets`

* use actual ckpts. (#42)

* flux3 klein example: rename from "small" and fix caption dropout for long prompts (#41)

* examples/flux3: rename small trainer, tests and readme files to klein

* examples/flux3: fix caption dropout shape mismatch for prompts longer than the pad window

* examples/flux3: add caption dropout to the i2i trainer

* examples/flux3: complete the klein rename (test class names, example output dirs)

* examples/flux3: guard against train_batch_size>1 with variable-length captions

* update flux3.md (#43)

* examples/flux3: support train_batch_size>1 and prior preservation with variable-length captions (#44)

* doc

* autodocstrings.

---------

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com>
Co-authored-by: Pauline Bailly-Masson <155966238+paulinebm@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: yiyixuxu <yixu310@gmail.com>
Co-authored-by: linoytsaban <linoy.tsaban@gmail.com>
Co-authored-by: apolinário <joaopaulo.passos@gmail.com>
Co-authored-by: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com>
Co-authored-by: Aritra Roy Gosthipaty <aritra.born2fly@gmail.com>
Co-authored-by: Yumeng Li <109459218+YumengLi007@users.noreply.github.com>
This commit is contained in:
Sayak Paul
2026-10-15 22:00:50 +05:30
committed by GitHub
parent 62cde2b011
commit 7ecb1e1b4f
81 changed files with 10638 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
.github/workflows/benchmark.yml vendored
View File

@@ -33,7 +33,7 @@ jobs:
options: --shm-size "16gb" --ipc host --gpus all
steps:
- name: Checkout diffusers
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
- name: NVIDIA-SMI
@@ -63,7 +63,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: benchmark_test_reports
path: benchmarks/${{ env.BASE_PATH }}

4
.github/workflows/bot_pytest.yml vendored
View File

@@ -75,7 +75,7 @@ jobs:
shell: bash
steps:
- name: Checkout PR head
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
# Works for forks too — no fork credentials needed.
ref: refs/pull/${{ github.event.issue.number }}/head
@@ -116,7 +116,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: bot_gpu_test_reports
path: reports

12
.github/workflows/build_docker_images.yml vendored
View File

@@ -31,10 +31,10 @@ jobs:
pull-requests: read
steps:
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
- name: Check out code
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Find Changed Dockerfiles
id: file_changes
@@ -107,16 +107,16 @@ jobs:
steps:
- name: Checkout repository
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
- name: Login to Docker Hub
uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3.7.0
uses: docker/login-action@dbcb813823bdd20940b903addbd779551569679f # v4.6.0
with:
username: ${{ env.REGISTRY }}
password: ${{ secrets.DOCKERHUB_TOKEN }}
- name: Build and push
uses: docker/build-push-action@10e90e3645eae34f1e60eeb005ba3a3d33f178e8 # v6.19.2
uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a # v7.3.0
with:
no-cache: true
context: ./docker/${{ matrix.image-name }}

4
.github/workflows/build_pr_documentation.yml vendored
View File

@@ -22,10 +22,10 @@ jobs:
steps:
- name: Checkout repository
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Set up Python
uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v6
with:
python-version: '3.10'

4
.github/workflows/claude_review.yml vendored
View File

@@ -36,7 +36,7 @@ jobs:
cancel-in-progress: false
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd #v6.0.2
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 #v7.0.1
with:
fetch-depth: 1
@@ -86,7 +86,7 @@ jobs:
git clone --local --bare . /tmp/local-origin.git
git config url."file:///tmp/local-origin.git".insteadOf "$(git remote get-url origin)"
- uses: anthropics/claude-code-action@80b31826338489861333dc17217865dfe8085cdc # v1.0.155
- uses: anthropics/claude-code-action@9d7150bc8a3dae8149739a88019d192b579ad90c # v1.0.193
env:
CLAUDE_SYSTEM_PROMPT: |
You are a strict code reviewer for the diffusers library (huggingface/diffusers).

2
.github/workflows/codeql.yml vendored
View File

@@ -10,7 +10,7 @@ on:
jobs:
codeql:
name: CodeQL Analysis
uses: huggingface/security-workflows/.github/workflows/codeql-reusable.yml@dc6ca34688e6876c2dd18750719b44d177586c17 # v1
uses: huggingface/security-workflows/.github/workflows/codeql-reusable.yml@e0d40aaf3c566a5e820787840453f45f909e337f # v1.2.0
permissions:
security-events: write
packages: read

2
.github/workflows/issue_labeler.yml vendored
View File

@@ -14,7 +14,7 @@ jobs:
issues: write
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Install dependencies
run: pip install huggingface_hub
- name: Get labels from LLM

4
.github/workflows/mirror_community_pipeline.yml vendored
View File

@@ -71,13 +71,13 @@ jobs:
run: |
echo "CHECKOUT_REF: ${{ env.CHECKOUT_REF }}"
echo "PATH_IN_REPO: ${{ env.PATH_IN_REPO }}"
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
ref: ${{ env.CHECKOUT_REF }}
# Setup + install dependencies
- name: Set up Python
uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
with:
python-version: "3.10"
- name: Install dependencies

46
.github/workflows/nightly_tests.yml vendored
View File

@@ -38,7 +38,7 @@ jobs:
pipeline_test_matrix: ${{ steps.fetch_pipeline_matrix.outputs.pipeline_test_matrix }}
steps:
- name: Checkout diffusers
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
- name: Install dependencies
@@ -54,7 +54,7 @@ jobs:
- name: Pipeline Tests Artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: test-pipelines.json
path: reports
@@ -76,7 +76,7 @@ jobs:
options: --shm-size "16gb" --ipc host --gpus all
steps:
- name: Checkout diffusers
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
- name: NVIDIA-SMI
@@ -108,7 +108,7 @@ jobs:
cat reports/tests_pipeline_${{ matrix.module }}_cuda_failures_short.txt
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: pipeline_${{ matrix.module }}_test_reports
path: reports
@@ -132,7 +132,7 @@ jobs:
module: [models, schedulers, hooks, lora, others, single_file, examples]
steps:
- name: Checkout diffusers
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -180,7 +180,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch_${{ matrix.module }}_cuda_test_reports
path: reports
@@ -199,7 +199,7 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -225,7 +225,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch_compile_test_reports
path: reports
@@ -244,7 +244,7 @@ jobs:
options: --shm-size "16gb" --ipc host --gpus all
steps:
- name: Checkout diffusers
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
- name: NVIDIA-SMI
@@ -279,7 +279,7 @@ jobs:
cat reports/tests_big_gpu_torch_cuda_failures_short.txt
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch_cuda_big_gpu_test_reports
path: reports
@@ -298,7 +298,7 @@ jobs:
shell: bash
steps:
- name: Checkout diffusers
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -337,7 +337,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch_minimum_version_cuda_test_reports
path: reports
@@ -370,7 +370,7 @@ jobs:
options: --shm-size "20gb" --ipc host --gpus all
steps:
- name: Checkout diffusers
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
- name: NVIDIA-SMI
@@ -409,7 +409,7 @@ jobs:
cat reports/tests_${{ matrix.config.backend }}_torch_cuda_failures_short.txt
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch_cuda_${{ matrix.config.backend }}_reports
path: reports
@@ -428,7 +428,7 @@ jobs:
options: --shm-size "20gb" --ipc host --gpus all
steps:
- name: Checkout diffusers
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
- name: NVIDIA-SMI
@@ -462,7 +462,7 @@ jobs:
cat reports/tests_pipeline_level_quant_torch_cuda_failures_short.txt
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch_cuda_pipeline_level_quant_reports
path: reports
@@ -487,7 +487,7 @@ jobs:
image: diffusers/diffusers-pytorch-cpu
steps:
- name: Checkout diffusers
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -495,7 +495,7 @@ jobs:
run: mkdir -p combined_reports
- name: Download all test reports
uses: actions/download-artifact@37930b1c2abaa49bbe596cd826c3c89aef350131 # v7.0.0
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
with:
path: artifacts
@@ -521,7 +521,7 @@ jobs:
cat $CONSOLIDATED_REPORT_PATH >> $GITHUB_STEP_SUMMARY
- name: Upload consolidated report
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: consolidated_test_report
path: ${{ env.CONSOLIDATED_REPORT_PATH }}
@@ -535,7 +535,7 @@ jobs:
#
# steps:
# - name: Checkout diffusers
# uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
# uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
# with:
# fetch-depth: 2
#
@@ -575,7 +575,7 @@ jobs:
#
# - name: Test suite reports artifacts
# if: ${{ always() }}
# uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6
# uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
# with:
# name: torch_mps_test_reports
# path: reports
@@ -591,7 +591,7 @@ jobs:
#
# steps:
# - name: Checkout diffusers
# uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
# uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
# with:
# fetch-depth: 2
#
@@ -631,7 +631,7 @@ jobs:
#
# - name: Test suite reports artifacts
# if: ${{ always() }}
# uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6
# uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
# with:
# name: torch_mps_test_reports
# path: reports

4
.github/workflows/notify_slack_about_release.yml vendored
View File

@@ -15,10 +15,10 @@ jobs:
runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Setup Python
uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
with:
python-version: '3.10'

4
.github/workflows/pr_dependency_test.yml vendored
View File

@@ -24,9 +24,9 @@ jobs:
contents: read
runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Set up Python
uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
with:
python-version: "3.11"
- name: Install dependencies

4
.github/workflows/pr_labeler.yml vendored
View File

@@ -20,12 +20,12 @@ jobs:
REPO: ${{ github.repository }}
steps:
# label
- uses: actions/labeler@8558fd74291d67161a8a78ce36a881fa63b766a9 # v5
- uses: actions/labeler@bf12e9b00b37c5c0ca2b87b79b2daf7891dbda13 # v7.0.0
with:
sync-labels: true
# missing-tests
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
id: checkout
if: ${{ !cancelled() }}
with:

4
.github/workflows/pr_link_issue_reminder.yml vendored
View File

@@ -26,10 +26,10 @@ jobs:
# (Sayak Paul, YiYi Xu, Dhruv Nair).
SLACK_MENTION_IDS: ${{ secrets.SLACK_PR_LINK_ISSUE_MENTION_IDS }}
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Setup Python
uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v6
with:
python-version: "3.10"

14
.github/workflows/pr_modular_tests.yml vendored
View File

@@ -46,9 +46,9 @@ jobs:
contents: read
runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Set up Python
uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
with:
python-version: "3.11"
- name: Install dependencies
@@ -68,9 +68,9 @@ jobs:
needs: check_code_quality
runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Set up Python
uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
with:
python-version: "3.11"
- name: Install dependencies
@@ -99,7 +99,7 @@ jobs:
container:
image: diffusers/diffusers-pytorch-cpu
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Install dependencies
run: |
printf 'torch==2.10.0\ntorchvision==0.25.0\ntorchaudio==2.10.0\n' > "$UV_OVERRIDE"
@@ -131,7 +131,7 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -160,7 +160,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: pr_pytorch_pipelines_torch_cpu_modular_pipelines_test_reports
path: reports

12
.github/workflows/pr_test_fetcher.yml vendored
View File

@@ -33,7 +33,7 @@ jobs:
test_map: ${{ steps.set_matrix.outputs.test_map }}
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 0
- name: Install dependencies
@@ -47,7 +47,7 @@ jobs:
run: |
python utils/tests_fetcher.py | tee test_preparation.txt
- name: Report fetched tests
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: test_fetched
path: test_preparation.txt
@@ -90,7 +90,7 @@ jobs:
shell: bash
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -116,7 +116,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: ${{ matrix.modules }}_test_reports
path: reports
@@ -147,7 +147,7 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -173,7 +173,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: pr_${{ matrix.config.report }}_test_reports
path: reports

20
.github/workflows/pr_tests.yml vendored
View File

@@ -41,9 +41,9 @@ jobs:
contents: read
runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Set up Python
uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
with:
python-version: "3.11"
- name: Install dependencies
@@ -63,9 +63,9 @@ jobs:
needs: check_code_quality
runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Set up Python
uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
with:
python-version: "3.11"
- name: Install dependencies
@@ -126,7 +126,7 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -170,7 +170,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: pr_${{ matrix.config.framework }}_${{ matrix.config.report }}_test_reports
path: reports
@@ -204,7 +204,7 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -232,7 +232,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: pr_${{ matrix.config.report }}_test_reports
path: reports
@@ -260,7 +260,7 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -301,7 +301,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: pr_lora_test_reports
path: reports

24
.github/workflows/pr_tests_gpu.yml vendored
View File

@@ -43,9 +43,9 @@ jobs:
contents: read
runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Set up Python
uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
with:
python-version: "3.11"
- name: Install dependencies
@@ -65,9 +65,9 @@ jobs:
needs: check_code_quality
runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Set up Python
uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
with:
python-version: "3.11"
- name: Install dependencies
@@ -100,7 +100,7 @@ jobs:
pipeline_test_matrix: ${{ steps.fetch_pipeline_matrix.outputs.pipeline_test_matrix }}
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
- name: Install dependencies
@@ -118,7 +118,7 @@ jobs:
echo "pipeline_test_matrix=$matrix" >> $GITHUB_OUTPUT
- name: Pipeline Tests Artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: test-pipelines.json
path: reports
@@ -140,7 +140,7 @@ jobs:
options: --shm-size "16gb" --ipc host --gpus all
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -189,7 +189,7 @@ jobs:
cat reports/tests_pipeline_${{ matrix.module }}_cuda_failures_short.txt
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: pipeline_${{ matrix.module }}_test_reports
path: reports
@@ -214,7 +214,7 @@ jobs:
module: [models, schedulers, lora, others]
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -260,7 +260,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch_cuda_test_reports_${{ matrix.module }}
path: reports
@@ -278,7 +278,7 @@ jobs:
options: --gpus all --shm-size "16gb" --ipc host
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -310,7 +310,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: examples_test_reports
path: reports

4
.github/workflows/pr_torch_dependency_test.yml vendored
View File

@@ -24,9 +24,9 @@ jobs:
contents: read
runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Set up Python
uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
with:
python-version: "3.11"
- name: Install dependencies

24
.github/workflows/push_tests.yml vendored
View File

@@ -39,7 +39,7 @@ jobs:
pipeline_test_matrix: ${{ steps.fetch_pipeline_matrix.outputs.pipeline_test_matrix }}
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
- name: Install dependencies
@@ -57,7 +57,7 @@ jobs:
echo "pipeline_test_matrix=$matrix" >> $GITHUB_OUTPUT
- name: Pipeline Tests Artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: test-pipelines.json
path: reports
@@ -79,7 +79,7 @@ jobs:
options: --shm-size "16gb" --ipc host --gpus all
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
- name: NVIDIA-SMI
@@ -110,7 +110,7 @@ jobs:
cat reports/tests_pipeline_${{ matrix.module }}_cuda_failures_short.txt
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: pipeline_${{ matrix.module }}_test_reports
path: reports
@@ -134,7 +134,7 @@ jobs:
module: [models, schedulers, hooks, lora, others, single_file]
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -169,7 +169,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch_cuda_test_reports_${{ matrix.module }}
path: reports
@@ -188,7 +188,7 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -214,7 +214,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch_compile_test_reports
path: reports
@@ -233,7 +233,7 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -258,7 +258,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch_xformers_test_reports
path: reports
@@ -276,7 +276,7 @@ jobs:
options: --gpus all --shm-size "16gb" --ipc host
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -307,7 +307,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: examples_test_reports
path: reports

4
.github/workflows/push_tests_fast.yml vendored
View File

@@ -62,7 +62,7 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -95,7 +95,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: pr_${{ matrix.config.report }}_test_reports
path: reports

4
.github/workflows/push_tests_mps.yml vendored
View File

@@ -28,7 +28,7 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -70,7 +70,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: pr_torch_mps_test_reports
path: reports

12
.github/workflows/pypi_publish.yaml vendored
View File

@@ -18,10 +18,10 @@ jobs:
runs-on: ubuntu-22.04
steps:
- name: Checkout repo
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
with:
python-version: "3.11"
@@ -51,7 +51,7 @@ jobs:
python -c "from diffusers import *"
- name: Upload build artifacts
uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: python-dist
path: dist/
@@ -65,15 +65,15 @@ jobs:
id-token: write
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Download build artifacts
uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
with:
name: python-dist
path: dist/
- name: Publish package distributions to TestPyPI
uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # v1.13.0
uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # release/v1
with:
verbose: true

28
.github/workflows/release_tests_fast.yml vendored
View File

@@ -38,7 +38,7 @@ jobs:
pipeline_test_matrix: ${{ steps.fetch_pipeline_matrix.outputs.pipeline_test_matrix }}
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
- name: Install dependencies
@@ -57,7 +57,7 @@ jobs:
echo "pipeline_test_matrix=$matrix" >> $GITHUB_OUTPUT
- name: Pipeline Tests Artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: test-pipelines.json
path: reports
@@ -79,7 +79,7 @@ jobs:
options: --shm-size "16gb" --ipc host --gpus all
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
- name: NVIDIA-SMI
@@ -110,7 +110,7 @@ jobs:
cat reports/tests_pipeline_${{ matrix.module }}_cuda_failures_short.txt
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: pipeline_${{ matrix.module }}_test_reports
path: reports
@@ -134,7 +134,7 @@ jobs:
module: [models, schedulers, lora, others, single_file]
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -169,7 +169,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch_cuda_${{ matrix.module }}_test_reports
path: reports
@@ -188,7 +188,7 @@ jobs:
shell: bash
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -227,7 +227,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch_minimum_version_cuda_test_reports
path: reports
@@ -246,7 +246,7 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -272,7 +272,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch_compile_test_reports
path: reports
@@ -291,7 +291,7 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -317,7 +317,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: torch_xformers_test_reports
path: reports
@@ -336,7 +336,7 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
@@ -369,7 +369,7 @@ jobs:
- name: Test suite reports artifacts
if: ${{ always() }}
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6.0.0
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: examples_test_reports
path: reports

4
.github/workflows/ssh-pr-runner.yml vendored
View File

@@ -32,12 +32,12 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2
- name: Tailscale # In order to be able to SSH when a test fails
uses: huggingface/tailscale-action@7d53c9737e53934c30290b5524d1c9b4a7c98c8a # main
uses: huggingface/tailscale-action@07d36e4c3b68707bd3fe070dbb96f03f993a6efd # main
with:
authkey: ${{ secrets.TAILSCALE_SSH_AUTHKEY }}
slackChannel: ${{ secrets.SLACK_CIFEEDBACK_CHANNEL }}

2
.github/workflows/ssh-runner.yml vendored
View File

@@ -40,7 +40,7 @@ jobs:
steps:
- name: Checkout diffusers
uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 2

4
.github/workflows/stale.yml vendored
View File

@@ -15,10 +15,10 @@ jobs:
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Setup Python
uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
with:
python-version: 3.10

4
.github/workflows/trufflehog.yml vendored
View File

@@ -13,11 +13,11 @@ jobs:
runs-on: ubuntu-22.04
steps:
- name: Checkout code
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
fetch-depth: 0
- name: Secret Scanning
uses: trufflesecurity/trufflehog@6bd2d14f7a4bc1e569fa3550efa7ec632a4fa67b # main
uses: trufflesecurity/trufflehog@bcfcf73aaf4759d4dadc2783177c245a02792318 # main
with:
extra_args: --results=verified,unknown

4
.github/workflows/typos.yml vendored
View File

@@ -13,7 +13,7 @@ jobs:
runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: typos-action
uses: crate-ci/typos@65120634e79d8374d1aa2f27e54baa0c364fff5a # v1.42.1
uses: crate-ci/typos@8a48f81b6c64dcfea44b3633223084c4be58ac5f # v1.49.0

2
.github/workflows/update_metadata.yml vendored
View File

@@ -20,7 +20,7 @@ jobs:
shell: bash -l {0}
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Setup environment
run: |

6
docs/source/en/_toctree.yml
View File

@@ -85,6 +85,8 @@
title: Getting started
- local: quantization/bitsandbytes
title: bitsandbytes
- local: quantization/bfl
title: BFL (fp8r / nvfp4)
- local: quantization/gguf
title: gguf
- local: quantization/nunchaku
@@ -349,6 +351,8 @@
title: ErnieImageTransformer2DModel
- local: api/models/flux2_transformer
title: Flux2Transformer2DModel
- local: api/models/flux3_transformer
title: Flux3Transformer2DModel
- local: api/models/flux_transformer
title: FluxTransformer2DModel
- local: api/models/glm_image_transformer2d
@@ -587,6 +591,8 @@
title: Flux
- local: api/pipelines/flux2
title: Flux2
- local: api/pipelines/flux3
title: Flux3
- local: api/pipelines/control_flux_inpaint
title: FluxControlInpaint
- local: api/pipelines/glm_image

4
docs/source/en/api/loaders/lora.md
View File

@@ -68,6 +68,10 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi
[[autodoc]] loaders.lora_pipeline.Flux2LoraLoaderMixin
## Flux3LoraLoaderMixin
[[autodoc]] loaders.lora_pipeline.Flux3LoraLoaderMixin
## ErnieImageLoraLoaderMixin
[[autodoc]] loaders.lora_pipeline.ErnieImageLoraLoaderMixin

34
docs/source/en/api/models/flux3_transformer.md Normal file
View File

@@ -0,0 +1,34 @@
<!-- Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License. -->
# Flux3Transformer2DModel
A Transformer model for image-like data from [FLUX.3](https://bfl.ai/), for text-to-image and image-to-image generation. Three token modes share one joint sequence: the noisy image tokens, the VAE-encoded reference image tokens and the text context. Each mode first passes through its own per-mode transformer blocks, then the modes are concatenated for the joint single blocks, and the image tokens pass through the output layers. Timesteps are per token: the image tokens carry the flow timestep, the reference tokens and the text context sit at `0`.
```python
import torch
from diffusers import Flux3Transformer2DModel
transformer = Flux3Transformer2DModel.from_pretrained(
"black-forest-labs/flux-3-image-klein-base", subfolder="transformer", dtype=torch.bfloat16
)
```
## Flux3Transformer2DModel
[[autodoc]] Flux3Transformer2DModel
## Transformer2DModelOutput
[[autodoc]] models.modeling_outputs.Transformer2DModelOutput

275
docs/source/en/api/pipelines/flux3.md Normal file
View File

@@ -0,0 +1,275 @@
<!-- Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License. -->
# FLUX.3
FLUX.3 is an image generation model from Black Forest Labs: a single transformer ([`Flux3Transformer2DModel`]) generates images from a text prompt, alone or together with one or more reference images. Text conditioning comes from Qwen3-VL, and images use the FLUX.2 packed autoencoder ([`AutoencoderKLFlux2`]).
FLUX.3 ships as a modular pipeline. Prompting matters: the model responds best to full scene captions that describe the subject, composition, lighting, and atmosphere. Black Forest Labs provides a [prompt upsampler](https://huggingface.co/black-forest-labs/flux-3-image-klein-pup) that expands short prompts automatically; see [Prompt upsampling](#prompt-upsampling). Get started with a text-to-image generation:
```python
import torch
from diffusers import ModularPipeline
pipe = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein-base")
pipe.load_components(dtype={"vae": torch.float32, "default": torch.bfloat16})
pipe.to("cuda")
image = pipe(
prompt=(
"A wide-angle wildlife photograph of a red fox standing alert in a snow-covered forest "
"clearing at golden hour. The fox is centered in the frame, its thick orange winter coat "
"contrasting with the fresh white snow, breath visible in the cold air. Snow-laden pine "
"trees fill the softly blurred background, and low warm sunlight rakes across the scene "
"from the left, casting long blue shadows."
),
generator=torch.Generator().manual_seed(0),
output="images",
)[0]
image.save("fox.png")
```
`prompt` is the only required input: the pipeline carries the recommended sampling settings (50 steps, guidance scale 3.5, scheduler shift 6.93) and generates 1024 x 1024 images by default. It takes one prompt per call; use `num_images_per_prompt` for several images of the same prompt.
## Prompt upsampling
The prompt upsampler, [black-forest-labs/flux-3-image-klein-pup](https://huggingface.co/black-forest-labs/flux-3-image-klein-pup), turns the user prompt, together with the reference images when editing, into a full scene caption. It is a separate model; run it with Transformers and pass its output as `prompt`:
```python
import json
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration
model_id = "black-forest-labs/flux-3-image-klein-pup"
processor = AutoProcessor.from_pretrained(model_id)
upsampler = Qwen3_5ForConditionalGeneration.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto").eval()
system_prompts = json.load(open(hf_hub_download(model_id, "system_prompts.json")))
def upsample(prompt, images=()):
if images:
text = (
"Reference images follow in order (the first image, the second image, ...).\n"
f"User request: {prompt}\nOutput the caption."
)
content = [{"type": "image", "image": image} for image in images] + [{"type": "text", "text": text}]
messages = [{"role": "system", "content": system_prompts["i2i"]}, {"role": "user", "content": content}]
else:
messages = [{"role": "system", "content": system_prompts["t2i"]}, {"role": "user", "content": prompt}]
inputs = processor.apply_chat_template(
messages,
add_generation_prompt=True,
tokenize=True,
return_dict=True,
return_tensors="pt",
enable_thinking=False,
).to(upsampler.device)
with torch.no_grad():
output = upsampler.generate(**inputs, max_new_tokens=3584, do_sample=False)
return processor.decode(output[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True).strip()
image = pipe(prompt=upsample("An astronaut riding a horse"), output="images")[0]
```
For an edit, pass the same reference images to both: `pipe(prompt=upsample(instruction, [image]), image=image, ...)`. The upsampler also writes the [layout](#layout-control) captions from a list of boxes; see its model card.
## Image-to-image
Pass one or more reference images with `image`. Every reference keeps its own resolution (center-cropped to the 16-pixel grid; references above 16 megapixels are downscaled first), and the generation size defaults to one megapixel at the first reference's aspect ratio. Describe the *target* image as a complete caption of the reference with your edit folded in, rather than a one-line instruction, or let the [prompt upsampler](#prompt-upsampling) write it from the instruction and the references:
```python
night = pipe(
prompt=(
"The same red fox in the same snow-covered forest clearing, now at night under a full moon: "
"cool blue moonlight, deep shadows, stars visible through the pine branches, the fox's eyes "
"catching the light."
),
image=image,
generator=torch.Generator().manual_seed(0),
output="images",
)[0]
```
With several references, refer to them in the prompt as `<ref_image_0>`, `<ref_image_1>`, ... in the order they are passed to `image`:
```python
combined = pipe(
prompt="The fox from <ref_image_0> sitting on the wooden bench from <ref_image_1>, in the same snowy clearing.",
image=[image, bench],
height=768,
width=1360,
output="images",
)[0]
```
Pass `height` and/or `width` to override the derived size.
## Layout control
To place content explicitly, tag regions in the caption with `<id>` and append a JSON list of regions to the prompt. Each region has an `id`, a `bbox` of `[y0, x0, y1, x1]` (top-left and bottom-right corners, origin top-left) normalized to 0-1000 on both axes so it is independent of the output size, and a `desc`:
```python
prompt = (
"Minimalist graphic illustration featuring a black silhouette of a person <silhouette_1> centered against a "
"solid, vibrant chartreuse background <background_1>. "
'[{"id": "background_1", "bbox": [0, 0, 1000, 1000], "desc": "A flat field of neon yellow-green."}, '
'{"id": "silhouette_1", "bbox": [150, 150, 850, 850], "desc": "A black silhouette of a person in motion."}]'
)
```
When editing, each region names its source reference with `from` (for example `"ref_image_0"`), where it is in that image with `src_bbox`, and where it should end up with `tgt_bbox`. Equal boxes keep a region in place:
```python
prompt = (
"In <ref_image_0>, move the knight figure <knight_1> upwards and to the left. Keep the rest of the image "
"unchanged, preserving the backdrop <backdrop_1>. "
'[{"id": "knight_1", "from": "ref_image_0", "src_bbox": [500, 150, 850, 350], "tgt_bbox": [194, 55, 544, 255], '
'"desc": "A miniature amigurumi knight crocheted from grey yarn."}, '
'{"id": "backdrop_1", "from": "ref_image_0", "src_bbox": [0, 0, 1000, 1000], "tgt_bbox": [0, 0, 1000, 1000], '
'"desc": "A soft, out-of-focus backdrop of knitted yarn."}]'
)
image = pipe(prompt=prompt, image=knight, output="images")[0]
```
## Guidance scale
The pipeline applies classifier-free guidance through its `guider` component and does not take a `guidance_scale` argument at call time. The default guider is [`ClassifierFreeGuidance`] with `guidance_scale=3.5`, the recommended setting for both workflows. To change it, create an updated copy of the guider:
```python
pipe.guider = pipe.guider.new(guidance_scale=5.0)
```
See the [guiders guide](../../using-diffusers/guiders) to switch to a different guidance method.
## Step-distilled model
The step-distilled checkpoint generates in 16 steps without classifier-free guidance. Its transformer is a flow map: at every step it is told the current timestep and the one the step jumps to. It loads like the base model (the repo's `modular_model_index.json` resolves to [`Flux3DistilledModularPipeline`] and [`Flux3DistilledAutoBlocks`]); the pipeline has no `guider` and does not take `negative_prompt`:
```python
pipe = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein")
pipe.load_components(dtype=torch.bfloat16)
pipe.to("cuda")
image = pipe(prompt=prompt, generator=torch.Generator().manual_seed(0), output="images")[0]
```
It serves the same two workflows as the base model. `num_inference_steps` defaults to 16; the model is trained for 4 to 16 steps.
## Reproducible generation
A `generator` fixes the noise the pipeline draws (see [Reproducibility](../../using-diffusers/reusing_seeds)). To reuse noise drawn elsewhere, pass it as `latents` with shape `(batch_size, 128, height // 16, width // 16)`. For example, the [BFL implementation](https://github.com/black-forest-labs/flux-3) draws its noise as image tokens on the CPU and casts it to bf16; this reproduces its sample for a given seed:
```python
height, width = 1024, 1024
h, w = height // 16, width // 16
noise = torch.randn((h * w, 128), generator=torch.Generator().manual_seed(0)).to(torch.bfloat16)
latents = noise.T.reshape(1, 128, h, w)
image = pipe(prompt=prompt, latents=latents, height=height, width=width, output="images")[0]
```
## Workflows
The pipeline serves two workflows, `text2image` and `image_conditioned`; the one that runs is picked by the inputs (`image` selects image-to-image). [`~ModularPipeline.from_pretrained`] can also build a pipeline for one workflow with the `workflow` argument, which then only expects and loads the components that workflow needs:
```python
from diffusers.modular_pipelines import Flux3AutoBlocks
print(Flux3AutoBlocks().available_workflows)
# ["text2image", "image_conditioned"]
t2i = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein-base", workflow="text2image")
print(t2i.blocks.doc)
```
## Memory
Loading everything in bf16 takes about 20 GB of weights (the transformer is 5.5B parameters, the Qwen3-VL-4B text encoder 4B), with the image VAE in fp32. For GPUs with less memory, enable component offloading:
```python
pipe = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein-base")
pipe.load_components(dtype={"vae": torch.float32, "default": torch.bfloat16})
pipe.enable_auto_cpu_offload(device="cuda")
```
## Single-file loading
[`Flux3Transformer2DModel`] loads the transformer straight from a Black Forest Labs model package with [`~loaders.FromOriginalModelMixin.from_single_file`]. The package `config.json` is not read; pass the diffusers config with `config` and `subfolder`.
```python
import torch
from diffusers import Flux3Transformer2DModel, ModularPipeline
transformer = Flux3Transformer2DModel.from_single_file(
"https://huggingface.co/black-forest-labs/flux-3-image-klein-base/blob/main/model.safetensors",
config="black-forest-labs/flux-3-image-klein-base",
subfolder="transformer",
dtype=torch.bfloat16,
)
pipe = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein-base")
pipe.load_components(dtype={"vae": torch.float32, "default": torch.bfloat16})
pipe.update_components(transformer=transformer)
pipe.to("cuda")
```
Load the other components by name first: [`~ModularPipeline.load_components`] without `names` would also load the transformer from the pipeline repo.
Loading the FP8r and NVFP4 single-file checkpoints are also supported. Both load through the same call with a [`BFLQuantizationConfig`]. See the [BFL quantization guide](../../quantization/bfl) for the kernel requirements of each scheme.
```python
from diffusers import BFLQuantizationConfig
transformer = Flux3Transformer2DModel.from_single_file(
"https://huggingface.co/black-forest-labs/flux-3-image-klein-base/blob/main/variants/nvfp4/model.safetensors",
config="black-forest-labs/flux-3-image-klein-base",
subfolder="transformer",
quantization_config=BFLQuantizationConfig(compute_dtype=torch.bfloat16),
dtype=torch.bfloat16,
)
```
When loading the step-distilled checkpoints, we need their
configs (`flow_map=True`):
```python
transformer = Flux3Transformer2DModel.from_single_file(
"https://huggingface.co/black-forest-labs/flux-3-image-klein/blob/main/variants/nvfp4/model.safetensors",
config="black-forest-labs/flux-3-image-klein",
subfolder="transformer",
quantization_config=BFLQuantizationConfig(compute_dtype=torch.bfloat16),
dtype=torch.bfloat16,
)
```
## Flux3ModularPipeline
[[autodoc]] Flux3ModularPipeline
## Flux3AutoBlocks
[[autodoc]] Flux3AutoBlocks
## Flux3DistilledAutoBlocks
[[autodoc]] Flux3DistilledAutoBlocks
## Flux3DistilledModularPipeline
[[autodoc]] Flux3DistilledModularPipeline

4
docs/source/en/api/quantization.md
View File

@@ -22,6 +22,10 @@ Quantization techniques reduce memory and computational costs by representing we
[[autodoc]] quantizers.PipelineQuantizationConfig
## BFLQuantizationConfig
[[autodoc]] quantizers.quantization_config.BFLQuantizationConfig
## BitsAndBytesConfig
[[autodoc]] quantizers.quantization_config.BitsAndBytesConfig

63
docs/source/en/quantization/bfl.md Normal file
View File

@@ -0,0 +1,63 @@
<!--Copyright 2026 The HuggingFace Team. All rights reserved.
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
specific language governing permissions and limitations under the License.
-->
# BFL (fp8r / nvfp4)
[Black Forest Labs](https://huggingface.co/black-forest-labs) ships quantized variants of its [Flux3](../api/pipelines/flux3) transformers in two schemes, stored as single-file checkpoints next to the original bf16 weights:
- `fp8r`: rowwise FP8. Each quantized linear stores an E4M3 `weight` and an fp32 per-row `weight_scale`. Activations are quantized to E4M3 per row at runtime.
- `nvfp4`: NVFP4 in the [ModelOpt](https://github.com/NVIDIA/Model-Optimizer) layout. Each quantized linear stores a packed E2M1 `weight` (two values per byte), E4M3 block scales of 16 (`weight_scale`), a global fp32 `weight_scale_2`, and a calibrated fp32 activation scale `input_scale`.
Diffusers loads both through [`~FromSingleFileMixin.from_single_file`] with a [`BFLQuantizationConfig`]. The scheme is read from the checkpoint, so the same config works for either file. Only loading prequantized checkpoints is supported; pipeline-level loading is not.
The example below loads the `fp8r` checkpoint of [FLUX 3 Image Klein Base](https://huggingface.co/black-forest-labs/flux-3-image-klein-base), which lives at `variants/fp8r/model.safetensors` in the model repository. Pass `config` so that the model config is read from the Diffusers folder of that repository, then load the other components by name (see [Single-file loading](../api/pipelines/flux3#single-file-loading)).
```python
import torch
from diffusers import BFLQuantizationConfig, Flux3Transformer2DModel, ModularPipeline
ckpt_path = "https://huggingface.co/black-forest-labs/flux-3-image-klein-base/blob/main/variants/fp8r/model.safetensors"
transformer = Flux3Transformer2DModel.from_single_file(
ckpt_path,
config="black-forest-labs/flux-3-image-klein-base",
subfolder="transformer",
quantization_config=BFLQuantizationConfig(compute_dtype=torch.bfloat16),
dtype=torch.bfloat16,
)
pipe = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein-base")
pipe.load_components(dtype=torch.bfloat16)
pipe.update_components(transformer=transformer)
pipe.to("cuda")
prompt = "A wide-angle wildlife photograph of a red fox standing alert in a snow-covered forest clearing at golden hour."
image = pipe(prompt=prompt, generator=torch.Generator().manual_seed(0), output="images")[0]
image.save("flux3-klein-base-fp8r.png")
```
Replace `variants/fp8r` with `variants/nvfp4` to load the `nvfp4` variant. The step-distilled [FLUX 3 Image Klein](https://huggingface.co/black-forest-labs/flux-3-image-klein) ships the same two variants; load them with its repository id as `config` and pipeline.
## Kernels
The quantized weights stay in their low-precision storage dtype. Which GEMM runs depends on the device:
| Scheme | Fused kernel | Requirements |
|---|---|---|
| `fp8r` | `torch._scaled_mm` with rowwise scales | CUDA GPU with compute capability 8.9 or higher |
| `nvfp4` | `flashinfer.mm_fp4` (CUTLASS) with `flashinfer.nvfp4_quantize` for the activations | Blackwell GPU (compute capability 10.0 or higher), a CUDA toolkit with `nvcc`, and `flashinfer-python` with `flashinfer-cubin` at the same version |
Install the two FlashInfer packages at the same version, for example `pip install flashinfer-python==0.6.13 flashinfer-cubin==0.6.13`. The `flashinfer-cubin` release can lag behind `flashinfer-python`, and FlashInfer refuses to import when the versions differ.
FlashInfer compiles its NVFP4 kernels on first use with the CUDA toolkit of the environment and picks the target architecture from the visible GPU (or from `FLASHINFER_CUDA_ARCH_LIST` when set). Consumer and workstation Blackwell GPUs (compute capability 12.0) need a CUDA toolkit of 12.9 or newer. With an older toolkit FlashInfer finds no eligible architecture and the first NVFP4 call fails with `FlashInfer requires GPUs with sm75 or higher`.
On any other device the weights are dequantized to `compute_dtype` in each forward pass and the layer runs as a regular linear. In that case the activations are not quantized, so outputs differ slightly from the fused kernels.

1
docs/source/en/quantization/overview.md
View File

@@ -146,6 +146,7 @@ The following backends support loading prequantized checkpoints out of the box.
| [bitsandbytes](./bitsandbytes) | Config is saved in `config.json`; no extra arguments needed. Supports quantizing on the fly as well as loading prequantized checkpoints. |
| [torchao](./torchao) | Same as above. |
| [GGUF](./gguf) | Uses `from_single_file` with Model classes; pipeline-level loading is not supported. |
| [BFL](./bfl) | Loads Black Forest Labs `fp8r` and `nvfp4` single-file checkpoints with `from_single_file`; pipeline-level loading is not supported. |
| [AutoRound](./autoround) | Only loading is supported; quantize first with the AutoRound CLI or Python API. |
| [Nunchaku Lite](./nunchaku) | Config is saved in `config.json`; requires the `kernels` package. Only loading is supported. |
| [ModelOpt](./modelopt) | Only loading prequantized models is supported. |

224
examples/dreambooth/README_flux3_klein.md Normal file
View File

@@ -0,0 +1,224 @@
# DreamBooth LoRA training for the FLUX.3 image model
These scripts train LoRA adapters for the FLUX.3 image model (`Flux3Transformer2DModel`, served by
`Flux3ModularPipeline`) with the DreamBooth recipe:
- `train_dreambooth_lora_flux3_klein.py` — text-to-image (t2i)
- `train_dreambooth_lora_flux3_klein_img2img.py` — image-to-image / instruction editing (i2i)
> [!NOTE]
> LoRA **loading and inference** for this model live in the pipeline/loader PR (`Flux3LoraLoaderMixin`); these
> scripts only cover **training**. Adapters are saved through `Flux3LoraLoaderMixin.save_lora_weights`, so they load
> back with `Flux3ModularPipeline.load_lora_weights`.
A few things that are specific to FLUX.3 and differ from the FLUX.1/FLUX.2 trainers:
- **Text encoder is Qwen3-VL.** The text embedding is the stack of several Qwen3-VL hidden layers
(`--text_encoder_out_layers`, default `4 8 12 16 20 24 28 32`), concatenated along the channel dim. The prompt goes
through the Qwen chat template and the text window is right-padded to a multiple of 80 tokens (capped at
`--max_sequence_length`), matching the pipeline's `Flux3TextEncoderStep`.
- **No guidance distillation.** The model is not guidance-distilled, so there is no guidance argument in training;
classifier-free guidance happens at inference. Use `--caption_dropout` for unconditional training (t2i and i2i).
- **VAE normalization is a batch-norm.** Latents are 2×2-patchified to 128 channels and normalized with the VAE's
stored batch-norm statistics (not a scalar `scaling_factor`/`shift_factor`).
- **Flow matching with a fixed shift of 6.93** (no resolution-dependent dynamic shifting); the velocity target is
`noise - image`.
- **Validation reuses the in-memory trained transformer** via `Flux3AutoBlocks().init_pipeline()` +
`update_components(...)`, so the DiT is never loaded a second time.
## Running locally with PyTorch
### Installing the dependencies
Install `diffusers` from source and keep it up to date:
```bash
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install -e .
```
Then `cd` into `examples/dreambooth` and install the requirements:
```bash
pip install -r requirements_flux3_klein.txt
```
Initialize an [🤗 Accelerate](https://github.com/huggingface/accelerate/) environment with:
```bash
accelerate config
```
Or, for a default configuration:
```bash
accelerate config default
```
We use PEFT as the LoRA backend, so make sure `peft>=0.11.1` is installed.
## Text-to-image (t2i)
### Style LoRA example
This trains a *yarn-art* style LoRA on [Norod78/Yarn-art-style](https://huggingface.co/datasets/Norod78/Yarn-art-style),
whose captions already contain the trigger phrase `yarn art style`.
```bash
export MODEL_NAME="black-forest-labs/flux-3-image-klein-base"
export OUTPUT_DIR="trained-flux3-klein-yarn-lora"
accelerate launch train_dreambooth_lora_flux3_klein.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--dataset_name="Norod78/Yarn-art-style" \
--caption_column="text" \
--instance_prompt="yarn art style" \
--output_dir=$OUTPUT_DIR \
--mixed_precision="bf16" \
--resolution=1024 \
--train_batch_size=1 \
--gradient_checkpointing \
--cache_latents \
--rank=32 \
--lora_alpha=32 \
--optimizer="adamW" \
--use_8bit_adam \
--learning_rate=1e-4 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--max_train_steps=2000 \
--validation_prompt="a cat, yarn art style" \
--validation_epochs=25 \
--seed="0" \
--push_to_hub
```
For a subject/object LoRA, point `--instance_data_dir` at a local folder of images and set `--instance_prompt` to a
caption with a trigger word (e.g. `a photo of sks dog`) instead of `--dataset_name`/`--caption_column`.
> [!NOTE]
> For long prompts, raise `--max_sequence_length` (the Qwen3-VL token cap, default 512). Longer windows use more
> memory and can slow down training.
## Image-to-image (i2i)
The i2i script trains an instruction-editing / relighting-style LoRA on **(condition image, target image,
instruction)** triplets. The reference image is VAE-encoded and fed to the model as a separate, clean conditioning
stream (its own position ids and timestep 0); the loss is computed on the generated image.
This example uses [kontext-community/relighting](https://huggingface.co/datasets/kontext-community/relighting)
(`file_name` = source, `output` = relit target, `instruction` = edit prompt):
```bash
export MODEL_NAME="black-forest-labs/flux-3-image-klein-base"
export OUTPUT_DIR="trained-flux3-klein-relight-lora"
accelerate launch train_dreambooth_lora_flux3_klein_img2img.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--dataset_name="kontext-community/relighting" \
--image_column="output" \
--cond_image_column="file_name" \
--caption_column="instruction" \
--output_dir=$OUTPUT_DIR \
--mixed_precision="bf16" \
--resolution=1024 \
--train_batch_size=1 \
--gradient_checkpointing \
--cache_latents \
--rank=16 \
--lora_alpha=16 \
--optimizer="adamW" \
--use_8bit_adam \
--learning_rate=1e-4 \
--lr_scheduler="constant" \
--max_train_steps=1500 \
--seed="0" \
--validation_prompt="relight the scene with warm sunset light" \
--validation_image="path/to/a/source/image.png" \
--validation_epochs=25 \
--push_to_hub
```
i2i reads its data through `--dataset_name` (a dataset with an image column, a conditioning-image column and a caption
column), not a plain image folder.
### Addressing the reference image in the prompt
FLUX.3 refers to each conditioning image passed to `image=` as `<ref_image_0>`, `<ref_image_1>`, ... in the order
they are passed. i2i training uses a single conditioning image, so captions address it as `<ref_image_0>`, e.g.
`"Relight <ref_image_0> with warm sunset light."` Matching this grammar in your captions keeps the adapter aligned
with how the base model reads references.
For explicit placement you can tag regions with an `<id>` and append a JSON list of regions to the caption. Each
region has an `id`, a `bbox` of `[y0, x0, y1, x1]` (top-left and bottom-right corners, origin top-left) normalized to
0-1000 on both axes, and a `desc`. When editing, a region also names its source reference with `from` (e.g.
`"ref_image_0"`), where it starts with `src_bbox`, and where it should end up with `tgt_bbox` (equal boxes keep it in
place):
```text
In <ref_image_0>, move the knight <knight_1> up and to the left, keeping everything else.
[{"id": "knight_1", "from": "ref_image_0", "src_bbox": [500, 150, 850, 350], "tgt_bbox": [194, 55, 544, 255],
"desc": "A miniature knight crocheted from grey yarn."}]
```
See the [FLUX.3 pipeline docs](https://github.com/huggingface/diffusers/blob/main/docs/source/en/api/pipelines/flux3.md)
for the full prompt grammar.
## Tracking
Pass `--report_to="wandb"` (or `tensorboard`) and `--validation_prompt` (plus `--validation_image` for i2i) to log
validation samples during training.
## Memory optimizations
- **CPU offloading** — `--offload` moves the VAE and text encoder to CPU except while they are used.
- **Latent caching** — `--cache_latents` pre-encodes images to latents once (the VAE is then freed), at the cost of no
random image augmentations.
- **Low-precision training** — `--do_fp8_training` (torchao FP8) or `--bnb_quantization_config_path` (bitsandbytes
QLoRA); the two are mutually exclusive.
- **Gradient checkpointing / accumulation** — `--gradient_checkpointing`, `--gradient_accumulation_steps`.
- **8-bit Adam** — `--optimizer="adamW" --use_8bit_adam` (or `--optimizer="prodigy"`).
- **Saved-LoRA precision** — `--upcast_before_saving` saves the adapter in fp32 instead of the training dtype.
## LoRA rank, alpha and target modules
`--rank` / `--lora_alpha` / `--lora_dropout` control the adapter. By default the scripts adapt the attention and
feed-forward projections across every block:
```
to_q, to_k, to_v, to_out.0, ff.linear_in, ff.linear_out
```
These are exactly the leaf modules the non-diffusers (ostris / ai-toolkit) FLUX.3 LoRA converter maps, so adapters
trained here stay interchangeable with that layout. These names match across the per-mode blocks (`mode_blocks.image`,
`mode_blocks.image_cond`, `txt_mode_blocks`) and the joint `single_transformer_blocks`, so a single target set reaches
every stream. Use `--lora_layers` (a comma-separated list) to target a different set, e.g. to also adapt
`context_embedder`, `x_embedders.image` or `proj_out`.
## Aspect-ratio bucketing and caption dropout
- `--aspect_ratio_buckets` / `--use_aspect_ratio_buckets` group images of similar aspect ratio so a batch shares a
resolution (t2i).
- `--caption_dropout` replaces a fraction of captions (t2i) or instructions (i2i) with the empty prompt during
training, which improves classifier-free guidance at inference. For i2i the conditioning image is kept.
## Inference
Load a trained adapter onto the modular pipeline and generate:
```python
import torch
from diffusers.modular_pipelines import Flux3AutoBlocks
pipe = Flux3AutoBlocks().init_pipeline("black-forest-labs/flux-3-image-klein-base")
pipe.load_components(dtype=torch.bfloat16)
pipe.to("cuda")
pipe.load_lora_weights("trained-flux3-klein-yarn-lora", weight_name="pytorch_lora_weights.safetensors")
image = pipe(prompt="a cat, yarn art style", output="images")[0]
image.save("yarn_cat.png")
```
For i2i, pass the reference image through the pipeline's `image=` input.

11
examples/dreambooth/requirements_flux3_klein.txt Normal file
View File

@@ -0,0 +1,11 @@
accelerate>=0.31.0
torchvision
transformers>=4.57.0
ftfy
tensorboard
Jinja2
peft>=0.11.1
sentencepiece
bitsandbytes
prodigyopt
datasets>=3.1.0

194
examples/dreambooth/test_dreambooth_lora_flux3_klein.py Normal file
View File

@@ -0,0 +1,194 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import os
import sys
import tempfile
import safetensors
sys.path.append("..")
from test_examples_utils import ExamplesTestsAccelerate, run_command # noqa: E402
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger()
stream_handler = logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)
# Tiny randomly-initialized FLUX.3 image modular repo used for fast CI. The env var is an escape hatch for
# running against a local build (see `integration-notes/make_tiny_flux3_repo.py`).
FLUX3_TINY_REPO = os.environ.get("FLUX3_TINY_REPO", "hf-internal-testing/tiny-flux3-image-modular")
class DreamBoothLoRAFlux3Klein(ExamplesTestsAccelerate):
instance_data_dir = "docs/source/en/imgs"
instance_prompt = "dog"
pretrained_model_name_or_path = FLUX3_TINY_REPO
script_path = "examples/dreambooth/train_dreambooth_lora_flux3_klein.py"
transformer_layer_type = "single_transformer_blocks.0.attn.to_k"
def test_dreambooth_lora_flux3_klein(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_args = f"""
{self.script_path}
--pretrained_model_name_or_path {self.pretrained_model_name_or_path}
--instance_data_dir {self.instance_data_dir}
--instance_prompt {self.instance_prompt}
--resolution 64
--train_batch_size 1
--gradient_accumulation_steps 1
--max_train_steps 2
--learning_rate 5.0e-04
--scale_lr
--lr_scheduler constant
--lr_warmup_steps 0
--max_sequence_length 8
--text_encoder_out_layers 1
--output_dir {tmpdir}
""".split()
run_command(self._launch_args + test_args)
# save_pretrained smoke test
self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")))
# make sure the state_dict has the correct naming in the parameters.
lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
is_lora = all("lora" in k for k in lora_state_dict.keys())
self.assertTrue(is_lora)
# when not training the text encoder, all the parameters in the state dict should start
# with `"transformer"` in their names.
starts_with_transformer = all(key.startswith("transformer") for key in lora_state_dict.keys())
self.assertTrue(starts_with_transformer)
def test_dreambooth_lora_flux3_klein_latent_caching(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_args = f"""
{self.script_path}
--pretrained_model_name_or_path {self.pretrained_model_name_or_path}
--instance_data_dir {self.instance_data_dir}
--instance_prompt {self.instance_prompt}
--resolution 64
--train_batch_size 1
--gradient_accumulation_steps 1
--max_train_steps 2
--cache_latents
--learning_rate 5.0e-04
--scale_lr
--lr_scheduler constant
--lr_warmup_steps 0
--max_sequence_length 8
--text_encoder_out_layers 1
--output_dir {tmpdir}
""".split()
run_command(self._launch_args + test_args)
# save_pretrained smoke test
self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")))
# make sure the state_dict has the correct naming in the parameters.
lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
is_lora = all("lora" in k for k in lora_state_dict.keys())
self.assertTrue(is_lora)
starts_with_transformer = all(key.startswith("transformer") for key in lora_state_dict.keys())
self.assertTrue(starts_with_transformer)
def test_dreambooth_lora_flux3_klein_layers(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_args = f"""
{self.script_path}
--pretrained_model_name_or_path {self.pretrained_model_name_or_path}
--instance_data_dir {self.instance_data_dir}
--instance_prompt {self.instance_prompt}
--resolution 64
--train_batch_size 1
--gradient_accumulation_steps 1
--max_train_steps 2
--cache_latents
--learning_rate 5.0e-04
--scale_lr
--lora_layers {self.transformer_layer_type}
--lr_scheduler constant
--lr_warmup_steps 0
--max_sequence_length 8
--text_encoder_out_layers 1
--output_dir {tmpdir}
""".split()
run_command(self._launch_args + test_args)
# save_pretrained smoke test
self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")))
# make sure the state_dict has the correct naming in the parameters.
lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
is_lora = all("lora" in k for k in lora_state_dict.keys())
self.assertTrue(is_lora)
# only the layers named by `--lora_layers` should be trained, so every key sits under that path
starts_with_transformer = all(
key.startswith(f"transformer.{self.transformer_layer_type}") for key in lora_state_dict.keys()
)
self.assertTrue(starts_with_transformer)
def test_dreambooth_lora_flux3_klein_checkpointing_checkpoints_total_limit(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_args = f"""
{self.script_path}
--pretrained_model_name_or_path={self.pretrained_model_name_or_path}
--instance_data_dir={self.instance_data_dir}
--output_dir={tmpdir}
--instance_prompt={self.instance_prompt}
--resolution=64
--train_batch_size=1
--gradient_accumulation_steps=1
--max_train_steps=6
--checkpoints_total_limit=2
--max_sequence_length 8
--checkpointing_steps=2
--text_encoder_out_layers 1
""".split()
run_command(self._launch_args + test_args)
self.assertEqual(
{x for x in os.listdir(tmpdir) if "checkpoint" in x},
{"checkpoint-4", "checkpoint-6"},
)
def test_dreambooth_lora_flux3_klein_checkpointing_checkpoints_total_limit_removes_multiple_checkpoints(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_args = f"""
{self.script_path}
--pretrained_model_name_or_path={self.pretrained_model_name_or_path}
--instance_data_dir={self.instance_data_dir}
--output_dir={tmpdir}
--instance_prompt={self.instance_prompt}
--resolution=64
--train_batch_size=1
--gradient_accumulation_steps=1
--max_train_steps=4
--checkpointing_steps=2
--max_sequence_length 8
--text_encoder_out_layers 1
""".split()
run_command(self._launch_args + test_args)
self.assertEqual({x for x in os.listdir(tmpdir) if "checkpoint" in x}, {"checkpoint-2", "checkpoint-4"})

121
examples/dreambooth/test_dreambooth_lora_flux3_klein_img2img.py Normal file
View File

@@ -0,0 +1,121 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import os
import sys
import tempfile
import safetensors
sys.path.append("..")
from test_examples_utils import ExamplesTestsAccelerate, run_command # noqa: E402
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger()
stream_handler = logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)
# Tiny randomly-initialized FLUX.3 image modular repo used for fast CI. The env var is an escape hatch for
# running against a local build (see `integration-notes/make_tiny_flux3_repo.py`).
FLUX3_TINY_REPO = os.environ.get("FLUX3_TINY_REPO", "hf-internal-testing/tiny-flux3-image-modular")
class DreamBoothLoRAFlux3KleinImg2Img(ExamplesTestsAccelerate):
# image-to-image training needs a paired conditioning image, so it reads a dataset with an image +
# conditioning_image column rather than a plain folder
dataset_name = "hf-internal-testing/fill10"
image_column = "image"
cond_image_column = "conditioning_image"
caption_column = "text"
instance_prompt = "dog"
pretrained_model_name_or_path = FLUX3_TINY_REPO
script_path = "examples/dreambooth/train_dreambooth_lora_flux3_klein_img2img.py"
def test_dreambooth_lora_flux3_klein_img2img(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_args = f"""
{self.script_path}
--pretrained_model_name_or_path {self.pretrained_model_name_or_path}
--dataset_name {self.dataset_name}
--image_column {self.image_column}
--cond_image_column {self.cond_image_column}
--caption_column {self.caption_column}
--instance_prompt {self.instance_prompt}
--resolution 64
--train_batch_size 1
--gradient_accumulation_steps 1
--max_train_steps 2
--learning_rate 5.0e-04
--lr_scheduler constant
--lr_warmup_steps 0
--max_sequence_length 8
--text_encoder_out_layers 1
--output_dir {tmpdir}
""".split()
run_command(self._launch_args + test_args)
# save_pretrained smoke test
self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")))
# make sure the state_dict has the correct naming in the parameters.
lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
is_lora = all("lora" in k for k in lora_state_dict.keys())
self.assertTrue(is_lora)
starts_with_transformer = all(key.startswith("transformer") for key in lora_state_dict.keys())
self.assertTrue(starts_with_transformer)
# the reference image is consumed by its own ("image_cond") mode blocks, so training an img2img adapter
# with the default targets reaches both the generated-image and the reference mode blocks
self.assertTrue(any("mode_blocks.image." in key for key in lora_state_dict.keys()))
self.assertTrue(any("mode_blocks.image_cond." in key for key in lora_state_dict.keys()))
def test_dreambooth_lora_flux3_klein_img2img_latent_caching(self):
with tempfile.TemporaryDirectory() as tmpdir:
test_args = f"""
{self.script_path}
--pretrained_model_name_or_path {self.pretrained_model_name_or_path}
--dataset_name {self.dataset_name}
--image_column {self.image_column}
--cond_image_column {self.cond_image_column}
--caption_column {self.caption_column}
--instance_prompt {self.instance_prompt}
--resolution 64
--train_batch_size 1
--gradient_accumulation_steps 1
--max_train_steps 2
--cache_latents
--learning_rate 5.0e-04
--lr_scheduler constant
--lr_warmup_steps 0
--max_sequence_length 8
--text_encoder_out_layers 1
--output_dir {tmpdir}
""".split()
run_command(self._launch_args + test_args)
# save_pretrained smoke test
self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")))
lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))
is_lora = all("lora" in k for k in lora_state_dict.keys())
self.assertTrue(is_lora)
# the cached conditioning latents must still produce an adapter on the reference mode blocks
self.assertTrue(any("image_cond" in key for key in lora_state_dict.keys()))

2210
examples/dreambooth/train_dreambooth_lora_flux3_klein.py Normal file
View File

File diff suppressed because it is too large Load Diff

2080
examples/dreambooth/train_dreambooth_lora_flux3_klein_img2img.py Normal file
View File

File diff suppressed because it is too large Load Diff

112
scripts/convert_flux3_to_diffusers.py Normal file
View File

@@ -0,0 +1,112 @@
import argparse
import json
import safetensors.torch
import torch
from huggingface_hub import hf_hub_download
from diffusers import Flux3Transformer2DModel
from diffusers.loaders.single_file_utils import convert_flux3_transformer_checkpoint_to_diffusers
"""
Converts the FLUX.3 image transformer from a BFL model package (`config.json` + `model.safetensors`) to diffusers.
# from the Hub package
python scripts/convert_flux3_to_diffusers.py \
--original_state_dict_repo_id <release repo id> \
--output_path flux3-image-diffusers
# from local files
python scripts/convert_flux3_to_diffusers.py \
--checkpoint_path /path/to/model.safetensors --config_path /path/to/config.json \
--output_path flux3-image-diffusers
The image VAE (`ae.safetensors`) is the FLUX.2 autoencoder; convert it with the FLUX.2 script:
python scripts/convert_flux2_to_diffusers.py --vae --checkpoint_path /path/to/ae.safetensors --output_path flux3-image-diffusers
The text encoder (`text_encoder/`, `tokenizer/`) is Qwen3-VL in transformers format and needs no conversion.
"""
parser = argparse.ArgumentParser()
parser.add_argument("--original_state_dict_repo_id", default=None, type=str)
parser.add_argument("--dit_filename", default="model.safetensors", type=str)
parser.add_argument("--config_filename", default="config.json", type=str)
parser.add_argument("--checkpoint_path", default=None, type=str)
parser.add_argument(
"--config_path",
default=None,
type=str,
help="the package config.json; the default architecture is used when omitted",
)
parser.add_argument("--dit_dtype", type=str, default="bf16", choices=["bf16", "fp32"])
parser.add_argument("--output_path", type=str, required=True)
def load_original_checkpoint(args):
if args.original_state_dict_repo_id is not None:
ckpt_path = hf_hub_download(repo_id=args.original_state_dict_repo_id, filename=args.dit_filename)
elif args.checkpoint_path is not None:
ckpt_path = args.checkpoint_path
else:
raise ValueError("please provide either `original_state_dict_repo_id` or a local `checkpoint_path`")
return safetensors.torch.load_file(ckpt_path)
def load_package_config(args):
if args.config_path is not None:
config_path = args.config_path
elif args.original_state_dict_repo_id is not None:
config_path = hf_hub_download(repo_id=args.original_state_dict_repo_id, filename=args.config_filename)
else:
return None
with open(config_path) as f:
return json.load(f)
def transformer_config_from_package(package_config):
"""The diffusers config of the package's `dit` section."""
dit = package_config["dit"]
if dit.get("vec_in_dim") is not None:
raise ValueError(
"this package has a pooled-vector embedder (`vec_in_dim`), which Flux3Transformer2DModel does not have"
)
if list(dit["out_modes"]) != ["image"]:
raise ValueError(f"expected an image-only package (`out_modes` == ['image']), got {dit['out_modes']}")
return {
"in_channels": dit["in_channels"]["image"],
"cond_in_channels": dit["in_channels"]["image_cond"],
"num_layers": dit["depth"],
"num_single_layers": dit["depth_single"],
"attention_head_dim": dit["hidden"] // dit["heads"],
"num_attention_heads": dit["heads"],
"joint_attention_dim": dit["in_channels"]["text"],
"mlp_ratio": dit["mlp_ratio"],
"axes_dims_rope": tuple(dit["axes_dim"]),
"rope_theta": dit["theta"],
"flow_map": dit.get("flow_map", False),
}
def main(args):
original_ckpt = load_original_checkpoint(args)
converted_transformer_state_dict = convert_flux3_transformer_checkpoint_to_diffusers(original_ckpt)
package_config = load_package_config(args)
transformer_config = transformer_config_from_package(package_config) if package_config is not None else {}
transformer = Flux3Transformer2DModel(**transformer_config)
transformer.load_state_dict(converted_transformer_state_dict, strict=True)
dtype = torch.bfloat16 if args.dit_dtype == "bf16" else torch.float32
transformer.to(dtype).save_pretrained(f"{args.output_path}/transformer")
# confirm the checkpoint round-trips
Flux3Transformer2DModel.from_pretrained(f"{args.output_path}/transformer")
print(f"saved transformer to {args.output_path}/transformer")
if __name__ == "__main__":
args = parser.parse_args()
main(args)

30
src/diffusers/__init__.py
View File

@@ -134,6 +134,18 @@ except OptionalDependencyNotAvailable:
else:
_import_structure["quantizers.quantization_config"].append("NunchakuLiteQuantizationConfig")
try:
if not is_torch_available():
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils import dummy_bfl_objects
_import_structure["utils.dummy_bfl_objects"] = [
name for name in dir(dummy_bfl_objects) if not name.startswith("_")
]
else:
_import_structure["quantizers.quantization_config"].append("BFLQuantizationConfig")
try:
if not is_auto_round_available():
raise OptionalDependencyNotAvailable()
@@ -288,6 +300,7 @@ else:
"EasyAnimateTransformer3DModel",
"ErnieImageTransformer2DModel",
"Flux2Transformer2DModel",
"Flux3Transformer2DModel",
"FluxControlNetModel",
"FluxMultiControlNetModel",
"FluxTransformer2DModel",
@@ -540,6 +553,10 @@ else:
"Flux2KleinBaseModularPipeline",
"Flux2KleinModularPipeline",
"Flux2ModularPipeline",
"Flux3AutoBlocks",
"Flux3DistilledAutoBlocks",
"Flux3DistilledModularPipeline",
"Flux3ModularPipeline",
"FluxAutoBlocks",
"FluxKontextAutoBlocks",
"FluxKontextModularPipeline",
@@ -1057,6 +1074,14 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
else:
from .quantizers.quantization_config import NunchakuLiteQuantizationConfig
try:
if not is_torch_available():
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from .utils.dummy_bfl_objects import *
else:
from .quantizers.quantization_config import BFLQuantizationConfig
try:
if not is_auto_round_available():
raise OptionalDependencyNotAvailable()
@@ -1190,6 +1215,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
EasyAnimateTransformer3DModel,
ErnieImageTransformer2DModel,
Flux2Transformer2DModel,
Flux3Transformer2DModel,
FluxControlNetModel,
FluxMultiControlNetModel,
FluxTransformer2DModel,
@@ -1421,6 +1447,10 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
Flux2KleinBaseModularPipeline,
Flux2KleinModularPipeline,
Flux2ModularPipeline,
Flux3AutoBlocks,
Flux3DistilledAutoBlocks,
Flux3DistilledModularPipeline,
Flux3ModularPipeline,
FluxAutoBlocks,
FluxKontextAutoBlocks,
FluxKontextModularPipeline,

2
src/diffusers/loaders/__init__.py
View File

@@ -41,6 +41,7 @@ if is_torch_available():
"Krea2LoraLoaderMixin",
"ZImageLoraLoaderMixin",
"Flux2LoraLoaderMixin",
"Flux3LoraLoaderMixin",
"Ideogram4LoraLoaderMixin",
"ErnieImageLoraLoaderMixin",
"CosmosLoraLoaderMixin",
@@ -82,6 +83,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
CosmosLoraLoaderMixin,
ErnieImageLoraLoaderMixin,
Flux2LoraLoaderMixin,
Flux3LoraLoaderMixin,
FluxLoraLoaderMixin,
HeliosLoraLoaderMixin,
HiDreamImageLoraLoaderMixin,

20
src/diffusers/loaders/lora_conversion_utils.py
View File

@@ -2551,6 +2551,26 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict):
return converted_state_dict
def _convert_non_diffusers_flux3_lora_to_diffusers(state_dict):
block_mapping = {
".q_proj.": ".attn.to_q.",
".k_proj.": ".attn.to_k.",
".v_proj.": ".attn.to_v.",
".attn_out.": ".attn.to_out.0.",
".mlp_in.": ".ff.linear_in.",
".mlp_out.": ".ff.linear_out.",
}
converted_state_dict = {}
for key, value in state_dict.items():
key = key.removeprefix("diffusion_model.")
key = key.replace("content_mode_blocks.", "mode_blocks.")
key = key.replace("single_blocks.", "single_transformer_blocks.")
for original_name, diffusers_name in block_mapping.items():
key = key.replace(original_name, diffusers_name)
converted_state_dict[f"transformer.{key}"] = value
return converted_state_dict
def _convert_kohya_flux2_lora_to_diffusers(state_dict):
def _convert_to_ai_toolkit(sds_sd, ait_sd, sds_key, ait_key):
if sds_key + ".lora_down.weight" not in sds_sd:

183
src/diffusers/loaders/lora_pipeline.py
View File

@@ -51,6 +51,7 @@ from .lora_conversion_utils import (
_convert_non_diffusers_anima_lora_to_diffusers,
_convert_non_diffusers_flux2_lokr_to_diffusers,
_convert_non_diffusers_flux2_lora_to_diffusers,
_convert_non_diffusers_flux3_lora_to_diffusers,
_convert_non_diffusers_hidream_lora_to_diffusers,
_convert_non_diffusers_ideogram4_lora_to_diffusers,
_convert_non_diffusers_krea2_lora_to_diffusers,
@@ -5980,6 +5981,188 @@ class Flux2LoraLoaderMixin(LoraBaseMixin):
super().unfuse_lora(components=components, **kwargs)
class Flux3LoraLoaderMixin(LoraBaseMixin):
_lora_loadable_modules = ["transformer"]
transformer_name = TRANSFORMER_NAME
@classmethod
@validate_hf_hub_args
def lora_state_dict(
cls,
pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
**kwargs,
):
cache_dir = kwargs.pop("cache_dir", None)
force_download = kwargs.pop("force_download", False)
proxies = kwargs.pop("proxies", None)
local_files_only = kwargs.pop("local_files_only", None)
token = kwargs.pop("token", None)
revision = kwargs.pop("revision", None)
subfolder = kwargs.pop("subfolder", None)
weight_name = kwargs.pop("weight_name", None)
use_safetensors = kwargs.pop("use_safetensors", None)
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
allow_pickle = False
if use_safetensors is None:
use_safetensors = True
allow_pickle = True
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
state_dict, metadata = _fetch_state_dict(
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
weight_name=weight_name,
use_safetensors=use_safetensors,
local_files_only=local_files_only,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
token=token,
revision=revision,
subfolder=subfolder,
user_agent=user_agent,
allow_pickle=allow_pickle,
)
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
if is_dora_scale_present:
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
logger.warning(warn_msg)
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
if any(k.startswith("diffusion_model.") for k in state_dict):
state_dict = _convert_non_diffusers_flux3_lora_to_diffusers(state_dict)
out = (state_dict, metadata) if return_lora_metadata else state_dict
return out
@require_peft_backend
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
def load_lora_weights(
self,
pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
adapter_name: str | None = None,
hotswap: bool = False,
**kwargs,
):
"""
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details.
"""
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
# if a dict is passed, copy it instead of modifying it inplace
if isinstance(pretrained_model_name_or_path_or_dict, dict):
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
kwargs["return_lora_metadata"] = True
state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
is_correct_format = all("lora" in key for key in state_dict.keys())
if not is_correct_format:
raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.")
self.load_lora_into_transformer(
state_dict,
transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
adapter_name=adapter_name,
metadata=metadata,
_pipeline=self,
low_cpu_mem_usage=low_cpu_mem_usage,
hotswap=hotswap,
)
@classmethod
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer
def load_lora_into_transformer(
cls,
state_dict,
transformer,
adapter_name=None,
_pipeline=None,
low_cpu_mem_usage=False,
hotswap: bool = False,
metadata=None,
):
"""
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details.
"""
# Load the layers corresponding to transformer.
logger.info(f"Loading {cls.transformer_name}.")
transformer.load_lora_adapter(
state_dict,
network_alphas=None,
adapter_name=adapter_name,
metadata=metadata,
_pipeline=_pipeline,
low_cpu_mem_usage=low_cpu_mem_usage,
hotswap=hotswap,
)
@classmethod
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
def save_lora_weights(
cls,
save_directory: str | os.PathLike,
transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
is_main_process: bool = True,
weight_name: str = None,
save_function: Callable = None,
safe_serialization: bool = True,
transformer_lora_adapter_metadata: dict | None = None,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
"""
lora_layers = {}
lora_metadata = {}
if transformer_lora_layers:
lora_layers[cls.transformer_name] = transformer_lora_layers
lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata
if not lora_layers:
raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.")
cls._save_lora_weights(
save_directory=save_directory,
lora_layers=lora_layers,
lora_metadata=lora_metadata,
is_main_process=is_main_process,
weight_name=weight_name,
save_function=save_function,
safe_serialization=safe_serialization,
)
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
def fuse_lora(
self,
components: list[str] = ["transformer"],
lora_scale: float = 1.0,
safe_fusing: bool = False,
adapter_names: list[str] | None = None,
**kwargs,
):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details.
"""
super().fuse_lora(
components=components,
lora_scale=lora_scale,
safe_fusing=safe_fusing,
adapter_names=adapter_names,
**kwargs,
)
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
r"""
See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
"""
super().unfuse_lora(components=components, **kwargs)
class Ideogram4LoraLoaderMixin(LoraBaseMixin):
r"""
Load LoRA layers into [`Ideogram4Transformer2DModel`]. Specific to [`Ideogram4Pipeline`].

8
src/diffusers/loaders/single_file_model.py
View File

@@ -39,6 +39,7 @@ from .single_file_utils import (
convert_cosmos_transformer_checkpoint_to_diffusers,
convert_ernie_image_transformer_checkpoint_to_diffusers,
convert_flux2_transformer_checkpoint_to_diffusers,
convert_flux3_transformer_checkpoint_to_diffusers,
convert_flux_transformer_checkpoint_to_diffusers,
convert_hidream_transformer_to_diffusers,
convert_hunyuan_video_transformer_to_diffusers,
@@ -208,6 +209,10 @@ SINGLE_FILE_LOADABLE_CLASSES = {
"checkpoint_mapping_fn": convert_flux2_transformer_checkpoint_to_diffusers,
"default_subfolder": "transformer",
},
"Flux3Transformer2DModel": {
"checkpoint_mapping_fn": convert_flux3_transformer_checkpoint_to_diffusers,
"default_subfolder": "transformer",
},
"ZImageTransformer2DModel": {
"checkpoint_mapping_fn": convert_z_image_transformer_checkpoint_to_diffusers,
"default_subfolder": "transformer",
@@ -409,9 +414,12 @@ class FromOriginalModelMixin:
user_agent=user_agent,
)
if quantization_config is not None:
if not low_cpu_mem_usage:
raise ValueError("`low_cpu_mem_usage` cannot be False when using quantization.")
hf_quantizer = DiffusersAutoQuantizer.from_config(quantization_config)
hf_quantizer.validate_environment()
torch_dtype = hf_quantizer.update_torch_dtype(torch_dtype)
checkpoint = hf_quantizer.maybe_update_state_dict(checkpoint)
else:
hf_quantizer = None

77
src/diffusers/loaders/single_file_utils.py
View File

@@ -150,6 +150,11 @@ CHECKPOINT_KEY_NAMES = {
"net.pos_embedder.dim_spatial_range",
],
"flux2": ["model.diffusion_model.single_stream_modulation.lin.weight", "single_stream_modulation.lin.weight"],
"flux3": [
"model.diffusion_model.early_stream_modulations.image.lin.weight",
"early_stream_modulations.image.lin.weight",
],
"flux3-distilled": ["model.diffusion_model.ref_time_in.in_layer.weight", "ref_time_in.in_layer.weight"],
"chroma": [
"model.diffusion_model.distilled_guidance_layer.in_proj.bias",
"distilled_guidance_layer.in_proj.bias",
@@ -211,6 +216,8 @@ DIFFUSERS_DEFAULT_PIPELINE_PATHS = {
"flux-depth": {"pretrained_model_name_or_path": "black-forest-labs/FLUX.1-Depth-dev"},
"flux-schnell": {"pretrained_model_name_or_path": "black-forest-labs/FLUX.1-schnell"},
"flux-2-dev": {"pretrained_model_name_or_path": "black-forest-labs/FLUX.2-dev"},
"flux-3-image": {"pretrained_model_name_or_path": "black-forest-labs/flux-3-image-klein-base"},
"flux-3-image-distilled": {"pretrained_model_name_or_path": "black-forest-labs/flux-3-image-klein"},
"chroma": {"pretrained_model_name_or_path": "lodestones/Chroma1-HD"},
"ltx-video": {"pretrained_model_name_or_path": "diffusers/LTX-Video-0.9.0"},
"ltx-video-0.9.1": {"pretrained_model_name_or_path": "diffusers/LTX-Video-0.9.1"},
@@ -694,6 +701,13 @@ def infer_diffusers_model_type(checkpoint):
elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["flux2"]):
model_type = "flux-2-dev"
elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["flux3"]):
# the step-distilled (flow-map) checkpoint carries a second timestep embedder
if any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["flux3-distilled"]):
model_type = "flux-3-image-distilled"
else:
model_type = "flux-3-image"
elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["chroma"]):
model_type = "chroma"
@@ -4000,6 +4014,69 @@ def convert_flux2_transformer_checkpoint_to_diffusers(checkpoint, **kwargs):
return converted_state_dict
def convert_flux3_transformer_checkpoint_to_diffusers(checkpoint, **kwargs):
# inner-module renames shared by the per-mode blocks and the joint single blocks
FLUX3_TRANSFORMER_BLOCK_KEY_MAP = {
"q_proj.weight": "attn.to_q.weight",
"k_proj.weight": "attn.to_k.weight",
"v_proj.weight": "attn.to_v.weight",
"attn_out.weight": "attn.to_out.0.weight",
"norm.query_norm.scale": "attn.norm_q.weight",
"norm.key_norm.scale": "attn.norm_k.weight",
"mlp_in.weight": "ff.linear_in.weight",
"mlp_out.weight": "ff.linear_out.weight",
}
# module-prefix renames
FLUX3_TRANSFORMER_PREFIX_MAP = {
"emb_in.": "x_embedders.",
"txt_in.": "context_embedder.",
"content_mode_blocks.": "mode_blocks.",
"single_blocks.": "single_transformer_blocks.",
}
# exact-key renames
FLUX3_TRANSFORMER_KEYS_RENAME_DICT = {
"time_in.in_layer.weight": "timestep_embedder.linear_1.weight",
"time_in.out_layer.weight": "timestep_embedder.linear_2.weight",
# flow map (step-distilled packages): the embedder of the timestep each step jumps to
"ref_time_in.in_layer.weight": "target_timestep_embedder.linear_1.weight",
"ref_time_in.out_layer.weight": "target_timestep_embedder.linear_2.weight",
"final_layer.image.linear.weight": "proj_out.weight",
# original adaLN weights produce [shift, scale]; Flux3AdaLayerNorm expects [scale, shift]
"final_layer.image.adaLN_modulation.1.weight": "norm_out.linear.weight",
}
converted_state_dict = {}
for key in list(checkpoint.keys()):
# training-only REPA projection heads, not part of the inference model
if key.startswith("repa_projectors."):
checkpoint.pop(key)
continue
value = checkpoint.pop(key)
if key in FLUX3_TRANSFORMER_KEYS_RENAME_DICT:
if key.endswith("adaLN_modulation.1.weight"):
value = swap_scale_shift(value, dim=0)
converted_state_dict[FLUX3_TRANSFORMER_KEYS_RENAME_DICT[key]] = value
continue
new_key = key
for old_prefix, new_prefix in FLUX3_TRANSFORMER_PREFIX_MAP.items():
if new_key.startswith(old_prefix):
new_key = new_prefix + new_key[len(old_prefix) :]
break
new_key = new_key.replace(".lin.weight", ".linear.weight")
for old_suffix, new_suffix in FLUX3_TRANSFORMER_BLOCK_KEY_MAP.items():
if new_key.endswith(old_suffix):
new_key = new_key[: -len(old_suffix)] + new_suffix
break
converted_state_dict[new_key] = value
return converted_state_dict
def convert_z_image_transformer_checkpoint_to_diffusers(checkpoint, **kwargs):
Z_IMAGE_KEYS_RENAME_DICT = {
"final_layer.": "all_final_layer.2-1.",

2
src/diffusers/models/__init__.py
View File

@@ -123,6 +123,7 @@ if is_torch_available():
_import_structure["transformers.transformer_ernie_image"] = ["ErnieImageTransformer2DModel"]
_import_structure["transformers.transformer_flux"] = ["FluxTransformer2DModel"]
_import_structure["transformers.transformer_flux2"] = ["Flux2Transformer2DModel"]
_import_structure["transformers.transformer_flux3"] = ["Flux3Transformer2DModel"]
_import_structure["transformers.transformer_glm_image"] = ["GlmImageTransformer2DModel"]
_import_structure["transformers.transformer_helios"] = ["HeliosTransformer3DModel"]
_import_structure["transformers.transformer_hidream_image"] = ["HiDreamImageTransformer2DModel"]
@@ -270,6 +271,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
EasyAnimateTransformer3DModel,
ErnieImageTransformer2DModel,
Flux2Transformer2DModel,
Flux3Transformer2DModel,
FluxTransformer2DModel,
GlmImageTransformer2DModel,
HeliosTransformer3DModel,

1
src/diffusers/models/transformers/__init__.py
View File

@@ -34,6 +34,7 @@ if is_torch_available():
from .transformer_ernie_image import ErnieImageTransformer2DModel
from .transformer_flux import FluxTransformer2DModel
from .transformer_flux2 import Flux2Transformer2DModel
from .transformer_flux3 import Flux3Transformer2DModel
from .transformer_glm_image import GlmImageTransformer2DModel
from .transformer_helios import HeliosTransformer3DModel
from .transformer_hidream_image import HiDreamImageTransformer2DModel

656
src/diffusers/models/transformers/transformer_flux3.py Normal file
View File

@@ -0,0 +1,656 @@
# Copyright 2026 Black Forest Labs and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any
import torch
import torch.nn as nn
import torch.nn.functional as F
from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
from ...utils import apply_lora_scale, logging
from ...utils.torch_utils import maybe_adjust_dtype_for_device
from ..attention import AttentionMixin, AttentionModuleMixin
from ..attention_dispatch import dispatch_attention_fn
from ..embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed
from ..modeling_outputs import Transformer2DModelOutput
from ..modeling_utils import ModelMixin
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# Copied from diffusers.models.transformers.transformer_flux2.Flux2SwiGLU with Flux2->Flux3
class Flux3SwiGLU(nn.Module):
"""
Flux 2 uses a SwiGLU-style activation in the transformer feedforward sub-blocks, but with the linear projection
layer fused into the first linear layer of the FF sub-block. Thus, this module has no trainable parameters.
"""
def __init__(self):
super().__init__()
self.gate_fn = nn.SiLU()
def forward(self, x: torch.Tensor) -> torch.Tensor:
half = x.shape[-1] // 2
x = self.gate_fn(x[..., :half]) * x[..., half:]
return x
# Copied from diffusers.models.transformers.transformer_flux2.Flux2FeedForward with Flux2->Flux3
class Flux3FeedForward(nn.Module):
def __init__(
self,
dim: int,
dim_out: int | None = None,
mult: float = 3.0,
inner_dim: int | None = None,
bias: bool = False,
):
super().__init__()
if inner_dim is None:
inner_dim = int(dim * mult)
dim_out = dim_out or dim
# Flux3SwiGLU will reduce the dimension by half
self.linear_in = nn.Linear(dim, inner_dim * 2, bias=bias)
self.act_fn = Flux3SwiGLU()
self.linear_out = nn.Linear(inner_dim, dim_out, bias=bias)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.linear_in(x)
x = self.act_fn(x)
x = self.linear_out(x)
return x
class Flux3Modulation(nn.Module):
"""Produces one set of per-token shift/scale/gate modulation parameters from a pre-activated per-token
conditioning embedding of shape `(batch_size, seq_len, dim)`. The per-mode and the joint heads of a token mode
share one activated embedding (activated once by the caller), so this module only holds the projection."""
def __init__(self, dim: int, bias: bool = False):
super().__init__()
self.linear = nn.Linear(dim, dim * 3, bias=bias)
def forward(self, temb: torch.Tensor) -> torch.Tensor:
return self.linear(temb)
def _get_projections(attn: "Flux3Attention", hidden_states: torch.Tensor):
return attn.to_q(hidden_states), attn.to_k(hidden_states), attn.to_v(hidden_states)
def _get_fused_projections(attn: "Flux3Attention", hidden_states: torch.Tensor):
return attn.to_qkv(hidden_states).chunk(3, dim=-1)
def _get_qkv_projections(attn: "Flux3Attention", hidden_states: torch.Tensor):
if attn.fused_projections:
return _get_fused_projections(attn, hidden_states)
return _get_projections(attn, hidden_states)
class Flux3AttnProcessor:
_attention_backend = None
_parallel_config = None
def __call__(
self,
attn: "Flux3Attention",
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
query, key, value = _get_qkv_projections(attn, hidden_states)
query = query.unflatten(-1, (attn.heads, -1))
key = key.unflatten(-1, (attn.heads, -1))
value = value.unflatten(-1, (attn.heads, -1))
query = attn.norm_q(query.float())
key = attn.norm_k(key.float())
if image_rotary_emb is not None:
query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
query, key = query.to(value.dtype), key.to(value.dtype)
hidden_states = dispatch_attention_fn(
query,
key,
value,
attn_mask=attention_mask,
backend=self._attention_backend,
parallel_config=self._parallel_config,
)
hidden_states = hidden_states.flatten(2, 3)
hidden_states = hidden_states.to(query.dtype)
hidden_states = attn.to_out[0](hidden_states)
hidden_states = attn.to_out[1](hidden_states)
return hidden_states
class Flux3Attention(torch.nn.Module, AttentionModuleMixin):
_default_processor_cls = Flux3AttnProcessor
_available_processors = [Flux3AttnProcessor]
def __init__(
self,
query_dim: int,
heads: int = 8,
dim_head: int = 64,
dropout: float = 0.0,
eps: float = 1e-6,
processor=None,
):
super().__init__()
self.head_dim = dim_head
self.inner_dim = dim_head * heads
self.query_dim = query_dim
self.heads = heads
self.dropout = dropout
self.use_bias = False
self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=False)
self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=False)
self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=False)
self.norm_q = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=True)
self.norm_k = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=True)
self.to_out = torch.nn.ModuleList([])
self.to_out.append(torch.nn.Linear(self.inner_dim, query_dim, bias=False))
self.to_out.append(torch.nn.Dropout(dropout))
if processor is None:
processor = self._default_processor_cls()
self.set_processor(processor)
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
return self.processor(self, hidden_states, attention_mask, image_rotary_emb)
class Flux3ModeTransformerBlock(nn.Module):
"""
Per-mode transformer block: parallel attention + feedforward computed from the same modulated input, with
self-attention restricted to a single token mode (the image, the reference images, or the text context).
"""
def __init__(
self,
dim: int,
num_attention_heads: int,
attention_head_dim: int,
mlp_ratio: float = 3.0,
eps: float = 1e-6,
):
super().__init__()
self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
self.attn = Flux3Attention(
query_dim=dim,
heads=num_attention_heads,
dim_head=attention_head_dim,
eps=eps,
)
self.ff = Flux3FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio)
def forward(
self,
hidden_states: torch.Tensor,
temb_mod: torch.Tensor,
image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
) -> torch.Tensor:
# modulation is split here rather than in Flux3Modulation, to avoid passing tuples into checkpoints https://github.com/huggingface/diffusers/issues/12776
mod_shift, mod_scale, mod_gate = torch.chunk(temb_mod, 3, dim=-1)
norm_hidden_states = self.norm(hidden_states)
norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift
attn_output = self.attn(hidden_states=norm_hidden_states, image_rotary_emb=image_rotary_emb)
ff_output = self.ff(norm_hidden_states)
return hidden_states + mod_gate * (attn_output + ff_output)
class Flux3SingleTransformerBlock(nn.Module):
"""
Joint transformer block: text and content tokens are concatenated into a single sequence for parallel attention +
feedforward, with separate per-token modulation for the text and content tokens.
"""
def __init__(
self,
dim: int,
num_attention_heads: int,
attention_head_dim: int,
mlp_ratio: float = 3.0,
eps: float = 1e-6,
):
super().__init__()
self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
self.attn = Flux3Attention(
query_dim=dim,
heads=num_attention_heads,
dim_head=attention_head_dim,
eps=eps,
)
self.ff = Flux3FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio)
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb_mod: torch.Tensor,
temb_mod_context: torch.Tensor,
image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
) -> tuple[torch.Tensor, torch.Tensor]:
# modulation is split here rather than in Flux3Modulation, to avoid passing tuples into checkpoints https://github.com/huggingface/diffusers/issues/12776
mod_shift, mod_scale, mod_gate = torch.chunk(temb_mod, 3, dim=-1)
c_mod_shift, c_mod_scale, c_mod_gate = torch.chunk(temb_mod_context, 3, dim=-1)
norm_hidden_states = self.norm(hidden_states)
norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift
norm_encoder_hidden_states = self.norm(encoder_hidden_states)
norm_encoder_hidden_states = (1 + c_mod_scale) * norm_encoder_hidden_states + c_mod_shift
combined = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1)
attn_output = self.attn(hidden_states=combined, image_rotary_emb=image_rotary_emb)
ff_output = self.ff(combined)
output = attn_output + ff_output
context_output, content_output = torch.split(
output, [encoder_hidden_states.shape[1], hidden_states.shape[1]], dim=1
)
encoder_hidden_states = encoder_hidden_states + c_mod_gate * context_output
hidden_states = hidden_states + mod_gate * content_output
return encoder_hidden_states, hidden_states
# Copied from diffusers.models.transformers.transformer_flux2.Flux2PosEmbed with Flux2->Flux3
class Flux3PosEmbed(nn.Module):
# modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
def __init__(self, theta: int, axes_dim: list[int]):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim
def forward(self, ids: torch.Tensor) -> torch.Tensor:
# Expected ids shape: [S, len(self.axes_dim)]
cos_out = []
sin_out = []
pos = ids.float()
freqs_dtype = maybe_adjust_dtype_for_device(torch.float64, ids.device)
# Unlike Flux 1, loop over len(self.axes_dim) rather than ids.shape[-1]
for i in range(len(self.axes_dim)):
cos, sin = get_1d_rotary_pos_embed(
self.axes_dim[i],
pos[..., i],
theta=self.theta,
repeat_interleave_real=True,
use_real=True,
freqs_dtype=freqs_dtype,
)
cos_out.append(cos)
sin_out.append(sin)
freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
return freqs_cos, freqs_sin
class Flux3AdaLayerNorm(nn.Module):
"""Output adaptive layer norm conditioned on a per-token embedding of shape `(batch_size, seq_len, dim)`."""
def __init__(self, dim: int, eps: float = 1e-6, bias: bool = False):
super().__init__()
self.act_fn = nn.SiLU()
self.linear = nn.Linear(dim, dim * 2, bias=bias)
self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
emb = self.linear(self.act_fn(temb))
scale, shift = torch.chunk(emb, 2, dim=-1)
return self.norm(hidden_states) * (1 + scale) + shift
class Flux3Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, AttentionMixin):
"""
The Transformer model introduced in FLUX.3 for text-to-image and image-to-image generation.
Three token modes share one joint sequence: the noisy image tokens (`hidden_states`), the optional VAE-encoded
reference image tokens (`image_cond_hidden_states`) and the text context (`encoder_hidden_states`). Each mode first
passes through its own per-mode transformer blocks (`mode_blocks` for the image modes, `txt_mode_blocks` for the
text context), then the text context and the active image modes are concatenated for the joint single blocks, and
only the image tokens pass through the output layers. The modulation is per mode: the image tokens are modulated by
the flow timestep, the reference image tokens and the text context by timestep `0`. A flow-map checkpoint
(`flow_map=True`, the step-distilled models) additionally embeds the timestep each denoising step jumps to and adds
it to every mode's modulation.
Args:
in_channels (`int`, defaults to `128`):
Channel width of the image tokens (the packed VAE latents, `latent_channels * 4`).
cond_in_channels (`int`, defaults to `128`):
Channel width of the reference image tokens.
num_layers (`int`, defaults to `5`):
The number of per-mode transformer blocks.
num_single_layers (`int`, defaults to `28`):
The number of joint single transformer blocks.
attention_head_dim (`int`, defaults to `128`):
The number of dimensions to use for each attention head.
num_attention_heads (`int`, defaults to `24`):
The number of attention heads to use.
joint_attention_dim (`int`, defaults to `20480`):
The embedding/channel dimension of `encoder_hidden_states` (stacked Qwen3-VL hidden states).
timestep_channels (`int`, defaults to `256`):
The number of channels of the sinusoidal timestep embedding.
mlp_ratio (`float`, defaults to `3.0`):
The expansion ratio of the feedforward sub-blocks.
axes_dims_rope (`tuple[int]`, defaults to `(32, 32, 32, 32)`):
The dimensions to use for the rotary positional embeddings, per `(t, h, w, l)` position axis.
rope_theta (`int`, defaults to `10000`):
The base frequency of the rotary positional embeddings.
eps (`float`, defaults to `1e-6`):
The epsilon of the layer norms and the QK norms.
flow_map (`bool`, defaults to `False`):
Whether the model is a flow map (step-distilled): adds the `target_timestep_embedder` and requires
`target_timestep` in the forward pass.
"""
_supports_gradient_checkpointing = True
_no_split_modules = ["Flux3ModeTransformerBlock", "Flux3SingleTransformerBlock"]
# QK norm and RoPE run in fp32 with a single cast back, matching the BFL implementation
_keep_in_fp32_modules = ["norm_q", "norm_k"]
_skip_layerwise_casting_patterns = ["pos_embed", "norm"]
_repeated_blocks = ["Flux3ModeTransformerBlock", "Flux3SingleTransformerBlock"]
@register_to_config
def __init__(
self,
in_channels: int = 128,
cond_in_channels: int = 128,
num_layers: int = 5,
num_single_layers: int = 28,
attention_head_dim: int = 128,
num_attention_heads: int = 24,
joint_attention_dim: int = 20480,
timestep_channels: int = 256,
mlp_ratio: float = 3.0,
axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32),
rope_theta: int = 10000,
eps: float = 1e-6,
flow_map: bool = False,
):
super().__init__()
self.inner_dim = num_attention_heads * attention_head_dim
in_channels_per_mode = {"image": in_channels, "image_cond": cond_in_channels}
modes = list(in_channels_per_mode.keys())
# 1. Sinusoidal positional embedding for RoPE on all token modes
self.pos_embed = Flux3PosEmbed(theta=rope_theta, axes_dim=list(axes_dims_rope))
# 2. Per-token timestep embedding
self.time_proj = Timesteps(num_channels=timestep_channels, flip_sin_to_cos=True, downscale_freq_shift=0)
self.timestep_embedder = TimestepEmbedding(
in_channels=timestep_channels, time_embed_dim=self.inner_dim, sample_proj_bias=False
)
# flow map (step-distilled checkpoints): a second embedder for the timestep each step jumps to
self.target_timestep_embedder = (
TimestepEmbedding(in_channels=timestep_channels, time_embed_dim=self.inner_dim, sample_proj_bias=False)
if flow_map
else None
)
# 3. Modulation: one head per mode (+ text) for the per-mode blocks and for the joint blocks
self.early_stream_modulations = nn.ModuleDict(
{name: Flux3Modulation(self.inner_dim, bias=False) for name in [*modes, "txt"]}
)
self.single_stream_modulations = nn.ModuleDict(
{name: Flux3Modulation(self.inner_dim, bias=False) for name in [*modes, "txt"]}
)
# 4. Input projections
self.x_embedders = nn.ModuleDict(
{name: nn.Linear(channels, self.inner_dim, bias=False) for name, channels in in_channels_per_mode.items()}
)
self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim, bias=False)
# 5. Per-mode transformer blocks
self.mode_blocks = nn.ModuleDict(
{
name: nn.ModuleList(
[
Flux3ModeTransformerBlock(
dim=self.inner_dim,
num_attention_heads=num_attention_heads,
attention_head_dim=attention_head_dim,
mlp_ratio=mlp_ratio,
eps=eps,
)
for _ in range(num_layers)
]
)
for name in modes
}
)
self.txt_mode_blocks = nn.ModuleList(
[
Flux3ModeTransformerBlock(
dim=self.inner_dim,
num_attention_heads=num_attention_heads,
attention_head_dim=attention_head_dim,
mlp_ratio=mlp_ratio,
eps=eps,
)
for _ in range(num_layers)
]
)
# 6. Joint single transformer blocks
self.single_transformer_blocks = nn.ModuleList(
[
Flux3SingleTransformerBlock(
dim=self.inner_dim,
num_attention_heads=num_attention_heads,
attention_head_dim=attention_head_dim,
mlp_ratio=mlp_ratio,
eps=eps,
)
for _ in range(num_single_layers)
]
)
# 7. Output layers (image tokens only)
self.norm_out = Flux3AdaLayerNorm(self.inner_dim, eps=eps)
self.proj_out = nn.Linear(self.inner_dim, in_channels, bias=False)
self.gradient_checkpointing = False
@apply_lora_scale("attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
timestep: torch.Tensor,
img_ids: torch.Tensor,
txt_ids: torch.Tensor,
image_cond_hidden_states: torch.Tensor | None = None,
image_cond_ids: torch.Tensor | None = None,
target_timestep: torch.Tensor | None = None,
attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
) -> Transformer2DModelOutput | tuple[torch.Tensor]:
"""
The [`Flux3Transformer2DModel`] forward method.
Args:
hidden_states (`torch.Tensor` of shape `(batch_size, image_seq_len, in_channels)`):
Noisy image latent tokens.
encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_seq_len, joint_attention_dim)`):
Text context embeddings (stacked Qwen3-VL hidden states).
timestep (`torch.Tensor` of shape `(batch_size,)`):
Flow timestep in `[0, 1]` of the image tokens. The reference image tokens and the text context sit at
timestep `0`.
img_ids (`torch.Tensor` of shape `(image_seq_len, 4)`):
Position ids `(t, h, w, l)` of the image tokens, used for the rotary positional embeddings and shared
by all samples in the batch.
txt_ids (`torch.Tensor` of shape `(text_seq_len, 4)`):
Position ids of the text context tokens.
image_cond_hidden_states (`torch.Tensor` of shape `(batch_size, cond_seq_len, cond_in_channels)`, *optional*):
VAE-encoded reference image tokens (clean conditioning content, for image-to-image generation).
image_cond_ids (`torch.Tensor` of shape `(cond_seq_len, 4)`, *optional*):
Position ids of the reference image tokens. Required with `image_cond_hidden_states`.
target_timestep (`torch.Tensor` of shape `(batch_size,)`, *optional*):
Flow timestep in `[0, 1]` that the current denoising step jumps to. Required by, and only by, a
flow-map model (`flow_map=True`).
attention_kwargs (`dict`, *optional*):
A kwargs dictionary whose `scale` entry, if specified, scales the LoRA layers for this forward pass.
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
tuple.
Returns:
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
`tuple` where the first element is the predicted velocity of the image tokens, of the same shape as
`hidden_states`.
"""
if image_cond_hidden_states is not None and image_cond_ids is None:
raise ValueError("`image_cond_hidden_states` requires `image_cond_ids`.")
if image_cond_hidden_states is None and image_cond_ids is not None:
raise ValueError("`image_cond_ids` was passed without `image_cond_hidden_states`.")
if (target_timestep is not None) != bool(self.config.flow_map):
raise ValueError(
"`target_timestep` is required by, and only by, a flow-map transformer (`flow_map=True`)."
)
batch_size, image_seq_len = hidden_states.shape[:2]
dtype = hidden_states.dtype
# 1. Timestep embeddings: the image tokens at `timestep`, the text context and the reference tokens at 0.
# Each is a per-token embedding (batch_size, seq_len, dim) since the modulation is per token.
image_temb = self.time_proj(timestep.to(dtype) * 1000)
image_temb = self.timestep_embedder(image_temb.to(dtype))[:, None, :]
zero_temb = self.time_proj(torch.zeros(batch_size, device=hidden_states.device, dtype=dtype))
zero_temb = self.timestep_embedder(zero_temb.to(dtype))[:, None, :]
if target_timestep is not None:
# flow map: the embedding of the timestep this step jumps to is added to every mode's embedding, so it
# reaches the per-mode blocks, the joint blocks and the output layer
target_temb = self.time_proj(target_timestep.to(dtype) * 1000)
target_temb = self.target_timestep_embedder(target_temb.to(dtype))[:, None, :]
image_temb = image_temb + target_temb
zero_temb = zero_temb + target_temb
image_temb = image_temb.expand(batch_size, image_seq_len, -1)
zero_temb_act = F.silu(zero_temb)
# 2. Text context: the two modulation sets, input projection and RoPE
txt_temb_act = zero_temb_act.expand(batch_size, encoder_hidden_states.shape[1], -1)
txt_early_mod = self.early_stream_modulations["txt"](txt_temb_act)
txt_single_mod = self.single_stream_modulations["txt"](txt_temb_act)
encoder_hidden_states = self.context_embedder(encoder_hidden_states)
txt_rotary_emb = self.pos_embed(txt_ids)
# 3. Image tokens
image_temb_act = F.silu(image_temb)
image_early_mod = self.early_stream_modulations["image"](image_temb_act)
image_single_mod = self.single_stream_modulations["image"](image_temb_act)
hidden_states = self.x_embedders["image"](hidden_states)
image_rotary_emb = self.pos_embed(img_ids)
# 4. Reference image tokens (optional)
if image_cond_hidden_states is not None:
cond_temb_act = zero_temb_act.expand(batch_size, image_cond_hidden_states.shape[1], -1)
cond_early_mod = self.early_stream_modulations["image_cond"](cond_temb_act)
cond_single_mod = self.single_stream_modulations["image_cond"](cond_temb_act)
image_cond_hidden_states = self.x_embedders["image_cond"](image_cond_hidden_states)
cond_rotary_emb = self.pos_embed(image_cond_ids)
# 5. Per-mode transformer blocks (each mode self-attends within itself only)
for block in self.txt_mode_blocks:
if torch.is_grad_enabled() and self.gradient_checkpointing:
encoder_hidden_states = self._gradient_checkpointing_func(
block, encoder_hidden_states, txt_early_mod, txt_rotary_emb
)
else:
encoder_hidden_states = block(encoder_hidden_states, txt_early_mod, txt_rotary_emb)
for block in self.mode_blocks["image"]:
if torch.is_grad_enabled() and self.gradient_checkpointing:
hidden_states = self._gradient_checkpointing_func(
block, hidden_states, image_early_mod, image_rotary_emb
)
else:
hidden_states = block(hidden_states, image_early_mod, image_rotary_emb)
if image_cond_hidden_states is not None:
for block in self.mode_blocks["image_cond"]:
if torch.is_grad_enabled() and self.gradient_checkpointing:
image_cond_hidden_states = self._gradient_checkpointing_func(
block, image_cond_hidden_states, cond_early_mod, cond_rotary_emb
)
else:
image_cond_hidden_states = block(image_cond_hidden_states, cond_early_mod, cond_rotary_emb)
# 6. Joint single transformer blocks over [txt, image, image_cond]
if image_cond_hidden_states is not None:
content_states = torch.cat([hidden_states, image_cond_hidden_states], dim=1)
content_mod = torch.cat([image_single_mod, cond_single_mod], dim=1)
joint_rotary_emb = (
torch.cat([txt_rotary_emb[0], image_rotary_emb[0], cond_rotary_emb[0]], dim=0),
torch.cat([txt_rotary_emb[1], image_rotary_emb[1], cond_rotary_emb[1]], dim=0),
)
else:
content_states = hidden_states
content_mod = image_single_mod
joint_rotary_emb = (
torch.cat([txt_rotary_emb[0], image_rotary_emb[0]], dim=0),
torch.cat([txt_rotary_emb[1], image_rotary_emb[1]], dim=0),
)
for block in self.single_transformer_blocks:
if torch.is_grad_enabled() and self.gradient_checkpointing:
encoder_hidden_states, content_states = self._gradient_checkpointing_func(
block, content_states, encoder_hidden_states, content_mod, txt_single_mod, joint_rotary_emb
)
else:
encoder_hidden_states, content_states = block(
hidden_states=content_states,
encoder_hidden_states=encoder_hidden_states,
temb_mod=content_mod,
temb_mod_context=txt_single_mod,
image_rotary_emb=joint_rotary_emb,
)
# 7. Output layers: the image tokens sit first in the content sequence; the reference tokens produce no output
hidden_states = content_states[:, :image_seq_len]
hidden_states = self.norm_out(hidden_states, image_temb)
output = self.proj_out(hidden_states)
if not return_dict:
return (output,)
return Transformer2DModelOutput(sample=output)

12
src/diffusers/modular_pipelines/__init__.py
View File

@@ -81,6 +81,12 @@ else:
"Flux2KleinModularPipeline",
"Flux2KleinBaseModularPipeline",
]
_import_structure["flux3"] = [
"Flux3AutoBlocks",
"Flux3DistilledAutoBlocks",
"Flux3DistilledModularPipeline",
"Flux3ModularPipeline",
]
_import_structure["ideogram4"] = [
"Ideogram4AutoBlocks",
"Ideogram4ModularPipeline",
@@ -179,6 +185,12 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
Flux2KleinModularPipeline,
Flux2ModularPipeline,
)
from .flux3 import (
Flux3AutoBlocks,
Flux3DistilledAutoBlocks,
Flux3DistilledModularPipeline,
Flux3ModularPipeline,
)
from .helios import (
HeliosAutoBlocks,
HeliosModularPipeline,

49
src/diffusers/modular_pipelines/flux3/__init__.py Normal file
View File

@@ -0,0 +1,49 @@
from typing import TYPE_CHECKING
from ...utils import (
DIFFUSERS_SLOW_IMPORT,
OptionalDependencyNotAvailable,
_LazyModule,
get_objects_from_module,
is_torch_available,
is_transformers_available,
)
_dummy_objects = {}
_import_structure = {}
try:
if not (is_transformers_available() and is_torch_available()):
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from ...utils import dummy_torch_and_transformers_objects # noqa F403
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
else:
_import_structure["modular_blocks_flux3"] = ["Flux3AutoBlocks"]
_import_structure["modular_blocks_flux3_distilled"] = ["Flux3DistilledAutoBlocks"]
_import_structure["modular_pipeline"] = ["Flux3DistilledModularPipeline", "Flux3ModularPipeline"]
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
try:
if not (is_transformers_available() and is_torch_available()):
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
else:
from .modular_blocks_flux3 import Flux3AutoBlocks
from .modular_blocks_flux3_distilled import Flux3DistilledAutoBlocks
from .modular_pipeline import Flux3DistilledModularPipeline, Flux3ModularPipeline
else:
import sys
sys.modules[__name__] = _LazyModule(
__name__,
globals()["__file__"],
_import_structure,
module_spec=__spec__,
)
for name, value in _dummy_objects.items():
setattr(sys.modules[__name__], name, value)

436
src/diffusers/modular_pipelines/flux3/before_denoise.py Normal file
View File

@@ -0,0 +1,436 @@
# Copyright 2026 Black Forest Labs and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import numpy as np
import torch
from ...configuration_utils import FrozenDict
from ...schedulers import FlowMatchEulerDiscreteScheduler
from ...utils import logging
from ...utils.torch_utils import randn_tensor
from ..modular_pipeline import ModularPipelineBlocks, PipelineState
from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
from .modular_pipeline import Flux3ModularPipeline
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
num_inference_steps: int | None = None,
device: str | torch.device | None = None,
timesteps: list[int] | None = None,
sigmas: list[float] | None = None,
**kwargs,
):
r"""
Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
Args:
scheduler (`SchedulerMixin`):
The scheduler to get timesteps from.
num_inference_steps (`int`):
The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
timesteps (`list[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
sigmas (`list[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
`tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
if timesteps is not None:
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
if not accepts_timesteps:
raise ValueError(
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
f" timestep schedules. Please check whether you are using the correct scheduler."
)
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
timesteps = scheduler.timesteps
num_inference_steps = len(timesteps)
elif sigmas is not None:
accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
if not accept_sigmas:
raise ValueError(
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
f" sigmas schedules. Please check whether you are using the correct scheduler."
)
scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
timesteps = scheduler.timesteps
num_inference_steps = len(timesteps)
else:
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
timesteps = scheduler.timesteps
return timesteps, num_inference_steps
class Flux3SetTimestepsStep(ModularPipelineBlocks):
model_name = "flux3"
def __init__(self, shift: float = 6.93, num_inference_steps: int = 50):
# the constant timestep shift ("alpha") of the FLUX.3 sampling recipe, configured on the scheduler, and the
# default number of steps (50; 16 for the step-distilled checkpoints)
self.shift = shift
self.num_inference_steps = num_inference_steps
super().__init__()
@property
def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"scheduler",
FlowMatchEulerDiscreteScheduler,
config=FrozenDict({"shift": self.shift}),
default_creation_method="from_config",
),
]
@property
def description(self) -> str:
return (
"Step that sets the scheduler's timesteps for FLUX.3 inference (constant shift, configured on the "
"scheduler)."
)
@property
def inputs(self) -> list[InputParam]:
return [
InputParam.template("num_inference_steps", default=self.num_inference_steps),
InputParam.template("timesteps", note="overrides `num_inference_steps` when provided"),
InputParam.template("sigmas", note="overrides `num_inference_steps` when provided"),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"),
OutputParam(
"num_inference_steps",
type_hint=int,
description="The number of denoising steps to perform at inference time",
),
]
@torch.no_grad()
def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
device = components._execution_device
scheduler = components.scheduler
num_inference_steps = block_state.num_inference_steps
sigmas = block_state.sigmas
timesteps = block_state.timesteps
if timesteps is None and sigmas is None:
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
timesteps, num_inference_steps = retrieve_timesteps(
scheduler,
num_inference_steps,
device,
timesteps=timesteps,
sigmas=sigmas,
)
block_state.timesteps = timesteps
block_state.num_inference_steps = num_inference_steps
components.scheduler.set_begin_index(0)
self.set_block_state(state, block_state)
return components, state
def _prepare_image_ids(height: int, width: int, image_rope_scale: float, l_coord: int = 0) -> torch.Tensor:
"""4D position ids `(t, h, w, l)` of shape `(height * width, 4)` for an image token grid: `t=0`, the H/W axes
scaled by the FLUX.3 image RoPE scale, and a constant `l` (0 for the generated image, negative for references)."""
t = torch.arange(1)
h = torch.arange(height)
w = torch.arange(width)
l = torch.tensor([l_coord])
ids = torch.cartesian_prod(t, h, w, l).to(dtype=torch.float32)
ids[:, 1:3] = ids[:, 1:3] * image_rope_scale
return ids
class Flux3PrepareLatentsStep(ModularPipelineBlocks):
model_name = "flux3"
def __init__(self, default_height: int = 1024, default_width: int = 1024, image_rope_scale: float = 0.5):
# text-to-image size; image-to-image derives its size from the reference image upstream (its branch declares
# no default, each branch of the auto blocks resolves its own)
self.default_height = default_height
self.default_width = default_width
# the model was trained with the image H/W position ids scaled by this factor
self.image_rope_scale = image_rope_scale
super().__init__()
@property
def description(self) -> str:
return "Prepare latents step that prepares the initial noise image latents as tokens, with their position ids."
@property
def inputs(self) -> list[InputParam]:
return [
InputParam.template("height", default=self.default_height),
InputParam.template("width", default=self.default_width),
InputParam.template("latents"),
InputParam.template("num_images_per_prompt"),
InputParam.template("generator"),
InputParam(
"batch_size",
required=True,
type_hint=int,
description="Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.",
),
InputParam("dtype", type_hint=torch.dtype, description="The dtype of the model inputs"),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
),
OutputParam(
"img_ids",
kwargs_type="denoiser_input_fields",
type_hint=torch.Tensor,
description="Position IDs (T, H, W, L) for the image latent tokens (for RoPE)",
),
]
@staticmethod
def check_inputs(height: int, width: int, image_scale_factor: int):
if height % image_scale_factor != 0 or width % image_scale_factor != 0:
raise ValueError(
f"`height` and `width` have to be divisible by {image_scale_factor} but are {height} and {width}."
)
@torch.no_grad()
def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
device = components._execution_device
self.check_inputs(block_state.height, block_state.width, components.vae_scale_factor * 2)
batch_size = block_state.batch_size * block_state.num_images_per_prompt
# image tokens sit on a (height // 16, width // 16) grid with vae_latent_channels * 4 channels
# (the VAE downsamples 8x and its latents are packed 2x2)
latent_height = block_state.height // (components.vae_scale_factor * 2)
latent_width = block_state.width // (components.vae_scale_factor * 2)
latents = block_state.latents
if latents is None:
shape = (batch_size, components.num_channels_latents, latent_height, latent_width)
if isinstance(block_state.generator, list) and len(block_state.generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(block_state.generator)}, but requested an"
f" effective batch size of {batch_size}. Make sure the batch size matches the length of the"
" generators."
)
# draw in float32 then cast: a reduced-precision draw consumes the RNG stream differently,
# which would break seed-compatibility with the reference implementation
latents = randn_tensor(shape, generator=block_state.generator, device=device, dtype=torch.float32)
latents = latents.to(block_state.dtype)
else:
latents = latents.to(device=device, dtype=block_state.dtype)
block_state.latents = latents.flatten(2).transpose(1, 2) # pack to tokens: (B, seq_len, C)
block_state.img_ids = _prepare_image_ids(latent_height, latent_width, self.image_rope_scale).to(device)
self.set_block_state(state, block_state)
return components, state
class Flux3PrepareImageCondLatentsStep(ModularPipelineBlocks):
model_name = "flux3"
def __init__(self, image_rope_scale: float = 0.5):
self.image_rope_scale = image_rope_scale
super().__init__()
@property
def description(self) -> str:
return (
"Step that packs the reference image latents into one token sequence with their position ids for FLUX.3 "
"image-to-image generation. The references are told apart on the `l` position axis (-1, -2, ...)."
)
@property
def inputs(self) -> list[InputParam]:
return [
InputParam(
"image_cond_latents",
required=True,
type_hint=list[torch.Tensor],
description="Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16) per image",
),
InputParam("dtype", type_hint=torch.dtype, description="The dtype of the model inputs"),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"image_cond_hidden_states",
kwargs_type="denoiser_input_fields",
type_hint=torch.Tensor,
description="Reference image latent tokens, shape (batch_size, cond_seq_len, C)",
),
OutputParam(
"image_cond_ids",
kwargs_type="denoiser_input_fields",
type_hint=torch.Tensor,
description="Position IDs (T, H, W, L) of the reference image tokens, shape (cond_seq_len, 4)",
),
]
@torch.no_grad()
def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
device = components._execution_device
cond_tokens, cond_ids = [], []
for index, latents in enumerate(block_state.image_cond_latents):
latent_height, latent_width = latents.shape[-2:]
cond_tokens.append(latents.flatten(2).transpose(1, 2)) # pack to tokens: (B, seq_len, C)
cond_ids.append(
_prepare_image_ids(latent_height, latent_width, self.image_rope_scale, l_coord=-(1 + index))
)
block_state.image_cond_hidden_states = torch.cat(cond_tokens, dim=1).to(device=device, dtype=block_state.dtype)
block_state.image_cond_ids = torch.cat(cond_ids, dim=0).to(device)
self.set_block_state(state, block_state)
return components, state
class Flux3TextRoPEInputsStep(ModularPipelineBlocks):
model_name = "flux3"
@property
def description(self) -> str:
return (
"Step that prepares the 4D RoPE position IDs of the text tokens for FLUX.3 denoising. Place it after "
"the text encoder step."
)
@property
def inputs(self) -> list[InputParam]:
return [
InputParam.template(
"prompt_embeds",
description="Text embeddings used to guide the generation. Can be generated from text_encoder step.",
),
InputParam.template(
"negative_prompt_embeds",
description="Negative text embeddings used to guide the generation. Can be generated from text_encoder step.",
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="txt_ids",
kwargs_type="denoiser_input_fields",
type_hint=torch.Tensor,
description="4D position IDs (T, H, W, L) for text tokens, used for RoPE calculation.",
),
OutputParam(
name="negative_txt_ids",
kwargs_type="denoiser_input_fields",
type_hint=torch.Tensor,
description="4D position IDs (T, H, W, L) for negative text tokens, used for RoPE calculation.",
),
]
@staticmethod
def _prepare_text_ids(x: torch.Tensor):
"""Prepare 4D position IDs (T, H, W, L) of shape (seq_len, 4) for text tokens; text advances along the L axis."""
seq_len = x.shape[1]
t = torch.arange(1)
h = torch.arange(1)
w = torch.arange(1)
seq_l = torch.arange(seq_len)
return torch.cartesian_prod(t, h, w, seq_l).to(dtype=torch.float32)
def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
device = block_state.prompt_embeds.device
block_state.txt_ids = self._prepare_text_ids(block_state.prompt_embeds).to(device)
block_state.negative_txt_ids = None
if block_state.negative_prompt_embeds is not None:
block_state.negative_txt_ids = self._prepare_text_ids(block_state.negative_prompt_embeds).to(device)
self.set_block_state(state, block_state)
return components, state
class Flux3DistilledTextRoPEInputsStep(Flux3TextRoPEInputsStep):
model_name = "flux3-distilled"
@property
def description(self) -> str:
return (
"Step that prepares the 4D RoPE position IDs of the text tokens for the step-distilled FLUX.3 model (no "
"negative prompt). Place it after the text encoder step."
)
@property
def inputs(self) -> list[InputParam]:
return [
InputParam.template(
"prompt_embeds",
description="Text embeddings used to guide the generation. Can be generated from text_encoder step.",
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
name="txt_ids",
kwargs_type="denoiser_input_fields",
type_hint=torch.Tensor,
description="4D position IDs (T, H, W, L) for text tokens, used for RoPE calculation.",
),
]
def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
block_state.txt_ids = self._prepare_text_ids(block_state.prompt_embeds).to(block_state.prompt_embeds.device)
self.set_block_state(state, block_state)
return components, state

170
src/diffusers/modular_pipelines/flux3/decoders.py Normal file
View File

@@ -0,0 +1,170 @@
# Copyright 2026 Black Forest Labs and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import Any, Union
import numpy as np
import PIL
import torch
from ...configuration_utils import FrozenDict
from ...models import AutoencoderKLFlux2
from ...utils import logging
from ..modular_pipeline import ModularPipelineBlocks, PipelineState
from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
from .image_processor import Flux3ImageProcessor
from .modular_pipeline import Flux3ModularPipeline
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
class Flux3UnpackLatentsStep(ModularPipelineBlocks):
model_name = "flux3"
@property
def description(self) -> str:
return "Step that unpacks the image latent tokens from the denoising step back onto their 2D grid."
@property
def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam(
"latents",
required=True,
type_hint=torch.Tensor,
description="The denoised image latent tokens from the denoising step. Shape: (B, seq_len, C)",
),
# no default: the size is already baked into the latents, every workflow sets these upstream
InputParam(
"height",
type_hint=int,
description="The height in pixels of the generated image; must match the size the latents were prepared at",
),
InputParam(
"width",
type_hint=int,
description="The width in pixels of the generated image; must match the size the latents were prepared at",
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"latents",
type_hint=torch.Tensor,
description="The denoised latents, unpacked to (B, C, height, width).",
)
]
@torch.no_grad()
def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
latent_height = block_state.height // (components.vae_scale_factor * 2)
latent_width = block_state.width // (components.vae_scale_factor * 2)
if latent_height * latent_width != block_state.latents.shape[1]:
raise ValueError(
f"`latents` has {block_state.latents.shape[1]} tokens, which does not match "
f"height x width = {block_state.height} x {block_state.width} ({latent_height * latent_width} tokens)"
)
# unpack: (batch_size, latent_height * latent_width, num_channels) -> (batch_size, num_channels, latent_height, latent_width)
batch_size, _, num_channels = block_state.latents.shape
block_state.latents = block_state.latents.view(batch_size, latent_height, latent_width, num_channels).permute(
0, 3, 1, 2
)
self.set_block_state(state, block_state)
return components, state
class Flux3DecodeStep(ModularPipelineBlocks):
model_name = "flux3"
@property
def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKLFlux2),
ComponentSpec(
"image_processor",
Flux3ImageProcessor,
config=FrozenDict({"vae_scale_factor": 16, "vae_latent_channels": 32}),
default_creation_method="from_config",
),
]
@property
def description(self) -> str:
return (
"Step that decodes the denoised latents into images using the FLUX.3 image VAE (the FLUX.2 packed "
"autoencoder) with batch norm denormalization."
)
@property
def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam.template("output_type"),
InputParam(
"latents",
required=True,
type_hint=torch.Tensor,
description="The denoised latents from the denoising step",
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"images",
type_hint=Union[list[PIL.Image.Image], torch.Tensor, np.ndarray],
description="The generated images, can be a list of PIL.Image.Image, torch.Tensor or a numpy array",
)
]
@staticmethod
def _unpatchify_latents(latents):
"""Convert patchified latents back to regular format."""
batch_size, num_channels_latents, height, width = latents.shape
latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), 2, 2, height, width)
latents = latents.permute(0, 1, 4, 2, 5, 3)
latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), height * 2, width * 2)
return latents
@torch.no_grad()
def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
vae = components.vae
latents = block_state.latents
latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to(
latents.device, latents.dtype
)
latents = latents * latents_bn_std + latents_bn_mean
latents = self._unpatchify_latents(latents)
block_state.images = vae.decode(latents.to(vae.dtype), return_dict=False)[0]
block_state.images = components.image_processor.postprocess(
block_state.images, output_type=block_state.output_type
)
self.set_block_state(state, block_state)
return components, state

426
src/diffusers/modular_pipelines/flux3/denoise.py Normal file
View File

@@ -0,0 +1,426 @@
# Copyright 2026 Black Forest Labs and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
from typing import Any
import torch
from ...configuration_utils import FrozenDict
from ...guiders import ClassifierFreeGuidance
from ...models import Flux3Transformer2DModel
from ...schedulers import FlowMatchEulerDiscreteScheduler
from ...utils import is_torch_xla_available, logging
from ..modular_pipeline import (
BlockState,
LoopSequentialPipelineBlocks,
ModularPipelineBlocks,
PipelineState,
)
from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
from .modular_pipeline import Flux3ModularPipeline
if is_torch_xla_available():
import torch_xla.core.xla_model as xm
XLA_AVAILABLE = True
else:
XLA_AVAILABLE = False
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
class Flux3LoopDenoiser(ModularPipelineBlocks):
model_name = "flux3"
def __init__(self, guidance_scale: float = 3.5):
self.guidance_scale = guidance_scale
super().__init__()
@property
def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("transformer", Flux3Transformer2DModel),
ComponentSpec(
"guider",
ClassifierFreeGuidance,
config=FrozenDict({"guidance_scale": self.guidance_scale}),
default_creation_method="from_config",
),
]
@property
def description(self) -> str:
return (
"Step within the denoising loop that predicts the velocity of the image tokens. The static conditioning "
"fields (`img_ids`, the reference image tokens and their ids) arrive via `denoiser_input_fields`; "
"everything matching the transformer signature is forwarded. "
"This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
"object (e.g. `Flux3DenoiseStep`)"
)
@property
def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam(
"latents",
required=True,
type_hint=torch.Tensor,
description="The image latent tokens to denoise. Shape: (B, seq_len, C)",
),
InputParam.template("denoiser_input_fields"),
InputParam.template("attention_kwargs", kwargs_type="denoiser_input_fields"),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam("noise_pred", type_hint=torch.Tensor, description="The predicted velocity for this step")]
@torch.no_grad()
def __call__(
self, components: Flux3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
) -> PipelineState:
denoiser_input_fields = dict(block_state.denoiser_input_fields)
guider_inputs = {
"encoder_hidden_states": (
denoiser_input_fields.pop("prompt_embeds", None),
denoiser_input_fields.pop("negative_prompt_embeds", None),
),
"txt_ids": (
denoiser_input_fields.pop("txt_ids", None),
denoiser_input_fields.pop("negative_txt_ids", None),
),
}
# the remaining collected fields that match the forward signature go straight to the transformer (`img_ids`,
# and the `image_cond_*` fields of image-to-image generation)
transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys())
transformer_kwargs = {}
for field_name in list(denoiser_input_fields):
if field_name in transformer_args:
transformer_kwargs[field_name] = denoiser_input_fields.pop(field_name)
if denoiser_input_fields and i == 0:
logger.warning(f"`denoiser_input_fields` not passed to the transformer: {sorted(denoiser_input_fields)}")
latents = block_state.latents.to(components.transformer.dtype)
# the flow timestep in [0, 1] (scheduler timesteps are sigmas * 1000); the transformer applies it to every
# image token while the reference image tokens and the text context stay at 0
timestep = (t / 1000).expand(latents.shape[0]).to(latents.dtype)
components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
guider_state = components.guider.prepare_inputs(guider_inputs)
for guider_state_batch in guider_state:
components.guider.prepare_models(components.transformer)
cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()}
guider_state_batch.noise_pred = components.transformer(
hidden_states=latents,
timestep=timestep,
return_dict=False,
**transformer_kwargs,
**cond_kwargs,
)[0]
components.guider.cleanup_models(components.transformer)
block_state.noise_pred = components.guider(guider_state)[0]
return components, block_state
class Flux3DistilledLoopDenoiser(ModularPipelineBlocks):
model_name = "flux3-distilled"
@property
def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("transformer", Flux3Transformer2DModel)]
@property
def description(self) -> str:
return (
"Step within the denoising loop that predicts the velocity of the image tokens with a step-distilled "
"(flow-map) FLUX.3 transformer: no classifier-free guidance, and each step also passes the transformer the "
"timestep it jumps to. The static conditioning fields (`prompt_embeds`, `txt_ids`, `img_ids`, the "
"reference image tokens and their ids) arrive via `denoiser_input_fields`. "
"This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
"object (e.g. `Flux3DistilledDenoiseStep`)"
)
@property
def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam(
"latents",
required=True,
type_hint=torch.Tensor,
description="The image latent tokens to denoise. Shape: (B, seq_len, C)",
),
InputParam.template("denoiser_input_fields"),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam("noise_pred", type_hint=torch.Tensor, description="The predicted velocity for this step")]
@torch.no_grad()
def __call__(
self, components: Flux3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
) -> PipelineState:
denoiser_input_fields = dict(block_state.denoiser_input_fields)
prompt_embeds = denoiser_input_fields.pop("prompt_embeds")
# the remaining collected fields that match the forward signature go straight to the transformer (`img_ids`,
# `txt_ids`, and the `image_cond_*` fields of image-to-image generation)
transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys())
transformer_kwargs = {}
for field_name in list(denoiser_input_fields):
if field_name in transformer_args:
transformer_kwargs[field_name] = denoiser_input_fields.pop(field_name)
if denoiser_input_fields and i == 0:
logger.warning(f"`denoiser_input_fields` not passed to the transformer: {sorted(denoiser_input_fields)}")
latents = block_state.latents.to(components.transformer.dtype)
batch_size = latents.shape[0]
# the flow timestep in [0, 1] (scheduler timesteps are sigmas * 1000), and the one this step jumps to: the next
# scheduler timestep, 0 after the last
timesteps = block_state.timesteps
next_t = timesteps[i + 1] if i + 1 < len(timesteps) else torch.zeros_like(t)
timestep = (t / 1000).expand(batch_size).to(latents.dtype)
target_timestep = (next_t / 1000).expand(batch_size).to(latents.dtype)
block_state.noise_pred = components.transformer(
hidden_states=latents,
encoder_hidden_states=prompt_embeds,
timestep=timestep,
target_timestep=target_timestep,
return_dict=False,
**transformer_kwargs,
)[0]
return components, block_state
class Flux3LoopAfterDenoiser(ModularPipelineBlocks):
model_name = "flux3"
@property
def expected_components(self) -> list[ComponentSpec]:
# config-less: the scheduler config is owned by `Flux3SetTimestepsStep`, whose spec is combined first
return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler, default_creation_method="from_config")]
@property
def description(self) -> str:
return (
"Step within the denoising loop that updates the latents after denoising. "
"This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
"object (e.g. `Flux3DenoiseLoopWrapper`)"
)
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")]
@torch.no_grad()
def __call__(
self, components: Flux3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
) -> tuple[Flux3ModularPipeline, BlockState]:
latents_dtype = block_state.latents.dtype
block_state.latents = components.scheduler.step(
block_state.noise_pred,
t,
block_state.latents,
return_dict=False,
)[0]
if block_state.latents.dtype != latents_dtype:
if torch.backends.mps.is_available():
block_state.latents = block_state.latents.to(latents_dtype)
return components, block_state
class Flux3DenoiseLoopWrapper(LoopSequentialPipelineBlocks):
model_name = "flux3"
@property
def description(self) -> str:
return (
"Pipeline block that iteratively denoises the latents over `timesteps`. "
"The specific steps within each iteration can be customized with `sub_blocks` attribute"
)
@property
def loop_expected_components(self) -> list[ComponentSpec]:
return [
# config-less: the scheduler config is owned by `Flux3SetTimestepsStep`, whose spec is combined first
ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler, default_creation_method="from_config"),
ComponentSpec("transformer", Flux3Transformer2DModel),
]
@property
def loop_inputs(self) -> list[InputParam]:
return [
InputParam(
"timesteps",
required=True,
type_hint=torch.Tensor,
description="The timesteps to use for the denoising process.",
),
InputParam(
"num_inference_steps",
required=True,
type_hint=int,
description="The number of inference steps to use for the denoising process.",
),
]
@torch.no_grad()
def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
block_state.num_warmup_steps = max(
len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0
)
with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
for i, t in enumerate(block_state.timesteps):
components, block_state = self.loop_step(components, block_state, i=i, t=t)
if i == len(block_state.timesteps) - 1 or (
(i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
):
progress_bar.update()
if XLA_AVAILABLE:
xm.mark_step()
self.set_block_state(state, block_state)
return components, state
# auto_docstring
class Flux3DenoiseStep(Flux3DenoiseLoopWrapper):
"""
Denoise step that iteratively denoises the latents for FLUX.3.
Its loop logic is defined in `Flux3DenoiseLoopWrapper.__call__` method At each iteration, it runs blocks defined
in `sub_blocks` sequentially:
- `Flux3LoopDenoiser`
- `Flux3LoopAfterDenoiser`
This block supports text-to-image and image-to-image generation.
Components:
transformer (`Flux3Transformer2DModel`) guider (`ClassifierFreeGuidance`) scheduler
(`FlowMatchEulerDiscreteScheduler`)
Inputs:
timesteps (`Tensor`):
The timesteps to use for the denoising process.
num_inference_steps (`int`):
The number of inference steps to use for the denoising process.
latents (`Tensor`):
The image latent tokens to denoise. Shape: (B, seq_len, C)
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
Outputs:
latents (`Tensor`):
The denoised latents
"""
block_classes = [Flux3LoopDenoiser(), Flux3LoopAfterDenoiser()]
block_names = ["denoiser", "after_denoiser"]
@property
def description(self) -> str:
return (
"Denoise step that iteratively denoises the latents for FLUX.3. \n"
"Its loop logic is defined in `Flux3DenoiseLoopWrapper.__call__` method \n"
"At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
" - `Flux3LoopDenoiser`\n"
" - `Flux3LoopAfterDenoiser`\n"
"This block supports text-to-image and image-to-image generation."
)
# auto_docstring
class Flux3DistilledDenoiseStep(Flux3DenoiseLoopWrapper):
"""
Denoise step that iteratively denoises the latents for the step-distilled (flow-map) FLUX.3 model.
Its loop logic is defined in `Flux3DenoiseLoopWrapper.__call__` method At each iteration, it runs blocks defined
in `sub_blocks` sequentially:
- `Flux3DistilledLoopDenoiser`
- `Flux3LoopAfterDenoiser`
This block supports text-to-image and image-to-image generation, without classifier-free guidance.
Components:
transformer (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
Inputs:
timesteps (`Tensor`):
The timesteps to use for the denoising process.
num_inference_steps (`int`):
The number of inference steps to use for the denoising process.
latents (`Tensor`):
The image latent tokens to denoise. Shape: (B, seq_len, C)
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
Outputs:
latents (`Tensor`):
The denoised latents
"""
model_name = "flux3-distilled"
"""
Denoise step that iteratively denoises the latents for the step-distilled (flow-map) FLUX.3 model.
Its loop logic is defined in `Flux3DenoiseLoopWrapper.__call__` method At each iteration, it runs blocks defined
in `sub_blocks` sequentially:
- `Flux3DistilledLoopDenoiser`
- `Flux3LoopAfterDenoiser`
This block supports text-to-image and image-to-image generation, without classifier-free guidance.
Components:
transformer (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
Inputs:
timesteps (`Tensor`):
The timesteps to use for the denoising process.
num_inference_steps (`int`):
The number of inference steps to use for the denoising process.
latents (`Tensor`):
The image latent tokens to denoise. Shape: (B, seq_len, C)
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
Outputs:
latents (`Tensor`):
The denoised latents
"""
block_classes = [Flux3DistilledLoopDenoiser(), Flux3LoopAfterDenoiser()]
block_names = ["denoiser", "after_denoiser"]
@property
def description(self) -> str:
return (
"Denoise step that iteratively denoises the latents for the step-distilled (flow-map) FLUX.3 model. \n"
"Its loop logic is defined in `Flux3DenoiseLoopWrapper.__call__` method \n"
"At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
" - `Flux3DistilledLoopDenoiser`\n"
" - `Flux3LoopAfterDenoiser`\n"
"This block supports text-to-image and image-to-image generation, without classifier-free guidance."
)

397
src/diffusers/modular_pipelines/flux3/encoders.py Normal file
View File

@@ -0,0 +1,397 @@
# Copyright 2026 Black Forest Labs and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import math
import PIL.Image
import torch
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
from ...configuration_utils import FrozenDict
from ...models import AutoencoderKLFlux2
from ...utils import logging
from ..modular_pipeline import ModularPipelineBlocks, PipelineState
from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
from .image_processor import Flux3ImageProcessor
from .modular_pipeline import Flux3ModularPipeline
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
class Flux3ProcessImagesInputStep(ModularPipelineBlocks):
model_name = "flux3"
@property
def description(self) -> str:
return (
"Image preprocess step for FLUX.3 image-to-image generation. Each reference image keeps its own "
"resolution: it is converted to RGB, downscaled only when it exceeds 4096 x 4096 pixels, and "
"center-cropped to a multiple of 16 pixels."
)
@property
def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"image_processor",
Flux3ImageProcessor,
config=FrozenDict({"vae_scale_factor": 16, "vae_latent_channels": 32}),
default_creation_method="from_config",
),
]
@property
def inputs(self) -> list[InputParam]:
return [
InputParam.template(
"image",
description="The reference image or images to condition the generation on; each keeps its own "
"resolution (center-cropped to the 16-pixel grid)",
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"condition_images",
type_hint=list[torch.Tensor],
description="The preprocessed reference images, one tensor of shape (1, 3, height_i, width_i) in [-1, 1] per image",
)
]
@torch.no_grad()
def __call__(
self, components: Flux3ModularPipeline, state: PipelineState
) -> tuple[Flux3ModularPipeline, PipelineState]:
block_state = self.get_block_state(state)
images = block_state.image if isinstance(block_state.image, list) else [block_state.image]
condition_images = []
for image in images:
if not isinstance(image, PIL.Image.Image):
raise ValueError(f"`image` must be a `PIL.Image.Image` or a list of them, got {type(image)}")
# references keep their own resolution: downscaled only above 4096 x 4096 pixels, then center-cropped to
# the 16-pixel grid (`resize_mode="crop"` is a pure crop for this processor)
image = components.image_processor._resize_if_exceeds_area(image, 4096 * 4096)
condition_images.append(components.image_processor.preprocess(image, resize_mode="crop"))
block_state.condition_images = condition_images
self.set_block_state(state, block_state)
return components, state
class Flux3VaeEncoderStep(ModularPipelineBlocks):
model_name = "flux3"
@property
def description(self) -> str:
return (
"VAE Encoder step that encodes the preprocessed reference images into normalized packed latents, one "
"latent tensor per image."
)
@property
def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("vae", AutoencoderKLFlux2)]
@property
def inputs(self) -> list[InputParam]:
return [
InputParam(
"condition_images",
required=True,
type_hint=list[torch.Tensor],
description="The preprocessed reference images, one tensor of shape (1, 3, height_i, width_i) in [-1, 1] per image",
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"image_cond_latents",
type_hint=list[torch.Tensor],
description="Normalized packed latents of the reference images, one tensor of shape (1, 128, height_i//16, width_i//16) per image",
)
]
@staticmethod
def _patchify_latents(latents):
"""Convert latents to patchified format: (B, C, H, W) -> (B, C*4, H//2, W//2)."""
batch_size, num_channels_latents, height, width = latents.shape
latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
latents = latents.permute(0, 1, 3, 5, 2, 4)
latents = latents.reshape(batch_size, num_channels_latents * 4, height // 2, width // 2)
return latents
@torch.no_grad()
def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
vae = components.vae
device = components._execution_device
latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(device, vae.dtype)
latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to(
device, vae.dtype
)
image_cond_latents = []
for image in block_state.condition_images:
latents = vae.encode(image.to(device=device, dtype=vae.dtype)).latent_dist.mode()
latents = self._patchify_latents(latents)
image_cond_latents.append((latents - latents_bn_mean) / latents_bn_std)
block_state.image_cond_latents = image_cond_latents
self.set_block_state(state, block_state)
return components, state
class Flux3TextEncoderStep(ModularPipelineBlocks):
model_name = "flux3"
def __init__(self, pad_multiple: int = 80):
# the text window is right-padded to a multiple of this many tokens (capped at `max_sequence_length`); the
# padding is masked inside Qwen3-VL but not in the transformer, so the window size is part of the model recipe
self.pad_multiple = pad_multiple
super().__init__()
@property
def description(self) -> str:
return (
"Text Encoder step that generates text embeddings using Qwen3-VL to guide the generation. Takes one "
f"prompt per call; it goes through the chat template and is right-padded to a multiple of "
f"{self.pad_multiple} tokens."
)
@property
def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("text_encoder", Qwen3VLForConditionalGeneration),
ComponentSpec("tokenizer", AutoProcessor),
]
@property
def inputs(self) -> list[InputParam]:
return [
InputParam.template("prompt", description="The prompt to guide the generation (one prompt per call)"),
# no `default=""`: a materialized default would be indistinguishable from a user-passed value,
# and the guider-disabled warning depends on that distinction (the block falls back to "" itself)
InputParam.template("negative_prompt", description="The prompt not to guide the generation"),
InputParam.template(
"max_sequence_length",
default=8192,
description="Maximum number of prompt tokens; longer prompts are truncated.",
),
InputParam(
"text_encoder_out_layers",
type_hint=tuple[int],
default=(4, 8, 12, 16, 20, 24, 28, 32),
required=False,
description="Indices of the text-encoder hidden layers whose states are stacked to form the text embeddings",
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"prompt_embeds",
kwargs_type="denoiser_input_fields",
type_hint=torch.Tensor,
description="Text embeddings (stacked Qwen3-VL hidden states) used to guide the image generation",
),
OutputParam(
"negative_prompt_embeds",
kwargs_type="denoiser_input_fields",
type_hint=torch.Tensor,
description="Negative text embeddings used to guide the image generation",
),
]
@staticmethod
def check_inputs(block_state):
# one prompt per call (the text window is sized per prompt, and the transformer attends to its padding);
# `num_images_per_prompt` generates several images from it
if not isinstance(block_state.prompt, str):
raise ValueError(
f"`prompt` must be a str; batched prompts are not supported, got {type(block_state.prompt).__name__}."
)
if block_state.negative_prompt is not None and not isinstance(block_state.negative_prompt, str):
raise ValueError(
"`negative_prompt` must be a str or None; batched prompts are not supported, "
f"got {type(block_state.negative_prompt).__name__}."
)
@staticmethod
def _get_qwen3_vl_prompt_embeds(
text_encoder: Qwen3VLForConditionalGeneration,
tokenizer: AutoProcessor,
prompt: str,
device: torch.device | None = None,
max_sequence_length: int = 8192,
pad_multiple: int = 80,
hidden_states_layers: tuple[int, ...] = (4, 8, 12, 16, 20, 24, 28, 32),
):
dtype = text_encoder.dtype
device = text_encoder.device if device is None else device
text = tokenizer.apply_chat_template(
[{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True
)
# the window is the prompt's length rounded up to the pad multiple, capped at `max_sequence_length`
real_length = tokenizer.tokenizer(text, return_tensors="pt", truncation=True, max_length=max_sequence_length)[
"input_ids"
].shape[1]
padded_length = min(math.ceil(real_length / pad_multiple) * pad_multiple, max_sequence_length)
inputs = tokenizer.tokenizer(
text,
return_tensors="pt",
padding="max_length",
padding_side="right",
truncation=True,
max_length=padded_length,
)
# Qwen3-VL computes the position id of the padding tokens differently across transformers versions (1 in
# 4.57, the training environment; 0 in 5.x). FLUX.3 attends to the padding tokens, so the position ids are
# passed explicitly to match the training environment.
input_ids, attention_mask = inputs["input_ids"].to(device), inputs["attention_mask"].to(device)
position_ids = torch.arange(padded_length, device=device)[None].masked_fill(attention_mask == 0, 1)
output = text_encoder(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
output_hidden_states=True,
use_cache=False,
)
# Only use outputs from intermediate layers and stack them along the channel dimension
out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1)
out = out.to(dtype=dtype, device=device)
batch_size, num_channels, seq_len, hidden_dim = out.shape
prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim)
return prompt_embeds
@torch.no_grad()
def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
self.check_inputs(block_state)
device = components._execution_device
block_state.prompt_embeds = self._get_qwen3_vl_prompt_embeds(
text_encoder=components.text_encoder,
tokenizer=components.tokenizer,
prompt=block_state.prompt,
device=device,
max_sequence_length=block_state.max_sequence_length,
pad_multiple=self.pad_multiple,
hidden_states_layers=block_state.text_encoder_out_layers,
)
# The guider is not a component of this block. Combined with the denoiser, the pipeline's guider decides
# through `requires_unconditional_embeds` whether unconditional embeddings are needed (the negative prompt
# then defaults to ""); with the guider disabled a provided `negative_prompt` is ignored with a warning.
# Standalone (no guider), the negative prompt is encoded only when the caller passes one.
has_guider = getattr(components, "guider", None) is not None
if components.requires_unconditional_embeds or (block_state.negative_prompt is not None and not has_guider):
block_state.negative_prompt_embeds = self._get_qwen3_vl_prompt_embeds(
text_encoder=components.text_encoder,
tokenizer=components.tokenizer,
prompt=block_state.negative_prompt or "",
device=device,
max_sequence_length=block_state.max_sequence_length,
pad_multiple=self.pad_multiple,
hidden_states_layers=block_state.text_encoder_out_layers,
)
else:
if block_state.negative_prompt is not None:
logger.warning("`negative_prompt` is provided but the guider is disabled, it will be ignored.")
block_state.negative_prompt_embeds = None
self.set_block_state(state, block_state)
return components, state
class Flux3DistilledTextEncoderStep(Flux3TextEncoderStep):
model_name = "flux3-distilled"
@property
def description(self) -> str:
return (
"Text Encoder step that generates text embeddings using Qwen3-VL to guide the generation with the "
"step-distilled FLUX.3 model, which samples without classifier-free guidance (no negative prompt). Takes "
f"one prompt per call; it goes through the chat template and is right-padded to a multiple of "
f"{self.pad_multiple} tokens."
)
@property
def inputs(self) -> list[InputParam]:
return [
InputParam.template("prompt", description="The prompt to guide the generation (one prompt per call)"),
InputParam.template(
"max_sequence_length",
default=8192,
description="Maximum number of prompt tokens; longer prompts are truncated.",
),
InputParam(
"text_encoder_out_layers",
type_hint=tuple[int],
default=(4, 8, 12, 16, 20, 24, 28, 32),
required=False,
description="Indices of the text-encoder hidden layers whose states are stacked to form the text embeddings",
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"prompt_embeds",
kwargs_type="denoiser_input_fields",
type_hint=torch.Tensor,
description="Text embeddings (stacked Qwen3-VL hidden states) used to guide the image generation",
),
]
@staticmethod
def check_inputs(block_state):
if not isinstance(block_state.prompt, str):
raise ValueError(
f"`prompt` must be a str; batched prompts are not supported, got {type(block_state.prompt).__name__}."
)
@torch.no_grad()
def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
self.check_inputs(block_state)
block_state.prompt_embeds = self._get_qwen3_vl_prompt_embeds(
text_encoder=components.text_encoder,
tokenizer=components.tokenizer,
prompt=block_state.prompt,
device=components._execution_device,
max_sequence_length=block_state.max_sequence_length,
pad_multiple=self.pad_multiple,
hidden_states_layers=block_state.text_encoder_out_layers,
)
self.set_block_state(state, block_state)
return components, state

116
src/diffusers/modular_pipelines/flux3/image_processor.py Normal file
View File

@@ -0,0 +1,116 @@
# Copyright 2026 Black Forest Labs and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import math
import PIL.Image
from ...image_processor import VaeImageProcessor
class Flux3ImageProcessor(VaeImageProcessor):
r"""
Image processor for FLUX.3: preprocesses the reference images of image-to-image generation and postprocesses the
generated images.
Args:
do_resize (`bool`, *optional*, defaults to `True`):
Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. Can accept
`height` and `width` arguments from [`image_processor.VaeImageProcessor.preprocess`] method.
vae_scale_factor (`int`, *optional*, defaults to `16`):
VAE (spatial) scale factor, including the 2x2 packing of the latents. If `do_resize` is `True`, the image
is automatically resized to multiples of this factor.
vae_latent_channels (`int`, *optional*, defaults to `32`):
VAE latent channels.
do_normalize (`bool`, *optional*, defaults to `True`):
Whether to normalize the image to [-1,1].
do_binarize (`bool`, *optional*, defaults to `False`):
Whether to binarize the image to 0/1.
do_convert_rgb (`bool`, *optional*, defaults to be `True`):
Whether to convert the images to RGB format.
do_convert_grayscale (`bool`, *optional*, defaults to be `False`):
Whether to convert the images to grayscale format.
"""
# Copied from diffusers.pipelines.flux2.image_processor.Flux2ImageProcessor.__init__
def __init__(
self,
do_resize: bool = True,
vae_scale_factor: int = 16,
vae_latent_channels: int = 32,
do_normalize: bool = True,
do_binarize: bool = False,
do_convert_rgb: bool = True,
do_convert_grayscale: bool = False,
):
super().__init__(
do_resize=do_resize,
vae_scale_factor=vae_scale_factor,
vae_latent_channels=vae_latent_channels,
do_normalize=do_normalize,
do_binarize=do_binarize,
do_convert_rgb=do_convert_rgb,
do_convert_grayscale=do_convert_grayscale,
)
@staticmethod
# Copied from diffusers.pipelines.flux2.image_processor.Flux2ImageProcessor._resize_to_target_area
def _resize_to_target_area(image: PIL.Image.Image, target_area: int = 1024 * 1024) -> PIL.Image.Image:
image_width, image_height = image.size
scale = math.sqrt(target_area / (image_width * image_height))
width = int(image_width * scale)
height = int(image_height * scale)
return image.resize((width, height), PIL.Image.Resampling.LANCZOS)
@staticmethod
# Copied from diffusers.pipelines.flux2.image_processor.Flux2ImageProcessor._resize_if_exceeds_area with Flux2->Flux3
def _resize_if_exceeds_area(image, target_area=1024 * 1024) -> PIL.Image.Image:
image_width, image_height = image.size
pixel_count = image_width * image_height
if pixel_count <= target_area:
return image
return Flux3ImageProcessor._resize_to_target_area(image, target_area)
# Copied from diffusers.pipelines.flux2.image_processor.Flux2ImageProcessor._resize_and_crop
def _resize_and_crop(
self,
image: PIL.Image.Image,
width: int,
height: int,
) -> PIL.Image.Image:
r"""
center crop the image to the specified width and height.
Args:
image (`PIL.Image.Image`):
The image to resize and crop.
width (`int`):
The width to resize the image to.
height (`int`):
The height to resize the image to.
Returns:
`PIL.Image.Image`:
The resized and cropped image.
"""
image_width, image_height = image.size
left = (image_width - width) // 2
top = (image_height - height) // 2
right = left + width
bottom = top + height
return image.crop((left, top, right, bottom))

319
src/diffusers/modular_pipelines/flux3/inputs.py Normal file
View File

@@ -0,0 +1,319 @@
# Copyright 2026 Black Forest Labs and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import math
import torch
from ...models import Flux3Transformer2DModel
from ...utils import logging
from ..modular_pipeline import ModularPipelineBlocks, PipelineState
from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
from .modular_pipeline import Flux3ModularPipeline
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# Copied from diffusers.modular_pipelines.qwenimage.inputs.repeat_tensor_to_batch_size
def repeat_tensor_to_batch_size(
input_name: str,
input_tensor: torch.Tensor,
batch_size: int,
num_images_per_prompt: int = 1,
) -> torch.Tensor:
"""Repeat tensor elements to match the final batch size.
This function expands a tensor's batch dimension to match the final batch size (batch_size * num_images_per_prompt)
by repeating each element along dimension 0.
The input tensor must have batch size 1 or batch_size. The function will:
- If batch size is 1: repeat each element (batch_size * num_images_per_prompt) times
- If batch size equals batch_size: repeat each element num_images_per_prompt times
Args:
input_name (str): Name of the input tensor (used for error messages)
input_tensor (torch.Tensor): The tensor to repeat. Must have batch size 1 or batch_size.
batch_size (int): The base batch size (number of prompts)
num_images_per_prompt (int, optional): Number of images to generate per prompt. Defaults to 1.
Returns:
torch.Tensor: The repeated tensor with final batch size (batch_size * num_images_per_prompt)
Raises:
ValueError: If input_tensor is not a torch.Tensor or has invalid batch size
Examples:
tensor = torch.tensor([[1, 2, 3]]) # shape: [1, 3] repeated = repeat_tensor_to_batch_size("image", tensor,
batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]) - shape:
[4, 3]
tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]) # shape: [2, 3] repeated = repeat_tensor_to_batch_size("image",
tensor, batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [4, 5, 6], [4, 5, 6]])
- shape: [4, 3]
"""
# make sure input is a tensor
if not isinstance(input_tensor, torch.Tensor):
raise ValueError(f"`{input_name}` must be a tensor")
# make sure input tensor e.g. image_latents has batch size 1 or batch_size same as prompts
if input_tensor.shape[0] == 1:
repeat_by = batch_size * num_images_per_prompt
elif input_tensor.shape[0] == batch_size:
repeat_by = num_images_per_prompt
else:
raise ValueError(f"`{input_name}` must have batch size 1 or {batch_size}, but got {input_tensor.shape[0]}")
# expand the tensor to match the batch_size * num_images_per_prompt
input_tensor = input_tensor.repeat_interleave(repeat_by, dim=0)
return input_tensor
class Flux3TextInputStep(ModularPipelineBlocks):
model_name = "flux3"
@property
def description(self) -> str:
return (
"This step:\n"
" 1. Determines `batch_size` from `prompt_embeds` and the compute `dtype` from the transformer\n"
" 2. Ensures all text embeddings have consistent batch sizes (batch_size * num_images_per_prompt)"
)
@property
def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("transformer", Flux3Transformer2DModel)]
@property
def inputs(self) -> list[InputParam]:
return [
InputParam.template("num_images_per_prompt"),
InputParam(
"prompt_embeds",
required=True,
kwargs_type="denoiser_input_fields",
type_hint=torch.Tensor,
description="Pre-generated text embeddings. Can be generated from text_encoder step.",
),
InputParam(
"negative_prompt_embeds",
required=False,
kwargs_type="denoiser_input_fields",
type_hint=torch.Tensor,
description="Pre-generated negative text embeddings. Can be generated from text_encoder step.",
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"batch_size",
type_hint=int,
description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt",
),
OutputParam(
"dtype",
type_hint=torch.dtype,
description="Data type of model tensor inputs (the transformer's dtype)",
),
OutputParam(
"prompt_embeds",
type_hint=torch.Tensor,
kwargs_type="denoiser_input_fields",
description="Text embeddings used to guide the image generation",
),
OutputParam(
"negative_prompt_embeds",
type_hint=torch.Tensor,
kwargs_type="denoiser_input_fields",
description="Negative text embeddings used to guide the image generation",
),
]
@torch.no_grad()
def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
block_state.batch_size = block_state.prompt_embeds.shape[0]
# the denoiser consumes every tensor this chain prepares, so its dtype is the compute dtype:
# the text encoder may run in a different precision (e.g. a bf16 encoder with an fp32 transformer)
block_state.dtype = components.transformer.dtype
block_state.prompt_embeds = repeat_tensor_to_batch_size(
"prompt_embeds", block_state.prompt_embeds, block_state.batch_size, block_state.num_images_per_prompt
).to(block_state.dtype)
if block_state.negative_prompt_embeds is not None:
block_state.negative_prompt_embeds = repeat_tensor_to_batch_size(
"negative_prompt_embeds",
block_state.negative_prompt_embeds,
block_state.batch_size,
block_state.num_images_per_prompt,
).to(block_state.dtype)
elif components.requires_unconditional_embeds:
raise ValueError(
"The guider needs unconditional embeddings but `negative_prompt_embeds` is None. Pass `negative_prompt` "
"(an empty string for plain classifier-free guidance) to the text encoder step, or disable the guider."
)
self.set_block_state(state, block_state)
return components, state
class Flux3DistilledTextInputStep(ModularPipelineBlocks):
model_name = "flux3-distilled"
@property
def description(self) -> str:
return (
"This step:\n"
" 1. Determines `batch_size` from `prompt_embeds` and the compute `dtype` from the transformer\n"
" 2. Repeats the text embeddings to the batch size (batch_size * num_images_per_prompt)\n"
"For the step-distilled FLUX.3 model, which samples without classifier-free guidance (no negative "
"embeddings)."
)
@property
def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("transformer", Flux3Transformer2DModel)]
@property
def inputs(self) -> list[InputParam]:
return [
InputParam.template("num_images_per_prompt"),
InputParam(
"prompt_embeds",
required=True,
kwargs_type="denoiser_input_fields",
type_hint=torch.Tensor,
description="Pre-generated text embeddings. Can be generated from text_encoder step.",
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"batch_size",
type_hint=int,
description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt",
),
OutputParam(
"dtype",
type_hint=torch.dtype,
description="Data type of model tensor inputs (the transformer's dtype)",
),
OutputParam(
"prompt_embeds",
type_hint=torch.Tensor,
kwargs_type="denoiser_input_fields",
description="Text embeddings used to guide the image generation",
),
]
@torch.no_grad()
def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
block_state.batch_size = block_state.prompt_embeds.shape[0]
# the denoiser consumes every tensor this chain prepares, so its dtype is the compute dtype
block_state.dtype = components.transformer.dtype
block_state.prompt_embeds = repeat_tensor_to_batch_size(
"prompt_embeds", block_state.prompt_embeds, block_state.batch_size, block_state.num_images_per_prompt
).to(block_state.dtype)
self.set_block_state(state, block_state)
return components, state
class Flux3ImageCondInputStep(ModularPipelineBlocks):
model_name = "flux3"
def __init__(self, default_height: int = 1024, default_width: int = 1024):
# the generation size defaults to this area at the first reference image's aspect ratio
self.default_height = default_height
self.default_width = default_width
super().__init__()
@property
def description(self) -> str:
return (
"Step that expands the reference image latents to the final batch size and sets the generation "
f"`height`/`width` when not provided: {self.default_height} x {self.default_width} pixels of area at the "
"first reference image's aspect ratio, on the 16-pixel grid (a given side overrides its default). The "
"reference images keep their own resolution, the generation size is independent of them."
)
@property
def inputs(self) -> list[InputParam]:
return [
InputParam(
"image_cond_latents",
required=True,
type_hint=list[torch.Tensor],
description="Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16) per image",
),
InputParam(
"height",
type_hint=int,
description="The height in pixels of the generated image; defaults to the first reference image's "
"aspect ratio at the default area",
),
InputParam(
"width",
type_hint=int,
description="The width in pixels of the generated image; defaults to the first reference image's "
"aspect ratio at the default area",
),
InputParam.template("num_images_per_prompt"),
InputParam("batch_size", required=True, type_hint=int),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"image_cond_latents",
type_hint=list[torch.Tensor],
description="Normalized packed latents of the reference images, expanded to the final batch size",
),
OutputParam("height", type_hint=int, description="The height in pixels of the generated image"),
OutputParam("width", type_hint=int, description="The width in pixels of the generated image"),
]
@torch.no_grad()
def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
# the generation size defaults to the default area at the first reference's aspect ratio, rounded to a
# multiple of 16 pixels (one token); a given `height` or `width` overrides its side
if block_state.height is None or block_state.width is None:
multiple_of = components.vae_scale_factor * 2
latent_height, latent_width = block_state.image_cond_latents[0].shape[-2:]
scale = math.sqrt(self.default_height * self.default_width / (latent_height * latent_width)) / multiple_of
block_state.height = block_state.height or max(multiple_of, round(latent_height * scale) * multiple_of)
block_state.width = block_state.width or max(multiple_of, round(latent_width * scale) * multiple_of)
block_state.image_cond_latents = [
repeat_tensor_to_batch_size(
"image_cond_latents", latents, block_state.batch_size, block_state.num_images_per_prompt
)
for latents in block_state.image_cond_latents
]
self.set_block_state(state, block_state)
return components, state

406
src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py Normal file
View File

@@ -0,0 +1,406 @@
# Copyright 2026 Black Forest Labs and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from ...utils import logging
from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
from ..modular_pipeline_utils import InsertableDict, OutputParam
from .before_denoise import (
Flux3PrepareImageCondLatentsStep,
Flux3PrepareLatentsStep,
Flux3SetTimestepsStep,
Flux3TextRoPEInputsStep,
)
from .decoders import Flux3DecodeStep, Flux3UnpackLatentsStep
from .denoise import Flux3DenoiseStep
from .encoders import Flux3ProcessImagesInputStep, Flux3TextEncoderStep, Flux3VaeEncoderStep
from .inputs import Flux3ImageCondInputStep, Flux3TextInputStep
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
###
### VAE encoder (reference images)
###
# auto_docstring
class Flux3VaeEncoderSequentialStep(SequentialPipelineBlocks):
"""
VAE encoder step that preprocesses and encodes the reference images into their latent representations.
Components:
image_processor (`Flux3ImageProcessor`) vae (`AutoencoderKLFlux2`)
Inputs:
image (`Image | list`):
The reference image or images to condition the generation on; each keeps its own resolution
(center-cropped to the 16-pixel grid)
Outputs:
condition_images (`list`):
The preprocessed reference images, one tensor of shape (1, 3, height_i, width_i) in [-1, 1] per image
image_cond_latents (`list`):
Normalized packed latents of the reference images, one tensor of shape (1, 128, height_i//16,
width_i//16) per image
"""
model_name = "flux3"
block_classes = [Flux3ProcessImagesInputStep(), Flux3VaeEncoderStep()]
block_names = ["preprocess", "encode"]
@property
def description(self) -> str:
return "VAE encoder step that preprocesses and encodes the reference images into their latent representations."
# auto_docstring
class Flux3AutoVaeEncoderStep(AutoPipelineBlocks):
"""
VAE encoder step that encodes the reference images into their latent representations.
- `Flux3VaeEncoderSequentialStep` is used when `image` is provided (image-to-image).
- If `image` is not provided, step will be skipped.
Components:
image_processor (`Flux3ImageProcessor`) vae (`AutoencoderKLFlux2`)
Inputs:
image (`Image | list`, *optional*):
The reference image or images to condition the generation on; each keeps its own resolution
(center-cropped to the 16-pixel grid)
Outputs:
condition_images (`list`):
The preprocessed reference images, one tensor of shape (1, 3, height_i, width_i) in [-1, 1] per image
image_cond_latents (`list`):
Normalized packed latents of the reference images, one tensor of shape (1, 128, height_i//16,
width_i//16) per image
"""
model_name = "flux3"
block_classes = [Flux3VaeEncoderSequentialStep]
block_names = ["image_conditioning"]
block_trigger_inputs = ["image"]
@property
def description(self):
return (
"VAE encoder step that encodes the reference images into their latent representations.\n"
" - `Flux3VaeEncoderSequentialStep` is used when `image` is provided (image-to-image).\n"
" - If `image` is not provided, step will be skipped."
)
###
### Core denoise
###
Flux3CoreDenoiseBlocks = InsertableDict(
[
("input", Flux3TextInputStep()),
("prepare_latents", Flux3PrepareLatentsStep()),
("set_timesteps", Flux3SetTimestepsStep()),
("prepare_rope_inputs", Flux3TextRoPEInputsStep()),
("denoise", Flux3DenoiseStep()),
("after_denoise", Flux3UnpackLatentsStep()),
]
)
# auto_docstring
class Flux3CoreDenoiseStep(SequentialPipelineBlocks):
"""
Core denoise step that performs the denoising process for FLUX.3, for text-to-image generation.
Components:
transformer (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
(`ClassifierFreeGuidance`)
Inputs:
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
prompt_embeds (`Tensor`):
Pre-generated text embeddings. Can be generated from text_encoder step.
negative_prompt_embeds (`Tensor`, *optional*):
Pre-generated negative text embeddings. Can be generated from text_encoder step.
height (`int`, *optional*, defaults to 1024):
The height in pixels of the generated image.
width (`int`, *optional*, defaults to 1024):
The width in pixels of the generated image.
latents (`Tensor`, *optional*):
Pre-generated noisy latents for image generation.
generator (`Generator`, *optional*):
Torch generator for deterministic generation.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
timesteps (`Tensor`, *optional*):
Timesteps for the denoising process. (overrides `num_inference_steps` when provided)
sigmas (`list`, *optional*):
Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided)
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
Outputs:
latents (`Tensor`):
Denoised latents.
"""
model_name = "flux3"
block_classes = Flux3CoreDenoiseBlocks.values()
block_names = Flux3CoreDenoiseBlocks.keys()
@property
def description(self):
return "Core denoise step that performs the denoising process for FLUX.3, for text-to-image generation."
@property
def outputs(self):
return [
OutputParam.template("latents"),
]
Flux3ImageConditionedCoreDenoiseBlocks = InsertableDict(
[
("input", Flux3TextInputStep()),
("cond_input", Flux3ImageCondInputStep()),
("prepare_latents", Flux3PrepareLatentsStep()),
("prepare_image_cond_latents", Flux3PrepareImageCondLatentsStep()),
("set_timesteps", Flux3SetTimestepsStep()),
("prepare_rope_inputs", Flux3TextRoPEInputsStep()),
("denoise", Flux3DenoiseStep()),
("after_denoise", Flux3UnpackLatentsStep()),
]
)
# auto_docstring
class Flux3ImageConditionedCoreDenoiseStep(SequentialPipelineBlocks):
"""
Core denoise step that performs the denoising process for FLUX.3 with reference images (image-to-image).
Components:
transformer (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
(`ClassifierFreeGuidance`)
Inputs:
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
prompt_embeds (`Tensor`):
Pre-generated text embeddings. Can be generated from text_encoder step.
negative_prompt_embeds (`Tensor`, *optional*):
Pre-generated negative text embeddings. Can be generated from text_encoder step.
image_cond_latents (`list`):
Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16)
per image
height (`int`, *optional*):
The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the
default area
width (`int`, *optional*):
The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the
default area
latents (`Tensor`, *optional*):
Pre-generated noisy latents for image generation.
generator (`Generator`, *optional*):
Torch generator for deterministic generation.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps.
timesteps (`Tensor`, *optional*):
Timesteps for the denoising process. (overrides `num_inference_steps` when provided)
sigmas (`list`, *optional*):
Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided)
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
Outputs:
latents (`Tensor`):
Denoised latents.
"""
model_name = "flux3"
block_classes = Flux3ImageConditionedCoreDenoiseBlocks.values()
block_names = Flux3ImageConditionedCoreDenoiseBlocks.keys()
@property
def description(self):
return (
"Core denoise step that performs the denoising process for FLUX.3 with reference images (image-to-image)."
)
@property
def outputs(self):
return [
OutputParam.template("latents"),
]
# auto_docstring
class Flux3AutoCoreDenoiseStep(AutoPipelineBlocks):
"""
Core denoise step that selects the FLUX.3 denoising workflow by the conditioning inputs.
- `Flux3ImageConditionedCoreDenoiseStep` is used when `image_cond_latents` is provided (image-to-image).
- `Flux3CoreDenoiseStep` is used otherwise (text-to-image).
Components:
transformer (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
(`ClassifierFreeGuidance`)
Inputs:
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
prompt_embeds (`Tensor`):
Pre-generated text embeddings. Can be generated from text_encoder step.
negative_prompt_embeds (`Tensor`, *optional*):
Pre-generated negative text embeddings. Can be generated from text_encoder step.
image_cond_latents (`list`, *optional*):
Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16)
per image
height (`int`, *optional*, defaults to None or 1024, depending on the workflow):
The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the
default area
width (`int`, *optional*, defaults to None or 1024, depending on the workflow):
The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the
default area
latents (`Tensor`):
Pre-generated noisy latents for image generation.
generator (`Generator`, *optional*):
Torch generator for deterministic generation.
num_inference_steps (`int`):
The number of denoising steps.
timesteps (`Tensor`):
Timesteps for the denoising process. (overrides `num_inference_steps` when provided)
sigmas (`list`, *optional*):
Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided)
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
Outputs:
latents (`Tensor`):
Denoised latents.
"""
model_name = "flux3"
block_classes = [Flux3ImageConditionedCoreDenoiseStep, Flux3CoreDenoiseStep]
block_names = ["image_conditioned", "text2image"]
block_trigger_inputs = ["image_cond_latents", None]
@property
def description(self):
return (
"Core denoise step that selects the FLUX.3 denoising workflow by the conditioning inputs.\n"
" - `Flux3ImageConditionedCoreDenoiseStep` is used when `image_cond_latents` is provided (image-to-image).\n"
" - `Flux3CoreDenoiseStep` is used otherwise (text-to-image)."
)
###
### Auto blocks
###
AUTO_BLOCKS = InsertableDict(
[
("text_encoder", Flux3TextEncoderStep()),
("vae_encoder", Flux3AutoVaeEncoderStep()),
("denoise", Flux3AutoCoreDenoiseStep()),
("decode", Flux3DecodeStep()),
]
)
# auto_docstring
class Flux3AutoBlocks(SequentialPipelineBlocks):
"""
Auto blocks that perform text-to-image and image-to-image generation using FLUX.3.
Supported workflows:
- `text2image`: requires `prompt`
- `image_conditioned`: requires `prompt`, `image`
Components:
text_encoder (`Qwen3VLForConditionalGeneration`) tokenizer (`AutoProcessor`) image_processor
(`Flux3ImageProcessor`) vae (`AutoencoderKLFlux2`) transformer (`Flux3Transformer2DModel`) scheduler
(`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`)
Inputs:
prompt (`str`):
The prompt to guide the generation (one prompt per call)
negative_prompt (`str`, *optional*):
The prompt not to guide the generation
max_sequence_length (`int`, *optional*, defaults to 8192):
Maximum number of prompt tokens; longer prompts are truncated.
text_encoder_out_layers (`tuple`, *optional*, defaults to (4, 8, 12, 16, 20, 24, 28, 32)):
Indices of the text-encoder hidden layers whose states are stacked to form the text embeddings
image (`Image | list`, *optional*):
The reference image or images to condition the generation on; each keeps its own resolution
(center-cropped to the 16-pixel grid)
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
image_cond_latents (`list`, *optional*):
Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16)
per image
height (`int`, *optional*, defaults to None or 1024, depending on the workflow):
The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the
default area
width (`int`, *optional*, defaults to None or 1024, depending on the workflow):
The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the
default area
latents (`Tensor`):
Pre-generated noisy latents for image generation.
generator (`Generator`, *optional*):
Torch generator for deterministic generation.
num_inference_steps (`int`):
The number of denoising steps.
timesteps (`Tensor`):
Timesteps for the denoising process. (overrides `num_inference_steps` when provided)
sigmas (`list`, *optional*):
Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided)
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
output_type (`str`, *optional*, defaults to pil):
Output format: 'pil', 'np', 'pt'.
Outputs:
images (`list`):
Generated images.
latents (`Tensor`):
Denoised latents.
"""
model_name = "flux3"
block_classes = AUTO_BLOCKS.values()
block_names = AUTO_BLOCKS.keys()
_workflow_map = {
"text2image": {"prompt": True},
"image_conditioned": {"prompt": True, "image": True},
}
@property
def description(self):
return "Auto blocks that perform text-to-image and image-to-image generation using FLUX.3."
@property
def outputs(self):
return [
OutputParam.template("images"),
OutputParam.template("latents"),
]

333
src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py Normal file
View File

@@ -0,0 +1,333 @@
# Copyright 2026 Black Forest Labs and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from ...utils import logging
from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
from ..modular_pipeline_utils import InsertableDict, OutputParam
from .before_denoise import (
Flux3DistilledTextRoPEInputsStep,
Flux3PrepareImageCondLatentsStep,
Flux3PrepareLatentsStep,
Flux3SetTimestepsStep,
)
from .decoders import Flux3DecodeStep, Flux3UnpackLatentsStep
from .denoise import Flux3DistilledDenoiseStep
from .encoders import Flux3DistilledTextEncoderStep
from .inputs import Flux3DistilledTextInputStep, Flux3ImageCondInputStep
from .modular_blocks_flux3 import Flux3AutoVaeEncoderStep
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# The step-distilled (flow-map) FLUX.3 checkpoints: 16 steps and no classifier-free guidance. The text encoder,
# text input, text RoPE and denoise steps have distilled versions without the negative prompt; the denoise loop
# passes the transformer the timestep each step jumps to and runs a single pass. Everything else is shared with the
# base model.
###
### Core denoise
###
Flux3DistilledCoreDenoiseBlocks = InsertableDict(
[
("input", Flux3DistilledTextInputStep()),
("prepare_latents", Flux3PrepareLatentsStep()),
("set_timesteps", Flux3SetTimestepsStep(num_inference_steps=16)),
("prepare_rope_inputs", Flux3DistilledTextRoPEInputsStep()),
("denoise", Flux3DistilledDenoiseStep()),
("after_denoise", Flux3UnpackLatentsStep()),
]
)
# auto_docstring
class Flux3DistilledCoreDenoiseStep(SequentialPipelineBlocks):
"""
Core denoise step that performs the denoising process for the step-distilled FLUX.3 model, for text-to-image
generation.
Components:
transformer (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
Inputs:
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
prompt_embeds (`Tensor`):
Pre-generated text embeddings. Can be generated from text_encoder step.
height (`int`, *optional*, defaults to 1024):
The height in pixels of the generated image.
width (`int`, *optional*, defaults to 1024):
The width in pixels of the generated image.
latents (`Tensor`, *optional*):
Pre-generated noisy latents for image generation.
generator (`Generator`, *optional*):
Torch generator for deterministic generation.
num_inference_steps (`int`, *optional*, defaults to 16):
The number of denoising steps.
timesteps (`Tensor`, *optional*):
Timesteps for the denoising process. (overrides `num_inference_steps` when provided)
sigmas (`list`, *optional*):
Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided)
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
Outputs:
latents (`Tensor`):
Denoised latents.
"""
model_name = "flux3-distilled"
block_classes = Flux3DistilledCoreDenoiseBlocks.values()
block_names = Flux3DistilledCoreDenoiseBlocks.keys()
@property
def description(self):
return (
"Core denoise step that performs the denoising process for the step-distilled FLUX.3 model, for "
"text-to-image generation."
)
@property
def outputs(self):
return [
OutputParam.template("latents"),
]
Flux3DistilledImageConditionedCoreDenoiseBlocks = InsertableDict(
[
("input", Flux3DistilledTextInputStep()),
("cond_input", Flux3ImageCondInputStep()),
("prepare_latents", Flux3PrepareLatentsStep()),
("prepare_image_cond_latents", Flux3PrepareImageCondLatentsStep()),
("set_timesteps", Flux3SetTimestepsStep(num_inference_steps=16)),
("prepare_rope_inputs", Flux3DistilledTextRoPEInputsStep()),
("denoise", Flux3DistilledDenoiseStep()),
("after_denoise", Flux3UnpackLatentsStep()),
]
)
# auto_docstring
class Flux3DistilledImageConditionedCoreDenoiseStep(SequentialPipelineBlocks):
"""
Core denoise step that performs the denoising process for the step-distilled FLUX.3 model with reference images
(image-to-image).
Components:
transformer (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
Inputs:
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
prompt_embeds (`Tensor`):
Pre-generated text embeddings. Can be generated from text_encoder step.
image_cond_latents (`list`):
Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16)
per image
height (`int`, *optional*):
The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the
default area
width (`int`, *optional*):
The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the
default area
latents (`Tensor`, *optional*):
Pre-generated noisy latents for image generation.
generator (`Generator`, *optional*):
Torch generator for deterministic generation.
num_inference_steps (`int`, *optional*, defaults to 16):
The number of denoising steps.
timesteps (`Tensor`, *optional*):
Timesteps for the denoising process. (overrides `num_inference_steps` when provided)
sigmas (`list`, *optional*):
Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided)
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
Outputs:
latents (`Tensor`):
Denoised latents.
"""
model_name = "flux3-distilled"
block_classes = Flux3DistilledImageConditionedCoreDenoiseBlocks.values()
block_names = Flux3DistilledImageConditionedCoreDenoiseBlocks.keys()
@property
def description(self):
return (
"Core denoise step that performs the denoising process for the step-distilled FLUX.3 model with reference "
"images (image-to-image)."
)
@property
def outputs(self):
return [
OutputParam.template("latents"),
]
# auto_docstring
class Flux3DistilledAutoCoreDenoiseStep(AutoPipelineBlocks):
"""
Core denoise step that selects the step-distilled FLUX.3 denoising workflow by the conditioning inputs.
- `Flux3DistilledImageConditionedCoreDenoiseStep` is used when `image_cond_latents` is provided
(image-to-image).
- `Flux3DistilledCoreDenoiseStep` is used otherwise (text-to-image).
Components:
transformer (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
Inputs:
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
prompt_embeds (`Tensor`):
Pre-generated text embeddings. Can be generated from text_encoder step.
image_cond_latents (`list`, *optional*):
Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16)
per image
height (`int`, *optional*, defaults to None or 1024, depending on the workflow):
The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the
default area
width (`int`, *optional*, defaults to None or 1024, depending on the workflow):
The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the
default area
latents (`Tensor`):
Pre-generated noisy latents for image generation.
generator (`Generator`, *optional*):
Torch generator for deterministic generation.
num_inference_steps (`int`):
The number of denoising steps.
timesteps (`Tensor`):
Timesteps for the denoising process. (overrides `num_inference_steps` when provided)
sigmas (`list`, *optional*):
Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided)
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
Outputs:
latents (`Tensor`):
Denoised latents.
"""
model_name = "flux3-distilled"
block_classes = [Flux3DistilledImageConditionedCoreDenoiseStep, Flux3DistilledCoreDenoiseStep]
block_names = ["image_conditioned", "text2image"]
block_trigger_inputs = ["image_cond_latents", None]
@property
def description(self):
return (
"Core denoise step that selects the step-distilled FLUX.3 denoising workflow by the conditioning inputs.\n"
" - `Flux3DistilledImageConditionedCoreDenoiseStep` is used when `image_cond_latents` is provided "
"(image-to-image).\n"
" - `Flux3DistilledCoreDenoiseStep` is used otherwise (text-to-image)."
)
###
### Auto blocks
###
DISTILLED_AUTO_BLOCKS = InsertableDict(
[
("text_encoder", Flux3DistilledTextEncoderStep()),
("vae_encoder", Flux3AutoVaeEncoderStep()),
("denoise", Flux3DistilledAutoCoreDenoiseStep()),
("decode", Flux3DecodeStep()),
]
)
# auto_docstring
class Flux3DistilledAutoBlocks(SequentialPipelineBlocks):
"""
Auto blocks that perform text-to-image and image-to-image generation using the step-distilled FLUX.3 model (16
steps, no classifier-free guidance).
Supported workflows:
- `text2image`: requires `prompt`
- `image_conditioned`: requires `prompt`, `image`
Components:
text_encoder (`Qwen3VLForConditionalGeneration`) tokenizer (`AutoProcessor`) image_processor
(`Flux3ImageProcessor`) vae (`AutoencoderKLFlux2`) transformer (`Flux3Transformer2DModel`) scheduler
(`FlowMatchEulerDiscreteScheduler`)
Inputs:
prompt (`str`):
The prompt to guide the generation (one prompt per call)
max_sequence_length (`int`, *optional*, defaults to 8192):
Maximum number of prompt tokens; longer prompts are truncated.
text_encoder_out_layers (`tuple`, *optional*, defaults to (4, 8, 12, 16, 20, 24, 28, 32)):
Indices of the text-encoder hidden layers whose states are stacked to form the text embeddings
image (`Image | list`, *optional*):
The reference image or images to condition the generation on; each keeps its own resolution
(center-cropped to the 16-pixel grid)
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
image_cond_latents (`list`, *optional*):
Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16)
per image
height (`int`, *optional*, defaults to None or 1024, depending on the workflow):
The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the
default area
width (`int`, *optional*, defaults to None or 1024, depending on the workflow):
The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the
default area
latents (`Tensor`):
Pre-generated noisy latents for image generation.
generator (`Generator`, *optional*):
Torch generator for deterministic generation.
num_inference_steps (`int`):
The number of denoising steps.
timesteps (`Tensor`):
Timesteps for the denoising process. (overrides `num_inference_steps` when provided)
sigmas (`list`, *optional*):
Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided)
**denoiser_input_fields (`None`, *optional*):
conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
output_type (`str`, *optional*, defaults to pil):
Output format: 'pil', 'np', 'pt'.
Outputs:
images (`list`):
Generated images.
latents (`Tensor`):
Denoised latents.
"""
model_name = "flux3-distilled"
block_classes = DISTILLED_AUTO_BLOCKS.values()
block_names = DISTILLED_AUTO_BLOCKS.keys()
_workflow_map = {
"text2image": {"prompt": True},
"image_conditioned": {"prompt": True, "image": True},
}
@property
def description(self):
return (
"Auto blocks that perform text-to-image and image-to-image generation using the step-distilled FLUX.3 "
"model (16 steps, no classifier-free guidance)."
)
@property
def outputs(self):
return [
OutputParam.template("images"),
OutputParam.template("latents"),
]

57
src/diffusers/modular_pipelines/flux3/modular_pipeline.py Normal file
View File

@@ -0,0 +1,57 @@
# Copyright 2026 Black Forest Labs and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from ...loaders import Flux3LoraLoaderMixin
from ...utils import logging
from ..modular_pipeline import ModularPipeline
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
class Flux3ModularPipeline(ModularPipeline, Flux3LoraLoaderMixin):
"""
A ModularPipeline for FLUX.3.
"""
default_blocks_name = "Flux3AutoBlocks"
@property
def vae_scale_factor(self):
vae_scale_factor = 8
if getattr(self, "vae", None) is not None:
vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
return vae_scale_factor
@property
def num_channels_latents(self):
num_channels_latents = 128
if getattr(self, "transformer", None):
num_channels_latents = self.transformer.config.in_channels
return num_channels_latents
@property
def requires_unconditional_embeds(self):
if getattr(self, "guider", None) is None:
return False
return self.guider._enabled and self.guider.num_conditions > 1
class Flux3DistilledModularPipeline(Flux3ModularPipeline):
"""
A ModularPipeline for the step-distilled FLUX.3 model (16 steps, no classifier-free guidance).
"""
default_blocks_name = "Flux3DistilledAutoBlocks"

2
src/diffusers/modular_pipelines/modular_pipeline.py
View File

@@ -195,6 +195,8 @@ MODULAR_PIPELINE_MAPPING = OrderedDict(
("flux-kontext", _create_default_map_fn("FluxKontextModularPipeline")),
("flux2", _create_default_map_fn("Flux2ModularPipeline")),
("flux2-klein", _flux2_klein_map_fn),
("flux3", _create_default_map_fn("Flux3ModularPipeline")),
("flux3-distilled", _create_default_map_fn("Flux3DistilledModularPipeline")),
("ideogram4", _create_default_map_fn("Ideogram4ModularPipeline")),
("krea2", _krea2_map_fn),
("qwenimage", _create_default_map_fn("QwenImageModularPipeline")),

4
src/diffusers/quantizers/auto.py
View File

@@ -19,12 +19,14 @@ https://github.com/huggingface/transformers/blob/c409cd81777fb27aadc043ed3d8339d
import warnings
from .autoround import AutoRoundQuantizer
from .bfl import BFLQuantizer
from .bitsandbytes import BnB4BitDiffusersQuantizer, BnB8BitDiffusersQuantizer
from .gguf import GGUFQuantizer
from .modelopt import NVIDIAModelOptQuantizer
from .nunchaku import NunchakuLiteQuantizer
from .quantization_config import (
AutoRoundConfig,
BFLQuantizationConfig,
BitsAndBytesConfig,
GGUFQuantizationConfig,
NunchakuLiteQuantizationConfig,
@@ -41,6 +43,7 @@ from .torchao import TorchAoHfQuantizer
AUTO_QUANTIZER_MAPPING = {
"bfl": BFLQuantizer,
"bitsandbytes_4bit": BnB4BitDiffusersQuantizer,
"bitsandbytes_8bit": BnB8BitDiffusersQuantizer,
"gguf": GGUFQuantizer,
@@ -53,6 +56,7 @@ AUTO_QUANTIZER_MAPPING = {
}
AUTO_QUANTIZATION_CONFIG_MAPPING = {
"bfl": BFLQuantizationConfig,
"bitsandbytes_4bit": BitsAndBytesConfig,
"bitsandbytes_8bit": BitsAndBytesConfig,
"gguf": GGUFQuantizationConfig,

1
src/diffusers/quantizers/bfl/__init__.py Normal file
View File

@@ -0,0 +1 @@
from .bfl_quantizer import BFLQuantizer

152
src/diffusers/quantizers/bfl/bfl_quantizer.py Normal file
View File

@@ -0,0 +1,152 @@
# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from ...utils import get_module_from_name, is_accelerate_available, is_accelerate_version, is_torch_available
from ..base import DiffusersQuantizer
if TYPE_CHECKING:
from ...models.modeling_utils import ModelMixin
if is_torch_available():
import torch
import torch.nn as nn
from .utils import (
BFLFP8RLinear,
BFLNVFP4Linear,
BFLQuantizedParameter,
_dequantize_bfl_and_restore_linear,
_dequantize_weight,
_replace_with_bfl_linear,
)
class BFLQuantizer(DiffusersQuantizer):
use_keep_in_fp32_modules = True
def __init__(self, quantization_config, **kwargs):
super().__init__(quantization_config, **kwargs)
self.compute_dtype = quantization_config.compute_dtype
self.keep_in_fp32_modules = []
def validate_environment(self, *args, **kwargs):
if not is_accelerate_available() or is_accelerate_version("<", "0.26.0"):
raise ImportError(
"Loading BFL quantized checkpoints requires `accelerate` installed in your environment: `pip install 'accelerate>=0.26.0'`"
)
def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
return torch.uint8
def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
if torch_dtype is None:
torch_dtype = self.compute_dtype
return torch_dtype
def maybe_update_state_dict(self, state_dict: dict[str, Any]) -> dict[str, Any]:
prefixes = [key[: -len(".weight_scale")] for key in state_dict if key.endswith(".weight_scale")]
if not prefixes:
raise ValueError("The checkpoint has no `weight_scale` tensors, so it is not an fp8r or nvfp4 checkpoint.")
for prefix in prefixes:
weight_scale = state_dict.pop(f"{prefix}.weight_scale")
weight_scale_2 = state_dict.pop(f"{prefix}.weight_scale_2", None)
input_scale = state_dict.pop(f"{prefix}.input_scale", None)
quant_type = "nvfp4" if weight_scale_2 is not None else "fp8r"
weight = state_dict[f"{prefix}.weight"]
expected_dtype = torch.uint8 if quant_type == "nvfp4" else torch.float8_e4m3fn
if weight.dtype != expected_dtype:
raise ValueError(
f"{prefix}.weight is {weight.dtype}, but the {quant_type} layout expects {expected_dtype}."
)
if quant_type == "nvfp4" and input_scale is None:
raise ValueError(
f"{prefix} has `weight_scale_2` but no `input_scale`, which the nvfp4 layout requires."
)
state_dict[f"{prefix}.weight"] = BFLQuantizedParameter(
weight, quant_type, weight_scale, weight_scale_2, input_scale
)
return state_dict
def check_if_quantized_param(
self,
model: "ModelMixin",
param_value: "torch.Tensor",
param_name: str,
state_dict: dict[str, Any],
**kwargs,
) -> bool:
return isinstance(param_value, BFLQuantizedParameter)
def create_quantized_param(
self,
model: "ModelMixin",
param_value: "BFLQuantizedParameter",
param_name: str,
target_device: "torch.device",
state_dict: dict[str, Any] | None = None,
unexpected_keys: list[str] | None = None,
**kwargs,
):
module, tensor_name = get_module_from_name(model, param_name)
if isinstance(module, (BFLFP8RLinear, BFLNVFP4Linear)):
module.weight = nn.Parameter(param_value.as_tensor().to(target_device), requires_grad=False)
module.weight_scale = param_value.weight_scale.to(target_device)
if isinstance(module, BFLNVFP4Linear):
module.weight_scale_2 = param_value.weight_scale_2.to(target_device)
module.input_scale = param_value.input_scale.to(target_device)
return
# The module was excluded from quantization (for example by `_keep_in_fp32_modules`), so it stays a
# regular `nn.Linear` and receives the dequantized weight.
keep_in_fp32 = any(m in param_name.split(".") for m in self.keep_in_fp32_modules)
dtype = torch.float32 if keep_in_fp32 else self.compute_dtype
module._parameters[tensor_name] = nn.Parameter(
_dequantize_weight(param_value).to(target_device, dtype), requires_grad=False
)
def _process_model_before_weight_loading(
self,
model: "ModelMixin",
device_map,
keep_in_fp32_modules: list[str] = [],
**kwargs,
):
self.keep_in_fp32_modules = [module for module in keep_in_fp32_modules if module is not None]
_replace_with_bfl_linear(
model, self.compute_dtype, kwargs.get("state_dict"), modules_to_not_convert=self.keep_in_fp32_modules
)
def _process_model_after_weight_loading(self, model: "ModelMixin", **kwargs):
return model
@property
def is_serializable(self):
return False
@property
def is_trainable(self) -> bool:
return False
@property
def is_compileable(self) -> bool:
return True
def _dequantize(self, model):
return _dequantize_bfl_and_restore_linear(model)

310
src/diffusers/quantizers/bfl/utils.py Normal file
View File

@@ -0,0 +1,310 @@
# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import torch
import torch.nn as nn
import torch.nn.functional as F
from ...utils import is_accelerate_available, is_flashinfer_available
if is_accelerate_available():
import accelerate
from accelerate import init_empty_weights
from accelerate.hooks import add_hook_to_module, remove_hook_from_module
if is_flashinfer_available():
import flashinfer
FP8_MAX = 448.0
E2M1_MAX = 6.0
NVFP4_BLOCK = 16
# E2M1 values indexed by their 4-bit code: bit 3 is the sign, bits 2-0 the magnitude.
E2M1_LUT = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0])
def quantize_fp8r(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""`[N, K]` float weight -> E4M3 weight `[N, K]` and fp32 per-row scales `[N]`."""
weight = weight.float()
weight_scale = (weight.abs().amax(dim=1) / FP8_MAX).clamp(min=1e-12)
weight_q = (weight / weight_scale[:, None]).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
return weight_q, weight_scale
def dequantize_fp8r(weight: torch.Tensor, weight_scale: torch.Tensor) -> torch.Tensor:
return weight.float() * weight_scale[:, None]
def quantize_nvfp4(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""`[N, K]` float weight -> packed E2M1 weight `[N, K/2]` (even element in the low nibble), E4M3 block scales
`[N, K/16]` and the fp32 global scale `weight_scale_2 = amax / (448 * 6)`, following the ModelOpt NVFP4 layout."""
weight = weight.float()
out_features, in_features = weight.shape
weight_scale_2 = weight.abs().amax() / (FP8_MAX * E2M1_MAX)
blocks = weight.reshape(out_features, in_features // NVFP4_BLOCK, NVFP4_BLOCK)
weight_scale = (blocks.abs().amax(dim=-1) / E2M1_MAX / weight_scale_2).to(torch.float8_e4m3fn)
scale = (weight_scale.float() * weight_scale_2)[..., None]
scaled = torch.where(scale > 0, blocks / scale, torch.zeros_like(blocks))
magnitude = scaled.abs().clamp(max=E2M1_MAX)
# Round to nearest even on the E2M1 grid, whose spacing is 0.5 below 2, 1 below 4 and 2 above.
magnitude = torch.where(
magnitude < 2,
torch.round(magnitude * 2) / 2,
torch.where(magnitude < 4, torch.round(magnitude), torch.round(magnitude / 2) * 2),
)
codes = torch.where(magnitude < 2, magnitude * 2, torch.where(magnitude < 4, magnitude + 2, magnitude / 2 + 4))
codes = codes.to(torch.uint8) | ((scaled < 0).to(torch.uint8) << 3)
codes = codes.reshape(out_features, in_features)
weight_q = codes[:, 0::2] | (codes[:, 1::2] << 4)
return weight_q, weight_scale, weight_scale_2
def dequantize_nvfp4(weight: torch.Tensor, weight_scale: torch.Tensor, weight_scale_2: torch.Tensor) -> torch.Tensor:
codes = torch.stack([weight & 0x0F, weight >> 4], dim=-1).reshape(weight.shape[0], -1)
values = E2M1_LUT.to(weight.device)[codes.long()]
scale = (weight_scale.float() * weight_scale_2).repeat_interleave(NVFP4_BLOCK, dim=1)
return values * scale
class BFLQuantizedParameter(nn.Parameter):
"""A quantized linear weight bundled with its scales.
Single-file checkpoint converters chunk, split and concatenate weights along the output dimension (fused QKV,
swapped scale/shift). Bundling keeps the per-row scales aligned with the rows through those operations. A cast to a
floating point dtype dequantizes the weight into a plain tensor (the loader casts the parameters of modules kept in
fp32); any other shape or dtype change is refused.
"""
def __new__(cls, data, quant_type, weight_scale, weight_scale_2=None, input_scale=None):
self = torch.Tensor._make_subclass(cls, data, False)
self.quant_type = quant_type
self.weight_scale = weight_scale
self.weight_scale_2 = weight_scale_2
self.input_scale = input_scale
return self
def as_tensor(self):
return torch.Tensor._make_subclass(torch.Tensor, self, False)
def _like(self, data, weight_scale):
return BFLQuantizedParameter(data, self.quant_type, weight_scale, self.weight_scale_2, self.input_scale)
def __repr__(self):
return f"BFLQuantizedParameter(quant_type={self.quant_type!r}, data={self.as_tensor()!r})"
@classmethod
def __torch_function__(cls, func, types, args=(), kwargs=None):
kwargs = kwargs or {}
if func in (torch.chunk, torch.Tensor.chunk, torch.split, torch.Tensor.split):
source = args[0]
dim = kwargs.get("dim", args[2] if len(args) > 2 else 0)
if dim != 0:
raise NotImplementedError(f"{func.__name__} along dim={dim} is not supported on quantized weights.")
with torch._C.DisableTorchFunctionSubclass():
weights = func(*args, **kwargs)
scales = func(source.weight_scale, *args[1:], **kwargs)
return tuple(source._like(w, s) for w, s in zip(weights, scales))
if func in (torch.cat, torch.concat, torch.concatenate):
tensors = args[0]
dim = kwargs.get("dim", args[1] if len(args) > 1 else 0)
source = tensors[0]
compatible = all(
isinstance(t, cls)
and t.quant_type == source.quant_type
and all(
(getattr(t, name) is None and getattr(source, name) is None)
or torch.equal(getattr(t, name), getattr(source, name))
for name in ("weight_scale_2", "input_scale")
)
for t in tensors
)
if dim != 0 or not compatible:
raise NotImplementedError(
"Only quantized weights of the same scheme and global scales can be concatenated along dim=0."
)
with torch._C.DisableTorchFunctionSubclass():
weight = torch.cat(tensors, dim=0)
weight_scale = torch.cat([t.weight_scale for t in tensors], dim=0)
return source._like(weight, weight_scale)
source = next(a for a in torch.utils._pytree.tree_leaves(args) if isinstance(a, cls))
with torch._C.DisableTorchFunctionSubclass():
result = func(*args, **kwargs)
if not isinstance(result, torch.Tensor) or func in torch.overrides.get_default_nowrap_functions():
return result
if result.shape == source.shape and result.dtype != source.dtype and result.is_floating_point():
return _dequantize_weight(source).to(result.device, result.dtype)
if result.shape != source.shape or result.dtype != source.dtype:
raise NotImplementedError(
f"{getattr(func, '__name__', func)} changes the shape or dtype of a quantized weight, which is not supported."
)
weight_scale = source.weight_scale.to(result.device)
return source._like(result, weight_scale)
class BFLFP8RLinear(nn.Linear):
"""Linear over an E4M3 weight `[N, K]` with fp32 per-row scales `[N]`.
On CUDA devices with FP8 support (compute capability 8.9+) activations are quantized per row and the GEMM runs
through `torch._scaled_mm`; elsewhere the weight is dequantized to `compute_dtype`.
"""
def __init__(self, in_features, out_features, bias=True, compute_dtype=None, device=None):
super().__init__(in_features, out_features, bias, device=device)
self.weight = nn.Parameter(
torch.empty(out_features, in_features, dtype=torch.float8_e4m3fn, device=device), requires_grad=False
)
self.register_buffer("weight_scale", torch.empty(out_features, dtype=torch.float32, device=device))
self.compute_dtype = compute_dtype
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
if (
inputs.is_cuda
and self.in_features % 16 == 0
and self.out_features % 16 == 0
and torch.cuda.get_device_capability(inputs.device) >= (8, 9)
):
x = inputs.reshape(-1, self.in_features).float()
x_scale = (x.abs().amax(dim=1, keepdim=True) / FP8_MAX).clamp(min=1e-12)
x_q = (x / x_scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
output = torch._scaled_mm(
x_q, self.weight.T, x_scale, self.weight_scale[None, :], out_dtype=torch.bfloat16, use_fast_accum=True
)
output = output.reshape(*inputs.shape[:-1], self.out_features).to(self.compute_dtype)
if self.bias is not None:
output = output + self.bias
return output
weight = dequantize_fp8r(self.weight, self.weight_scale).to(self.compute_dtype)
return F.linear(inputs, weight, self.bias)
class BFLNVFP4Linear(nn.Linear):
"""Linear over a packed E2M1 weight `[N, K/2]`, E4M3 block scales `[N, K/16]`, the fp32 global weight scale
`weight_scale_2` and the calibrated fp32 activation scale `input_scale`.
On Blackwell GPUs with FlashInfer installed activations are quantized to NVFP4 with the static `input_scale` and
the GEMM runs through `flashinfer.mm_fp4`; elsewhere the weight is dequantized to `compute_dtype`.
"""
def __init__(self, in_features, out_features, bias=True, compute_dtype=None, device=None):
super().__init__(in_features, out_features, bias, device=device)
self.weight = nn.Parameter(
torch.empty(out_features, in_features // 2, dtype=torch.uint8, device=device), requires_grad=False
)
self.register_buffer(
"weight_scale",
torch.empty(out_features, in_features // NVFP4_BLOCK, dtype=torch.float8_e4m3fn, device=device),
)
self.register_buffer("weight_scale_2", torch.empty((), dtype=torch.float32, device=device))
self.register_buffer("input_scale", torch.empty((), dtype=torch.float32, device=device))
self.compute_dtype = compute_dtype
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
if inputs.is_cuda and is_flashinfer_available() and torch.cuda.get_device_capability(inputs.device)[0] >= 10:
x = inputs.reshape(-1, self.in_features).to(torch.bfloat16).contiguous()
x_q, x_scale = flashinfer.nvfp4_quantize(
x, (1.0 / self.input_scale).reshape(1), sfLayout=flashinfer.SfLayout.layout_128x4, do_shuffle=False
)
rows, cols = self.weight_scale.shape
weight_scale = flashinfer.block_scale_interleave(self.weight_scale.view(torch.uint8).contiguous())
weight_scale = weight_scale.reshape(-(-rows // 128) * 128, -(-cols // 4) * 4).view(torch.float8_e4m3fn)
output = flashinfer.mm_fp4(
x_q,
self.weight.T,
x_scale,
weight_scale.T,
alpha=(self.input_scale * self.weight_scale_2).reshape(1),
out_dtype=torch.bfloat16,
backend="cutlass",
)
output = output.reshape(*inputs.shape[:-1], self.out_features).to(self.compute_dtype)
if self.bias is not None:
output = output + self.bias
return output
weight = dequantize_nvfp4(self.weight, self.weight_scale, self.weight_scale_2).to(self.compute_dtype)
return F.linear(inputs, weight, self.bias)
QUANTIZED_LINEARS = {"fp8r": BFLFP8RLinear, "nvfp4": BFLNVFP4Linear}
def _dequantize_weight(param: BFLQuantizedParameter) -> torch.Tensor:
if param.quant_type == "nvfp4":
return dequantize_nvfp4(param.as_tensor(), param.weight_scale, param.weight_scale_2)
return dequantize_fp8r(param.as_tensor(), param.weight_scale)
def _replace_with_bfl_linear(model, compute_dtype, state_dict, prefix="", modules_to_not_convert=[]):
for name, module in model.named_children():
if name in modules_to_not_convert:
continue
module_prefix = prefix + name + "."
_replace_with_bfl_linear(module, compute_dtype, state_dict, module_prefix, modules_to_not_convert)
weight = state_dict.get(module_prefix + "weight")
if isinstance(module, nn.Linear) and isinstance(weight, BFLQuantizedParameter):
with init_empty_weights():
model._modules[name] = QUANTIZED_LINEARS[weight.quant_type](
module.in_features, module.out_features, module.bias is not None, compute_dtype=compute_dtype
)
return model
# Copied from diffusers.quantizers.bitsandbytes.utils._create_accelerate_new_hook
def _create_accelerate_new_hook(old_hook):
r"""
Creates a new hook based on the old hook. Use it only if you know what you are doing ! This method is a copy of:
https://github.com/huggingface/peft/blob/748f7968f3a31ec06a1c2b0328993319ad9a150a/src/peft/utils/other.py#L245 with
some changes
"""
old_hook_cls = getattr(accelerate.hooks, old_hook.__class__.__name__)
old_hook_attr = old_hook.__dict__
filtered_old_hook_attr = {}
old_hook_init_signature = inspect.signature(old_hook_cls.__init__)
for k in old_hook_attr.keys():
if k in old_hook_init_signature.parameters:
filtered_old_hook_attr[k] = old_hook_attr[k]
new_hook = old_hook_cls(**filtered_old_hook_attr)
return new_hook
def _dequantize_bfl_and_restore_linear(model):
for name, module in model.named_children():
if isinstance(module, (BFLFP8RLinear, BFLNVFP4Linear)):
if isinstance(module, BFLNVFP4Linear):
weight = dequantize_nvfp4(module.weight, module.weight_scale, module.weight_scale_2)
else:
weight = dequantize_fp8r(module.weight, module.weight_scale)
device = module.weight.device
with init_empty_weights():
new_module = nn.Linear(module.in_features, module.out_features, module.bias is not None)
new_module.weight = nn.Parameter(weight.to(module.compute_dtype), requires_grad=False)
if module.bias is not None:
new_module.bias = module.bias
if hasattr(module, "_hf_hook"):
new_hook = _create_accelerate_new_hook(module._hf_hook)
remove_hook_from_module(module)
add_hook_to_module(new_module, new_hook)
model._modules[name] = new_module.to(device)
else:
_dequantize_bfl_and_restore_linear(module)
return model

25
src/diffusers/quantizers/quantization_config.py
View File

@@ -44,6 +44,7 @@ logger = logging.get_logger(__name__)
class QuantizationMethod(str, Enum):
BITS_AND_BYTES = "bitsandbytes"
BFL = "bfl"
GGUF = "gguf"
NUNCHAKU_LITE = "nunchaku_lite"
TORCHAO = "torchao"
@@ -431,6 +432,30 @@ class GGUFQuantizationConfig(QuantizationConfigMixin):
self.compute_dtype = torch.float32
@dataclass
class BFLQuantizationConfig(QuantizationConfigMixin):
"""Config for loading checkpoints quantized with the Black Forest Labs `fp8r` and `nvfp4` schemes.
Both schemes are read from the checkpoint through [`~FromSingleFileMixin.from_single_file`]; quantizing a model on
the fly is not supported. A linear layer is quantized when its weight comes with scale tensors:
- `fp8r`: `weight` is `float8_e4m3fn` with shape `[N, K]` and `weight_scale` is an fp32 per-row scale `[N]`.
- `nvfp4`: `weight` is `uint8` with shape `[N, K/2]` (two E2M1 values per byte), `weight_scale` is a
`float8_e4m3fn` block scale `[N, K/16]`, and `weight_scale_2` and `input_scale` are fp32 scalars holding the
global weight scale and the calibrated activation scale (the ModelOpt NVFP4 layout).
Args:
compute_dtype (`torch.dtype`, defaults to `torch.bfloat16`):
The dtype of the activations and of the unquantized layers. Devices without fused FP8 or NVFP4 kernels
dequantize the weights to this dtype in every forward pass.
"""
def __init__(self, compute_dtype: "torch.dtype" | None = None):
self.quant_method = QuantizationMethod.BFL
self.compute_dtype = compute_dtype if compute_dtype is not None else torch.bfloat16
self.pre_quantized = True
@dataclass
class NunchakuLiteQuantizationConfig(QuantizationConfigMixin):
"""Configuration for loading Nunchaku Lite checkpoints.

1
src/diffusers/utils/__init__.py
View File

@@ -77,6 +77,7 @@ from .import_utils import (
is_flash_attn_3_available,
is_flash_attn_available,
is_flash_attn_version,
is_flashinfer_available,
is_flashpack_available,
is_ftfy_available,
is_gguf_available,

17
src/diffusers/utils/dummy_bfl_objects.py Normal file
View File

@@ -0,0 +1,17 @@
# This file is autogenerated by the command `make fix-copies`, do not edit.
from ..utils import DummyObject, requires_backends
class BFLQuantizationConfig(metaclass=DummyObject):
_backends = ["torch"]
def __init__(self, *args, **kwargs):
requires_backends(self, ["torch"])
@classmethod
def from_config(cls, *args, **kwargs):
requires_backends(cls, ["torch"])
@classmethod
def from_pretrained(cls, *args, **kwargs):
requires_backends(cls, ["torch"])

15
src/diffusers/utils/dummy_pt_objects.py
View File

@@ -1354,6 +1354,21 @@ class Flux2Transformer2DModel(metaclass=DummyObject):
requires_backends(cls, ["torch"])
class Flux3Transformer2DModel(metaclass=DummyObject):
_backends = ["torch"]
def __init__(self, *args, **kwargs):
requires_backends(self, ["torch"])
@classmethod
def from_config(cls, *args, **kwargs):
requires_backends(cls, ["torch"])
@classmethod
def from_pretrained(cls, *args, **kwargs):
requires_backends(cls, ["torch"])
class FluxControlNetModel(metaclass=DummyObject):
_backends = ["torch"]

60
src/diffusers/utils/dummy_torch_and_transformers_objects.py
View File

@@ -242,6 +242,66 @@ class Flux2ModularPipeline(metaclass=DummyObject):
requires_backends(cls, ["torch", "transformers"])
class Flux3AutoBlocks(metaclass=DummyObject):
_backends = ["torch", "transformers"]
def __init__(self, *args, **kwargs):
requires_backends(self, ["torch", "transformers"])
@classmethod
def from_config(cls, *args, **kwargs):
requires_backends(cls, ["torch", "transformers"])
@classmethod
def from_pretrained(cls, *args, **kwargs):
requires_backends(cls, ["torch", "transformers"])
class Flux3DistilledAutoBlocks(metaclass=DummyObject):
_backends = ["torch", "transformers"]
def __init__(self, *args, **kwargs):
requires_backends(self, ["torch", "transformers"])
@classmethod
def from_config(cls, *args, **kwargs):
requires_backends(cls, ["torch", "transformers"])
@classmethod
def from_pretrained(cls, *args, **kwargs):
requires_backends(cls, ["torch", "transformers"])
class Flux3DistilledModularPipeline(metaclass=DummyObject):
_backends = ["torch", "transformers"]
def __init__(self, *args, **kwargs):
requires_backends(self, ["torch", "transformers"])
@classmethod
def from_config(cls, *args, **kwargs):
requires_backends(cls, ["torch", "transformers"])
@classmethod
def from_pretrained(cls, *args, **kwargs):
requires_backends(cls, ["torch", "transformers"])
class Flux3ModularPipeline(metaclass=DummyObject):
_backends = ["torch", "transformers"]
def __init__(self, *args, **kwargs):
requires_backends(self, ["torch", "transformers"])
@classmethod
def from_config(cls, *args, **kwargs):
requires_backends(cls, ["torch", "transformers"])
@classmethod
def from_pretrained(cls, *args, **kwargs):
requires_backends(cls, ["torch", "transformers"])
class FluxAutoBlocks(metaclass=DummyObject):
_backends = ["torch", "transformers"]

5
src/diffusers/utils/import_utils.py
View File

@@ -214,6 +214,7 @@ _nltk_available, _nltk_version = _is_package_available("nltk")
_cosmos_guardrail_available, _cosmos_guardrail_version = _is_package_available("cosmos_guardrail")
_sageattention_available, _sageattention_version = _is_package_available("sageattention")
_flash_attn_available, _flash_attn_version = _is_package_available("flash_attn")
_flashinfer_available, _flashinfer_version = _is_package_available("flashinfer", get_dist_name=True)
_flash_attn_3_available, _flash_attn_3_version = _is_package_available("flash_attn_3")
_kornia_available, _kornia_version = _is_package_available("kornia")
_nvidia_modelopt_available, _nvidia_modelopt_version = _is_package_available("modelopt", get_dist_name=True)
@@ -359,6 +360,10 @@ def is_gguf_available():
return _gguf_available
def is_flashinfer_available():
return _flashinfer_available
def is_flashpack_available():
return _flashpack_available

1
tests/conftest.py
View File

@@ -47,6 +47,7 @@ def pytest_configure(config):
config.addinivalue_line("markers", "gguf: marks tests for GGUF quantization functionality")
config.addinivalue_line("markers", "modelopt: marks tests for NVIDIA ModelOpt quantization functionality")
config.addinivalue_line("markers", "sdnq: marks tests for SDNQ quantization functionality")
config.addinivalue_line("markers", "bfl: marks tests for BFL (fp8r / nvfp4) quantization functionality")
config.addinivalue_line("markers", "context_parallel: marks tests for context parallel inference functionality")
config.addinivalue_line("markers", "tensor_parallel: marks tests for tensor parallel inference functionality")
config.addinivalue_line("markers", "slow: mark test as slow")

6
tests/models/testing_utils/__init__.py
View File

@@ -31,6 +31,9 @@ from .quantization import (
AutoRoundCompileTesterMixin,
AutoRoundConfigMixin,
AutoRoundTesterMixin,
BFLCompileTesterMixin,
BFLConfigMixin,
BFLTesterMixin,
BitsAndBytesCompileTesterMixin,
BitsAndBytesConfigMixin,
BitsAndBytesTesterMixin,
@@ -62,6 +65,9 @@ __all__ = [
"AttentionTesterMixin",
"AutoRoundConfigMixin",
"AutoRoundTesterMixin",
"BFLCompileTesterMixin",
"BFLConfigMixin",
"BFLTesterMixin",
"BaseModelTesterConfig",
"BitsAndBytesCompileTesterMixin",
"BitsAndBytesConfigMixin",

287
tests/models/testing_utils/quantization.py
View File

@@ -15,6 +15,7 @@
import gc
import re
import tempfile
import pytest
import safetensors.torch
@@ -22,6 +23,7 @@ import torch
from diffusers import (
AutoRoundConfig,
BFLQuantizationConfig,
BitsAndBytesConfig,
GGUFQuantizationConfig,
NunchakuLiteQuantizationConfig,
@@ -29,8 +31,10 @@ from diffusers import (
SDNQConfig,
TorchAoConfig,
)
from diffusers.loaders.single_file_model import SINGLE_FILE_LOADABLE_CLASSES, _get_single_file_loadable_mapping_class
from diffusers.utils.import_utils import (
is_bitsandbytes_available,
is_flashinfer_available,
is_gguf_available,
is_kernels_available,
is_nvidia_modelopt_available,
@@ -42,6 +46,7 @@ from ...testing_utils import (
assert_tensors_close,
backend_empty_cache,
is_autoround,
is_bfl,
is_bitsandbytes,
is_gguf,
is_modelopt,
@@ -59,6 +64,7 @@ from ...testing_utils import (
require_torchao_version_greater_or_equal,
torch_device,
)
from .single_file import fetch_checkpoint_metadata
if is_peft_available():
@@ -1823,3 +1829,284 @@ class AutoRoundCompileTesterMixin(AutoRoundConfigMixin, QuantizationCompileTeste
def test_autoround_torch_compile_with_group_offload(self):
self._test_torch_compile_with_group_offload(self.config_dict)
class BFLConfigMixin:
"""
Base mixin providing BFL (`fp8r` / `nvfp4`) quantization config and model creation.
BFL checkpoints are prequantized single files, one per scheme, and the tests never download their weights. The
safetensors header of each checkpoint says which linears the exporter quantized; the tests quantize those same
linears of the small test model with the library's own `quantize_fp8r` / `quantize_nvfp4` and load the result
through [`~FromOriginalModelMixin.from_single_file`], exactly as the real checkpoint is loaded.
Expected class attributes:
- model_class: The model class to test
Expected methods to be implemented by subclasses:
- get_init_dict(): Returns the init kwargs of the small test model
- get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass
Required properties (must be implemented by subclasses):
- bfl_checkpoints: Dict mapping a scheme (`"fp8r"`, `"nvfp4"`) to the URL or path of its single-file checkpoint
"""
BFL_CONFIGS = {
"fp8r": {"scheme": "fp8r", "compute_dtype": torch.bfloat16},
"nvfp4": {"scheme": "nvfp4", "compute_dtype": torch.bfloat16},
}
@property
def bfl_checkpoints(self) -> dict[str, str]:
"""Scheme -> single-file checkpoint. Must be implemented by subclasses."""
raise NotImplementedError("Subclasses must implement the `bfl_checkpoints` property.")
@staticmethod
def _bfl_linear_cls(scheme):
from diffusers.quantizers.bfl.utils import BFLFP8RLinear, BFLNVFP4Linear
return {"fp8r": BFLFP8RLinear, "nvfp4": BFLNVFP4Linear}[scheme]
def _bfl_quantized_modules(self, scheme) -> set[str]:
"""The diffusers names of the linears the checkpoint quantizes, read from its safetensors header."""
metadata = fetch_checkpoint_metadata(self.bfl_checkpoints[scheme])
if metadata is None:
pytest.skip(f"{self.bfl_checkpoints[scheme]} is not a safetensors checkpoint")
# one `weight_scale` per quantized linear; the model's single-file mapping renames them to diffusers keys
weights = {
f"{name[: -len('.weight_scale')]}.weight" for name in metadata.tensors if name.endswith(".weight_scale")
}
checkpoint = {
name: torch.empty(tuple(metadata.tensors[name].shape), device="meta") for name in sorted(weights)
}
mapping_fn = SINGLE_FILE_LOADABLE_CLASSES[_get_single_file_loadable_mapping_class(self.model_class)][
"checkpoint_mapping_fn"
]
converted = mapping_fn(checkpoint=checkpoint, config=None)
return {name[: -len(".weight")] for name in converted}
def _bfl_quantized_linears(self, model, scheme) -> set[str]:
"""The checkpoint's quantized linears that the small model has (it has fewer blocks than the checkpoint)."""
linears = {name for name, module in model.named_modules() if isinstance(module, torch.nn.Linear)}
return self._bfl_quantized_modules(scheme) & linears
def _load_unquantized_model(self):
torch.manual_seed(0)
return self.model_class(**self.get_init_dict()).eval()
@torch.no_grad()
def _bfl_quantized_state_dict(self, model, scheme):
from diffusers.quantizers.bfl.utils import E2M1_MAX, FP8_MAX, quantize_fp8r, quantize_nvfp4
quantized_modules = self._bfl_quantized_linears(model, scheme)
assert quantized_modules, "the small model has none of the linears the checkpoint quantizes"
# nvfp4 stores a calibrated activation scale per linear: record the input amax on the dummy inputs
amax = {}
if scheme == "nvfp4":
handles = [
model.get_submodule(name).register_forward_pre_hook(
lambda module, inputs, name=name: amax.__setitem__(name, inputs[0].abs().amax().item())
)
for name in quantized_modules
]
model.to(torch_device)(**self.get_dummy_inputs(), return_dict=False)
model.cpu()
for handle in handles:
handle.remove()
state_dict = {}
for name, tensor in model.state_dict().items():
module_name = name[: -len(".weight")] if name.endswith(".weight") else None
if module_name not in quantized_modules:
state_dict[name] = tensor
elif scheme == "fp8r":
state_dict[name], state_dict[f"{module_name}.weight_scale"] = quantize_fp8r(tensor)
else:
weight_q, weight_scale, weight_scale_2 = quantize_nvfp4(tensor)
state_dict[name] = weight_q
state_dict[f"{module_name}.weight_scale"] = weight_scale
state_dict[f"{module_name}.weight_scale_2"] = weight_scale_2
# linears the dummy inputs do not reach (e.g. the reference-image blocks) take the largest scale seen
input_amax = amax.get(module_name, max(amax.values(), default=1.0))
state_dict[f"{module_name}.input_scale"] = torch.tensor(input_amax / (FP8_MAX * E2M1_MAX))
return state_dict
def _create_quantized_model(self, config_kwargs, **extra_kwargs):
config_kwargs = dict(config_kwargs)
scheme = config_kwargs.pop("scheme")
compute_dtype = config_kwargs.get("compute_dtype", torch.bfloat16)
model = self._load_unquantized_model().to(compute_dtype)
state_dict = self._bfl_quantized_state_dict(model, scheme)
kwargs = {
"quantization_config": BFLQuantizationConfig(**config_kwargs),
"dtype": compute_dtype,
"device": str(torch_device),
}
kwargs.update(extra_kwargs)
with tempfile.TemporaryDirectory() as config_dir:
model.save_config(config_dir)
return self.model_class.from_single_file(state_dict, config=config_dir, **kwargs)
def _verify_if_layer_quantized(self, name, module, config_kwargs):
cls = self._bfl_linear_cls(config_kwargs["scheme"])
assert isinstance(module, cls), f"Layer {name} is not a {cls.__name__}, got {type(module)}"
if config_kwargs["scheme"] == "fp8r":
assert module.weight.dtype == torch.float8_e4m3fn, f"{name} weight dtype should be float8_e4m3fn"
assert module.weight_scale.shape == (module.out_features,), f"{name} weight_scale is not per row"
else:
assert module.weight.dtype == torch.uint8, f"{name} weight dtype should be uint8 (packed E2M1)"
assert module.weight.shape == (module.out_features, module.in_features // 2), (
f"{name} weight is not packed"
)
assert module.weight_scale.dtype == torch.float8_e4m3fn, f"{name} block scales should be float8_e4m3fn"
assert module.weight_scale_2.ndim == 0 and module.input_scale.ndim == 0, f"{name} global scales"
def _is_module_quantized(self, module, config_kwargs=None):
from diffusers.quantizers.bfl.utils import BFLFP8RLinear, BFLNVFP4Linear
return isinstance(module, (BFLFP8RLinear, BFLNVFP4Linear))
@is_quantization
@is_bfl
@require_accelerator
@require_accelerate
class BFLTesterMixin(BFLConfigMixin, QuantizationTesterMixin):
"""
Mixin class for testing BFL (`fp8r` / `nvfp4`) quantization on models.
Expected class attributes:
- model_class: The model class to test
Expected methods to be implemented by subclasses:
- get_init_dict(): Returns the init kwargs of the small test model
- get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass
Required properties (must be implemented by subclasses):
- bfl_checkpoints: Dict mapping a scheme (`"fp8r"`, `"nvfp4"`) to its single-file checkpoint; only the
safetensors header is fetched
Pytest mark: bfl
Use `pytest -m "not bfl"` to skip these tests
"""
@pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS))
def test_bfl_quantization_inference(self, config_name):
self._test_quantization_inference(BFLConfigMixin.BFL_CONFIGS[config_name])
@pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS))
def test_bfl_quantization_memory_footprint(self, config_name):
self._test_quantization_memory_footprint(BFLConfigMixin.BFL_CONFIGS[config_name])
def _test_quantized_layers(self, config_kwargs):
# The exporter quantizes a subset of the linears (embedders, output layers and some blocks stay in bf16),
# so the expected set is the one the checkpoint header describes, restricted to the small model's layers.
model = self._create_quantized_model(config_kwargs)
expected = self._bfl_quantized_linears(model, config_kwargs["scheme"])
quantized = set()
for name, module in model.named_modules():
if self._is_module_quantized(module):
self._verify_if_layer_quantized(name, module, config_kwargs)
quantized.add(name)
elif isinstance(module, torch.nn.Linear):
assert module.weight.dtype == config_kwargs["compute_dtype"], f"{name} is neither quantized nor bf16"
assert quantized, "No quantized layers found in model"
assert quantized == expected, f"Quantized layers differ from the checkpoint: {quantized ^ expected}"
@pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS))
def test_bfl_quantized_layers(self, config_name):
self._test_quantized_layers(BFLConfigMixin.BFL_CONFIGS[config_name])
@pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS))
@torch.no_grad()
def test_bfl_quantization_matches_unquantized(self, config_name):
"""The quantized model tracks the model it was quantized from, on the fused kernel or the fallback."""
config_kwargs = BFLConfigMixin.BFL_CONFIGS[config_name]
reference = self._load_unquantized_model().to(torch_device, config_kwargs["compute_dtype"])
expected = reference(**self.get_dummy_inputs(), return_dict=False)[0]
model = self._create_quantized_model(config_kwargs)
model.to(torch_device)
output = model(**self.get_dummy_inputs(), return_dict=False)[0]
assert output.shape == expected.shape and not torch.isnan(output).any()
cosine = torch.nn.functional.cosine_similarity(output.float().flatten(), expected.float().flatten(), dim=0)
assert cosine > {"fp8r": 0.99, "nvfp4": 0.9}[config_name], f"cosine similarity {cosine.item():.4f}"
@pytest.mark.parametrize("config_name", ["fp8r"], ids=["fp8r"])
def test_bfl_quantization_dtype_assignment(self, config_name):
self._test_quantization_dtype_assignment(BFLConfigMixin.BFL_CONFIGS[config_name])
@pytest.mark.parametrize("config_name", ["fp8r"], ids=["fp8r"])
def test_bfl_quantization_lora_inference(self, config_name):
self._test_quantization_lora_inference(BFLConfigMixin.BFL_CONFIGS[config_name])
@pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS))
def test_bfl_keep_modules_in_fp32(self, config_name):
# `to_q` is quantized in the checkpoint, so this exercises the quantizer's fp32 exclusion path rather than
# a layer the exporter already left in bf16.
_keep_in_fp32_modules = self.model_class._keep_in_fp32_modules
self.model_class._keep_in_fp32_modules = ["to_q"]
try:
self._test_keep_modules_in_fp32(BFLConfigMixin.BFL_CONFIGS[config_name])
finally:
self.model_class._keep_in_fp32_modules = _keep_in_fp32_modules
@pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS))
def test_bfl_dequantize(self, config_name):
self._test_dequantize(BFLConfigMixin.BFL_CONFIGS[config_name])
def test_bfl_rejects_unquantized_checkpoint(self):
model = self._load_unquantized_model().to(torch.bfloat16)
with tempfile.TemporaryDirectory() as config_dir:
model.save_config(config_dir)
with pytest.raises(ValueError, match="weight_scale"):
self.model_class.from_single_file(
model.state_dict(), config=config_dir, quantization_config=BFLQuantizationConfig()
)
@is_quantization
@is_bfl
@is_torch_compile
@require_accelerator
@require_accelerate
class BFLCompileTesterMixin(BFLConfigMixin, QuantizationCompileTesterMixin):
"""
Mixin class for testing torch.compile with BFL (`fp8r` / `nvfp4`) quantized models.
Expected class attributes:
- model_class: The model class to test
Expected methods to be implemented by subclasses:
- get_init_dict(): Returns the init kwargs of the small test model
- get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass
Required properties (must be implemented by subclasses):
- bfl_checkpoints: Dict mapping a scheme (`"fp8r"`, `"nvfp4"`) to its single-file checkpoint; only the
safetensors header is fetched
Pytest mark: bfl
Use `pytest -m "not bfl"` to skip these tests
"""
@pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS))
def test_bfl_torch_compile(self, config_name):
if (
config_name == "nvfp4"
and torch_device == "cuda"
and torch.cuda.get_device_capability()[0] >= 10
and is_flashinfer_available()
):
# Dynamo cannot trace FlashInfer's Python wrappers (they take a thread lock), and with graph breaks around
# them the remaining fragments recompile between calls. Expected to pass once the fused nvfp4 path wraps
# the FlashInfer calls as custom ops.
pytest.xfail("the fused nvfp4 path calls FlashInfer, which torch.compile cannot trace")
self._test_torch_compile(BFLConfigMixin.BFL_CONFIGS[config_name])
@pytest.mark.parametrize("config_name", ["fp8r"], ids=["fp8r"])
def test_bfl_torch_compile_with_group_offload(self, config_name):
self._test_torch_compile_with_group_offload(BFLConfigMixin.BFL_CONFIGS[config_name])

200
tests/models/transformers/test_models_transformer_flux3.py Normal file
View File

@@ -0,0 +1,200 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pytest
import torch
from diffusers import Flux3Transformer2DModel
from diffusers.utils.torch_utils import randn_tensor
from ...testing_utils import enable_full_determinism, torch_device
from ..testing_utils import (
AttentionTesterMixin,
BaseModelTesterConfig,
BFLCompileTesterMixin,
BFLTesterMixin,
MemoryTesterMixin,
ModelTesterMixin,
SingleFileTesterMixin,
TorchCompileTesterMixin,
TrainingTesterMixin,
)
enable_full_determinism()
class Flux3TransformerTesterConfig(BaseModelTesterConfig):
@property
def model_class(self):
return Flux3Transformer2DModel
@property
def output_shape(self) -> tuple[int, ...]:
return (16, 4)
@property
def input_shape(self) -> tuple[int, ...]:
return (16, 4)
@property
def main_input_name(self) -> str:
return "hidden_states"
@property
def uses_custom_attn_processor(self) -> bool:
return True
@property
def generator(self):
return torch.Generator("cpu").manual_seed(0)
def get_init_dict(self) -> dict[str, int | list[int]]:
return {
"in_channels": 4,
"cond_in_channels": 4,
"num_layers": 1,
"num_single_layers": 1,
"attention_head_dim": 16,
"num_attention_heads": 2,
"joint_attention_dim": 32,
"timestep_channels": 32,
"axes_dims_rope": [4, 4, 4, 4],
}
def get_dummy_inputs(
self, height: int = 4, width: int = 4, batch_size: int = 1, num_cond_tokens: int | None = None
) -> dict[str, torch.Tensor]:
# all per-mode blocks share one code object; equal text/image lengths keep compiled shapes uniform
num_tokens = height * width
generator = self.generator
inputs = {
"hidden_states": randn_tensor((batch_size, num_tokens, 4), generator=generator, device=torch_device),
"encoder_hidden_states": randn_tensor(
(batch_size, num_tokens, 32), generator=generator, device=torch_device
),
"timestep": torch.full((batch_size,), 0.5, device=torch_device),
"img_ids": torch.zeros(num_tokens, 4, device=torch_device),
"txt_ids": torch.zeros(num_tokens, 4, device=torch_device),
}
if num_cond_tokens is not None:
inputs["image_cond_hidden_states"] = randn_tensor(
(batch_size, num_cond_tokens, 4), generator=generator, device=torch_device
)
inputs["image_cond_ids"] = torch.zeros(num_cond_tokens, 4, device=torch_device)
return inputs
class TestFlux3TransformerModel(Flux3TransformerTesterConfig, ModelTesterMixin):
@torch.no_grad()
def test_image_cond_inputs(self):
torch.manual_seed(0)
model = self.model_class(**self.get_init_dict()).to(torch_device).eval()
# a distinct reference token count: the output is sliced back to the image tokens only
inputs = self.get_dummy_inputs(num_cond_tokens=10)
output = model(**inputs, return_dict=False)[0]
assert output.shape == inputs["hidden_states"].shape
# the references change the prediction
text_only = model(**{k: v for k, v in inputs.items() if not k.startswith("image_cond")}, return_dict=False)[0]
assert not torch.allclose(output, text_only)
@torch.no_grad()
def test_flow_map(self):
torch.manual_seed(0)
model = self.model_class(**self.get_init_dict(), flow_map=True).to(torch_device).eval()
inputs = self.get_dummy_inputs()
target_timestep = torch.full_like(inputs["timestep"], 0.25)
output = model(**inputs, target_timestep=target_timestep, return_dict=False)[0]
assert output.shape == inputs["hidden_states"].shape
# the target timestep changes the prediction
other = model(**inputs, target_timestep=torch.zeros_like(target_timestep), return_dict=False)[0]
assert not torch.allclose(output, other)
# a flow map requires the target timestep; a plain model refuses it
with pytest.raises(ValueError, match="target_timestep"):
model(**inputs, return_dict=False)
plain = self.model_class(**self.get_init_dict()).to(torch_device).eval()
assert plain.target_timestep_embedder is None
with pytest.raises(ValueError, match="target_timestep"):
plain(**inputs, target_timestep=target_timestep, return_dict=False)
class TestFlux3TransformerMemory(Flux3TransformerTesterConfig, MemoryTesterMixin):
pass
class TestFlux3TransformerTorchCompile(Flux3TransformerTesterConfig, TorchCompileTesterMixin):
@property
def different_shapes_for_compilation(self):
return [(4, 4), (4, 8), (8, 8)]
class TestFlux3TransformerTraining(Flux3TransformerTesterConfig, TrainingTesterMixin):
pass
class TestFlux3TransformerAttention(Flux3TransformerTesterConfig, AttentionTesterMixin):
pass
class TestFlux3SingleFile(Flux3TransformerTesterConfig, SingleFileTesterMixin):
@property
def ckpt_path(self):
return "https://huggingface.co/black-forest-labs/flux-3-image-klein-base/blob/main/model.safetensors"
@property
def pretrained_model_name_or_path(self):
return "black-forest-labs/flux-3-image-klein-base"
@property
def pretrained_model_kwargs(self):
return {"subfolder": "transformer"}
@property
def torch_dtype(self):
return torch.bfloat16
class Flux3TransformerBFLTesterConfig(Flux3TransformerTesterConfig):
"""Shared config for the BFL quantization tests: which linears to quantize comes from the real checkpoints'
headers; the model itself is the small test model in the compute dtype."""
@property
def bfl_checkpoints(self):
return {
"fp8r": "https://huggingface.co/black-forest-labs/flux-3-image-klein-base/blob/main/variants/fp8r/model.safetensors",
"nvfp4": "https://huggingface.co/black-forest-labs/flux-3-image-klein-base/blob/main/variants/nvfp4/model.safetensors",
}
@property
def torch_dtype(self):
return torch.bfloat16
def get_dummy_inputs(self, **kwargs):
"""Override to build inputs in the quantizer compute dtype (unquantized linears do strict-dtype matmuls)."""
inputs = super().get_dummy_inputs(**kwargs)
return {k: v.to(self.torch_dtype) if torch.is_floating_point(v) else v for k, v in inputs.items()}
class TestFlux3TransformerBFL(Flux3TransformerBFLTesterConfig, BFLTesterMixin):
"""BFL fp8r / nvfp4 quantization tests for the FLUX.3 image transformer."""
class TestFlux3TransformerBFLCompile(Flux3TransformerBFLTesterConfig, BFLCompileTesterMixin):
"""BFL fp8r / nvfp4 + torch.compile tests for the FLUX.3 image transformer."""

0
tests/modular_pipelines/flux3/__init__.py Normal file
View File

373
tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py Normal file
View File

@@ -0,0 +1,373 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
import PIL.Image
import pytest
import torch
from diffusers.modular_pipelines import (
Flux3AutoBlocks,
Flux3DistilledAutoBlocks,
Flux3DistilledModularPipeline,
Flux3ModularPipeline,
)
from diffusers.modular_pipelines.flux3.encoders import Flux3TextEncoderStep
from diffusers.modular_pipelines.flux3.inputs import Flux3ImageCondInputStep
from diffusers.utils import logging
from ...testing_utils import CaptureLogger, torch_device
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularGuiderTesterMixin,
ModularLoadingTesterMixin,
ModularLoraTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
FLUX3_TINY_REPO = "hf-internal-testing/tiny-flux3-image-modular"
FLUX3_DISTILLED_TINY_REPO = "hf-internal-testing/tiny-flux3-image-distilled-modular"
FLUX3_WORKFLOWS = {
"text2image": [
("text_encoder", "Flux3TextEncoderStep"),
("denoise.input", "Flux3TextInputStep"),
("denoise.prepare_latents", "Flux3PrepareLatentsStep"),
("denoise.set_timesteps", "Flux3SetTimestepsStep"),
("denoise.prepare_rope_inputs", "Flux3TextRoPEInputsStep"),
("denoise.denoise", "Flux3DenoiseStep"),
("denoise.after_denoise", "Flux3UnpackLatentsStep"),
("decode", "Flux3DecodeStep"),
],
"image_conditioned": [
("text_encoder", "Flux3TextEncoderStep"),
("vae_encoder.preprocess", "Flux3ProcessImagesInputStep"),
("vae_encoder.encode", "Flux3VaeEncoderStep"),
("denoise.input", "Flux3TextInputStep"),
("denoise.cond_input", "Flux3ImageCondInputStep"),
("denoise.prepare_latents", "Flux3PrepareLatentsStep"),
("denoise.prepare_image_cond_latents", "Flux3PrepareImageCondLatentsStep"),
("denoise.set_timesteps", "Flux3SetTimestepsStep"),
("denoise.prepare_rope_inputs", "Flux3TextRoPEInputsStep"),
("denoise.denoise", "Flux3DenoiseStep"),
("denoise.after_denoise", "Flux3UnpackLatentsStep"),
("decode", "Flux3DecodeStep"),
],
}
# `expected_workflow_defaults`: the components, inputs and component configs each workflow's blocks carry
FLUX3_COMPONENTS = {
"text_encoder": "Qwen3VLForConditionalGeneration",
"tokenizer": "AutoProcessor",
"transformer": "Flux3Transformer2DModel",
"scheduler": "FlowMatchEulerDiscreteScheduler",
"guider": "ClassifierFreeGuidance",
"vae": "AutoencoderKLFlux2",
"image_processor": "Flux3ImageProcessor",
}
# the recommended sampling settings the blocks pin on the scheduler / guider they create
FLUX3_SAMPLING = {"scheduler": {"shift": 6.93}, "guider": {"guidance_scale": 3.5}}
# every optional input with its declared default. text2image declares the 1024 x 1024 size; image_conditioned derives the
# size from the first reference image, so its branch declares none
FLUX3_INPUTS = {
"negative_prompt": None,
"max_sequence_length": 8192,
"text_encoder_out_layers": (4, 8, 12, 16, 20, 24, 28, 32),
"num_images_per_prompt": 1,
"latents": None,
"attention_kwargs": None,
"generator": None,
"num_inference_steps": 50,
"timesteps": None,
"sigmas": None,
"output_type": "pil",
}
FLUX3_TEXT2IMAGE_INPUTS = {**FLUX3_INPUTS, "height": 1024, "width": 1024}
FLUX3_IMAGE_CONDITIONED_INPUTS = {**FLUX3_INPUTS, "height": None, "width": None}
FLUX3_WORKFLOW_DEFAULTS = {
"text2image": {
"components": FLUX3_COMPONENTS,
"required_inputs": ["prompt"],
"inputs": FLUX3_TEXT2IMAGE_INPUTS,
"component_configs": FLUX3_SAMPLING,
},
"image_conditioned": {
"components": FLUX3_COMPONENTS,
"required_inputs": ["prompt", "image"],
"inputs": FLUX3_IMAGE_CONDITIONED_INPUTS,
"component_configs": FLUX3_SAMPLING,
},
}
# the step-distilled model: the same graph with distilled text encoder / text input / text RoPE / denoise steps (no
# guider, no negative prompt, 16 steps by default)
FLUX3_DISTILLED_WORKFLOWS = {
"text2image": [
("text_encoder", "Flux3DistilledTextEncoderStep"),
("denoise.input", "Flux3DistilledTextInputStep"),
("denoise.prepare_latents", "Flux3PrepareLatentsStep"),
("denoise.set_timesteps", "Flux3SetTimestepsStep"),
("denoise.prepare_rope_inputs", "Flux3DistilledTextRoPEInputsStep"),
("denoise.denoise", "Flux3DistilledDenoiseStep"),
("denoise.after_denoise", "Flux3UnpackLatentsStep"),
("decode", "Flux3DecodeStep"),
],
"image_conditioned": [
("text_encoder", "Flux3DistilledTextEncoderStep"),
("vae_encoder.preprocess", "Flux3ProcessImagesInputStep"),
("vae_encoder.encode", "Flux3VaeEncoderStep"),
("denoise.input", "Flux3DistilledTextInputStep"),
("denoise.cond_input", "Flux3ImageCondInputStep"),
("denoise.prepare_latents", "Flux3PrepareLatentsStep"),
("denoise.prepare_image_cond_latents", "Flux3PrepareImageCondLatentsStep"),
("denoise.set_timesteps", "Flux3SetTimestepsStep"),
("denoise.prepare_rope_inputs", "Flux3DistilledTextRoPEInputsStep"),
("denoise.denoise", "Flux3DistilledDenoiseStep"),
("denoise.after_denoise", "Flux3UnpackLatentsStep"),
("decode", "Flux3DecodeStep"),
],
}
FLUX3_DISTILLED_COMPONENTS = {name: cls for name, cls in FLUX3_COMPONENTS.items() if name != "guider"}
FLUX3_DISTILLED_SAMPLING = {"scheduler": {"shift": 6.93}}
# no negative prompt and no attention kwargs, 16 steps by default
FLUX3_DISTILLED_INPUTS = {
name: value for name, value in FLUX3_INPUTS.items() if name not in ("negative_prompt", "attention_kwargs")
}
FLUX3_DISTILLED_INPUTS["num_inference_steps"] = 16
FLUX3_DISTILLED_WORKFLOW_DEFAULTS = {
"text2image": {
"components": FLUX3_DISTILLED_COMPONENTS,
"required_inputs": ["prompt"],
"inputs": {**FLUX3_DISTILLED_INPUTS, "height": 1024, "width": 1024},
"component_configs": FLUX3_DISTILLED_SAMPLING,
},
"image_conditioned": {
"components": FLUX3_DISTILLED_COMPONENTS,
"required_inputs": ["prompt", "image"],
"inputs": {**FLUX3_DISTILLED_INPUTS, "height": None, "width": None},
"component_configs": FLUX3_DISTILLED_SAMPLING,
},
}
class Flux3ModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = Flux3ModularPipeline
pipeline_blocks_class = Flux3AutoBlocks
pretrained_model_name_or_path = FLUX3_TINY_REPO
params = frozenset(["prompt", "height", "width"])
batch_params = frozenset()
expected_workflow_blocks = FLUX3_WORKFLOWS
expected_workflow_defaults = FLUX3_WORKFLOW_DEFAULTS
@staticmethod
def get_reference_images(sizes=((64, 64),), seed=0):
rng = np.random.RandomState(seed)
return [PIL.Image.fromarray(np.uint8(rng.rand(height, width, 3) * 255)) for width, height in sizes]
def get_dummy_inputs(self, seed=0):
return {
"prompt": "a tiny fox in the snow",
"generator": self.get_generator(seed),
"num_inference_steps": 2,
"height": 64,
"width": 64,
"max_sequence_length": 8,
"text_encoder_out_layers": (1,),
"output_type": "pt",
}
class TestFlux3ModularPipelineFast(Flux3ModularPipelineTesterConfig, ModularPipelineTesterMixin):
@pytest.mark.skip(reason="FLUX.3 takes one prompt per call; batched prompts are not supported.")
def test_inference_batch_consistent(self):
pass
@pytest.mark.skip(reason="FLUX.3 takes one prompt per call; batched prompts are not supported.")
def test_inference_batch_single_identical(self):
pass
def test_num_images_per_prompt(self):
# one prompt per call, several images from it
super().test_num_images_per_prompt(batch_sizes=[1])
def test_multiple_reference_images(self):
# references keep their own resolution; several of them go into one conditioning sequence
pipe = self.get_pipeline().to(torch_device)
inputs = {**self.get_dummy_inputs(), "num_inference_steps": 1}
references = self.get_reference_images(sizes=((64, 48), (32, 32)))
images = pipe(**inputs, image=references, output="images")
assert images.shape == (1, 3, 64, 64)
# only using the first reference changes the result
single = pipe(**inputs, image=references[:1], output="images")
assert not torch.allclose(images, single)
def test_image_to_image_default_size(self):
# without a size, image-to-image generates at the default area (1024 x 1024 for the real model) at the first
# reference's aspect ratio, rounded to the token grid; a given side overrides its default. The default area is
# shrunk to 64 x 64 for the tiny model.
blocks = self.pipeline_blocks_class().get_workflow("image_conditioned")
blocks.sub_blocks["denoise.cond_input"] = Flux3ImageCondInputStep(default_height=64, default_width=64)
pipe = blocks.init_pipeline(self.pretrained_model_name_or_path)
pipe.load_components(dtype=torch.float32)
pipe.to(torch_device)
inputs = {k: v for k, v in self.get_dummy_inputs().items() if k not in ("height", "width")}
inputs["num_inference_steps"] = 1
inputs["image"] = self.get_reference_images(sizes=((64, 48),)) # 64 x 48 pixels, 4 pixels per token
# 64 x 64 of area at the 4:3 aspect ratio: 72 x 56 on the 4-pixel grid
assert pipe(**inputs, output="images").shape[-2:] == (56, 72)
assert pipe(**inputs, height=32, width=32, output="images").shape[-2:] == (32, 32)
def test_text_window_pads_to_multiple(self):
# the prompt is right-padded to the block's pad multiple (80 for the real model), capped at `max_sequence_length`
text_node = Flux3TextEncoderStep(pad_multiple=20).init_pipeline(self.pretrained_model_name_or_path)
text_node.load_components(dtype=torch.float32)
text_node.to(torch_device)
short = text_node(prompt="a tiny fox", text_encoder_out_layers=(1,), output="prompt_embeds")
long = text_node(prompt="a tiny fox in the snow " * 10, text_encoder_out_layers=(1,), output="prompt_embeds")
assert short.shape[1] % 20 == 0 and long.shape[1] % 20 == 0
assert long.shape[1] > short.shape[1]
# the cap wins over the pad multiple
capped = text_node(
prompt="a tiny fox", max_sequence_length=8, text_encoder_out_layers=(1,), output="prompt_embeds"
)
assert capped.shape[1] == 8
# TODO: move to `ModularGuiderTesterMixin` -- guider-disabled negative-prompt handling is generic
def test_guider_disabled_skips_negative_prompt_with_warning(self):
pipe = self.get_pipeline()
inputs = {**self.get_dummy_inputs(), "negative_prompt": "blurry", "num_inference_steps": 1}
assert pipe(**inputs, output="negative_prompt_embeds") is not None
pipe.guider.disable()
logger = logging.get_logger("diffusers.modular_pipelines.flux3.encoders")
with CaptureLogger(logger) as cap:
negative_prompt_embeds = pipe(**inputs, output="negative_prompt_embeds")
assert negative_prompt_embeds is None
assert "guider is disabled" in cap.out
# TODO: generic too (huggingface/diffusers#14469): a standalone text encoder node has no guider, so it
# encodes the negative prompt only when given one, and never a "" placeholder
def test_text_encoder_standalone_encodes_negative_prompt_only_when_given(self):
text_node = (
self.pipeline_blocks_class().sub_blocks["text_encoder"].init_pipeline(self.pretrained_model_name_or_path)
)
text_node.load_components(dtype=torch.float32)
text_node.to(torch_device)
assert "guider" not in text_node.component_names
inputs = {
k: v
for k, v in self.get_dummy_inputs().items()
if k in ("prompt", "max_sequence_length", "text_encoder_out_layers")
}
assert text_node(**inputs, output="negative_prompt_embeds") is None
assert text_node(**inputs, negative_prompt="blurry", output="negative_prompt_embeds") is not None
def test_text_input_requires_negative_prompt_embeds_with_cfg(self):
# node graph as in the docs: embeddings from a standalone text node (no negative prompt) fed to a
# generation node whose guider needs unconditional embeddings -> a clear error, not a crash in the transformer
blocks = self.pipeline_blocks_class()
text_node = blocks.sub_blocks.pop("text_encoder").init_pipeline(self.pretrained_model_name_or_path)
blocks.sub_blocks.pop("decode")
generate_node = blocks.init_pipeline(self.pretrained_model_name_or_path)
for node in (text_node, generate_node):
node.load_components(dtype=torch.float32)
node.to(torch_device)
inputs = self.get_dummy_inputs()
text_inputs = {k: inputs[k] for k in ("prompt", "max_sequence_length", "text_encoder_out_layers")}
embeddings = text_node(**text_inputs, output=["prompt_embeds", "negative_prompt_embeds"])
assert embeddings["negative_prompt_embeds"] is None
generate_inputs = {k: v for k, v in inputs.items() if k not in text_inputs}
with pytest.raises(ValueError, match="unconditional embeddings"):
generate_node(**embeddings, **generate_inputs, output="latents")
# with the negative prompt encoded (an empty string for plain CFG) the graph runs
embeddings = text_node(**text_inputs, negative_prompt="", output=["prompt_embeds", "negative_prompt_embeds"])
assert generate_node(**embeddings, **generate_inputs, output="latents") is not None
class TestFlux3ModularPipelineGuider(Flux3ModularPipelineTesterConfig, ModularGuiderTesterMixin):
def test_guider_cfg(self):
# the tiny random model's cond/uncond predictions differ only slightly
super().test_guider_cfg(expected_max_diff=1e-3)
class TestFlux3ModularPipelineLoading(Flux3ModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestFlux3ModularPipelineMemory(Flux3ModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
class TestFlux3ModularPipelineWorkflow(Flux3ModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class Flux3DistilledModularPipelineTesterConfig(Flux3ModularPipelineTesterConfig):
pipeline_class = Flux3DistilledModularPipeline
pipeline_blocks_class = Flux3DistilledAutoBlocks
pretrained_model_name_or_path = FLUX3_DISTILLED_TINY_REPO
expected_workflow_blocks = FLUX3_DISTILLED_WORKFLOWS
expected_workflow_defaults = FLUX3_DISTILLED_WORKFLOW_DEFAULTS
class TestFlux3DistilledModularPipelineFast(Flux3DistilledModularPipelineTesterConfig, ModularPipelineTesterMixin):
@pytest.mark.skip(reason="FLUX.3 takes one prompt per call; batched prompts are not supported.")
def test_inference_batch_consistent(self):
pass
@pytest.mark.skip(reason="FLUX.3 takes one prompt per call; batched prompts are not supported.")
def test_inference_batch_single_identical(self):
pass
def test_num_images_per_prompt(self):
super().test_num_images_per_prompt(batch_sizes=[1])
def test_no_guider_no_negative_prompt(self):
# the step-distilled model samples without classifier-free guidance: no guider, and the block set neither
# takes a negative prompt nor produces negative embeddings
pipe = self.get_pipeline()
assert "guider" not in pipe.component_names
blocks = self.pipeline_blocks_class()
assert {"negative_prompt", "negative_prompt_embeds"}.isdisjoint({p.name for p in blocks.inputs})
assert {"negative_prompt_embeds", "negative_txt_ids"}.isdisjoint({p.name for p in blocks.intermediate_outputs})
def test_image_to_image(self):
pipe = self.get_pipeline()
inputs = {**self.get_dummy_inputs(), "image": self.get_reference_images(sizes=((64, 48), (32, 32)))}
images = pipe(**inputs, output="images")
assert images.shape == (1, 3, 64, 64)
class TestFlux3DistilledModularPipelineLoading(Flux3DistilledModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestFlux3DistilledModularPipelineWorkflow(Flux3DistilledModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestFlux3ModularPipelineLoRA(Flux3ModularPipelineTesterConfig, ModularLoraTesterMixin):
pass

8
tests/testing_utils.py
View File

@@ -470,6 +470,14 @@ def is_sdnq(test_case):
return pytest.mark.sdnq(test_case)
def is_bfl(test_case):
"""
Decorator marking a test as a BFL (fp8r / nvfp4) quantization test. These tests can be filtered using:
pytest -m "not bfl" to skip pytest -m bfl to run only these tests
"""
return pytest.mark.bfl(test_case)
def is_context_parallel(test_case):
"""
Decorator marking a test as a context parallel inference test. These tests can be filtered using: