Merge remote-tracking branch 'refs/remotes/mflux/main' into easy-import-dx
This commit is contained in:
filipstrand
commit
e007402cfa
2
.gitignore
vendored
2
.gitignore
vendored
@ -13,3 +13,5 @@
|
||||
*.pyc
|
||||
*.safetensors
|
||||
*.json
|
||||
|
||||
*.egg-info
|
||||
84
README.md
84
README.md
@ -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:
|
||||
@ -111,7 +117,7 @@ from mflux import Flux1, Config
|
||||
flux = Flux1.from_alias(
|
||||
alias="schnell", # "schnell" or "dev"
|
||||
quantize=8, # 4 or 8
|
||||
)
|
||||
)
|
||||
|
||||
# Generate an image
|
||||
image = flux.generate_image(
|
||||
@ -131,7 +137,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 \
|
||||
@ -155,20 +161,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.
|
||||
|
||||
---
|
||||
@ -207,10 +213,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 \
|
||||
@ -220,7 +226,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 pond’s 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 pond’s calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier"
|
||||
```
|
||||
|
||||
|
||||
@ -231,7 +237,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:
|
||||
|
||||
@ -256,7 +262,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 \
|
||||
@ -266,21 +272,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 pond’s 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 pond’s 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`:
|
||||
|
||||
```
|
||||
@ -289,10 +295,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:
|
||||
@ -323,14 +329,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"
|
||||
@ -366,12 +372,12 @@ mflux-generate \
|
||||
```
|
||||
|
||||
|
||||
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:
|
||||
|
||||
@ -393,4 +399,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?)
|
||||
|
||||
@ -7,24 +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",
|
||||
"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]
|
||||
@ -33,7 +41,8 @@ homepage = "https://github.com/filipstrand/mflux"
|
||||
[project.scripts]
|
||||
mflux-generate = "mflux.generate:main"
|
||||
mflux-save = "mflux.save:main"
|
||||
mflux-generate-controlnet = "mflux.generate_controlnet:main"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
include = ["mflux*"]
|
||||
include = ["mflux*"]
|
||||
|
||||
@ -1,10 +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
|
||||
# direct habitual `pip install -r requirements.txt`
|
||||
# users to `pyproject.toml` defined req list
|
||||
-e .
|
||||
|
||||
@ -21,3 +21,16 @@ class Config:
|
||||
self.height = 16 * (height // 16)
|
||||
self.num_inference_steps = num_inference_steps
|
||||
self.guidance = guidance
|
||||
|
||||
|
||||
class ConfigControlnet(Config):
|
||||
def __init__(
|
||||
self,
|
||||
num_inference_steps: int = 4,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
guidance: float = 4.0,
|
||||
controlnet_strength: float = 1.0,
|
||||
):
|
||||
super().__init__(num_inference_steps, width, height, guidance)
|
||||
self.controlnet_strength = controlnet_strength
|
||||
|
||||
@ -1,13 +1,13 @@
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
|
||||
from mflux.config.config import Config
|
||||
from mflux.config.config import Config, ConfigControlnet
|
||||
from mflux.config.model_config import ModelConfig
|
||||
|
||||
|
||||
class RuntimeConfig:
|
||||
|
||||
def __init__(self, config: Config, model_config: ModelConfig):
|
||||
def __init__(self, config: Config | ConfigControlnet, model_config: ModelConfig):
|
||||
self.config = config
|
||||
self.model_config = model_config
|
||||
self.sigmas = self._create_sigmas(config, model_config)
|
||||
@ -35,6 +35,13 @@ class RuntimeConfig:
|
||||
@property
|
||||
def num_train_steps(self) -> int:
|
||||
return self.model_config.num_train_steps
|
||||
|
||||
@property
|
||||
def controlnet_strength(self) -> float:
|
||||
if isinstance(self.config, ConfigControlnet):
|
||||
return self.config.controlnet_strength
|
||||
else:
|
||||
raise NotImplementedError("Controlnet conditioning scale is only available for ConfigControlnet")
|
||||
|
||||
@staticmethod
|
||||
def _create_sigmas(config, model) -> mx.array:
|
||||
|
||||
0
src/mflux/controlnet/__init__.py
Normal file
0
src/mflux/controlnet/__init__.py
Normal file
14
src/mflux/controlnet/controlnet_util.py
Normal file
14
src/mflux/controlnet/controlnet_util.py
Normal file
@ -0,0 +1,14 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import PIL
|
||||
|
||||
|
||||
class ControlnetUtil:
|
||||
|
||||
@staticmethod
|
||||
def preprocess_canny(img: PIL.Image) -> PIL.Image:
|
||||
image_to_canny = np.array(img)
|
||||
image_to_canny = cv2.Canny(image_to_canny, 100, 200)
|
||||
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)
|
||||
189
src/mflux/controlnet/flux_controlnet.py
Normal file
189
src/mflux/controlnet/flux_controlnet.py
Normal file
@ -0,0 +1,189 @@
|
||||
import logging
|
||||
|
||||
import PIL.Image
|
||||
import mlx.core as mx
|
||||
from mlx import nn
|
||||
from tqdm import tqdm
|
||||
|
||||
from mflux.config.config import ConfigControlnet
|
||||
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.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
|
||||
from mflux.tokenizer.tokenizer_handler import TokenizerHandler
|
||||
from mflux.weights.model_saver import ModelSaver
|
||||
from mflux.weights.weight_handler import WeightHandler
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny"
|
||||
|
||||
|
||||
class Flux1Controlnet:
|
||||
def __init__(
|
||||
self,
|
||||
model_config: ModelConfig,
|
||||
quantize: int | None = None,
|
||||
local_path: str | None = None,
|
||||
lora_paths: list[str] | None = None,
|
||||
lora_scales: list[float] | None = None,
|
||||
controlnet_path: str | None = None,
|
||||
):
|
||||
self.lora_paths = lora_paths
|
||||
self.lora_scales = lora_scales
|
||||
self.model_config = model_config
|
||||
|
||||
# Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface
|
||||
tokenizers = TokenizerHandler(model_config.model_name, self.model_config.max_sequence_length, local_path)
|
||||
self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=self.model_config.max_sequence_length)
|
||||
self.clip_tokenizer = TokenizerCLIP(tokenizers.clip)
|
||||
|
||||
# Initialize the models
|
||||
self.vae = VAE()
|
||||
self.transformer = Transformer(model_config)
|
||||
self.t5_text_encoder = T5Encoder()
|
||||
self.clip_text_encoder = CLIPEncoder()
|
||||
|
||||
# Load the weights from disk, huggingface cache, or download from huggingface
|
||||
weights = WeightHandler(
|
||||
repo_id=model_config.model_name,
|
||||
local_path=local_path,
|
||||
lora_paths=lora_paths,
|
||||
lora_scales=lora_scales
|
||||
)
|
||||
|
||||
# Set the loaded weights if they are not quantized
|
||||
if weights.quantization_level is None:
|
||||
self._set_model_weights(weights)
|
||||
|
||||
# Optionally quantize the model here at initialization (also required if about to load quantized weights)
|
||||
self.bits = None
|
||||
if quantize is not None or weights.quantization_level is not None:
|
||||
self.bits = weights.quantization_level if weights.quantization_level is not None else quantize
|
||||
nn.quantize(self.vae, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=self.bits)
|
||||
nn.quantize(self.transformer, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 64, group_size=64, bits=self.bits)
|
||||
nn.quantize(self.t5_text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=self.bits)
|
||||
nn.quantize(self.clip_text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=self.bits)
|
||||
|
||||
# If loading previously saved quantized weights, the weights must be set after modules have been quantized
|
||||
if weights.quantization_level is not None:
|
||||
self._set_model_weights(weights)
|
||||
|
||||
weights_controlnet, ctrlnet_quantization_level, controlnet_config = WeightHandler.load_controlnet_transformer(controlnet_id=CONTROLNET_ID)
|
||||
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:
|
||||
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
|
||||
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)
|
||||
|
||||
if ctrlnet_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:
|
||||
# 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)
|
||||
|
||||
# 1. Create the initial latents
|
||||
latents = mx.random.normal(
|
||||
shape=[1, (config.height // 16) * (config.width // 16), 64],
|
||||
key=mx.random.key(seed)
|
||||
)
|
||||
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)
|
||||
clip_tokens = self.clip_tokenizer.tokenize(prompt)
|
||||
prompt_embeds = self.t5_text_encoder.forward(t5_tokens)
|
||||
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(
|
||||
t=t,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
hidden_states=latents,
|
||||
controlnet_cond=controlnet_cond,
|
||||
config=config,
|
||||
)
|
||||
# 3.t Predict the noise
|
||||
noise = self.transformer.predict(
|
||||
t=t,
|
||||
prompt_embeds=prompt_embeds,
|
||||
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,
|
||||
)
|
||||
|
||||
# 4.t Take one denoise step
|
||||
dt = config.sigmas[t + 1] - config.sigmas[t]
|
||||
latents += noise * dt
|
||||
|
||||
# Evaluate to enable progress tracking
|
||||
mx.eval(latents)
|
||||
|
||||
# 5. Decode the latent array and return the image
|
||||
latents = Flux1Controlnet._unpack_latents(latents, config.height, config.width)
|
||||
decoded = self.vae.decode(latents)
|
||||
return ImageUtil.to_image(
|
||||
decoded_latents=decoded,
|
||||
seed=seed,
|
||||
prompt=prompt,
|
||||
quantization=self.bits,
|
||||
generation_time=time_steps.format_dict['elapsed'],
|
||||
lora_paths=self.lora_paths,
|
||||
lora_scales=self.lora_scales,
|
||||
config=config,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _unpack_latents(latents: mx.array, height: int, width: int) -> mx.array:
|
||||
latents = mx.reshape(latents, (1, height // 16, width // 16, 16, 2, 2))
|
||||
latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5))
|
||||
latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2))
|
||||
return latents
|
||||
|
||||
@staticmethod
|
||||
def _pack_latents(latents: mx.array, height: int, width: int) -> mx.array:
|
||||
latents = mx.reshape(latents, (1, 16, height // 16, 2, width // 16, 2))
|
||||
latents = mx.transpose(latents, (0, 2, 4, 1, 3, 5))
|
||||
latents = mx.reshape(latents, (1, (width // 16) * (height // 16), 64))
|
||||
return latents
|
||||
|
||||
def _set_model_weights(self, weights):
|
||||
self.vae.update(weights.vae)
|
||||
self.transformer.update(weights.transformer)
|
||||
self.t5_text_encoder.update(weights.t5_encoder)
|
||||
self.clip_text_encoder.update(weights.clip_encoder)
|
||||
|
||||
def save_model(self, base_path: str) -> None:
|
||||
ModelSaver.save_model(self, self.bits, base_path)
|
||||
ModelSaver.save_weights(base_path, self.bits, self.transformer_controlnet, "transformer_controlnet")
|
||||
96
src/mflux/controlnet/transformer_controlnet.py
Normal file
96
src/mflux/controlnet/transformer_controlnet.py
Normal file
@ -0,0 +1,96 @@
|
||||
from typing import Tuple
|
||||
|
||||
import mlx.core as mx
|
||||
from mlx import nn
|
||||
|
||||
from mflux.config.model_config import ModelConfig
|
||||
from mflux.config.runtime_config import RuntimeConfig
|
||||
from mflux.models.transformer.embed_nd import EmbedND
|
||||
from mflux.models.transformer.joint_transformer_block import JointTransformerBlock
|
||||
from mflux.models.transformer.single_transformer_block import SingleTransformerBlock
|
||||
from mflux.models.transformer.time_text_embed import TimeTextEmbed
|
||||
from mflux.models.transformer.transformer import Transformer
|
||||
|
||||
|
||||
class TransformerControlnet(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_config: ModelConfig,
|
||||
num_blocks: int,
|
||||
num_single_blocks: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.pos_embed = EmbedND()
|
||||
self.x_embedder = nn.Linear(64, 3072)
|
||||
self.time_text_embed = TimeTextEmbed(model_config=model_config)
|
||||
self.context_embedder = nn.Linear(4096, 3072)
|
||||
self.transformer_blocks = [JointTransformerBlock(i) for i in range(num_blocks)]
|
||||
self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(num_single_blocks)]
|
||||
|
||||
zero_init = nn.init.constant(0)
|
||||
self.controlnet_x_embedder = nn.Linear(64, 3072).apply(zero_init)
|
||||
self.controlnet_blocks = [nn.Linear(3072, 3072).apply(zero_init) for _ in range(num_blocks)]
|
||||
|
||||
self.controlnet_single_blocks = [nn.Linear(3072, 3072) for _ in range(num_single_blocks)]
|
||||
|
||||
def forward(
|
||||
self,
|
||||
t: int,
|
||||
prompt_embeds: mx.array,
|
||||
pooled_prompt_embeds: mx.array,
|
||||
hidden_states: mx.array,
|
||||
controlnet_cond: mx.array,
|
||||
config: RuntimeConfig,
|
||||
) -> (list[mx.array], list[mx.array]):
|
||||
time_step = config.sigmas[t] * config.num_train_steps
|
||||
time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision)
|
||||
hidden_states = self.x_embedder(hidden_states)
|
||||
hidden_states = hidden_states + self.controlnet_x_embedder(controlnet_cond)
|
||||
conditioning_scale = config.config.controlnet_strength
|
||||
|
||||
guidance = mx.broadcast_to(config.guidance * config.num_train_steps, (1,)).astype(config.precision)
|
||||
text_embeddings = self.time_text_embed.forward(time_step, pooled_prompt_embeds, guidance)
|
||||
encoder_hidden_states = self.context_embedder(prompt_embeds)
|
||||
txt_ids = Transformer.prepare_text_ids(seq_len=prompt_embeds.shape[1])
|
||||
img_ids = Transformer.prepare_latent_image_ids(config.height, config.width)
|
||||
ids = mx.concatenate((txt_ids, img_ids), axis=1)
|
||||
image_rotary_emb = self.pos_embed.forward(ids)
|
||||
|
||||
block_samples = ()
|
||||
for block in self.transformer_blocks:
|
||||
encoder_hidden_states, hidden_states = block.forward(
|
||||
hidden_states=hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
text_embeddings=text_embeddings,
|
||||
rotary_embeddings=image_rotary_emb
|
||||
)
|
||||
block_samples = block_samples + (hidden_states,)
|
||||
|
||||
hidden_states = mx.concatenate([encoder_hidden_states, hidden_states], axis=1)
|
||||
|
||||
# controlnet block
|
||||
controlnet_block_samples = ()
|
||||
for block_sample, controlnet_block in zip(block_samples, self.controlnet_blocks):
|
||||
block_sample = controlnet_block(block_sample)
|
||||
controlnet_block_samples = controlnet_block_samples + (block_sample,)
|
||||
|
||||
single_block_samples = ()
|
||||
for block in self.single_transformer_blocks:
|
||||
ctrlnet_hidden_states = block.forward(
|
||||
hidden_states=ctrlnet_hidden_states,
|
||||
text_embeddings=text_embeddings,
|
||||
rotary_embeddings=image_rotary_emb
|
||||
)
|
||||
single_block_samples = single_block_samples + (ctrlnet_hidden_states[:, encoder_hidden_states.shape[1] :],)
|
||||
|
||||
controlnet_single_block_samples = ()
|
||||
for single_block_sample, controlnet_block in zip(single_block_samples, self.controlnet_single_blocks):
|
||||
single_block_sample = controlnet_block(single_block_sample)
|
||||
controlnet_single_block_samples = controlnet_single_block_samples + (single_block_sample,)
|
||||
|
||||
# scaling
|
||||
controlnet_block_samples = [sample * conditioning_scale for sample in controlnet_block_samples]
|
||||
controlnet_single_block_samples = [sample * conditioning_scale for sample in controlnet_single_block_samples]
|
||||
|
||||
return controlnet_block_samples, controlnet_single_block_samples
|
||||
@ -1,8 +1,5 @@
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
from mlx import nn
|
||||
from mlx.utils import tree_flatten
|
||||
from tqdm import tqdm
|
||||
|
||||
from mflux.config.config import Config
|
||||
@ -12,11 +9,12 @@ 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.image import Image
|
||||
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
|
||||
from mflux.tokenizer.tokenizer_handler import TokenizerHandler
|
||||
from mflux.weights.model_saver import ModelSaver
|
||||
from mflux.weights.weight_handler import WeightHandler
|
||||
|
||||
|
||||
@ -70,7 +68,7 @@ class Flux1:
|
||||
if weights.quantization_level is not None:
|
||||
self._set_model_weights(weights)
|
||||
|
||||
def generate_image(self, seed: int, prompt: str, config: Config = Config()) -> Image:
|
||||
def generate_image(self, seed: int, prompt: str, config: Config = Config()) -> GeneratedImage:
|
||||
# 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))
|
||||
@ -138,39 +136,5 @@ class Flux1:
|
||||
self.t5_text_encoder.update(weights.t5_encoder)
|
||||
self.clip_text_encoder.update(weights.clip_encoder)
|
||||
|
||||
def save_model(self, base_path: str):
|
||||
def _save_tokenizer(tokenizer, subdir: str):
|
||||
path = Path(base_path) / subdir
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
tokenizer.save_pretrained(path)
|
||||
|
||||
def _save_weights(model, subdir: str):
|
||||
path = Path(base_path) / subdir
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
weights = _split_weights(dict(tree_flatten(model.parameters())))
|
||||
for i, weight in enumerate(weights):
|
||||
mx.save_safetensors(str(path / f"{i}.safetensors"), weight, {"quantization_level": str(self.bits)})
|
||||
|
||||
def _split_weights(weights: dict, max_file_size_gb: int = 2) -> list:
|
||||
# Copied from mlx-examples repo
|
||||
max_file_size_bytes = max_file_size_gb << 30
|
||||
shards = []
|
||||
shard, shard_size = {}, 0
|
||||
for k, v in weights.items():
|
||||
if shard_size + v.nbytes > max_file_size_bytes:
|
||||
shards.append(shard)
|
||||
shard, shard_size = {}, 0
|
||||
shard[k] = v
|
||||
shard_size += v.nbytes
|
||||
shards.append(shard)
|
||||
return shards
|
||||
|
||||
# Save the tokenizers
|
||||
_save_tokenizer(self.clip_tokenizer.tokenizer, "tokenizer")
|
||||
_save_tokenizer(self.t5_tokenizer.tokenizer, "tokenizer_2")
|
||||
|
||||
# Save the models
|
||||
_save_weights(self.vae, "vae")
|
||||
_save_weights(self.transformer, "transformer")
|
||||
_save_weights(self.clip_text_encoder, "text_encoder")
|
||||
_save_weights(self.t5_text_encoder, "text_encoder_2")
|
||||
def save_model(self, base_path: str) -> None:
|
||||
ModelSaver.save_model(self, self.bits, base_path)
|
||||
|
||||
@ -14,7 +14,7 @@ def main():
|
||||
parser.add_argument('--seed', type=int, default=None, help='Entropy Seed (Default is time-based random-seed)')
|
||||
parser.add_argument('--height', type=int, default=1024, help='Image height (Default is 1024)')
|
||||
parser.add_argument('--width', type=int, default=1024, help='Image width (Default is 1024)')
|
||||
parser.add_argument('--steps', type=int, default=4, help='Inference Steps')
|
||||
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('--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')
|
||||
@ -27,6 +27,9 @@ def main():
|
||||
if args.path and args.model is None:
|
||||
parser.error("--model must be specified when using --path")
|
||||
|
||||
if args.steps is None:
|
||||
args.steps = 4 if args.model == "schnell" else 14
|
||||
|
||||
# Load the model
|
||||
flux = Flux1(
|
||||
model_config=ModelConfig.from_alias(args.model),
|
||||
|
||||
68
src/mflux/generate_controlnet.py
Normal file
68
src/mflux/generate_controlnet.py
Normal file
@ -0,0 +1,68 @@
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
from mflux.config.model_config import ModelConfig
|
||||
from mflux.config.config import ConfigControlnet
|
||||
from mflux.controlnet.flux_controlnet import Flux1Controlnet
|
||||
from mflux.post_processing.image_util import ImageUtil
|
||||
|
||||
|
||||
def main():
|
||||
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('--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)')
|
||||
parser.add_argument('--height', type=int, default=1024, help='Image height (Default is 1024)')
|
||||
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')
|
||||
parser.add_argument('--lora-scales', type=float, nargs='*', default=None, help='Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.')
|
||||
parser.add_argument('--metadata', action='store_true', help='Export image metadata as a JSON file.')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.path and args.model is None:
|
||||
parser.error("--model must be specified when using --path")
|
||||
|
||||
if args.steps is None:
|
||||
args.steps = 4 if args.model == "schnell" else 14
|
||||
|
||||
# Load the model
|
||||
flux = Flux1Controlnet(
|
||||
model_config=ModelConfig.from_alias(args.model),
|
||||
quantize=args.quantize,
|
||||
local_path=args.path,
|
||||
lora_paths=args.lora_paths,
|
||||
lora_scales=args.lora_scales
|
||||
)
|
||||
|
||||
# Generate an image
|
||||
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),
|
||||
config=ConfigControlnet(
|
||||
num_inference_steps=args.steps,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
guidance=args.guidance,
|
||||
controlnet_strength=args.controlnet_strength
|
||||
)
|
||||
)
|
||||
|
||||
# Save the image
|
||||
image.save(path=args.output, export_json_metadata=args.metadata)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@ -6,8 +6,8 @@ class FeedForward(nn.Module):
|
||||
|
||||
def __init__(self, activation_function):
|
||||
super().__init__()
|
||||
self.linear1 = nn.Linear(3072, 6144)
|
||||
self.linear2 = nn.Linear(6144, 3072)
|
||||
self.linear1 = nn.Linear(3072, 12288)
|
||||
self.linear2 = nn.Linear(12288, 3072)
|
||||
self.activation_function = activation_function
|
||||
|
||||
def forward(self, hidden_states: mx.array) -> mx.array:
|
||||
|
||||
@ -1,3 +1,5 @@
|
||||
import math
|
||||
|
||||
import mlx.core as mx
|
||||
from mlx import nn
|
||||
|
||||
@ -30,6 +32,8 @@ class Transformer(nn.Module):
|
||||
pooled_prompt_embeds: mx.array,
|
||||
hidden_states: mx.array,
|
||||
config: RuntimeConfig,
|
||||
controlnet_block_samples: list[mx.array] | None = None,
|
||||
controlnet_single_block_samples: list[mx.array] | None = None,
|
||||
) -> mx.array:
|
||||
time_step = config.sigmas[t] * config.num_train_steps
|
||||
time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision)
|
||||
@ -37,27 +41,38 @@ class Transformer(nn.Module):
|
||||
guidance = mx.broadcast_to(config.guidance * config.num_train_steps, (1,)).astype(config.precision)
|
||||
text_embeddings = self.time_text_embed.forward(time_step, pooled_prompt_embeds, guidance)
|
||||
encoder_hidden_states = self.context_embedder(prompt_embeds)
|
||||
txt_ids = Transformer._prepare_text_ids(seq_len=prompt_embeds.shape[1])
|
||||
img_ids = Transformer._prepare_latent_image_ids(config.height, config.width)
|
||||
txt_ids = Transformer.prepare_text_ids(seq_len=prompt_embeds.shape[1])
|
||||
img_ids = Transformer.prepare_latent_image_ids(config.height, config.width)
|
||||
ids = mx.concatenate((txt_ids, img_ids), axis=1)
|
||||
image_rotary_emb = self.pos_embed.forward(ids)
|
||||
|
||||
for block in self.transformer_blocks:
|
||||
for idx, block in enumerate(self.transformer_blocks):
|
||||
encoder_hidden_states, hidden_states = block.forward(
|
||||
hidden_states=hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
text_embeddings=text_embeddings,
|
||||
rotary_embeddings=image_rotary_emb
|
||||
)
|
||||
if controlnet_block_samples is not None and len(controlnet_block_samples) > 0:
|
||||
interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
|
||||
interval_control = int(math.ceil(interval_control))
|
||||
hidden_states = hidden_states + controlnet_block_samples[idx // interval_control]
|
||||
|
||||
hidden_states = mx.concatenate([encoder_hidden_states, hidden_states], axis=1)
|
||||
|
||||
for block in self.single_transformer_blocks:
|
||||
for idx, block in enumerate(self.single_transformer_blocks):
|
||||
hidden_states = block.forward(
|
||||
hidden_states=hidden_states,
|
||||
text_embeddings=text_embeddings,
|
||||
rotary_embeddings=image_rotary_emb
|
||||
)
|
||||
if controlnet_single_block_samples is not None and len(controlnet_single_block_samples) > 0:
|
||||
interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
|
||||
interval_control = int(math.ceil(interval_control))
|
||||
hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
|
||||
hidden_states[:, encoder_hidden_states.shape[1] :, ...]
|
||||
+ controlnet_single_block_samples[idx // interval_control]
|
||||
)
|
||||
|
||||
hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...]
|
||||
hidden_states = self.norm_out.forward(hidden_states, text_embeddings)
|
||||
@ -66,7 +81,7 @@ class Transformer(nn.Module):
|
||||
return noise
|
||||
|
||||
@staticmethod
|
||||
def _prepare_latent_image_ids(height: int, width: int) -> mx.array:
|
||||
def prepare_latent_image_ids(height: int, width: int) -> mx.array:
|
||||
latent_width = width // 16
|
||||
latent_height = height // 16
|
||||
latent_image_ids = mx.zeros((latent_height, latent_width, 3))
|
||||
@ -77,5 +92,5 @@ class Transformer(nn.Module):
|
||||
return latent_image_ids
|
||||
|
||||
@staticmethod
|
||||
def _prepare_text_ids(seq_len: mx.array) -> mx.array:
|
||||
def prepare_text_ids(seq_len: mx.array) -> mx.array:
|
||||
return mx.zeros((1, seq_len, 3))
|
||||
|
||||
@ -11,7 +11,7 @@ from mflux.config.model_config import ModelConfig
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Image:
|
||||
class GeneratedImage:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@ -26,6 +26,7 @@ class Image:
|
||||
generation_time: float,
|
||||
lora_paths: list[str],
|
||||
lora_scales: list[float],
|
||||
controlnet_strength: float | None = None,
|
||||
):
|
||||
self.image = image
|
||||
self.model_config = model_config
|
||||
@ -38,6 +39,7 @@ class Image:
|
||||
self.generation_time = generation_time
|
||||
self.lora_paths = lora_paths
|
||||
self.lora_scales = lora_scales
|
||||
self.controlnet_strength = controlnet_strength
|
||||
|
||||
def save(self, path: str, export_json_metadata: bool = False) -> None:
|
||||
file_path = Path(path)
|
||||
@ -123,4 +125,5 @@ class Image:
|
||||
'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 '',
|
||||
'prompt': self.prompt,
|
||||
'controlnet_strength': "None" if self.controlnet_strength is None else f"{self.controlnet_strength:.2f}",
|
||||
}
|
||||
@ -3,8 +3,9 @@ from PIL import Image
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
|
||||
from mflux.config.config import ConfigControlnet
|
||||
from mflux.config.runtime_config import RuntimeConfig
|
||||
from mflux.post_processing.image import Image
|
||||
from mflux.post_processing.generated_image import GeneratedImage
|
||||
|
||||
|
||||
class ImageUtil:
|
||||
@ -19,11 +20,11 @@ class ImageUtil:
|
||||
lora_paths: list[str],
|
||||
lora_scales: list[float],
|
||||
config: RuntimeConfig,
|
||||
) -> Image:
|
||||
) -> GeneratedImage:
|
||||
normalized = ImageUtil._denormalize(decoded_latents)
|
||||
normalized_numpy = ImageUtil._to_numpy(normalized)
|
||||
image = ImageUtil._numpy_to_pil(normalized_numpy)
|
||||
return Image(
|
||||
return GeneratedImage(
|
||||
image=image,
|
||||
model_config=config.model_config,
|
||||
seed=seed,
|
||||
@ -35,6 +36,7 @@ class ImageUtil:
|
||||
generation_time=generation_time,
|
||||
lora_paths=lora_paths,
|
||||
lora_scales=lora_scales,
|
||||
controlnet_strength=config.controlnet_strength if isinstance(config.config, ConfigControlnet) else None,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
@ -66,14 +68,12 @@ class ImageUtil:
|
||||
|
||||
@staticmethod
|
||||
def to_array(image: PIL.Image.Image) -> mx.array:
|
||||
image = ImageUtil._resize(image)
|
||||
image = ImageUtil._pil_to_numpy(image)
|
||||
array = mx.array(image)
|
||||
array = mx.transpose(array, (0, 3, 1, 2))
|
||||
array = ImageUtil._normalize(array)
|
||||
return array
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _resize(image):
|
||||
image = image.resize((1024, 1024), resample=PIL.Image.LANCZOS)
|
||||
return image
|
||||
def load_image(path: str) -> Image.Image:
|
||||
return Image.open(path)
|
||||
|
||||
50
src/mflux/weights/model_saver.py
Normal file
50
src/mflux/weights/model_saver.py
Normal file
@ -0,0 +1,50 @@
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
from mlx import nn
|
||||
from mlx.utils import tree_flatten
|
||||
from transformers import CLIPTokenizer, T5Tokenizer
|
||||
|
||||
|
||||
class ModelSaver:
|
||||
|
||||
@staticmethod
|
||||
def save_model(model, bits: int, base_path: str):
|
||||
# Save the tokenizers
|
||||
ModelSaver._save_tokenizer(base_path, model.clip_tokenizer.tokenizer, "tokenizer")
|
||||
ModelSaver._save_tokenizer(base_path, model.t5_tokenizer.tokenizer, "tokenizer_2")
|
||||
|
||||
# Save the models
|
||||
ModelSaver.save_weights(base_path, bits, model.vae, "vae")
|
||||
ModelSaver.save_weights(base_path, bits, model.transformer, "transformer")
|
||||
ModelSaver.save_weights(base_path, bits, model.clip_text_encoder, "text_encoder")
|
||||
ModelSaver.save_weights(base_path, bits, model.t5_text_encoder, "text_encoder_2")
|
||||
|
||||
@staticmethod
|
||||
def _save_tokenizer(base_path: str, tokenizer: CLIPTokenizer | T5Tokenizer, subdir: str):
|
||||
path = Path(base_path) / subdir
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
tokenizer.save_pretrained(path)
|
||||
|
||||
@staticmethod
|
||||
def save_weights(base_path: str, bits: int, model: nn.Module, subdir: str):
|
||||
path = Path(base_path) / subdir
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
weights = ModelSaver._split_weights(base_path, dict(tree_flatten(model.parameters())))
|
||||
for i, weight in enumerate(weights):
|
||||
mx.save_safetensors(str(path / f"{i}.safetensors"), weight, {"quantization_level": str(bits)})
|
||||
|
||||
@staticmethod
|
||||
def _split_weights(base_path: str, weights: dict, max_file_size_gb: int = 2) -> list:
|
||||
# Copied from mlx-examples repo
|
||||
max_file_size_bytes = max_file_size_gb << 30
|
||||
shards = []
|
||||
shard, shard_size = {}, 0
|
||||
for k, v in weights.items():
|
||||
if shard_size + v.nbytes > max_file_size_bytes:
|
||||
shards.append(shard)
|
||||
shard, shard_size = {}, 0
|
||||
shard[k] = v
|
||||
shard_size += v.nbytes
|
||||
shards.append(shard)
|
||||
return shards
|
||||
@ -1,3 +1,4 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
@ -87,6 +88,39 @@ class WeightHandler:
|
||||
"linear2": block["ff_context"]["net"][2]
|
||||
}
|
||||
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):
|
||||
|
||||
Loading…
Reference in New Issue
Block a user