From 7ecb1e1b4fe500be7fb7381a7bb79a24d7fa3941 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 15 Oct 2026 22:00:50 +0530 Subject: [PATCH] Flux3 image (#15046) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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](https://github.com/docker/setup-buildx-action/compare/8d2750c68a42422c14e847fe6c8ac0403b4cbd6f...bb05f3f5519dd87d3ba754cc423b652a5edd6d2c) Updates `docker/login-action` from 3.7.0 to 4.6.0 - [Release notes](https://github.com/docker/login-action/releases) - [Commits](https://github.com/docker/login-action/compare/c94ce9fb468520275223c153574b00df6fe4bcc9...dbcb813823bdd20940b903addbd779551569679f) Updates `docker/build-push-action` from 6.19.2 to 7.3.0 - [Release notes](https://github.com/docker/build-push-action/releases) - [Commits](https://github.com/docker/build-push-action/compare/10e90e3645eae34f1e60eeb005ba3a3d33f178e8...53b7df96c91f9c12dcc8a07bcb9ccacbed38856a) Updates `huggingface/doc-builder/.github/workflows/build_main_documentation.yml` from e60a538eea9817ab312196d0d233604b01697265 to af49b8dc3859c3d80c7220c8c0ba146431132e49 - [Release notes](https://github.com/huggingface/doc-builder/releases) - [Commits](https://github.com/huggingface/doc-builder/compare/e60a538eea9817ab312196d0d233604b01697265...af49b8dc3859c3d80c7220c8c0ba146431132e49) 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](https://github.com/huggingface/doc-builder/compare/e60a538eea9817ab312196d0d233604b01697265...af49b8dc3859c3d80c7220c8c0ba146431132e49) Updates `anthropics/claude-code-action` from 1.0.155 to 1.0.193 - [Release notes](https://github.com/anthropics/claude-code-action/releases) - [Commits](https://github.com/anthropics/claude-code-action/compare/80b31826338489861333dc17217865dfe8085cdc...9d7150bc8a3dae8149739a88019d192b579ad90c) 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](https://github.com/huggingface/security-workflows/compare/dc6ca34688e6876c2dd18750719b44d177586c17...e0d40aaf3c566a5e820787840453f45f909e337f) 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](https://github.com/actions/labeler/compare/8558fd74291d67161a8a78ce36a881fa63b766a9...bf12e9b00b37c5c0ca2b87b79b2daf7891dbda13) 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](https://github.com/pypa/gh-action-pypi-publish/compare/ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e...dc37677b2e1c63e2034f94d8a5b11f265b73ba33) Updates `huggingface/tailscale-action` from 7d53c9737e53934c30290b5524d1c9b4a7c98c8a to 07d36e4c3b68707bd3fe070dbb96f03f993a6efd - [Release notes](https://github.com/huggingface/tailscale-action/releases) - [Commits](https://github.com/huggingface/tailscale-action/compare/7d53c9737e53934c30290b5524d1c9b4a7c98c8a...07d36e4c3b68707bd3fe070dbb96f03f993a6efd) Updates `trufflesecurity/trufflehog` from 3.94.2 to 3.97.0 - [Release notes](https://github.com/trufflesecurity/trufflehog/releases) - [Commits](https://github.com/trufflesecurity/trufflehog/compare/6bd2d14f7a4bc1e569fa3550efa7ec632a4fa67b...bcfcf73aaf4759d4dadc2783177c245a02792318) 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](https://github.com/crate-ci/typos/compare/65120634e79d8374d1aa2f27e54baa0c364fff5a...8a48f81b6c64dcfea44b3633223084c4be58ac5f) Updates `huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml` from 9ad2de8582b56c017cb530c1165116d40433f1c6 to af49b8dc3859c3d80c7220c8c0ba146431132e49 - [Release notes](https://github.com/huggingface/doc-builder/releases) - [Commits](https://github.com/huggingface/doc-builder/compare/9ad2de8582b56c017cb530c1165116d40433f1c6...af49b8dc3859c3d80c7220c8c0ba146431132e49) --- 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] 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 * Bump doc-builder pin to 9fc8a41 in build_pr_documentation.yml Co-Authored-By: Claude Opus 4.8 * Bump doc-builder pin to 9fc8a41 in upload_pr_documentation.yml Co-Authored-By: Claude Opus 4.8 --------- Co-authored-by: Claude Opus 4.8 * 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 * integration-notes: ship the knight example image with the doc examples Co-Authored-By: Claude Fable 5.1 * integration-notes: drop the tokenizer config workaround (not needed with transformers 5.x) Co-Authored-By: Claude Fable 5.1 * 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 * 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 * 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 * 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 * Document reproducible generation with user-provided latents Co-Authored-By: Claude Fable 5.1 * Run the QK norm and RoPE in fp32 with a single cast back, matching the BFL implementation Co-Authored-By: Claude Fable 5.1 * parity script: run against the Hub repos and print the imported diffusers source Co-Authored-By: Claude Fable 5.1 * Wrap a forward docstring line for the doc-builder style check Co-Authored-By: Claude Fable 5.1 * 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 * 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 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 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] 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 Co-authored-by: yiyixuxu Co-authored-by: linoytsaban Co-authored-by: apolinário Co-authored-by: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com> Co-authored-by: Aritra Roy Gosthipaty Co-authored-by: Yumeng Li <109459218+YumengLi007@users.noreply.github.com> --- .github/workflows/benchmark.yml | 4 +- .github/workflows/bot_pytest.yml | 4 +- .github/workflows/build_docker_images.yml | 12 +- .github/workflows/build_pr_documentation.yml | 4 +- .github/workflows/claude_review.yml | 4 +- .github/workflows/codeql.yml | 2 +- .github/workflows/issue_labeler.yml | 2 +- .../workflows/mirror_community_pipeline.yml | 4 +- .github/workflows/nightly_tests.yml | 46 +- .../workflows/notify_slack_about_release.yml | 4 +- .github/workflows/pr_dependency_test.yml | 4 +- .github/workflows/pr_labeler.yml | 4 +- .github/workflows/pr_link_issue_reminder.yml | 4 +- .github/workflows/pr_modular_tests.yml | 14 +- .github/workflows/pr_test_fetcher.yml | 12 +- .github/workflows/pr_tests.yml | 20 +- .github/workflows/pr_tests_gpu.yml | 24 +- .../workflows/pr_torch_dependency_test.yml | 4 +- .github/workflows/push_tests.yml | 24 +- .github/workflows/push_tests_fast.yml | 4 +- .github/workflows/push_tests_mps.yml | 4 +- .github/workflows/pypi_publish.yaml | 12 +- .github/workflows/release_tests_fast.yml | 28 +- .github/workflows/ssh-pr-runner.yml | 4 +- .github/workflows/ssh-runner.yml | 2 +- .github/workflows/stale.yml | 4 +- .github/workflows/trufflehog.yml | 4 +- .github/workflows/typos.yml | 4 +- .github/workflows/update_metadata.yml | 2 +- docs/source/en/_toctree.yml | 6 + docs/source/en/api/loaders/lora.md | 4 + .../source/en/api/models/flux3_transformer.md | 34 + docs/source/en/api/pipelines/flux3.md | 275 ++ docs/source/en/api/quantization.md | 4 + docs/source/en/quantization/bfl.md | 63 + docs/source/en/quantization/overview.md | 1 + examples/dreambooth/README_flux3_klein.md | 224 ++ .../dreambooth/requirements_flux3_klein.txt | 11 + .../test_dreambooth_lora_flux3_klein.py | 194 ++ ...est_dreambooth_lora_flux3_klein_img2img.py | 121 + .../train_dreambooth_lora_flux3_klein.py | 2210 +++++++++++++++++ ...ain_dreambooth_lora_flux3_klein_img2img.py | 2080 ++++++++++++++++ scripts/convert_flux3_to_diffusers.py | 112 + src/diffusers/__init__.py | 30 + src/diffusers/loaders/__init__.py | 2 + .../loaders/lora_conversion_utils.py | 20 + src/diffusers/loaders/lora_pipeline.py | 183 ++ src/diffusers/loaders/single_file_model.py | 8 + src/diffusers/loaders/single_file_utils.py | 77 + src/diffusers/models/__init__.py | 2 + src/diffusers/models/transformers/__init__.py | 1 + .../models/transformers/transformer_flux3.py | 656 +++++ src/diffusers/modular_pipelines/__init__.py | 12 + .../modular_pipelines/flux3/__init__.py | 49 + .../modular_pipelines/flux3/before_denoise.py | 436 ++++ .../modular_pipelines/flux3/decoders.py | 170 ++ .../modular_pipelines/flux3/denoise.py | 426 ++++ .../modular_pipelines/flux3/encoders.py | 397 +++ .../flux3/image_processor.py | 116 + .../modular_pipelines/flux3/inputs.py | 319 +++ .../flux3/modular_blocks_flux3.py | 406 +++ .../flux3/modular_blocks_flux3_distilled.py | 333 +++ .../flux3/modular_pipeline.py | 57 + .../modular_pipelines/modular_pipeline.py | 2 + src/diffusers/quantizers/auto.py | 4 + src/diffusers/quantizers/bfl/__init__.py | 1 + src/diffusers/quantizers/bfl/bfl_quantizer.py | 152 ++ src/diffusers/quantizers/bfl/utils.py | 310 +++ .../quantizers/quantization_config.py | 25 + src/diffusers/utils/__init__.py | 1 + src/diffusers/utils/dummy_bfl_objects.py | 17 + src/diffusers/utils/dummy_pt_objects.py | 15 + .../dummy_torch_and_transformers_objects.py | 60 + src/diffusers/utils/import_utils.py | 5 + tests/conftest.py | 1 + tests/models/testing_utils/__init__.py | 6 + tests/models/testing_utils/quantization.py | 287 +++ .../test_models_transformer_flux3.py | 200 ++ tests/modular_pipelines/flux3/__init__.py | 0 .../flux3/test_modular_pipeline_flux3.py | 373 +++ tests/testing_utils.py | 8 + 81 files changed, 10638 insertions(+), 132 deletions(-) create mode 100644 docs/source/en/api/models/flux3_transformer.md create mode 100644 docs/source/en/api/pipelines/flux3.md create mode 100644 docs/source/en/quantization/bfl.md create mode 100644 examples/dreambooth/README_flux3_klein.md create mode 100644 examples/dreambooth/requirements_flux3_klein.txt create mode 100644 examples/dreambooth/test_dreambooth_lora_flux3_klein.py create mode 100644 examples/dreambooth/test_dreambooth_lora_flux3_klein_img2img.py create mode 100644 examples/dreambooth/train_dreambooth_lora_flux3_klein.py create mode 100644 examples/dreambooth/train_dreambooth_lora_flux3_klein_img2img.py create mode 100644 scripts/convert_flux3_to_diffusers.py create mode 100644 src/diffusers/models/transformers/transformer_flux3.py create mode 100644 src/diffusers/modular_pipelines/flux3/__init__.py create mode 100644 src/diffusers/modular_pipelines/flux3/before_denoise.py create mode 100644 src/diffusers/modular_pipelines/flux3/decoders.py create mode 100644 src/diffusers/modular_pipelines/flux3/denoise.py create mode 100644 src/diffusers/modular_pipelines/flux3/encoders.py create mode 100644 src/diffusers/modular_pipelines/flux3/image_processor.py create mode 100644 src/diffusers/modular_pipelines/flux3/inputs.py create mode 100644 src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py create mode 100644 src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py create mode 100644 src/diffusers/modular_pipelines/flux3/modular_pipeline.py create mode 100644 src/diffusers/quantizers/bfl/__init__.py create mode 100644 src/diffusers/quantizers/bfl/bfl_quantizer.py create mode 100644 src/diffusers/quantizers/bfl/utils.py create mode 100644 src/diffusers/utils/dummy_bfl_objects.py create mode 100644 tests/models/transformers/test_models_transformer_flux3.py create mode 100644 tests/modular_pipelines/flux3/__init__.py create mode 100644 tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index a1ea5e80fe..454cdff564 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -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 }} diff --git a/.github/workflows/bot_pytest.yml b/.github/workflows/bot_pytest.yml index 68ed8c5ca1..bd4a159335 100644 --- a/.github/workflows/bot_pytest.yml +++ b/.github/workflows/bot_pytest.yml @@ -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 diff --git a/.github/workflows/build_docker_images.yml b/.github/workflows/build_docker_images.yml index b582d1aef2..f143f6de3e 100644 --- a/.github/workflows/build_docker_images.yml +++ b/.github/workflows/build_docker_images.yml @@ -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 }} diff --git a/.github/workflows/build_pr_documentation.yml b/.github/workflows/build_pr_documentation.yml index df1e4139f1..4a58425df0 100644 --- a/.github/workflows/build_pr_documentation.yml +++ b/.github/workflows/build_pr_documentation.yml @@ -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' diff --git a/.github/workflows/claude_review.yml b/.github/workflows/claude_review.yml index f49d2ef9a5..afabdb49f7 100644 --- a/.github/workflows/claude_review.yml +++ b/.github/workflows/claude_review.yml @@ -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). diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index 587d168ca3..7da9ad364c 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -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 diff --git a/.github/workflows/issue_labeler.yml b/.github/workflows/issue_labeler.yml index a875559676..2fd7faea16 100644 --- a/.github/workflows/issue_labeler.yml +++ b/.github/workflows/issue_labeler.yml @@ -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 diff --git a/.github/workflows/mirror_community_pipeline.yml b/.github/workflows/mirror_community_pipeline.yml index a2c6bb15f4..128eb61c5c 100644 --- a/.github/workflows/mirror_community_pipeline.yml +++ b/.github/workflows/mirror_community_pipeline.yml @@ -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 diff --git a/.github/workflows/nightly_tests.yml b/.github/workflows/nightly_tests.yml index 536d46611d..5a6cf75a95 100644 --- a/.github/workflows/nightly_tests.yml +++ b/.github/workflows/nightly_tests.yml @@ -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 diff --git a/.github/workflows/notify_slack_about_release.yml b/.github/workflows/notify_slack_about_release.yml index d014528e01..c7f35bd6fc 100644 --- a/.github/workflows/notify_slack_about_release.yml +++ b/.github/workflows/notify_slack_about_release.yml @@ -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' diff --git a/.github/workflows/pr_dependency_test.yml b/.github/workflows/pr_dependency_test.yml index cd84b61b48..44653458e3 100644 --- a/.github/workflows/pr_dependency_test.yml +++ b/.github/workflows/pr_dependency_test.yml @@ -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 diff --git a/.github/workflows/pr_labeler.yml b/.github/workflows/pr_labeler.yml index 9df5dc517a..691dc2b6f0 100644 --- a/.github/workflows/pr_labeler.yml +++ b/.github/workflows/pr_labeler.yml @@ -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: diff --git a/.github/workflows/pr_link_issue_reminder.yml b/.github/workflows/pr_link_issue_reminder.yml index 675df095af..79d9cce9e3 100644 --- a/.github/workflows/pr_link_issue_reminder.yml +++ b/.github/workflows/pr_link_issue_reminder.yml @@ -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" diff --git a/.github/workflows/pr_modular_tests.yml b/.github/workflows/pr_modular_tests.yml index 31d9c60b2b..117a6321e5 100644 --- a/.github/workflows/pr_modular_tests.yml +++ b/.github/workflows/pr_modular_tests.yml @@ -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 diff --git a/.github/workflows/pr_test_fetcher.yml b/.github/workflows/pr_test_fetcher.yml index 2315e25cd4..b247aee41e 100644 --- a/.github/workflows/pr_test_fetcher.yml +++ b/.github/workflows/pr_test_fetcher.yml @@ -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 diff --git a/.github/workflows/pr_tests.yml b/.github/workflows/pr_tests.yml index fa915ee179..b25907d343 100644 --- a/.github/workflows/pr_tests.yml +++ b/.github/workflows/pr_tests.yml @@ -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 diff --git a/.github/workflows/pr_tests_gpu.yml b/.github/workflows/pr_tests_gpu.yml index 232ff7f3a5..1e9c2d0eff 100644 --- a/.github/workflows/pr_tests_gpu.yml +++ b/.github/workflows/pr_tests_gpu.yml @@ -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 diff --git a/.github/workflows/pr_torch_dependency_test.yml b/.github/workflows/pr_torch_dependency_test.yml index 1a45cb00f9..b8454009f7 100644 --- a/.github/workflows/pr_torch_dependency_test.yml +++ b/.github/workflows/pr_torch_dependency_test.yml @@ -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 diff --git a/.github/workflows/push_tests.yml b/.github/workflows/push_tests.yml index 31f73046be..d93ca5dbc2 100644 --- a/.github/workflows/push_tests.yml +++ b/.github/workflows/push_tests.yml @@ -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 diff --git a/.github/workflows/push_tests_fast.yml b/.github/workflows/push_tests_fast.yml index 5d7fe6f1e1..5a0dbcd2c8 100644 --- a/.github/workflows/push_tests_fast.yml +++ b/.github/workflows/push_tests_fast.yml @@ -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 diff --git a/.github/workflows/push_tests_mps.yml b/.github/workflows/push_tests_mps.yml index 5ce3028bc6..a0e9d7f13d 100644 --- a/.github/workflows/push_tests_mps.yml +++ b/.github/workflows/push_tests_mps.yml @@ -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 diff --git a/.github/workflows/pypi_publish.yaml b/.github/workflows/pypi_publish.yaml index 0c9b50f4a6..3c65e552c0 100644 --- a/.github/workflows/pypi_publish.yaml +++ b/.github/workflows/pypi_publish.yaml @@ -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 diff --git a/.github/workflows/release_tests_fast.yml b/.github/workflows/release_tests_fast.yml index 2a1f32fdef..83aaf48790 100644 --- a/.github/workflows/release_tests_fast.yml +++ b/.github/workflows/release_tests_fast.yml @@ -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 diff --git a/.github/workflows/ssh-pr-runner.yml b/.github/workflows/ssh-pr-runner.yml index 8d3018a781..836aff8c66 100644 --- a/.github/workflows/ssh-pr-runner.yml +++ b/.github/workflows/ssh-pr-runner.yml @@ -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 }} diff --git a/.github/workflows/ssh-runner.yml b/.github/workflows/ssh-runner.yml index c739e236d7..118a0ea20b 100644 --- a/.github/workflows/ssh-runner.yml +++ b/.github/workflows/ssh-runner.yml @@ -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 diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml index 984c8a090f..4c7f46722a 100644 --- a/.github/workflows/stale.yml +++ b/.github/workflows/stale.yml @@ -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 diff --git a/.github/workflows/trufflehog.yml b/.github/workflows/trufflehog.yml index 871cf74f76..a520a1cf97 100644 --- a/.github/workflows/trufflehog.yml +++ b/.github/workflows/trufflehog.yml @@ -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 diff --git a/.github/workflows/typos.yml b/.github/workflows/typos.yml index ef0d333c72..1d061eac70 100644 --- a/.github/workflows/typos.yml +++ b/.github/workflows/typos.yml @@ -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 diff --git a/.github/workflows/update_metadata.yml b/.github/workflows/update_metadata.yml index c1c8e546c6..6e47cd7bd3 100644 --- a/.github/workflows/update_metadata.yml +++ b/.github/workflows/update_metadata.yml @@ -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: | diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index ab98a52430..67854093f2 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -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 diff --git a/docs/source/en/api/loaders/lora.md b/docs/source/en/api/loaders/lora.md index 4722d08282..4c8eb20b29 100644 --- a/docs/source/en/api/loaders/lora.md +++ b/docs/source/en/api/loaders/lora.md @@ -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 diff --git a/docs/source/en/api/models/flux3_transformer.md b/docs/source/en/api/models/flux3_transformer.md new file mode 100644 index 0000000000..e87b20ef48 --- /dev/null +++ b/docs/source/en/api/models/flux3_transformer.md @@ -0,0 +1,34 @@ + + +# 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 diff --git a/docs/source/en/api/pipelines/flux3.md b/docs/source/en/api/pipelines/flux3.md new file mode 100644 index 0000000000..b33b86c479 --- /dev/null +++ b/docs/source/en/api/pipelines/flux3.md @@ -0,0 +1,275 @@ + + +# 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 ``, ``, ... in the order they are passed to `image`: + +```python +combined = pipe( + prompt="The fox from sitting on the wooden bench from , 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 `` 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 centered against a " + "solid, vibrant chartreuse background . " + '[{"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 , move the knight figure upwards and to the left. Keep the rest of the image " + "unchanged, preserving the backdrop . " + '[{"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 diff --git a/docs/source/en/api/quantization.md b/docs/source/en/api/quantization.md index eab5fc663b..415762fb1d 100644 --- a/docs/source/en/api/quantization.md +++ b/docs/source/en/api/quantization.md @@ -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 diff --git a/docs/source/en/quantization/bfl.md b/docs/source/en/quantization/bfl.md new file mode 100644 index 0000000000..fae0ccaf1e --- /dev/null +++ b/docs/source/en/quantization/bfl.md @@ -0,0 +1,63 @@ + + +# 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. diff --git a/docs/source/en/quantization/overview.md b/docs/source/en/quantization/overview.md index 9ff0de9a47..06a6b55ce6 100644 --- a/docs/source/en/quantization/overview.md +++ b/docs/source/en/quantization/overview.md @@ -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. | diff --git a/examples/dreambooth/README_flux3_klein.md b/examples/dreambooth/README_flux3_klein.md new file mode 100644 index 0000000000..2649356e29 --- /dev/null +++ b/examples/dreambooth/README_flux3_klein.md @@ -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 ``, ``, ... in the order +they are passed. i2i training uses a single conditioning image, so captions address it as ``, e.g. +`"Relight 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 `` 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 , move the knight 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. diff --git a/examples/dreambooth/requirements_flux3_klein.txt b/examples/dreambooth/requirements_flux3_klein.txt new file mode 100644 index 0000000000..f7791342a0 --- /dev/null +++ b/examples/dreambooth/requirements_flux3_klein.txt @@ -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 diff --git a/examples/dreambooth/test_dreambooth_lora_flux3_klein.py b/examples/dreambooth/test_dreambooth_lora_flux3_klein.py new file mode 100644 index 0000000000..31260859ab --- /dev/null +++ b/examples/dreambooth/test_dreambooth_lora_flux3_klein.py @@ -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"}) diff --git a/examples/dreambooth/test_dreambooth_lora_flux3_klein_img2img.py b/examples/dreambooth/test_dreambooth_lora_flux3_klein_img2img.py new file mode 100644 index 0000000000..845d956ea3 --- /dev/null +++ b/examples/dreambooth/test_dreambooth_lora_flux3_klein_img2img.py @@ -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())) diff --git a/examples/dreambooth/train_dreambooth_lora_flux3_klein.py b/examples/dreambooth/train_dreambooth_lora_flux3_klein.py new file mode 100644 index 0000000000..cd1f4ba779 --- /dev/null +++ b/examples/dreambooth/train_dreambooth_lora_flux3_klein.py @@ -0,0 +1,2210 @@ +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. 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. + +# /// script +# dependencies = [ +# "diffusers @ git+https://github.com/huggingface/diffusers.git", +# "torch>=2.0.0", +# "accelerate>=0.31.0", +# "transformers>=4.41.2", +# "ftfy", +# "tensorboard", +# "Jinja2", +# "peft>=0.11.1", +# "sentencepiece", +# "torchvision", +# "datasets", +# "bitsandbytes", +# "prodigyopt", +# ] +# /// + +import argparse +import copy +import itertools +import json +import logging +import math +import os +import random +import shutil +import warnings +from contextlib import nullcontext +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import transformers +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed +from huggingface_hub import create_repo, upload_folder +from huggingface_hub.utils import insecure_hashlib +from peft import LoraConfig, prepare_model_for_kbit_training, set_peft_model_state_dict +from peft.utils import get_peft_model_state_dict +from PIL import Image +from PIL.ImageOps import exif_transpose +from torch.utils.data import Dataset +from torch.utils.data.sampler import BatchSampler +from torchvision import transforms +from torchvision.transforms import functional as TF +from tqdm.auto import tqdm +from transformers import AutoProcessor, Qwen3VLForConditionalGeneration + +import diffusers +from diffusers import ( + AutoencoderKLFlux2, + BitsAndBytesConfig, + FlowMatchEulerDiscreteScheduler, + Flux3Transformer2DModel, + ModularPipeline, +) +from diffusers.loaders import Flux3LoraLoaderMixin +from diffusers.modular_pipelines.flux3 import Flux3AutoBlocks +from diffusers.optimization import get_scheduler +from diffusers.training_utils import ( + _collate_lora_metadata, + _to_cpu_contiguous, + cast_training_params, + compute_density_for_timestep_sampling, + compute_loss_weighting_for_sd3, + find_nearest_bucket, + free_memory, + generate_aspect_ratio_buckets, + get_fsdp_kwargs_from_accelerator, + offload_models, + parse_buckets_string, + wrap_with_fsdp, +) +from diffusers.utils import ( + check_min_version, + convert_unet_state_dict_to_peft, + is_wandb_available, +) +from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card +from diffusers.utils.import_utils import is_torch_npu_available +from diffusers.utils.torch_utils import is_compiled_module + + +if getattr(torch, "distributed", None) is not None: + import torch.distributed as dist + +if is_wandb_available(): + import wandb + +# Will error if the minimal version of diffusers is not installed. Remove at your own risks. +check_min_version("0.40.0.dev0") + +logger = get_logger(__name__) + +# FLUX.3 ships no `scheduler` subfolder: the modular pipeline builds its scheduler `from_config` with a fixed +# shift of 6.93 (the `Flux3SetTimestepsStep` default, and the reference sampler's `alpha`). There is no +# resolution-dependent dynamic shifting. Constructed here for the same reason, rather than loaded. +FLOW_SHIFT = 6.93 + +# FLUX.3 image RoPE scale applied to the H/W position-id axes (see the flux3 modular pipeline). +IMAGE_ROPE_SCALE = 0.5 + +# The text window is right-padded to a multiple of this many tokens (the `Flux3TextEncoderStep` default). +TEXT_PAD_MULTIPLE = 80 + + +# --- FLUX.3 latent / text helpers (bit-identical to the modular pipeline blocks) --- +def patchify_latents(latents): + """(B, C, H, W) -> (B, C*4, H//2, W//2).""" + b, c, h, w = latents.shape + latents = latents.view(b, c, h // 2, 2, w // 2, 2).permute(0, 1, 3, 5, 2, 4) + return latents.reshape(b, c * 4, h // 2, w // 2) + + +def pack_latents(latents): + """(B, C, H, W) -> (B, H*W, C).""" + b, c, h, w = latents.shape + return latents.reshape(b, c, h * w).permute(0, 2, 1) + + +def prepare_latent_ids(height, width, device, image_rope_scale=IMAGE_ROPE_SCALE): + """(H*W, 4) position ids (t, h, w, l); H/W axes scaled by the FLUX.3 image RoPE scale.""" + ids = torch.cartesian_prod(torch.arange(1), torch.arange(height), torch.arange(width), torch.arange(1)) + ids = ids.to(dtype=torch.float32) + ids[:, 1:3] = ids[:, 1:3] * image_rope_scale + return ids.to(device) + + +def prepare_text_ids(seq_len, device): + """(seq_len, 4) position ids for text tokens; text advances along the L axis.""" + ids = torch.cartesian_prod(torch.arange(1), torch.arange(1), torch.arange(1), torch.arange(seq_len)) + return ids.to(dtype=torch.float32).to(device) + + +@torch.no_grad() +def get_qwen3vl_prompt_embeds( + text_encoder, tokenizer, prompt, device, dtype, text_encoder_out_layers, max_sequence_length, pad_to_length=None +): + """Stack the FLUX.3 Qwen3-VL hidden layers -> (B, seq, len(layers) * hidden_dim). + + Mirrors `Flux3TextEncoderStep._get_qwen3_vl_prompt_embeds`: each prompt goes through the chat template and is + right-padded to a multiple of `TEXT_PAD_MULTIPLE` (capped at `max_sequence_length`). Qwen3-VL masks the padding + (so a prompt's real-token embeddings are independent of the window size), but the transformer attends to it, so + the window is part of the recipe -- for a single prompt this reproduces the pipeline window exactly. When a batch + of differing prompts is encoded together, they share the batch's largest window so the embeddings can be stacked; + with `train_batch_size=1` or a shared instance prompt every prompt keeps its own pipeline-exact window. + The position ids of the padding tokens are passed explicitly to match the Qwen3-VL 4.57 training environment. + """ + prompt = [prompt] if isinstance(prompt, str) else prompt + texts = [ + tokenizer.apply_chat_template([{"role": "user", "content": p}], tokenize=False, add_generation_prompt=True) + for p in prompt + ] + real_lengths = [ + tokenizer.tokenizer(t, return_tensors="pt", truncation=True, max_length=max_sequence_length)[ + "input_ids" + ].shape[1] + for t in texts + ] + padded_length = min(math.ceil(max(real_lengths) / TEXT_PAD_MULTIPLE) * TEXT_PAD_MULTIPLE, max_sequence_length) + if pad_to_length is not None: + # Force a specific window (caption dropout encodes the empty prompt at a dropped caption's window). Never + # shrink below the prompt's own real length, never exceed max_sequence_length. + padded_length = min(max(padded_length, pad_to_length), max_sequence_length) + enc = tokenizer.tokenizer( + texts, + return_tensors="pt", + padding="max_length", + padding_side="right", + truncation=True, + max_length=padded_length, + ) + input_ids, attn = enc["input_ids"].to(device), enc["attention_mask"].to(device) + position_ids = torch.arange(padded_length, device=device)[None].masked_fill(attn == 0, 1) + out = text_encoder( + input_ids=input_ids, + attention_mask=attn, + position_ids=position_ids, + output_hidden_states=True, + use_cache=False, + ) + stacked = torch.stack([out.hidden_states[k] for k in text_encoder_out_layers], dim=1).to( + dtype=dtype, device=device + ) + b, n, s, d = stacked.shape + return stacked.permute(0, 2, 1, 3).reshape(b, s, n * d) + + +def save_model_card( + repo_id: str, + images=None, + base_model: str = None, + instance_prompt=None, + validation_prompt=None, + repo_folder=None, + quant_training=None, +): + widget_dict = [] + if images is not None: + for i, image in enumerate(images): + image.save(os.path.join(repo_folder, f"image_{i}.png")) + widget_dict.append( + {"text": validation_prompt if validation_prompt else " ", "output": {"url": f"image_{i}.png"}} + ) + + model_description = f""" +# FLUX.3 DreamBooth LoRA - {repo_id} + + + +## Model description + +These are {repo_id} DreamBooth LoRA weights for {base_model}. + +The weights were trained using [DreamBooth](https://dreambooth.github.io/) with the [FLUX.3 diffusers trainer](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_flux3_klein.md). + +Quant training? {quant_training} + +## Trigger words + +You should use `{instance_prompt}` to trigger the image generation. + +## Download model + +[Download the *.safetensors LoRA]({repo_id}/tree/main) in the Files & versions tab. + +## Use it with the [🧨 diffusers library](https://github.com/huggingface/diffusers) + +```py +from diffusers import ModularPipeline +import torch +pipeline = ModularPipeline.from_pretrained("{base_model}") +pipeline.load_components(torch_dtype=torch.bfloat16) +pipeline.to('cuda') +pipeline.load_lora_weights('{repo_id}', weight_name='pytorch_lora_weights.safetensors') +image = pipeline(prompt='{validation_prompt if validation_prompt else instance_prompt}', output_type="pil").values["images"][0] +``` + +For more details, including weighting, merging and fusing LoRAs, check the [documentation on loading LoRAs in diffusers](https://huggingface.co/docs/diffusers/main/en/using-diffusers/loading_adapters) + +## License + +Please adhere to the licensing terms as described in the model repository. +""" + model_card = load_or_create_model_card( + repo_id_or_path=repo_id, + from_training=True, + license="other", + base_model=base_model, + prompt=instance_prompt, + model_description=model_description, + widget=widget_dict, + ) + tags = [ + "text-to-image", + "diffusers-training", + "diffusers", + "lora", + "flux3", + "flux3-diffusers", + "template:sd-lora", + ] + + model_card = populate_model_card(model_card, tags=tags) + model_card.save(os.path.join(repo_folder, "README.md")) + + +def log_validation( + pipeline, + args, + accelerator, + pipeline_args, + epoch, + torch_dtype, + is_final_validation=False, +): + args.num_validation_images = args.num_validation_images if args.num_validation_images else 1 + logger.info( + f"Running validation... \n Generating {args.num_validation_images} images with prompt:" + f" {args.validation_prompt}." + ) + + # run inference + generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed is not None else None + autocast_ctx = torch.autocast(accelerator.device.type) if not is_final_validation else nullcontext() + + images = [] + for _ in range(args.num_validation_images): + with autocast_ctx: + output = pipeline(generator=generator, output_type="pil", **pipeline_args) + images.append(output.values["images"][0]) + + for tracker in accelerator.trackers: + phase_name = "test" if is_final_validation else "validation" + if tracker.name == "tensorboard": + np_images = np.stack([np.asarray(img) for img in images]) + tracker.writer.add_images(phase_name, np_images, epoch, dataformats="NHWC") + if tracker.name == "wandb": + tracker.log( + { + phase_name: [ + wandb.Image(image, caption=f"{i}: {args.validation_prompt}") for i, image in enumerate(images) + ] + } + ) + + del pipeline + free_memory() + + return images + + +def make_validation_pipeline(args, transformer, weight_dtype, device): + # Build a modular pipeline for validation that REUSES the in-memory trained transformer, so validation runs + # on exactly the weights being trained and the DiT is never loaded a second time. Only the small auxiliary + # components are (re)loaded (they are freed after latent caching). + vae = AutoencoderKLFlux2.from_pretrained( + args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision, torch_dtype=torch.float32 + ) + text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( + args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision, torch_dtype=weight_dtype + ) + tokenizer = AutoProcessor.from_pretrained( + args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision + ) + scheduler = FlowMatchEulerDiscreteScheduler(shift=FLOW_SHIFT, use_dynamic_shifting=False) + pipeline = Flux3AutoBlocks().init_pipeline() + pipeline.update_components( + transformer=transformer, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, scheduler=scheduler + ) + pipeline.to(device) + return pipeline + + +def module_filter_fn(mod: torch.nn.Module, fqn: str): + # don't convert the output module (FLUX.3's output projection) + if fqn.endswith("proj_out"): + return False + # don't convert linear modules with weight dimensions not divisible by 16 + if isinstance(mod, torch.nn.Linear): + if mod.in_features % 16 != 0 or mod.out_features % 16 != 0: + return False + return True + + +def parse_args(input_args=None): + parser = argparse.ArgumentParser(description="Simple example of a training script.") + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + required=True, + help="Path to pretrained model or model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--revision", + type=str, + default=None, + required=False, + help="Revision of pretrained model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--bnb_quantization_config_path", + type=str, + default=None, + help="Quantization config in a JSON file that will be used to define the bitsandbytes quant config of the DiT.", + ) + parser.add_argument( + "--do_fp8_training", + action="store_true", + help="if we are doing FP8 training.", + ) + parser.add_argument( + "--variant", + type=str, + default=None, + help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16", + ) + parser.add_argument( + "--dataset_name", + type=str, + default=None, + help=( + "The name of the Dataset (from the HuggingFace hub) containing the training data of instance images (could be your own, possibly private," + " dataset). It can also be a path pointing to a local copy of a dataset in your filesystem," + " or to a folder containing files that 🤗 Datasets can understand." + ), + ) + parser.add_argument( + "--dataset_config_name", + type=str, + default=None, + help="The config of the Dataset, leave as None if there's only one config.", + ) + parser.add_argument( + "--instance_data_dir", + type=str, + default=None, + help=("A folder containing the training data. "), + ) + + parser.add_argument( + "--cache_dir", + type=str, + default=None, + help="The directory where the downloaded models and datasets will be stored.", + ) + + parser.add_argument( + "--image_column", + type=str, + default="image", + help="The column of the dataset containing the target image. By " + "default, the standard Image Dataset maps out 'file_name' " + "to 'image'.", + ) + parser.add_argument( + "--caption_column", + type=str, + default="text", + help="The column of the dataset containing the instance prompt for each image", + ) + + parser.add_argument("--repeats", type=int, default=1, help="How many times to repeat the training data.") + + parser.add_argument( + "--class_data_dir", + type=str, + default=None, + required=False, + help="A folder containing the training data of class images.", + ) + parser.add_argument( + "--instance_prompt", + type=str, + default=None, + required=True, + help="The prompt with identifier specifying the instance, e.g. 'photo of a TOK dog', 'in the style of TOK'", + ) + parser.add_argument( + "--class_prompt", + type=str, + default=None, + help="The prompt to specify images in the same class as provided instance images.", + ) + parser.add_argument( + "--max_sequence_length", + type=int, + default=512, + help=( + "Cap on the number of Qwen3-VL prompt tokens; longer prompts are truncated. The text window is the " + f"prompt's length rounded up to a multiple of {TEXT_PAD_MULTIPLE}, capped at this value." + ), + ) + parser.add_argument( + "--text_encoder_out_layers", + type=int, + nargs="+", + default=[4, 8, 12, 16, 20, 24, 28, 32], + help="Text encoder hidden layers to compute the final text embeddings.", + ) + parser.add_argument( + "--validation_prompt", + type=str, + default=None, + help="A prompt that is used during validation to verify that the model is learning.", + ) + parser.add_argument( + "--skip_final_inference", + default=False, + action="store_true", + help="Whether to skip the final inference step with loaded lora weights upon training completion. This will run intermediate validation inference if `validation_prompt` is provided. Specify to reduce memory.", + ) + parser.add_argument( + "--final_validation_prompt", + type=str, + default=None, + help="A prompt that is used during a final validation to verify that the model is learning. Ignored if `--validation_prompt` is provided.", + ) + parser.add_argument( + "--num_validation_images", + type=int, + default=4, + help="Number of images that should be generated during validation with `validation_prompt`.", + ) + parser.add_argument( + "--validation_epochs", + type=int, + default=50, + help=( + "Run dreambooth validation every X epochs. Dreambooth validation consists of running the prompt" + " `args.validation_prompt` multiple times: `args.num_validation_images`." + ), + ) + parser.add_argument( + "--rank", + type=int, + default=4, + help=("The dimension of the LoRA update matrices."), + ) + parser.add_argument( + "--lora_alpha", + type=int, + default=4, + help="LoRA alpha to be used for additional scaling.", + ) + parser.add_argument("--lora_dropout", type=float, default=0.0, help="Dropout probability for LoRA layers") + + parser.add_argument( + "--with_prior_preservation", + default=False, + action="store_true", + help="Flag to add prior preservation loss.", + ) + parser.add_argument("--prior_loss_weight", type=float, default=1.0, help="The weight of prior preservation loss.") + parser.add_argument( + "--num_class_images", + type=int, + default=100, + help=( + "Minimal class images for prior preservation loss. If there are not enough images already present in" + " class_data_dir, additional images will be sampled with class_prompt." + ), + ) + parser.add_argument( + "--output_dir", + type=str, + default="flux-dreambooth-lora", + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + parser.add_argument( + "--resolution", + type=int, + default=512, + help=( + "The resolution for input images, all the images in the train/validation dataset will be resized to this" + " resolution" + ), + ) + parser.add_argument( + "--aspect_ratio_buckets", + type=str, + default=None, + help=( + "Aspect ratio buckets to use for training. Define as a string of 'h1,w1;h2,w2;...'. " + "e.g. '1024,1024;768,1360;1360,768;880,1168;1168,880;1248,832;832,1248'. " + "Requires --use_aspect_ratio_buckets. Images are resized to cover and cropped to the nearest " + "listed bucket (smaller images are upscaled). When set, --resolution is ignored." + ), + ) + parser.add_argument( + "--use_aspect_ratio_buckets", + action="store_true", + help=( + "Enable aspect-ratio bucketing. Without --aspect_ratio_buckets, the buckets are computed on the " + "fly from --resolution and capped to each image's own resolution, so smaller images are assigned " + "to a smaller bucket instead of being upscaled. Provide --aspect_ratio_buckets to use an explicit list." + ), + ) + parser.add_argument( + "--center_crop", + default=False, + action="store_true", + help=( + "Whether to center crop the input images to the resolution. If not set, the images will be randomly" + " cropped. The images will be resized to the resolution first before cropping." + ), + ) + parser.add_argument( + "--random_flip", + action="store_true", + help="whether to randomly flip images horizontally", + ) + parser.add_argument( + "--caption_dropout", + type=float, + default=0.0, + help=( + "Probability of replacing an instance image's caption with an empty string during training, so that" + " fraction of samples is trained unconditionally. Improves classifier-free guidance. A common value is" + " 0.1. Class/prior-preservation captions are never dropped." + ), + ) + parser.add_argument( + "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader." + ) + parser.add_argument( + "--sample_batch_size", type=int, default=4, help="Batch size (per device) for sampling images." + ) + parser.add_argument("--num_train_epochs", type=int, default=1) + parser.add_argument( + "--max_train_steps", + type=int, + default=None, + help="Total number of training steps to perform. If provided, overrides num_train_epochs.", + ) + parser.add_argument( + "--checkpointing_steps", + type=int, + default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" + " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument( + "--checkpoints_total_limit", + type=int, + default=None, + help=("Max number of checkpoints to store."), + ) + parser.add_argument( + "--resume_from_checkpoint", + type=str, + default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument( + "--gradient_accumulation_steps", + type=int, + default=1, + help="Number of updates steps to accumulate before performing a backward/update pass.", + ) + parser.add_argument( + "--gradient_checkpointing", + action="store_true", + help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + + parser.add_argument( + "--scale_lr", + action="store_true", + default=False, + help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", + ) + parser.add_argument( + "--lr_scheduler", + type=str, + default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument( + "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." + ) + parser.add_argument( + "--lr_num_cycles", + type=int, + default=1, + help="Number of hard resets of the lr in cosine_with_restarts scheduler.", + ) + parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.") + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=0, + help=( + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ), + ) + parser.add_argument( + "--weighting_scheme", + type=str, + default="none", + choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "none"], + help=('We default to the "none" weighting scheme for uniform sampling and uniform loss'), + ) + parser.add_argument( + "--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--mode_scale", + type=float, + default=1.29, + help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.", + ) + parser.add_argument( + "--optimizer", + type=str, + default="AdamW", + help=('The optimizer type to use. Choose between ["AdamW", "prodigy"]'), + ) + + parser.add_argument( + "--use_8bit_adam", + action="store_true", + help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW", + ) + + parser.add_argument( + "--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers." + ) + parser.add_argument( + "--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers." + ) + parser.add_argument( + "--prodigy_beta3", + type=float, + default=None, + help="coefficients for computing the Prodigy stepsize using running averages. If set to None, " + "uses the value of square root of beta2. Ignored if optimizer is adamW", + ) + parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay") + parser.add_argument("--adam_weight_decay", type=float, default=1e-04, help="Weight decay to use for unet params") + + parser.add_argument( + "--lora_layers", + type=str, + default=None, + help=( + 'The transformer modules to apply LoRA training on. Please specify the layers in a comma separated. E.g. - "to_k,to_q,to_v,to_out.0" will result in lora training of attention layers only' + ), + ) + + parser.add_argument( + "--adam_epsilon", + type=float, + default=1e-08, + help="Epsilon value for the Adam optimizer and Prodigy optimizers.", + ) + + parser.add_argument( + "--prodigy_use_bias_correction", + type=bool, + default=True, + help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW", + ) + parser.add_argument( + "--prodigy_safeguard_warmup", + type=bool, + default=True, + help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. " + "Ignored if optimizer is adamW", + ) + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") + parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") + parser.add_argument( + "--hub_model_id", + type=str, + default=None, + help="The name of the repository to keep in sync with the local `output_dir`.", + ) + parser.add_argument( + "--logging_dir", + type=str, + default="logs", + help=( + "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." + ), + ) + parser.add_argument( + "--allow_tf32", + action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument( + "--cache_latents", + action="store_true", + default=False, + help="Cache the VAE latents", + ) + parser.add_argument( + "--report_to", + type=str, + default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + parser.add_argument( + "--mixed_precision", + type=str, + default=None, + choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + parser.add_argument( + "--upcast_before_saving", + action="store_true", + default=False, + help=( + "Whether to upcast the trained transformer layers to float32 before saving (at the end of training). " + "Defaults to precision dtype used for training to save memory" + ), + ) + parser.add_argument( + "--offload", + action="store_true", + help="Whether to offload the VAE and the text encoder to CPU when they are not used.", + ) + parser.add_argument( + "--prior_generation_precision", + type=str, + default=None, + choices=["no", "fp32", "fp16", "bf16"], + help=( + "Choose prior generation precision between fp32, fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to fp16 if a GPU is available else fp32." + ), + ) + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument("--enable_npu_flash_attention", action="store_true", help="Enabla Flash Attention for NPU") + parser.add_argument("--fsdp_text_encoder", action="store_true", help="Use FSDP for text encoder") + + if input_args is not None: + args = parser.parse_args(input_args) + else: + args = parser.parse_args() + + if args.dataset_name is None and args.instance_data_dir is None: + raise ValueError("Specify either `--dataset_name` or `--instance_data_dir`") + + if args.dataset_name is not None and args.instance_data_dir is not None: + raise ValueError("Specify only one of `--dataset_name` or `--instance_data_dir`") + if args.do_fp8_training and args.bnb_quantization_config_path: + raise ValueError("Both `do_fp8_training` and `bnb_quantization_config_path` cannot be passed.") + + env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) + if env_local_rank != -1 and env_local_rank != args.local_rank: + args.local_rank = env_local_rank + + if args.with_prior_preservation: + if args.class_data_dir is None: + raise ValueError("You must specify a data directory for class images.") + if args.class_prompt is None: + raise ValueError("You must specify prompt for class images.") + else: + # logger is not available yet + if args.class_data_dir is not None: + warnings.warn("You need not use --class_data_dir without --with_prior_preservation.") + if args.class_prompt is not None: + warnings.warn("You need not use --class_prompt without --with_prior_preservation.") + + return args + + +class DreamBoothDataset(Dataset): + """ + A dataset to prepare the instance and class images with the prompts for fine-tuning the model. + It pre-processes the images. + """ + + def __init__( + self, + instance_data_root, + instance_prompt, + class_prompt, + class_data_root=None, + class_num=None, + size=1024, + repeats=1, + center_crop=False, + buckets=None, + use_aspect_ratio_buckets=False, + bucket_divisibility=16, + bucket_base_resolutions=None, + ): + self.size = size + self.resolution = size + self.center_crop = center_crop + + self.instance_prompt = instance_prompt + self.custom_instance_prompts = None + self.class_prompt = class_prompt + + # Explicit user-provided bucket list (or None). The concrete list of buckets actually used is + # built from the data in `self.buckets` during preprocessing below. + self._explicit_buckets = buckets + self.use_aspect_ratio_buckets = use_aspect_ratio_buckets + self.bucket_divisibility = bucket_divisibility + self.bucket_base_resolutions = bucket_base_resolutions + + # if --dataset_name is provided or a metadata jsonl file is provided in the local --instance_data directory, + # we load the training data using load_dataset + if args.dataset_name is not None: + try: + from datasets import load_dataset + except ImportError: + raise ImportError( + "You are trying to load your data using the datasets library. If you wish to train using custom " + "captions please install the datasets library: `pip install datasets`. If you wish to load a " + "local folder containing images only, specify --instance_data_dir instead." + ) + # Downloading and loading a dataset from the hub. + # See more about loading custom images at + # https://huggingface.co/docs/datasets/v2.0.0/en/dataset_script + dataset = load_dataset( + args.dataset_name, + args.dataset_config_name, + cache_dir=args.cache_dir, + ) + # Preprocessing the datasets. + column_names = dataset["train"].column_names + + # 6. Get the column names for input/target. + if args.image_column is None: + image_column = column_names[0] + logger.info(f"image column defaulting to {image_column}") + else: + image_column = args.image_column + if image_column not in column_names: + raise ValueError( + f"`--image_column` value '{args.image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" + ) + instance_images = dataset["train"][image_column] + + if args.caption_column is None: + logger.info( + "No caption column provided, defaulting to instance_prompt for all images. If your dataset " + "contains captions/prompts for the images, make sure to specify the " + "column as --caption_column" + ) + self.custom_instance_prompts = None + else: + if args.caption_column not in column_names: + raise ValueError( + f"`--caption_column` value '{args.caption_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" + ) + custom_instance_prompts = dataset["train"][args.caption_column] + # create final list of captions according to --repeats + self.custom_instance_prompts = [] + for caption in custom_instance_prompts: + self.custom_instance_prompts.extend(itertools.repeat(caption, repeats)) + else: + self.instance_data_root = Path(instance_data_root) + if not self.instance_data_root.exists(): + raise ValueError("Instance images root doesn't exists.") + + instance_images = [Image.open(path) for path in list(Path(instance_data_root).iterdir())] + self.custom_instance_prompts = None + + self.instance_images = [] + for img in instance_images: + self.instance_images.extend(itertools.repeat(img, repeats)) + + self.pixel_values = [] + self.buckets = [] + bucket_to_idx = {} + for i, image in enumerate(self.instance_images): + image = exif_transpose(image) + if not image.mode == "RGB": + image = image.convert("RGB") + + width, height = image.size + + # Assign the image to a bucket. + target = self._bucket_for_image(height, width) + if target not in bucket_to_idx: + bucket_to_idx[target] = len(self.buckets) + self.buckets.append(target) + bucket_idx = bucket_to_idx[target] + + # based on the bucket assignment, define the transformations + image = self.train_transform( + image, + size=target, + center_crop=args.center_crop, + random_flip=args.random_flip, + ) + self.pixel_values.append((image, bucket_idx)) + + self.num_instance_images = len(self.instance_images) + self._length = self.num_instance_images + + if class_data_root is not None: + self.class_data_root = Path(class_data_root) + self.class_data_root.mkdir(parents=True, exist_ok=True) + self.class_images_path = list(self.class_data_root.iterdir()) + if class_num is not None: + self.num_class_images = min(len(self.class_images_path), class_num) + else: + self.num_class_images = len(self.class_images_path) + self._length = max(self.num_class_images, self.num_instance_images) + else: + self.class_data_root = None + + def __len__(self): + return self._length + + def __getitem__(self, index): + example = {} + instance_image, bucket_idx = self.pixel_values[index % self.num_instance_images] + example["index"] = index + example["instance_images"] = instance_image + example["bucket_idx"] = bucket_idx + if self.custom_instance_prompts: + caption = self.custom_instance_prompts[index % self.num_instance_images] + if caption: + example["instance_prompt"] = caption + else: + example["instance_prompt"] = self.instance_prompt + + else: # custom prompts were provided, but length does not match size of image dataset + example["instance_prompt"] = self.instance_prompt + + if self.class_data_root: + class_image = Image.open(self.class_images_path[index % self.num_class_images]) + class_image = exif_transpose(class_image) + + if not class_image.mode == "RGB": + class_image = class_image.convert("RGB") + # Match the class image to the paired instance image's bucket so they can be stacked into one batch. + example["class_images"] = self.train_transform( + class_image, size=self.buckets[bucket_idx], center_crop=self.center_crop + ) + example["class_prompt"] = self.class_prompt + + return example + + def _bucket_for_image(self, height, width): + # An explicit bucket list takes priority: pick the nearest, upscaling smaller images to cover it. + if self._explicit_buckets is not None: + return self._explicit_buckets[find_nearest_bucket(height, width, self._explicit_buckets)] + # On-the-fly bucketing: cap the ladder to the image's own resolution so smaller images are + # assigned to a smaller bucket rather than being upscaled (mirrors ostris' bucketing). + if self.use_aspect_ratio_buckets: + resolution = min(self.resolution, round((height * width) ** 0.5)) + ladder = generate_aspect_ratio_buckets( + resolution, + divisibility=self.bucket_divisibility, + base_resolutions=self.bucket_base_resolutions, + ) + return ladder[find_nearest_bucket(height, width, ladder)] + # No bucketing: a single square bucket reproduces the fixed-size resize + crop. + return (self.resolution, self.resolution) + + def train_transform(self, image, size, center_crop=False, random_flip=False): + # Resize preserving aspect ratio so the image covers the bucket, then crop to the bucket size. + target_height, target_width = size + width, height = image.size + scale = max(target_height / height, target_width / width) + new_height, new_width = round(height * scale), round(width * scale) + image = TF.resize(image, [new_height, new_width], interpolation=transforms.InterpolationMode.BILINEAR) + if center_crop: + image = TF.center_crop(image, size) + else: + i, j, h, w = transforms.RandomCrop.get_params(image, output_size=size) + image = TF.crop(image, i, j, h, w) + if random_flip and random.random() < 0.5: + image = TF.hflip(image) + return TF.normalize(TF.to_tensor(image), [0.5], [0.5]) + + +def collate_fn(examples, with_prior_preservation=False): + indices = [example["index"] for example in examples] + pixel_values = [example["instance_images"] for example in examples] + # Keep instance_prompts unchanged for prompt cache precompute; prompts may be extended with class prompts below. + instance_prompts = [example["instance_prompt"] for example in examples] + prompts = [example["instance_prompt"] for example in examples] + + # Concat class and instance examples for prior preservation. + # We do this to avoid doing two forward passes. + if with_prior_preservation: + pixel_values += [example["class_images"] for example in examples] + prompts += [example["class_prompt"] for example in examples] + + pixel_values = torch.stack(pixel_values) + pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float() + + batch = { + "indices": indices, + "pixel_values": pixel_values, + "instance_prompts": instance_prompts, + "prompts": prompts, + } + return batch + + +class BucketBatchSampler(BatchSampler): + def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False, seed: int = None): + if not isinstance(batch_size, int) or batch_size <= 0: + raise ValueError("batch_size should be a positive integer value, but got batch_size={}".format(batch_size)) + if not isinstance(drop_last, bool): + raise ValueError("drop_last should be a boolean value, but got drop_last={}".format(drop_last)) + + self.dataset = dataset + self.batch_size = batch_size + self.drop_last = drop_last + self.generator = random.Random(seed) if seed is not None else random + + # Group indices by bucket + self.bucket_indices = [[] for _ in range(len(self.dataset.buckets))] + for idx, (_, bucket_idx) in enumerate(self.dataset.pixel_values): + self.bucket_indices[bucket_idx].append(idx) + + self.sampler_len = 0 + for indices_in_bucket in self.bucket_indices: + num_batches, remainder = divmod(len(indices_in_bucket), self.batch_size) + self.sampler_len += num_batches + if remainder > 0 and not self.drop_last: + self.sampler_len += 1 + + def __iter__(self): + batches = [] + for indices_in_bucket in self.bucket_indices: + shuffled_indices = indices_in_bucket.copy() + self.generator.shuffle(shuffled_indices) + for i in range(0, len(shuffled_indices), self.batch_size): + batch = shuffled_indices[i : i + self.batch_size] + if len(batch) < self.batch_size and self.drop_last: + continue + batches.append(batch) + + self.generator.shuffle(batches) + for batch in batches: + yield batch + + def __len__(self): + return self.sampler_len + + +class PromptDataset(Dataset): + "A simple dataset to prepare the prompts to generate class images on multiple GPUs." + + def __init__(self, prompt, num_samples): + self.prompt = prompt + self.num_samples = num_samples + + def __len__(self): + return self.num_samples + + def __getitem__(self, index): + example = {} + example["prompt"] = self.prompt + example["index"] = index + return example + + +def main(args): + if args.report_to == "wandb" and args.hub_token is not None: + raise ValueError( + "You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token." + " Please use `hf auth login` to authenticate with the Hub." + ) + + if torch.backends.mps.is_available() and args.mixed_precision == "bf16": + # due to pytorch#99272, MPS does not yet support bfloat16. + raise ValueError( + "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." + ) + if args.do_fp8_training: + from torchao.float8 import Float8LinearConfig, convert_to_float8_training + + logging_dir = Path(args.output_dir, args.logging_dir) + + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + kwargs_handlers=[kwargs], + ) + + # Disable AMP for MPS. + if torch.backends.mps.is_available(): + accelerator.native_amp = False + + if args.report_to == "wandb": + if not is_wandb_available(): + raise ImportError("Make sure to install wandb if you want to use it for logging during training.") + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed) + + # Generate class images if prior preservation is enabled. + if args.with_prior_preservation: + class_images_dir = Path(args.class_data_dir) + if not class_images_dir.exists(): + class_images_dir.mkdir(parents=True) + cur_class_images = len(list(class_images_dir.iterdir())) + + if cur_class_images < args.num_class_images: + has_supported_fp16_accelerator = torch.cuda.is_available() or torch.backends.mps.is_available() + torch_dtype = torch.float16 if has_supported_fp16_accelerator else torch.float32 + if args.prior_generation_precision == "fp32": + torch_dtype = torch.float32 + elif args.prior_generation_precision == "fp16": + torch_dtype = torch.float16 + elif args.prior_generation_precision == "bf16": + torch_dtype = torch.bfloat16 + + pipeline = ModularPipeline.from_pretrained(args.pretrained_model_name_or_path) + pipeline.load_components(torch_dtype=torch_dtype) + + num_new_images = args.num_class_images - cur_class_images + logger.info(f"Number of class images to sample: {num_new_images}.") + + sample_dataset = PromptDataset(args.class_prompt, num_new_images) + sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size) + + sample_dataloader = accelerator.prepare(sample_dataloader) + pipeline.to(accelerator.device) + + for example in tqdm( + sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process + ): + with torch.autocast(device_type=accelerator.device.type, dtype=torch_dtype): + images = pipeline(prompt=example["prompt"], output_type="pil").values["images"] + + for i, image in enumerate(images): + hash_image = insecure_hashlib.sha1(image.tobytes()).hexdigest() + image_filename = class_images_dir / f"{example['index'][i] + cur_class_images}-{hash_image}.jpg" + image.save(image_filename) + + del pipeline + free_memory() + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + if args.push_to_hub: + repo_id = create_repo( + repo_id=args.hub_model_id or Path(args.output_dir).name, + exist_ok=True, + ).repo_id + + # Load the tokenizers + tokenizer = AutoProcessor.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="tokenizer", + revision=args.revision, + ) + + # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + + # Load scheduler and models + noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=FLOW_SHIFT, use_dynamic_shifting=False) + noise_scheduler_copy = copy.deepcopy(noise_scheduler) + vae = AutoencoderKLFlux2.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="vae", + revision=args.revision, + variant=args.variant, + ) + latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(accelerator.device) + latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to( + accelerator.device + ) + + quantization_config = None + if args.bnb_quantization_config_path is not None: + with open(args.bnb_quantization_config_path, "r") as f: + config_kwargs = json.load(f) + if "load_in_4bit" in config_kwargs and config_kwargs["load_in_4bit"]: + config_kwargs["bnb_4bit_compute_dtype"] = weight_dtype + quantization_config = BitsAndBytesConfig(**config_kwargs) + + transformer = Flux3Transformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + revision=args.revision, + variant=args.variant, + quantization_config=quantization_config, + torch_dtype=weight_dtype, + ) + if args.bnb_quantization_config_path is not None: + transformer = prepare_model_for_kbit_training(transformer, use_gradient_checkpointing=False) + + text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="text_encoder", + revision=args.revision, + variant=args.variant, + torch_dtype=weight_dtype, + ) + text_encoder.requires_grad_(False) + + # We only train the additional adapter LoRA layers + transformer.requires_grad_(False) + vae.requires_grad_(False) + + if args.enable_npu_flash_attention: + if is_torch_npu_available(): + logger.info("npu flash attention enabled.") + transformer.set_attention_backend("_native_npu") + else: + raise ValueError("npu flash attention requires torch_npu extensions and is supported only on npu device ") + + if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16: + # due to pytorch#99272, MPS does not yet support bfloat16. + raise ValueError( + "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." + ) + + to_kwargs = {"dtype": weight_dtype, "device": accelerator.device} if not args.offload else {"dtype": weight_dtype} + # flux vae is stable in bf16 so load it in weight_dtype to reduce memory + vae.to(**to_kwargs) + # we never offload the transformer to CPU, so we can just use the accelerator device + transformer_to_kwargs = ( + {"device": accelerator.device} + if args.bnb_quantization_config_path is not None + else {"device": accelerator.device, "dtype": weight_dtype} + ) + + is_fsdp = getattr(accelerator.state, "fsdp_plugin", None) is not None + if not is_fsdp: + transformer.to(**transformer_to_kwargs) + + if args.do_fp8_training: + convert_to_float8_training( + transformer, module_filter_fn=module_filter_fn, config=Float8LinearConfig(pad_inner_dim=True) + ) + + text_encoder.to(**to_kwargs) + + if args.gradient_checkpointing: + transformer.enable_gradient_checkpointing() + + if args.lora_layers is not None: + target_modules = [layer.strip() for layer in args.lora_layers.split(",")] + else: + # Attention + feed-forward projections across every block (mode blocks and single blocks). These are exactly + # the leaf modules the non-diffusers (ostris) FLUX.3 LoRA converter maps, so adapters trained here stay + # interchangeable with that layout. Pass `--lora_layers` to target a different set. + target_modules = [ + "to_q", + "to_k", + "to_v", + "to_out.0", + "ff.linear_in", + "ff.linear_out", + ] + + # now we will add new LoRA weights the transformer layers + transformer_lora_config = LoraConfig( + r=args.rank, + lora_alpha=args.lora_alpha, + lora_dropout=args.lora_dropout, + init_lora_weights="gaussian", + target_modules=target_modules, + ) + transformer.add_adapter(transformer_lora_config) + + def unwrap_model(model): + model = accelerator.unwrap_model(model) + model = model._orig_mod if is_compiled_module(model) else model + return model + + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + transformer_cls = type(unwrap_model(transformer)) + + # 1) Validate and pick the transformer model + modules_to_save: dict[str, Any] = {} + transformer_model = None + + for model in models: + if isinstance(unwrap_model(model), transformer_cls): + transformer_model = model + modules_to_save["transformer"] = model + else: + raise ValueError(f"unexpected save model: {model.__class__}") + + if transformer_model is None: + raise ValueError("No transformer model found in 'models'") + + # 2) Optionally gather FSDP state dict once + state_dict = accelerator.get_state_dict(model) if is_fsdp else None + + # 3) Only main process materializes the LoRA state dict + transformer_lora_layers_to_save = None + if accelerator.is_main_process: + peft_kwargs = {} + if is_fsdp: + peft_kwargs["state_dict"] = state_dict + + transformer_lora_layers_to_save = get_peft_model_state_dict( + unwrap_model(transformer_model) if is_fsdp else transformer_model, + **peft_kwargs, + ) + + if is_fsdp: + transformer_lora_layers_to_save = _to_cpu_contiguous(transformer_lora_layers_to_save) + + # make sure to pop weight so that corresponding model is not saved again + if weights: + weights.pop() + + Flux3LoraLoaderMixin.save_lora_weights( + output_dir, + transformer_lora_layers=transformer_lora_layers_to_save, + **_collate_lora_metadata(modules_to_save), + ) + + def load_model_hook(models, input_dir): + transformer_ = None + + if not is_fsdp: + while len(models) > 0: + model = models.pop() + + if isinstance(unwrap_model(model), type(unwrap_model(transformer))): + transformer_ = unwrap_model(model) + else: + raise ValueError(f"unexpected save model: {model.__class__}") + else: + transformer_ = Flux3Transformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + ) + transformer_.add_adapter(transformer_lora_config) + + lora_state_dict = Flux3LoraLoaderMixin.lora_state_dict(input_dir) + + transformer_state_dict = { + f"{k.replace('transformer.', '')}": v for k, v in lora_state_dict.items() if k.startswith("transformer.") + } + transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict) + incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default") + if incompatible_keys is not None: + # check only for unexpected keys + unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) + if unexpected_keys: + logger.warning( + f"Loading adapter weights from state_dict led to unexpected keys not found in the model: " + f" {unexpected_keys}. " + ) + + # Make sure the trainable params are in float32. This is again needed since the base models + # are in `weight_dtype`. More details: + # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804 + if args.mixed_precision == "fp16": + models = [transformer_] + # only upcast trainable parameters (LoRA) into fp32 + cast_training_params(models) + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32 and torch.cuda.is_available(): + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + # Make sure the trainable params are in float32. + if args.mixed_precision == "fp16": + models = [transformer] + # only upcast trainable parameters (LoRA) into fp32 + cast_training_params(models, dtype=torch.float32) + + transformer_lora_parameters = list(filter(lambda p: p.requires_grad, transformer.parameters())) + + # Optimization parameters + transformer_parameters_with_lr = {"params": transformer_lora_parameters, "lr": args.learning_rate} + params_to_optimize = [transformer_parameters_with_lr] + + # Optimizer creation + if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"): + logger.warning( + f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]." + "Defaulting to adamW" + ) + args.optimizer = "adamw" + + if args.use_8bit_adam and not args.optimizer.lower() == "adamw": + logger.warning( + f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was " + f"set to {args.optimizer.lower()}" + ) + + if args.optimizer.lower() == "adamw": + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." + ) + + optimizer_class = bnb.optim.AdamW8bit + else: + optimizer_class = torch.optim.AdamW + + optimizer = optimizer_class( + params_to_optimize, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + if args.optimizer.lower() == "prodigy": + try: + import prodigyopt + except ImportError: + raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`") + + optimizer_class = prodigyopt.Prodigy + + if args.learning_rate <= 0.1: + logger.warning( + "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0" + ) + + optimizer = optimizer_class( + params_to_optimize, + betas=(args.adam_beta1, args.adam_beta2), + beta3=args.prodigy_beta3, + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + decouple=args.prodigy_decouple, + use_bias_correction=args.prodigy_use_bias_correction, + safeguard_warmup=args.prodigy_safeguard_warmup, + ) + + # Resolve the bucketing mode. Bucketing must be enabled explicitly with --use_aspect_ratio_buckets; + # a bucket list without that flag is an error. With the flag, an explicit --aspect_ratio_buckets list + # drives assignment, otherwise buckets are computed on the fly inside the dataset. Without the flag a + # single square bucket reproduces the fixed-size resize + crop. + if args.aspect_ratio_buckets is not None and not args.use_aspect_ratio_buckets: + raise ValueError("--aspect_ratio_buckets requires --use_aspect_ratio_buckets to be set.") + if args.aspect_ratio_buckets is not None: + buckets = parse_buckets_string(args.aspect_ratio_buckets) + use_aspect_ratio_buckets = False + logger.info(f"Using explicit aspect ratio buckets: {buckets}") + elif args.use_aspect_ratio_buckets: + buckets = None + use_aspect_ratio_buckets = True + logger.info( + "No --aspect_ratio_buckets provided; auto-computing aspect ratio buckets on the fly from --resolution." + ) + else: + buckets = [(args.resolution, args.resolution)] + use_aspect_ratio_buckets = False + + # Dataset and DataLoaders creation: + train_dataset = DreamBoothDataset( + instance_data_root=args.instance_data_dir, + instance_prompt=args.instance_prompt, + class_prompt=args.class_prompt, + class_data_root=args.class_data_dir if args.with_prior_preservation else None, + class_num=args.num_class_images, + size=args.resolution, + repeats=args.repeats, + center_crop=args.center_crop, + buckets=buckets, + use_aspect_ratio_buckets=use_aspect_ratio_buckets, + ) + precompute_latents = args.cache_latents or train_dataset.custom_instance_prompts + batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=True, seed=args.seed) + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation), + num_workers=args.dataloader_num_workers, + ) + + def compute_text_embeddings(prompt, text_encoder, tokenizer, pad_to_length=None): + # FLUX.3 has no classic text-encoding pipeline; encode directly with the frozen Qwen3-VL encoder. + with torch.no_grad(): + prompt_embeds = get_qwen3vl_prompt_embeds( + text_encoder, + tokenizer, + prompt, + accelerator.device, + weight_dtype, + text_encoder_out_layers=args.text_encoder_out_layers, + max_sequence_length=args.max_sequence_length, + pad_to_length=pad_to_length, + ) + # FLUX.3 text ids are positional only (identical for every sample of a given length). We carry them + # batched here so the [instance, class] concat / repeat / caption-dropout plumbing below is unchanged; + # the transformer consumes a single (seq, 4) slice at the forward call. + text_ids = prepare_text_ids(prompt_embeds.shape[1], accelerator.device) + text_ids = text_ids.unsqueeze(0).expand(prompt_embeds.shape[0], -1, -1) + return prompt_embeds, text_ids + + # Prompt embeddings get stacked across a batch (train_batch_size > 1) and concatenated with the class prompt + # (prior preservation). FLUX.3 pads each prompt to its own window, so stacking requires a shared window. Compute + # one dataset-wide window here and pad all custom captions (and the class prompt) to it. bs=1 without prior + # preservation never stacks, so it keeps each prompt's exact per-prompt window. + uniform_window = None + if train_dataset.custom_instance_prompts and (args.train_batch_size > 1 or args.with_prior_preservation): + windowed_prompts = list(train_dataset.custom_instance_prompts) + if args.with_prior_preservation: + windowed_prompts.append(args.class_prompt) + real_lengths = [ + len( + tokenizer.tokenizer( + tokenizer.apply_chat_template( + [{"role": "user", "content": p}], tokenize=False, add_generation_prompt=True + ), + truncation=True, + max_length=args.max_sequence_length, + )["input_ids"] + ) + for p in windowed_prompts + ] + uniform_window = min( + math.ceil(max(real_lengths) / TEXT_PAD_MULTIPLE) * TEXT_PAD_MULTIPLE, args.max_sequence_length + ) + + # If no type of tuning is done on the text_encoder and custom instance prompts are NOT + # provided (i.e. the --instance_prompt is used for all images), we encode the instance prompt once to avoid + # the redundant encoding. + if not train_dataset.custom_instance_prompts: + with offload_models(text_encoder, device=accelerator.device, offload=args.offload): + instance_prompt_hidden_states, instance_text_ids = compute_text_embeddings( + args.instance_prompt, text_encoder, tokenizer + ) + + # Handle class prompt for prior-preservation. + if args.with_prior_preservation: + with offload_models(text_encoder, device=accelerator.device, offload=args.offload): + class_prompt_hidden_states, class_text_ids = compute_text_embeddings( + args.class_prompt, text_encoder, tokenizer, pad_to_length=uniform_window + ) + + # FLUX.3 validation runs through the modular pipeline (which encodes prompts itself), so we only assemble + # the generation kwargs here rather than precomputing text embeddings. + validation_pipeline_args = {} + if args.validation_prompt is not None: + validation_pipeline_args = { + "prompt": args.validation_prompt, + "negative_prompt": "", + "height": args.resolution, + "width": args.resolution, + "num_inference_steps": 28, + # Encode the validation prompt exactly as training did (same window cap and hidden layers), so the + # modular pipeline's text embeddings match the ones the adapter was trained against. + "max_sequence_length": args.max_sequence_length, + "text_encoder_out_layers": tuple(args.text_encoder_out_layers), + } + + # Init FSDP for text encoder + if args.fsdp_text_encoder: + fsdp_kwargs = get_fsdp_kwargs_from_accelerator(accelerator) + text_encoder = wrap_with_fsdp( + model=text_encoder, + device=accelerator.device, + offload=args.offload, + limit_all_gathers=True, + use_orig_params=True, + fsdp_kwargs=fsdp_kwargs, + ) + dist.barrier() + + # If custom instance prompts are NOT provided (i.e. the instance prompt is used for all images), + # pack the statically computed variables appropriately here. This is so that we don't + # have to pass them to the dataloader. + if not train_dataset.custom_instance_prompts: + prompt_embeds = instance_prompt_hidden_states + text_ids = instance_text_ids + if args.with_prior_preservation: + prompt_embeds = torch.cat([prompt_embeds, class_prompt_hidden_states], dim=0) + text_ids = torch.cat([text_ids, class_text_ids], dim=0) + static_prompt_embeds = prompt_embeds + static_text_ids = text_ids + + # if cache_latents is set to True, we encode images to latents and store them. + # Similar to pre-encoding in the case of a single instance prompt, if custom prompts are provided + # we encode them in advance as well. + if args.cache_latents: + instance_latents_cache = [None] * train_dataset.num_instance_images + class_latents_cache = [None] * train_dataset.num_instance_images if args.with_prior_preservation else None + if train_dataset.custom_instance_prompts: + prompt_embeds_cache = [None] * train_dataset.num_instance_images + text_ids_cache = [None] * train_dataset.num_instance_images + if precompute_latents: + cache_batch_sampler = BucketBatchSampler( + train_dataset, batch_size=args.train_batch_size, drop_last=False, seed=args.seed + ) + cache_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=cache_batch_sampler, + collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation), + num_workers=args.dataloader_num_workers, + ) + for batch in tqdm(cache_dataloader, desc="Caching latents"): + with torch.no_grad(): + sample_indices = batch["indices"] + if args.cache_latents: + with offload_models(vae, device=accelerator.device, offload=args.offload): + batch["pixel_values"] = batch["pixel_values"].to( + accelerator.device, non_blocking=True, dtype=vae.dtype + ) + latents = vae.encode(batch["pixel_values"]).latent_dist.mode() + if args.with_prior_preservation: + instance_latents, class_latents = torch.chunk(latents, 2, dim=0) + else: + instance_latents = latents + for i, idx in enumerate(sample_indices): + instance_latents_cache[idx] = instance_latents[i : i + 1] + if args.with_prior_preservation: + class_latents_cache[idx] = class_latents[i : i + 1] + if train_dataset.custom_instance_prompts: + if args.fsdp_text_encoder: + prompt_embeds, text_ids = compute_text_embeddings( + batch["instance_prompts"], text_encoder, tokenizer, pad_to_length=uniform_window + ) + else: + with offload_models(text_encoder, device=accelerator.device, offload=args.offload): + prompt_embeds, text_ids = compute_text_embeddings( + batch["instance_prompts"], text_encoder, tokenizer, pad_to_length=uniform_window + ) + for i, idx in enumerate(sample_indices): + prompt_embeds_cache[idx] = prompt_embeds[i : i + 1] + text_ids_cache[idx] = text_ids[i : i + 1] + + if args.cache_latents: + assert all(latents is not None for latents in instance_latents_cache), "Latent cache has unfilled entries." + if args.with_prior_preservation: + assert all(latents is not None for latents in class_latents_cache), ( + "Class latent cache has unfilled entries." + ) + if train_dataset.custom_instance_prompts: + assert all(embeds is not None for embeds in prompt_embeds_cache), ( + "Prompt embedding cache has unfilled entries." + ) + assert all(ids is not None for ids in text_ids_cache), "Text ID cache has unfilled entries." + + # Caption dropout swaps the empty-prompt embedding in for a dropped caption. FLUX.3 pads each prompt to its own + # window (a multiple of TEXT_PAD_MULTIPLE) and the transformer attends to it, so the empty prompt must match the + # dropped sample's window. Precompute it for every window in use here, while the text encoder is still loaded + # (it is deleted just below, so this cannot be done in the training loop). + empty_prompt_by_window = {} + if args.caption_dropout > 0: + dropout_windows = ( + {embeds.shape[1] for embeds in prompt_embeds_cache} + if train_dataset.custom_instance_prompts + else {static_prompt_embeds.shape[1]} + ) + for window in sorted(dropout_windows): + with offload_models(text_encoder, device=accelerator.device, offload=args.offload): + empty_prompt_by_window[window] = compute_text_embeddings( + "", text_encoder, tokenizer, pad_to_length=window + ) + + # move back to cpu before deleting to ensure memory is freed see: https://github.com/huggingface/diffusers/issues/11376#issue-3008144624 + if args.cache_latents: + vae = vae.to("cpu") + del vae + + # move back to cpu before deleting to ensure memory is freed see: https://github.com/huggingface/diffusers/issues/11376#issue-3008144624 + text_encoder = text_encoder.to("cpu") + del text_encoder, tokenizer + free_memory() + + # Scheduler and math around the number of training steps. + # Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation. + num_warmup_steps_for_scheduler = args.lr_warmup_steps * accelerator.num_processes + if args.max_train_steps is None: + len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes) + num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps) + num_training_steps_for_scheduler = ( + args.num_train_epochs * accelerator.num_processes * num_update_steps_per_epoch + ) + else: + num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=num_warmup_steps_for_scheduler, + num_training_steps=num_training_steps_for_scheduler, + num_cycles=args.lr_num_cycles, + power=args.lr_power, + ) + + # Prepare everything with our `accelerator`. + transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + transformer, optimizer, train_dataloader, lr_scheduler + ) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + if num_training_steps_for_scheduler != args.max_train_steps: + logger.warning( + f"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match " + f"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. " + f"This inconsistency may result in the learning rate scheduler not functioning properly." + ) + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if accelerator.is_main_process: + tracker_name = "dreambooth-flux3-lora" + args_cp = vars(args).copy() + args_cp["text_encoder_out_layers"] = str(args_cp["text_encoder_out_layers"]) + accelerator.init_trackers(tracker_name, config=args_cp) + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + + logger.info("***** Running training *****") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num batches each epoch = {len(train_dataloader)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + global_step = 0 + first_epoch = 0 + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the mos recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(os.path.join(args.output_dir, path)) + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + first_epoch = global_step // num_update_steps_per_epoch + + else: + initial_global_step = 0 + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + + def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): + sigmas = noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype) + schedule_timesteps = noise_scheduler_copy.timesteps.to(accelerator.device) + timesteps = timesteps.to(accelerator.device) + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + for epoch in range(first_epoch, args.num_train_epochs): + transformer.train() + + for batch in train_dataloader: + models_to_accumulate = [transformer] + sample_indices = batch["indices"] + prompts = batch["prompts"] + + with accelerator.accumulate(models_to_accumulate): + if train_dataset.custom_instance_prompts: + prompt_embeds = torch.cat([prompt_embeds_cache[idx] for idx in sample_indices], dim=0) + text_ids = torch.cat([text_ids_cache[idx] for idx in sample_indices], dim=0) + if args.with_prior_preservation: + prompt_embeds = torch.cat( + [prompt_embeds, class_prompt_hidden_states.repeat(len(sample_indices), 1, 1)], dim=0 + ) + text_ids = torch.cat([text_ids, class_text_ids.repeat(len(sample_indices), 1, 1)], dim=0) + else: + # With prior preservation, prompt_embeds/text_ids already contain [instance, class] entries, + # while collate_fn orders batches as [inst1..instB, class1..classB]. Repeat each entry along + # dim 0 to preserve that grouping instead of interleaving [inst, class, inst, class, ...]. + num_repeat_elements = len(prompts) // 2 if args.with_prior_preservation else len(prompts) + prompt_embeds = static_prompt_embeds.repeat_interleave(num_repeat_elements, dim=0) + text_ids = static_text_ids.repeat_interleave(num_repeat_elements, dim=0) + + # Caption dropout: replace a sample's caption embedding with the empty-prompt embedding so it + # trains unconditionally. Both paths order embeddings as [instance..., class...], so only the + # leading instance rows are dropped, never class/prior captions. + if args.caption_dropout > 0: + n_inst = len(sample_indices) + drop_mask = torch.rand(n_inst, device=prompt_embeds.device) < args.caption_dropout + if drop_mask.any(): + empty_hidden_states, empty_text_ids = empty_prompt_by_window[prompt_embeds.shape[1]] + prompt_embeds[:n_inst][drop_mask] = empty_hidden_states.to( + device=prompt_embeds.device, dtype=prompt_embeds.dtype + ) + text_ids[:n_inst][drop_mask] = empty_text_ids.to(device=text_ids.device, dtype=text_ids.dtype) + + # Convert images to latent space + if args.cache_latents: + model_input = torch.cat([instance_latents_cache[idx] for idx in sample_indices], dim=0) + if args.with_prior_preservation: + model_input = torch.cat( + [model_input, torch.cat([class_latents_cache[idx] for idx in sample_indices], dim=0)], + dim=0, + ) + else: + with offload_models(vae, device=accelerator.device, offload=args.offload): + pixel_values = batch["pixel_values"].to(device=accelerator.device, dtype=vae.dtype) + model_input = vae.encode(pixel_values).latent_dist.mode() + + model_input = patchify_latents(model_input) + model_input = (model_input - latents_bn_mean) / latents_bn_std + + _, _, lh, lw = model_input.shape + # Sample noise that we'll add to the latents + noise = torch.randn_like(model_input) + bsz = model_input.shape[0] + + # Sample a random timestep for each image + # for weighting schemes where we sample timesteps non-uniformly + u = compute_density_for_timestep_sampling( + weighting_scheme=args.weighting_scheme, + batch_size=bsz, + logit_mean=args.logit_mean, + logit_std=args.logit_std, + mode_scale=args.mode_scale, + ) + indices = (u * noise_scheduler_copy.config.num_train_timesteps).long() + timesteps = noise_scheduler_copy.timesteps[indices].to(device=model_input.device) + + # Add noise according to flow matching. + # zt = (1 - texp) * x + texp * z1 + sigmas = get_sigmas(timesteps, n_dim=model_input.ndim, dtype=model_input.dtype) + noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise + + # [B, C, H, W] -> [B, H*W, C] + packed_noisy_model_input = pack_latents(noisy_model_input) + image_ids = prepare_latent_ids(lh, lw, accelerator.device) # (image_seq_len, 4) + # FLUX.3 is not guidance-distilled: no guidance argument. The flow timestep is a per-sample scalar in + # [0, 1] (the transformer scales it by 1000 and broadcasts it over the image tokens internally); the + # text context sits at timestep 0. + timestep = sigmas.reshape(bsz) + + # Predict the noise residual. The transformer returns only the generated image tokens, so there is + # nothing to slice/unpack afterwards. + with accelerator.autocast(): + model_pred = transformer( + hidden_states=packed_noisy_model_input, # (B, image_seq_len, C) + img_ids=image_ids, # (image_seq_len, 4) + timestep=timestep, # (B,) + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids[0], # (text_seq_len, 4), shared across the batch + return_dict=False, + )[0] + + # these weighting schemes use a uniform timestep sampling + # and instead post-weight the loss + weighting = compute_loss_weighting_for_sd3( + weighting_scheme=args.weighting_scheme, sigmas=sigmas.reshape(bsz, 1, 1) + ) + + # flow matching loss (velocity target) on the packed image tokens + target = pack_latents(noise) - pack_latents(model_input) + + if args.with_prior_preservation: + # Chunk the noise and model_pred into two parts and compute the loss on each part separately. + model_pred, model_pred_prior = torch.chunk(model_pred, 2, dim=0) + target, target_prior = torch.chunk(target, 2, dim=0) + weighting, weighting_prior = torch.chunk(weighting, 2, dim=0) + + # Compute prior loss + prior_loss = torch.mean( + (weighting_prior.float() * (model_pred_prior.float() - target_prior.float()) ** 2).reshape( + target_prior.shape[0], -1 + ), + 1, + ) + prior_loss = prior_loss.mean() + + # Compute regular loss. + loss = torch.mean( + (weighting.float() * (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1), + 1, + ) + loss = loss.mean() + + if args.with_prior_preservation: + # Add the prior loss to the instance loss. + loss = loss + args.prior_loss_weight * prior_loss + + accelerator.backward(loss) + if accelerator.sync_gradients: + params_to_clip = transformer.parameters() + accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) + + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # Checks if the accelerator has performed an optimization step behind the scenes + if accelerator.sync_gradients: + progress_bar.update(1) + global_step += 1 + + if accelerator.is_main_process or is_fsdp: + if global_step % args.checkpointing_steps == 0: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + accelerator.log(logs, step=global_step) + + if global_step >= args.max_train_steps: + break + + if accelerator.is_main_process: + if args.validation_prompt is not None and epoch % args.validation_epochs == 0: + transformer_ = unwrap_model(transformer) + transformer_.eval() + pipeline = make_validation_pipeline(args, transformer_, weight_dtype, accelerator.device) + images = log_validation( + pipeline=pipeline, + args=args, + accelerator=accelerator, + pipeline_args=validation_pipeline_args, + epoch=epoch, + torch_dtype=weight_dtype, + ) + transformer_.train() + free_memory() + + # Save the lora layers + accelerator.wait_for_everyone() + + if is_fsdp: + transformer = unwrap_model(transformer) + state_dict = accelerator.get_state_dict(transformer) + if accelerator.is_main_process: + modules_to_save = {} + if is_fsdp: + if args.bnb_quantization_config_path is None: + if args.upcast_before_saving: + state_dict = { + k: v.to(torch.float32) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items() + } + else: + state_dict = { + k: v.to(weight_dtype) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items() + } + + transformer_lora_layers = get_peft_model_state_dict( + transformer, + state_dict=state_dict, + ) + transformer_lora_layers = { + k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v + for k, v in transformer_lora_layers.items() + } + + else: + transformer = unwrap_model(transformer) + if args.bnb_quantization_config_path is None: + if args.upcast_before_saving: + transformer.to(torch.float32) + else: + transformer = transformer.to(weight_dtype) + transformer_lora_layers = get_peft_model_state_dict(transformer) + + modules_to_save["transformer"] = transformer + + Flux3LoraLoaderMixin.save_lora_weights( + save_directory=args.output_dir, + transformer_lora_layers=transformer_lora_layers, + **_collate_lora_metadata(modules_to_save), + ) + + images = [] + run_validation = (args.validation_prompt and args.num_validation_images > 0) or (args.final_validation_prompt) + should_run_final_inference = not args.skip_final_inference and run_validation + if should_run_final_inference: + pipeline = make_validation_pipeline(args, transformer, weight_dtype, accelerator.device) + + # run inference + images = [] + if args.validation_prompt and args.num_validation_images > 0: + images = log_validation( + pipeline=pipeline, + args=args, + accelerator=accelerator, + pipeline_args=validation_pipeline_args, + epoch=epoch, + is_final_validation=True, + torch_dtype=weight_dtype, + ) + del pipeline + free_memory() + + validation_prompt = args.validation_prompt if args.validation_prompt else args.final_validation_prompt + quant_training = None + if args.do_fp8_training: + quant_training = "FP8 TorchAO" + elif args.bnb_quantization_config_path: + quant_training = "BitsandBytes" + save_model_card( + (args.hub_model_id or Path(args.output_dir).name) if not args.push_to_hub else repo_id, + images=images, + base_model=args.pretrained_model_name_or_path, + instance_prompt=args.instance_prompt, + validation_prompt=validation_prompt, + repo_folder=args.output_dir, + quant_training=quant_training, + ) + + if args.push_to_hub: + upload_folder( + repo_id=repo_id, + folder_path=args.output_dir, + commit_message="End of training", + ignore_patterns=["step_*", "epoch_*"], + ) + + accelerator.end_training() + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/examples/dreambooth/train_dreambooth_lora_flux3_klein_img2img.py b/examples/dreambooth/train_dreambooth_lora_flux3_klein_img2img.py new file mode 100644 index 0000000000..0fec7b2ffc --- /dev/null +++ b/examples/dreambooth/train_dreambooth_lora_flux3_klein_img2img.py @@ -0,0 +1,2080 @@ +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. 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. + +# /// script +# dependencies = [ +# "diffusers @ git+https://github.com/huggingface/diffusers.git", +# "torch>=2.0.0", +# "accelerate>=0.31.0", +# "transformers>=4.41.2", +# "ftfy", +# "tensorboard", +# "Jinja2", +# "peft>=0.11.1", +# "sentencepiece", +# "torchvision", +# "datasets", +# "bitsandbytes", +# "prodigyopt", +# ] +# /// + +import argparse +import copy +import itertools +import json +import logging +import math +import os +import random +import shutil +from contextlib import nullcontext +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import transformers +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed +from huggingface_hub import create_repo, upload_folder +from peft import LoraConfig, prepare_model_for_kbit_training, set_peft_model_state_dict +from peft.utils import get_peft_model_state_dict +from PIL import Image +from PIL.ImageOps import exif_transpose +from torch.utils.data import Dataset +from torch.utils.data.sampler import BatchSampler +from torchvision import transforms +from torchvision.transforms import functional as TF +from tqdm.auto import tqdm +from transformers import AutoProcessor, Qwen3VLForConditionalGeneration + +import diffusers +from diffusers import ( + AutoencoderKLFlux2, + BitsAndBytesConfig, + FlowMatchEulerDiscreteScheduler, + Flux3Transformer2DModel, +) +from diffusers.loaders import Flux3LoraLoaderMixin +from diffusers.modular_pipelines.flux3 import Flux3AutoBlocks +from diffusers.optimization import get_scheduler +from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor +from diffusers.training_utils import ( + _collate_lora_metadata, + _to_cpu_contiguous, + cast_training_params, + compute_density_for_timestep_sampling, + compute_loss_weighting_for_sd3, + find_nearest_bucket, + free_memory, + generate_aspect_ratio_buckets, + get_fsdp_kwargs_from_accelerator, + offload_models, + parse_buckets_string, + wrap_with_fsdp, +) +from diffusers.utils import ( + check_min_version, + convert_unet_state_dict_to_peft, + is_wandb_available, + load_image, +) +from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card +from diffusers.utils.import_utils import is_torch_npu_available +from diffusers.utils.torch_utils import is_compiled_module + + +if getattr(torch, "distributed", None) is not None: + import torch.distributed as dist + +if is_wandb_available(): + import wandb + +# Will error if the minimal version of diffusers is not installed. Remove at your own risks. +check_min_version("0.40.0.dev0") + +logger = get_logger(__name__) + +# FLUX.3 ships no `scheduler` subfolder: the modular pipeline builds its scheduler `from_config` with a fixed +# shift of 6.93 (the `Flux3SetTimestepsStep` default, and the reference sampler's `alpha`). There is no +# resolution-dependent dynamic shifting. Constructed here for the same reason, rather than loaded. +FLOW_SHIFT = 6.93 + +# FLUX.3 image RoPE scale applied to the H/W position-id axes (see the flux3 modular pipeline). +IMAGE_ROPE_SCALE = 0.5 + +# The text window is right-padded to a multiple of this many tokens (the `Flux3TextEncoderStep` default). +TEXT_PAD_MULTIPLE = 80 + + +# --- FLUX.3 latent / text helpers (bit-identical to the modular pipeline blocks) --- +def patchify_latents(latents): + """(B, C, H, W) -> (B, C*4, H//2, W//2).""" + b, c, h, w = latents.shape + latents = latents.view(b, c, h // 2, 2, w // 2, 2).permute(0, 1, 3, 5, 2, 4) + return latents.reshape(b, c * 4, h // 2, w // 2) + + +def pack_latents(latents): + """(B, C, H, W) -> (B, H*W, C).""" + b, c, h, w = latents.shape + return latents.reshape(b, c, h * w).permute(0, 2, 1) + + +def prepare_latent_ids(height, width, device, image_rope_scale=IMAGE_ROPE_SCALE): + """(H*W, 4) position ids (t, h, w, l) for the generated image; H/W axes scaled by the image RoPE scale.""" + ids = torch.cartesian_prod(torch.arange(1), torch.arange(height), torch.arange(width), torch.arange(1)) + ids = ids.to(dtype=torch.float32) + ids[:, 1:3] = ids[:, 1:3] * image_rope_scale + return ids.to(device) + + +def prepare_image_cond_ids(height, width, device, image_rope_scale=IMAGE_ROPE_SCALE): + """(H*W, 4) position ids for the conditioning image: reference sits at the L axis position -1; H/W scaled.""" + ids = torch.cartesian_prod(torch.arange(1), torch.arange(height), torch.arange(width), torch.tensor([-1])) + ids = ids.to(dtype=torch.float32) + ids[:, 1:3] = ids[:, 1:3] * image_rope_scale + return ids.to(device) + + +def prepare_text_ids(seq_len, device): + """(seq_len, 4) position ids for text tokens; text advances along the L axis.""" + ids = torch.cartesian_prod(torch.arange(1), torch.arange(1), torch.arange(1), torch.arange(seq_len)) + return ids.to(dtype=torch.float32).to(device) + + +@torch.no_grad() +def get_qwen3vl_prompt_embeds( + text_encoder, tokenizer, prompt, device, dtype, text_encoder_out_layers, max_sequence_length, pad_to_length=None +): + """Stack the FLUX.3 Qwen3-VL hidden layers -> (B, seq, len(layers) * hidden_dim). + + Mirrors `Flux3TextEncoderStep._get_qwen3_vl_prompt_embeds`: each prompt goes through the chat template and is + right-padded to a multiple of `TEXT_PAD_MULTIPLE` (capped at `max_sequence_length`). Qwen3-VL masks the padding + (so a prompt's real-token embeddings are independent of the window size), but the transformer attends to it, so + the window is part of the recipe -- for a single prompt this reproduces the pipeline window exactly. When a batch + of differing prompts is encoded together, they share the batch's largest window so the embeddings can be stacked; + with `train_batch_size=1` or a shared instance prompt every prompt keeps its own pipeline-exact window. + The position ids of the padding tokens are passed explicitly to match the Qwen3-VL 4.57 training environment. + """ + prompt = [prompt] if isinstance(prompt, str) else prompt + texts = [ + tokenizer.apply_chat_template([{"role": "user", "content": p}], tokenize=False, add_generation_prompt=True) + for p in prompt + ] + real_lengths = [ + tokenizer.tokenizer(t, return_tensors="pt", truncation=True, max_length=max_sequence_length)[ + "input_ids" + ].shape[1] + for t in texts + ] + padded_length = min(math.ceil(max(real_lengths) / TEXT_PAD_MULTIPLE) * TEXT_PAD_MULTIPLE, max_sequence_length) + if pad_to_length is not None: + # Force a specific window (caption dropout encodes the empty prompt at a dropped instruction's window). Never + # shrink below the prompt's own real length, never exceed max_sequence_length. + padded_length = min(max(padded_length, pad_to_length), max_sequence_length) + enc = tokenizer.tokenizer( + texts, + return_tensors="pt", + padding="max_length", + padding_side="right", + truncation=True, + max_length=padded_length, + ) + input_ids, attn = enc["input_ids"].to(device), enc["attention_mask"].to(device) + position_ids = torch.arange(padded_length, device=device)[None].masked_fill(attn == 0, 1) + out = text_encoder( + input_ids=input_ids, + attention_mask=attn, + position_ids=position_ids, + output_hidden_states=True, + use_cache=False, + ) + stacked = torch.stack([out.hidden_states[k] for k in text_encoder_out_layers], dim=1).to( + dtype=dtype, device=device + ) + b, n, s, d = stacked.shape + return stacked.permute(0, 2, 1, 3).reshape(b, s, n * d) + + +def save_model_card( + repo_id: str, + images=None, + base_model: str = None, + instance_prompt=None, + validation_prompt=None, + repo_folder=None, + quant_training=None, +): + widget_dict = [] + if images is not None: + for i, image in enumerate(images): + image.save(os.path.join(repo_folder, f"image_{i}.png")) + widget_dict.append( + {"text": validation_prompt if validation_prompt else " ", "output": {"url": f"image_{i}.png"}} + ) + + model_description = f""" +# FLUX.3 DreamBooth LoRA - {repo_id} + + + +## Model description + +These are {repo_id} DreamBooth LoRA weights for {base_model}. + +The weights were trained using [DreamBooth](https://dreambooth.github.io/) with the [FLUX.3 diffusers trainer](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_flux3_klein.md). + +Quant training? {quant_training} + +## Trigger words + +You should use `{instance_prompt}` to trigger the image generation. + +## Download model + +[Download the *.safetensors LoRA]({repo_id}/tree/main) in the Files & versions tab. + +## Use it with the [🧨 diffusers library](https://github.com/huggingface/diffusers) + +```py +from diffusers import ModularPipeline +import torch +pipeline = ModularPipeline.from_pretrained("{base_model}") +pipeline.load_components(torch_dtype=torch.bfloat16) +pipeline.to('cuda') +pipeline.load_lora_weights('{repo_id}', weight_name='pytorch_lora_weights.safetensors') +image = pipeline(prompt='{validation_prompt if validation_prompt else instance_prompt}', output_type="pil").values["images"][0] +``` + +For more details, including weighting, merging and fusing LoRAs, check the [documentation on loading LoRAs in diffusers](https://huggingface.co/docs/diffusers/main/en/using-diffusers/loading_adapters) + +## License + +Please adhere to the licensing terms as described in the model repository. +""" + model_card = load_or_create_model_card( + repo_id_or_path=repo_id, + from_training=True, + license="other", + base_model=base_model, + prompt=instance_prompt, + model_description=model_description, + widget=widget_dict, + ) + tags = [ + "text-to-image", + "diffusers-training", + "diffusers", + "lora", + "flux3", + "flux3-diffusers", + "template:sd-lora", + ] + + model_card = populate_model_card(model_card, tags=tags) + model_card.save(os.path.join(repo_folder, "README.md")) + + +def log_validation( + pipeline, + args, + accelerator, + pipeline_args, + epoch, + torch_dtype, + is_final_validation=False, +): + args.num_validation_images = args.num_validation_images if args.num_validation_images else 1 + logger.info( + f"Running validation... \n Generating {args.num_validation_images} images with prompt:" + f" {args.validation_prompt}." + ) + + # run inference + generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed is not None else None + autocast_ctx = torch.autocast(accelerator.device.type) if not is_final_validation else nullcontext() + + images = [] + for _ in range(args.num_validation_images): + with autocast_ctx: + output = pipeline(generator=generator, output_type="pil", **pipeline_args) + images.append(output.values["images"][0]) + + for tracker in accelerator.trackers: + phase_name = "test" if is_final_validation else "validation" + if tracker.name == "tensorboard": + np_images = np.stack([np.asarray(img) for img in images]) + tracker.writer.add_images(phase_name, np_images, epoch, dataformats="NHWC") + if tracker.name == "wandb": + tracker.log( + { + phase_name: [ + wandb.Image(image, caption=f"{i}: {args.validation_prompt}") for i, image in enumerate(images) + ] + } + ) + + del pipeline + free_memory() + + return images + + +def make_validation_pipeline(args, transformer, weight_dtype, device): + # Build a modular pipeline for validation that REUSES the in-memory trained transformer, so validation runs + # on exactly the weights being trained and the DiT is never loaded a second time. Only the small auxiliary + # components are (re)loaded (they are freed after latent caching). + vae = AutoencoderKLFlux2.from_pretrained( + args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision, torch_dtype=torch.float32 + ) + text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( + args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision, torch_dtype=weight_dtype + ) + tokenizer = AutoProcessor.from_pretrained( + args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision + ) + scheduler = FlowMatchEulerDiscreteScheduler(shift=FLOW_SHIFT, use_dynamic_shifting=False) + pipeline = Flux3AutoBlocks().init_pipeline() + pipeline.update_components( + transformer=transformer, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, scheduler=scheduler + ) + pipeline.to(device) + return pipeline + + +def module_filter_fn(mod: torch.nn.Module, fqn: str): + # don't convert the output module (FLUX.3's output projection) + if fqn.endswith("proj_out"): + return False + # don't convert linear modules with weight dimensions not divisible by 16 + if isinstance(mod, torch.nn.Linear): + if mod.in_features % 16 != 0 or mod.out_features % 16 != 0: + return False + return True + + +def parse_args(input_args=None): + parser = argparse.ArgumentParser(description="Simple example of a training script.") + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + required=True, + help="Path to pretrained model or model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--revision", + type=str, + default=None, + required=False, + help="Revision of pretrained model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--bnb_quantization_config_path", + type=str, + default=None, + help="Quantization config in a JSON file that will be used to define the bitsandbytes quant config of the DiT.", + ) + parser.add_argument( + "--do_fp8_training", + action="store_true", + help="if we are doing FP8 training.", + ) + parser.add_argument( + "--variant", + type=str, + default=None, + help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16", + ) + parser.add_argument( + "--dataset_name", + type=str, + default=None, + help=( + "The name of the Dataset (from the HuggingFace hub) containing the training data of instance images (could be your own, possibly private," + " dataset). It can also be a path pointing to a local copy of a dataset in your filesystem," + " or to a folder containing files that 🤗 Datasets can understand." + ), + ) + parser.add_argument( + "--dataset_config_name", + type=str, + default=None, + help="The config of the Dataset, leave as None if there's only one config.", + ) + parser.add_argument( + "--instance_data_dir", + type=str, + default=None, + help=("A folder containing the training data. "), + ) + + parser.add_argument( + "--cache_dir", + type=str, + default=None, + help="The directory where the downloaded models and datasets will be stored.", + ) + + parser.add_argument( + "--image_column", + type=str, + default="image", + help="The column of the dataset containing the target image. By " + "default, the standard Image Dataset maps out 'file_name' " + "to 'image'.", + ) + parser.add_argument( + "--cond_image_column", + type=str, + default=None, + help="Column in the dataset containing the condition image. Must be specified when performing I2I fine-tuning", + ) + parser.add_argument( + "--caption_column", + type=str, + default=None, + help="The column of the dataset containing the instance prompt for each image", + ) + + parser.add_argument("--repeats", type=int, default=1, help="How many times to repeat the training data.") + parser.add_argument( + "--caption_dropout", + type=float, + default=0.0, + help=( + "Probability of replacing a sample's instruction with an empty string during training, so that fraction" + " of samples is trained unconditionally on the text (the conditioning image is kept). Improves" + " classifier-free guidance. A common value is 0.1." + ), + ) + + parser.add_argument( + "--instance_prompt", + type=str, + default=None, + required=False, + help="The prompt with identifier specifying the instance, e.g. 'photo of a TOK dog', 'in the style of TOK'", + ) + parser.add_argument( + "--max_sequence_length", + type=int, + default=512, + help=( + "Cap on the number of Qwen3-VL prompt tokens; longer prompts are truncated. The text window is the " + f"prompt's length rounded up to a multiple of {TEXT_PAD_MULTIPLE}, capped at this value." + ), + ) + parser.add_argument( + "--text_encoder_out_layers", + type=int, + nargs="+", + default=[4, 8, 12, 16, 20, 24, 28, 32], + help="Text encoder hidden layers to compute the final text embeddings.", + ) + parser.add_argument( + "--validation_prompt", + type=str, + default=None, + help="A prompt that is used during validation to verify that the model is learning.", + ) + parser.add_argument( + "--validation_image", + type=str, + default=None, + help="path to an image that is used during validation as the condition image to verify that the model is learning.", + ) + parser.add_argument( + "--skip_final_inference", + default=False, + action="store_true", + help="Whether to skip the final inference step with loaded lora weights upon training completion. This will run intermediate validation inference if `validation_prompt` is provided. Specify to reduce memory.", + ) + parser.add_argument( + "--final_validation_prompt", + type=str, + default=None, + help="A prompt that is used during a final validation to verify that the model is learning. Ignored if `--validation_prompt` is provided.", + ) + parser.add_argument( + "--num_validation_images", + type=int, + default=4, + help="Number of images that should be generated during validation with `validation_prompt`.", + ) + parser.add_argument( + "--validation_epochs", + type=int, + default=50, + help=( + "Run dreambooth validation every X epochs. Dreambooth validation consists of running the prompt" + " `args.validation_prompt` multiple times: `args.num_validation_images`." + ), + ) + parser.add_argument( + "--rank", + type=int, + default=4, + help=("The dimension of the LoRA update matrices."), + ) + parser.add_argument( + "--lora_alpha", + type=int, + default=4, + help="LoRA alpha to be used for additional scaling.", + ) + parser.add_argument("--lora_dropout", type=float, default=0.0, help="Dropout probability for LoRA layers") + + parser.add_argument( + "--output_dir", + type=str, + default="flux-dreambooth-lora", + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + parser.add_argument( + "--resolution", + type=int, + default=512, + help=( + "The resolution for input images, all the images in the train/validation dataset will be resized to this" + " resolution" + ), + ) + parser.add_argument( + "--aspect_ratio_buckets", + type=str, + default=None, + help=( + "Aspect ratio buckets to use for training. Define as a string of 'h1,w1;h2,w2;...'. " + "e.g. '1024,1024;768,1360;1360,768;880,1168;1168,880;1248,832;832,1248'. " + "Requires --use_aspect_ratio_buckets. Images are resized to cover and cropped to the nearest " + "listed bucket (smaller images are upscaled). When set, --resolution is ignored." + ), + ) + parser.add_argument( + "--use_aspect_ratio_buckets", + action="store_true", + help=( + "Enable aspect-ratio bucketing. Without --aspect_ratio_buckets, the buckets are computed on the " + "fly from --resolution and capped to each image's own resolution, so smaller images are assigned " + "to a smaller bucket instead of being upscaled. Provide --aspect_ratio_buckets to use an explicit list." + ), + ) + parser.add_argument( + "--center_crop", + default=False, + action="store_true", + help=( + "Whether to center crop the input images to the resolution. If not set, the images will be randomly" + " cropped. The images will be resized to the resolution first before cropping." + ), + ) + parser.add_argument( + "--random_flip", + action="store_true", + help="whether to randomly flip images horizontally", + ) + parser.add_argument( + "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader." + ) + parser.add_argument("--num_train_epochs", type=int, default=1) + parser.add_argument( + "--max_train_steps", + type=int, + default=None, + help="Total number of training steps to perform. If provided, overrides num_train_epochs.", + ) + parser.add_argument( + "--checkpointing_steps", + type=int, + default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" + " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument( + "--checkpoints_total_limit", + type=int, + default=None, + help=("Max number of checkpoints to store."), + ) + parser.add_argument( + "--resume_from_checkpoint", + type=str, + default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument( + "--gradient_accumulation_steps", + type=int, + default=1, + help="Number of updates steps to accumulate before performing a backward/update pass.", + ) + parser.add_argument( + "--gradient_checkpointing", + action="store_true", + help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + + parser.add_argument( + "--scale_lr", + action="store_true", + default=False, + help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", + ) + parser.add_argument( + "--lr_scheduler", + type=str, + default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument( + "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." + ) + parser.add_argument( + "--lr_num_cycles", + type=int, + default=1, + help="Number of hard resets of the lr in cosine_with_restarts scheduler.", + ) + parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.") + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=0, + help=( + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ), + ) + parser.add_argument( + "--weighting_scheme", + type=str, + default="none", + choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "none"], + help=('We default to the "none" weighting scheme for uniform sampling and uniform loss'), + ) + parser.add_argument( + "--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--mode_scale", + type=float, + default=1.29, + help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.", + ) + parser.add_argument( + "--optimizer", + type=str, + default="AdamW", + help=('The optimizer type to use. Choose between ["AdamW", "prodigy"]'), + ) + + parser.add_argument( + "--use_8bit_adam", + action="store_true", + help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW", + ) + + parser.add_argument( + "--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers." + ) + parser.add_argument( + "--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers." + ) + parser.add_argument( + "--prodigy_beta3", + type=float, + default=None, + help="coefficients for computing the Prodigy stepsize using running averages. If set to None, " + "uses the value of square root of beta2. Ignored if optimizer is adamW", + ) + parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay") + parser.add_argument("--adam_weight_decay", type=float, default=1e-04, help="Weight decay to use for unet params") + + parser.add_argument( + "--lora_layers", + type=str, + default=None, + help=( + 'The transformer modules to apply LoRA training on. Please specify the layers in a comma separated. E.g. - "to_k,to_q,to_v,to_out.0" will result in lora training of attention layers only' + ), + ) + + parser.add_argument( + "--adam_epsilon", + type=float, + default=1e-08, + help="Epsilon value for the Adam optimizer and Prodigy optimizers.", + ) + + parser.add_argument( + "--prodigy_use_bias_correction", + type=bool, + default=True, + help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW", + ) + parser.add_argument( + "--prodigy_safeguard_warmup", + type=bool, + default=True, + help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. " + "Ignored if optimizer is adamW", + ) + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") + parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") + parser.add_argument( + "--hub_model_id", + type=str, + default=None, + help="The name of the repository to keep in sync with the local `output_dir`.", + ) + parser.add_argument( + "--logging_dir", + type=str, + default="logs", + help=( + "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." + ), + ) + parser.add_argument( + "--allow_tf32", + action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument( + "--cache_latents", + action="store_true", + default=False, + help="Cache the VAE latents", + ) + parser.add_argument( + "--report_to", + type=str, + default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + parser.add_argument( + "--mixed_precision", + type=str, + default=None, + choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + parser.add_argument( + "--upcast_before_saving", + action="store_true", + default=False, + help=( + "Whether to upcast the trained transformer layers to float32 before saving (at the end of training). " + "Defaults to precision dtype used for training to save memory" + ), + ) + parser.add_argument( + "--offload", + action="store_true", + help="Whether to offload the VAE and the text encoder to CPU when they are not used.", + ) + + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument("--enable_npu_flash_attention", action="store_true", help="Enabla Flash Attention for NPU") + parser.add_argument("--fsdp_text_encoder", action="store_true", help="Use FSDP for text encoder") + + if input_args is not None: + args = parser.parse_args(input_args) + else: + args = parser.parse_args() + + if args.cond_image_column is None: + raise ValueError( + "you must provide --cond_image_column for image-to-image training. Otherwise use the FLUX.3 text-to-image training script (train_dreambooth_lora_flux3_klein.py)." + ) + else: + assert args.image_column is not None + assert args.caption_column is not None + + if args.dataset_name is None and args.instance_data_dir is None: + raise ValueError("Specify either `--dataset_name` or `--instance_data_dir`") + + if args.dataset_name is not None and args.instance_data_dir is not None: + raise ValueError("Specify only one of `--dataset_name` or `--instance_data_dir`") + + env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) + if env_local_rank != -1 and env_local_rank != args.local_rank: + args.local_rank = env_local_rank + + return args + + +class DreamBoothDataset(Dataset): + """ + A dataset to prepare the instance and class images with the prompts for fine-tuning the model. + It pre-processes the images. + """ + + def __init__( + self, + instance_data_root, + instance_prompt, + size=1024, + repeats=1, + center_crop=False, + buckets=None, + use_aspect_ratio_buckets=False, + bucket_divisibility=16, + bucket_base_resolutions=None, + ): + self.size = size + self.resolution = size + self.center_crop = center_crop + + self.instance_prompt = instance_prompt + self.custom_instance_prompts = None + + # Explicit user-provided bucket list (or None). The concrete list of buckets actually used is + # built from the data in `self.buckets` during preprocessing below. + self._explicit_buckets = buckets + self.use_aspect_ratio_buckets = use_aspect_ratio_buckets + self.bucket_divisibility = bucket_divisibility + self.bucket_base_resolutions = bucket_base_resolutions + + # if --dataset_name is provided or a metadata jsonl file is provided in the local --instance_data directory, + # we load the training data using load_dataset + if args.dataset_name is not None: + try: + from datasets import load_dataset + except ImportError: + raise ImportError( + "You are trying to load your data using the datasets library. If you wish to train using custom " + "captions please install the datasets library: `pip install datasets`. If you wish to load a " + "local folder containing images only, specify --instance_data_dir instead." + ) + # Downloading and loading a dataset from the hub. + # See more about loading custom images at + # https://huggingface.co/docs/datasets/v2.0.0/en/dataset_script + dataset = load_dataset( + args.dataset_name, + args.dataset_config_name, + cache_dir=args.cache_dir, + ) + # Preprocessing the datasets. + column_names = dataset["train"].column_names + + # 6. Get the column names for input/target. + if args.cond_image_column is not None and args.cond_image_column not in column_names: + raise ValueError( + f"`--cond_image_column` value '{args.cond_image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" + ) + if args.image_column is None: + image_column = column_names[0] + logger.info(f"image column defaulting to {image_column}") + else: + image_column = args.image_column + if image_column not in column_names: + raise ValueError( + f"`--image_column` value '{args.image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" + ) + instance_images = dataset["train"][image_column] + cond_images = None + cond_image_column = args.cond_image_column + if cond_image_column is not None: + cond_images = [dataset["train"][i][cond_image_column] for i in range(len(dataset["train"]))] + assert len(instance_images) == len(cond_images) + + if args.caption_column is None: + logger.info( + "No caption column provided, defaulting to instance_prompt for all images. If your dataset " + "contains captions/prompts for the images, make sure to specify the " + "column as --caption_column" + ) + self.custom_instance_prompts = None + else: + if args.caption_column not in column_names: + raise ValueError( + f"`--caption_column` value '{args.caption_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" + ) + custom_instance_prompts = dataset["train"][args.caption_column] + # create final list of captions according to --repeats + self.custom_instance_prompts = [] + for caption in custom_instance_prompts: + self.custom_instance_prompts.extend(itertools.repeat(caption, repeats)) + else: + self.instance_data_root = Path(instance_data_root) + if not self.instance_data_root.exists(): + raise ValueError("Instance images root doesn't exists.") + + instance_images = [Image.open(path) for path in list(Path(instance_data_root).iterdir())] + self.custom_instance_prompts = None + + self.instance_images = [] + self.cond_images = [] + for i, img in enumerate(instance_images): + self.instance_images.extend(itertools.repeat(img, repeats)) + if args.dataset_name is not None and cond_images is not None: + self.cond_images.extend(itertools.repeat(cond_images[i], repeats)) + + self.pixel_values = [] + self.cond_pixel_values = [] + self.buckets = [] + bucket_to_idx = {} + for i, image in enumerate(self.instance_images): + image = exif_transpose(image) + if not image.mode == "RGB": + image = image.convert("RGB") + dest_image = None + if self.cond_images: + dest_image = self.cond_images[i] + image_width, image_height = dest_image.size + if image_width * image_height > 1024 * 1024: + dest_image = Flux2ImageProcessor._resize_to_target_area(dest_image, 1024 * 1024) + image_width, image_height = dest_image.size + + # VAE spatial compression: 2 ** (len(vae.config.block_out_channels) - 1) = 8 for the Flux2 VAE + multiple_of = 2 ** (4 - 1) + image_width = (image_width // multiple_of) * multiple_of + image_height = (image_height // multiple_of) * multiple_of + image_processor = Flux2ImageProcessor() + dest_image = image_processor.preprocess( + dest_image, height=image_height, width=image_width, resize_mode="crop" + ) + # Convert back to PIL + dest_image = dest_image.squeeze(0) + if dest_image.min() < 0: + dest_image = (dest_image + 1) / 2 + dest_image = (torch.clamp(dest_image, 0, 1) * 255).byte().cpu() + + if dest_image.shape[0] == 1: + # Gray scale image + dest_image = Image.fromarray(dest_image.squeeze().numpy(), mode="L") + else: + # RGB scale image: (C, H, W) -> (H, W, C) + dest_image = TF.to_pil_image(dest_image) + + dest_image = exif_transpose(dest_image) + if not dest_image.mode == "RGB": + dest_image = dest_image.convert("RGB") + + width, height = image.size + + # Assign the image to a bucket. + target = self._bucket_for_image(height, width) + if target not in bucket_to_idx: + bucket_to_idx[target] = len(self.buckets) + self.buckets.append(target) + bucket_idx = bucket_to_idx[target] + + # based on the bucket assignment, define the transformations + image, dest_image = self.paired_transform( + image, + dest_image=dest_image, + size=target, + center_crop=args.center_crop, + random_flip=args.random_flip, + ) + self.pixel_values.append((image, bucket_idx)) + if dest_image is not None: + self.cond_pixel_values.append((dest_image, bucket_idx)) + + self.num_instance_images = len(self.instance_images) + self._length = self.num_instance_images + + def __len__(self): + return self._length + + def __getitem__(self, index): + example = {} + instance_image, bucket_idx = self.pixel_values[index % self.num_instance_images] + example["instance_images"] = instance_image + example["bucket_idx"] = bucket_idx + if self.cond_pixel_values: + dest_image, _ = self.cond_pixel_values[index % self.num_instance_images] + example["cond_images"] = dest_image + + if self.custom_instance_prompts: + caption = self.custom_instance_prompts[index % self.num_instance_images] + if caption: + example["instance_prompt"] = caption + else: + example["instance_prompt"] = self.instance_prompt + + else: # custom prompts were provided, but length does not match size of image dataset + example["instance_prompt"] = self.instance_prompt + + return example + + def _bucket_for_image(self, height, width): + # An explicit bucket list takes priority: pick the nearest, upscaling smaller images to cover it. + if self._explicit_buckets is not None: + return self._explicit_buckets[find_nearest_bucket(height, width, self._explicit_buckets)] + # On-the-fly bucketing: cap the ladder to the image's own resolution so smaller images are + # assigned to a smaller bucket rather than being upscaled (mirrors ostris' bucketing). + if self.use_aspect_ratio_buckets: + resolution = min(self.resolution, round((height * width) ** 0.5)) + ladder = generate_aspect_ratio_buckets( + resolution, + divisibility=self.bucket_divisibility, + base_resolutions=self.bucket_base_resolutions, + ) + return ladder[find_nearest_bucket(height, width, ladder)] + # No bucketing: a single square bucket reproduces the fixed-size resize + crop. + return (self.resolution, self.resolution) + + def paired_transform(self, image, dest_image=None, size=(224, 224), center_crop=False, random_flip=False): + # Resize preserving aspect ratio so the image covers the bucket, then crop to the bucket size. + # The same geometry is applied to the conditioning image so the pair stays aligned. + target_height, target_width = size + width, height = image.size + scale = max(target_height / height, target_width / width) + new_size = [round(height * scale), round(width * scale)] + image = TF.resize(image, new_size, interpolation=transforms.InterpolationMode.BILINEAR) + if dest_image is not None: + dest_image = TF.resize(dest_image, new_size, interpolation=transforms.InterpolationMode.BILINEAR) + if center_crop: + image = TF.center_crop(image, size) + if dest_image is not None: + dest_image = TF.center_crop(dest_image, size) + else: + i, j, h, w = transforms.RandomCrop.get_params(image, output_size=size) + image = TF.crop(image, i, j, h, w) + if dest_image is not None: + dest_image = TF.crop(dest_image, i, j, h, w) + if random_flip and random.random() < 0.5: + image = TF.hflip(image) + if dest_image is not None: + dest_image = TF.hflip(dest_image) + image = TF.normalize(TF.to_tensor(image), [0.5], [0.5]) + if dest_image is not None: + dest_image = TF.normalize(TF.to_tensor(dest_image), [0.5], [0.5]) + return (image, dest_image) if dest_image is not None else (image, None) + + +def collate_fn(examples): + pixel_values = [example["instance_images"] for example in examples] + prompts = [example["instance_prompt"] for example in examples] + + pixel_values = torch.stack(pixel_values) + pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float() + + batch = {"pixel_values": pixel_values, "prompts": prompts} + if any("cond_images" in example for example in examples): + cond_pixel_values = [example["cond_images"] for example in examples] + cond_pixel_values = torch.stack(cond_pixel_values) + cond_pixel_values = cond_pixel_values.to(memory_format=torch.contiguous_format).float() + batch.update({"cond_pixel_values": cond_pixel_values}) + return batch + + +class BucketBatchSampler(BatchSampler): + def __init__( + self, + dataset: DreamBoothDataset, + batch_size: int, + drop_last: bool = False, + shuffle_batches_each_epoch: bool = True, + ): + if not isinstance(batch_size, int) or batch_size <= 0: + raise ValueError("batch_size should be a positive integer value, but got batch_size={}".format(batch_size)) + if not isinstance(drop_last, bool): + raise ValueError("drop_last should be a boolean value, but got drop_last={}".format(drop_last)) + + self.dataset = dataset + self.batch_size = batch_size + self.drop_last = drop_last + self.shuffle_batches_each_epoch = shuffle_batches_each_epoch + + # Group indices by bucket + self.bucket_indices = [[] for _ in range(len(self.dataset.buckets))] + for idx, (_, bucket_idx) in enumerate(self.dataset.pixel_values): + self.bucket_indices[bucket_idx].append(idx) + + self.sampler_len = 0 + self.batches = [] + + # Pre-generate batches for each bucket + for indices_in_bucket in self.bucket_indices: + # Shuffle indices within the bucket + random.shuffle(indices_in_bucket) + # Create batches + for i in range(0, len(indices_in_bucket), self.batch_size): + batch = indices_in_bucket[i : i + self.batch_size] + if len(batch) < self.batch_size and self.drop_last: + continue # Skip partial batch if drop_last is True + self.batches.append(batch) + self.sampler_len += 1 # Count the number of batches + + if not self.shuffle_batches_each_epoch: + # Shuffle the precomputed batches once to mix buckets while keeping + # the order stable across epochs for step-indexed caches. + random.shuffle(self.batches) + + def __iter__(self): + if self.shuffle_batches_each_epoch: + random.shuffle(self.batches) + for batch in self.batches: + yield batch + + def __len__(self): + return self.sampler_len + + +def main(args): + if args.report_to == "wandb" and args.hub_token is not None: + raise ValueError( + "You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token." + " Please use `hf auth login` to authenticate with the Hub." + ) + + if torch.backends.mps.is_available() and args.mixed_precision == "bf16": + # due to pytorch#99272, MPS does not yet support bfloat16. + raise ValueError( + "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." + ) + if args.do_fp8_training: + from torchao.float8 import Float8LinearConfig, convert_to_float8_training + + logging_dir = Path(args.output_dir, args.logging_dir) + + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + kwargs_handlers=[kwargs], + ) + + # Disable AMP for MPS. + if torch.backends.mps.is_available(): + accelerator.native_amp = False + + if args.report_to == "wandb": + if not is_wandb_available(): + raise ImportError("Make sure to install wandb if you want to use it for logging during training.") + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed) + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + if args.push_to_hub: + repo_id = create_repo( + repo_id=args.hub_model_id or Path(args.output_dir).name, + exist_ok=True, + ).repo_id + + # Load the tokenizers + tokenizer = AutoProcessor.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="tokenizer", + revision=args.revision, + ) + + # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + + # Load scheduler and models + noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=FLOW_SHIFT, use_dynamic_shifting=False) + noise_scheduler_copy = copy.deepcopy(noise_scheduler) + vae = AutoencoderKLFlux2.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="vae", + revision=args.revision, + variant=args.variant, + ) + latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(accelerator.device) + latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to( + accelerator.device + ) + + quantization_config = None + if args.bnb_quantization_config_path is not None: + with open(args.bnb_quantization_config_path, "r") as f: + config_kwargs = json.load(f) + if "load_in_4bit" in config_kwargs and config_kwargs["load_in_4bit"]: + config_kwargs["bnb_4bit_compute_dtype"] = weight_dtype + quantization_config = BitsAndBytesConfig(**config_kwargs) + + transformer = Flux3Transformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + revision=args.revision, + variant=args.variant, + quantization_config=quantization_config, + torch_dtype=weight_dtype, + ) + if args.bnb_quantization_config_path is not None: + transformer = prepare_model_for_kbit_training(transformer, use_gradient_checkpointing=False) + + text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="text_encoder", + revision=args.revision, + variant=args.variant, + torch_dtype=weight_dtype, + ) + text_encoder.requires_grad_(False) + + # We only train the additional adapter LoRA layers + transformer.requires_grad_(False) + vae.requires_grad_(False) + + if args.enable_npu_flash_attention: + if is_torch_npu_available(): + logger.info("npu flash attention enabled.") + transformer.set_attention_backend("_native_npu") + else: + raise ValueError("npu flash attention requires torch_npu extensions and is supported only on npu device ") + + if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16: + # due to pytorch#99272, MPS does not yet support bfloat16. + raise ValueError( + "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." + ) + + to_kwargs = {"dtype": weight_dtype, "device": accelerator.device} if not args.offload else {"dtype": weight_dtype} + # flux vae is stable in bf16 so load it in weight_dtype to reduce memory + vae.to(**to_kwargs) + # we never offload the transformer to CPU, so we can just use the accelerator device + transformer_to_kwargs = ( + {"device": accelerator.device} + if args.bnb_quantization_config_path is not None + else {"device": accelerator.device, "dtype": weight_dtype} + ) + + is_fsdp = getattr(accelerator.state, "fsdp_plugin", None) is not None + if not is_fsdp: + transformer.to(**transformer_to_kwargs) + + if args.do_fp8_training: + convert_to_float8_training( + transformer, module_filter_fn=module_filter_fn, config=Float8LinearConfig(pad_inner_dim=True) + ) + + text_encoder.to(**to_kwargs) + + if args.gradient_checkpointing: + transformer.enable_gradient_checkpointing() + + if args.lora_layers is not None: + target_modules = [layer.strip() for layer in args.lora_layers.split(",")] + else: + # Attention + feed-forward projections across every block (the generated-image, reference and text mode + # blocks, and the joint single blocks). These are exactly the leaf modules the non-diffusers (ostris) + # FLUX.3 LoRA converter maps, so adapters trained here stay interchangeable with that layout. The reference + # ("image_cond") mode blocks are adapted too, so the LoRA can adjust how the condition image is consumed. + # Pass `--lora_layers` to target a different set. + target_modules = [ + "to_q", + "to_k", + "to_v", + "to_out.0", + "ff.linear_in", + "ff.linear_out", + ] + + # now we will add new LoRA weights the transformer layers + transformer_lora_config = LoraConfig( + r=args.rank, + lora_alpha=args.lora_alpha, + lora_dropout=args.lora_dropout, + init_lora_weights="gaussian", + target_modules=target_modules, + ) + transformer.add_adapter(transformer_lora_config) + + def unwrap_model(model): + model = accelerator.unwrap_model(model) + model = model._orig_mod if is_compiled_module(model) else model + return model + + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + transformer_cls = type(unwrap_model(transformer)) + + # 1) Validate and pick the transformer model + modules_to_save: dict[str, Any] = {} + transformer_model = None + + for model in models: + if isinstance(unwrap_model(model), transformer_cls): + transformer_model = model + modules_to_save["transformer"] = model + else: + raise ValueError(f"unexpected save model: {model.__class__}") + + if transformer_model is None: + raise ValueError("No transformer model found in 'models'") + + # 2) Optionally gather FSDP state dict once + state_dict = accelerator.get_state_dict(model) if is_fsdp else None + + # 3) Only main process materializes the LoRA state dict + transformer_lora_layers_to_save = None + if accelerator.is_main_process: + peft_kwargs = {} + if is_fsdp: + peft_kwargs["state_dict"] = state_dict + + transformer_lora_layers_to_save = get_peft_model_state_dict( + unwrap_model(transformer_model) if is_fsdp else transformer_model, + **peft_kwargs, + ) + + if is_fsdp: + transformer_lora_layers_to_save = _to_cpu_contiguous(transformer_lora_layers_to_save) + + # make sure to pop weight so that corresponding model is not saved again + if weights: + weights.pop() + + Flux3LoraLoaderMixin.save_lora_weights( + output_dir, + transformer_lora_layers=transformer_lora_layers_to_save, + **_collate_lora_metadata(modules_to_save), + ) + + def load_model_hook(models, input_dir): + transformer_ = None + + if not is_fsdp: + while len(models) > 0: + model = models.pop() + + if isinstance(unwrap_model(model), type(unwrap_model(transformer))): + transformer_ = unwrap_model(model) + else: + raise ValueError(f"unexpected save model: {model.__class__}") + else: + transformer_ = Flux3Transformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + ) + transformer_.add_adapter(transformer_lora_config) + + lora_state_dict = Flux3LoraLoaderMixin.lora_state_dict(input_dir) + + transformer_state_dict = { + f"{k.replace('transformer.', '')}": v for k, v in lora_state_dict.items() if k.startswith("transformer.") + } + transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict) + incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default") + if incompatible_keys is not None: + # check only for unexpected keys + unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) + if unexpected_keys: + logger.warning( + f"Loading adapter weights from state_dict led to unexpected keys not found in the model: " + f" {unexpected_keys}. " + ) + + # Make sure the trainable params are in float32. This is again needed since the base models + # are in `weight_dtype`. More details: + # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804 + if args.mixed_precision == "fp16": + models = [transformer_] + # only upcast trainable parameters (LoRA) into fp32 + cast_training_params(models) + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32 and torch.cuda.is_available(): + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + # Make sure the trainable params are in float32. + if args.mixed_precision == "fp16": + models = [transformer] + # only upcast trainable parameters (LoRA) into fp32 + cast_training_params(models, dtype=torch.float32) + + transformer_lora_parameters = list(filter(lambda p: p.requires_grad, transformer.parameters())) + + # Optimization parameters + transformer_parameters_with_lr = {"params": transformer_lora_parameters, "lr": args.learning_rate} + params_to_optimize = [transformer_parameters_with_lr] + + # Optimizer creation + if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"): + logger.warning( + f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]." + "Defaulting to adamW" + ) + args.optimizer = "adamw" + + if args.use_8bit_adam and not args.optimizer.lower() == "adamw": + logger.warning( + f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was " + f"set to {args.optimizer.lower()}" + ) + + if args.optimizer.lower() == "adamw": + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." + ) + + optimizer_class = bnb.optim.AdamW8bit + else: + optimizer_class = torch.optim.AdamW + + optimizer = optimizer_class( + params_to_optimize, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + if args.optimizer.lower() == "prodigy": + try: + import prodigyopt + except ImportError: + raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`") + + optimizer_class = prodigyopt.Prodigy + + if args.learning_rate <= 0.1: + logger.warning( + "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0" + ) + + optimizer = optimizer_class( + params_to_optimize, + betas=(args.adam_beta1, args.adam_beta2), + beta3=args.prodigy_beta3, + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + decouple=args.prodigy_decouple, + use_bias_correction=args.prodigy_use_bias_correction, + safeguard_warmup=args.prodigy_safeguard_warmup, + ) + + # Resolve the bucketing mode. Bucketing must be enabled explicitly with --use_aspect_ratio_buckets; + # a bucket list without that flag is an error. With the flag, an explicit --aspect_ratio_buckets list + # drives assignment, otherwise buckets are computed on the fly inside the dataset. Without the flag a + # single square bucket reproduces the fixed-size resize + crop. + if args.aspect_ratio_buckets is not None and not args.use_aspect_ratio_buckets: + raise ValueError("--aspect_ratio_buckets requires --use_aspect_ratio_buckets to be set.") + if args.aspect_ratio_buckets is not None: + buckets = parse_buckets_string(args.aspect_ratio_buckets) + use_aspect_ratio_buckets = False + logger.info(f"Using explicit aspect ratio buckets: {buckets}") + elif args.use_aspect_ratio_buckets: + buckets = None + use_aspect_ratio_buckets = True + logger.info( + "No --aspect_ratio_buckets provided; auto-computing aspect ratio buckets on the fly from --resolution." + ) + else: + buckets = [(args.resolution, args.resolution)] + use_aspect_ratio_buckets = False + + # Dataset and DataLoaders creation: + train_dataset = DreamBoothDataset( + instance_data_root=args.instance_data_dir, + instance_prompt=args.instance_prompt, + size=args.resolution, + repeats=args.repeats, + center_crop=args.center_crop, + buckets=buckets, + use_aspect_ratio_buckets=use_aspect_ratio_buckets, + ) + has_step_indexed_caches = precompute_latents = args.cache_latents or train_dataset.custom_instance_prompts + batch_sampler = BucketBatchSampler( + train_dataset, + batch_size=args.train_batch_size, + drop_last=True, + shuffle_batches_each_epoch=not has_step_indexed_caches, + ) + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + collate_fn=lambda examples: collate_fn(examples), + num_workers=args.dataloader_num_workers, + ) + + def compute_text_embeddings(prompt, text_encoder, tokenizer, pad_to_length=None): + # FLUX.3 has no classic text-encoding pipeline; encode directly with the frozen Qwen3-VL encoder. + with torch.no_grad(): + prompt_embeds = get_qwen3vl_prompt_embeds( + text_encoder, + tokenizer, + prompt, + accelerator.device, + weight_dtype, + text_encoder_out_layers=args.text_encoder_out_layers, + max_sequence_length=args.max_sequence_length, + pad_to_length=pad_to_length, + ) + # FLUX.3 text ids are positional only (identical for every sample of a given length). We carry them + # batched here so the repeat plumbing below is unchanged; the transformer consumes a single (seq, 4) + # slice at the forward call. + text_ids = prepare_text_ids(prompt_embeds.shape[1], accelerator.device) + text_ids = text_ids.unsqueeze(0).expand(prompt_embeds.shape[0], -1, -1) + return prompt_embeds, text_ids + + # If no type of tuning is done on the text_encoder and custom instance prompts are NOT + # provided (i.e. the --instance_prompt is used for all images), we encode the instance prompt once to avoid + # the redundant encoding. + if not train_dataset.custom_instance_prompts: + with offload_models(text_encoder, device=accelerator.device, offload=args.offload): + instance_prompt_hidden_states, instance_text_ids = compute_text_embeddings( + args.instance_prompt, text_encoder, tokenizer + ) + + # FLUX.3 validation runs through the modular pipeline (which encodes prompts itself), so we only assemble + # the generation kwargs here. The reference/condition image is routed to the pipeline's `image` input. + validation_pipeline_args = {} + if args.validation_prompt is not None: + validation_image = load_image(args.validation_image).convert("RGB") + validation_pipeline_args = { + "prompt": args.validation_prompt, + "image": validation_image, + "negative_prompt": "", + "height": args.resolution, + "width": args.resolution, + "num_inference_steps": 28, + # Encode the validation prompt exactly as training did (same window cap and hidden layers), so the + # modular pipeline's text embeddings match the ones the adapter was trained against. + "max_sequence_length": args.max_sequence_length, + "text_encoder_out_layers": tuple(args.text_encoder_out_layers), + } + + # Init FSDP for text encoder + if args.fsdp_text_encoder: + fsdp_kwargs = get_fsdp_kwargs_from_accelerator(accelerator) + text_encoder = wrap_with_fsdp( + model=text_encoder, + device=accelerator.device, + offload=args.offload, + limit_all_gathers=True, + use_orig_params=True, + fsdp_kwargs=fsdp_kwargs, + ) + dist.barrier() + + # If custom instance prompts are NOT provided (i.e. the instance prompt is used for all images), + # pack the statically computed variables appropriately here. This is so that we don't + # have to pass them to the dataloader. + if not train_dataset.custom_instance_prompts: + prompt_embeds = instance_prompt_hidden_states + text_ids = instance_text_ids + + # if cache_latents is set to True, we encode images to latents and store them. + # Similar to pre-encoding in the case of a single instance prompt, if custom prompts are provided + # we encode them in advance as well. + if precompute_latents: + prompt_embeds_cache = [] + text_ids_cache = [] + latents_cache = [] + cond_latents_cache = [] + for batch in tqdm(train_dataloader, desc="Caching latents"): + with torch.no_grad(): + if args.cache_latents: + with offload_models(vae, device=accelerator.device, offload=args.offload): + batch["pixel_values"] = batch["pixel_values"].to( + accelerator.device, non_blocking=True, dtype=vae.dtype + ) + latents_cache.append(vae.encode(batch["pixel_values"]).latent_dist) + batch["cond_pixel_values"] = batch["cond_pixel_values"].to( + accelerator.device, non_blocking=True, dtype=vae.dtype + ) + cond_latents_cache.append(vae.encode(batch["cond_pixel_values"]).latent_dist) + if train_dataset.custom_instance_prompts: + if args.fsdp_text_encoder: + prompt_embeds, text_ids = compute_text_embeddings(batch["prompts"], text_encoder, tokenizer) + else: + with offload_models(text_encoder, device=accelerator.device, offload=args.offload): + prompt_embeds, text_ids = compute_text_embeddings( + batch["prompts"], text_encoder, tokenizer + ) + prompt_embeds_cache.append(prompt_embeds) + text_ids_cache.append(text_ids) + + # Caption dropout replaces the instruction with the empty prompt for a fraction of samples (the conditioning + # image is kept). FLUX.3 pads each prompt to its own window (a multiple of TEXT_PAD_MULTIPLE) and the transformer + # attends to it, so the empty prompt must match the dropped sample's window. Precompute it for every window in + # use here, while the text encoder is still loaded (it is deleted below, so this cannot be done in the loop). + empty_prompt_by_window = {} + if args.caption_dropout > 0: + dropout_windows = ( + {embeds.shape[1] for embeds in prompt_embeds_cache} + if train_dataset.custom_instance_prompts + else {instance_prompt_hidden_states.shape[1]} + ) + for window in sorted(dropout_windows): + with offload_models(text_encoder, device=accelerator.device, offload=args.offload): + empty_prompt_by_window[window] = compute_text_embeddings( + "", text_encoder, tokenizer, pad_to_length=window + ) + + # move back to cpu before deleting to ensure memory is freed see: https://github.com/huggingface/diffusers/issues/11376#issue-3008144624 + if args.cache_latents: + vae = vae.to("cpu") + del vae + + # move back to cpu before deleting to ensure memory is freed see: https://github.com/huggingface/diffusers/issues/11376#issue-3008144624 + text_encoder = text_encoder.to("cpu") + del text_encoder, tokenizer + free_memory() + + # Scheduler and math around the number of training steps. + # Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation. + num_warmup_steps_for_scheduler = args.lr_warmup_steps * accelerator.num_processes + if args.max_train_steps is None: + len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes) + num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps) + num_training_steps_for_scheduler = ( + args.num_train_epochs * accelerator.num_processes * num_update_steps_per_epoch + ) + else: + num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=num_warmup_steps_for_scheduler, + num_training_steps=num_training_steps_for_scheduler, + num_cycles=args.lr_num_cycles, + power=args.lr_power, + ) + + # Prepare everything with our `accelerator`. + transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + transformer, optimizer, train_dataloader, lr_scheduler + ) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + if num_training_steps_for_scheduler != args.max_train_steps: + logger.warning( + f"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match " + f"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. " + f"This inconsistency may result in the learning rate scheduler not functioning properly." + ) + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if accelerator.is_main_process: + tracker_name = "dreambooth-flux3-image2img-lora" + args_cp = vars(args).copy() + args_cp["text_encoder_out_layers"] = str(args_cp["text_encoder_out_layers"]) + accelerator.init_trackers(tracker_name, config=args_cp) + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + + logger.info("***** Running training *****") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num batches each epoch = {len(train_dataloader)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + global_step = 0 + first_epoch = 0 + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the mos recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(os.path.join(args.output_dir, path)) + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + first_epoch = global_step // num_update_steps_per_epoch + + else: + initial_global_step = 0 + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + + def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): + sigmas = noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype) + schedule_timesteps = noise_scheduler_copy.timesteps.to(accelerator.device) + timesteps = timesteps.to(accelerator.device) + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + for epoch in range(first_epoch, args.num_train_epochs): + transformer.train() + + for step, batch in enumerate(train_dataloader): + models_to_accumulate = [transformer] + prompts = batch["prompts"] + + with accelerator.accumulate(models_to_accumulate): + if train_dataset.custom_instance_prompts: + prompt_embeds = prompt_embeds_cache[step] + text_ids = text_ids_cache[step] + else: + num_repeat_elements = len(prompts) + prompt_embeds = prompt_embeds.repeat(num_repeat_elements, 1, 1) + text_ids = text_ids.repeat(num_repeat_elements, 1, 1) + + # Caption dropout: replace the instruction with the empty prompt for a fraction of samples so they + # train unconditionally on the text (the conditioning image is kept). Clone first so the in-place + # write never mutates the cached embeddings. + if args.caption_dropout > 0: + drop_mask = torch.rand(prompt_embeds.shape[0], device=prompt_embeds.device) < args.caption_dropout + if drop_mask.any(): + empty_hidden_states, empty_text_ids = empty_prompt_by_window[prompt_embeds.shape[1]] + prompt_embeds = prompt_embeds.clone() + text_ids = text_ids.clone() + prompt_embeds[drop_mask] = empty_hidden_states.to( + device=prompt_embeds.device, dtype=prompt_embeds.dtype + ) + text_ids[drop_mask] = empty_text_ids.to(device=text_ids.device, dtype=text_ids.dtype) + + # Convert images to latent space + if args.cache_latents: + model_input = latents_cache[step].mode() + cond_model_input = cond_latents_cache[step].mode() + else: + with offload_models(vae, device=accelerator.device, offload=args.offload): + pixel_values = batch["pixel_values"].to(device=accelerator.device, dtype=vae.dtype) + cond_pixel_values = batch["cond_pixel_values"].to(device=accelerator.device, dtype=vae.dtype) + model_input = vae.encode(pixel_values).latent_dist.mode() + cond_model_input = vae.encode(cond_pixel_values).latent_dist.mode() + + model_input = patchify_latents(model_input) + model_input = (model_input - latents_bn_mean) / latents_bn_std + + cond_model_input = patchify_latents(cond_model_input) + cond_model_input = (cond_model_input - latents_bn_mean) / latents_bn_std + + _, _, lh, lw = model_input.shape + _, _, ch, cw = cond_model_input.shape + + # Sample noise that we'll add to the latents + noise = torch.randn_like(model_input) + bsz = model_input.shape[0] + + # Sample a random timestep for each image + # for weighting schemes where we sample timesteps non-uniformly + u = compute_density_for_timestep_sampling( + weighting_scheme=args.weighting_scheme, + batch_size=bsz, + logit_mean=args.logit_mean, + logit_std=args.logit_std, + mode_scale=args.mode_scale, + ) + indices = (u * noise_scheduler_copy.config.num_train_timesteps).long() + timesteps = noise_scheduler_copy.timesteps[indices].to(device=model_input.device) + + # Add noise according to flow matching. + # zt = (1 - texp) * x + texp * z1 + sigmas = get_sigmas(timesteps, n_dim=model_input.ndim, dtype=model_input.dtype) + noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise + + # [B, C, H, W] -> [B, H*W, C]: the noisy generated-image tokens. + packed_noisy_model_input = pack_latents(noisy_model_input) + image_ids = prepare_latent_ids(lh, lw, accelerator.device) # (image_seq_len, 4) + # FLUX.3 is not guidance-distilled: no guidance argument. The flow timestep is a per-sample scalar in + # [0, 1] (the transformer scales it by 1000 and broadcasts it over the image tokens internally). + timestep = sigmas.reshape(bsz) + + # The clean reference image is fed as a SEPARATE conditioning stream: its own packed tokens and + # position ids at the l=-1 axis. The transformer holds the reference (and the text) at timestep 0 + # internally, so no cond timestep is passed. It does not concatenate the reference onto the image + # stream and returns only the generated image tokens, so there is nothing to slice off afterwards. + image_cond_hidden_states = pack_latents(cond_model_input) + image_cond_ids = prepare_image_cond_ids(ch, cw, accelerator.device) # (cond_seq_len, 4) + + # Predict the noise residual. The transformer returns only the generated image tokens, so there is + # nothing to slice/unpack afterwards. + with accelerator.autocast(): + model_pred = transformer( + hidden_states=packed_noisy_model_input, # (B, image_seq_len, C) + img_ids=image_ids, # (image_seq_len, 4) + timestep=timestep, # (B,) + image_cond_hidden_states=image_cond_hidden_states, # (B, cond_seq_len, C) + image_cond_ids=image_cond_ids, # (cond_seq_len, 4) + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids[0], # (text_seq_len, 4), shared across the batch + return_dict=False, + )[0] + + # these weighting schemes use a uniform timestep sampling + # and instead post-weight the loss + weighting = compute_loss_weighting_for_sd3( + weighting_scheme=args.weighting_scheme, sigmas=sigmas.reshape(bsz, 1, 1) + ) + + # flow matching loss (velocity target) on the packed image tokens + target = pack_latents(noise) - pack_latents(model_input) + + # Compute regular loss. + loss = torch.mean( + (weighting.float() * (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1), + 1, + ) + loss = loss.mean() + + accelerator.backward(loss) + if accelerator.sync_gradients: + params_to_clip = transformer.parameters() + accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) + + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # Checks if the accelerator has performed an optimization step behind the scenes + if accelerator.sync_gradients: + progress_bar.update(1) + global_step += 1 + + if accelerator.is_main_process or is_fsdp: + if global_step % args.checkpointing_steps == 0: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + accelerator.log(logs, step=global_step) + + if global_step >= args.max_train_steps: + break + + if accelerator.is_main_process: + if args.validation_prompt is not None and epoch % args.validation_epochs == 0: + transformer_ = unwrap_model(transformer) + transformer_.eval() + pipeline = make_validation_pipeline(args, transformer_, weight_dtype, accelerator.device) + images = log_validation( + pipeline=pipeline, + args=args, + accelerator=accelerator, + pipeline_args=validation_pipeline_args, + epoch=epoch, + torch_dtype=weight_dtype, + ) + transformer_.train() + free_memory() + + # Save the lora layers + accelerator.wait_for_everyone() + + if is_fsdp: + transformer = unwrap_model(transformer) + state_dict = accelerator.get_state_dict(transformer) + if accelerator.is_main_process: + modules_to_save = {} + if is_fsdp: + if args.bnb_quantization_config_path is None: + if args.upcast_before_saving: + state_dict = { + k: v.to(torch.float32) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items() + } + else: + state_dict = { + k: v.to(weight_dtype) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items() + } + + transformer_lora_layers = get_peft_model_state_dict( + transformer, + state_dict=state_dict, + ) + transformer_lora_layers = { + k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v + for k, v in transformer_lora_layers.items() + } + + else: + transformer = unwrap_model(transformer) + if args.bnb_quantization_config_path is None: + if args.upcast_before_saving: + transformer.to(torch.float32) + else: + transformer = transformer.to(weight_dtype) + transformer_lora_layers = get_peft_model_state_dict(transformer) + + modules_to_save["transformer"] = transformer + + Flux3LoraLoaderMixin.save_lora_weights( + save_directory=args.output_dir, + transformer_lora_layers=transformer_lora_layers, + **_collate_lora_metadata(modules_to_save), + ) + + images = [] + run_validation = (args.validation_prompt and args.num_validation_images > 0) or (args.final_validation_prompt) + should_run_final_inference = not args.skip_final_inference and run_validation + if should_run_final_inference: + pipeline = make_validation_pipeline(args, transformer, weight_dtype, accelerator.device) + + # run inference + images = [] + if args.validation_prompt and args.num_validation_images > 0: + images = log_validation( + pipeline=pipeline, + args=args, + accelerator=accelerator, + pipeline_args=validation_pipeline_args, + epoch=epoch, + is_final_validation=True, + torch_dtype=weight_dtype, + ) + del pipeline + free_memory() + + validation_prompt = args.validation_prompt if args.validation_prompt else args.final_validation_prompt + quant_training = None + if args.do_fp8_training: + quant_training = "FP8 TorchAO" + elif args.bnb_quantization_config_path: + quant_training = "BitsandBytes" + save_model_card( + (args.hub_model_id or Path(args.output_dir).name) if not args.push_to_hub else repo_id, + images=images, + base_model=args.pretrained_model_name_or_path, + instance_prompt=args.instance_prompt, + validation_prompt=validation_prompt, + repo_folder=args.output_dir, + quant_training=quant_training, + ) + + if args.push_to_hub: + upload_folder( + repo_id=repo_id, + folder_path=args.output_dir, + commit_message="End of training", + ignore_patterns=["step_*", "epoch_*"], + ) + + accelerator.end_training() + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/scripts/convert_flux3_to_diffusers.py b/scripts/convert_flux3_to_diffusers.py new file mode 100644 index 0000000000..f57763aa87 --- /dev/null +++ b/scripts/convert_flux3_to_diffusers.py @@ -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 \ + --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) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index f6ec5c1653..4f6790e5e4 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -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, diff --git a/src/diffusers/loaders/__init__.py b/src/diffusers/loaders/__init__.py index 5edfa6c01d..72a24ba380 100644 --- a/src/diffusers/loaders/__init__.py +++ b/src/diffusers/loaders/__init__.py @@ -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, diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index bfaec3b0b5..e3e9359168 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -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: diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 46867b2257..22d596c6c6 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -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`]. diff --git a/src/diffusers/loaders/single_file_model.py b/src/diffusers/loaders/single_file_model.py index 1556227673..610e377d3a 100644 --- a/src/diffusers/loaders/single_file_model.py +++ b/src/diffusers/loaders/single_file_model.py @@ -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 diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index 3e160be02b..e65e0be50f 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -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.", diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index cbe283b739..9e27e27e39 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -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, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 079c4a37a2..1076b2b33f 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -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 diff --git a/src/diffusers/models/transformers/transformer_flux3.py b/src/diffusers/models/transformers/transformer_flux3.py new file mode 100644 index 0000000000..2c5fe7a40b --- /dev/null +++ b/src/diffusers/models/transformers/transformer_flux3.py @@ -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) diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 85fef3a1e0..48ccfcc981 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -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, diff --git a/src/diffusers/modular_pipelines/flux3/__init__.py b/src/diffusers/modular_pipelines/flux3/__init__.py new file mode 100644 index 0000000000..2eeae56207 --- /dev/null +++ b/src/diffusers/modular_pipelines/flux3/__init__.py @@ -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) diff --git a/src/diffusers/modular_pipelines/flux3/before_denoise.py b/src/diffusers/modular_pipelines/flux3/before_denoise.py new file mode 100644 index 0000000000..bf6428dc58 --- /dev/null +++ b/src/diffusers/modular_pipelines/flux3/before_denoise.py @@ -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 diff --git a/src/diffusers/modular_pipelines/flux3/decoders.py b/src/diffusers/modular_pipelines/flux3/decoders.py new file mode 100644 index 0000000000..32f3d5a72e --- /dev/null +++ b/src/diffusers/modular_pipelines/flux3/decoders.py @@ -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 diff --git a/src/diffusers/modular_pipelines/flux3/denoise.py b/src/diffusers/modular_pipelines/flux3/denoise.py new file mode 100644 index 0000000000..ed91f8569c --- /dev/null +++ b/src/diffusers/modular_pipelines/flux3/denoise.py @@ -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." + ) diff --git a/src/diffusers/modular_pipelines/flux3/encoders.py b/src/diffusers/modular_pipelines/flux3/encoders.py new file mode 100644 index 0000000000..18c6a454a6 --- /dev/null +++ b/src/diffusers/modular_pipelines/flux3/encoders.py @@ -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 diff --git a/src/diffusers/modular_pipelines/flux3/image_processor.py b/src/diffusers/modular_pipelines/flux3/image_processor.py new file mode 100644 index 0000000000..08b8df17fd --- /dev/null +++ b/src/diffusers/modular_pipelines/flux3/image_processor.py @@ -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)) diff --git a/src/diffusers/modular_pipelines/flux3/inputs.py b/src/diffusers/modular_pipelines/flux3/inputs.py new file mode 100644 index 0000000000..f4ab66553a --- /dev/null +++ b/src/diffusers/modular_pipelines/flux3/inputs.py @@ -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 diff --git a/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py new file mode 100644 index 0000000000..2c7889a5cd --- /dev/null +++ b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py @@ -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"), + ] diff --git a/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py new file mode 100644 index 0000000000..aa08bf6ae2 --- /dev/null +++ b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py @@ -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"), + ] diff --git a/src/diffusers/modular_pipelines/flux3/modular_pipeline.py b/src/diffusers/modular_pipelines/flux3/modular_pipeline.py new file mode 100644 index 0000000000..fbd29f24ac --- /dev/null +++ b/src/diffusers/modular_pipelines/flux3/modular_pipeline.py @@ -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" diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 337d4ea2a1..28e485eddb 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -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")), diff --git a/src/diffusers/quantizers/auto.py b/src/diffusers/quantizers/auto.py index ea6caf91ab..377b46b7b3 100644 --- a/src/diffusers/quantizers/auto.py +++ b/src/diffusers/quantizers/auto.py @@ -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, diff --git a/src/diffusers/quantizers/bfl/__init__.py b/src/diffusers/quantizers/bfl/__init__.py new file mode 100644 index 0000000000..d5294d03a4 --- /dev/null +++ b/src/diffusers/quantizers/bfl/__init__.py @@ -0,0 +1 @@ +from .bfl_quantizer import BFLQuantizer diff --git a/src/diffusers/quantizers/bfl/bfl_quantizer.py b/src/diffusers/quantizers/bfl/bfl_quantizer.py new file mode 100644 index 0000000000..8802144fa5 --- /dev/null +++ b/src/diffusers/quantizers/bfl/bfl_quantizer.py @@ -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) diff --git a/src/diffusers/quantizers/bfl/utils.py b/src/diffusers/quantizers/bfl/utils.py new file mode 100644 index 0000000000..e75b79c4ee --- /dev/null +++ b/src/diffusers/quantizers/bfl/utils.py @@ -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 diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index 30e89f53f9..5777ab45fa 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -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. diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 5f89eb318d..c25edb7f7c 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -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, diff --git a/src/diffusers/utils/dummy_bfl_objects.py b/src/diffusers/utils/dummy_bfl_objects.py new file mode 100644 index 0000000000..77fd7880e9 --- /dev/null +++ b/src/diffusers/utils/dummy_bfl_objects.py @@ -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"]) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 9e1f5f027c..84fd853d17 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -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"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 212626998b..c066286de2 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -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"] diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index af765021f4..c0752f36d5 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -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 diff --git a/tests/conftest.py b/tests/conftest.py index d0401281e3..eaebf335c9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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") diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index a630a2c8b2..dabdd518e9 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -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", diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 65f904bd49..66b9b66c32 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -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]) diff --git a/tests/models/transformers/test_models_transformer_flux3.py b/tests/models/transformers/test_models_transformer_flux3.py new file mode 100644 index 0000000000..66f9f81496 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_flux3.py @@ -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.""" diff --git a/tests/modular_pipelines/flux3/__init__.py b/tests/modular_pipelines/flux3/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py b/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py new file mode 100644 index 0000000000..4787485e19 --- /dev/null +++ b/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py @@ -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 diff --git a/tests/testing_utils.py b/tests/testing_utils.py index a22da6de08..cda114cbde 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -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: