Merge pull request #63 from filipstrand/controlnet-additions
Controlnet additions
This commit is contained in:
Filip Strand
GitHub
commit
c05329e73b
141
README.md
141
README.md
@ -2,30 +2,55 @@
|
||||
|
||||
*A MLX port of FLUX based on the Huggingface Diffusers implementation.*
|
||||
|
||||
|
||||
### About
|
||||
|
||||
Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black Forest Labs](https://blackforestlabs.ai) locally on your Mac!
|
||||
|
||||
### Table of contents
|
||||
|
||||
<!-- TOC start (generated with https://github.com/derlin/bitdowntoc) -->
|
||||
|
||||
- [Philosophy](#philosophy)
|
||||
- [💿 Installation](#-installation)
|
||||
- [🖼️ Generating an image](#%EF%B8%8F-generating-an-image)
|
||||
* [📜 Full list of Command-Line Arguments](#-full-list-of-command-line-arguments)
|
||||
- [⏱️ Image generation speed (updated)](#%EF%B8%8F-image-generation-speed-updated)
|
||||
- [↔️ Equivalent to Diffusers implementation](#%EF%B8%8F-equivalent-to-diffusers-implementation)
|
||||
- [🗜️ Quantization](#%EF%B8%8F-quantization)
|
||||
* [📊 Size comparisons for quantized models](#-size-comparisons-for-quantized-models)
|
||||
* [💾 Saving a quantized version to disk](#-saving-a-quantized-version-to-disk)
|
||||
* [💽 Loading and running a quantized version from disk](#-loading-and-running-a-quantized-version-from-disk)
|
||||
- [💽 Running a non-quantized model directly from disk](#-running-a-non-quantized-model-directly-from-disk)
|
||||
- [🔌 LoRA](#-lora)
|
||||
* [Multi-LoRA](#multi-lora)
|
||||
* [Supported LoRA formats (updated)](#supported-lora-formats-updated)
|
||||
- [🕹️ Controlnet](#%EF%B8%8F-controlnet)
|
||||
- [🚧 Current limitations](#-current-limitations)
|
||||
- [✅ TODO](#-todo)
|
||||
|
||||
<!-- TOC end -->
|
||||
|
||||
### 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 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
|
||||
(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).
|
||||
(thereby avoiding too many abstractions). While MFLUX priorities readability over generality and performance, [it can still be quite fast](#%EF%B8%8F-image-generation-speed-updated), [and even faster quantized](#%EF%B8%8F-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
|
||||
like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) for simple image post-processing.
|
||||
|
||||
### Models
|
||||
|
||||
- [x] FLUX.1-Scnhell
|
||||
- [x] FLUX.1-Dev
|
||||
### 💿 Installation
|
||||
For users, the easiest way to install MFLUX is to use `uv tool`: If you have [installed `uv`](https://github.com/astral-sh/uv?tab=readme-ov-file#installation), simply:
|
||||
|
||||
### Installation
|
||||
For users, the easiest way to install MFLUX is to use `uv tool`:
|
||||
```sh
|
||||
uv tool install mflux
|
||||
```
|
||||
|
||||
If you have [installed `uv`](https://github.com/astral-sh/uv?tab=readme-ov-file#installation), simply: `uv tool install mflux` to get the `mflux-generate` and related command line executables. You can skip to the usage guides below.
|
||||
to get the `mflux-generate` and related command line executables. You can skip to the usage guides below.
|
||||
|
||||
<details>
|
||||
<summary>For the classic way to create a user virtual environment:</summary>
|
||||
@ -46,11 +71,19 @@ pip install -U mflux
|
||||
<summary>For contributors (click to expand)</summary>
|
||||
|
||||
1. Clone the repo:
|
||||
```sh
|
||||
```sh
|
||||
git clone git@github.com:filipstrand/mflux.git
|
||||
```
|
||||
2. `make install` and `make test` (and `make clean` for venv resets)
|
||||
3. Follow format and lint checks prior to submitting Pull Requests. The recommended `make lint` and `make format` installs and uses [`ruff`](https://github.com/astral-sh/ruff). You can setup your editor/IDE to lint/format automatically, or use our provided `make` helpers:
|
||||
2. Install the application
|
||||
|
||||
```sh
|
||||
make install
|
||||
```
|
||||
3. To run the test suite
|
||||
```sh
|
||||
make test
|
||||
```
|
||||
4. Follow format and lint checks prior to submitting Pull Requests. The recommended `make lint` and `make format` installs and uses [`ruff`](https://github.com/astral-sh/ruff). You can setup your editor/IDE to lint/format automatically, or use our provided `make` helpers:
|
||||
- `make format` - formats your code
|
||||
- `make lint` - shows your lint errors and warnings, but does not auto fix
|
||||
- `make check` - via `pre-commit` hooks, formats your code **and** attempts to auto fix lint errors
|
||||
@ -58,7 +91,7 @@ pip install -U mflux
|
||||
|
||||
</details>
|
||||
|
||||
### Generating an image
|
||||
### 🖼️ Generating an image
|
||||
|
||||
Run the command `mflux-generate` by specifying a prompt and the model and some optional arguments. For example, here we use a quantized version of the `schnell` model for 2 steps:
|
||||
|
||||
@ -72,7 +105,7 @@ This example uses the more powerful `dev` model with 25 time steps:
|
||||
mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2 -q 8
|
||||
```
|
||||
|
||||
⚠️ *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.* ⚠️
|
||||
⚠️ *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](#%EF%B8%8F-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:*
|
||||
|
||||
@ -84,7 +117,7 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
|
||||
|
||||
🔒 [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.
|
||||
|
||||
@ -104,14 +137,20 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
|
||||
|
||||
- **`--path`** (optional, `str`, default: `None`): Path to a local model on disk.
|
||||
|
||||
- **`--quantize`** or **`-q`** (optional, `int`, default: `None`): [Quantization](#quantization) (choose between `4` or `8`).
|
||||
- **`--quantize`** or **`-q`** (optional, `int`, default: `None`): [Quantization](#%EF%B8%8F-quantization) (choose between `4` or `8`).
|
||||
|
||||
- **`--lora-paths`** (optional, `[str]`, default: `None`): The paths to the [LoRA](#LoRA) weights.
|
||||
- **`--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.)
|
||||
- **`--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`)
|
||||
|
||||
- **`--controlnet-image-path`** (required, `str`): Path to the local image used by ControlNet to guide output generation.
|
||||
|
||||
- **`--controlnet-strength`** (optional, `float`, default: `0.4`): Degree of influence the control image has on the output. Ranges from `0.0` (no influence) to `1.0` (full influence).
|
||||
|
||||
- **`--controlnet-save-canny`** (optional, bool, default: False): If set, saves the Canny edge detection reference image used by ControlNet.
|
||||
|
||||
Or, with the correct python environment active, create and run a separate script like the following:
|
||||
|
||||
```python
|
||||
@ -139,7 +178,7 @@ image.save(path="image.png")
|
||||
|
||||
For more options on how to configure MFLUX, please see [generate.py](src/mflux/generate.py).
|
||||
|
||||
### Image generation speed (updated)
|
||||
### ⏱️ Image generation speed (updated)
|
||||
|
||||
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:
|
||||
@ -168,7 +207,7 @@ time mflux-generate \
|
||||
*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.
|
||||
@ -217,10 +256,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).
|
||||
|
||||
```sh
|
||||
mflux-generate \
|
||||
@ -234,14 +273,14 @@ mflux-generate \
|
||||
```
|
||||
|
||||
|
||||
*In this example, weights are quantized at **runtime** - this is convenient if you don't want to [save a quantized copy of the weights to disk](#saving-a-quantized-version-to-disk), but still want to benefit from the potential speedup and RAM reduction quantization might bring.*
|
||||
*In this example, weights are quantized at **runtime** - this is convenient if you don't want to [save a quantized copy of the weights to disk](#-saving-a-quantized-version-to-disk), but still want to benefit from the potential speedup and RAM reduction quantization might bring.*
|
||||
|
||||
|
||||
By selecting the `--quantize` or `-q` flag to be `4`, `8`, or removing it entirely, we get all 3 images above. As can be seen, there is very little difference between the images (especially between the 8-bit, and the non-quantized result).
|
||||
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:
|
||||
|
||||
@ -251,7 +290,7 @@ The model sizes for both `schnell` and `dev` at various quantization levels are
|
||||
|
||||
The reason weights sizes are not fully cut in half is because a small number of weights are not quantized and kept at full precision.
|
||||
|
||||
#### Saving a quantized version to disk
|
||||
#### 💾 Saving a quantized version to disk
|
||||
|
||||
To save a local copy of the quantized weights, run the `mflux-save` command like so:
|
||||
|
||||
@ -264,7 +303,7 @@ mflux-save \
|
||||
|
||||
*Note that when saving a quantized version, you will need the original huggingface weights.*
|
||||
|
||||
It is also possible to specify [LoRA](#lora) adapters when saving the model, e.g
|
||||
It is also possible to specify [LoRA](#-lora) adapters when saving the model, e.g
|
||||
|
||||
```sh
|
||||
mflux-save \
|
||||
@ -278,7 +317,7 @@ mflux-save \
|
||||
When generating images with a model like this, no LoRA adapter is needed to be specified since
|
||||
it is already baked into the saved quantized weights.
|
||||
|
||||
#### Loading and running a quantized version from disk
|
||||
#### 💽 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:
|
||||
|
||||
@ -295,14 +334,14 @@ mflux-generate \
|
||||
|
||||
*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](#%EF%B8%8F-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:*
|
||||
- [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
|
||||
### 💽 Running a non-quantized model directly from disk
|
||||
|
||||
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`:
|
||||
@ -349,8 +388,9 @@ This mirrors how the resources are placed in the [HuggingFace Repo](https://hugg
|
||||
*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.*
|
||||
|
||||
---
|
||||
|
||||
### LoRA
|
||||
### 🔌 LoRA
|
||||
|
||||
MFLUX support loading trained [LoRA](https://huggingface.co/docs/diffusers/en/training/lora) adapters (actual training support is coming).
|
||||
|
||||
@ -407,14 +447,53 @@ The following table show the current supported formats:
|
||||
|
||||
To report additional formats, examples or other any suggestions related to LoRA format support, please see [issue #47](https://github.com/filipstrand/mflux/issues/47).
|
||||
|
||||
### Current limitations
|
||||
---
|
||||
|
||||
### 🕹️ Controlnet
|
||||
|
||||
MFLUX has [Controlnet](https://huggingface.co/docs/diffusers/en/using-diffusers/controlnet) support for an even more fine-grained control
|
||||
of the image generation. By providing a reference image via `--controlnet-image-path` and a strength parameter via `--controlnet-strength`, you can guide the generation toward the reference image.
|
||||
|
||||
```sh
|
||||
mflux-generate-controlnet \
|
||||
--prompt "A comic strip with a joker in a purple suit" \
|
||||
--model dev \
|
||||
--steps 20 \
|
||||
--seed 1727047657 \
|
||||
--height 1066 \
|
||||
--width 692 \
|
||||
-q 8 \
|
||||
--lora-paths "Dark Comic - s0_8 g4.safetensors" \
|
||||
--controlnet-image-path "reference.png" \
|
||||
--controlnet-strength 0.5 \
|
||||
--controlnet-save-canny
|
||||
```
|
||||
|
||||
|
||||
*This example combines the controlnet reference image with the LoRA [Dark Comic Flux](https://civitai.com/models/742916/dark-comic-flux)*.
|
||||
|
||||
⚠️ *Note: Controlnet requires an additional one-time download of ~3.58GB of weights from Huggingface. This happens automatically the first time you run the `generate-controlnet` command.
|
||||
At the moment, the Controlnet used is [InstantX/FLUX.1-dev-Controlnet-Canny](https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Canny), which was trained for the `dev` model.
|
||||
It can work well with `schnell`, but performance is not guaranteed.*
|
||||
|
||||
⚠️ *Note: The output can be highly sensitive to the controlnet strength and is very much dependent on the reference image.
|
||||
Too high settings will corrupt the image. A recommended starting point a value like 0.4 and to play around with the strength.*
|
||||
|
||||
|
||||
Controlnet can also work well together with [LoRA adapters](#-lora). In the example below the same reference image is used as a controlnet input
|
||||
with different prompts and LoRA adapters active.
|
||||
|
||||
|
||||
|
||||
### 🚧 Current limitations
|
||||
|
||||
- Images are generated one by one.
|
||||
- Negative prompts not supported.
|
||||
- LoRA weights are only supported for the transformer part of the network.
|
||||
- Some LoRA adapters does not work.
|
||||
- Currently, the supported controlnet is the [canny-only version](https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Canny).
|
||||
|
||||
### TODO
|
||||
### ✅ TODO
|
||||
|
||||
- [ ] Establish unit test suite
|
||||
- [ ] LoRA fine-tuning
|
||||
|
||||
BIN
src/mflux/assets/controlnet1.jpg
Normal file
BIN
src/mflux/assets/controlnet1.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 712 KiB |
BIN
src/mflux/assets/controlnet2.jpg
Normal file
BIN
src/mflux/assets/controlnet2.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 598 KiB |
@ -1,6 +1,12 @@
|
||||
import logging
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import PIL
|
||||
import PIL.Image
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ControlnetUtil:
|
||||
@ -11,3 +17,18 @@ class ControlnetUtil:
|
||||
image_to_canny = np.array(image_to_canny[:, :, None])
|
||||
image_to_canny = np.concatenate([image_to_canny, image_to_canny, image_to_canny], axis=2)
|
||||
return PIL.Image.fromarray(image_to_canny)
|
||||
|
||||
@staticmethod
|
||||
def scale_image(height: int, width: int, img: PIL.Image) -> PIL.Image:
|
||||
if height != img.height or width != img.width:
|
||||
log.warning(f"Control image has different dimensions than the model. Resizing to {width}x{height}")
|
||||
img = img.resize((width, height), PIL.Image.LANCZOS)
|
||||
return img
|
||||
|
||||
@staticmethod
|
||||
def save_canny_image(control_image: PIL.Image, path: str):
|
||||
from mflux import ImageUtil
|
||||
|
||||
base, ext = os.path.splitext(path)
|
||||
new_filename = f"{base}_controlnet_canny{ext}"
|
||||
ImageUtil.save_image(control_image, new_filename)
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
import logging
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import PIL.Image
|
||||
import mlx.core as mx
|
||||
from mlx import nn
|
||||
from tqdm import tqdm
|
||||
@ -10,11 +10,11 @@ from mflux.config.model_config import ModelConfig
|
||||
from mflux.config.runtime_config import RuntimeConfig
|
||||
from mflux.controlnet.controlnet_util import ControlnetUtil
|
||||
from mflux.controlnet.transformer_controlnet import TransformerControlnet
|
||||
from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet
|
||||
from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
|
||||
from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
|
||||
from mflux.models.transformer.transformer import Transformer
|
||||
from mflux.models.vae.vae import VAE
|
||||
from mflux.post_processing.generated_image import GeneratedImage
|
||||
from mflux.post_processing.image_util import ImageUtil
|
||||
from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
|
||||
from mflux.tokenizer.t5_tokenizer import TokenizerT5
|
||||
@ -22,6 +22,10 @@ from mflux.tokenizer.tokenizer_handler import TokenizerHandler
|
||||
from mflux.weights.model_saver import ModelSaver
|
||||
from mflux.weights.weight_handler import WeightHandler
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mflux.post_processing.generated_image import GeneratedImage
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny"
|
||||
@ -79,51 +83,55 @@ class Flux1Controlnet:
|
||||
if weights.quantization_level is not None:
|
||||
self._set_model_weights(weights)
|
||||
|
||||
weights_controlnet, ctrlnet_quantization_level, controlnet_config = WeightHandler.load_controlnet_transformer(
|
||||
weights_controlnet, controlnet_quantization_level, controlnet_config = WeightHandlerControlnet.load_controlnet_transformer(
|
||||
controlnet_id=CONTROLNET_ID
|
||||
)
|
||||
) # fmt: off
|
||||
self.transformer_controlnet = TransformerControlnet(
|
||||
model_config=model_config,
|
||||
num_blocks=controlnet_config["num_layers"],
|
||||
num_single_blocks=controlnet_config["num_single_layers"],
|
||||
)
|
||||
|
||||
if ctrlnet_quantization_level is None:
|
||||
if controlnet_quantization_level is None:
|
||||
self.transformer_controlnet.update(weights_controlnet)
|
||||
|
||||
self.bits = None
|
||||
if quantize is not None or ctrlnet_quantization_level is not None:
|
||||
self.bits = ctrlnet_quantization_level if ctrlnet_quantization_level is not None else quantize
|
||||
# fmt: off
|
||||
nn.quantize(self.transformer_controlnet, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 128, group_size=128, bits=self.bits)
|
||||
# fmt: on
|
||||
if quantize is not None or controlnet_quantization_level is not None:
|
||||
self.bits = controlnet_quantization_level if controlnet_quantization_level is not None else quantize
|
||||
nn.quantize(self.transformer_controlnet, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 128, group_size=128, bits=self.bits) # fmt: off
|
||||
|
||||
if ctrlnet_quantization_level is not None:
|
||||
if controlnet_quantization_level is not None:
|
||||
self.transformer_controlnet.update(weights_controlnet)
|
||||
|
||||
def generate_image(self, seed: int, prompt: str, control_image: PIL.Image.Image, config: ConfigControlnet = ConfigControlnet()) -> GeneratedImage: # fmt: off
|
||||
def generate_image(
|
||||
self,
|
||||
seed: int,
|
||||
prompt: str,
|
||||
output: str,
|
||||
controlnet_image_path: str,
|
||||
controlnet_save_canny: bool = False,
|
||||
config: ConfigControlnet = ConfigControlnet()
|
||||
) -> "GeneratedImage": # fmt: off
|
||||
# Create a new runtime config based on the model type and input parameters
|
||||
config = RuntimeConfig(config, self.model_config)
|
||||
time_steps = tqdm(range(config.num_inference_steps))
|
||||
|
||||
if config.height != control_image.height or config.width != control_image.width:
|
||||
log.warning(
|
||||
f"Control image has different dimensions than the model. Resizing to {config.width}x{config.height}"
|
||||
)
|
||||
control_image = control_image.resize((config.width, config.height), PIL.Image.LANCZOS)
|
||||
# Embedd the controlnet reference image
|
||||
control_image = ImageUtil.load_image(controlnet_image_path)
|
||||
control_image = ControlnetUtil.scale_image(config.height, config.width, control_image)
|
||||
control_image = ControlnetUtil.preprocess_canny(control_image)
|
||||
if controlnet_save_canny:
|
||||
ControlnetUtil.save_canny_image(control_image, output)
|
||||
controlnet_cond = ImageUtil.to_array(control_image)
|
||||
controlnet_cond = self.vae.encode(controlnet_cond)
|
||||
controlnet_cond = (controlnet_cond / self.vae.scaling_factor) + self.vae.shift_factor
|
||||
controlnet_cond = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width)
|
||||
|
||||
# 1. Create the initial latents
|
||||
latents = mx.random.normal(
|
||||
shape=[1, (config.height // 16) * (config.width // 16), 64],
|
||||
key=mx.random.key(seed)
|
||||
) # fmt: off
|
||||
control_image = ControlnetUtil.preprocess_canny(control_image)
|
||||
controlnet_cond = ImageUtil.to_array(control_image)
|
||||
controlnet_cong = self.vae.encode(controlnet_cond)
|
||||
# the rescaling in the next line is not in the huggingface code, but without it the images from
|
||||
# the chosen controlnet model are very bad
|
||||
controlnet_cond = (controlnet_cong / self.vae.scaling_factor) + self.vae.shift_factor
|
||||
controlnet_cond = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width)
|
||||
|
||||
# 2. Embedd the prompt
|
||||
t5_tokens = self.t5_tokenizer.tokenize(prompt)
|
||||
@ -132,7 +140,8 @@ class Flux1Controlnet:
|
||||
pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens)
|
||||
|
||||
for t in time_steps:
|
||||
ctrlnet_block_samples, ctrlnet_single_block_samples = self.transformer_controlnet.forward(
|
||||
# Compute controlnet samples
|
||||
controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet.forward(
|
||||
t=t,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
@ -140,6 +149,7 @@ class Flux1Controlnet:
|
||||
controlnet_cond=controlnet_cond,
|
||||
config=config,
|
||||
)
|
||||
|
||||
# 3.t Predict the noise
|
||||
noise = self.transformer.predict(
|
||||
t=t,
|
||||
@ -147,8 +157,8 @@ class Flux1Controlnet:
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
hidden_states=latents,
|
||||
config=config,
|
||||
controlnet_block_samples=ctrlnet_block_samples,
|
||||
controlnet_single_block_samples=ctrlnet_single_block_samples,
|
||||
controlnet_block_samples=controlnet_block_samples,
|
||||
controlnet_single_block_samples=controlnet_single_block_samples,
|
||||
)
|
||||
|
||||
# 4.t Take one denoise step
|
||||
@ -170,6 +180,7 @@ class Flux1Controlnet:
|
||||
lora_paths=self.lora_paths,
|
||||
lora_scales=self.lora_scales,
|
||||
config=config,
|
||||
controlnet_image_path=controlnet_image_path,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
||||
45
src/mflux/controlnet/weight_handler_controlnet.py
Normal file
45
src/mflux/controlnet/weight_handler_controlnet.py
Normal file
@ -0,0 +1,45 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
from huggingface_hub import snapshot_download
|
||||
from mlx.utils import tree_unflatten
|
||||
|
||||
from mflux.weights.weight_util import WeightUtil
|
||||
|
||||
|
||||
class WeightHandlerControlnet:
|
||||
@staticmethod
|
||||
def load_controlnet_transformer(controlnet_id: str) -> (dict, int):
|
||||
controlnet_path = Path(
|
||||
snapshot_download(repo_id=controlnet_id, allow_patterns=["*.safetensors", "config.json"])
|
||||
)
|
||||
file = next(controlnet_path.glob("diffusion_pytorch_model.safetensors"))
|
||||
quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level")
|
||||
weights = list(mx.load(str(file)).items())
|
||||
|
||||
if quantization_level is not None:
|
||||
return tree_unflatten(weights), quantization_level
|
||||
|
||||
weights = [WeightUtil.reshape_weights(k, v) for k, v in weights]
|
||||
weights = WeightUtil.flatten(weights)
|
||||
weights = tree_unflatten(weights)
|
||||
|
||||
# Quantized weights (i.e. ones exported from this project) don't need any post-processing.
|
||||
if quantization_level is not None:
|
||||
return weights, quantization_level
|
||||
|
||||
# Reshape and process the huggingface weights
|
||||
if "transformer_blocks" in weights:
|
||||
for block in weights["transformer_blocks"]:
|
||||
block["ff"] = {
|
||||
"linear1": block["ff"]["net"][0]["proj"],
|
||||
"linear2": block["ff"]["net"][2]
|
||||
} # fmt: off
|
||||
if block.get("ff_context") is not None:
|
||||
block["ff_context"] = {
|
||||
"linear1": block["ff_context"]["net"][0]["proj"],
|
||||
"linear2": block["ff_context"]["net"][2],
|
||||
}
|
||||
config = json.load(open(controlnet_path / "config.json"))
|
||||
return weights, quantization_level, config
|
||||
@ -1,14 +1,16 @@
|
||||
import argparse
|
||||
import time
|
||||
|
||||
from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig, ImageUtil
|
||||
from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig
|
||||
|
||||
|
||||
def main():
|
||||
# fmt: off
|
||||
parser = argparse.ArgumentParser(description="Generate an image based on a prompt.")
|
||||
parser.add_argument("--prompt", type=str, required=True, help="The textual description of the image to generate.")
|
||||
parser.add_argument("--control-image-path", type=str, required=True, help="Local path of the image to use as input for controlnet.")
|
||||
parser.add_argument("--controlnet-image-path", type=str, required=True, help="Local path of the image to use as input for controlnet.")
|
||||
parser.add_argument("--controlnet-strength", type=float, default=0.4, help="Controls how strongly the control image influences the output image. A value of 0.0 means no influence. (Default is 0.4)")
|
||||
parser.add_argument("--controlnet-save-canny", action="store_true", help="If set, save the Canny edge detection reference input image.")
|
||||
parser.add_argument("--output", type=str, default="image.png", help="The filename for the output image. Default is \"image.png\".")
|
||||
parser.add_argument("--model", "-m", type=str, required=True, choices=["dev", "schnell"], help="The model to use (\"schnell\" or \"dev\").")
|
||||
parser.add_argument("--seed", type=int, default=None, help="Entropy Seed (Default is time-based random-seed)")
|
||||
@ -16,7 +18,6 @@ def main():
|
||||
parser.add_argument("--width", type=int, default=1024, help="Image width (Default is 1024)")
|
||||
parser.add_argument("--steps", type=int, default=None, help="Inference Steps")
|
||||
parser.add_argument("--guidance", type=float, default=3.5, help="Guidance Scale (Default is 3.5)")
|
||||
parser.add_argument("--controlnet-strength", type=float, default=0.7, help="Controls how strongly the control image influences the output image. A value of 0.0 means no influence. (Default is 0.7)")
|
||||
parser.add_argument("--quantize", "-q", type=int, choices=[4, 8], default=None, help="Quantize the model (4 or 8, Default is None)")
|
||||
parser.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk")
|
||||
parser.add_argument("--lora-paths", type=str, nargs="*", default=None, help="Local safetensors for applying LORA from disk")
|
||||
@ -45,7 +46,9 @@ def main():
|
||||
image = flux.generate_image(
|
||||
seed=int(time.time()) if args.seed is None else args.seed,
|
||||
prompt=args.prompt,
|
||||
control_image=ImageUtil.load_image(args.control_image_path),
|
||||
output=args.output,
|
||||
controlnet_image_path=args.controlnet_image_path,
|
||||
controlnet_save_canny=args.controlnet_save_canny,
|
||||
config=ConfigControlnet(
|
||||
num_inference_steps=args.steps,
|
||||
height=args.height,
|
||||
|
||||
@ -1,15 +1,10 @@
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
import importlib
|
||||
|
||||
import PIL.Image
|
||||
import mlx.core as mx
|
||||
import piexif
|
||||
|
||||
from mflux.config.model_config import ModelConfig
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class GeneratedImage:
|
||||
def __init__(
|
||||
@ -25,6 +20,7 @@ class GeneratedImage:
|
||||
generation_time: float,
|
||||
lora_paths: list[str],
|
||||
lora_scales: list[float],
|
||||
controlnet_image_path: str | None = None,
|
||||
controlnet_strength: float | None = None,
|
||||
):
|
||||
self.image = image
|
||||
@ -38,67 +34,17 @@ class GeneratedImage:
|
||||
self.generation_time = generation_time
|
||||
self.lora_paths = lora_paths
|
||||
self.lora_scales = lora_scales
|
||||
self.controlnet_image = controlnet_image_path
|
||||
self.controlnet_strength = controlnet_strength
|
||||
|
||||
def save(self, path: str, export_json_metadata: bool = False) -> None:
|
||||
file_path = Path(path)
|
||||
file_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
file_name = file_path.stem
|
||||
file_extension = file_path.suffix
|
||||
from mflux import ImageUtil
|
||||
|
||||
# If a file already exists, create a new name with a counter
|
||||
counter = 1
|
||||
while file_path.exists():
|
||||
new_name = f"{file_name}_{counter}{file_extension}"
|
||||
file_path = file_path.with_name(new_name)
|
||||
counter += 1
|
||||
|
||||
try:
|
||||
# Save image without metadata first
|
||||
self.image.save(file_path)
|
||||
log.info(f"Image saved successfully at: {file_path}")
|
||||
|
||||
# Optionally save json metadata file
|
||||
if export_json_metadata:
|
||||
with open(f"{file_path.with_suffix('.json')}", "w") as json_file:
|
||||
json.dump(self._get_metadata(), json_file, indent=4)
|
||||
|
||||
# Embed metadata
|
||||
self._embed_metadata(file_path)
|
||||
log.info(f"Metadata embedded successfully at: {file_path}")
|
||||
except Exception as e:
|
||||
log.error(f"Error saving image: {e}")
|
||||
|
||||
def _embed_metadata(self, path: str) -> None:
|
||||
try:
|
||||
# Prepare metadata
|
||||
metadata = self._get_metadata()
|
||||
|
||||
# Convert metadata dictionary to a string
|
||||
metadata_str = str(metadata)
|
||||
|
||||
# Convert the string to bytes (using UTF-8 encoding)
|
||||
user_comment_bytes = metadata_str.encode("utf-8")
|
||||
|
||||
# Define the UserComment tag ID
|
||||
USER_COMMENT_TAG_ID = 0x9286
|
||||
|
||||
# Create a piexif-compatible dictionary structure
|
||||
exif_piexif_dict = {"Exif": {USER_COMMENT_TAG_ID: user_comment_bytes}}
|
||||
|
||||
# Load the image and embed the EXIF data
|
||||
image = PIL.Image.open(path)
|
||||
exif_bytes = piexif.dump(exif_piexif_dict)
|
||||
image.info["exif"] = exif_bytes
|
||||
|
||||
# Save the image with metadata
|
||||
image.save(path, exif=exif_bytes)
|
||||
|
||||
except Exception as e:
|
||||
log.error(f"Error embedding metadata: {e}")
|
||||
ImageUtil.save_image(self.image, path, self._get_metadata(), export_json_metadata)
|
||||
|
||||
def _get_metadata(self) -> dict:
|
||||
return {
|
||||
"mflux_version": str(GeneratedImage.get_version()),
|
||||
"model": str(self.model_config.alias),
|
||||
"seed": str(self.seed),
|
||||
"steps": str(self.steps),
|
||||
@ -106,8 +52,16 @@ class GeneratedImage:
|
||||
"precision": f"{self.precision}",
|
||||
"quantization": "None" if self.quantization is None else f"{self.quantization} bit",
|
||||
"generation_time": f"{self.generation_time:.2f} seconds",
|
||||
"lora_paths": ", ".join(self.lora_paths) if self.lora_paths else "",
|
||||
"lora_scales": ", ".join([f"{scale:.2f}" for scale in self.lora_scales]) if self.lora_scales else "",
|
||||
"lora_paths": ", ".join(self.lora_paths) if self.lora_paths else "None",
|
||||
"lora_scales": ", ".join([f"{scale:.2f}" for scale in self.lora_scales]) if self.lora_scales else "None",
|
||||
"prompt": self.prompt,
|
||||
"controlnet_image": "None" if self.controlnet_image is None else self.controlnet_image,
|
||||
"controlnet_strength": "None" if self.controlnet_strength is None else f"{self.controlnet_strength:.2f}",
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_version():
|
||||
try:
|
||||
return importlib.metadata.version("mflux")
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
return "unknown"
|
||||
|
||||
@ -1,12 +1,19 @@
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import PIL
|
||||
from PIL import Image
|
||||
import PIL.Image
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
|
||||
from PIL import Image
|
||||
import piexif
|
||||
from mflux.config.config import ConfigControlnet
|
||||
from mflux.config.runtime_config import RuntimeConfig
|
||||
from mflux.post_processing.generated_image import GeneratedImage
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ImageUtil:
|
||||
@staticmethod
|
||||
@ -19,6 +26,7 @@ class ImageUtil:
|
||||
lora_paths: list[str],
|
||||
lora_scales: list[float],
|
||||
config: RuntimeConfig,
|
||||
controlnet_image_path: str | None = None,
|
||||
) -> GeneratedImage:
|
||||
normalized = ImageUtil._denormalize(decoded_latents)
|
||||
normalized_numpy = ImageUtil._to_numpy(normalized)
|
||||
@ -35,6 +43,7 @@ class ImageUtil:
|
||||
generation_time=generation_time,
|
||||
lora_paths=lora_paths,
|
||||
lora_scales=lora_scales,
|
||||
controlnet_image_path=controlnet_image_path,
|
||||
controlnet_strength=config.controlnet_strength if isinstance(config.config, ConfigControlnet) else None,
|
||||
)
|
||||
|
||||
@ -76,3 +85,65 @@ class ImageUtil:
|
||||
@staticmethod
|
||||
def load_image(path: str) -> Image.Image:
|
||||
return Image.open(path)
|
||||
|
||||
@staticmethod
|
||||
def save_image(
|
||||
image: PIL.Image.Image,
|
||||
path: str,
|
||||
metadata: dict | None = None,
|
||||
export_json_metadata: bool = False
|
||||
) -> None: # fmt: off
|
||||
file_path = Path(path)
|
||||
file_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
file_name = file_path.stem
|
||||
file_extension = file_path.suffix
|
||||
|
||||
# If a file already exists, create a new name with a counter
|
||||
counter = 1
|
||||
while file_path.exists():
|
||||
new_name = f"{file_name}_{counter}{file_extension}"
|
||||
file_path = file_path.with_name(new_name)
|
||||
counter += 1
|
||||
|
||||
try:
|
||||
# Save image without metadata first
|
||||
image.save(file_path)
|
||||
log.info(f"Image saved successfully at: {file_path}")
|
||||
|
||||
# Optionally save json metadata file
|
||||
if export_json_metadata:
|
||||
with open(f"{file_path.with_suffix('.json')}", "w") as json_file:
|
||||
json.dump(metadata, json_file, indent=4)
|
||||
|
||||
# Embed metadata
|
||||
if metadata is not None:
|
||||
ImageUtil._embed_metadata(metadata, file_path)
|
||||
log.info(f"Metadata embedded successfully at: {file_path}")
|
||||
except Exception as e:
|
||||
log.error(f"Error saving image: {e}")
|
||||
|
||||
@staticmethod
|
||||
def _embed_metadata(metadata: dict, path: str) -> None:
|
||||
try:
|
||||
# Convert metadata dictionary to a string
|
||||
metadata_str = str(metadata)
|
||||
|
||||
# Convert the string to bytes (using UTF-8 encoding)
|
||||
user_comment_bytes = metadata_str.encode("utf-8")
|
||||
|
||||
# Define the UserComment tag ID
|
||||
USER_COMMENT_TAG_ID = 0x9286
|
||||
|
||||
# Create a piexif-compatible dictionary structure
|
||||
exif_piexif_dict = {"Exif": {USER_COMMENT_TAG_ID: user_comment_bytes}}
|
||||
|
||||
# Load the image and embed the EXIF data
|
||||
image = PIL.Image.open(path)
|
||||
exif_bytes = piexif.dump(exif_piexif_dict)
|
||||
image.info["exif"] = exif_bytes
|
||||
|
||||
# Save the image with metadata
|
||||
image.save(path, exif=exif_bytes)
|
||||
|
||||
except Exception as e:
|
||||
log.error(f"Error embedding metadata: {e}")
|
||||
|
||||
@ -1,4 +1,3 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
@ -88,44 +87,6 @@ class WeightHandler:
|
||||
}
|
||||
return weights, quantization_level
|
||||
|
||||
@staticmethod
|
||||
def load_controlnet_transformer(controlnet_id: str) -> (dict, int):
|
||||
controlnet_path = Path(
|
||||
snapshot_download(
|
||||
repo_id=controlnet_id,
|
||||
allow_patterns=["*.safetensors", "config.json"],
|
||||
)
|
||||
)
|
||||
file = next(controlnet_path.glob("diffusion_pytorch_model.safetensors"))
|
||||
quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level")
|
||||
weights = list(mx.load(str(file)).items())
|
||||
|
||||
if quantization_level is not None:
|
||||
return tree_unflatten(weights), quantization_level
|
||||
|
||||
weights = [WeightUtil.reshape_weights(k, v) for k, v in weights]
|
||||
weights = WeightUtil.flatten(weights)
|
||||
weights = tree_unflatten(weights)
|
||||
|
||||
# Quantized weights (i.e. ones exported from this project) don't need any post-processing.
|
||||
if quantization_level is not None:
|
||||
return weights, quantization_level
|
||||
|
||||
# Reshape and process the huggingface weights
|
||||
if "transformer_blocks" in weights:
|
||||
for block in weights["transformer_blocks"]:
|
||||
block["ff"] = {
|
||||
"linear1": block["ff"]["net"][0]["proj"],
|
||||
"linear2": block["ff"]["net"][2],
|
||||
}
|
||||
if block.get("ff_context") is not None:
|
||||
block["ff_context"] = {
|
||||
"linear1": block["ff_context"]["net"][0]["proj"],
|
||||
"linear2": block["ff_context"]["net"][2],
|
||||
}
|
||||
config = json.load(open(controlnet_path / "config.json"))
|
||||
return weights, quantization_level, config
|
||||
|
||||
@staticmethod
|
||||
def load_vae(root_path: Path) -> (dict, int):
|
||||
weights, quantization_level = WeightHandler._get_weights("vae", root_path)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user