pyproject.toml to be authoritative source of deps, soft introduce ruff linter/formatter

This commit is contained in:
Anthony Wu 2024-09-19 11:20:06 -07:00
parent d7fe524c92
commit c74279d394
3 changed files with 69 additions and 64 deletions

View File

@ -8,13 +8,13 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
### Philosophy
MFLUX is a line-by-line port of the FLUX implementation in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) library to [Apple MLX](https://github.com/ml-explore/mlx).
MFLUX is a line-by-line port of the FLUX implementation in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) library to [Apple MLX](https://github.com/ml-explore/mlx).
MFLUX is purposefully kept minimal and explicit - Network architectures are hardcoded and no config files are used
except for the tokenizers. The aim is to have a tiny codebase with the single purpose of expressing these models
except for the tokenizers. The aim is to have a tiny codebase with the single purpose of expressing these models
(thereby avoiding too many abstractions). While MFLUX priorities readability over generality and performance, [it can still be quite fast](#image-generation-speed-updated), [and even faster quantized](#quantization).
All models are implemented from scratch in MLX and only the tokenizers are used via the
[Huggingface Transformers](https://github.com/huggingface/transformers) library. Other than that, there are only minimal dependencies
All models are implemented from scratch in MLX and only the tokenizers are used via the
[Huggingface Transformers](https://github.com/huggingface/transformers) library. Other than that, there are only minimal dependencies
like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) for simple image post-processing.
### Models
@ -40,12 +40,18 @@ This creates and activates a virtual environment in the `mflux` folder. After th
```
2. Navigate to the project and set up a virtual environment:
```
cd mflux && python3 -m venv .venv && source .venv/bin/activate
cd mflux
python3 -m venv .venv # alternative: `uv venv`
source .venv/bin/activate
```
3. Install the required dependencies:
3. Install the required dependencies in [development mode](https://setuptools.pypa.io/en/latest/userguide/development_mode.html) as defined in `pyproject.toml`:
```
pip install -r requirements.txt
pip install -e .
```
4. Follow format and lint checks prior to submitting Pull Requests. Install the `ruff` tool as a common resource somewhere else in your system external to this project.
- [`uv tool install ruff`](https://github.com/astral-sh/uv) OR [`pipx install ruff`](https://github.com/pypa/pipx).
- [`ruff check` and `ruff check --fix`](https://github.com/astral-sh/ruff?tab=readme-ov-file#usage) any linter warnings generated from your PR diff.
- [`ruff format`](https://github.com/astral-sh/ruff?tab=readme-ov-file#usage) any added or modified files in your PR diff.
</details>
### Generating an image
@ -64,7 +70,7 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
⚠️ *If the specific model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell or Dev model respectively). See the [quantization](#quantization) section for running compressed versions of the model.* ⚠️
*By default, model files are downloaded to the `.cache` folder within your home directory. For example, in my setup, the path looks like this:*
*By default, model files are downloaded to the `.cache` folder within your home directory. For example, in my setup, the path looks like this:*
```
/Users/filipstrand/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-dev
@ -72,9 +78,9 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
*To change this default behavior, you can do so by modifying the `HF_HOME` environment variable. For more details on how to adjust this setting, please refer to the [Hugging Face documentation](https://huggingface.co/docs/huggingface_hub/en/package_reference/environment_variables)*.
🔒 [FLUX.1-dev currently requires granted access to its Huggingface repo. For troubleshooting, see the issue tracker](https://github.com/filipstrand/mflux/issues/14) 🔒
🔒 [FLUX.1-dev currently requires granted access to its Huggingface repo. For troubleshooting, see the issue tracker](https://github.com/filipstrand/mflux/issues/14) 🔒
#### Full list of Command-Line Arguments
#### Full list of Command-Line Arguments
- **`--prompt`** (required, `str`): Text description of the image to generate.
@ -99,7 +105,7 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
- **`--lora-paths`** (optional, `[str]`, default: `None`): The paths to the [LoRA](#LoRA) weights.
- **`--lora-scales`** (optional, `[float]`, default: `None`): The scale for each respective [LoRA](#LoRA) (will default to `1.0` if not specified and only one LoRA weight is loaded.)
- **`--metadata`** (optional): Exports a `.json` file containing the metadata for the image with the same name. (Even without this flag, the image metadata is saved and can be viewed using `exiftool image.png`)
Or, with the correct python environment active, create and run a separate script like the following:
@ -112,7 +118,7 @@ from mflux.config.config import Config
flux = Flux1.from_alias(
alias="schnell", # "schnell" or "dev"
quantize=8, # 4 or 8
)
)
# Generate an image
image = flux.generate_image(
@ -132,7 +138,7 @@ For more options on how to configure MFLUX, please see [generate.py](src/mflux/g
### Image generation speed (updated)
These numbers are based on the non-quantized `schnell` model, with the configuration provided in the code snippet below.
These numbers are based on the non-quantized `schnell` model, with the configuration provided in the code snippet below.
To time your machine, run the following:
```
time mflux-generate \
@ -156,20 +162,20 @@ time mflux-generate \
| 2021 M1 Pro (32GB) | @filipstrand | ~160s | |
| 2023 M2 Max (32GB) | @filipstrand | ~70s | |
*Note that these numbers includes starting the application from scratch, which means doing model i/o, setting/quantizing weights etc.
*Note that these numbers includes starting the application from scratch, which means doing model i/o, setting/quantizing weights etc.
If we assume that the model is already loaded, you can inspect the image metadata using `exiftool image.png` and see the total duration of the denoising loop (excluding text embedding).*
### Equivalent to Diffusers implementation
### Equivalent to Diffusers implementation
There is only a single source of randomness when generating an image: The initial latent array.
In this implementation, this initial latent is fully deterministically controlled by the input `seed` parameter.
There is only a single source of randomness when generating an image: The initial latent array.
In this implementation, this initial latent is fully deterministically controlled by the input `seed` parameter.
However, if we were to import a fixed instance of this latent array saved from the Diffusers implementation, then MFLUX will produce an identical image to the Diffusers implementation (assuming a fixed prompt and using the default parameter settings in the Diffusers setup).
The images below illustrate this equivalence.
In all cases the Schnell model was run for 2 time steps.
The Diffusers implementation ran in CPU mode.
The precision for MFLUX can be set in the [Config](src/mflux/config/config.py) class.
The images below illustrate this equivalence.
In all cases the Schnell model was run for 2 time steps.
The Diffusers implementation ran in CPU mode.
The precision for MFLUX can be set in the [Config](src/mflux/config/config.py) class.
There is typically a noticeable but very small difference in the final image when switching between 16bit and 32bit precision.
---
@ -208,10 +214,10 @@ Luxury food photograph of an italian Linguine pasta alle vongole dish with lots
---
### Quantization
### Quantization
MFLUX supports running FLUX in 4-bit or 8-bit quantized mode. Running a quantized version can greatly speed up the
generation process and reduce the memory consumption by several gigabytes. [Quantized models also take up less disk space](#size-comparisons-for-quantized-models).
generation process and reduce the memory consumption by several gigabytes. [Quantized models also take up less disk space](#size-comparisons-for-quantized-models).
```
mflux-generate \
@ -221,7 +227,7 @@ mflux-generate \
--quantize 8 \
--height 1920 \
--width 1024 \
--prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the ponds calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier"
--prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the ponds calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier"
```
![image](src/mflux/assets/comparison6.jpg)
@ -232,7 +238,7 @@ By selecting the `--quantize` or `-q` flag to be `4`, `8`, or removing it entire
Image generation times in this example are based on a 2021 M1 Pro (32GB) machine. Even though the images are almost identical, there is a ~2x speedup by
running the 8-bit quantized version on this particular machine. Unlike the non-quantized version, for the 8-bit version the swap memory usage is drastically reduced and GPU utilization is close to 100% during the whole generation. Results here can vary across different machines.
#### Size comparisons for quantized models
#### Size comparisons for quantized models
The model sizes for both `schnell` and `dev` at various quantization levels are as follows:
@ -257,7 +263,7 @@ mflux-save \
#### Loading and running a quantized version from disk
To generate a new image from the quantized model, simply provide a `--path` to where it was saved:
To generate a new image from the quantized model, simply provide a `--path` to where it was saved:
```
mflux-generate \
@ -267,21 +273,21 @@ mflux-generate \
--seed 2 \
--height 1920 \
--width 1024 \
--prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the ponds calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier"
--prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the ponds calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier"
```
*Note: When loading a quantized model from disk, there is no need to pass in `-q` flag, since we can infer this from the weight metadata.*
*Also Note: Once we have a local model (quantized [or not](#running-a-non-quantized-model-directly-from-disk)) specified via the `--path` argument, the huggingface cache models are not required to launch the model.
In other words, you can reclaim the 34GB diskspace (per model) by deleting the full 16-bit model from the [Huggingface cache](#generating-an-image) if you choose.*
*Also Note: Once we have a local model (quantized [or not](#running-a-non-quantized-model-directly-from-disk)) specified via the `--path` argument, the huggingface cache models are not required to launch the model.
In other words, you can reclaim the 34GB diskspace (per model) by deleting the full 16-bit model from the [Huggingface cache](#generating-an-image) if you choose.*
*If you don't want to download the full models and quantize them yourself, the 4-bit weights are available here for a direct download:*
*If you don't want to download the full models and quantize them yourself, the 4-bit weights are available here for a direct download:*
- [madroid/flux.1-schnell-mflux-4bit](https://huggingface.co/madroid/flux.1-schnell-mflux-4bit)
- [madroid/flux.1-dev-mflux-4bit](https://huggingface.co/madroid/flux.1-dev-mflux-4bit)
### Running a non-quantized model directly from disk
MFLUX also supports running a non-quantized model directly from a custom location.
MFLUX also supports running a non-quantized model directly from a custom location.
In the example below, the model is placed in `/Users/filipstrand/Desktop/schnell`:
```
@ -290,10 +296,10 @@ mflux-generate \
--model schnell \
--steps 2 \
--seed 2 \
--prompt "Luxury food photograph"
--prompt "Luxury food photograph"
```
Note that the `--model` flag must be set when loading a model from disk.
Note that the `--model` flag must be set when loading a model from disk.
Also note that unlike when using the typical `alias` way of initializing the model (which internally handles that the required resources are downloaded),
when loading a model directly from disk, we require the downloaded models to look like the following:
@ -324,14 +330,14 @@ when loading a model directly from disk, we require the downloaded models to loo
```
This mirrors how the resources are placed in the [HuggingFace Repo](https://huggingface.co/black-forest-labs/FLUX.1-schnell/tree/main) for FLUX.1.
*Huggingface weights, unlike quantized ones exported directly from this project, have to be
processed a bit differently, which is why we require this structure above.*
processed a bit differently, which is why we require this structure above.*
### LoRA
MFLUX support loading trained [LoRA](https://huggingface.co/docs/diffusers/en/training/lora) adapters (actual training support is coming).
The following example [The_Hound](https://huggingface.co/TheLastBen/The_Hound) LoRA from [@TheLastBen](https://github.com/TheLastBen):
The following example [The_Hound](https://huggingface.co/TheLastBen/The_Hound) LoRA from [@TheLastBen](https://github.com/TheLastBen):
```
mflux-generate --prompt "sandor clegane" --model dev --steps 20 --seed 43 -q 8 --lora-paths "sandor_clegane_single_layer.safetensors"
@ -367,12 +373,12 @@ mflux-generate \
```
![image](src/mflux/assets/lora3.jpg)
Just to see the difference, this image displays the four cases: One of having both adapters fully active, partially active and no LoRA at all.
The example above also show the usage of `--lora-scales` flag.
Just to see the difference, this image displays the four cases: One of having both adapters fully active, partially active and no LoRA at all.
The example above also show the usage of `--lora-scales` flag.
#### Supported LoRA formats (updated)
Since different fine-tuning services can use different implementations of FLUX, the corresponding
Since different fine-tuning services can use different implementations of FLUX, the corresponding
LoRA weights trained on these services can be different from one another. The aim of MFLUX is to support the most common ones.
The following table show the current supported formats:
@ -394,4 +400,4 @@ To report additional formats, examples or other any suggestions related to LoRA
### TODO
- [ ] LoRA fine-tuning
- [ ] Frontend support (Gradio/Streamlit/Other?)
- [ ] Frontend support (Gradio/Streamlit/Other?)

View File

@ -7,25 +7,32 @@ name = "mflux"
version = "0.2.1"
description = "A MLX port of FLUX based on the Huggingface Diffusers implementation."
readme = "README.md"
authors = [
{ name = "Filip Strand", email = "strand.filip@gmail.com" }
]
classifiers = [
"Programming Language :: Python :: 3",
"Operating System :: MacOS",
]
keywords = ["diffusers", "flux", "mlx"]
authors = [{ name = "Filip Strand", email = "strand.filip@gmail.com" }]
maintainers = [{ name = "Filip Strand", email = "strand.filip@gmail.com" }]
requires-python = ">=3.10"
dependencies = [
"huggingface-hub>=0.24.5",
"mlx>=0.16.0",
"numpy>=2.0.0",
"numpy>=2.0.1",
"opencv-python>=4.10.0",
"piexif>=1.1.3",
"pillow>=10.4.0",
"transformers>=4.44.0",
"safetensors>=0.4.4",
"sentencepiece>=0.2.0",
"torch>=2.3.1",
"tqdm>=4.66.5",
"huggingface-hub>=0.24.5",
"safetensors>=0.4.4",
"piexif>=1.1.3",
"opencv-python>=4.10.0",
"transformers>=4.44.0",
]
classifiers = [
"Framework :: MLX",
"Intended Audience :: Developers",
"Operating System :: MacOS",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
]
[project.urls]
@ -38,4 +45,4 @@ mflux-generate-controlnet = "mflux.generate_controlnet:main"
[tool.setuptools.packages.find]
where = ["src"]
include = ["mflux*"]
include = ["mflux*"]

View File

@ -1,11 +1,3 @@
mlx>=0.16.0
numpy>=2.0.1
pillow>=10.4.0
transformers>=4.44.0
sentencepiece>=0.2.0
torch>=2.3.1
tqdm>=4.66.5
huggingface-hub>=0.24.5
safetensors>=0.4.4
piexif>=1.1.3
opencv-python>=4.10.0
# direct habitual `pip install -r requirements.txt`
# users to `pyproject.toml` defined req list
-e .