Upload 45 files

.gitattributes

@@ -33,3 +33,11 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/flux/assets/cup.png filter=lfs diff=lfs merge=lfs -text
+data/flux/assets/dev_grid.jpg filter=lfs diff=lfs merge=lfs -text
+data/flux/assets/docs/canny.png filter=lfs diff=lfs merge=lfs -text
+data/flux/assets/docs/depth.png filter=lfs diff=lfs merge=lfs -text
+data/flux/assets/docs/inpainting.png filter=lfs diff=lfs merge=lfs -text
+data/flux/assets/docs/outpainting.png filter=lfs diff=lfs merge=lfs -text
+data/flux/assets/grid.jpg filter=lfs diff=lfs merge=lfs -text
+data/flux/assets/schnell_grid.jpg filter=lfs diff=lfs merge=lfs -text

data/README.md

@@ -0,0 +1,87 @@
+# FLUX
+by Black Forest Labs: https://blackforestlabs.ai. Documentation for our API can be found here: [docs.bfl.ml](https://docs.bfl.ml/).
+
+![grid](assets/grid.jpg)
+
+This repo contains minimal inference code to run image generation & editing with our Flux models.
+
+## Local installation
+
+```bash
+cd $HOME && git clone https://github.com/black-forest-labs/flux
+cd $HOME/flux
+python3.10 -m venv .venv
+source .venv/bin/activate
+pip install -e ".[all]"
+```
+
+### Models
+
+We are offering an extensive suite of models. For more information about the invidual models, please refer to the link under **Usage**.
+
+| Name                        | Usage                                                      | HuggingFace repo                                               | License                                                               |
+| --------------------------- | ---------------------------------------------------------- |  ------------------------------------------------------------- | --------------------------------------------------------------------- |
+| `FLUX.1 [schnell]`          | [Text to Image](docs/text-to-image.md)                     | https://huggingface.co/black-forest-labs/FLUX.1-schnell        | [apache-2.0](model_licenses/LICENSE-FLUX1-schnell)                    |
+| `FLUX.1 [dev]`              | [Text to Image](docs/text-to-image.md)                     | https://huggingface.co/black-forest-labs/FLUX.1-dev            | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) |
+| `FLUX.1 Fill [dev]`         | [In/Out-painting](docs/fill.md)                            | https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev       | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) |
+| `FLUX.1 Canny [dev]`        | [Structural Conditioning](docs/structural-conditioning.md) | https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev      | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) |
+| `FLUX.1 Depth [dev]`        | [Structural Conditioning](docs/structural-conditioning.md) | https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev      | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) |
+| `FLUX.1 Canny [dev] LoRA`   | [Structural Conditioning](docs/structural-conditioning.md) | https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev-lora | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) |
+| `FLUX.1 Depth [dev] LoRA`   | [Structural Conditioning](docs/structural-conditioning.md) | https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev-lora | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) |
+| `FLUX.1 Redux [dev]`        | [Image variation](docs/image-variation.md)                 | https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev      | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) |
+| `FLUX.1 [pro]`              | [Text to Image](docs/text-to-image.md)                     | [Available in our API.](https://docs.bfl.ml/)             |
+| `FLUX1.1 [pro]`             | [Text to Image](docs/text-to-image.md)                     | [Available in our API.](https://docs.bfl.ml/)             |
+| `FLUX1.1 [pro] Ultra/raw`   | [Text to Image](docs/text-to-image.md)                     | [Available in our API.](https://docs.bfl.ml/)             |
+| `FLUX.1 Fill [pro]`         | [In/Out-painting](docs/fill.md)                            | [Available in our API.](https://docs.bfl.ml/)             |
+| `FLUX.1 Canny [pro]`        | [Structural Conditioning](docs/controlnet.md)              | [Available in our API.](https://docs.bfl.ml/)             |
+| `FLUX.1 Depth [pro]`        | [Structural Conditioning](docs/controlnet.md)              | [Available in our API.](https://docs.bfl.ml/)             |
+| `FLUX1.1 Redux [pro]`       | [Image variation](docs/image-variation.md)                 | [Available in our API.](https://docs.bfl.ml/)             |
+| `FLUX1.1 Redux [pro] Ultra` | [Image variation](docs/image-variation.md)                 | [Available in our API.](https://docs.bfl.ml/)             |
+
+The weights of the autoencoder are also released under [apache-2.0](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/apache-2.0.md) and can be found in the HuggingFace repos above.
+
+## API usage
+
+Our API offers access to our models. It is documented here:
+[docs.bfl.ml](https://docs.bfl.ml/).
+
+In this repository we also offer an easy python interface. To use this, you
+first need to register with the API on [api.bfl.ml](https://api.bfl.ml/), and
+create a new API key.
+
+To use the API key either run `export BFL_API_KEY=<your_key_here>` or provide
+it via the `api_key=<your_key_here>` parameter. It is also expected that you
+have installed the package as above.
+
+Usage from python:
+
+```python
+from flux.api import ImageRequest
+
+# this will create an api request directly but not block until the generation is finished
+request = ImageRequest("A beautiful beach", name="flux.1.1-pro")
+# or: request = ImageRequest("A beautiful beach", name="flux.1.1-pro", api_key="your_key_here")
+
+# any of the following will block until the generation is finished
+request.url
+# -> https:<...>/sample.jpg
+request.bytes
+# -> b"..." bytes for the generated image
+request.save("outputs/api.jpg")
+# saves the sample to local storage
+request.image
+# -> a PIL image
+```
+
+Usage from the command line:
+
+```bash
+$ python -m flux.api --prompt="A beautiful beach" url
+https:<...>/sample.jpg
+
+# generate and save the result
+$ python -m flux.api --prompt="A beautiful beach" save outputs/api
+
+# open the image directly
+$ python -m flux.api --prompt="A beautiful beach" image show
+```

data/flux/.github/workflows/ci.yaml

@@ -0,0 +1,20 @@
+name: CI
+on: push
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions/setup-python@v2
+        with:
+          python-version: "3.10"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install ruff==0.6.8
+      - name: Run Ruff
+        run: ruff check --output-format=github .
+      - name: Check imports
+        run: ruff check --select I --output-format=github .
+      - name: Check formatting
+        run: ruff format --check .

data/flux/.gitignore

@@ -0,0 +1,230 @@
+# Created by https://www.toptal.com/developers/gitignore/api/linux,windows,macos,visualstudiocode,python
+# Edit at https://www.toptal.com/developers/gitignore?templates=linux,windows,macos,visualstudiocode,python
+
+### Linux ###
+*~
+
+# temporary files which can be created if a process still has a handle open of a deleted file
+.fuse_hidden*
+
+# KDE directory preferences
+.directory
+
+# Linux trash folder which might appear on any partition or disk
+.Trash-*
+
+# .nfs files are created when an open file is removed but is still being accessed
+.nfs*
+
+### macOS ###
+# General
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# Icon must end with two \r
+Icon
+
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+### VisualStudioCode ###
+.vscode/*
+!.vscode/settings.json
+!.vscode/tasks.json
+!.vscode/launch.json
+!.vscode/extensions.json
+*.code-workspace
+
+# Local History for Visual Studio Code
+.history/
+
+### VisualStudioCode Patch ###
+# Ignore all local history of files
+.history
+.ionide
+
+### Windows ###
+# Windows thumbnail cache files
+Thumbs.db
+Thumbs.db:encryptable
+ehthumbs.db
+ehthumbs_vista.db
+
+# Dump file
+*.stackdump
+
+# Folder config file
+[Dd]esktop.ini
+
+# Recycle Bin used on file shares
+$RECYCLE.BIN/
+
+# Windows Installer files
+*.cab
+*.msi
+*.msix
+*.msm
+*.msp
+
+# Windows shortcuts
+*.lnk
+
+# End of https://www.toptal.com/developers/gitignore/api/linux,windows,macos,visualstudiocode,python

data/flux/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

data/flux/README.md

@@ -0,0 +1,87 @@
+# FLUX
+by Black Forest Labs: https://blackforestlabs.ai. Documentation for our API can be found here: [docs.bfl.ml](https://docs.bfl.ml/).
+
+![grid](assets/grid.jpg)
+
+This repo contains minimal inference code to run image generation & editing with our Flux models.
+
+## Local installation
+
+```bash
+cd $HOME && git clone https://github.com/black-forest-labs/flux
+cd $HOME/flux
+python3.10 -m venv .venv
+source .venv/bin/activate
+pip install -e ".[all]"
+```
+
+### Models
+
+We are offering an extensive suite of models. For more information about the invidual models, please refer to the link under **Usage**.
+
+| Name                        | Usage                                                      | HuggingFace repo                                               | License                                                               |
+| --------------------------- | ---------------------------------------------------------- |  ------------------------------------------------------------- | --------------------------------------------------------------------- |
+| `FLUX.1 [schnell]`          | [Text to Image](docs/text-to-image.md)                     | https://huggingface.co/black-forest-labs/FLUX.1-schnell        | [apache-2.0](model_licenses/LICENSE-FLUX1-schnell)                    |
+| `FLUX.1 [dev]`              | [Text to Image](docs/text-to-image.md)                     | https://huggingface.co/black-forest-labs/FLUX.1-dev            | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) |
+| `FLUX.1 Fill [dev]`         | [In/Out-painting](docs/fill.md)                            | https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev       | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) |
+| `FLUX.1 Canny [dev]`        | [Structural Conditioning](docs/structural-conditioning.md) | https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev      | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) |
+| `FLUX.1 Depth [dev]`        | [Structural Conditioning](docs/structural-conditioning.md) | https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev      | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) |
+| `FLUX.1 Canny [dev] LoRA`   | [Structural Conditioning](docs/structural-conditioning.md) | https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev-lora | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) |
+| `FLUX.1 Depth [dev] LoRA`   | [Structural Conditioning](docs/structural-conditioning.md) | https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev-lora | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) |
+| `FLUX.1 Redux [dev]`        | [Image variation](docs/image-variation.md)                 | https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev      | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) |
+| `FLUX.1 [pro]`              | [Text to Image](docs/text-to-image.md)                     | [Available in our API.](https://docs.bfl.ml/)             |
+| `FLUX1.1 [pro]`             | [Text to Image](docs/text-to-image.md)                     | [Available in our API.](https://docs.bfl.ml/)             |
+| `FLUX1.1 [pro] Ultra/raw`   | [Text to Image](docs/text-to-image.md)                     | [Available in our API.](https://docs.bfl.ml/)             |
+| `FLUX.1 Fill [pro]`         | [In/Out-painting](docs/fill.md)                            | [Available in our API.](https://docs.bfl.ml/)             |
+| `FLUX.1 Canny [pro]`        | [Structural Conditioning](docs/controlnet.md)              | [Available in our API.](https://docs.bfl.ml/)             |
+| `FLUX.1 Depth [pro]`        | [Structural Conditioning](docs/controlnet.md)              | [Available in our API.](https://docs.bfl.ml/)             |
+| `FLUX1.1 Redux [pro]`       | [Image variation](docs/image-variation.md)                 | [Available in our API.](https://docs.bfl.ml/)             |
+| `FLUX1.1 Redux [pro] Ultra` | [Image variation](docs/image-variation.md)                 | [Available in our API.](https://docs.bfl.ml/)             |
+
+The weights of the autoencoder are also released under [apache-2.0](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/apache-2.0.md) and can be found in the HuggingFace repos above.
+
+## API usage
+
+Our API offers access to our models. It is documented here:
+[docs.bfl.ml](https://docs.bfl.ml/).
+
+In this repository we also offer an easy python interface. To use this, you
+first need to register with the API on [api.bfl.ml](https://api.bfl.ml/), and
+create a new API key.
+
+To use the API key either run `export BFL_API_KEY=<your_key_here>` or provide
+it via the `api_key=<your_key_here>` parameter. It is also expected that you
+have installed the package as above.
+
+Usage from python:
+
+```python
+from flux.api import ImageRequest
+
+# this will create an api request directly but not block until the generation is finished
+request = ImageRequest("A beautiful beach", name="flux.1.1-pro")
+# or: request = ImageRequest("A beautiful beach", name="flux.1.1-pro", api_key="your_key_here")
+
+# any of the following will block until the generation is finished
+request.url
+# -> https:<...>/sample.jpg
+request.bytes
+# -> b"..." bytes for the generated image
+request.save("outputs/api.jpg")
+# saves the sample to local storage
+request.image
+# -> a PIL image
+```
+
+Usage from the command line:
+
+```bash
+$ python -m flux.api --prompt="A beautiful beach" url
+https:<...>/sample.jpg
+
+# generate and save the result
+$ python -m flux.api --prompt="A beautiful beach" save outputs/api
+
+# open the image directly
+$ python -m flux.api --prompt="A beautiful beach" image show
+```

data/flux/demo_gr.py

@@ -0,0 +1,247 @@
+import os
+import time
+import uuid
+
+import gradio as gr
+import numpy as np
+import torch
+from einops import rearrange
+from PIL import ExifTags, Image
+from transformers import pipeline
+
+from flux.cli import SamplingOptions
+from flux.sampling import denoise, get_noise, get_schedule, prepare, unpack
+from flux.util import configs, embed_watermark, load_ae, load_clip, load_flow_model, load_t5
+
+NSFW_THRESHOLD = 0.85
+
+
+def get_models(name: str, device: torch.device, offload: bool, is_schnell: bool):
+    t5 = load_t5(device, max_length=256 if is_schnell else 512)
+    clip = load_clip(device)
+    model = load_flow_model(name, device="cpu" if offload else device)
+    ae = load_ae(name, device="cpu" if offload else device)
+    nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=device)
+    return model, ae, t5, clip, nsfw_classifier
+
+
+class FluxGenerator:
+    def __init__(self, model_name: str, device: str, offload: bool):
+        self.device = torch.device(device)
+        self.offload = offload
+        self.model_name = model_name
+        self.is_schnell = model_name == "flux-schnell"
+        self.model, self.ae, self.t5, self.clip, self.nsfw_classifier = get_models(
+            model_name,
+            device=self.device,
+            offload=self.offload,
+            is_schnell=self.is_schnell,
+        )
+
+    @torch.inference_mode()
+    def generate_image(
+        self,
+        width,
+        height,
+        num_steps,
+        guidance,
+        seed,
+        prompt,
+        init_image=None,
+        image2image_strength=0.0,
+        add_sampling_metadata=True,
+    ):
+        seed = int(seed)
+        if seed == -1:
+            seed = None
+
+        opts = SamplingOptions(
+            prompt=prompt,
+            width=width,
+            height=height,
+            num_steps=num_steps,
+            guidance=guidance,
+            seed=seed,
+        )
+
+        if opts.seed is None:
+            opts.seed = torch.Generator(device="cpu").seed()
+        print(f"Generating '{opts.prompt}' with seed {opts.seed}")
+        t0 = time.perf_counter()
+
+        if init_image is not None:
+            if isinstance(init_image, np.ndarray):
+                init_image = torch.from_numpy(init_image).permute(2, 0, 1).float() / 255.0
+                init_image = init_image.unsqueeze(0)
+            init_image = init_image.to(self.device)
+            init_image = torch.nn.functional.interpolate(init_image, (opts.height, opts.width))
+            if self.offload:
+                self.ae.encoder.to(self.device)
+            init_image = self.ae.encode(init_image.to())
+            if self.offload:
+                self.ae = self.ae.cpu()
+                torch.cuda.empty_cache()
+
+        # prepare input
+        x = get_noise(
+            1,
+            opts.height,
+            opts.width,
+            device=self.device,
+            dtype=torch.bfloat16,
+            seed=opts.seed,
+        )
+        timesteps = get_schedule(
+            opts.num_steps,
+            x.shape[-1] * x.shape[-2] // 4,
+            shift=(not self.is_schnell),
+        )
+        if init_image is not None:
+            t_idx = int((1 - image2image_strength) * num_steps)
+            t = timesteps[t_idx]
+            timesteps = timesteps[t_idx:]
+            x = t * x + (1.0 - t) * init_image.to(x.dtype)
+
+        if self.offload:
+            self.t5, self.clip = self.t5.to(self.device), self.clip.to(self.device)
+        inp = prepare(t5=self.t5, clip=self.clip, img=x, prompt=opts.prompt)
+
+        # offload TEs to CPU, load model to gpu
+        if self.offload:
+            self.t5, self.clip = self.t5.cpu(), self.clip.cpu()
+            torch.cuda.empty_cache()
+            self.model = self.model.to(self.device)
+
+        # denoise initial noise
+        x = denoise(self.model, **inp, timesteps=timesteps, guidance=opts.guidance)
+
+        # offload model, load autoencoder to gpu
+        if self.offload:
+            self.model.cpu()
+            torch.cuda.empty_cache()
+            self.ae.decoder.to(x.device)
+
+        # decode latents to pixel space
+        x = unpack(x.float(), opts.height, opts.width)
+        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16):
+            x = self.ae.decode(x)
+
+        if self.offload:
+            self.ae.decoder.cpu()
+            torch.cuda.empty_cache()
+
+        t1 = time.perf_counter()
+
+        print(f"Done in {t1 - t0:.1f}s.")
+        # bring into PIL format
+        x = x.clamp(-1, 1)
+        x = embed_watermark(x.float())
+        x = rearrange(x[0], "c h w -> h w c")
+
+        img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
+        nsfw_score = [x["score"] for x in self.nsfw_classifier(img) if x["label"] == "nsfw"][0]
+
+        if nsfw_score < NSFW_THRESHOLD:
+            filename = f"output/gradio/{uuid.uuid4()}.jpg"
+            os.makedirs(os.path.dirname(filename), exist_ok=True)
+            exif_data = Image.Exif()
+            if init_image is None:
+                exif_data[ExifTags.Base.Software] = "AI generated;txt2img;flux"
+            else:
+                exif_data[ExifTags.Base.Software] = "AI generated;img2img;flux"
+            exif_data[ExifTags.Base.Make] = "Black Forest Labs"
+            exif_data[ExifTags.Base.Model] = self.model_name
+            if add_sampling_metadata:
+                exif_data[ExifTags.Base.ImageDescription] = prompt
+
+            img.save(filename, format="jpeg", exif=exif_data, quality=95, subsampling=0)
+
+            return img, str(opts.seed), filename, None
+        else:
+            return None, str(opts.seed), None, "Your generated image may contain NSFW content."
+
+
+def create_demo(
+    model_name: str, device: str = "cuda" if torch.cuda.is_available() else "cpu", offload: bool = False
+):
+    generator = FluxGenerator(model_name, device, offload)
+    is_schnell = model_name == "flux-schnell"
+
+    with gr.Blocks() as demo:
+        gr.Markdown(f"# Flux Image Generation Demo - Model: {model_name}")
+
+        with gr.Row():
+            with gr.Column():
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    value='a photo of a forest with mist swirling around the tree trunks. The word "FLUX" is painted over it in big, red brush strokes with visible texture',
+                )
+                do_img2img = gr.Checkbox(label="Image to Image", value=False, interactive=not is_schnell)
+                init_image = gr.Image(label="Input Image", visible=False)
+                image2image_strength = gr.Slider(
+                    0.0, 1.0, 0.8, step=0.1, label="Noising strength", visible=False
+                )
+
+                with gr.Accordion("Advanced Options", open=False):
+                    width = gr.Slider(128, 8192, 1360, step=16, label="Width")
+                    height = gr.Slider(128, 8192, 768, step=16, label="Height")
+                    num_steps = gr.Slider(1, 50, 4 if is_schnell else 50, step=1, label="Number of steps")
+                    guidance = gr.Slider(
+                        1.0, 10.0, 3.5, step=0.1, label="Guidance", interactive=not is_schnell
+                    )
+                    seed = gr.Textbox(-1, label="Seed (-1 for random)")
+                    add_sampling_metadata = gr.Checkbox(
+                        label="Add sampling parameters to metadata?", value=True
+                    )
+
+                generate_btn = gr.Button("Generate")
+
+            with gr.Column():
+                output_image = gr.Image(label="Generated Image")
+                seed_output = gr.Number(label="Used Seed")
+                warning_text = gr.Textbox(label="Warning", visible=False)
+                download_btn = gr.File(label="Download full-resolution")
+
+        def update_img2img(do_img2img):
+            return {
+                init_image: gr.update(visible=do_img2img),
+                image2image_strength: gr.update(visible=do_img2img),
+            }
+
+        do_img2img.change(update_img2img, do_img2img, [init_image, image2image_strength])
+
+        generate_btn.click(
+            fn=generator.generate_image,
+            inputs=[
+                width,
+                height,
+                num_steps,
+                guidance,
+                seed,
+                prompt,
+                init_image,
+                image2image_strength,
+                add_sampling_metadata,
+            ],
+            outputs=[output_image, seed_output, download_btn, warning_text],
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Flux")
+    parser.add_argument(
+        "--name", type=str, default="flux-schnell", choices=list(configs.keys()), help="Model name"
+    )
+    parser.add_argument(
+        "--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device to use"
+    )
+    parser.add_argument("--offload", action="store_true", help="Offload model to CPU when not in use")
+    parser.add_argument("--share", action="store_true", help="Create a public link to your demo")
+    args = parser.parse_args()
+
+    demo = create_demo(args.name, args.device, args.offload)
+    demo.launch(share=args.share)

data/flux/demo_st.py

@@ -0,0 +1,293 @@
+import os
+import re
+import time
+from glob import iglob
+from io import BytesIO
+
+import streamlit as st
+import torch
+from einops import rearrange
+from fire import Fire
+from PIL import ExifTags, Image
+from st_keyup import st_keyup
+from torchvision import transforms
+from transformers import pipeline
+
+from flux.cli import SamplingOptions
+from flux.sampling import denoise, get_noise, get_schedule, prepare, unpack
+from flux.util import (
+    configs,
+    embed_watermark,
+    load_ae,
+    load_clip,
+    load_flow_model,
+    load_t5,
+)
+
+NSFW_THRESHOLD = 0.85
+
+
+@st.cache_resource()
+def get_models(name: str, device: torch.device, offload: bool, is_schnell: bool):
+    t5 = load_t5(device, max_length=256 if is_schnell else 512)
+    clip = load_clip(device)
+    model = load_flow_model(name, device="cpu" if offload else device)
+    ae = load_ae(name, device="cpu" if offload else device)
+    nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=device)
+    return model, ae, t5, clip, nsfw_classifier
+
+
+def get_image() -> torch.Tensor | None:
+    image = st.file_uploader("Input", type=["jpg", "JPEG", "png"])
+    if image is None:
+        return None
+    image = Image.open(image).convert("RGB")
+
+    transform = transforms.Compose(
+        [
+            transforms.ToTensor(),
+            transforms.Lambda(lambda x: 2.0 * x - 1.0),
+        ]
+    )
+    img: torch.Tensor = transform(image)
+    return img[None, ...]
+
+
+@torch.inference_mode()
+def main(
+    device: str = "cuda" if torch.cuda.is_available() else "cpu",
+    offload: bool = False,
+    output_dir: str = "output",
+):
+    torch_device = torch.device(device)
+    names = list(configs.keys())
+    name = st.selectbox("Which model to load?", names)
+    if name is None or not st.checkbox("Load model", False):
+        return
+
+    is_schnell = name == "flux-schnell"
+    model, ae, t5, clip, nsfw_classifier = get_models(
+        name,
+        device=torch_device,
+        offload=offload,
+        is_schnell=is_schnell,
+    )
+
+    do_img2img = (
+        st.checkbox(
+            "Image to Image",
+            False,
+            disabled=is_schnell,
+            help="Partially noise an image and denoise again to get variations.\n\nOnly works for flux-dev",
+        )
+        and not is_schnell
+    )
+    if do_img2img:
+        init_image = get_image()
+        if init_image is None:
+            st.warning("Please add an image to do image to image")
+        image2image_strength = st.number_input("Noising strength", min_value=0.0, max_value=1.0, value=0.8)
+        if init_image is not None:
+            h, w = init_image.shape[-2:]
+            st.write(f"Got image of size {w}x{h} ({h*w/1e6:.2f}MP)")
+        resize_img = st.checkbox("Resize image", False) or init_image is None
+    else:
+        init_image = None
+        resize_img = True
+        image2image_strength = 0.0
+
+    # allow for packing and conversion to latent space
+    width = int(
+        16 * (st.number_input("Width", min_value=128, value=1360, step=16, disabled=not resize_img) // 16)
+    )
+    height = int(
+        16 * (st.number_input("Height", min_value=128, value=768, step=16, disabled=not resize_img) // 16)
+    )
+    num_steps = int(st.number_input("Number of steps", min_value=1, value=(4 if is_schnell else 50)))
+    guidance = float(st.number_input("Guidance", min_value=1.0, value=3.5, disabled=is_schnell))
+    seed_str = st.text_input("Seed", disabled=is_schnell)
+    if seed_str.isdecimal():
+        seed = int(seed_str)
+    else:
+        st.info("No seed set, set to positive integer to enable")
+        seed = None
+    save_samples = st.checkbox("Save samples?", not is_schnell)
+    add_sampling_metadata = st.checkbox("Add sampling parameters to metadata?", True)
+
+    default_prompt = (
+        "a photo of a forest with mist swirling around the tree trunks. The word "
+        '"FLUX" is painted over it in big, red brush strokes with visible texture'
+    )
+    prompt = st_keyup("Enter a prompt", value=default_prompt, debounce=300, key="interactive_text")
+
+    output_name = os.path.join(output_dir, "img_{idx}.jpg")
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+        idx = 0
+    else:
+        fns = [fn for fn in iglob(output_name.format(idx="*")) if re.search(r"img_[0-9]+\.jpg$", fn)]
+        if len(fns) > 0:
+            idx = max(int(fn.split("_")[-1].split(".")[0]) for fn in fns) + 1
+        else:
+            idx = 0
+
+    rng = torch.Generator(device="cpu")
+
+    if "seed" not in st.session_state:
+        st.session_state.seed = rng.seed()
+
+    def increment_counter():
+        st.session_state.seed += 1
+
+    def decrement_counter():
+        if st.session_state.seed > 0:
+            st.session_state.seed -= 1
+
+    opts = SamplingOptions(
+        prompt=prompt,
+        width=width,
+        height=height,
+        num_steps=num_steps,
+        guidance=guidance,
+        seed=seed,
+    )
+
+    if name == "flux-schnell":
+        cols = st.columns([5, 1, 1, 5])
+        with cols[1]:
+            st.button("↩", on_click=increment_counter)
+        with cols[2]:
+            st.button("↪", on_click=decrement_counter)
+    if is_schnell or st.button("Sample"):
+        if is_schnell:
+            opts.seed = st.session_state.seed
+        elif opts.seed is None:
+            opts.seed = rng.seed()
+        print(f"Generating '{opts.prompt}' with seed {opts.seed}")
+        t0 = time.perf_counter()
+
+        if init_image is not None:
+            if resize_img:
+                init_image = torch.nn.functional.interpolate(init_image, (opts.height, opts.width))
+            else:
+                h, w = init_image.shape[-2:]
+                init_image = init_image[..., : 16 * (h // 16), : 16 * (w // 16)]
+                opts.height = init_image.shape[-2]
+                opts.width = init_image.shape[-1]
+            if offload:
+                ae.encoder.to(torch_device)
+            init_image = ae.encode(init_image.to(torch_device))
+            if offload:
+                ae = ae.cpu()
+                torch.cuda.empty_cache()
+
+        # prepare input
+        x = get_noise(
+            1,
+            opts.height,
+            opts.width,
+            device=torch_device,
+            dtype=torch.bfloat16,
+            seed=opts.seed,
+        )
+        # divide pixel space by 16**2 to account for latent space conversion
+        timesteps = get_schedule(
+            opts.num_steps,
+            (x.shape[-1] * x.shape[-2]) // 4,
+            shift=(not is_schnell),
+        )
+        if init_image is not None:
+            t_idx = int((1 - image2image_strength) * num_steps)
+            t = timesteps[t_idx]
+            timesteps = timesteps[t_idx:]
+            x = t * x + (1.0 - t) * init_image.to(x.dtype)
+
+        if offload:
+            t5, clip = t5.to(torch_device), clip.to(torch_device)
+        inp = prepare(t5=t5, clip=clip, img=x, prompt=opts.prompt)
+
+        # offload TEs to CPU, load model to gpu
+        if offload:
+            t5, clip = t5.cpu(), clip.cpu()
+            torch.cuda.empty_cache()
+            model = model.to(torch_device)
+
+        # denoise initial noise
+        x = denoise(model, **inp, timesteps=timesteps, guidance=opts.guidance)
+
+        # offload model, load autoencoder to gpu
+        if offload:
+            model.cpu()
+            torch.cuda.empty_cache()
+            ae.decoder.to(x.device)
+
+        # decode latents to pixel space
+        x = unpack(x.float(), opts.height, opts.width)
+        with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
+            x = ae.decode(x)
+
+        if offload:
+            ae.decoder.cpu()
+            torch.cuda.empty_cache()
+
+        t1 = time.perf_counter()
+
+        fn = output_name.format(idx=idx)
+        print(f"Done in {t1 - t0:.1f}s.")
+        # bring into PIL format and save
+        x = x.clamp(-1, 1)
+        x = embed_watermark(x.float())
+        x = rearrange(x[0], "c h w -> h w c")
+
+        img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
+        nsfw_score = [x["score"] for x in nsfw_classifier(img) if x["label"] == "nsfw"][0]
+
+        if nsfw_score < NSFW_THRESHOLD:
+            buffer = BytesIO()
+            exif_data = Image.Exif()
+            if init_image is None:
+                exif_data[ExifTags.Base.Software] = "AI generated;txt2img;flux"
+            else:
+                exif_data[ExifTags.Base.Software] = "AI generated;img2img;flux"
+            exif_data[ExifTags.Base.Make] = "Black Forest Labs"
+            exif_data[ExifTags.Base.Model] = name
+            if add_sampling_metadata:
+                exif_data[ExifTags.Base.ImageDescription] = prompt
+            img.save(buffer, format="jpeg", exif=exif_data, quality=95, subsampling=0)
+
+            img_bytes = buffer.getvalue()
+            if save_samples:
+                print(f"Saving {fn}")
+                with open(fn, "wb") as file:
+                    file.write(img_bytes)
+                idx += 1
+
+            st.session_state["samples"] = {
+                "prompt": opts.prompt,
+                "img": img,
+                "seed": opts.seed,
+                "bytes": img_bytes,
+            }
+            opts.seed = None
+        else:
+            st.warning("Your generated image may contain NSFW content.")
+            st.session_state["samples"] = None
+
+    samples = st.session_state.get("samples", None)
+    if samples is not None:
+        st.image(samples["img"], caption=samples["prompt"])
+        st.download_button(
+            "Download full-resolution",
+            samples["bytes"],
+            file_name="generated.jpg",
+            mime="image/jpg",
+        )
+        st.write(f"Seed: {samples['seed']}")
+
+
+def app():
+    Fire(main)
+
+
+if __name__ == "__main__":
+    app()

data/flux/demo_st_fill.py

@@ -0,0 +1,487 @@
+import os
+import re
+import tempfile
+import time
+from glob import iglob
+from io import BytesIO
+
+import numpy as np
+import streamlit as st
+import torch
+from einops import rearrange
+from PIL import ExifTags, Image
+from st_keyup import st_keyup
+from streamlit_drawable_canvas import st_canvas
+from transformers import pipeline
+
+from flux.sampling import denoise, get_noise, get_schedule, prepare_fill, unpack
+from flux.util import embed_watermark, load_ae, load_clip, load_flow_model, load_t5
+
+NSFW_THRESHOLD = 0.85
+
+
+def add_border_and_mask(image, zoom_all=1.0, zoom_left=0, zoom_right=0, zoom_up=0, zoom_down=0, overlap=0):
+    """Adds a black border around the image with individual side control and mask overlap"""
+    orig_width, orig_height = image.size
+
+    # Calculate padding for each side (in pixels)
+    left_pad = int(orig_width * zoom_left)
+    right_pad = int(orig_width * zoom_right)
+    top_pad = int(orig_height * zoom_up)
+    bottom_pad = int(orig_height * zoom_down)
+
+    # Calculate overlap in pixels
+    overlap_left = int(orig_width * overlap)
+    overlap_right = int(orig_width * overlap)
+    overlap_top = int(orig_height * overlap)
+    overlap_bottom = int(orig_height * overlap)
+
+    # If using the all-sides zoom, add it to each side
+    if zoom_all > 1.0:
+        extra_each_side = (zoom_all - 1.0) / 2
+        left_pad += int(orig_width * extra_each_side)
+        right_pad += int(orig_width * extra_each_side)
+        top_pad += int(orig_height * extra_each_side)
+        bottom_pad += int(orig_height * extra_each_side)
+
+    # Calculate new dimensions (ensure they're multiples of 32)
+    new_width = 32 * round((orig_width + left_pad + right_pad) / 32)
+    new_height = 32 * round((orig_height + top_pad + bottom_pad) / 32)
+
+    # Create new image with black border
+    bordered_image = Image.new("RGB", (new_width, new_height), (0, 0, 0))
+    # Paste original image in position
+    paste_x = left_pad
+    paste_y = top_pad
+    bordered_image.paste(image, (paste_x, paste_y))
+
+    # Create mask (white where the border is, black where the original image was)
+    mask = Image.new("L", (new_width, new_height), 255)  # White background
+    # Paste black rectangle with overlap adjustment
+    mask.paste(
+        0,
+        (
+            paste_x + overlap_left,  # Left edge moves right
+            paste_y + overlap_top,  # Top edge moves down
+            paste_x + orig_width - overlap_right,  # Right edge moves left
+            paste_y + orig_height - overlap_bottom,  # Bottom edge moves up
+        ),
+    )
+
+    return bordered_image, mask
+
+
+@st.cache_resource()
+def get_models(name: str, device: torch.device, offload: bool):
+    t5 = load_t5(device, max_length=128)
+    clip = load_clip(device)
+    model = load_flow_model(name, device="cpu" if offload else device)
+    ae = load_ae(name, device="cpu" if offload else device)
+    nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=device)
+    return model, ae, t5, clip, nsfw_classifier
+
+
+def resize(img: Image.Image, min_mp: float = 0.5, max_mp: float = 2.0) -> Image.Image:
+    width, height = img.size
+    mp = (width * height) / 1_000_000  # Current megapixels
+
+    if min_mp <= mp <= max_mp:
+        # Even if MP is in range, ensure dimensions are multiples of 32
+        new_width = int(32 * round(width / 32))
+        new_height = int(32 * round(height / 32))
+        if new_width != width or new_height != height:
+            return img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+        return img
+
+    # Calculate scaling factor
+    if mp < min_mp:
+        scale = (min_mp / mp) ** 0.5
+    else:  # mp > max_mp
+        scale = (max_mp / mp) ** 0.5
+
+    new_width = int(32 * round(width * scale / 32))
+    new_height = int(32 * round(height * scale / 32))
+
+    return img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+
+def clear_canvas_state():
+    """Clear all canvas-related state"""
+    keys_to_clear = ["canvas", "last_image_dims"]
+    for key in keys_to_clear:
+        if key in st.session_state:
+            del st.session_state[key]
+
+
+def set_new_image(img: Image.Image):
+    """Safely set a new image and clear relevant state"""
+    st.session_state["current_image"] = img
+    clear_canvas_state()
+    st.rerun()
+
+
+def downscale_image(img: Image.Image, scale_factor: float) -> Image.Image:
+    """Downscale image by a given factor while maintaining 32-pixel multiple dimensions"""
+    if scale_factor >= 1.0:
+        return img
+
+    width, height = img.size
+    new_width = int(32 * round(width * scale_factor / 32))
+    new_height = int(32 * round(height * scale_factor / 32))
+
+    # Ensure minimum dimensions
+    new_width = max(64, new_width)  # minimum 64 pixels
+    new_height = max(64, new_height)  # minimum 64 pixels
+
+    return img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+
+@torch.inference_mode()
+def main(
+    device: str = "cuda" if torch.cuda.is_available() else "cpu",
+    offload: bool = False,
+    output_dir: str = "output",
+):
+    torch_device = torch.device(device)
+    st.title("Flux Fill: Inpainting & Outpainting")
+
+    # Model selection and loading
+    name = "flux-dev-fill"
+    if not st.checkbox("Load model", False):
+        return
+
+    try:
+        model, ae, t5, clip, nsfw_classifier = get_models(
+            name,
+            device=torch_device,
+            offload=offload,
+        )
+    except Exception as e:
+        st.error(f"Error loading models: {e}")
+        return
+
+    # Mode selection
+    mode = st.radio("Select Mode", ["Inpainting", "Outpainting"])
+
+    # Image handling - either from previous generation or new upload
+    if "input_image" in st.session_state:
+        image = st.session_state["input_image"]
+        del st.session_state["input_image"]
+        set_new_image(image)
+        st.write("Continuing from previous result")
+    else:
+        uploaded_image = st.file_uploader("Upload image", type=["jpg", "jpeg", "png"])
+        if uploaded_image is None:
+            st.warning("Please upload an image")
+            return
+
+        if (
+            "current_image_name" not in st.session_state
+            or st.session_state["current_image_name"] != uploaded_image.name
+        ):
+            try:
+                image = Image.open(uploaded_image).convert("RGB")
+                st.session_state["current_image_name"] = uploaded_image.name
+                set_new_image(image)
+            except Exception as e:
+                st.error(f"Error loading image: {e}")
+                return
+        else:
+            image = st.session_state.get("current_image")
+            if image is None:
+                st.error("Error: Image state is invalid. Please reupload the image.")
+                clear_canvas_state()
+                return
+
+    # Add downscale control
+    with st.expander("Image Size Control"):
+        current_mp = (image.size[0] * image.size[1]) / 1_000_000
+        st.write(f"Current image size: {image.size[0]}x{image.size[1]} ({current_mp:.1f}MP)")
+
+        scale_factor = st.slider(
+            "Downscale Factor",
+            min_value=0.1,
+            max_value=1.0,
+            value=1.0,
+            step=0.1,
+            help="1.0 = original size, 0.5 = half size, etc.",
+        )
+
+        if scale_factor < 1.0 and st.button("Apply Downscaling"):
+            image = downscale_image(image, scale_factor)
+            set_new_image(image)
+            st.rerun()
+
+    # Resize image with validation
+    try:
+        original_mp = (image.size[0] * image.size[1]) / 1_000_000
+        image = resize(image)
+        width, height = image.size
+        current_mp = (width * height) / 1_000_000
+
+        if width % 32 != 0 or height % 32 != 0:
+            st.error("Error: Image dimensions must be multiples of 32")
+            return
+
+        st.write(f"Image dimensions: {width}x{height} pixels")
+        if original_mp != current_mp:
+            st.write(
+                f"Image has been resized from {original_mp:.1f}MP to {current_mp:.1f}MP to stay within bounds (0.5MP - 2MP)"
+            )
+    except Exception as e:
+        st.error(f"Error processing image: {e}")
+        return
+
+    if mode == "Outpainting":
+        # Outpainting controls
+        zoom_all = st.slider("Zoom Out Amount (All Sides)", min_value=1.0, max_value=3.0, value=1.0, step=0.1)
+
+        with st.expander("Advanced Zoom Controls"):
+            st.info("These controls add additional zoom to specific sides")
+            col1, col2 = st.columns(2)
+            with col1:
+                zoom_left = st.slider("Left", min_value=0.0, max_value=1.0, value=0.0, step=0.1)
+                zoom_right = st.slider("Right", min_value=0.0, max_value=1.0, value=0.0, step=0.1)
+            with col2:
+                zoom_up = st.slider("Up", min_value=0.0, max_value=1.0, value=0.0, step=0.1)
+                zoom_down = st.slider("Down", min_value=0.0, max_value=1.0, value=0.0, step=0.1)
+
+        overlap = st.slider("Overlap", min_value=0.01, max_value=0.25, value=0.01, step=0.01)
+
+        # Generate bordered image and mask
+        image_for_generation, mask = add_border_and_mask(
+            image,
+            zoom_all=zoom_all,
+            zoom_left=zoom_left,
+            zoom_right=zoom_right,
+            zoom_up=zoom_up,
+            zoom_down=zoom_down,
+            overlap=overlap,
+        )
+        width, height = image_for_generation.size
+
+        # Show preview
+        col1, col2 = st.columns(2)
+        with col1:
+            st.image(image_for_generation, caption="Image with Border")
+        with col2:
+            st.image(mask, caption="Mask (white areas will be generated)")
+
+    else:  # Inpainting mode
+        # Canvas setup with dimension tracking
+        canvas_key = f"canvas_{width}_{height}"
+        if "last_image_dims" not in st.session_state:
+            st.session_state.last_image_dims = (width, height)
+        elif st.session_state.last_image_dims != (width, height):
+            clear_canvas_state()
+            st.session_state.last_image_dims = (width, height)
+            st.rerun()
+
+        try:
+            canvas_result = st_canvas(
+                fill_color="rgba(255, 255, 255, 0.0)",
+                stroke_width=st.slider("Brush size", 1, 500, 50),
+                stroke_color="#fff",
+                background_image=image,
+                height=height,
+                width=width,
+                drawing_mode="freedraw",
+                key=canvas_key,
+                display_toolbar=True,
+            )
+        except Exception as e:
+            st.error(f"Error creating canvas: {e}")
+            clear_canvas_state()
+            st.rerun()
+            return
+
+    # Sampling parameters
+    num_steps = int(st.number_input("Number of steps", min_value=1, value=50))
+    guidance = float(st.number_input("Guidance", min_value=1.0, value=30.0))
+    seed_str = st.text_input("Seed")
+    if seed_str.isdecimal():
+        seed = int(seed_str)
+    else:
+        st.info("No seed set, using random seed")
+        seed = None
+
+    save_samples = st.checkbox("Save samples?", True)
+    add_sampling_metadata = st.checkbox("Add sampling parameters to metadata?", True)
+
+    # Prompt input
+    prompt = st_keyup("Enter a prompt", value="", debounce=300, key="interactive_text")
+
+    # Setup output path
+    output_name = os.path.join(output_dir, "img_{idx}.jpg")
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+        idx = 0
+    else:
+        fns = [fn for fn in iglob(output_name.format(idx="*")) if re.search(r"img_[0-9]+\.jpg$", fn)]
+        idx = len(fns)
+
+    if st.button("Generate"):
+        valid_input = False
+
+        if mode == "Inpainting" and canvas_result.image_data is not None:
+            valid_input = True
+            # Create mask from canvas
+            try:
+                mask = Image.fromarray(canvas_result.image_data)
+                mask = mask.getchannel("A")  # Get alpha channel
+                mask_array = np.array(mask)
+                mask_array = (mask_array > 0).astype(np.uint8) * 255
+                mask = Image.fromarray(mask_array)
+                image_for_generation = image
+            except Exception as e:
+                st.error(f"Error creating mask: {e}")
+                return
+
+        elif mode == "Outpainting":
+            valid_input = True
+            # image_for_generation and mask are already set above
+
+        if not valid_input:
+            st.error("Please draw a mask or configure outpainting settings")
+            return
+
+        # Create temporary files
+        with (
+            tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_img,
+            tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_mask,
+        ):
+            try:
+                image_for_generation.save(tmp_img.name)
+                mask.save(tmp_mask.name)
+            except Exception as e:
+                st.error(f"Error saving temporary files: {e}")
+                return
+
+            try:
+                # Generate inpainting/outpainting
+                rng = torch.Generator(device="cpu")
+                if seed is None:
+                    seed = rng.seed()
+
+                print(f"Generating with seed {seed}:\n{prompt}")
+                t0 = time.perf_counter()
+
+                x = get_noise(
+                    1,
+                    height,
+                    width,
+                    device=torch_device,
+                    dtype=torch.bfloat16,
+                    seed=seed,
+                )
+
+                if offload:
+                    t5, clip, ae = t5.to(torch_device), clip.to(torch_device), ae.to(torch_device)
+
+                inp = prepare_fill(
+                    t5,
+                    clip,
+                    x,
+                    prompt=prompt,
+                    ae=ae,
+                    img_cond_path=tmp_img.name,
+                    mask_path=tmp_mask.name,
+                )
+
+                timesteps = get_schedule(num_steps, inp["img"].shape[1], shift=True)
+
+                if offload:
+                    t5, clip, ae = t5.cpu(), clip.cpu(), ae.cpu()
+                    torch.cuda.empty_cache()
+                    model = model.to(torch_device)
+
+                x = denoise(model, **inp, timesteps=timesteps, guidance=guidance)
+
+                if offload:
+                    model.cpu()
+                    torch.cuda.empty_cache()
+                    ae.decoder.to(x.device)
+
+                x = unpack(x.float(), height, width)
+                with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
+                    x = ae.decode(x)
+
+                t1 = time.perf_counter()
+                print(f"Done in {t1 - t0:.1f}s")
+
+                # Process and display result
+                x = x.clamp(-1, 1)
+                x = embed_watermark(x.float())
+                x = rearrange(x[0], "c h w -> h w c")
+                img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
+
+                nsfw_score = [x["score"] for x in nsfw_classifier(img) if x["label"] == "nsfw"][0]
+
+                if nsfw_score < NSFW_THRESHOLD:
+                    buffer = BytesIO()
+                    exif_data = Image.Exif()
+                    exif_data[ExifTags.Base.Software] = "AI generated;inpainting;flux"
+                    exif_data[ExifTags.Base.Make] = "Black Forest Labs"
+                    exif_data[ExifTags.Base.Model] = name
+                    if add_sampling_metadata:
+                        exif_data[ExifTags.Base.ImageDescription] = prompt
+                    img.save(buffer, format="jpeg", exif=exif_data, quality=95, subsampling=0)
+
+                    img_bytes = buffer.getvalue()
+                    if save_samples:
+                        fn = output_name.format(idx=idx)
+                        print(f"Saving {fn}")
+                        with open(fn, "wb") as file:
+                            file.write(img_bytes)
+
+                    st.session_state["samples"] = {
+                        "prompt": prompt,
+                        "img": img,
+                        "seed": seed,
+                        "bytes": img_bytes,
+                    }
+                else:
+                    st.warning("Your generated image may contain NSFW content.")
+                    st.session_state["samples"] = None
+
+            except Exception as e:
+                st.error(f"Error during generation: {e}")
+                return
+            finally:
+                # Clean up temporary files
+                try:
+                    os.unlink(tmp_img.name)
+                    os.unlink(tmp_mask.name)
+                except Exception as e:
+                    print(f"Error cleaning up temporary files: {e}")
+
+    # Display results
+    samples = st.session_state.get("samples", None)
+    if samples is not None:
+        st.image(samples["img"], caption=samples["prompt"])
+        col1, col2 = st.columns(2)
+        with col1:
+            st.download_button(
+                "Download full-resolution",
+                samples["bytes"],
+                file_name="generated.jpg",
+                mime="image/jpg",
+            )
+        with col2:
+            if st.button("Continue from this image"):
+                # Store the generated image
+                new_image = samples["img"]
+                # Clear ALL canvas state
+                clear_canvas_state()
+                if "samples" in st.session_state:
+                    del st.session_state["samples"]
+                # Set as current image
+                st.session_state["current_image"] = new_image
+                st.rerun()
+
+        st.write(f"Seed: {samples['seed']}")
+
+
+if __name__ == "__main__":
+    st.set_page_config(layout="wide")
+    main()

data/flux/docs/fill.md

@@ -0,0 +1,44 @@
+## Models
+
+FLUX.1 Fill introduces advanced inpainting and outpainting capabilities. It allows for seamless edits that integrate naturally with existing images.
+
+| Name                | HuggingFace repo                                         | License                                                               | sha256sum                                                        |
+| ------------------- | -------------------------------------------------------- | --------------------------------------------------------------------- | ---------------------------------------------------------------- |
+| `FLUX.1 Fill [dev]` | https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) | 03e289f530df51d014f48e675a9ffa2141bc003259bf5f25d75b957e920a41ca |
+| `FLUX.1 Fill [pro]` | Only available in our API.                               |
+
+## Examples
+
+![inpainting](../assets/docs/inpainting.png)
+![outpainting](../assets/docs/outpainting.png)
+
+## Open-weights usage
+
+The weights will be downloaded automatically from HuggingFace once you start one of the demos. To download `FLUX.1 Fill [dev]`, you will need to be logged in, see [here](https://huggingface.co/docs/huggingface_hub/guides/cli#huggingface-cli-login). Alternatively, if you have downloaded the model weights manually from [here](https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev), you can specify the downloaded paths via environment variables:
+
+```bash
+export FLUX_DEV_FILL=<path_to_flux_dev_fill_sft_file>
+export AE=<path_to_ae_sft_file>
+```
+
+For interactive sampling run
+
+```bash
+python -m src.flux.cli_fill --loop
+```
+
+Or to generate a single sample run
+
+```bash
+python -m src.flux.cli_fill \
+  --img_cond_path <path_to_input_image> \
+  --img_cond_mask <path_to_input_mask>
+```
+
+The input_mask should be an image of the same size as the conditioning image that only contains black and white pixels; see [an example mask](../assets/cup_mask.png) for [this image](../assets/cup.png).
+
+We also provide an interactive streamlit demo. The demo can be run via
+
+```bash
+streamlit run demo_st_fill.py
+```

data/flux/docs/image-variation.md

@@ -0,0 +1,33 @@
+## Models
+
+FLUX.1 Redux is an adapter for the FLUX.1 text-to-image base models, FLUX.1 [dev] and FLUX.1 [schnell], which can be used to generate image variations. 
+In addition, FLUX.1 Redux [pro] is available in our API and, augmenting the [dev] adapter, the API endpoint allows users to modify an image given a textual description. The feature is supported in our latest model FLUX1.1 [pro] Ultra, allowing for combining input images and text prompts to create high-quality 4-megapixel outputs with flexible aspect ratios.
+
+| Name                        | HuggingFace repo                                                                                | License                                                               | sha256sum                                                        |
+| --------------------------- | ----------------------------------------------------------------------------------------------- | --------------------------------------------------------------------- | ---------------------------------------------------------------- |
+| `FLUX.1 Redux [dev]`        | https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev                                       | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) | a1b3bdcb4bdc58ce04874b9ca776d61fc3e914bb6beab41efb63e4e2694dca45 |
+| `FLUX.1 Redux [pro]`        | [Available in our API.](https://docs.bfl.ml/) Supports image variations.                        |
+| `FLUX1.1 Redux [pro] Ultra` | [Available in our API.](https://docs.bfl.ml/) Supports image variations based on a text prompt. |
+
+## Examples
+
+![redux](../assets/docs/redux.png)
+
+## Open-weights usage
+
+The text-to-image base model weights and the autoencoder weights will be downloaded automatically from HuggingFace once you start the demo. To download `FLUX.1 [dev]`, you will need to be logged in, see [here](https://huggingface.co/docs/huggingface_hub/guides/cli#huggingface-cli-login). You need to manually download the adapter weights from [here](https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev) and specify them via an environment variable `export FLUX_REDUX=<path_to_flux_redux_sft_file>`. In general, you may specify any manually downloaded weights via environment variables:
+
+```bash
+export FLUX_REDUX=<path_to_flux_redux_sft_file>
+export FLUX_SCHNELL=<path_to_flux_schnell_sft_file>
+export FLUX_DEV=<path_to_flux_dev_sft_file>
+export AE=<path_to_ae_sft_file>
+```
+
+For interactive sampling run
+
+```bash
+python -m src.flux.cli_redux --loop --name <name>
+```
+
+where `name` is one of `flux-dev` or `flux-schnell`.

data/flux/docs/structural-conditioning.md

@@ -0,0 +1,40 @@
+## Models
+
+Structural conditioning uses canny edge or depth detection to maintain precise control during image transformations. By preserving the original image's structure through edge or depth maps, users can make text-guided edits while keeping the core composition intact. This is particularly effective for retexturing images. We release four variations: two based on edge maps (full model and LoRA for FLUX.1 [dev]) and two based on depth maps (full model and LoRA for FLUX.1 [dev]).
+
+| Name                      | HuggingFace repo                                               | License                                                               | sha256sum                                                        |
+| ------------------------- | -------------------------------------------------------------- | --------------------------------------------------------------------- | ---------------------------------------------------------------- |
+| `FLUX.1 Canny [dev]`      | https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev      | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) | 996876670169591cb412b937fbd46ea14cbed6933aef17c48a2dcd9685c98cdb |
+| `FLUX.1 Depth [dev]`      | https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev      | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) | 41360d1662f44ca45bc1b665fe6387e91802f53911001630d970a4f8be8dac21 |
+| `FLUX.1 Canny [dev] LoRA` | https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev-lora | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) | 8eaa21b9c43d5e7242844deb64b8cf22ae9010f813f955ca8c05f240b8a98f7e |
+| `FLUX.1 Depth [dev] LoRA` | https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev-lora | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) | 1938b38ea0fdd98080fa3e48beb2bedfbc7ad102d8b65e6614de704a46d8b907 | 
+| `FLUX.1 Canny [pro]`      | [Available in our API](https://docs.bfl.ml/).                  |
+| `FLUX.1 Depth [pro]`      | [Available in our API](https://docs.bfl.ml/).                  |
+
+## Examples
+
+![canny](../assets/docs/canny.png)
+![depth](../assets/docs/depth.png)
+
+## Open-weights usage
+
+The full model weights (`FLUX.1 Canny [dev], Flux.1 Depth [dev], FLUX.1 [dev], and the autoencoder) will be downloaded automatically from HuggingFace once you start one of the demos. To download them, you will need to be logged in, see [here](https://huggingface.co/docs/huggingface_hub/guides/cli#huggingface-cli-login). The LoRA weights are not downloaded automatically, but can be downloaded manually [here (Canny)](https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev-lora) and [here (Depth)](https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev-lora). You may specify any manually downloaded weights via environment variables: (**necessary for LoRAs**):
+
+```bash
+export FLUX_DEV_DEPTH=<path_to_flux_dev_depth_sft_file>
+export FLUX_DEV_CANNY=<path_to_flux_dev_canny_sft_file>
+export FLUX_DEV_DEPTH_LORA=<path_to_flux_dev_depth_lora_sft_file>
+export FLUX_DEV_CANNY_LORA=<path_to_flux_dev_canny_lora_sft_file>
+export FLUX_REDUX=<path_to_flux_redux_sft_file>
+export FLUX_SCHNELL=<path_to_flux_schnell_sft_file>
+export FLUX_DEV=<path_to_flux_dev_sft_file>
+export AE=<path_to_ae_sft_file>
+```
+
+For interactive sampling run
+
+```bash
+python -m src.flux.cli_control --loop --name <name>
+```
+
+where `name` is one of `flux-dev-canny`, `flux-dev-depth`, `flux-dev-canny-lora`, or `flux-dev-depth-lora`.

data/flux/docs/text-to-image.md

@@ -0,0 +1,93 @@
+## Models
+
+We currently offer four text-to-image models. `FLUX1.1 [pro]` is our most capable model which can generate images at up to 4MP while maintaining an impressive generation time of only 10 seconds per sample.
+
+| Name                      | HuggingFace repo                                        | License                                                               | sha256sum                                                        |
+| ------------------------- | ------------------------------------------------------- | --------------------------------------------------------------------- | ---------------------------------------------------------------- |
+| `FLUX.1 [schnell]`        | https://huggingface.co/black-forest-labs/FLUX.1-schnell | [apache-2.0](model_licenses/LICENSE-FLUX1-schnell)                    | 9403429e0052277ac2a87ad800adece5481eecefd9ed334e1f348723621d2a0a |
+| `FLUX.1 [dev]`            | https://huggingface.co/black-forest-labs/FLUX.1-dev     | [FLUX.1-dev Non-Commercial License](model_licenses/LICENSE-FLUX1-dev) | 4610115bb0c89560703c892c59ac2742fa821e60ef5871b33493ba544683abd7 |
+| `FLUX.1 [pro]`            | [Available in our API](https://docs.bfl.ml/).           |
+| `FLUX1.1 [pro]`           | [Available in our API](https://docs.bfl.ml/).           |
+| `FLUX1.1 [pro] Ultra/raw` | [Available in our API](https://docs.bfl.ml/).           |
+
+## Open-weights usage
+
+The weights will be downloaded automatically from HuggingFace once you start one of the demos. To download `FLUX.1 [dev]`, you will need to be logged in, see [here](https://huggingface.co/docs/huggingface_hub/guides/cli#huggingface-cli-login).
+If you have downloaded the model weights manually, you can specify the downloaded paths via environment-variables:
+
+```bash
+export FLUX_SCHNELL=<path_to_flux_schnell_sft_file>
+export FLUX_DEV=<path_to_flux_dev_sft_file>
+export AE=<path_to_ae_sft_file>
+```
+
+For interactive sampling run
+
+```bash
+python -m flux --name <name> --loop
+```
+
+Or to generate a single sample run
+
+```bash
+python -m flux --name <name> \
+  --height <height> --width <width> \
+  --prompt "<prompt>"
+```
+
+We also provide a streamlit demo that does both text-to-image and image-to-image. The demo can be run via
+
+```bash
+streamlit run demo_st.py
+```
+
+We also offer a Gradio-based demo for an interactive experience. To run the Gradio demo:
+
+```bash
+python demo_gr.py --name flux-schnell --device cuda
+```
+
+Options:
+
+- `--name`: Choose the model to use (options: "flux-schnell", "flux-dev")
+- `--device`: Specify the device to use (default: "cuda" if available, otherwise "cpu")
+- `--offload`: Offload model to CPU when not in use
+- `--share`: Create a public link to your demo
+
+To run the demo with the dev model and create a public link:
+
+```bash
+python demo_gr.py --name flux-dev --share
+```
+
+## Diffusers integration
+
+`FLUX.1 [schnell]` and `FLUX.1 [dev]` are integrated with the [🧨 diffusers](https://github.com/huggingface/diffusers) library. To use it with diffusers, install it:
+
+```shell
+pip install git+https://github.com/huggingface/diffusers.git
+```
+
+Then you can use `FluxPipeline` to run the model
+
+```python
+import torch
+from diffusers import FluxPipeline
+
+model_id = "black-forest-labs/FLUX.1-schnell" #you can also use `black-forest-labs/FLUX.1-dev`
+
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
+pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
+
+prompt = "A cat holding a sign that says hello world"
+seed = 42
+image = pipe(
+    prompt,
+    output_type="pil",
+    num_inference_steps=4, #use a larger number if you are using [dev]
+    generator=torch.Generator("cpu").manual_seed(seed)
+).images[0]
+image.save("flux-schnell.png")
+```
+
+To learn more check out the [diffusers](https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux) documentation

data/flux/model_cards/FLUX.1-dev.md

@@ -0,0 +1,46 @@
+![FLUX.1 [dev] Grid](../assets/dev_grid.jpg)
+
+`FLUX.1 [dev]` is a 12 billion parameter rectified flow transformer capable of generating images from text descriptions.
+For more information, please read our [blog post](https://blackforestlabs.ai/announcing-black-forest-labs/).
+
+# Key Features
+1. Cutting-edge output quality, second only to our state-of-the-art model `FLUX.1 [pro]`.
+2. Competitive prompt following, matching the performance of closed source alternatives.
+3. Trained using guidance distillation, making `FLUX.1 [dev]` more efficient.
+4. Open weights to drive new scientific research, and empower artists to develop innovative workflows.
+5. Generated outputs can be used for personal, scientific, and commercial purposes, as described in the [flux-1-dev-non-commercial-license](./licence.md).
+
+# Usage
+We provide a reference implementation of `FLUX.1 [dev]`, as well as sampling code, in a dedicated [github repository](https://github.com/black-forest-labs/flux).
+Developers and creatives looking to build on top of `FLUX.1 [dev]` are encouraged to use this as a starting point.
+
+## API Endpoints
+The FLUX.1 models are also available via API from the following sources
+1. [bfl.ml](https://docs.bfl.ml/) (currently `FLUX.1 [pro]`)
+2. [replicate.com](https://replicate.com/collections/flux)
+3. [fal.ai](https://fal.ai/models/fal-ai/flux/dev)
+
+## ComfyUI
+`FLUX.1 [dev]` is also available in [Comfy UI](https://github.com/comfyanonymous/ComfyUI) for local inference with a node-based workflow.
+
+---
+# Limitations
+- This model is not intended or able to provide factual information.
+- As a statistical model this checkpoint might amplify existing societal biases.
+- The model may fail to generate output that matches the prompts.
+- Prompt following is heavily influenced by the prompting-style.
+
+# Out-of-Scope Use
+The model and its derivatives may not be used
+
+- In any way that violates any applicable national, federal, state, local or international law or regulation.
+- For the purpose of exploiting, harming or attempting to exploit or harm minors in any way; including but not limited to the solicitation, creation, acquisition, or dissemination of child exploitative content.
+- To generate or disseminate verifiably false information and/or content with the purpose of harming others.
+- To generate or disseminate personal identifiable information that can be used to harm an individual.
+- To harass, abuse, threaten, stalk, or bully individuals or groups of individuals.
+- To create non-consensual nudity or illegal pornographic content.
+- For fully automated decision making that adversely impacts an individual's legal rights or otherwise creates or modifies a binding, enforceable obligation.
+- Generating or facilitating large-scale disinformation campaigns.
+
+# License
+This model falls under the [`FLUX.1 [dev]` Non-Commercial License](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md).

data/flux/model_cards/FLUX.1-schnell.md

@@ -0,0 +1,41 @@
+![FLUX.1 [schnell] Grid](../assets/schnell_grid.jpg)
+
+`FLUX.1 [schnell]` is a 12 billion parameter rectified flow transformer capable of generating images from text descriptions.
+For more information, please read our [blog post](https://blackforestlabs.ai/announcing-black-forest-labs/).
+
+# Key Features
+1. Cutting-edge output quality and competitive prompt following, matching the performance of closed source alternatives.
+2. Trained using latent adversarial diffusion distillation, `FLUX.1 [schnell]` can generate high-quality images in only 1 to 4 steps.
+3. Released under the `apache-2.0` licence, the model can be used for personal, scientific, and commercial purposes.
+
+# Usage
+We provide a reference implementation of `FLUX.1 [schnell]`, as well as sampling code, in a dedicated [github repository](https://github.com/black-forest-labs/flux).
+Developers and creatives looking to build on top of `FLUX.1 [schnell]` are encouraged to use this as a starting point.
+
+## API Endpoints
+The FLUX.1 models are also available via API from the following sources
+1. [bfl.ml](https://docs.bfl.ml/) (currently `FLUX.1 [pro]`)
+2. [replicate.com](https://replicate.com/collections/flux)
+3. [fal.ai](https://fal.ai/models/fal-ai/flux/schnell)
+
+## ComfyUI
+`FLUX.1 [schnell]` is also available in [Comfy UI](https://github.com/comfyanonymous/ComfyUI) for local inference with a node-based workflow.
+
+---
+# Limitations
+- This model is not intended or able to provide factual information.
+- As a statistical model this checkpoint might amplify existing societal biases.
+- The model may fail to generate output that matches the prompts.
+- Prompt following is heavily influenced by the prompting-style.
+
+# Out-of-Scope Use
+The model and its derivatives may not be used
+
+- In any way that violates any applicable national, federal, state, local or international law or regulation.
+- For the purpose of exploiting, harming or attempting to exploit or harm minors in any way; including but not limited to the solicitation, creation, acquisition, or dissemination of child exploitative content.
+- To generate or disseminate verifiably false information and/or content with the purpose of harming others.
+- To generate or disseminate personal identifiable information that can be used to harm an individual.
+- To harass, abuse, threaten, stalk, or bully individuals or groups of individuals.
+- To create non-consensual nudity or illegal pornographic content.
+- For fully automated decision making that adversely impacts an individual's legal rights or otherwise creates or modifies a binding, enforceable obligation.
+- Generating or facilitating large-scale disinformation campaigns.

data/flux/model_licenses/LICENSE-FLUX1-dev

@@ -0,0 +1,42 @@
+FLUX.1 [dev] Non-Commercial License 
+Black Forest Labs, Inc. (“we” or “our” or “Company”) is pleased to make available the weights, parameters and inference code for the FLUX.1 [dev] Model (as defined below) freely available for your non-commercial and non-production use as set forth in this FLUX.1 [dev] Non-Commercial License (“License”).  The “FLUX.1 [dev] Model” means the FLUX.1 [dev] AI models, including FLUX.1 [dev], FLUX.1 Fill [dev], FLUX.1 Depth [dev], FLUX.1 Canny [dev], FLUX.1 Redux [dev], FLUX.1 Canny [dev] LoRA and FLUX.1 Depth [dev] LoRA, and their elements which includes algorithms, software, checkpoints, parameters, source code (inference code, evaluation code, and if applicable, fine-tuning code) and any other materials associated with the FLUX.1 [dev] AI models made available by Company under this License, including if any, the technical documentation, manuals and instructions for the use and operation thereof (collectively, “FLUX.1 [dev] Model”).
+By downloading, accessing, use, Distributing (as defined below), or creating a Derivative (as defined below) of the FLUX.1 [dev] Model, you agree to the terms of this License. If you do not agree to this License, then you do not have any rights to access, use, Distribute or create a Derivative of the FLUX.1 [dev] Model and you must immediately cease using the FLUX.1 [dev] Model. If you are agreeing to be bound by the terms of this License on behalf of your employer or other entity, you represent and warrant to us that you have full legal authority to bind your employer or such entity to this License. If you do not have the requisite authority, you may not accept the License or access the FLUX.1 [dev] Model on behalf of your employer or other entity.
+    1. Definitions. Capitalized terms used in this License but not defined herein have the following meanings:
+        a. “Derivative”  means any (i) modified version of the FLUX.1 [dev] Model (including but not limited to any customized or fine-tuned version thereof), (ii) work based on the FLUX.1 [dev] Model, or (iii) any other derivative work thereof. For the avoidance of doubt, Outputs are not considered Derivatives under this License. 
+        b. “Distribution” or “Distribute” or “Distributing” means providing or making available, by any means, a copy of the FLUX.1 [dev] Models and/or the Derivatives as the case may be. 
+        c. “Non-Commercial Purpose” means any of the following uses, but only so far as you do not receive any direct or indirect payment arising from the use of the model or its output: (i) personal use for research, experiment, and testing for the benefit of public knowledge, personal study, private entertainment, hobby projects, or otherwise not directly or indirectly connected to any commercial activities, business operations, or employment responsibilities; (ii) use by commercial or for-profit entities for testing, evaluation, or non-commercial research and development in a non-production environment, (iii) use by any charitable organization for charitable purposes, or for testing or evaluation. For clarity, use for revenue-generating activity or direct interactions with or impacts on end users, or use to train, fine tune or distill other models for commercial use is not a Non-Commercial purpose.
+        d. “Outputs” means any content generated by the operation of the FLUX.1 [dev] Models or the Derivatives from a prompt (i.e., text instructions) provided by users. For the avoidance of doubt, Outputs do not include any components of a FLUX.1 [dev] Models, such as any fine-tuned versions of the FLUX.1 [dev] Models, the weights, or parameters. 
+        e.   “you” or “your” means the individual or entity entering into this License with Company.
+    2. License Grant.
+        a. License. Subject to your compliance with this License, Company grants you a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable, royalty free and limited license to access, use, create Derivatives of, and Distribute the FLUX.1 [dev] Models solely for your Non-Commercial Purposes. The foregoing license is personal to you, and you may not assign or sublicense this License or any other rights or obligations under this License without Company’s prior written consent; any such assignment or sublicense will be void and will automatically and immediately terminate this License.  Any restrictions set forth herein in regarding the FLUX.1 [dev] Model also applies to any Derivative you create or that are created on your behalf.
+        b. Non-Commercial Use Only.  You may only access, use, Distribute, or creative Derivatives of or the FLUX.1 [dev] Model or Derivatives for Non-Commercial Purposes.  If You want to use a FLUX.1 [dev] Model a Derivative for any purpose that is not expressly authorized under this License, such as for a commercial activity, you must request a license from Company, which Company may grant to you in Company’s sole discretion and which additional use may be subject to a fee, royalty or other revenue share. Please contact Company at the following e-mail address if you want to discuss such a license: info@blackforestlabs.ai. 
+        c. Reserved Rights. The grant of rights expressly set forth in this License are the complete grant of rights to you in the FLUX.1 [dev] Model, and no other licenses are granted, whether by waiver, estoppel, implication, equity or otherwise. Company and its licensors reserve all rights not expressly granted by this License. 
+        d. Outputs. We claim no ownership rights in and to the Outputs. You are solely responsible for the Outputs you generate and their subsequent uses in accordance with this License.  You may use Output for any purpose (including for commercial purposes), except as expressly prohibited herein.  You may not use the Output to train, fine-tune or distill a model that is competitive with the FLUX.1 [dev] Model.
+    3. Distribution. Subject to this License, you may Distribute copies of the FLUX.1 [dev] Model and/or Derivatives made by you, under the following conditions: 
+        a. you must make available a copy of this License to third-party recipients of the FLUX.1 [dev] Models and/or Derivatives you Distribute, and specify that any rights to use the FLUX.1 [dev] Models and/or Derivatives shall be directly granted by Company to said third-party recipients pursuant to this License; 
+        b. you must make prominently display the following notice alongside the Distribution of the FLUX.1 [dev] Model or Derivative (such as via a “Notice” text file distributed as part of such FLUX.1 [dev] Model or Derivative) (the “Attribution Notice”): 
+“The FLUX.1 [dev] Model is licensed by Black Forest Labs. Inc. under the FLUX.1 [dev] Non-Commercial License. Copyright Black Forest Labs. Inc. 
+IN NO EVENT SHALL BLACK FOREST LABS, INC. BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH USE OF THIS MODEL.”
+        c. in the case of Distribution of Derivatives made by you, you must also include in the Attribution Notice a statement that you have modified the applicable FLUX.1 [dev] Model; and
+        d. in the case of Distribution of Derivatives made by you, any terms and conditions you impose on any third-party recipients relating to Derivatives made by or for you shall neither limit such third-party recipients’ use of the FLUX.1 [dev] Model or any Derivatives made by or for Company in accordance with this License nor conflict with any of its terms and conditions. 
+        e. In the case of Distribution of Derivatives made by you, you must not misrepresent or imply, through any means, that the Derivatives made by or for you and/or any modified version of the FLUX.1 [dev] Model you Distribute under your name and responsibility is an official product of the Company or has been endorsed, approved or validated by the Company, unless you are authorized by Company to do so in writing.
+    4. Restrictions.  You will not, and will not permit, assist or cause any third party to 
+        a. use, modify, copy, reproduce, create Derivatives of, or Distribute the FLUX.1 [dev] Model (or any Derivative thereof, or any data produced by the FLUX.1 [dev] Model), in whole or in part, for (i) any commercial or production purposes, (ii) military purposes, (iii) purposes of surveillance, including any research or development relating to surveillance, (iv) biometric processing, (v) in any manner that infringes, misappropriates, or otherwise violates any third-party rights, or (vi) in any manner that violates any applicable law and violating any privacy or security laws, rules, regulations, directives, or governmental requirements (including the General Data Privacy Regulation (Regulation (EU) 2016/679), the California Consumer Privacy Act, and any and all laws governing the processing of biometric information), as well as all amendments and successor laws to any of the foregoing;
+        b. alter or remove copyright and other proprietary notices which appear on or in any portion of the FLUX.1 [dev] Model;
+        c. utilize any equipment, device, software, or other means to circumvent or remove any security or protection used by Company in connection with the FLUX.1 [dev] Model, or to circumvent or remove any usage restrictions, or to enable functionality disabled by FLUX.1 [dev] Model; or
+        d. offer or impose any terms on the FLUX.1 [dev] Model that alter, restrict, or are inconsistent with the terms of this License.
+        e. violate any applicable U.S. and non-U.S. export control and trade sanctions laws (“Export Laws”) in connection with your use or Distribution of any FLUX.1 [dev] Model;
+        f. directly or indirectly Distribute, export, or otherwise transfer FLUX.1 [dev] Model  (a) to any individual, entity, or country prohibited by Export Laws; (b) to anyone on U.S. or non-U.S. government restricted parties lists; or (c) for any purpose prohibited by Export Laws, including nuclear, chemical or biological weapons, or missile technology applications; 3) use or download FLUX.1 [dev] Model if you or they are  (a) located in a comprehensively sanctioned jurisdiction, (b) currently listed on any U.S. or non-U.S. restricted parties list, or (c) for any purpose prohibited by Export Laws; and (4) will not disguise your location through IP proxying or other methods.
+    5. DISCLAIMERS.  THE FLUX.1 [dev] MODEL IS PROVIDED “AS IS” AND “WITH ALL FAULTS” WITH NO WARRANTY OF ANY KIND, EXPRESS OR IMPLIED. COMPANY EXPRESSLY DISCLAIMS ALL REPRESENTATIONS AND WARRANTIES, EXPRESS OR IMPLIED, WHETHER BY STATUTE, CUSTOM, USAGE OR OTHERWISE AS TO ANY MATTERS RELATED TO THE FLUX.1 [dev] MODEL, INCLUDING BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE, SATISFACTORY QUALITY, OR NON-INFRINGEMENT. COMPANY MAKES NO WARRANTIES OR REPRESENTATIONS THAT THE FLUX.1 [dev] MODEL WILL BE ERROR FREE OR FREE OF VIRUSES OR OTHER HARMFUL COMPONENTS, OR PRODUCE ANY PARTICULAR RESULTS.
+    6. LIMITATION OF LIABILITY.  TO THE FULLEST EXTENT PERMITTED BY LAW, IN NO EVENT WILL COMPANY BE LIABLE TO YOU OR YOUR EMPLOYEES, AFFILIATES, USERS, OFFICERS OR DIRECTORS (A) UNDER ANY THEORY OF LIABILITY, WHETHER BASED IN CONTRACT, TORT, NEGLIGENCE, STRICT LIABILITY, WARRANTY, OR OTHERWISE UNDER THIS LICENSE, OR (B) FOR ANY INDIRECT, CONSEQUENTIAL, EXEMPLARY, INCIDENTAL, PUNITIVE OR SPECIAL DAMAGES OR LOST PROFITS, EVEN IF COMPANY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. THE FLUX.1 [dev] MODEL, ITS CONSTITUENT COMPONENTS, AND ANY OUTPUT (COLLECTIVELY, “MODEL MATERIALS”) ARE NOT DESIGNED OR INTENDED FOR USE IN ANY APPLICATION OR SITUATION WHERE FAILURE OR FAULT OF THE MODEL MATERIALS COULD REASONABLY BE ANTICIPATED TO LEAD TO SERIOUS INJURY OF ANY PERSON, INCLUDING POTENTIAL DISCRIMINATION OR VIOLATION OF AN INDIVIDUAL’S PRIVACY RIGHTS, OR TO SEVERE PHYSICAL, PROPERTY, OR ENVIRONMENTAL DAMAGE (EACH, A “HIGH-RISK USE”). IF YOU ELECT TO USE ANY OF THE MODEL MATERIALS FOR A HIGH-RISK USE, YOU DO SO AT YOUR OWN RISK. YOU AGREE TO DESIGN AND IMPLEMENT APPROPRIATE DECISION-MAKING AND RISK-MITIGATION PROCEDURES AND POLICIES IN CONNECTION WITH A HIGH-RISK USE SUCH THAT EVEN IF THERE IS A FAILURE OR FAULT IN ANY OF THE MODEL MATERIALS, THE SAFETY OF PERSONS OR PROPERTY AFFECTED BY THE ACTIVITY STAYS AT A LEVEL THAT IS REASONABLE, APPROPRIATE, AND LAWFUL FOR THE FIELD OF THE HIGH-RISK USE.
+    7. INDEMNIFICATION
+
+You will indemnify, defend and hold harmless Company and our subsidiaries and affiliates, and each of our respective shareholders, directors, officers, employees, agents, successors, and assigns (collectively, the “Company Parties”) from and against any losses, liabilities, damages, fines, penalties, and expenses (including reasonable attorneys’ fees) incurred by any Company Party in connection with any claim, demand, allegation, lawsuit, proceeding, or investigation (collectively, “Claims”) arising out of or related to  (a) your access to or use of the FLUX.1 [dev] Model (as well as any Output, results or data generated from such access or use), including any High-Risk Use (defined below); (b) your violation of this License; or (c) your violation, misappropriation or infringement of any rights of another (including intellectual property or other proprietary rights and privacy rights). You will promptly notify the Company Parties of any such Claims, and cooperate with Company Parties in defending such Claims. You will also grant the Company Parties sole control of the defense or settlement, at Company’s sole option, of any Claims. This indemnity is in addition to, and not in lieu of, any other indemnities or remedies set forth in a written agreement between you and Company or the other Company Parties.
+    8. Termination; Survival.
+        a. This License will automatically terminate upon any breach by you of the terms of this License.
+        b. We may terminate this License, in whole or in part, at any time upon notice (including electronic) to you.
+        c. If You initiate any legal action or proceedings against Company or any other entity (including a cross-claim or counterclaim in a lawsuit), alleging that the FLUX.1 [dev] Model or any Derivative, or any part thereof, infringe upon intellectual property or other rights owned or licensable by you, then any licenses granted to you under this License will immediately terminate as of the date such legal action or claim is filed or initiated.
+        d. Upon termination of this License, you must cease all use, access or Distribution of the FLUX.1 [dev] Model and any Derivatives.  The following sections survive termination of this License  2(c), 2(d), 4-11.  
+    9. Third Party Materials. The FLUX.1 [dev] Model may contain third-party software or other components (including free and open source software) (all of the foregoing, “Third Party Materials”), which are subject to the license terms of the respective third-party licensors. Your dealings or correspondence with third parties and your use of or interaction with any Third Party Materials are solely between you and the third party. Company does not control or endorse, and makes no representations or warranties regarding, any Third Party Materials, and your access to and use of such Third Party Materials are at your own risk.
+    10. Trademarks. You have not been granted any trademark license as part of this License and may not use any name or mark associated with Company without the prior written permission of Company, except to the extent necessary to make the reference required in the Attribution Notice as specified above or as is reasonably necessary in describing the FLUX.1 [dev] Model and its creators.  
+    11. General. This License will be governed and construed under the laws of the State of Delaware without regard to conflicts of law provisions. If any provision or part of a provision of this License is unlawful, void or unenforceable, that provision or part of the provision is deemed severed from this License, and will not affect the validity and enforceability of any remaining provisions. The failure of Company to exercise or enforce any right or provision of this License will not operate as a waiver of such right or provision. This License does not confer any third-party beneficiary rights upon any other person or entity. This License, together with the Documentation, contains the entire understanding between you and Company regarding the subject matter of this License, and supersedes all other written or oral agreements and understandings between you and Company regarding such subject matter. No change or addition to any provision of this License will be binding unless it is in writing and signed by an authorized representative of both you and Company.

data/flux/model_licenses/LICENSE-FLUX1-schnell

@@ -0,0 +1,54 @@
+
+
+Apache License
+Version 2.0, January 2004
+http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+1. Definitions.
+
+"License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document.
+
+"Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License.
+
+"Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
+
+"You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License.
+
+"Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.
+
+"Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.
+
+"Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below).
+
+"Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof.
+
+"Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
+
+"Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work.
+
+2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form.
+
+3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
+
+4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions:
+
+    You must give any other recipients of the Work or Derivative Works a copy of this License; and
+    You must cause any modified files to carry prominent notices stating that You changed the files; and
+    You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and
+    If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License.
+
+You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License.
+
+5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions.
+
+6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
+
+7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages.
+
+9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability.
+
+END OF TERMS AND CONDITIONS

data/flux/pyproject.toml

@@ -0,0 +1,99 @@
+[project]
+name = "flux"
+authors = [
+  { name = "Black Forest Labs", email = "support@blackforestlabs.ai" },
+]
+description = "Inference codebase for FLUX"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { file = "LICENSE.md" }
+dynamic = ["version"]
+dependencies = [
+  "torch == 2.5.1",
+  "torchvision",
+  "einops",
+  "fire >= 0.6.0",
+  "huggingface-hub",
+  "safetensors",
+  "sentencepiece",
+  "transformers",
+  "tokenizers",
+  "protobuf",
+  "requests",
+  "invisible-watermark",
+  "ruff == 0.6.8",
+]
+
+[project.optional-dependencies]
+streamlit = [
+  "streamlit",
+  "streamlit-drawable-canvas",
+  "streamlit-keyup",
+]
+gradio = [
+  "gradio",
+]
+all = [
+  "flux[streamlit]",
+  "flux[gradio]",
+]
+
+[project.scripts]
+flux = "flux.cli:app"
+
+[build-system]
+build-backend = "setuptools.build_meta"
+requires = ["setuptools>=64", "wheel", "setuptools_scm>=8"]
+
+[tool.ruff]
+line-length = 110
+target-version = "py310"
+extend-exclude = ["/usr/lib/*"]
+
+[tool.ruff.lint]
+ignore = [
+  "E501", # line too long - will be fixed in format
+]
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+line-ending = "auto"
+skip-magic-trailing-comma = false
+docstring-code-format = true
+exclude = [
+  "src/flux/_version.py", # generated by setuptools_scm
+]
+
+[tool.ruff.lint.isort]
+combine-as-imports = true
+force-wrap-aliases = true
+known-local-folder = ["src"]
+known-first-party = ["flux"]
+
+[tool.pyright]
+include = ["src"]
+exclude = [
+  "**/__pycache__", # cache directories
+  "./typings",      # generated type stubs
+]
+stubPath = "./typings"
+
+[tool.tomlsort]
+in_place = true
+no_sort_tables = true
+spaces_before_inline_comment = 1
+spaces_indent_inline_array = 2
+trailing_comma_inline_array = true
+sort_first = [
+  "project",
+  "build-system",
+  "tool.setuptools",
+]
+
+# needs to be last for CI reasons
+[tool.setuptools_scm]
+write_to = "src/flux/_version.py"
+parentdir_prefix_version = "flux-"
+fallback_version = "0.0.0"
+version_scheme = "post-release"

data/flux/setup.py

@@ -0,0 +1,3 @@
+import setuptools
+
+setuptools.setup()

data/flux/src/flux/__init__.py

@@ -0,0 +1,13 @@
+try:
+    from ._version import (
+        version as __version__,  # type: ignore
+        version_tuple,
+    )
+except ImportError:
+    __version__ = "unknown (no version information available)"
+    version_tuple = (0, 0, "unknown", "noinfo")
+
+from pathlib import Path
+
+PACKAGE = __package__.replace("_", "-")
+PACKAGE_ROOT = Path(__file__).parent

data/flux/src/flux/__main__.py

@@ -0,0 +1,4 @@
+from .cli import app
+
+if __name__ == "__main__":
+    app()

data/flux/src/flux/api.py

@@ -0,0 +1,225 @@
+import io
+import os
+import time
+from pathlib import Path
+
+import requests
+from PIL import Image
+
+API_URL = "https://api.bfl.ml"
+API_ENDPOINTS = {
+    "flux.1-pro": "flux-pro",
+    "flux.1-dev": "flux-dev",
+    "flux.1.1-pro": "flux-pro-1.1",
+}
+
+
+class ApiException(Exception):
+    def __init__(self, status_code: int, detail: str | list[dict] | None = None):
+        super().__init__()
+        self.detail = detail
+        self.status_code = status_code
+
+    def __str__(self) -> str:
+        return self.__repr__()
+
+    def __repr__(self) -> str:
+        if self.detail is None:
+            message = None
+        elif isinstance(self.detail, str):
+            message = self.detail
+        else:
+            message = "[" + ",".join(d["msg"] for d in self.detail) + "]"
+        return f"ApiException({self.status_code=}, {message=}, detail={self.detail})"
+
+
+class ImageRequest:
+    def __init__(
+        self,
+        # api inputs
+        prompt: str,
+        name: str = "flux.1.1-pro",
+        width: int | None = None,
+        height: int | None = None,
+        num_steps: int | None = None,
+        prompt_upsampling: bool | None = None,
+        seed: int | None = None,
+        guidance: float | None = None,
+        interval: float | None = None,
+        safety_tolerance: int | None = None,
+        # behavior of this class
+        validate: bool = True,
+        launch: bool = True,
+        api_key: str | None = None,
+    ):
+        """
+        Manages an image generation request to the API.
+
+        All parameters not specified will use the API defaults.
+
+        Args:
+            prompt: Text prompt for image generation.
+            width: Width of the generated image in pixels. Must be a multiple of 32.
+            height: Height of the generated image in pixels. Must be a multiple of 32.
+            name: Which model version to use
+            num_steps: Number of steps for the image generation process.
+            prompt_upsampling: Whether to perform upsampling on the prompt.
+            seed: Optional seed for reproducibility.
+            guidance: Guidance scale for image generation.
+            safety_tolerance: Tolerance level for input and output moderation.
+                 Between 0 and 6, 0 being most strict, 6 being least strict.
+            validate: Run input validation
+            launch: Directly launches request
+            api_key: Your API key if not provided by the environment
+
+        Raises:
+            ValueError: For invalid input, when `validate`
+            ApiException: For errors raised from the API
+        """
+        if validate:
+            if name not in API_ENDPOINTS.keys():
+                raise ValueError(f"Invalid model {name}")
+            elif width is not None and width % 32 != 0:
+                raise ValueError(f"width must be divisible by 32, got {width}")
+            elif width is not None and not (256 <= width <= 1440):
+                raise ValueError(f"width must be between 256 and 1440, got {width}")
+            elif height is not None and height % 32 != 0:
+                raise ValueError(f"height must be divisible by 32, got {height}")
+            elif height is not None and not (256 <= height <= 1440):
+                raise ValueError(f"height must be between 256 and 1440, got {height}")
+            elif num_steps is not None and not (1 <= num_steps <= 50):
+                raise ValueError(f"steps must be between 1 and 50, got {num_steps}")
+            elif guidance is not None and not (1.5 <= guidance <= 5.0):
+                raise ValueError(f"guidance must be between 1.5 and 4, got {guidance}")
+            elif interval is not None and not (1.0 <= interval <= 4.0):
+                raise ValueError(f"interval must be between 1 and 4, got {interval}")
+            elif safety_tolerance is not None and not (0 <= safety_tolerance <= 6.0):
+                raise ValueError(f"safety_tolerance must be between 0 and 6, got {interval}")
+
+            if name == "flux.1-dev":
+                if interval is not None:
+                    raise ValueError("Interval is not supported for flux.1-dev")
+            if name == "flux.1.1-pro":
+                if interval is not None or num_steps is not None or guidance is not None:
+                    raise ValueError("Interval, num_steps and guidance are not supported for " "flux.1.1-pro")
+
+        self.name = name
+        self.request_json = {
+            "prompt": prompt,
+            "width": width,
+            "height": height,
+            "steps": num_steps,
+            "prompt_upsampling": prompt_upsampling,
+            "seed": seed,
+            "guidance": guidance,
+            "interval": interval,
+            "safety_tolerance": safety_tolerance,
+        }
+        self.request_json = {key: value for key, value in self.request_json.items() if value is not None}
+
+        self.request_id: str | None = None
+        self.result: dict | None = None
+        self._image_bytes: bytes | None = None
+        self._url: str | None = None
+        if api_key is None:
+            self.api_key = os.environ.get("BFL_API_KEY")
+        else:
+            self.api_key = api_key
+
+        if launch:
+            self.request()
+
+    def request(self):
+        """
+        Request to generate the image.
+        """
+        if self.request_id is not None:
+            return
+        response = requests.post(
+            f"{API_URL}/v1/{API_ENDPOINTS[self.name]}",
+            headers={
+                "accept": "application/json",
+                "x-key": self.api_key,
+                "Content-Type": "application/json",
+            },
+            json=self.request_json,
+        )
+        result = response.json()
+        if response.status_code != 200:
+            raise ApiException(status_code=response.status_code, detail=result.get("detail"))
+        self.request_id = response.json()["id"]
+
+    def retrieve(self) -> dict:
+        """
+        Wait for the generation to finish and retrieve response.
+        """
+        if self.request_id is None:
+            self.request()
+        while self.result is None:
+            response = requests.get(
+                f"{API_URL}/v1/get_result",
+                headers={
+                    "accept": "application/json",
+                    "x-key": self.api_key,
+                },
+                params={
+                    "id": self.request_id,
+                },
+            )
+            result = response.json()
+            if "status" not in result:
+                raise ApiException(status_code=response.status_code, detail=result.get("detail"))
+            elif result["status"] == "Ready":
+                self.result = result["result"]
+            elif result["status"] == "Pending":
+                time.sleep(0.5)
+            else:
+                raise ApiException(status_code=200, detail=f"API returned status '{result['status']}'")
+        return self.result
+
+    @property
+    def bytes(self) -> bytes:
+        """
+        Generated image as bytes.
+        """
+        if self._image_bytes is None:
+            response = requests.get(self.url)
+            if response.status_code == 200:
+                self._image_bytes = response.content
+            else:
+                raise ApiException(status_code=response.status_code)
+        return self._image_bytes
+
+    @property
+    def url(self) -> str:
+        """
+        Public url to retrieve the image from
+        """
+        if self._url is None:
+            result = self.retrieve()
+            self._url = result["sample"]
+        return self._url
+
+    @property
+    def image(self) -> Image.Image:
+        """
+        Load the image as a PIL Image
+        """
+        return Image.open(io.BytesIO(self.bytes))
+
+    def save(self, path: str):
+        """
+        Save the generated image to a local path
+        """
+        suffix = Path(self.url).suffix
+        if not path.endswith(suffix):
+            path = path + suffix
+        Path(path).resolve().parent.mkdir(parents=True, exist_ok=True)
+        with open(path, "wb") as file:
+            file.write(self.bytes)
+
+
+if __name__ == "__main__":
+    from fire import Fire
+
+    Fire(ImageRequest)

data/flux/src/flux/cli.py

@@ -0,0 +1,238 @@
+import os
+import re
+import time
+from dataclasses import dataclass
+from glob import iglob
+
+import torch
+from fire import Fire
+from transformers import pipeline
+
+from flux.sampling import denoise, get_noise, get_schedule, prepare, unpack
+from flux.util import configs, load_ae, load_clip, load_flow_model, load_t5, save_image
+
+NSFW_THRESHOLD = 0.85
+
+
+@dataclass
+class SamplingOptions:
+    prompt: str
+    width: int
+    height: int
+    num_steps: int
+    guidance: float
+    seed: int | None
+
+
+def parse_prompt(options: SamplingOptions) -> SamplingOptions | None:
+    user_question = "Next prompt (write /h for help, /q to quit and leave empty to repeat):\n"
+    usage = (
+        "Usage: Either write your prompt directly, leave this field empty "
+        "to repeat the prompt or write a command starting with a slash:\n"
+        "- '/w <width>' will set the width of the generated image\n"
+        "- '/h <height>' will set the height of the generated image\n"
+        "- '/s <seed>' sets the next seed\n"
+        "- '/g <guidance>' sets the guidance (flux-dev only)\n"
+        "- '/n <steps>' sets the number of steps\n"
+        "- '/q' to quit"
+    )
+
+    while (prompt := input(user_question)).startswith("/"):
+        if prompt.startswith("/w"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, width = prompt.split()
+            options.width = 16 * (int(width) // 16)
+            print(
+                f"Setting resolution to {options.width} x {options.height} "
+                f"({options.height *options.width/1e6:.2f}MP)"
+            )
+        elif prompt.startswith("/h"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, height = prompt.split()
+            options.height = 16 * (int(height) // 16)
+            print(
+                f"Setting resolution to {options.width} x {options.height} "
+                f"({options.height *options.width/1e6:.2f}MP)"
+            )
+        elif prompt.startswith("/g"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, guidance = prompt.split()
+            options.guidance = float(guidance)
+            print(f"Setting guidance to {options.guidance}")
+        elif prompt.startswith("/s"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, seed = prompt.split()
+            options.seed = int(seed)
+            print(f"Setting seed to {options.seed}")
+        elif prompt.startswith("/n"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, steps = prompt.split()
+            options.num_steps = int(steps)
+            print(f"Setting number of steps to {options.num_steps}")
+        elif prompt.startswith("/q"):
+            print("Quitting")
+            return None
+        else:
+            if not prompt.startswith("/h"):
+                print(f"Got invalid command '{prompt}'\n{usage}")
+            print(usage)
+    if prompt != "":
+        options.prompt = prompt
+    return options
+
+
+@torch.inference_mode()
+def main(
+    name: str = "flux-schnell",
+    width: int = 1360,
+    height: int = 768,
+    seed: int | None = None,
+    prompt: str = (
+        "a photo of a forest with mist swirling around the tree trunks. The word "
+        '"FLUX" is painted over it in big, red brush strokes with visible texture'
+    ),
+    device: str = "cuda" if torch.cuda.is_available() else "cpu",
+    num_steps: int | None = None,
+    loop: bool = False,
+    guidance: float = 3.5,
+    offload: bool = False,
+    output_dir: str = "output",
+    add_sampling_metadata: bool = True,
+):
+    """
+    Sample the flux model. Either interactively (set `--loop`) or run for a
+    single image.
+
+    Args:
+        name: Name of the model to load
+        height: height of the sample in pixels (should be a multiple of 16)
+        width: width of the sample in pixels (should be a multiple of 16)
+        seed: Set a seed for sampling
+        output_name: where to save the output image, `{idx}` will be replaced
+            by the index of the sample
+        prompt: Prompt used for sampling
+        device: Pytorch device
+        num_steps: number of sampling steps (default 4 for schnell, 50 for guidance distilled)
+        loop: start an interactive session and sample multiple times
+        guidance: guidance value used for guidance distillation
+        add_sampling_metadata: Add the prompt to the image Exif metadata
+    """
+    nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=device)
+
+    if name not in configs:
+        available = ", ".join(configs.keys())
+        raise ValueError(f"Got unknown model name: {name}, chose from {available}")
+
+    torch_device = torch.device(device)
+    if num_steps is None:
+        num_steps = 4 if name == "flux-schnell" else 50
+
+    # allow for packing and conversion to latent space
+    height = 16 * (height // 16)
+    width = 16 * (width // 16)
+
+    output_name = os.path.join(output_dir, "img_{idx}.jpg")
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+        idx = 0
+    else:
+        fns = [fn for fn in iglob(output_name.format(idx="*")) if re.search(r"img_[0-9]+\.jpg$", fn)]
+        if len(fns) > 0:
+            idx = max(int(fn.split("_")[-1].split(".")[0]) for fn in fns) + 1
+        else:
+            idx = 0
+
+    # init all components
+    t5 = load_t5(torch_device, max_length=256 if name == "flux-schnell" else 512)
+    clip = load_clip(torch_device)
+    model = load_flow_model(name, device="cpu" if offload else torch_device)
+    ae = load_ae(name, device="cpu" if offload else torch_device)
+
+    rng = torch.Generator(device="cpu")
+    opts = SamplingOptions(
+        prompt=prompt,
+        width=width,
+        height=height,
+        num_steps=num_steps,
+        guidance=guidance,
+        seed=seed,
+    )
+
+    if loop:
+        opts = parse_prompt(opts)
+
+    while opts is not None:
+        if opts.seed is None:
+            opts.seed = rng.seed()
+        print(f"Generating with seed {opts.seed}:\n{opts.prompt}")
+        t0 = time.perf_counter()
+
+        # prepare input
+        x = get_noise(
+            1,
+            opts.height,
+            opts.width,
+            device=torch_device,
+            dtype=torch.bfloat16,
+            seed=opts.seed,
+        )
+        opts.seed = None
+        if offload:
+            ae = ae.cpu()
+            torch.cuda.empty_cache()
+            t5, clip = t5.to(torch_device), clip.to(torch_device)
+        inp = prepare(t5, clip, x, prompt=opts.prompt)
+        timesteps = get_schedule(opts.num_steps, inp["img"].shape[1], shift=(name != "flux-schnell"))
+
+        # offload TEs to CPU, load model to gpu
+        if offload:
+            t5, clip = t5.cpu(), clip.cpu()
+            torch.cuda.empty_cache()
+            model = model.to(torch_device)
+
+        # denoise initial noise
+        x = denoise(model, **inp, timesteps=timesteps, guidance=opts.guidance)
+
+        # offload model, load autoencoder to gpu
+        if offload:
+            model.cpu()
+            torch.cuda.empty_cache()
+            ae.decoder.to(x.device)
+
+        # decode latents to pixel space
+        x = unpack(x.float(), opts.height, opts.width)
+        with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
+            x = ae.decode(x)
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        t1 = time.perf_counter()
+
+        fn = output_name.format(idx=idx)
+        print(f"Done in {t1 - t0:.1f}s. Saving {fn}")
+
+        idx = save_image(nsfw_classifier, name, output_name, idx, x, add_sampling_metadata, prompt)
+
+        if loop:
+            print("-" * 80)
+            opts = parse_prompt(opts)
+        else:
+            opts = None
+
+
+def app():
+    Fire(main)
+
+
+if __name__ == "__main__":
+    app()

data/flux/src/flux/cli_control.py

@@ -0,0 +1,347 @@
+import os
+import re
+import time
+from dataclasses import dataclass
+from glob import iglob
+
+import torch
+from fire import Fire
+from transformers import pipeline
+
+from flux.modules.image_embedders import CannyImageEncoder, DepthImageEncoder
+from flux.sampling import denoise, get_noise, get_schedule, prepare_control, unpack
+from flux.util import configs, load_ae, load_clip, load_flow_model, load_t5, save_image
+
+
+@dataclass
+class SamplingOptions:
+    prompt: str
+    width: int
+    height: int
+    num_steps: int
+    guidance: float
+    seed: int | None
+    img_cond_path: str
+    lora_scale: float | None
+
+
+def parse_prompt(options: SamplingOptions) -> SamplingOptions | None:
+    user_question = "Next prompt (write /h for help, /q to quit and leave empty to repeat):\n"
+    usage = (
+        "Usage: Either write your prompt directly, leave this field empty "
+        "to repeat the prompt or write a command starting with a slash:\n"
+        "- '/w <width>' will set the width of the generated image\n"
+        "- '/h <height>' will set the height of the generated image\n"
+        "- '/s <seed>' sets the next seed\n"
+        "- '/g <guidance>' sets the guidance (flux-dev only)\n"
+        "- '/n <steps>' sets the number of steps\n"
+        "- '/q' to quit"
+    )
+
+    while (prompt := input(user_question)).startswith("/"):
+        if prompt.startswith("/w"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, width = prompt.split()
+            options.width = 16 * (int(width) // 16)
+            print(
+                f"Setting resolution to {options.width} x {options.height} "
+                f"({options.height *options.width/1e6:.2f}MP)"
+            )
+        elif prompt.startswith("/h"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, height = prompt.split()
+            options.height = 16 * (int(height) // 16)
+            print(
+                f"Setting resolution to {options.width} x {options.height} "
+                f"({options.height *options.width/1e6:.2f}MP)"
+            )
+        elif prompt.startswith("/g"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, guidance = prompt.split()
+            options.guidance = float(guidance)
+            print(f"Setting guidance to {options.guidance}")
+        elif prompt.startswith("/s"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, seed = prompt.split()
+            options.seed = int(seed)
+            print(f"Setting seed to {options.seed}")
+        elif prompt.startswith("/n"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, steps = prompt.split()
+            options.num_steps = int(steps)
+            print(f"Setting number of steps to {options.num_steps}")
+        elif prompt.startswith("/q"):
+            print("Quitting")
+            return None
+        else:
+            if not prompt.startswith("/h"):
+                print(f"Got invalid command '{prompt}'\n{usage}")
+            print(usage)
+    if prompt != "":
+        options.prompt = prompt
+    return options
+
+
+def parse_img_cond_path(options: SamplingOptions | None) -> SamplingOptions | None:
+    if options is None:
+        return None
+
+    user_question = "Next conditioning image (write /h for help, /q to quit and leave empty to repeat):\n"
+    usage = (
+        "Usage: Either write your prompt directly, leave this field empty "
+        "to repeat the conditioning image or write a command starting with a slash:\n"
+        "- '/q' to quit"
+    )
+
+    while True:
+        img_cond_path = input(user_question)
+
+        if img_cond_path.startswith("/"):
+            if img_cond_path.startswith("/q"):
+                print("Quitting")
+                return None
+            else:
+                if not img_cond_path.startswith("/h"):
+                    print(f"Got invalid command '{img_cond_path}'\n{usage}")
+                print(usage)
+            continue
+
+        if img_cond_path == "":
+            break
+
+        if not os.path.isfile(img_cond_path) or not img_cond_path.lower().endswith(
+            (".jpg", ".jpeg", ".png", ".webp")
+        ):
+            print(f"File '{img_cond_path}' does not exist or is not a valid image file")
+            continue
+
+        options.img_cond_path = img_cond_path
+        break
+
+    return options
+
+
+def parse_lora_scale(options: SamplingOptions | None) -> tuple[SamplingOptions | None, bool]:
+    changed = False
+
+    if options is None:
+        return None, changed
+
+    user_question = "Next lora scale (write /h for help, /q to quit and leave empty to repeat):\n"
+    usage = (
+        "Usage: Either write your prompt directly, leave this field empty "
+        "to repeat the lora scale or write a command starting with a slash:\n"
+        "- '/q' to quit"
+    )
+
+    while (prompt := input(user_question)).startswith("/"):
+        if prompt.startswith("/q"):
+            print("Quitting")
+            return None, changed
+        else:
+            if not prompt.startswith("/h"):
+                print(f"Got invalid command '{prompt}'\n{usage}")
+            print(usage)
+    if prompt != "":
+        options.lora_scale = float(prompt)
+        changed = True
+    return options, changed
+
+
+@torch.inference_mode()
+def main(
+    name: str,
+    width: int = 1024,
+    height: int = 1024,
+    seed: int | None = None,
+    prompt: str = "a robot made out of gold",
+    device: str = "cuda" if torch.cuda.is_available() else "cpu",
+    num_steps: int = 50,
+    loop: bool = False,
+    guidance: float | None = None,
+    offload: bool = False,
+    output_dir: str = "output",
+    add_sampling_metadata: bool = True,
+    img_cond_path: str = "assets/robot.webp",
+    lora_scale: float | None = 0.85,
+):
+    """
+    Sample the flux model. Either interactively (set `--loop`) or run for a
+    single image.
+
+    Args:
+        height: height of the sample in pixels (should be a multiple of 16)
+        width: width of the sample in pixels (should be a multiple of 16)
+        seed: Set a seed for sampling
+        output_name: where to save the output image, `{idx}` will be replaced
+            by the index of the sample
+        prompt: Prompt used for sampling
+        device: Pytorch device
+        num_steps: number of sampling steps (default 4 for schnell, 50 for guidance distilled)
+        loop: start an interactive session and sample multiple times
+        guidance: guidance value used for guidance distillation
+        add_sampling_metadata: Add the prompt to the image Exif metadata
+        img_cond_path: path to conditioning image (jpeg/png/webp)
+    """
+    nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=device)
+
+    assert name in [
+        "flux-dev-canny",
+        "flux-dev-depth",
+        "flux-dev-canny-lora",
+        "flux-dev-depth-lora",
+    ], f"Got unknown model name: {name}"
+    if guidance is None:
+        if name in ["flux-dev-canny", "flux-dev-canny-lora"]:
+            guidance = 30.0
+        elif name in ["flux-dev-depth", "flux-dev-depth-lora"]:
+            guidance = 10.0
+        else:
+            raise NotImplementedError()
+
+    if name not in configs:
+        available = ", ".join(configs.keys())
+        raise ValueError(f"Got unknown model name: {name}, chose from {available}")
+
+    torch_device = torch.device(device)
+
+    output_name = os.path.join(output_dir, "img_{idx}.jpg")
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+        idx = 0
+    else:
+        fns = [fn for fn in iglob(output_name.format(idx="*")) if re.search(r"img_[0-9]+\.jpg$", fn)]
+        if len(fns) > 0:
+            idx = max(int(fn.split("_")[-1].split(".")[0]) for fn in fns) + 1
+        else:
+            idx = 0
+
+    # init all components
+    t5 = load_t5(torch_device, max_length=512)
+    clip = load_clip(torch_device)
+    model = load_flow_model(name, device="cpu" if offload else torch_device)
+    ae = load_ae(name, device="cpu" if offload else torch_device)
+
+    # set lora scale
+    if "lora" in name and lora_scale is not None:
+        for _, module in model.named_modules():
+            if hasattr(module, "set_scale"):
+                module.set_scale(lora_scale)
+
+    if name in ["flux-dev-depth", "flux-dev-depth-lora"]:
+        img_embedder = DepthImageEncoder(torch_device)
+    elif name in ["flux-dev-canny", "flux-dev-canny-lora"]:
+        img_embedder = CannyImageEncoder(torch_device)
+    else:
+        raise NotImplementedError()
+
+    rng = torch.Generator(device="cpu")
+    opts = SamplingOptions(
+        prompt=prompt,
+        width=width,
+        height=height,
+        num_steps=num_steps,
+        guidance=guidance,
+        seed=seed,
+        img_cond_path=img_cond_path,
+        lora_scale=lora_scale,
+    )
+
+    if loop:
+        opts = parse_prompt(opts)
+        opts = parse_img_cond_path(opts)
+        if "lora" in name:
+            opts, changed = parse_lora_scale(opts)
+            if changed:
+                # update the lora scale:
+                for _, module in model.named_modules():
+                    if hasattr(module, "set_scale"):
+                        module.set_scale(opts.lora_scale)
+
+    while opts is not None:
+        if opts.seed is None:
+            opts.seed = rng.seed()
+        print(f"Generating with seed {opts.seed}:\n{opts.prompt}")
+        t0 = time.perf_counter()
+
+        # prepare input
+        x = get_noise(
+            1,
+            opts.height,
+            opts.width,
+            device=torch_device,
+            dtype=torch.bfloat16,
+            seed=opts.seed,
+        )
+        opts.seed = None
+        if offload:
+            t5, clip, ae = t5.to(torch_device), clip.to(torch_device), ae.to(torch_device)
+        inp = prepare_control(
+            t5,
+            clip,
+            x,
+            prompt=opts.prompt,
+            ae=ae,
+            encoder=img_embedder,
+            img_cond_path=opts.img_cond_path,
+        )
+        timesteps = get_schedule(opts.num_steps, inp["img"].shape[1], shift=(name != "flux-schnell"))
+
+        # offload TEs and AE to CPU, load model to gpu
+        if offload:
+            t5, clip, ae = t5.cpu(), clip.cpu(), ae.cpu()
+            torch.cuda.empty_cache()
+            model = model.to(torch_device)
+
+        # denoise initial noise
+        x = denoise(model, **inp, timesteps=timesteps, guidance=opts.guidance)
+
+        # offload model, load autoencoder to gpu
+        if offload:
+            model.cpu()
+            torch.cuda.empty_cache()
+            ae.decoder.to(x.device)
+
+        # decode latents to pixel space
+        x = unpack(x.float(), opts.height, opts.width)
+        with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
+            x = ae.decode(x)
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        t1 = time.perf_counter()
+        print(f"Done in {t1 - t0:.1f}s")
+
+        idx = save_image(nsfw_classifier, name, output_name, idx, x, add_sampling_metadata, prompt)
+
+        if loop:
+            print("-" * 80)
+            opts = parse_prompt(opts)
+            opts = parse_img_cond_path(opts)
+            if "lora" in name:
+                opts, changed = parse_lora_scale(opts)
+                if changed:
+                    # update the lora scale:
+                    for _, module in model.named_modules():
+                        if hasattr(module, "set_scale"):
+                            module.set_scale(opts.lora_scale)
+        else:
+            opts = None
+
+
+def app():
+    Fire(main)
+
+
+if __name__ == "__main__":
+    app()

data/flux/src/flux/cli_fill.py

@@ -0,0 +1,334 @@
+import os
+import re
+import time
+from dataclasses import dataclass
+from glob import iglob
+
+import torch
+from fire import Fire
+from PIL import Image
+from transformers import pipeline
+
+from flux.sampling import denoise, get_noise, get_schedule, prepare_fill, unpack
+from flux.util import configs, load_ae, load_clip, load_flow_model, load_t5, save_image
+
+
+@dataclass
+class SamplingOptions:
+    prompt: str
+    width: int
+    height: int
+    num_steps: int
+    guidance: float
+    seed: int | None
+    img_cond_path: str
+    img_mask_path: str
+
+
+def parse_prompt(options: SamplingOptions) -> SamplingOptions | None:
+    user_question = "Next prompt (write /h for help, /q to quit and leave empty to repeat):\n"
+    usage = (
+        "Usage: Either write your prompt directly, leave this field empty "
+        "to repeat the prompt or write a command starting with a slash:\n"
+        "- '/s <seed>' sets the next seed\n"
+        "- '/g <guidance>' sets the guidance (flux-dev only)\n"
+        "- '/n <steps>' sets the number of steps\n"
+        "- '/q' to quit"
+    )
+
+    while (prompt := input(user_question)).startswith("/"):
+        if prompt.startswith("/g"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, guidance = prompt.split()
+            options.guidance = float(guidance)
+            print(f"Setting guidance to {options.guidance}")
+        elif prompt.startswith("/s"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, seed = prompt.split()
+            options.seed = int(seed)
+            print(f"Setting seed to {options.seed}")
+        elif prompt.startswith("/n"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, steps = prompt.split()
+            options.num_steps = int(steps)
+            print(f"Setting number of steps to {options.num_steps}")
+        elif prompt.startswith("/q"):
+            print("Quitting")
+            return None
+        else:
+            if not prompt.startswith("/h"):
+                print(f"Got invalid command '{prompt}'\n{usage}")
+            print(usage)
+    if prompt != "":
+        options.prompt = prompt
+    return options
+
+
+def parse_img_cond_path(options: SamplingOptions | None) -> SamplingOptions | None:
+    if options is None:
+        return None
+
+    user_question = "Next conditioning image (write /h for help, /q to quit and leave empty to repeat):\n"
+    usage = (
+        "Usage: Either write your prompt directly, leave this field empty "
+        "to repeat the conditioning image or write a command starting with a slash:\n"
+        "- '/q' to quit"
+    )
+
+    while True:
+        img_cond_path = input(user_question)
+
+        if img_cond_path.startswith("/"):
+            if img_cond_path.startswith("/q"):
+                print("Quitting")
+                return None
+            else:
+                if not img_cond_path.startswith("/h"):
+                    print(f"Got invalid command '{img_cond_path}'\n{usage}")
+                print(usage)
+            continue
+
+        if img_cond_path == "":
+            break
+
+        if not os.path.isfile(img_cond_path) or not img_cond_path.lower().endswith(
+            (".jpg", ".jpeg", ".png", ".webp")
+        ):
+            print(f"File '{img_cond_path}' does not exist or is not a valid image file")
+            continue
+        else:
+            with Image.open(img_cond_path) as img:
+                width, height = img.size
+
+            if width % 32 != 0 or height % 32 != 0:
+                print(f"Image dimensions must be divisible by 32, got {width}x{height}")
+                continue
+
+        options.img_cond_path = img_cond_path
+        break
+
+    return options
+
+
+def parse_img_mask_path(options: SamplingOptions | None) -> SamplingOptions | None:
+    if options is None:
+        return None
+
+    user_question = "Next conditioning mask (write /h for help, /q to quit and leave empty to repeat):\n"
+    usage = (
+        "Usage: Either write your prompt directly, leave this field empty "
+        "to repeat the conditioning mask or write a command starting with a slash:\n"
+        "- '/q' to quit"
+    )
+
+    while True:
+        img_mask_path = input(user_question)
+
+        if img_mask_path.startswith("/"):
+            if img_mask_path.startswith("/q"):
+                print("Quitting")
+                return None
+            else:
+                if not img_mask_path.startswith("/h"):
+                    print(f"Got invalid command '{img_mask_path}'\n{usage}")
+                print(usage)
+            continue
+
+        if img_mask_path == "":
+            break
+
+        if not os.path.isfile(img_mask_path) or not img_mask_path.lower().endswith(
+            (".jpg", ".jpeg", ".png", ".webp")
+        ):
+            print(f"File '{img_mask_path}' does not exist or is not a valid image file")
+            continue
+        else:
+            with Image.open(img_mask_path) as img:
+                width, height = img.size
+
+            if width % 32 != 0 or height % 32 != 0:
+                print(f"Image dimensions must be divisible by 32, got {width}x{height}")
+                continue
+            else:
+                with Image.open(options.img_cond_path) as img_cond:
+                    img_cond_width, img_cond_height = img_cond.size
+
+                if width != img_cond_width or height != img_cond_height:
+                    print(
+                        f"Mask dimensions must match conditioning image, got {width}x{height} and {img_cond_width}x{img_cond_height}"
+                    )
+                    continue
+
+        options.img_mask_path = img_mask_path
+        break
+
+    return options
+
+
+@torch.inference_mode()
+def main(
+    seed: int | None = None,
+    prompt: str = "a white paper cup",
+    device: str = "cuda" if torch.cuda.is_available() else "cpu",
+    num_steps: int = 50,
+    loop: bool = False,
+    guidance: float = 30.0,
+    offload: bool = False,
+    output_dir: str = "output",
+    add_sampling_metadata: bool = True,
+    img_cond_path: str = "assets/cup.png",
+    img_mask_path: str = "assets/cup_mask.png",
+):
+    """
+    Sample the flux model. Either interactively (set `--loop`) or run for a
+    single image. This demo assumes that the conditioning image and mask have
+    the same shape and that height and width are divisible by 32.
+
+    Args:
+        seed: Set a seed for sampling
+        output_name: where to save the output image, `{idx}` will be replaced
+            by the index of the sample
+        prompt: Prompt used for sampling
+        device: Pytorch device
+        num_steps: number of sampling steps (default 4 for schnell, 50 for guidance distilled)
+        loop: start an interactive session and sample multiple times
+        guidance: guidance value used for guidance distillation
+        add_sampling_metadata: Add the prompt to the image Exif metadata
+        img_cond_path: path to conditioning image (jpeg/png/webp)
+        img_mask_path: path to conditioning mask (jpeg/png/webp
+    """
+    nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=device)
+
+    name = "flux-dev-fill"
+    if name not in configs:
+        available = ", ".join(configs.keys())
+        raise ValueError(f"Got unknown model name: {name}, chose from {available}")
+
+    torch_device = torch.device(device)
+
+    output_name = os.path.join(output_dir, "img_{idx}.jpg")
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+        idx = 0
+    else:
+        fns = [fn for fn in iglob(output_name.format(idx="*")) if re.search(r"img_[0-9]+\.jpg$", fn)]
+        if len(fns) > 0:
+            idx = max(int(fn.split("_")[-1].split(".")[0]) for fn in fns) + 1
+        else:
+            idx = 0
+
+    # init all components
+    t5 = load_t5(torch_device, max_length=128)
+    clip = load_clip(torch_device)
+    model = load_flow_model(name, device="cpu" if offload else torch_device)
+    ae = load_ae(name, device="cpu" if offload else torch_device)
+
+    rng = torch.Generator(device="cpu")
+    with Image.open(img_cond_path) as img:
+        width, height = img.size
+    opts = SamplingOptions(
+        prompt=prompt,
+        width=width,
+        height=height,
+        num_steps=num_steps,
+        guidance=guidance,
+        seed=seed,
+        img_cond_path=img_cond_path,
+        img_mask_path=img_mask_path,
+    )
+
+    if loop:
+        opts = parse_prompt(opts)
+        opts = parse_img_cond_path(opts)
+
+        with Image.open(opts.img_cond_path) as img:
+            width, height = img.size
+        opts.height = height
+        opts.width = width
+
+        opts = parse_img_mask_path(opts)
+
+    while opts is not None:
+        if opts.seed is None:
+            opts.seed = rng.seed()
+        print(f"Generating with seed {opts.seed}:\n{opts.prompt}")
+        t0 = time.perf_counter()
+
+        # prepare input
+        x = get_noise(
+            1,
+            opts.height,
+            opts.width,
+            device=torch_device,
+            dtype=torch.bfloat16,
+            seed=opts.seed,
+        )
+        opts.seed = None
+        if offload:
+            t5, clip, ae = t5.to(torch_device), clip.to(torch_device), ae.to(torch.device)
+        inp = prepare_fill(
+            t5,
+            clip,
+            x,
+            prompt=opts.prompt,
+            ae=ae,
+            img_cond_path=opts.img_cond_path,
+            mask_path=opts.img_mask_path,
+        )
+
+        timesteps = get_schedule(opts.num_steps, inp["img"].shape[1], shift=(name != "flux-schnell"))
+
+        # offload TEs and AE to CPU, load model to gpu
+        if offload:
+            t5, clip, ae = t5.cpu(), clip.cpu(), ae.cpu()
+            torch.cuda.empty_cache()
+            model = model.to(torch_device)
+
+        # denoise initial noise
+        x = denoise(model, **inp, timesteps=timesteps, guidance=opts.guidance)
+
+        # offload model, load autoencoder to gpu
+        if offload:
+            model.cpu()
+            torch.cuda.empty_cache()
+            ae.decoder.to(x.device)
+
+        # decode latents to pixel space
+        x = unpack(x.float(), opts.height, opts.width)
+        with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
+            x = ae.decode(x)
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        t1 = time.perf_counter()
+        print(f"Done in {t1 - t0:.1f}s")
+
+        idx = save_image(nsfw_classifier, name, output_name, idx, x, add_sampling_metadata, prompt)
+
+        if loop:
+            print("-" * 80)
+            opts = parse_prompt(opts)
+            opts = parse_img_cond_path(opts)
+
+            with Image.open(opts.img_cond_path) as img:
+                width, height = img.size
+            opts.height = height
+            opts.width = width
+
+            opts = parse_img_mask_path(opts)
+        else:
+            opts = None
+
+
+def app():
+    Fire(main)
+
+
+if __name__ == "__main__":
+    app()

data/flux/src/flux/cli_redux.py

@@ -0,0 +1,279 @@
+import os
+import re
+import time
+from dataclasses import dataclass
+from glob import iglob
+
+import torch
+from fire import Fire
+from transformers import pipeline
+
+from flux.modules.image_embedders import ReduxImageEncoder
+from flux.sampling import denoise, get_noise, get_schedule, prepare_redux, unpack
+from flux.util import configs, load_ae, load_clip, load_flow_model, load_t5, save_image
+
+
+@dataclass
+class SamplingOptions:
+    prompt: str
+    width: int
+    height: int
+    num_steps: int
+    guidance: float
+    seed: int | None
+    img_cond_path: str
+
+
+def parse_prompt(options: SamplingOptions) -> SamplingOptions | None:
+    user_question = "Write /h for help, /q to quit and leave empty to repeat):\n"
+    usage = (
+        "Usage: Leave this field empty to do nothing "
+        "or write a command starting with a slash:\n"
+        "- '/w <width>' will set the width of the generated image\n"
+        "- '/h <height>' will set the height of the generated image\n"
+        "- '/s <seed>' sets the next seed\n"
+        "- '/g <guidance>' sets the guidance (flux-dev only)\n"
+        "- '/n <steps>' sets the number of steps\n"
+        "- '/q' to quit"
+    )
+
+    while (prompt := input(user_question)).startswith("/"):
+        if prompt.startswith("/w"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, width = prompt.split()
+            options.width = 16 * (int(width) // 16)
+            print(
+                f"Setting resolution to {options.width} x {options.height} "
+                f"({options.height *options.width/1e6:.2f}MP)"
+            )
+        elif prompt.startswith("/h"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, height = prompt.split()
+            options.height = 16 * (int(height) // 16)
+            print(
+                f"Setting resolution to {options.width} x {options.height} "
+                f"({options.height *options.width/1e6:.2f}MP)"
+            )
+        elif prompt.startswith("/g"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, guidance = prompt.split()
+            options.guidance = float(guidance)
+            print(f"Setting guidance to {options.guidance}")
+        elif prompt.startswith("/s"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, seed = prompt.split()
+            options.seed = int(seed)
+            print(f"Setting seed to {options.seed}")
+        elif prompt.startswith("/n"):
+            if prompt.count(" ") != 1:
+                print(f"Got invalid command '{prompt}'\n{usage}")
+                continue
+            _, steps = prompt.split()
+            options.num_steps = int(steps)
+            print(f"Setting number of steps to {options.num_steps}")
+        elif prompt.startswith("/q"):
+            print("Quitting")
+            return None
+        else:
+            if not prompt.startswith("/h"):
+                print(f"Got invalid command '{prompt}'\n{usage}")
+            print(usage)
+    return options
+
+
+def parse_img_cond_path(options: SamplingOptions | None) -> SamplingOptions | None:
+    if options is None:
+        return None
+
+    user_question = "Next conditioning image (write /h for help, /q to quit and leave empty to repeat):\n"
+    usage = (
+        "Usage: Either write your prompt directly, leave this field empty "
+        "to repeat the conditioning image or write a command starting with a slash:\n"
+        "- '/q' to quit"
+    )
+
+    while True:
+        img_cond_path = input(user_question)
+
+        if img_cond_path.startswith("/"):
+            if img_cond_path.startswith("/q"):
+                print("Quitting")
+                return None
+            else:
+                if not img_cond_path.startswith("/h"):
+                    print(f"Got invalid command '{img_cond_path}'\n{usage}")
+                print(usage)
+            continue
+
+        if img_cond_path == "":
+            break
+
+        if not os.path.isfile(img_cond_path) or not img_cond_path.lower().endswith(
+            (".jpg", ".jpeg", ".png", ".webp")
+        ):
+            print(f"File '{img_cond_path}' does not exist or is not a valid image file")
+            continue
+
+        options.img_cond_path = img_cond_path
+        break
+
+    return options
+
+
+@torch.inference_mode()
+def main(
+    name: str = "flux-dev",
+    width: int = 1360,
+    height: int = 768,
+    seed: int | None = None,
+    device: str = "cuda" if torch.cuda.is_available() else "cpu",
+    num_steps: int | None = None,
+    loop: bool = False,
+    guidance: float = 2.5,
+    offload: bool = False,
+    output_dir: str = "output",
+    add_sampling_metadata: bool = True,
+    img_cond_path: str = "assets/robot.webp",
+):
+    """
+    Sample the flux model. Either interactively (set `--loop`) or run for a
+    single image.
+
+    Args:
+        name: Name of the model to load
+        height: height of the sample in pixels (should be a multiple of 16)
+        width: width of the sample in pixels (should be a multiple of 16)
+        seed: Set a seed for sampling
+        output_name: where to save the output image, `{idx}` will be replaced
+            by the index of the sample
+        prompt: Prompt used for sampling
+        device: Pytorch device
+        num_steps: number of sampling steps (default 4 for schnell, 50 for guidance distilled)
+        loop: start an interactive session and sample multiple times
+        guidance: guidance value used for guidance distillation
+        add_sampling_metadata: Add the prompt to the image Exif metadata
+        img_cond_path: path to conditioning image (jpeg/png/webp)
+    """
+    nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=device)
+
+    if name not in configs:
+        available = ", ".join(configs.keys())
+        raise ValueError(f"Got unknown model name: {name}, chose from {available}")
+
+    torch_device = torch.device(device)
+    if num_steps is None:
+        num_steps = 4 if name == "flux-schnell" else 50
+
+    output_name = os.path.join(output_dir, "img_{idx}.jpg")
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+        idx = 0
+    else:
+        fns = [fn for fn in iglob(output_name.format(idx="*")) if re.search(r"img_[0-9]+\.jpg$", fn)]
+        if len(fns) > 0:
+            idx = max(int(fn.split("_")[-1].split(".")[0]) for fn in fns) + 1
+        else:
+            idx = 0
+
+    # init all components
+    t5 = load_t5(torch_device, max_length=256 if name == "flux-schnell" else 512)
+    clip = load_clip(torch_device)
+    model = load_flow_model(name, device="cpu" if offload else torch_device)
+    ae = load_ae(name, device="cpu" if offload else torch_device)
+    img_embedder = ReduxImageEncoder(torch_device)
+
+    rng = torch.Generator(device="cpu")
+    prompt = ""
+    opts = SamplingOptions(
+        prompt=prompt,
+        width=width,
+        height=height,
+        num_steps=num_steps,
+        guidance=guidance,
+        seed=seed,
+        img_cond_path=img_cond_path,
+    )
+
+    if loop:
+        opts = parse_prompt(opts)
+        opts = parse_img_cond_path(opts)
+
+    while opts is not None:
+        if opts.seed is None:
+            opts.seed = rng.seed()
+        print(f"Generating with seed {opts.seed}:\n{opts.prompt}")
+        t0 = time.perf_counter()
+
+        # prepare input
+        x = get_noise(
+            1,
+            opts.height,
+            opts.width,
+            device=torch_device,
+            dtype=torch.bfloat16,
+            seed=opts.seed,
+        )
+        opts.seed = None
+        if offload:
+            ae = ae.cpu()
+            torch.cuda.empty_cache()
+            t5, clip = t5.to(torch_device), clip.to(torch_device)
+        inp = prepare_redux(
+            t5,
+            clip,
+            x,
+            prompt=opts.prompt,
+            encoder=img_embedder,
+            img_cond_path=opts.img_cond_path,
+        )
+        timesteps = get_schedule(opts.num_steps, inp["img"].shape[1], shift=(name != "flux-schnell"))
+
+        # offload TEs to CPU, load model to gpu
+        if offload:
+            t5, clip = t5.cpu(), clip.cpu()
+            torch.cuda.empty_cache()
+            model = model.to(torch_device)
+
+        # denoise initial noise
+        x = denoise(model, **inp, timesteps=timesteps, guidance=opts.guidance)
+
+        # offload model, load autoencoder to gpu
+        if offload:
+            model.cpu()
+            torch.cuda.empty_cache()
+            ae.decoder.to(x.device)
+
+        # decode latents to pixel space
+        x = unpack(x.float(), opts.height, opts.width)
+        with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
+            x = ae.decode(x)
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        t1 = time.perf_counter()
+        print(f"Done in {t1 - t0:.1f}s")
+
+        idx = save_image(nsfw_classifier, name, output_name, idx, x, add_sampling_metadata, prompt)
+
+        if loop:
+            print("-" * 80)
+            opts = parse_prompt(opts)
+            opts = parse_img_cond_path(opts)
+        else:
+            opts = None
+
+
+def app():
+    Fire(main)
+
+
+if __name__ == "__main__":
+    app()

data/flux/src/flux/math.py

@@ -0,0 +1,30 @@
+import torch
+from einops import rearrange
+from torch import Tensor
+
+
+def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
+    q, k = apply_rope(q, k, pe)
+
+    x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+    x = rearrange(x, "B H L D -> B L (H D)")
+
+    return x
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    assert dim % 2 == 0
+    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
+    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+    return out.float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)

data/flux/src/flux/model.py

@@ -0,0 +1,143 @@
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor, nn
+
+from flux.modules.layers import (
+    DoubleStreamBlock,
+    EmbedND,
+    LastLayer,
+    MLPEmbedder,
+    SingleStreamBlock,
+    timestep_embedding,
+)
+from flux.modules.lora import LinearLora, replace_linear_with_lora
+
+
+@dataclass
+class FluxParams:
+    in_channels: int
+    out_channels: int
+    vec_in_dim: int
+    context_in_dim: int
+    hidden_size: int
+    mlp_ratio: float
+    num_heads: int
+    depth: int
+    depth_single_blocks: int
+    axes_dim: list[int]
+    theta: int
+    qkv_bias: bool
+    guidance_embed: bool
+
+
+class Flux(nn.Module):
+    """
+    Transformer model for flow matching on sequences.
+    """
+
+    def __init__(self, params: FluxParams):
+        super().__init__()
+
+        self.params = params
+        self.in_channels = params.in_channels
+        self.out_channels = params.out_channels
+        if params.hidden_size % params.num_heads != 0:
+            raise ValueError(
+                f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
+            )
+        pe_dim = params.hidden_size // params.num_heads
+        if sum(params.axes_dim) != pe_dim:
+            raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
+        self.hidden_size = params.hidden_size
+        self.num_heads = params.num_heads
+        self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
+        self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
+        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
+        self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
+        self.guidance_in = (
+            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity()
+        )
+        self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)
+
+        self.double_blocks = nn.ModuleList(
+            [
+                DoubleStreamBlock(
+                    self.hidden_size,
+                    self.num_heads,
+                    mlp_ratio=params.mlp_ratio,
+                    qkv_bias=params.qkv_bias,
+                )
+                for _ in range(params.depth)
+            ]
+        )
+
+        self.single_blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio)
+                for _ in range(params.depth_single_blocks)
+            ]
+        )
+
+        self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
+
+    def forward(
+        self,
+        img: Tensor,
+        img_ids: Tensor,
+        txt: Tensor,
+        txt_ids: Tensor,
+        timesteps: Tensor,
+        y: Tensor,
+        guidance: Tensor | None = None,
+    ) -> Tensor:
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        # running on sequences img
+        img = self.img_in(img)
+        vec = self.time_in(timestep_embedding(timesteps, 256))
+        if self.params.guidance_embed:
+            if guidance is None:
+                raise ValueError("Didn't get guidance strength for guidance distilled model.")
+            vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
+        vec = vec + self.vector_in(y)
+        txt = self.txt_in(txt)
+
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        pe = self.pe_embedder(ids)
+
+        for block in self.double_blocks:
+            img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
+
+        img = torch.cat((txt, img), 1)
+        for block in self.single_blocks:
+            img = block(img, vec=vec, pe=pe)
+        img = img[:, txt.shape[1] :, ...]
+
+        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
+        return img
+
+
+class FluxLoraWrapper(Flux):
+    def __init__(
+        self,
+        lora_rank: int = 128,
+        lora_scale: float = 1.0,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+
+        self.lora_rank = lora_rank
+
+        replace_linear_with_lora(
+            self,
+            max_rank=lora_rank,
+            scale=lora_scale,
+        )
+
+    def set_lora_scale(self, scale: float) -> None:
+        for module in self.modules():
+            if isinstance(module, LinearLora):
+                module.set_scale(scale=scale)

data/flux/src/flux/modules/autoencoder.py

@@ -0,0 +1,312 @@
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+
+
+@dataclass
+class AutoEncoderParams:
+    resolution: int
+    in_channels: int
+    ch: int
+    out_ch: int
+    ch_mult: list[int]
+    num_res_blocks: int
+    z_channels: int
+    scale_factor: float
+    shift_factor: float
+
+
+def swish(x: Tensor) -> Tensor:
+    return x * torch.sigmoid(x)
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        b, c, h, w = q.shape
+        q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
+        k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
+        v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
+        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
+
+        return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+
+        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if self.in_channels != self.out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = swish(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = swish(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        # no asymmetric padding in torch conv, must do it ourselves
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: Tensor):
+        pad = (0, 1, 0, 1)
+        x = nn.functional.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor):
+        x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        x = self.conv(x)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        # downsampling
+        self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        block_in = self.ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # end
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        # end
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        ch: int,
+        out_ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.ffactor = 2 ** (self.num_resolutions - 1)
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+
+        # z to block_in
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: Tensor) -> Tensor:
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        return h
+
+
+class DiagonalGaussian(nn.Module):
+    def __init__(self, sample: bool = True, chunk_dim: int = 1):
+        super().__init__()
+        self.sample = sample
+        self.chunk_dim = chunk_dim
+
+    def forward(self, z: Tensor) -> Tensor:
+        mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim)
+        if self.sample:
+            std = torch.exp(0.5 * logvar)
+            return mean + std * torch.randn_like(mean)
+        else:
+            return mean
+
+
+class AutoEncoder(nn.Module):
+    def __init__(self, params: AutoEncoderParams):
+        super().__init__()
+        self.encoder = Encoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.decoder = Decoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            out_ch=params.out_ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.reg = DiagonalGaussian()
+
+        self.scale_factor = params.scale_factor
+        self.shift_factor = params.shift_factor
+
+    def encode(self, x: Tensor) -> Tensor:
+        z = self.reg(self.encoder(x))
+        z = self.scale_factor * (z - self.shift_factor)
+        return z
+
+    def decode(self, z: Tensor) -> Tensor:
+        z = z / self.scale_factor + self.shift_factor
+        return self.decoder(z)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.decode(self.encode(x))

data/flux/src/flux/modules/conditioner.py

@@ -0,0 +1,37 @@
+from torch import Tensor, nn
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
+
+
+class HFEmbedder(nn.Module):
+    def __init__(self, version: str, max_length: int, **hf_kwargs):
+        super().__init__()
+        self.is_clip = version.startswith("openai")
+        self.max_length = max_length
+        self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
+
+        if self.is_clip:
+            self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length)
+            self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs)
+        else:
+            self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length)
+            self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs)
+
+        self.hf_module = self.hf_module.eval().requires_grad_(False)
+
+    def forward(self, text: list[str]) -> Tensor:
+        batch_encoding = self.tokenizer(
+            text,
+            truncation=True,
+            max_length=self.max_length,
+            return_length=False,
+            return_overflowing_tokens=False,
+            padding="max_length",
+            return_tensors="pt",
+        )
+
+        outputs = self.hf_module(
+            input_ids=batch_encoding["input_ids"].to(self.hf_module.device),
+            attention_mask=None,
+            output_hidden_states=False,
+        )
+        return outputs[self.output_key]

data/flux/src/flux/modules/image_embedders.py

@@ -0,0 +1,103 @@
+import os
+
+import cv2
+import numpy as np
+import torch
+from einops import rearrange, repeat
+from PIL import Image
+from safetensors.torch import load_file as load_sft
+from torch import nn
+from transformers import AutoModelForDepthEstimation, AutoProcessor, SiglipImageProcessor, SiglipVisionModel
+
+from flux.util import print_load_warning
+
+
+class DepthImageEncoder:
+    depth_model_name = "LiheYoung/depth-anything-large-hf"
+
+    def __init__(self, device):
+        self.device = device
+        self.depth_model = AutoModelForDepthEstimation.from_pretrained(self.depth_model_name).to(device)
+        self.processor = AutoProcessor.from_pretrained(self.depth_model_name)
+
+    def __call__(self, img: torch.Tensor) -> torch.Tensor:
+        hw = img.shape[-2:]
+
+        img = torch.clamp(img, -1.0, 1.0)
+        img_byte = ((img + 1.0) * 127.5).byte()
+
+        img = self.processor(img_byte, return_tensors="pt")["pixel_values"]
+        depth = self.depth_model(img.to(self.device)).predicted_depth
+        depth = repeat(depth, "b h w -> b 3 h w")
+        depth = torch.nn.functional.interpolate(depth, hw, mode="bicubic", antialias=True)
+
+        depth = depth / 127.5 - 1.0
+        return depth
+
+
+class CannyImageEncoder:
+    def __init__(
+        self,
+        device,
+        min_t: int = 50,
+        max_t: int = 200,
+    ):
+        self.device = device
+        self.min_t = min_t
+        self.max_t = max_t
+
+    def __call__(self, img: torch.Tensor) -> torch.Tensor:
+        assert img.shape[0] == 1, "Only batch size 1 is supported"
+
+        img = rearrange(img[0], "c h w -> h w c")
+        img = torch.clamp(img, -1.0, 1.0)
+        img_np = ((img + 1.0) * 127.5).numpy().astype(np.uint8)
+
+        # Apply Canny edge detection
+        canny = cv2.Canny(img_np, self.min_t, self.max_t)
+
+        # Convert back to torch tensor and reshape
+        canny = torch.from_numpy(canny).float() / 127.5 - 1.0
+        canny = rearrange(canny, "h w -> 1 1 h w")
+        canny = repeat(canny, "b 1 ... -> b 3 ...")
+        return canny.to(self.device)
+
+
+class ReduxImageEncoder(nn.Module):
+    siglip_model_name = "google/siglip-so400m-patch14-384"
+
+    def __init__(
+        self,
+        device,
+        redux_dim: int = 1152,
+        txt_in_features: int = 4096,
+        redux_path: str | None = os.getenv("FLUX_REDUX"),
+        dtype=torch.bfloat16,
+    ) -> None:
+        assert redux_path is not None, "Redux path must be provided"
+
+        super().__init__()
+
+        self.redux_dim = redux_dim
+        self.device = device if isinstance(device, torch.device) else torch.device(device)
+        self.dtype = dtype
+
+        with self.device:
+            self.redux_up = nn.Linear(redux_dim, txt_in_features * 3, dtype=dtype)
+            self.redux_down = nn.Linear(txt_in_features * 3, txt_in_features, dtype=dtype)
+
+            sd = load_sft(redux_path, device=str(device))
+            missing, unexpected = self.load_state_dict(sd, strict=False, assign=True)
+            print_load_warning(missing, unexpected)
+
+            self.siglip = SiglipVisionModel.from_pretrained(self.siglip_model_name).to(dtype=dtype)
+        self.normalize = SiglipImageProcessor.from_pretrained(self.siglip_model_name)
+
+    def __call__(self, x: Image.Image) -> torch.Tensor:
+        imgs = self.normalize.preprocess(images=[x], do_resize=True, return_tensors="pt", do_convert_rgb=True)
+
+        _encoded_x = self.siglip(**imgs.to(device=self.device, dtype=self.dtype)).last_hidden_state
+
+        projected_x = self.redux_down(nn.functional.silu(self.redux_up(_encoded_x)))
+
+        return projected_x

data/flux/src/flux/modules/layers.py

@@ -0,0 +1,253 @@
+import math
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+
+from flux.math import attention, rope
+
+
+class EmbedND(nn.Module):
+    def __init__(self, dim: int, theta: int, axes_dim: list[int]):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        n_axes = ids.shape[-1]
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
+            dim=-3,
+        )
+
+        return emb.unsqueeze(1)
+
+
+def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
+    """
+    Create sinusoidal timestep embeddings.
+    :param t: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an (N, D) Tensor of positional embeddings.
+    """
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
+        t.device
+    )
+
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+
+class MLPEmbedder(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int):
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+        self.silu = nn.SiLU()
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor):
+        x_dtype = x.dtype
+        x = x.float()
+        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
+        return (x * rrms).to(dtype=x_dtype) * self.scale
+
+
+class QKNorm(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        q = self.query_norm(q)
+        k = self.key_norm(k)
+        return q.to(v), k.to(v)
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.norm = QKNorm(head_dim)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x: Tensor, pe: Tensor) -> Tensor:
+        qkv = self.qkv(x)
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.norm(q, k, v)
+        x = attention(q, k, v, pe=pe)
+        x = self.proj(x)
+        return x
+
+
+@dataclass
+class ModulationOut:
+    shift: Tensor
+    scale: Tensor
+    gate: Tensor
+
+
+class Modulation(nn.Module):
+    def __init__(self, dim: int, double: bool):
+        super().__init__()
+        self.is_double = double
+        self.multiplier = 6 if double else 3
+        self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
+
+    def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]:
+        out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)
+
+        return (
+            ModulationOut(*out[:3]),
+            ModulationOut(*out[3:]) if self.is_double else None,
+        )
+
+
+class DoubleStreamBlock(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):
+        super().__init__()
+
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.img_mod = Modulation(hidden_size, double=True)
+        self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
+
+        self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
+        )
+
+        self.txt_mod = Modulation(hidden_size, double=True)
+        self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
+
+        self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.txt_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
+        )
+
+    def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:
+        img_mod1, img_mod2 = self.img_mod(vec)
+        txt_mod1, txt_mod2 = self.txt_mod(vec)
+
+        # prepare image for attention
+        img_modulated = self.img_norm1(img)
+        img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
+        img_qkv = self.img_attn.qkv(img_modulated)
+        img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
+
+        # prepare txt for attention
+        txt_modulated = self.txt_norm1(txt)
+        txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
+        txt_qkv = self.txt_attn.qkv(txt_modulated)
+        txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
+
+        # run actual attention
+        q = torch.cat((txt_q, img_q), dim=2)
+        k = torch.cat((txt_k, img_k), dim=2)
+        v = torch.cat((txt_v, img_v), dim=2)
+
+        attn = attention(q, k, v, pe=pe)
+        txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
+
+        # calculate the img bloks
+        img = img + img_mod1.gate * self.img_attn.proj(img_attn)
+        img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
+
+        # calculate the txt bloks
+        txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
+        txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
+        return img, txt
+
+
+class SingleStreamBlock(nn.Module):
+    """
+    A DiT block with parallel linear layers as described in
+    https://arxiv.org/abs/2302.05442 and adapted modulation interface.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qk_scale: float | None = None,
+    ):
+        super().__init__()
+        self.hidden_dim = hidden_size
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        # qkv and mlp_in
+        self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
+        # proj and mlp_out
+        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
+
+        self.norm = QKNorm(head_dim)
+
+        self.hidden_size = hidden_size
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+
+        self.mlp_act = nn.GELU(approximate="tanh")
+        self.modulation = Modulation(hidden_size, double=False)
+
+    def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
+        mod, _ = self.modulation(vec)
+        x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
+        qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
+
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.norm(q, k, v)
+
+        # compute attention
+        attn = attention(q, k, v, pe=pe)
+        # compute activation in mlp stream, cat again and run second linear layer
+        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
+        return x + mod.gate * output
+
+
+class LastLayer(nn.Module):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+
+    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+        shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
+        x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
+        x = self.linear(x)
+        return x

data/flux/src/flux/modules/lora.py

@@ -0,0 +1,94 @@
+import torch
+from torch import nn
+
+
+def replace_linear_with_lora(
+    module: nn.Module,
+    max_rank: int,
+    scale: float = 1.0,
+) -> None:
+    for name, child in module.named_children():
+        if isinstance(child, nn.Linear):
+            new_lora = LinearLora(
+                in_features=child.in_features,
+                out_features=child.out_features,
+                bias=child.bias,
+                rank=max_rank,
+                scale=scale,
+                dtype=child.weight.dtype,
+                device=child.weight.device,
+            )
+
+            new_lora.weight = child.weight
+            new_lora.bias = child.bias if child.bias is not None else None
+
+            setattr(module, name, new_lora)
+        else:
+            replace_linear_with_lora(
+                module=child,
+                max_rank=max_rank,
+                scale=scale,
+            )
+
+
+class LinearLora(nn.Linear):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool,
+        rank: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        lora_bias: bool = True,
+        scale: float = 1.0,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            in_features=in_features,
+            out_features=out_features,
+            bias=bias is not None,
+            device=device,
+            dtype=dtype,
+            *args,
+            **kwargs,
+        )
+
+        assert isinstance(scale, float), "scale must be a float"
+
+        self.scale = scale
+        self.rank = rank
+        self.lora_bias = lora_bias
+        self.dtype = dtype
+        self.device = device
+
+        if rank > (new_rank := min(self.out_features, self.in_features)):
+            self.rank = new_rank
+
+        self.lora_A = nn.Linear(
+            in_features=in_features,
+            out_features=self.rank,
+            bias=False,
+            dtype=dtype,
+            device=device,
+        )
+        self.lora_B = nn.Linear(
+            in_features=self.rank,
+            out_features=out_features,
+            bias=self.lora_bias,
+            dtype=dtype,
+            device=device,
+        )
+
+    def set_scale(self, scale: float) -> None:
+        assert isinstance(scale, float), "scalar value must be a float"
+        self.scale = scale
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        base_out = super().forward(input)
+
+        _lora_out_B = self.lora_B(self.lora_A(input))
+        lora_update = _lora_out_B * self.scale
+
+        return base_out + lora_update

data/flux/src/flux/sampling.py

@@ -0,0 +1,282 @@
+import math
+from typing import Callable
+
+import numpy as np
+import torch
+from einops import rearrange, repeat
+from PIL import Image
+from torch import Tensor
+
+from .model import Flux
+from .modules.autoencoder import AutoEncoder
+from .modules.conditioner import HFEmbedder
+from .modules.image_embedders import CannyImageEncoder, DepthImageEncoder, ReduxImageEncoder
+
+
+def get_noise(
+    num_samples: int,
+    height: int,
+    width: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    seed: int,
+):
+    return torch.randn(
+        num_samples,
+        16,
+        # allow for packing
+        2 * math.ceil(height / 16),
+        2 * math.ceil(width / 16),
+        device=device,
+        dtype=dtype,
+        generator=torch.Generator(device=device).manual_seed(seed),
+    )
+
+
+def prepare(t5: HFEmbedder, clip: HFEmbedder, img: Tensor, prompt: str | list[str]) -> dict[str, Tensor]:
+    bs, c, h, w = img.shape
+    if bs == 1 and not isinstance(prompt, str):
+        bs = len(prompt)
+
+    img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+    if img.shape[0] == 1 and bs > 1:
+        img = repeat(img, "1 ... -> bs ...", bs=bs)
+
+    img_ids = torch.zeros(h // 2, w // 2, 3)
+    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
+    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
+    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+
+    if isinstance(prompt, str):
+        prompt = [prompt]
+    txt = t5(prompt)
+    if txt.shape[0] == 1 and bs > 1:
+        txt = repeat(txt, "1 ... -> bs ...", bs=bs)
+    txt_ids = torch.zeros(bs, txt.shape[1], 3)
+
+    vec = clip(prompt)
+    if vec.shape[0] == 1 and bs > 1:
+        vec = repeat(vec, "1 ... -> bs ...", bs=bs)
+
+    return {
+        "img": img,
+        "img_ids": img_ids.to(img.device),
+        "txt": txt.to(img.device),
+        "txt_ids": txt_ids.to(img.device),
+        "vec": vec.to(img.device),
+    }
+
+
+def prepare_control(
+    t5: HFEmbedder,
+    clip: HFEmbedder,
+    img: Tensor,
+    prompt: str | list[str],
+    ae: AutoEncoder,
+    encoder: DepthImageEncoder | CannyImageEncoder,
+    img_cond_path: str,
+) -> dict[str, Tensor]:
+    # load and encode the conditioning image
+    bs, _, h, w = img.shape
+    if bs == 1 and not isinstance(prompt, str):
+        bs = len(prompt)
+
+    img_cond = Image.open(img_cond_path).convert("RGB")
+
+    width = w * 8
+    height = h * 8
+    img_cond = img_cond.resize((width, height), Image.LANCZOS)
+    img_cond = np.array(img_cond)
+    img_cond = torch.from_numpy(img_cond).float() / 127.5 - 1.0
+    img_cond = rearrange(img_cond, "h w c -> 1 c h w")
+
+    with torch.no_grad():
+        img_cond = encoder(img_cond)
+        img_cond = ae.encode(img_cond)
+
+    img_cond = img_cond.to(torch.bfloat16)
+    img_cond = rearrange(img_cond, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+    if img_cond.shape[0] == 1 and bs > 1:
+        img_cond = repeat(img_cond, "1 ... -> bs ...", bs=bs)
+
+    return_dict = prepare(t5, clip, img, prompt)
+    return_dict["img_cond"] = img_cond
+    return return_dict
+
+
+def prepare_fill(
+    t5: HFEmbedder,
+    clip: HFEmbedder,
+    img: Tensor,
+    prompt: str | list[str],
+    ae: AutoEncoder,
+    img_cond_path: str,
+    mask_path: str,
+) -> dict[str, Tensor]:
+    # load and encode the conditioning image and the mask
+    bs, _, _, _ = img.shape
+    if bs == 1 and not isinstance(prompt, str):
+        bs = len(prompt)
+
+    img_cond = Image.open(img_cond_path).convert("RGB")
+    img_cond = np.array(img_cond)
+    img_cond = torch.from_numpy(img_cond).float() / 127.5 - 1.0
+    img_cond = rearrange(img_cond, "h w c -> 1 c h w")
+
+    mask = Image.open(mask_path).convert("L")
+    mask = np.array(mask)
+    mask = torch.from_numpy(mask).float() / 255.0
+    mask = rearrange(mask, "h w -> 1 1 h w")
+
+    with torch.no_grad():
+        img_cond = img_cond.to(img.device)
+        mask = mask.to(img.device)
+        img_cond = img_cond * (1 - mask)
+        img_cond = ae.encode(img_cond)
+        mask = mask[:, 0, :, :]
+        mask = mask.to(torch.bfloat16)
+        mask = rearrange(
+            mask,
+            "b (h ph) (w pw) -> b (ph pw) h w",
+            ph=8,
+            pw=8,
+        )
+        mask = rearrange(mask, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+        if mask.shape[0] == 1 and bs > 1:
+            mask = repeat(mask, "1 ... -> bs ...", bs=bs)
+
+    img_cond = img_cond.to(torch.bfloat16)
+    img_cond = rearrange(img_cond, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+    if img_cond.shape[0] == 1 and bs > 1:
+        img_cond = repeat(img_cond, "1 ... -> bs ...", bs=bs)
+
+    img_cond = torch.cat((img_cond, mask), dim=-1)
+
+    return_dict = prepare(t5, clip, img, prompt)
+    return_dict["img_cond"] = img_cond.to(img.device)
+    return return_dict
+
+
+def prepare_redux(
+    t5: HFEmbedder,
+    clip: HFEmbedder,
+    img: Tensor,
+    prompt: str | list[str],
+    encoder: ReduxImageEncoder,
+    img_cond_path: str,
+) -> dict[str, Tensor]:
+    bs, _, h, w = img.shape
+    if bs == 1 and not isinstance(prompt, str):
+        bs = len(prompt)
+
+    img_cond = Image.open(img_cond_path).convert("RGB")
+    with torch.no_grad():
+        img_cond = encoder(img_cond)
+
+    img_cond = img_cond.to(torch.bfloat16)
+    if img_cond.shape[0] == 1 and bs > 1:
+        img_cond = repeat(img_cond, "1 ... -> bs ...", bs=bs)
+
+    img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+    if img.shape[0] == 1 and bs > 1:
+        img = repeat(img, "1 ... -> bs ...", bs=bs)
+
+    img_ids = torch.zeros(h // 2, w // 2, 3)
+    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
+    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
+    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+
+    if isinstance(prompt, str):
+        prompt = [prompt]
+    txt = t5(prompt)
+    txt = torch.cat((txt, img_cond.to(txt)), dim=-2)
+    if txt.shape[0] == 1 and bs > 1:
+        txt = repeat(txt, "1 ... -> bs ...", bs=bs)
+    txt_ids = torch.zeros(bs, txt.shape[1], 3)
+
+    vec = clip(prompt)
+    if vec.shape[0] == 1 and bs > 1:
+        vec = repeat(vec, "1 ... -> bs ...", bs=bs)
+
+    return {
+        "img": img,
+        "img_ids": img_ids.to(img.device),
+        "txt": txt.to(img.device),
+        "txt_ids": txt_ids.to(img.device),
+        "vec": vec.to(img.device),
+    }
+
+
+def time_shift(mu: float, sigma: float, t: Tensor):
+    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+
+def get_lin_function(
+    x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
+) -> Callable[[float], float]:
+    m = (y2 - y1) / (x2 - x1)
+    b = y1 - m * x1
+    return lambda x: m * x + b
+
+
+def get_schedule(
+    num_steps: int,
+    image_seq_len: int,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+    shift: bool = True,
+) -> list[float]:
+    # extra step for zero
+    timesteps = torch.linspace(1, 0, num_steps + 1)
+
+    # shifting the schedule to favor high timesteps for higher signal images
+    if shift:
+        # estimate mu based on linear estimation between two points
+        mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
+        timesteps = time_shift(mu, 1.0, timesteps)
+
+    return timesteps.tolist()
+
+
+def denoise(
+    model: Flux,
+    # model input
+    img: Tensor,
+    img_ids: Tensor,
+    txt: Tensor,
+    txt_ids: Tensor,
+    vec: Tensor,
+    # sampling parameters
+    timesteps: list[float],
+    guidance: float = 4.0,
+    # extra img tokens
+    img_cond: Tensor | None = None,
+):
+    # this is ignored for schnell
+    guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
+    for t_curr, t_prev in zip(timesteps[:-1], timesteps[1:]):
+        t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
+        pred = model(
+            img=torch.cat((img, img_cond), dim=-1) if img_cond is not None else img,
+            img_ids=img_ids,
+            txt=txt,
+            txt_ids=txt_ids,
+            y=vec,
+            timesteps=t_vec,
+            guidance=guidance_vec,
+        )
+
+        img = img + (t_prev - t_curr) * pred
+
+    return img
+
+
+def unpack(x: Tensor, height: int, width: int) -> Tensor:
+    return rearrange(
+        x,
+        "b (h w) (c ph pw) -> b c (h ph) (w pw)",
+        h=math.ceil(height / 16),
+        w=math.ceil(width / 16),
+        ph=2,
+        pw=2,
+    )

data/flux/src/flux/util.py

@@ -0,0 +1,447 @@
+import os
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from imwatermark import WatermarkEncoder
+from PIL import ExifTags, Image
+from safetensors.torch import load_file as load_sft
+
+from flux.model import Flux, FluxLoraWrapper, FluxParams
+from flux.modules.autoencoder import AutoEncoder, AutoEncoderParams
+from flux.modules.conditioner import HFEmbedder
+
+
+def save_image(
+    nsfw_classifier,
+    name: str,
+    output_name: str,
+    idx: int,
+    x: torch.Tensor,
+    add_sampling_metadata: bool,
+    prompt: str,
+    nsfw_threshold: float = 0.85,
+) -> int:
+    fn = output_name.format(idx=idx)
+    print(f"Saving {fn}")
+    # bring into PIL format and save
+    x = x.clamp(-1, 1)
+    x = embed_watermark(x.float())
+    x = rearrange(x[0], "c h w -> h w c")
+
+    img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
+    nsfw_score = [x["score"] for x in nsfw_classifier(img) if x["label"] == "nsfw"][0]
+
+    if nsfw_score < nsfw_threshold:
+        exif_data = Image.Exif()
+        exif_data[ExifTags.Base.Software] = "AI generated;txt2img;flux"
+        exif_data[ExifTags.Base.Make] = "Black Forest Labs"
+        exif_data[ExifTags.Base.Model] = name
+        if add_sampling_metadata:
+            exif_data[ExifTags.Base.ImageDescription] = prompt
+        img.save(fn, exif=exif_data, quality=95, subsampling=0)
+        idx += 1
+    else:
+        print("Your generated image may contain NSFW content.")
+
+    return idx
+
+
+@dataclass
+class ModelSpec:
+    params: FluxParams
+    ae_params: AutoEncoderParams
+    ckpt_path: str | None
+    lora_path: str | None
+    ae_path: str | None
+    repo_id: str | None
+    repo_flow: str | None
+    repo_ae: str | None
+
+
+configs = {
+    "flux-dev": ModelSpec(
+        repo_id="black-forest-labs/FLUX.1-dev",
+        repo_flow="flux1-dev.safetensors",
+        repo_ae="ae.safetensors",
+        ckpt_path=os.getenv("FLUX_DEV"),
+        lora_path=None,
+        params=FluxParams(
+            in_channels=64,
+            out_channels=64,
+            vec_in_dim=768,
+            context_in_dim=4096,
+            hidden_size=3072,
+            mlp_ratio=4.0,
+            num_heads=24,
+            depth=19,
+            depth_single_blocks=38,
+            axes_dim=[16, 56, 56],
+            theta=10_000,
+            qkv_bias=True,
+            guidance_embed=True,
+        ),
+        ae_path=os.getenv("AE"),
+        ae_params=AutoEncoderParams(
+            resolution=256,
+            in_channels=3,
+            ch=128,
+            out_ch=3,
+            ch_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            z_channels=16,
+            scale_factor=0.3611,
+            shift_factor=0.1159,
+        ),
+    ),
+    "flux-schnell": ModelSpec(
+        repo_id="black-forest-labs/FLUX.1-schnell",
+        repo_flow="flux1-schnell.safetensors",
+        repo_ae="ae.safetensors",
+        ckpt_path=os.getenv("FLUX_SCHNELL"),
+        lora_path=None,
+        params=FluxParams(
+            in_channels=64,
+            out_channels=64,
+            vec_in_dim=768,
+            context_in_dim=4096,
+            hidden_size=3072,
+            mlp_ratio=4.0,
+            num_heads=24,
+            depth=19,
+            depth_single_blocks=38,
+            axes_dim=[16, 56, 56],
+            theta=10_000,
+            qkv_bias=True,
+            guidance_embed=False,
+        ),
+        ae_path=os.getenv("AE"),
+        ae_params=AutoEncoderParams(
+            resolution=256,
+            in_channels=3,
+            ch=128,
+            out_ch=3,
+            ch_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            z_channels=16,
+            scale_factor=0.3611,
+            shift_factor=0.1159,
+        ),
+    ),
+    "flux-dev-canny": ModelSpec(
+        repo_id="black-forest-labs/FLUX.1-Canny-dev",
+        repo_flow="flux1-canny-dev.safetensors",
+        repo_ae="ae.safetensors",
+        ckpt_path=os.getenv("FLUX_DEV_CANNY"),
+        lora_path=None,
+        params=FluxParams(
+            in_channels=128,
+            out_channels=64,
+            vec_in_dim=768,
+            context_in_dim=4096,
+            hidden_size=3072,
+            mlp_ratio=4.0,
+            num_heads=24,
+            depth=19,
+            depth_single_blocks=38,
+            axes_dim=[16, 56, 56],
+            theta=10_000,
+            qkv_bias=True,
+            guidance_embed=True,
+        ),
+        ae_path=os.getenv("AE"),
+        ae_params=AutoEncoderParams(
+            resolution=256,
+            in_channels=3,
+            ch=128,
+            out_ch=3,
+            ch_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            z_channels=16,
+            scale_factor=0.3611,
+            shift_factor=0.1159,
+        ),
+    ),
+    "flux-dev-canny-lora": ModelSpec(
+        repo_id="black-forest-labs/FLUX.1-dev",
+        repo_flow="flux1-dev.safetensors",
+        repo_ae="ae.safetensors",
+        ckpt_path=os.getenv("FLUX_DEV"),
+        lora_path=os.getenv("FLUX_DEV_CANNY_LORA"),
+        params=FluxParams(
+            in_channels=128,
+            out_channels=64,
+            vec_in_dim=768,
+            context_in_dim=4096,
+            hidden_size=3072,
+            mlp_ratio=4.0,
+            num_heads=24,
+            depth=19,
+            depth_single_blocks=38,
+            axes_dim=[16, 56, 56],
+            theta=10_000,
+            qkv_bias=True,
+            guidance_embed=True,
+        ),
+        ae_path=os.getenv("AE"),
+        ae_params=AutoEncoderParams(
+            resolution=256,
+            in_channels=3,
+            ch=128,
+            out_ch=3,
+            ch_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            z_channels=16,
+            scale_factor=0.3611,
+            shift_factor=0.1159,
+        ),
+    ),
+    "flux-dev-depth": ModelSpec(
+        repo_id="black-forest-labs/FLUX.1-Depth-dev",
+        repo_flow="flux1-depth-dev.safetensors",
+        repo_ae="ae.safetensors",
+        ckpt_path=os.getenv("FLUX_DEV_DEPTH"),
+        lora_path=None,
+        params=FluxParams(
+            in_channels=128,
+            out_channels=64,
+            vec_in_dim=768,
+            context_in_dim=4096,
+            hidden_size=3072,
+            mlp_ratio=4.0,
+            num_heads=24,
+            depth=19,
+            depth_single_blocks=38,
+            axes_dim=[16, 56, 56],
+            theta=10_000,
+            qkv_bias=True,
+            guidance_embed=True,
+        ),
+        ae_path=os.getenv("AE"),
+        ae_params=AutoEncoderParams(
+            resolution=256,
+            in_channels=3,
+            ch=128,
+            out_ch=3,
+            ch_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            z_channels=16,
+            scale_factor=0.3611,
+            shift_factor=0.1159,
+        ),
+    ),
+    "flux-dev-depth-lora": ModelSpec(
+        repo_id="black-forest-labs/FLUX.1-dev",
+        repo_flow="flux1-dev.safetensors",
+        repo_ae="ae.safetensors",
+        ckpt_path=os.getenv("FLUX_DEV"),
+        lora_path=os.getenv("FLUX_DEV_DEPTH_LORA"),
+        params=FluxParams(
+            in_channels=128,
+            out_channels=64,
+            vec_in_dim=768,
+            context_in_dim=4096,
+            hidden_size=3072,
+            mlp_ratio=4.0,
+            num_heads=24,
+            depth=19,
+            depth_single_blocks=38,
+            axes_dim=[16, 56, 56],
+            theta=10_000,
+            qkv_bias=True,
+            guidance_embed=True,
+        ),
+        ae_path=os.getenv("AE"),
+        ae_params=AutoEncoderParams(
+            resolution=256,
+            in_channels=3,
+            ch=128,
+            out_ch=3,
+            ch_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            z_channels=16,
+            scale_factor=0.3611,
+            shift_factor=0.1159,
+        ),
+    ),
+    "flux-dev-fill": ModelSpec(
+        repo_id="black-forest-labs/FLUX.1-Fill-dev",
+        repo_flow="flux1-fill-dev.safetensors",
+        repo_ae="ae.safetensors",
+        ckpt_path=os.getenv("FLUX_DEV_FILL"),
+        lora_path=None,
+        params=FluxParams(
+            in_channels=384,
+            out_channels=64,
+            vec_in_dim=768,
+            context_in_dim=4096,
+            hidden_size=3072,
+            mlp_ratio=4.0,
+            num_heads=24,
+            depth=19,
+            depth_single_blocks=38,
+            axes_dim=[16, 56, 56],
+            theta=10_000,
+            qkv_bias=True,
+            guidance_embed=True,
+        ),
+        ae_path=os.getenv("AE"),
+        ae_params=AutoEncoderParams(
+            resolution=256,
+            in_channels=3,
+            ch=128,
+            out_ch=3,
+            ch_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            z_channels=16,
+            scale_factor=0.3611,
+            shift_factor=0.1159,
+        ),
+    ),
+}
+
+
+def print_load_warning(missing: list[str], unexpected: list[str]) -> None:
+    if len(missing) > 0 and len(unexpected) > 0:
+        print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+        print("\n" + "-" * 79 + "\n")
+        print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+    elif len(missing) > 0:
+        print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+    elif len(unexpected) > 0:
+        print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+
+
+def load_flow_model(
+    name: str, device: str | torch.device = "cuda", hf_download: bool = True, verbose: bool = False
+) -> Flux:
+    # Loading Flux
+    print("Init model")
+    ckpt_path = configs[name].ckpt_path
+    lora_path = configs[name].lora_path
+    if (
+        ckpt_path is None
+        and configs[name].repo_id is not None
+        and configs[name].repo_flow is not None
+        and hf_download
+    ):
+        ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow)
+
+    with torch.device("meta" if ckpt_path is not None else device):
+        if lora_path is not None:
+            model = FluxLoraWrapper(params=configs[name].params).to(torch.bfloat16)
+        else:
+            model = Flux(configs[name].params).to(torch.bfloat16)
+
+    if ckpt_path is not None:
+        print("Loading checkpoint")
+        # load_sft doesn't support torch.device
+        sd = load_sft(ckpt_path, device=str(device))
+        sd = optionally_expand_state_dict(model, sd)
+        missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
+        if verbose:
+            print_load_warning(missing, unexpected)
+
+    if configs[name].lora_path is not None:
+        print("Loading LoRA")
+        lora_sd = load_sft(configs[name].lora_path, device=str(device))
+        # loading the lora params + overwriting scale values in the norms
+        missing, unexpected = model.load_state_dict(lora_sd, strict=False, assign=True)
+        if verbose:
+            print_load_warning(missing, unexpected)
+    return model
+
+
+def load_t5(device: str | torch.device = "cuda", max_length: int = 512) -> HFEmbedder:
+    # max length 64, 128, 256 and 512 should work (if your sequence is short enough)
+    return HFEmbedder("google/t5-v1_1-xxl", max_length=max_length, torch_dtype=torch.bfloat16).to(device)
+
+
+def load_clip(device: str | torch.device = "cuda") -> HFEmbedder:
+    return HFEmbedder("openai/clip-vit-large-patch14", max_length=77, torch_dtype=torch.bfloat16).to(device)
+
+
+def load_ae(name: str, device: str | torch.device = "cuda", hf_download: bool = True) -> AutoEncoder:
+    ckpt_path = configs[name].ae_path
+    if (
+        ckpt_path is None
+        and configs[name].repo_id is not None
+        and configs[name].repo_ae is not None
+        and hf_download
+    ):
+        ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_ae)
+
+    # Loading the autoencoder
+    print("Init AE")
+    with torch.device("meta" if ckpt_path is not None else device):
+        ae = AutoEncoder(configs[name].ae_params)
+
+    if ckpt_path is not None:
+        sd = load_sft(ckpt_path, device=str(device))
+        missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
+        print_load_warning(missing, unexpected)
+    return ae
+
+
+def optionally_expand_state_dict(model: torch.nn.Module, state_dict: dict) -> dict:
+    """
+    Optionally expand the state dict to match the model's parameters shapes.
+    """
+    for name, param in model.named_parameters():
+        if name in state_dict:
+            if state_dict[name].shape != param.shape:
+                print(
+                    f"Expanding '{name}' with shape {state_dict[name].shape} to model parameter with shape {param.shape}."
+                )
+                # expand with zeros:
+                expanded_state_dict_weight = torch.zeros_like(param, device=state_dict[name].device)
+                slices = tuple(slice(0, dim) for dim in state_dict[name].shape)
+                expanded_state_dict_weight[slices] = state_dict[name]
+                state_dict[name] = expanded_state_dict_weight
+
+    return state_dict
+
+
+class WatermarkEmbedder:
+    def __init__(self, watermark):
+        self.watermark = watermark
+        self.num_bits = len(WATERMARK_BITS)
+        self.encoder = WatermarkEncoder()
+        self.encoder.set_watermark("bits", self.watermark)
+
+    def __call__(self, image: torch.Tensor) -> torch.Tensor:
+        """
+        Adds a predefined watermark to the input image
+
+        Args:
+            image: ([N,] B, RGB, H, W) in range [-1, 1]
+
+        Returns:
+            same as input but watermarked
+        """
+        image = 0.5 * image + 0.5
+        squeeze = len(image.shape) == 4
+        if squeeze:
+            image = image[None, ...]
+        n = image.shape[0]
+        image_np = rearrange((255 * image).detach().cpu(), "n b c h w -> (n b) h w c").numpy()[:, :, :, ::-1]
+        # torch (b, c, h, w) in [0, 1] -> numpy (b, h, w, c) [0, 255]
+        # watermarking libary expects input as cv2 BGR format
+        for k in range(image_np.shape[0]):
+            image_np[k] = self.encoder.encode(image_np[k], "dwtDct")
+        image = torch.from_numpy(rearrange(image_np[:, :, :, ::-1], "(n b) h w c -> n b c h w", n=n)).to(
+            image.device
+        )
+        image = torch.clamp(image / 255, min=0.0, max=1.0)
+        if squeeze:
+            image = image[0]
+        image = 2 * image - 1
+        return image
+
+
+# A fixed 48-bit message that was chosen at random
+WATERMARK_MESSAGE = 0b001010101111111010000111100111001111010100101110
+# bin(x)[2:] gives bits of x as str, use int to convert them to 0/1
+WATERMARK_BITS = [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]
+embed_watermark = WatermarkEmbedder(WATERMARK_BITS)