Dreambooth v1
This commit is contained in:
filipstrand
parent
627398f531
commit
bae49c8145
5
.gitignore
vendored
5
.gitignore
vendored
@ -18,3 +18,8 @@
|
||||
|
||||
# build/ generated from 'pip install -e .' in developer mode
|
||||
build/
|
||||
.gz
|
||||
.fdb_latexmk
|
||||
.fls
|
||||
.aux
|
||||
.log
|
||||
|
||||
146
README.md
146
README.md
@ -26,6 +26,13 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
|
||||
- [🔌 LoRA](#-lora)
|
||||
* [Multi-LoRA](#multi-lora)
|
||||
* [Supported LoRA formats (updated)](#supported-lora-formats-updated)
|
||||
- [🎛️ Dreambooth fine-tuning](#-dreambooth-fine-tuning)
|
||||
* [Training configuration](#training-configuration)
|
||||
* [Training example](#training-example)
|
||||
* [Resuming a training run](#resuming-a-training-run)
|
||||
* [Configuration details](#configuration-details)
|
||||
* [Memory issues](#memory-issues)
|
||||
* [Misc](#misc)
|
||||
- [🕹️ Controlnet](#%EF%B8%8F-controlnet)
|
||||
- [🚧 Current limitations](#-current-limitations)
|
||||
- [💡Workflow tips](#workflow-tips)
|
||||
@ -586,6 +593,140 @@ 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).
|
||||
|
||||
---
|
||||
|
||||
### 🎛 Dreambooth fine-tuning
|
||||
|
||||
As of release 5.0, MFLUX has support for fine-tuning your own LoRA adapters using the [Dreambooth](https://dreambooth.github.io) technique.
|
||||
|
||||
|
||||
|
||||
*This example shows the MFLUX training progression of the [included training example](#training-example) which is based on the [DreamBooth Dataset](https://github.com/google/dreambooth), also used in the [mlx-examples repo](https://github.com/ml-explore/mlx-examples/tree/main/flux#finetuning).*
|
||||
|
||||
#### Training configuration
|
||||
|
||||
To describe a training run, you need to provide a [training configuration](src/mflux/dreambooth/_example/train.json) file which specifies the details such as
|
||||
what training data to use and various parameters. To try it out, one of the easiest ways is to start from the
|
||||
provided [example configuration](src/mflux/dreambooth/_example/train.json) and simply use your own dataset and prompts by modifying the `examples` section of the json file.
|
||||
|
||||
#### Training example
|
||||
|
||||
A complete example ([training configuration](src/mflux/dreambooth/_example/train.json) + [dataset](src/mflux/dreambooth/_example/images)) is provided in this repository. To start a training run, go to the project folder `cd mflux`, and simply run:
|
||||
|
||||
```sh
|
||||
mflux-train --train-config src/mflux/dreambooth/_example/train.json
|
||||
```
|
||||
|
||||
By default, this will train an adapter with images of size `512x512` with a batch size of 1 and can take up to several hours to fully complete depending on your machine.
|
||||
If this task is too computationally demanding, see the section on [memory issues](#memory-issues) for tips on how to speed things up and what tradeoffs exists.
|
||||
|
||||
During training, MFLUX will output training checkpoints with artifacts (weights, states) according to what is specified in the configuration file.
|
||||
As specified in the file `train.json`, these files will be placed in a folder on the Desktop called `~/Desktop/train`, but this can of course be changed to any other path by adjusting the configuration.
|
||||
All training artifacts will be saved as self-contained zip file, which can later be pointed to [resume an existing training run](#resuming-a-training-run).
|
||||
To find the LoRA weights, simply unzip and look for the `adapter` safetensors file and [use it as you would with a regular downloaded LoRA adapter](#-lora).
|
||||
|
||||
#### Resuming a training run
|
||||
|
||||
The training process will continue to run until each training example has been used `num_epochs` times.
|
||||
For various reasons however, the user might choose to interrupt the process.
|
||||
To resume training for a given checkpoint, say `0001000_checkpoint.zip`, simply run:
|
||||
|
||||
```sh
|
||||
mflux-train --train-checkpoint 0001000_checkpoint.zip
|
||||
```
|
||||
|
||||
There are two nice properties of the training procedure:
|
||||
|
||||
- Fully deterministic (given a specified `seed` in the training configuration)
|
||||
- The complete training state (including optimizer state) is saved at each checkpoint.
|
||||
|
||||
Because of these, MFLUX has the ability to resume a training run from a previous checkpoint and have the results
|
||||
be *exactly* identical to a training run which was never interrupted in the first place.
|
||||
|
||||
*⚠️ Note: Everything but the dataset itself is contained within this zipfile, as the dataset can be quite large.
|
||||
The zip file will contain configuration files which point to the original dataset, so make sure that it is in the same place when resuming training*.
|
||||
|
||||
*⚠️ Note: One current limitation is that a training run can only be resumed if it has not yet been completed.
|
||||
In other words, only checkpoints that represents an interrupted training-run can be resumed and run until completion.*
|
||||
|
||||
#### Configuration details
|
||||
|
||||
Currently, MFLUX supports fine-tuning only for the transformer part of the model.
|
||||
In the training configuration, under `lora_layers`, you can specify which layers you want to train. The available ones are:
|
||||
|
||||
- `transformer_blocks`:
|
||||
- `attn.to_q`
|
||||
- `attn.to_k`
|
||||
- `attn.to_v`
|
||||
- `attn.add_q_proj`
|
||||
- `attn.add_k_proj`
|
||||
- `attn.add_v_proj`
|
||||
- `attn.to_out`
|
||||
- `attn.to_add_out`
|
||||
- `ff.linear1`
|
||||
- `ff.linear2`
|
||||
- `ff_context.linear1`
|
||||
- `ff_context.linear2`
|
||||
- `single_transformer_blocks`:
|
||||
- `proj_out`
|
||||
- `proj_mlp`
|
||||
- `attn.to_q`
|
||||
- `attn.to_k`
|
||||
- `attn.to_v`
|
||||
|
||||
The `block_range` under the respective layer category specifies which blocks to train.
|
||||
The maximum range available for the different layer categories are:
|
||||
|
||||
- `transformer_blocks`:
|
||||
- `start: 0`
|
||||
- `end: 19`
|
||||
- `single_transformer_blocks`:
|
||||
- `start: 0`
|
||||
- `end: 38`
|
||||
|
||||
*⚠️ Note: As the joint transformer blocks (`transformer_blocks`) - are placed earlier on in the sequence of computations, they will require more resources to train.
|
||||
In other words, training later layers, such as only the `single_transformer_blocks` should be faster. However, training too few / only later layers might result in a faster but unsuccessful training.*
|
||||
|
||||
*Under the `examples` section, there is an argument called `"path"` which specified where the images are located. This path is relative to the config file itself.*
|
||||
|
||||
#### Memory issues
|
||||
|
||||
Depending on the configuration of the training setup, fine-tuning can be quite memory intensive.
|
||||
In the worst case, if your Mac runs out of memory it might freeze completely and crash!
|
||||
|
||||
To avoid this, consider some of the following strategies to reduce memory requirements by adjusting the parameters in the training configuration:
|
||||
|
||||
- Use a quantized based model by setting `"quantize": 4` or `"quantize": 8`
|
||||
- For the `layer_types`, consider skipping some of the trainable layers (e.g. by not including `proj_out` etc.)
|
||||
- Use a lower `rank` value for the LoRA matrices.
|
||||
- Don't train all the `38` layers from `single_transformer_blocks` or all of the `19` layers from `transformer_blocks`
|
||||
- Use a smaller batch size, for example `"batch_size": 1`
|
||||
- Make sure your Mac is not busy with other background tasks that holds memory.
|
||||
|
||||
Applying some of these strategies, like how [train.json](src/mflux/dreambooth/_example/train.json) is set up by default,
|
||||
will allow a 32GB M1 Pro to perform a successful fine-tuning run.
|
||||
Note, however, that reducing the trainable parameters might lead to worse performance.
|
||||
|
||||
|
||||
*Additional techniques such as gradient checkpoint and other strategies might be implemented in the future.*
|
||||
|
||||
#### Misc
|
||||
|
||||
This feature is currently v1 and can be considered a bit experimental. Interfaces might change (configuration file setup etc.)
|
||||
The aim is to also gradually expand the scope of this feature with alternatives techniques, data augmentation etc.
|
||||
|
||||
- As with loading external LoRA adapters, the MFLUX training currently only supports training the transformer part of the network.
|
||||
- Sometimes, a model trained with the `dev` model might actually work better when applied to the `schnell` weights.
|
||||
- Currently, all training images are assumed to be in the resolution specified in the configuration file.
|
||||
- Loss curve can be a bit misleading/hard to read, sometimes it conveys little improvement over time, but actual image samples shows the real progress.
|
||||
- When plotting the loss during training, we label it as "validation loss" but is actually only the first 10 elements of the training examples for now. Future updates should support user inputs of separate validation images.
|
||||
- Training also works with the original model as quantized!
|
||||
- [For the curious, a motivation for the loss function can be found at here](src/mflux/dreambooth/optimization/_loss_derivation/dreambooth_loss.pdf).
|
||||
- Two great resources that heavily inspired this feature are the:
|
||||
- The fine-tuning script in [mlx-examples](https://github.com/ml-explore/mlx-examples/tree/main/flux#finetuning)
|
||||
- The original fine-tuning script in [Diffusers](https://huggingface.co/docs/diffusers/v0.11.0/en/training/dreambooth)
|
||||
|
||||
|
||||
---
|
||||
|
||||
### 🕹️ Controlnet
|
||||
@ -631,6 +772,7 @@ with different prompts and LoRA adapters active.
|
||||
- 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).
|
||||
- Dreambooth training currently does not support sending in training parameters as flags.
|
||||
|
||||
### 💡Workflow Tips
|
||||
|
||||
@ -642,14 +784,14 @@ with different prompts and LoRA adapters active.
|
||||
|
||||
### ✅ TODO
|
||||
|
||||
- [ ] LoRA fine-tuning (in progress, see [DreamBooth support #83](https://github.com/filipstrand/mflux/pull/83))
|
||||
- [ ] [FLUX.1 Tools](https://blackforestlabs.ai/flux-1-tools/)
|
||||
|
||||
### 🔬 Cool research / features to support
|
||||
- [ ] [PuLID](https://github.com/ToTheBeginning/PuLID)
|
||||
- [ ] [depth based controlnet](https://huggingface.co/InstantX/SD3-Controlnet-Depth) via [ml-depth-pro](https://github.com/apple/ml-depth-pro) or similar?
|
||||
- [ ] [RF-Inversion](https://github.com/filipstrand/mflux/issues/91)
|
||||
- [ ] [In-Context LoRA](https://github.com/ali-vilab/In-Context-LoRA)
|
||||
- [ ] [FLUX.1 Tools](https://blackforestlabs.ai/flux-1-tools/)
|
||||
|
||||
|
||||
### 🌱 Related projects
|
||||
|
||||
|
||||
@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "mflux"
|
||||
version = "0.4.1"
|
||||
version = "0.5.0"
|
||||
description = "A MLX port of FLUX based on the Huggingface Diffusers implementation."
|
||||
readme = "README.md"
|
||||
keywords = ["diffusers", "flux", "mlx"]
|
||||
@ -14,6 +14,7 @@ license = { file = "LICENSE" }
|
||||
requires-python = ">=3.10"
|
||||
dependencies = [
|
||||
"huggingface-hub>=0.24.5,<1.0",
|
||||
"matplotlib>=3.9.2,<4.0",
|
||||
"mlx>=0.20.0,<1.0",
|
||||
"numpy>=2.0.1,<3.0",
|
||||
"opencv-python>=4.10.0,<5.0",
|
||||
@ -26,6 +27,7 @@ dependencies = [
|
||||
# uv pip install https://github.com/anthonywu/sentencepiece/releases/download/0.2.1-py13dev/sentencepiece-0.2.1-cp313-cp313-macosx_11_0_arm64.whl
|
||||
"sentencepiece>=0.2.0,<1.0; python_version<'3.13'",
|
||||
"tokenizers>=0.20.3; python_version>='3.13'", # transformers -> tokenizers
|
||||
"toml>=0.10.2,<1.0",
|
||||
"torch>=2.3.1,<3.0; python_version<'3.13'",
|
||||
# torch dev builds: pip install --pre --index-url https://download.pytorch.org/whl/nightly
|
||||
"torch>=2.6.0.dev20241106; python_version>='3.13'",
|
||||
@ -53,8 +55,9 @@ 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"
|
||||
mflux-save = "mflux.save:main"
|
||||
mflux-train = "mflux.train:main"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
|
||||
BIN
src/mflux/assets/dreambooth.jpg
Normal file
BIN
src/mflux/assets/dreambooth.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 733 KiB |
@ -3,7 +3,6 @@ from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import mlx.core as mx
|
||||
from mlx import nn
|
||||
from tqdm import tqdm
|
||||
|
||||
from mflux.config.config import ConfigControlnet
|
||||
@ -26,6 +25,8 @@ 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
|
||||
from mflux.weights.weight_handler_lora import WeightHandlerLoRA
|
||||
from mflux.weights.weight_util import WeightUtil
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mflux.post_processing.generated_image import GeneratedImage
|
||||
@ -61,52 +62,29 @@ class Flux1Controlnet:
|
||||
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
|
||||
) # fmt: off
|
||||
|
||||
# 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
|
||||
# fmt: off
|
||||
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)
|
||||
# fmt: on
|
||||
|
||||
# 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, 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"],
|
||||
# Set the weights and quantize the model
|
||||
weights = WeightHandler.load_regular_weights(repo_id=model_config.model_name, local_path=local_path)
|
||||
self.bits = WeightUtil.set_weights_and_quantize(
|
||||
quantize_arg=quantize,
|
||||
weights=weights,
|
||||
vae=self.vae,
|
||||
transformer=self.transformer,
|
||||
t5_text_encoder=self.t5_text_encoder,
|
||||
clip_text_encoder=self.clip_text_encoder,
|
||||
)
|
||||
|
||||
if controlnet_quantization_level is None:
|
||||
self.transformer_controlnet.update(weights_controlnet)
|
||||
# Set LoRA weights
|
||||
lora_weights = WeightHandlerLoRA.load_lora_weights(transformer=self.transformer, lora_files=lora_paths, lora_scales=lora_scales) # fmt:off
|
||||
WeightHandlerLoRA.set_lora_weights(transformer=self.transformer, loras=lora_weights)
|
||||
|
||||
self.bits = None
|
||||
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 controlnet_quantization_level is not None:
|
||||
self.transformer_controlnet.update(weights_controlnet)
|
||||
# Set Controlnet weights
|
||||
weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer(controlnet_id=CONTROLNET_ID)
|
||||
self.transformer_controlnet = TransformerControlnet(model_config=model_config, num_blocks=weights_controlnet.config["num_layers"], num_single_blocks=weights_controlnet.config["num_single_layers"]) # fmt:off
|
||||
WeightUtil.set_controlnet_weights_and_quantize(
|
||||
quantize_arg=quantize,
|
||||
weights=weights_controlnet,
|
||||
transformer_controlnet=self.transformer_controlnet,
|
||||
)
|
||||
|
||||
def generate_image(
|
||||
self,
|
||||
@ -147,13 +125,13 @@ class Flux1Controlnet:
|
||||
# 2. Embed 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)
|
||||
prompt_embeds = self.t5_text_encoder(t5_tokens)
|
||||
pooled_prompt_embeds = self.clip_text_encoder(clip_tokens)
|
||||
|
||||
for t in time_steps:
|
||||
try:
|
||||
# Compute controlnet samples
|
||||
controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet.forward(
|
||||
controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet(
|
||||
t=t,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
@ -202,12 +180,6 @@ class Flux1Controlnet:
|
||||
controlnet_image_path=controlnet_image_path,
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
@ -35,7 +35,7 @@ class TransformerControlnet(nn.Module):
|
||||
|
||||
self.controlnet_single_blocks = [nn.Linear(3072, 3072) for _ in range(num_single_blocks)]
|
||||
|
||||
def forward(
|
||||
def __call__(
|
||||
self,
|
||||
t: int,
|
||||
prompt_embeds: mx.array,
|
||||
@ -51,16 +51,16 @@ class TransformerControlnet(nn.Module):
|
||||
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)
|
||||
text_embeddings = self.time_text_embed(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(height=config.height, width=config.width)
|
||||
ids = mx.concatenate((txt_ids, img_ids), axis=1)
|
||||
image_rotary_emb = self.pos_embed.forward(ids)
|
||||
image_rotary_emb = self.pos_embed(ids)
|
||||
|
||||
block_samples = ()
|
||||
for block in self.transformer_blocks:
|
||||
encoder_hidden_states, hidden_states = block.forward(
|
||||
encoder_hidden_states, hidden_states = block(
|
||||
hidden_states=hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
text_embeddings=text_embeddings,
|
||||
@ -78,7 +78,7 @@ class TransformerControlnet(nn.Module):
|
||||
|
||||
single_block_samples = ()
|
||||
for block in self.single_transformer_blocks:
|
||||
hidden_states = block.forward(
|
||||
hidden_states = block(
|
||||
hidden_states=hidden_states,
|
||||
text_embeddings=text_embeddings,
|
||||
rotary_embeddings=image_rotary_emb,
|
||||
|
||||
@ -5,21 +5,30 @@ import mlx.core as mx
|
||||
from huggingface_hub import snapshot_download
|
||||
from mlx.utils import tree_unflatten
|
||||
|
||||
from mflux.weights.weight_handler import MetaData
|
||||
from mflux.weights.weight_util import WeightUtil
|
||||
|
||||
|
||||
class WeightHandlerControlnet:
|
||||
def __init__(self, meta_data: MetaData, config: dict, controlnet_transformer: dict | None = None):
|
||||
self.meta_data = meta_data
|
||||
self.controlnet_transformer = controlnet_transformer
|
||||
self.config = config
|
||||
|
||||
@staticmethod
|
||||
def load_controlnet_transformer(controlnet_id: str) -> (dict, int):
|
||||
controlnet_path = Path(
|
||||
snapshot_download(repo_id=controlnet_id, allow_patterns=["*.safetensors", "config.json"])
|
||||
)
|
||||
def load_controlnet_transformer(controlnet_id: str) -> "WeightHandlerControlnet":
|
||||
controlnet_path = Path(snapshot_download(repo_id=controlnet_id, allow_patterns=["*.safetensors", "config.json"])) # fmt:off
|
||||
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())
|
||||
config = json.load(open(controlnet_path / "config.json"))
|
||||
|
||||
if quantization_level is not None:
|
||||
return tree_unflatten(weights), quantization_level
|
||||
return WeightHandlerControlnet(
|
||||
config=config,
|
||||
controlnet_transformer=tree_unflatten(weights),
|
||||
meta_data=MetaData(quantization_level=quantization_level),
|
||||
)
|
||||
|
||||
weights = [WeightUtil.reshape_weights(k, v) for k, v in weights]
|
||||
weights = WeightUtil.flatten(weights)
|
||||
@ -27,7 +36,11 @@ class WeightHandlerControlnet:
|
||||
|
||||
# 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
|
||||
return WeightHandlerControlnet(
|
||||
config=config,
|
||||
controlnet_transformer=weights,
|
||||
meta_data=MetaData(quantization_level=quantization_level)
|
||||
) # fmt:off
|
||||
|
||||
# Reshape and process the huggingface weights
|
||||
if "transformer_blocks" in weights:
|
||||
@ -41,5 +54,9 @@ class WeightHandlerControlnet:
|
||||
"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
|
||||
|
||||
return WeightHandlerControlnet(
|
||||
config=config,
|
||||
controlnet_transformer=weights,
|
||||
meta_data=MetaData(quantization_level=quantization_level)
|
||||
) # fmt:off
|
||||
|
||||
BIN
src/mflux/dreambooth/_example/images/01.jpeg
Normal file
BIN
src/mflux/dreambooth/_example/images/01.jpeg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 662 KiB |
BIN
src/mflux/dreambooth/_example/images/02.jpeg
Normal file
BIN
src/mflux/dreambooth/_example/images/02.jpeg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 1.3 MiB |
BIN
src/mflux/dreambooth/_example/images/03.jpeg
Normal file
BIN
src/mflux/dreambooth/_example/images/03.jpeg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
BIN
src/mflux/dreambooth/_example/images/04.jpeg
Normal file
BIN
src/mflux/dreambooth/_example/images/04.jpeg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
BIN
src/mflux/dreambooth/_example/images/05.jpeg
Normal file
BIN
src/mflux/dreambooth/_example/images/05.jpeg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
67
src/mflux/dreambooth/_example/train.json
Normal file
67
src/mflux/dreambooth/_example/train.json
Normal file
@ -0,0 +1,67 @@
|
||||
{
|
||||
"model": "dev",
|
||||
"seed": 42,
|
||||
"steps": 20,
|
||||
"guidance": 3.0,
|
||||
"quantize": 4,
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"training_loop": {
|
||||
"num_epochs": 100,
|
||||
"batch_size": 1
|
||||
},
|
||||
"optimizer": {
|
||||
"name": "AdamW",
|
||||
"learning_rate": 1e-4
|
||||
},
|
||||
"save": {
|
||||
"output_path": "~/Desktop/train",
|
||||
"checkpoint_frequency": 10
|
||||
},
|
||||
"instrumentation": {
|
||||
"plot_frequency": 1,
|
||||
"generate_image_frequency": 20,
|
||||
"validation_prompt": "photo of sks dog"
|
||||
},
|
||||
"lora_layers": {
|
||||
"single_transformer_blocks" : {
|
||||
"block_range": {
|
||||
"start": 0,
|
||||
"end": 38
|
||||
},
|
||||
"layer_types": [
|
||||
"proj_out",
|
||||
"proj_mlp",
|
||||
"attn.to_q",
|
||||
"attn.to_k",
|
||||
"attn.to_v"
|
||||
],
|
||||
"lora_rank": 4
|
||||
}
|
||||
},
|
||||
"examples": {
|
||||
"path": "images/",
|
||||
"images": [
|
||||
{
|
||||
"image": "01.jpeg",
|
||||
"prompt": "photo of sks dog"
|
||||
},
|
||||
{
|
||||
"image": "02.jpeg",
|
||||
"prompt": "photo of sks dog"
|
||||
},
|
||||
{
|
||||
"image": "03.jpeg",
|
||||
"prompt": "photo of sks dog"
|
||||
},
|
||||
{
|
||||
"image": "04.jpeg",
|
||||
"prompt": "photo of sks dog"
|
||||
},
|
||||
{
|
||||
"image": "05.jpeg",
|
||||
"prompt": "photo of sks dog"
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
0
src/mflux/dreambooth/dataset/__init__.py
Normal file
0
src/mflux/dreambooth/dataset/__init__.py
Normal file
28
src/mflux/dreambooth/dataset/batch.py
Normal file
28
src/mflux/dreambooth/dataset/batch.py
Normal file
@ -0,0 +1,28 @@
|
||||
from pathlib import Path
|
||||
from random import Random
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
|
||||
class Example:
|
||||
def __init__(
|
||||
self,
|
||||
example_id: int,
|
||||
prompt: str,
|
||||
image_path: str | Path,
|
||||
encoded_image: mx.array,
|
||||
prompt_embeds: mx.array,
|
||||
pooled_prompt_embeds: mx.array,
|
||||
):
|
||||
self.example_id = example_id
|
||||
self.prompt = prompt
|
||||
self.image_name = str(image_path)
|
||||
self.clean_latents = encoded_image
|
||||
self.prompt_embeds = prompt_embeds
|
||||
self.pooled_prompt_embeds = pooled_prompt_embeds
|
||||
|
||||
|
||||
class Batch:
|
||||
def __init__(self, examples: list[Example], rng: Random):
|
||||
self.rng = rng
|
||||
self.examples = examples
|
||||
80
src/mflux/dreambooth/dataset/dataset.py
Normal file
80
src/mflux/dreambooth/dataset/dataset.py
Normal file
@ -0,0 +1,80 @@
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
import PIL.Image
|
||||
from mlx import nn
|
||||
from tqdm import tqdm
|
||||
|
||||
from mflux import Flux1, ImageUtil
|
||||
from mflux.dreambooth.dataset.batch import Example
|
||||
from mflux.dreambooth.dataset.dreambooth_preprocessing import DreamBoothPreProcessing
|
||||
from mflux.dreambooth.state.training_spec import ExampleSpec
|
||||
from mflux.post_processing.array_util import ArrayUtil
|
||||
|
||||
|
||||
class Dataset:
|
||||
def __init__(self, examples: list[Example]):
|
||||
self.examples = examples
|
||||
|
||||
@staticmethod
|
||||
def prepare_dataset(
|
||||
flux: Flux1,
|
||||
raw_data: list[ExampleSpec],
|
||||
width: int,
|
||||
height: int,
|
||||
) -> "Dataset":
|
||||
# Encode the original examples (image and text)
|
||||
examples = Dataset._create_examples(flux, raw_data, width=width, height=height)
|
||||
|
||||
# Expend the original dataset to get more training data with variations
|
||||
augmented_examples = []
|
||||
for example in examples:
|
||||
[augmented_examples.append(variation) for variation in DreamBoothPreProcessing.augment(example)]
|
||||
|
||||
# Dataset is now prepared
|
||||
return Dataset(augmented_examples)
|
||||
|
||||
def size(self) -> int:
|
||||
return len(self.examples)
|
||||
|
||||
@staticmethod
|
||||
def _create_examples(
|
||||
flux: Flux1,
|
||||
raw_data: list[ExampleSpec],
|
||||
width: int,
|
||||
height: int,
|
||||
) -> list[Example]:
|
||||
examples = []
|
||||
for i, entry in enumerate(tqdm(raw_data, desc="Encoding original dataset")):
|
||||
# Encode the image
|
||||
encoded_image = Dataset._encode_image(flux.vae, entry.image, width=width, height=height)
|
||||
|
||||
# Encode the prompt
|
||||
prompt_embeds = flux.t5_text_encoder(flux.t5_tokenizer.tokenize(entry.prompt))
|
||||
pooled_prompt_embeds = flux.clip_text_encoder(flux.clip_tokenizer.tokenize(entry.prompt))
|
||||
|
||||
# Create the example object
|
||||
example = Example(
|
||||
example_id=i,
|
||||
prompt=entry.prompt,
|
||||
image_path=entry.image,
|
||||
encoded_image=encoded_image,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
)
|
||||
examples.append(example)
|
||||
|
||||
# Evaluate to enable progress tracking
|
||||
mx.eval(encoded_image)
|
||||
mx.eval(prompt_embeds)
|
||||
mx.eval(pooled_prompt_embeds)
|
||||
|
||||
return examples
|
||||
|
||||
@staticmethod
|
||||
def _encode_image(vae: nn.Module, image_path: Path, width: int, height: int) -> mx.array:
|
||||
image = PIL.Image.open(image_path.resolve()).convert("RGB")
|
||||
scaled_user_image = ImageUtil.scale_to_dimensions(image, target_width=width, target_height=height)
|
||||
encoded = vae.encode(ImageUtil.to_array(scaled_user_image))
|
||||
latents = ArrayUtil.pack_latents(encoded, width=width, height=height)
|
||||
return latents
|
||||
8
src/mflux/dreambooth/dataset/dreambooth_preprocessing.py
Normal file
8
src/mflux/dreambooth/dataset/dreambooth_preprocessing.py
Normal file
@ -0,0 +1,8 @@
|
||||
from mflux.dreambooth.dataset.batch import Example
|
||||
|
||||
|
||||
class DreamBoothPreProcessing:
|
||||
@staticmethod
|
||||
def augment(example: Example) -> list[Example]:
|
||||
# Currently this does nothing.
|
||||
return [example]
|
||||
162
src/mflux/dreambooth/dataset/iterator.py
Normal file
162
src/mflux/dreambooth/dataset/iterator.py
Normal file
@ -0,0 +1,162 @@
|
||||
import datetime
|
||||
import json
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from mflux.dreambooth.dataset.batch import Batch
|
||||
from mflux.dreambooth.dataset.dataset import Dataset
|
||||
from mflux.dreambooth.state.training_spec import TrainingSpec
|
||||
from mflux.dreambooth.state.zip_util import ZipUtil
|
||||
|
||||
|
||||
class Iterator:
|
||||
def __init__(
|
||||
self,
|
||||
dataset: Dataset,
|
||||
batch_size: int,
|
||||
num_epochs: int | None = None,
|
||||
seed: int | None = None,
|
||||
position: int = 0,
|
||||
epoch: int = 0,
|
||||
num_iterations: int = 0,
|
||||
current_permutation: list[int] | None = None,
|
||||
rng_state: tuple | None = None,
|
||||
start_date_time: datetime.datetime | None = None,
|
||||
):
|
||||
self.dataset = dataset
|
||||
self.batch_size = batch_size
|
||||
self.total_examples = dataset.size()
|
||||
self.num_epochs = num_epochs
|
||||
self.seed = seed
|
||||
self.rng = random.Random(seed)
|
||||
|
||||
self._position = position
|
||||
self._epoch = epoch
|
||||
self.num_iterations = num_iterations
|
||||
self.start_date_time = start_date_time or datetime.datetime.now()
|
||||
|
||||
if rng_state is not None:
|
||||
self.rng.setstate(rng_state)
|
||||
|
||||
if current_permutation is not None:
|
||||
self._current_permutation = current_permutation
|
||||
else:
|
||||
self._initialize_permutation()
|
||||
|
||||
@staticmethod
|
||||
def from_spec(training_spec: TrainingSpec, dataset: Dataset) -> "Iterator":
|
||||
if training_spec.training_loop.iterator_state_path is not None:
|
||||
return Iterator.from_json(
|
||||
training_spec=training_spec,
|
||||
iterator_path=training_spec.training_loop.iterator_state_path,
|
||||
dataset=dataset,
|
||||
)
|
||||
|
||||
return Iterator(
|
||||
seed=training_spec.seed,
|
||||
dataset=dataset,
|
||||
batch_size=training_spec.training_loop.batch_size,
|
||||
num_epochs=training_spec.training_loop.num_epochs,
|
||||
)
|
||||
|
||||
def _initialize_permutation(self):
|
||||
self._current_permutation = list(range(self.total_examples))
|
||||
self.rng.shuffle(self._current_permutation)
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
# Check if we've reached the specified number of epochs
|
||||
if self.num_epochs is not None and self._epoch >= self.num_epochs:
|
||||
raise StopIteration()
|
||||
|
||||
# If we've used all examples, start new epoch
|
||||
if self._position >= self.total_examples:
|
||||
self._position = 0
|
||||
self._epoch += 1
|
||||
|
||||
# Check again after incrementing epoch
|
||||
if self.num_epochs is not None and self._epoch >= self.num_epochs:
|
||||
raise StopIteration()
|
||||
|
||||
self._initialize_permutation()
|
||||
|
||||
# Calculate batch size for this iteration
|
||||
remaining = self.total_examples - self._position
|
||||
current_batch_size = min(self.batch_size, remaining)
|
||||
|
||||
# Get indices for current batch
|
||||
batch_indices = self._current_permutation[self._position : self._position + current_batch_size]
|
||||
|
||||
# Update position
|
||||
self._position += current_batch_size
|
||||
|
||||
# Increment iteration counter
|
||||
self.num_iterations += 1
|
||||
|
||||
# Get examples using indices
|
||||
examples = [self.dataset.examples[i] for i in batch_indices]
|
||||
|
||||
return Batch(examples=examples, rng=self.rng)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Returns the current state as a dictionary."""
|
||||
return {
|
||||
"position": self._position,
|
||||
"epoch": self._epoch,
|
||||
"num_iterations": self.num_iterations,
|
||||
"current_permutation": self._current_permutation,
|
||||
"rng_state": list(self.rng.getstate()),
|
||||
"batch_size": self.batch_size,
|
||||
"num_epochs": self.num_epochs,
|
||||
"start_date_time": self.start_date_time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, state_dict: dict[str, Any], dataset: Dataset) -> "Iterator":
|
||||
"""Creates an iterator from a state dictionary."""
|
||||
# Convert the RNG state elements to tuples
|
||||
rng_state = state_dict["rng_state"]
|
||||
rng_state = (rng_state[0], tuple(rng_state[1]), rng_state[2])
|
||||
|
||||
return cls(
|
||||
dataset=dataset,
|
||||
batch_size=state_dict["batch_size"],
|
||||
num_epochs=state_dict["num_epochs"],
|
||||
position=state_dict["position"],
|
||||
epoch=state_dict["epoch"],
|
||||
num_iterations=state_dict["num_iterations"],
|
||||
current_permutation=state_dict["current_permutation"],
|
||||
rng_state=rng_state,
|
||||
start_date_time=datetime.datetime.strptime(state_dict["start_date_time"], "%Y-%m-%d %H:%M:%S"),
|
||||
)
|
||||
|
||||
def to_json(self, filepath: str) -> None:
|
||||
"""Save iterator state to a JSON file."""
|
||||
with open(filepath, "w") as f:
|
||||
json.dump(self.to_dict(), f, indent=4)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, training_spec: TrainingSpec, iterator_path: str, dataset: Dataset) -> "Iterator":
|
||||
data = ZipUtil.unzip(
|
||||
zip_path=training_spec.checkpoint_path,
|
||||
filename=iterator_path,
|
||||
loader=lambda x: json.load(open(x, "r"))
|
||||
) # fmt: off
|
||||
return cls.from_dict(data, dataset)
|
||||
|
||||
def get_validation_batch(self) -> Batch:
|
||||
# This is of course a misleading label, and future updates should include actual validation image inputs.
|
||||
examples = self.dataset.examples
|
||||
if len(examples) > 10:
|
||||
examples = examples[0:10]
|
||||
|
||||
return Batch(examples=examples, rng=self.rng)
|
||||
|
||||
def total_number_of_steps(self) -> int:
|
||||
return self.total_examples * self.num_epochs
|
||||
|
||||
def save(self, path: Path) -> None:
|
||||
self.to_json(str(path))
|
||||
71
src/mflux/dreambooth/dreambooth.py
Normal file
71
src/mflux/dreambooth/dreambooth.py
Normal file
@ -0,0 +1,71 @@
|
||||
from mlx import nn
|
||||
from tqdm import tqdm
|
||||
|
||||
from mflux import Flux1
|
||||
from mflux.config.runtime_config import RuntimeConfig
|
||||
from mflux.dreambooth.optimization.dreambooth_loss import DreamBoothLoss
|
||||
from mflux.dreambooth.state.training_spec import TrainingSpec
|
||||
from mflux.dreambooth.state.training_state import TrainingState
|
||||
from mflux.dreambooth.statistics.plotter import Plotter
|
||||
from mflux.weights.weight_handler_lora import WeightHandlerLoRA
|
||||
|
||||
|
||||
class DreamBooth:
|
||||
@staticmethod
|
||||
def train(
|
||||
flux: Flux1,
|
||||
runtime_config: RuntimeConfig,
|
||||
training_spec: TrainingSpec,
|
||||
training_state: TrainingState
|
||||
): # fmt:off
|
||||
# Freeze the model and assign the LoRA layers to the model
|
||||
flux.freeze()
|
||||
WeightHandlerLoRA.set_lora_layers(
|
||||
transformer_module=flux.transformer,
|
||||
lora_layers=training_state.lora_layers
|
||||
) # fmt:off
|
||||
|
||||
# Define loss computation as a function of a batch 'b'
|
||||
train_step_function = nn.value_and_grad(
|
||||
model=flux,
|
||||
fn=lambda b: DreamBoothLoss.compute_loss(flux, runtime_config, b)
|
||||
) # fmt: off
|
||||
|
||||
# Setup progress bar
|
||||
batches = tqdm(
|
||||
training_state.iterator,
|
||||
total=training_state.iterator.total_number_of_steps(),
|
||||
initial=training_state.iterator.num_iterations,
|
||||
)
|
||||
|
||||
# Training loop
|
||||
for batch in batches:
|
||||
# Perform one gradient update on the LoRA the weights
|
||||
loss, grads = train_step_function(batch)
|
||||
training_state.optimizer.optimizer.update(model=flux, gradients=grads)
|
||||
del loss, grads
|
||||
|
||||
# Plot loss progress periodically
|
||||
if training_state.should_plot_loss(training_spec):
|
||||
validation_batch = training_state.iterator.get_validation_batch()
|
||||
validation_loss = DreamBoothLoss.compute_loss(flux, runtime_config, validation_batch)
|
||||
training_state.statistics.append_values(step=training_state.iterator.num_iterations, loss=validation_loss) # fmt: off
|
||||
Plotter.update_loss_plot(training_spec=training_spec, training_state=training_state)
|
||||
del validation_loss
|
||||
|
||||
# Generate a test image from the model periodically
|
||||
if training_state.should_generate_image(training_spec):
|
||||
image = flux.generate_image(
|
||||
seed=training_spec.seed,
|
||||
config=runtime_config.config,
|
||||
prompt=training_spec.instrumentation.validation_prompt,
|
||||
) # fmt: off
|
||||
image.save(path=training_state.get_current_validation_image_path(training_spec))
|
||||
del image
|
||||
|
||||
# Save checkpoint periodically
|
||||
if training_state.should_save(training_spec):
|
||||
training_state.save(training_spec)
|
||||
|
||||
# Save the final state
|
||||
training_state.save(training_spec)
|
||||
74
src/mflux/dreambooth/dreambooth_initializer.py
Normal file
74
src/mflux/dreambooth/dreambooth_initializer.py
Normal file
@ -0,0 +1,74 @@
|
||||
import mlx.core.random as random
|
||||
|
||||
from mflux import Config, Flux1, ModelConfig
|
||||
from mflux.config.runtime_config import RuntimeConfig
|
||||
from mflux.dreambooth.dataset.dataset import Dataset
|
||||
from mflux.dreambooth.dataset.iterator import Iterator
|
||||
from mflux.dreambooth.lora_layers.lora_layers import LoRALayers
|
||||
from mflux.dreambooth.optimization.optimizer import Optimizer
|
||||
from mflux.dreambooth.state.training_spec import TrainingSpec
|
||||
from mflux.dreambooth.state.training_state import TrainingState
|
||||
from mflux.dreambooth.statistics.statistics import Statistics
|
||||
|
||||
|
||||
class DreamBoothInitializer:
|
||||
@staticmethod
|
||||
def initialize(
|
||||
config_path: str | None,
|
||||
checkpoint_path: str | None,
|
||||
) -> (Flux1, RuntimeConfig, TrainingSpec, TrainingState):
|
||||
# The training specification describing the details of the training process. It is resolved
|
||||
# differently depending on if training starts from scratch or resumes from checkpoint.
|
||||
training_spec = TrainingSpec.resolve(
|
||||
config_path=config_path,
|
||||
checkpoint_path=checkpoint_path
|
||||
) # fmt: off
|
||||
|
||||
# Set global random seed to make training deterministic
|
||||
random.seed(training_spec.seed)
|
||||
|
||||
# Load the model
|
||||
flux = Flux1(
|
||||
model_config=ModelConfig.from_alias(training_spec.model),
|
||||
quantize=training_spec.quantize,
|
||||
)
|
||||
runtime_config = RuntimeConfig(
|
||||
model_config=ModelConfig.from_alias(training_spec.model),
|
||||
config=Config(
|
||||
num_inference_steps=training_spec.steps,
|
||||
width=training_spec.width,
|
||||
height=training_spec.height,
|
||||
guidance=training_spec.guidance,
|
||||
),
|
||||
)
|
||||
|
||||
# Create the optimizer
|
||||
optimizer = Optimizer.from_spec(training_spec)
|
||||
|
||||
# Create the LoRA layers by matching them against the corresponding Flux layers
|
||||
lora_layers = LoRALayers.from_spec(flux=flux, training_spec=training_spec)
|
||||
|
||||
# Prepare the fine-tuning dataset and create the iterator
|
||||
dataset = Dataset.prepare_dataset(
|
||||
flux=flux,
|
||||
raw_data=training_spec.examples,
|
||||
width=training_spec.width,
|
||||
height=training_spec.height,
|
||||
)
|
||||
iterator = Iterator.from_spec(
|
||||
training_spec=training_spec,
|
||||
dataset=dataset
|
||||
) # fmt: off
|
||||
|
||||
# Setup loss statistics
|
||||
statistics = Statistics.from_spec(training_spec=training_spec)
|
||||
|
||||
# The training state consisting of everything that moves during training
|
||||
training_state = TrainingState(
|
||||
optimizer=optimizer,
|
||||
lora_layers=lora_layers,
|
||||
iterator=iterator,
|
||||
statistics=statistics,
|
||||
)
|
||||
|
||||
return flux, runtime_config, training_spec, training_state
|
||||
0
src/mflux/dreambooth/lora_layers/__init__.py
Normal file
0
src/mflux/dreambooth/lora_layers/__init__.py
Normal file
20
src/mflux/dreambooth/lora_layers/fused_linear_lora_layer.py
Normal file
20
src/mflux/dreambooth/lora_layers/fused_linear_lora_layer.py
Normal file
@ -0,0 +1,20 @@
|
||||
import mlx.core as mx
|
||||
from mlx import nn
|
||||
|
||||
from mflux.dreambooth.lora_layers.linear_lora_layer import LoRALinear
|
||||
|
||||
|
||||
class FusedLoRALinear(nn.Module):
|
||||
def __init__(self, base_linear: nn.Linear | nn.QuantizedLinear, loras: list[LoRALinear]):
|
||||
super().__init__()
|
||||
self.base_linear = base_linear
|
||||
self.loras = loras
|
||||
|
||||
def __call__(self, x):
|
||||
base_out = self.base_linear(x)
|
||||
|
||||
lora_out = mx.zeros_like(base_out)
|
||||
for lora in self.loras:
|
||||
lora_out += lora.scale * mx.matmul(mx.matmul(x, lora.lora_A), lora.lora_B)
|
||||
|
||||
return base_out + lora_out
|
||||
53
src/mflux/dreambooth/lora_layers/linear_lora_layer.py
Normal file
53
src/mflux/dreambooth/lora_layers/linear_lora_layer.py
Normal file
@ -0,0 +1,53 @@
|
||||
import math
|
||||
|
||||
import mlx.core as mx
|
||||
from mlx import nn
|
||||
|
||||
|
||||
class LoRALinear(nn.Module):
|
||||
@staticmethod
|
||||
def from_linear(
|
||||
linear: nn.Linear | nn.QuantizedLinear,
|
||||
r: int = 16,
|
||||
scale: float = 1.0,
|
||||
):
|
||||
output_dims, input_dims = linear.weight.shape
|
||||
if isinstance(linear, nn.QuantizedLinear):
|
||||
input_dims *= 32 // linear.bits
|
||||
lora_lin = LoRALinear(
|
||||
input_dims=input_dims,
|
||||
output_dims=output_dims,
|
||||
r=r,
|
||||
scale=scale,
|
||||
)
|
||||
lora_lin.linear = linear
|
||||
return lora_lin
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_dims: int,
|
||||
output_dims: int,
|
||||
r: int = 8,
|
||||
scale: float = 1.0,
|
||||
bias: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.linear = nn.Linear(input_dims, output_dims, bias=bias)
|
||||
self.scale = scale
|
||||
scale = 1 / math.sqrt(input_dims)
|
||||
|
||||
self.lora_A = mx.random.uniform(
|
||||
low=-scale,
|
||||
high=scale,
|
||||
shape=(input_dims, r),
|
||||
)
|
||||
self.lora_B = mx.random.uniform(
|
||||
low=-scale,
|
||||
high=scale,
|
||||
shape=(r, output_dims),
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
base_out = self.linear(x)
|
||||
lora_out = mx.matmul(mx.matmul(x, self.lora_A), self.lora_B)
|
||||
return base_out + self.scale * lora_out
|
||||
239
src/mflux/dreambooth/lora_layers/lora_layers.py
Normal file
239
src/mflux/dreambooth/lora_layers/lora_layers.py
Normal file
@ -0,0 +1,239 @@
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import mlx
|
||||
import mlx.core as mx
|
||||
from mlx import nn
|
||||
from mlx.utils import tree_flatten
|
||||
|
||||
from mflux.dreambooth.lora_layers.linear_lora_layer import LoRALinear
|
||||
from mflux.dreambooth.state.training_spec import SingleTransformerBlocks, TrainingSpec, TransformerBlocks
|
||||
from mflux.dreambooth.state.zip_util import ZipUtil
|
||||
from mflux.models.transformer.joint_transformer_block import JointTransformerBlock
|
||||
from mflux.models.transformer.single_transformer_block import SingleTransformerBlock
|
||||
from mflux.post_processing.generated_image import GeneratedImage
|
||||
from mflux.weights.weight_handler import MetaData, WeightHandler
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mflux import Flux1
|
||||
|
||||
|
||||
class LoRALayers:
|
||||
def __init__(self, weights: "WeightHandler"):
|
||||
self.layers = weights
|
||||
|
||||
@staticmethod
|
||||
def from_spec(flux: "Flux1", training_spec: TrainingSpec) -> "LoRALayers":
|
||||
if training_spec.lora_layers.state_path is not None:
|
||||
# Load from state if present in the spec
|
||||
from mflux.weights.weight_handler_lora import WeightHandlerLoRA
|
||||
|
||||
weights = ZipUtil.unzip(
|
||||
zip_path=training_spec.checkpoint_path,
|
||||
filename=training_spec.lora_layers.state_path,
|
||||
loader=lambda x: WeightHandlerLoRA.load_lora_weights(
|
||||
transformer=flux.transformer, lora_files=[x], lora_scales=[1.0]
|
||||
),
|
||||
)
|
||||
return LoRALayers(weights=weights[0])
|
||||
else:
|
||||
# Construct the LoRA weights from the spec
|
||||
transformer_lora_layers = {}
|
||||
single_transformer_lora_layers = {}
|
||||
|
||||
if training_spec.lora_layers.transformer_blocks:
|
||||
transformer_lora_layers = LoRALayers._construct_layers(
|
||||
blocks=flux.transformer.transformer_blocks,
|
||||
block_spec=training_spec.lora_layers.transformer_blocks,
|
||||
block_prefix="transformer.transformer_blocks",
|
||||
)
|
||||
|
||||
if training_spec.lora_layers.single_transformer_blocks:
|
||||
single_transformer_lora_layers = LoRALayers._construct_layers(
|
||||
blocks=flux.transformer.single_transformer_blocks,
|
||||
block_spec=training_spec.lora_layers.single_transformer_blocks,
|
||||
block_prefix="transformer.single_transformer_blocks",
|
||||
)
|
||||
|
||||
lora_layers = {**transformer_lora_layers, **single_transformer_lora_layers}
|
||||
|
||||
weights = WeightHandler(
|
||||
meta_data=MetaData(is_mflux=True),
|
||||
transformer=mlx.utils.tree_unflatten(list(lora_layers.items()))['transformer'],
|
||||
) # fmt:off
|
||||
|
||||
return LoRALayers(weights=weights)
|
||||
|
||||
@staticmethod
|
||||
def _construct_layers(
|
||||
block_spec: TransformerBlocks | SingleTransformerBlocks,
|
||||
blocks: list[JointTransformerBlock] | list[SingleTransformerBlock],
|
||||
block_prefix: str,
|
||||
) -> dict:
|
||||
start = block_spec.block_range.start
|
||||
end = block_spec.block_range.end
|
||||
lora_layers = {}
|
||||
for i in range(start, end):
|
||||
block = blocks[i]
|
||||
|
||||
for layer_type in block_spec.layer_types:
|
||||
original_layer = LoRALayers._get_nested_attr(block, layer_type)
|
||||
is_list = isinstance(original_layer, list)
|
||||
|
||||
lora_layer = LoRALinear.from_linear(
|
||||
linear=original_layer[0] if is_list else original_layer,
|
||||
r=block_spec.lora_rank,
|
||||
)
|
||||
layer_path = f"{block_prefix}.{i}.{layer_type}"
|
||||
|
||||
lora_layers[layer_path] = [lora_layer] if is_list else lora_layer
|
||||
|
||||
return lora_layers
|
||||
|
||||
@staticmethod
|
||||
def transformer_dict_from_template(weights: dict, transformer: nn.Module, scale: float) -> dict:
|
||||
lora_layers = {}
|
||||
for key in weights.keys():
|
||||
if key.endswith(".lora_A"):
|
||||
base_path = key[: -len(".lora_A")]
|
||||
parts = base_path.split(".")
|
||||
|
||||
if parts[1] == "transformer_blocks":
|
||||
LoRALayers._handle_transformer_blocks(
|
||||
weights=weights,
|
||||
scale=scale,
|
||||
transformer=transformer,
|
||||
lora_layers=lora_layers,
|
||||
base_path=base_path,
|
||||
)
|
||||
|
||||
if parts[1] == "single_transformer_blocks":
|
||||
LoRALayers._handle_single_transformer_blocks(
|
||||
weights=weights,
|
||||
scale=scale,
|
||||
transformer=transformer,
|
||||
lora_layers=lora_layers,
|
||||
base_path=base_path,
|
||||
)
|
||||
|
||||
return lora_layers
|
||||
|
||||
@staticmethod
|
||||
def _handle_transformer_blocks(
|
||||
weights: dict, scale: float, transformer: nn.Module, lora_layers: dict, base_path: str
|
||||
):
|
||||
parts = base_path.split(".")
|
||||
block_idx = int(parts[2])
|
||||
module_name = parts[3]
|
||||
attr_name = parts[4]
|
||||
block = transformer.transformer_blocks[block_idx]
|
||||
module = getattr(block, module_name)
|
||||
original_layer = getattr(module, attr_name)
|
||||
|
||||
if len(parts) == 6:
|
||||
original_layer = original_layer[0] # Special case here
|
||||
|
||||
# Create LoRA layer
|
||||
lora_A = weights[f"{base_path}.lora_A"]
|
||||
rank = lora_A.shape[1]
|
||||
lora_layer = LoRALinear.from_linear(linear=original_layer, r=rank, scale=scale)
|
||||
|
||||
# Set the weights
|
||||
lora_layer.lora_A = weights[f"{base_path}.lora_A"]
|
||||
lora_layer.lora_B = weights[f"{base_path}.lora_B"]
|
||||
|
||||
# Store the layer
|
||||
lora_layers[base_path] = lora_layer
|
||||
|
||||
@staticmethod
|
||||
def _handle_single_transformer_blocks(
|
||||
weights: dict, scale: float, transformer: nn.Module, lora_layers: dict, base_path: str
|
||||
):
|
||||
parts = base_path.split(".")
|
||||
block_idx = int(parts[2])
|
||||
module_name = parts[3]
|
||||
if len(parts) == 4:
|
||||
original_layer = getattr(transformer.single_transformer_blocks[block_idx], module_name)
|
||||
elif len(parts) == 5:
|
||||
attr_name = parts[4]
|
||||
block = transformer.single_transformer_blocks[block_idx]
|
||||
module = getattr(block, module_name)
|
||||
original_layer = getattr(module, attr_name)
|
||||
|
||||
# Create LoRA layer
|
||||
lora_A = weights[f"{base_path}.lora_A"]
|
||||
rank = lora_A.shape[1]
|
||||
lora_layer = LoRALinear.from_linear(linear=original_layer, r=rank, scale=scale)
|
||||
|
||||
# Set the weights
|
||||
lora_layer.lora_A = weights[f"{base_path}.lora_A"]
|
||||
lora_layer.lora_B = weights[f"{base_path}.lora_B"]
|
||||
|
||||
# Store the layer
|
||||
lora_layers[base_path] = lora_layer
|
||||
|
||||
@staticmethod
|
||||
def set_transformer_block(transformer_block, dictionary: dict):
|
||||
for key, val in dictionary.items():
|
||||
if key == "attn":
|
||||
LoRALayers._set_attribute(transformer_block, key, val, "to_q")
|
||||
LoRALayers._set_attribute(transformer_block, key, val, "to_k")
|
||||
LoRALayers._set_attribute(transformer_block, key, val, "to_v")
|
||||
LoRALayers._set_attribute(transformer_block, key, val, "to_out")
|
||||
LoRALayers._set_attribute(transformer_block, key, val, "add_q_proj")
|
||||
LoRALayers._set_attribute(transformer_block, key, val, "add_k_proj")
|
||||
LoRALayers._set_attribute(transformer_block, key, val, "add_v_proj")
|
||||
LoRALayers._set_attribute(transformer_block, key, val, "to_add_out")
|
||||
elif key == "ff" or key == "ff_context":
|
||||
LoRALayers._set_attribute(transformer_block, key, val, "linear1")
|
||||
LoRALayers._set_attribute(transformer_block, key, val, "linear2")
|
||||
elif key == "norm1" or key == "norm1_context":
|
||||
LoRALayers._set_attribute(transformer_block, key, val, "linear")
|
||||
else:
|
||||
raise Exception("Could not set LoRA weights")
|
||||
|
||||
@staticmethod
|
||||
def set_single_transformer_block(single_transformer_block, dictionary: dict):
|
||||
for key, val in dictionary.items():
|
||||
if key == "attn":
|
||||
LoRALayers._set_attribute(single_transformer_block, key, val, "to_q")
|
||||
LoRALayers._set_attribute(single_transformer_block, key, val, "to_k")
|
||||
LoRALayers._set_attribute(single_transformer_block, key, val, "to_v")
|
||||
elif key == "norm":
|
||||
LoRALayers._set_attribute(single_transformer_block, key, val, "linear")
|
||||
elif key == "proj_mlp" or key == "proj_out":
|
||||
single_transformer_block[key] = val
|
||||
else:
|
||||
raise Exception("Could not set LoRA weights")
|
||||
|
||||
@staticmethod
|
||||
def _set_attribute(block, key: str, val: dict, name: str):
|
||||
if block[key].get(name, False) and val.get(name, False):
|
||||
block[key][name] = val[name]
|
||||
|
||||
@staticmethod
|
||||
def _get_nested_attr(obj, attr_path):
|
||||
attrs = attr_path.split(".")
|
||||
for attr in attrs:
|
||||
obj = getattr(obj, attr)
|
||||
return obj
|
||||
|
||||
def save(self, path: Path, training_spec: TrainingSpec) -> None:
|
||||
weights = {}
|
||||
for entry in tree_flatten(self.layers.transformer):
|
||||
name = entry[0]
|
||||
weight = entry[1]
|
||||
if name.endswith(".lora_A") or name.endswith(".lora_B"):
|
||||
weights[name] = weight
|
||||
|
||||
weights = {key: mx.transpose(val) for key, val in weights.items()}
|
||||
weights = {"transformer": weights}
|
||||
mx.save_safetensors(
|
||||
str(path),
|
||||
dict(tree_flatten(weights)),
|
||||
metadata={
|
||||
"mflux_version": GeneratedImage.get_version(),
|
||||
"transformer_blocks": str(training_spec.lora_layers.transformer_blocks),
|
||||
"single_transformer_blocks": str(training_spec.lora_layers.single_transformer_blocks),
|
||||
},
|
||||
)
|
||||
0
src/mflux/dreambooth/optimization/__init__.py
Normal file
0
src/mflux/dreambooth/optimization/__init__.py
Normal file
Binary file not shown.
@ -0,0 +1,145 @@
|
||||
\documentclass{article}
|
||||
\usepackage{amsmath, amssymb}
|
||||
\setlength{\parindent}{0pt} % Remove indentation
|
||||
\setlength{\parskip}{1em}
|
||||
|
||||
\begin{document}
|
||||
|
||||
\title{DreamBooth loss}
|
||||
\author{Filip Strand}
|
||||
\date{2024-10-27}
|
||||
\maketitle
|
||||
|
||||
Given a reference image $X$ we can encode it using the \texttt{vae} to get $l_T$. Thus, $l_T$ represents a (clean) latent without any added noise. In this derivation, we use the notation from \texttt{MFLUX}\footnote{Which might be contrary to standard terminology} and let $t$ move from $t=0$ to $t=T$ when generating an image from noise. In other words, when generating an image from pure noise, we sample a $l_0 \sim \mathcal{N}(0,1)$ (a pure noise latent) and progressively move to $l_T$, the clean image latent, (which would subsequently be passed to the \texttt{vae} for decoding)
|
||||
|
||||
\[
|
||||
l_0 \rightarrow l_1 \rightarrow l_2 \rightarrow ... \rightarrow l_T
|
||||
\]
|
||||
|
||||
The time steps $t=0, t=1, t=2, \dots$ etc. have a direct correspondence with a set of $\sigma$ values $\sigma_0, \sigma_1, \dots$. While $t$ progresses in discrete integer steps $0, 1, 2, \dots$, the $\sigma$ values are floats that progressively go down in value from $1.0$ to $0.0$. As an example, it could look like the following:
|
||||
|
||||
\[
|
||||
t_0=0, t_1=1, t_2=2, \dots t=T
|
||||
\]
|
||||
\[
|
||||
\sigma_0=1.0, \sigma_1=0.96, \sigma_2=0.92, \dots \sigma_T = 0.0
|
||||
\]
|
||||
|
||||
The exact value of the $\sigma_t$ depend on how many time steps $T$ we use.
|
||||
|
||||
Given this, we can get the equivalent noised version of the clean $l_T$ at a time step $t$ by linearly interpolating between a pure noise array $\epsilon$ (same thing as $l_0$) and the clean latent $l_T$:
|
||||
|
||||
|
||||
\[
|
||||
l_t = (1-\sigma_t) \cdot l_T + \sigma_t \epsilon
|
||||
\]
|
||||
|
||||
As can be seen, when $t=0$, i.e $\sigma_0 = 1.0$ then this gives
|
||||
|
||||
\[
|
||||
l_0 = (1-\sigma_0) \cdot l_T + \sigma_0 \epsilon = \epsilon
|
||||
\]
|
||||
|
||||
and, conversely, when $t=T$, i.e $\sigma_T = 0.0$ we get
|
||||
|
||||
\[
|
||||
l_T = (1-\sigma_T) \cdot l_T + \sigma_T \epsilon = l_T
|
||||
\]
|
||||
|
||||
Suppose we create two subsequent latents $l_t$ and $l_{t+1}$ by the linear interpolation method above:
|
||||
|
||||
\[
|
||||
l_t = (1-\sigma_t) \cdot l_T + \sigma_t \epsilon
|
||||
\]
|
||||
\[
|
||||
l_{t+1} = (1-\sigma_{t+1}) \cdot l_T + \sigma_{t+1} \epsilon
|
||||
\]
|
||||
|
||||
As noted above, since the transformer part of the diffusion model (denoted $T_\theta$ here, with tunable parameters $\theta$) takes us from a latent $l_t$ to $l_{t+1}$, we can feed it $l_t$ and the explicit time step $t$ and get out $\epsilon'_t$, for the predicted noise array
|
||||
|
||||
\[
|
||||
\epsilon'_t = T_\theta(l_t,t)
|
||||
\]
|
||||
|
||||
Then, we get the predicted next latent $l_{t+1}'$ by taking a $dt_t$ size step
|
||||
|
||||
|
||||
\[
|
||||
l_{t+1}' = l_{t} + \epsilon'_t \cdot dt_t = l_{t} + \epsilon'_t \cdot (\sigma_{t+1} - \sigma_t)
|
||||
\]
|
||||
|
||||
where $dt_t = (\sigma_{t+1} - \sigma_t)$. The idea now is to form the loss by trying to get $l_{t+1}'$ close to $l_{t+1}$. Lets write this, somewhat informally, as
|
||||
|
||||
|
||||
\[
|
||||
\boxed{l_{t+1}' \approx l_{t+1}}
|
||||
\]
|
||||
|
||||
Then we can start substituting in our expressions for the LHS and RHS to simplify this expression. Starting by replacing $l_{t+1}'$ on the the LHS, we get:
|
||||
|
||||
|
||||
\[
|
||||
l_{t} + \epsilon'_t \cdot (\sigma_{t+1} - \sigma_t) \approx l_{t+1}
|
||||
\]
|
||||
|
||||
Now replacing $l_{t+1}$ on the RHS we get
|
||||
|
||||
\[
|
||||
l_{t} + \epsilon'_t \cdot (\sigma_{t+1} - \sigma_t) \approx (1-\sigma_{t+1}) \cdot l_T + \sigma_{t+1} \epsilon
|
||||
\]
|
||||
|
||||
Solving for our predicted $\epsilon'_t$ on the left hand side, we get
|
||||
|
||||
\[
|
||||
\epsilon'_t \approx \dfrac{(1-\sigma_{t+1}) \cdot l_T + \sigma_{t+1} \epsilon - l_{t}}{(\sigma_{t+1} - \sigma_t)}
|
||||
\]
|
||||
|
||||
Further substituting in the expression for $l_t$ that we got from linear interpolation we get
|
||||
|
||||
|
||||
\[
|
||||
\epsilon'_t \approx \dfrac{(1-\sigma_{t+1}) \cdot l_T + \sigma_{t+1} \epsilon - ((1-\sigma_t) \cdot l_T + \sigma_t \epsilon)}{(\sigma_{t+1} - \sigma_t)}
|
||||
\]
|
||||
|
||||
Simplifying the numerator by collecting $\epsilon$ and $l_T$ terms we get
|
||||
|
||||
\[
|
||||
\epsilon'_t \approx \dfrac{
|
||||
(\sigma_{t+1}-\sigma_t) \cdot \epsilon
|
||||
-
|
||||
(\sigma_{t+1}-\sigma_t) \cdot l_T
|
||||
}
|
||||
{(\sigma_{t+1} - \sigma_t)}
|
||||
\]
|
||||
|
||||
Finally, we see that the common factor of $(\sigma_{t+1}-\sigma_t)$ cancels and we get
|
||||
|
||||
\[
|
||||
\epsilon'_t \approx \epsilon - l_T
|
||||
\]
|
||||
|
||||
In order to turn this into an optimization objective we simply take LHS minus RHS, squared, and try to make sure this difference goes to zero. This means that our final loss can be written as
|
||||
|
||||
|
||||
\[
|
||||
\boxed{ \texttt{loss} := ||\epsilon'_t - (\epsilon - l_T)||^2}
|
||||
\]
|
||||
|
||||
or, equivalently,
|
||||
|
||||
\[
|
||||
\boxed{ \texttt{loss} := ||\epsilon'_t + l_T - \epsilon||^2}
|
||||
\]
|
||||
|
||||
or, by emphasizing the transformer part of the network,
|
||||
|
||||
\[
|
||||
\boxed{ \texttt{loss} := ||T_\theta(l_t,t) + l_T - \epsilon||^2}
|
||||
\]
|
||||
|
||||
In this form, the loss can be interpreted as having a the clean image latent $l_T$, then at an arbitrary time step $t$, we let the transformer predict the noise and the resulting quantity should be close to $\epsilon$.
|
||||
|
||||
|
||||
One nice thing to note here is that this loss holds for arbitrary $t$ values within our range $[0, T)$. This property is computationally very efficient since we don't need to compute the full chain of diffusion steps during training.
|
||||
|
||||
\end{document}
|
||||
65
src/mflux/dreambooth/optimization/dreambooth_loss.py
Normal file
65
src/mflux/dreambooth/optimization/dreambooth_loss.py
Normal file
@ -0,0 +1,65 @@
|
||||
import random
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
from mflux import Config, Flux1
|
||||
from mflux.config.runtime_config import RuntimeConfig
|
||||
from mflux.dreambooth.dataset.batch import Batch
|
||||
from mflux.dreambooth.dataset.dataset import Example
|
||||
from mflux.latent_creator.latent_creator import LatentCreator
|
||||
|
||||
|
||||
class DreamBoothLoss:
|
||||
@staticmethod
|
||||
def compute_loss(flux: Flux1, config: RuntimeConfig, batch: Batch) -> mx.float16:
|
||||
losses = [
|
||||
DreamBoothLoss._single_example_loss(flux, config, example, batch.rng)
|
||||
for example in batch.examples
|
||||
] # fmt: off
|
||||
return mx.mean(mx.array(losses))
|
||||
|
||||
@staticmethod
|
||||
def _single_example_loss(flux: Flux1, config: RuntimeConfig, example: Example, rng: random.Random) -> mx.float16:
|
||||
# Must be a better way to handle the randomness than this, but we already
|
||||
# save/restore the random state via the iterator so this is a continent shortcut.
|
||||
time_seed = rng.randint(0, 2**32 - 1)
|
||||
noise_seed = rng.randint(0, 2**32 - 1)
|
||||
|
||||
# Draw a random timestep t from [0, num_inference_steps]
|
||||
t = int(
|
||||
mx.random.randint(
|
||||
low=0,
|
||||
high=config.num_inference_steps,
|
||||
shape=[],
|
||||
key=mx.random.key(time_seed)
|
||||
)
|
||||
) # fmt: off
|
||||
|
||||
# Get the clean image latent
|
||||
clean_image = example.clean_latents
|
||||
|
||||
# Generate pure noise
|
||||
pure_noise = mx.random.normal(
|
||||
shape=clean_image.shape,
|
||||
dtype=Config.precision,
|
||||
key=mx.random.key(noise_seed)
|
||||
) # fmt: off
|
||||
|
||||
# By linear interpolation between the clean image and pure noise, construct a latent at time t
|
||||
latents_t = LatentCreator.add_noise_by_interpolation(
|
||||
clean=clean_image,
|
||||
noise=pure_noise,
|
||||
sigma=config.sigmas[t]
|
||||
) # fmt: off
|
||||
|
||||
# Predict the noise from timestep t
|
||||
predicted_noise = flux.transformer.predict(
|
||||
t=t,
|
||||
prompt_embeds=example.prompt_embeds,
|
||||
pooled_prompt_embeds=example.pooled_prompt_embeds,
|
||||
hidden_states=latents_t,
|
||||
config=config,
|
||||
)
|
||||
|
||||
# Construct the loss (derivation in src/mflux/dreambooth/optimization/_loss_derivation)
|
||||
return (clean_image + predicted_noise - pure_noise).square().mean()
|
||||
58
src/mflux/dreambooth/optimization/optimizer.py
Normal file
58
src/mflux/dreambooth/optimization/optimizer.py
Normal file
@ -0,0 +1,58 @@
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.optimizers
|
||||
import mlx.optimizers as optim
|
||||
from mlx.utils import tree_flatten, tree_unflatten
|
||||
|
||||
from mflux.dreambooth.state.training_spec import TrainingSpec
|
||||
from mflux.dreambooth.state.zip_util import ZipUtil
|
||||
|
||||
|
||||
class Optimizers(Enum):
|
||||
ADAM = ("Adam", optim.Adam)
|
||||
ADAMW = ("AdamW", optim.AdamW)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
alias: str,
|
||||
optimizer: mlx.optimizers.Optimizer,
|
||||
):
|
||||
self.alias = alias
|
||||
self.optimizer = optimizer
|
||||
|
||||
@staticmethod
|
||||
def from_alias(alias: str) -> "mlx.optimizers.Optimizer":
|
||||
try:
|
||||
for opt in Optimizers:
|
||||
if opt.alias == alias:
|
||||
return opt.optimizer
|
||||
except KeyError:
|
||||
raise ValueError(f"'{alias}' is not a valid optimization")
|
||||
|
||||
|
||||
class Optimizer:
|
||||
def __init__(self, optimizer: mlx.optimizers.Optimizer):
|
||||
self.optimizer = optimizer
|
||||
|
||||
def save(self, path: Path) -> None:
|
||||
state = tree_flatten(self.optimizer.state)
|
||||
mx.save_safetensors(str(path), dict(state))
|
||||
|
||||
@staticmethod
|
||||
def from_spec(training_spec: TrainingSpec) -> "Optimizer":
|
||||
opt = Optimizers.from_alias(training_spec.optimizer.name)
|
||||
# noinspection PyCallingNonCallable
|
||||
opt = opt(learning_rate=training_spec.optimizer.learning_rate)
|
||||
|
||||
# Load from state if present in the spec
|
||||
if training_spec.optimizer.state_path is not None:
|
||||
state = ZipUtil.unzip(
|
||||
zip_path=training_spec.checkpoint_path,
|
||||
filename=training_spec.optimizer.state_path,
|
||||
loader=lambda x: tree_unflatten(list(mx.load(x).items())),
|
||||
)
|
||||
opt.state = state
|
||||
|
||||
return Optimizer(opt)
|
||||
0
src/mflux/dreambooth/state/__init__.py
Normal file
0
src/mflux/dreambooth/state/__init__.py
Normal file
224
src/mflux/dreambooth/state/training_spec.py
Normal file
224
src/mflux/dreambooth/state/training_spec.py
Normal file
@ -0,0 +1,224 @@
|
||||
import datetime
|
||||
import json
|
||||
import os
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
from mflux.dreambooth.state.zip_util import ZipUtil
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExampleSpec:
|
||||
image: Path
|
||||
prompt: str
|
||||
|
||||
@classmethod
|
||||
def create(cls, param: dict[str, str], absolute_or_relative_path: str, base_path: Path) -> "ExampleSpec":
|
||||
image_path = Path(param["image"])
|
||||
|
||||
# If the path is not absolute, resolve it relative to the base path
|
||||
if not Path(absolute_or_relative_path).is_absolute():
|
||||
image_path = Path(base_path).parent / absolute_or_relative_path / image_path
|
||||
else:
|
||||
image_path = Path(absolute_or_relative_path) / image_path
|
||||
|
||||
return cls(
|
||||
image=image_path,
|
||||
prompt=param["prompt"],
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainingLoopSpec:
|
||||
num_epochs: int
|
||||
batch_size: int
|
||||
iterator_state_path: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class OptimizerSpec:
|
||||
name: str
|
||||
learning_rate: float
|
||||
state_path: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class SaveSpec:
|
||||
checkpoint_frequency: int
|
||||
output_path: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class StatisticsSpec:
|
||||
state_path: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class BlockRange:
|
||||
start: int
|
||||
end: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class TransformerBlocks:
|
||||
block_range: BlockRange
|
||||
layer_types: List[str]
|
||||
lora_rank: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class SingleTransformerBlocks:
|
||||
block_range: BlockRange
|
||||
layer_types: List[str]
|
||||
lora_rank: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class LoraLayersSpec:
|
||||
transformer_blocks: TransformerBlocks
|
||||
single_transformer_blocks: SingleTransformerBlocks
|
||||
state_path: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class InstrumentationSpec:
|
||||
plot_frequency: int
|
||||
generate_image_frequency: int
|
||||
validation_prompt: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainingSpec:
|
||||
model: str
|
||||
seed: int
|
||||
steps: int
|
||||
guidance: float
|
||||
quantize: int
|
||||
width: int
|
||||
height: int
|
||||
training_loop: TrainingLoopSpec
|
||||
optimizer: OptimizerSpec
|
||||
saver: SaveSpec
|
||||
instrumentation: InstrumentationSpec
|
||||
lora_layers: LoraLayersSpec
|
||||
statistics: StatisticsSpec
|
||||
examples: List[ExampleSpec]
|
||||
config_path: str | None = None
|
||||
checkpoint_path: str | None = None
|
||||
|
||||
@staticmethod
|
||||
def resolve(config_path: str | None, checkpoint_path: str | None) -> "TrainingSpec":
|
||||
if not config_path and not checkpoint_path:
|
||||
raise ValueError("Either 'config_path' or 'checkpoint_path' must be provided.")
|
||||
|
||||
if checkpoint_path:
|
||||
return TrainingSpec._from_checkpoint(checkpoint_path, new_folder=False)
|
||||
|
||||
return TrainingSpec._from_config(config_path, new_folder=True)
|
||||
|
||||
@staticmethod
|
||||
def _from_config(path: str, new_folder: bool = True) -> "TrainingSpec":
|
||||
with open(Path(path), "r") as f:
|
||||
data = json.load(f)
|
||||
|
||||
return TrainingSpec.from_conf(data, path, new_folder)
|
||||
|
||||
@staticmethod
|
||||
def from_conf(config: dict, config_path: str | None, new_folder: bool = True) -> "TrainingSpec":
|
||||
absolute_config_path = None if config_path is None else Path(config_path).absolute()
|
||||
absolute_or_relative_path = config["examples"]["path"]
|
||||
examples = [ExampleSpec.create(ex, absolute_or_relative_path, absolute_config_path) for ex in config["examples"]["images"]] # fmt: off
|
||||
training_loop = TrainingLoopSpec(**config["training_loop"])
|
||||
optimizer = OptimizerSpec(**config["optimizer"])
|
||||
saver = SaveSpec(
|
||||
checkpoint_frequency=config["save"]["checkpoint_frequency"],
|
||||
output_path=TrainingSpec._resolve_output_path(config["save"]["output_path"], new_folder),
|
||||
)
|
||||
instrumentation = (
|
||||
None if config.get("instrumentation", None) is None else InstrumentationSpec(**config["instrumentation"])
|
||||
)
|
||||
statistics = (
|
||||
StatisticsSpec() if config.get("statistics", None) is None else StatisticsSpec(**config["statistics"])
|
||||
)
|
||||
|
||||
# Parse LoraLayers structure
|
||||
transformer_blocks = config["lora_layers"].get("transformer_blocks", None)
|
||||
if transformer_blocks:
|
||||
transformer_blocks = TransformerBlocks(
|
||||
block_range=BlockRange(**transformer_blocks["block_range"]),
|
||||
layer_types=transformer_blocks["layer_types"],
|
||||
lora_rank=transformer_blocks["lora_rank"],
|
||||
)
|
||||
|
||||
single_transformer_blocks = config["lora_layers"].get("single_transformer_blocks", None)
|
||||
if single_transformer_blocks:
|
||||
single_transformer_blocks = SingleTransformerBlocks(
|
||||
block_range=BlockRange(**single_transformer_blocks["block_range"]),
|
||||
layer_types=single_transformer_blocks["layer_types"],
|
||||
lora_rank=single_transformer_blocks["lora_rank"],
|
||||
)
|
||||
|
||||
lora_layers = LoraLayersSpec(
|
||||
transformer_blocks=transformer_blocks, single_transformer_blocks=single_transformer_blocks
|
||||
)
|
||||
|
||||
# Create the training specification
|
||||
return TrainingSpec(
|
||||
model=config["model"],
|
||||
seed=config["seed"],
|
||||
steps=config["steps"],
|
||||
guidance=config["guidance"],
|
||||
quantize=config["quantize"],
|
||||
width=config["width"],
|
||||
height=config["height"],
|
||||
training_loop=training_loop,
|
||||
optimizer=optimizer,
|
||||
saver=saver,
|
||||
instrumentation=instrumentation,
|
||||
lora_layers=lora_layers,
|
||||
statistics=statistics,
|
||||
examples=examples,
|
||||
config_path=None if absolute_config_path is None else str(absolute_config_path),
|
||||
)
|
||||
|
||||
def to_json(self) -> str:
|
||||
spec_dict = asdict(self)
|
||||
serialized_dict = TrainingSpec._custom_serializer(spec_dict)
|
||||
return json.dumps(serialized_dict, indent=4)
|
||||
|
||||
@staticmethod
|
||||
def _custom_serializer(obj):
|
||||
if isinstance(obj, Path):
|
||||
return str(obj)
|
||||
elif isinstance(obj, datetime.datetime):
|
||||
return obj.strftime("%Y-%m-%d %H:%M:%S")
|
||||
elif isinstance(obj, list):
|
||||
return [TrainingSpec._custom_serializer(item) for item in obj]
|
||||
elif isinstance(obj, dict):
|
||||
return {key: TrainingSpec._custom_serializer(value) for key, value in obj.items()}
|
||||
return obj
|
||||
|
||||
@staticmethod
|
||||
def _resolve_output_path(path: str, new_folder: bool) -> str:
|
||||
requested_path = os.path.expanduser(path)
|
||||
now_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
if new_folder and os.path.exists(requested_path):
|
||||
requested_path = f"{requested_path}_{now_str}"
|
||||
|
||||
os.makedirs(requested_path, exist_ok=True)
|
||||
return requested_path
|
||||
|
||||
@staticmethod
|
||||
def _from_checkpoint(path: str, new_folder: bool) -> "TrainingSpec":
|
||||
checkpoint = ZipUtil.unzip(path, "checkpoint.json", lambda x: json.load(open(x, "r")))
|
||||
config = ZipUtil.unzip(path, checkpoint["files"]["config"], lambda x: json.load(open(x, "r")))
|
||||
spec = TrainingSpec.from_conf(config, None, new_folder)
|
||||
spec.optimizer.state_path = checkpoint["files"]["optimizer"]
|
||||
spec.lora_layers.state_path = checkpoint["files"]["lora_adapter"]
|
||||
spec.training_loop.iterator_state_path = checkpoint["files"]["iterator"]
|
||||
spec.statistics.state_path = checkpoint["files"]["loss"]
|
||||
spec.checkpoint_path = path
|
||||
spec.config_path = None
|
||||
return spec
|
||||
162
src/mflux/dreambooth/state/training_state.py
Normal file
162
src/mflux/dreambooth/state/training_state.py
Normal file
@ -0,0 +1,162 @@
|
||||
import datetime
|
||||
import json
|
||||
import tempfile
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
|
||||
from mflux.dreambooth.dataset.iterator import Iterator
|
||||
from mflux.dreambooth.lora_layers.lora_layers import LoRALayers
|
||||
from mflux.dreambooth.optimization.optimizer import Optimizer
|
||||
from mflux.dreambooth.state.training_spec import TrainingSpec
|
||||
from mflux.dreambooth.state.zip_util import ZipUtil
|
||||
from mflux.dreambooth.statistics.statistics import Statistics
|
||||
|
||||
DREAMBOOTH_PATH_CHECKPOINTS = "_checkpoints"
|
||||
DREAMBOOTH_FILE_NAME_CHECKPOINT = "checkpoint"
|
||||
|
||||
DREAMBOOTH_FILE_NAME_LORA_ADAPTER = "adapter"
|
||||
DREAMBOOTH_FILE_NAME_OPTIMIZER = "optimizer"
|
||||
DREAMBOOTH_FILE_NAME_ITERATOR = "iterator"
|
||||
DREAMBOOTH_FILE_NAME_LOSS_FILE = "loss"
|
||||
DREAMBOOTH_FILE_NAME_CONFIG_FILE = "config"
|
||||
|
||||
DREAMBOOTH_PATH_VALIDATION_IMAGES = "_validation/images/"
|
||||
DREAMBOOTH_FILE_NAME_VALIDATION_IMAGE = "validation_image"
|
||||
|
||||
DREAMBOOTH_PATH_VALIDATION_PLOT = "_validation/plots/"
|
||||
DREAMBOOTH_FILE_NAME_VALIDATION_LOSS = "loss"
|
||||
|
||||
|
||||
class TrainingState:
|
||||
def __init__(
|
||||
self,
|
||||
iterator: Iterator,
|
||||
lora_layers: LoRALayers,
|
||||
optimizer: Optimizer,
|
||||
statistics: Statistics,
|
||||
):
|
||||
self.iterator = iterator
|
||||
self.optimizer = optimizer
|
||||
self.lora_layers = lora_layers
|
||||
self.statistics = statistics
|
||||
|
||||
def save(self, training_spec: TrainingSpec) -> None:
|
||||
# Create a temporary directory to store files before zipping
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
# Save files to temporary directory
|
||||
optimizer_path = Path(temp_dir) / f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_OPTIMIZER}.safetensors" # fmt:off
|
||||
lora_path = Path(temp_dir) / f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_LORA_ADAPTER}.safetensors" # fmt:off
|
||||
iterator_path = Path(temp_dir) / f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_ITERATOR}.json"
|
||||
loss_path = Path(temp_dir) / f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_LOSS_FILE}.json"
|
||||
config_path = Path(temp_dir) / f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_CONFIG_FILE}.json"
|
||||
checkpoint_path = Path(temp_dir) / f"{DREAMBOOTH_FILE_NAME_CHECKPOINT}.json" # fmt: off
|
||||
paths = [optimizer_path, lora_path, iterator_path, loss_path, config_path, checkpoint_path]
|
||||
|
||||
# Save individual files to temporary directory
|
||||
self.optimizer.save(optimizer_path)
|
||||
self.lora_layers.save(lora_path, training_spec)
|
||||
self.iterator.save(iterator_path)
|
||||
self.statistics.save(loss_path)
|
||||
self._save_train_config(config_path, training_spec) # For completeness, we also save the (initial) config
|
||||
|
||||
# Create checkpoint data
|
||||
checkpoint_data = self._create_checkpoint_data(training_spec, self.iterator.start_date_time)
|
||||
with open(checkpoint_path, "w") as json_file:
|
||||
json.dump(checkpoint_data, json_file, indent=4)
|
||||
|
||||
# Create zip file
|
||||
output_path = Path(training_spec.saver.output_path) / DREAMBOOTH_PATH_CHECKPOINTS
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
zip_path = output_path / f"{self.iterator.num_iterations:07d}_checkpoint.zip"
|
||||
|
||||
with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
|
||||
for file_path in paths:
|
||||
if file_path.exists():
|
||||
zipf.write(file_path, file_path.name)
|
||||
|
||||
def _create_checkpoint_data(self, training_spec: TrainingSpec, start_date_time: datetime) -> dict:
|
||||
now = datetime.datetime.now()
|
||||
return {
|
||||
"metadata": {
|
||||
"start": start_date_time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"end": now.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"duration": self._format_duration(start_date_time, now),
|
||||
"number_of_training_examples": self.iterator.dataset.size(),
|
||||
},
|
||||
"files": {
|
||||
"config": f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_CONFIG_FILE}.json",
|
||||
"optimizer": f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_OPTIMIZER}.safetensors",
|
||||
"lora_adapter": f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_LORA_ADAPTER}.safetensors",
|
||||
"iterator": f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_ITERATOR}.json",
|
||||
"loss": f"{self.iterator.num_iterations:07d}_{DREAMBOOTH_FILE_NAME_LOSS_FILE}.json",
|
||||
},
|
||||
}
|
||||
|
||||
def should_save(self, training_spec: TrainingSpec) -> bool:
|
||||
return self.iterator.num_iterations % training_spec.saver.checkpoint_frequency == 0
|
||||
|
||||
def should_plot_loss(self, training_spec: TrainingSpec) -> bool:
|
||||
if training_spec.instrumentation is None:
|
||||
return False
|
||||
return self.iterator.num_iterations % training_spec.instrumentation.plot_frequency == 0
|
||||
|
||||
def should_generate_image(self, training_spec: TrainingSpec) -> bool:
|
||||
if training_spec.instrumentation is None:
|
||||
return False
|
||||
return self.iterator.num_iterations % training_spec.instrumentation.generate_image_frequency == 0
|
||||
|
||||
def get_current_validation_image_path(self, training_spec: TrainingSpec) -> Path:
|
||||
output_path = Path(training_spec.saver.output_path) / DREAMBOOTH_PATH_VALIDATION_IMAGES
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
path = output_path / Path(f"{self.iterator.num_iterations :07d}_{DREAMBOOTH_FILE_NAME_VALIDATION_IMAGE}.png")
|
||||
return path
|
||||
|
||||
def get_current_loss_plot_path(self, training_spec: TrainingSpec) -> Path:
|
||||
output_path = Path(training_spec.saver.output_path) / DREAMBOOTH_PATH_VALIDATION_PLOT
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
path = output_path / Path(f"{self.iterator.num_iterations :07d}_{DREAMBOOTH_FILE_NAME_VALIDATION_LOSS}.pdf")
|
||||
return path
|
||||
|
||||
@staticmethod
|
||||
def _get_parent_path(path: str | None) -> str | None:
|
||||
return None if path is None else str(Path(path).resolve())
|
||||
|
||||
@staticmethod
|
||||
def _format_duration(start: datetime, end: datetime) -> str:
|
||||
# Calculate the duration
|
||||
duration = end - start
|
||||
total_seconds = int(duration.total_seconds())
|
||||
|
||||
# Calculate hours, minutes, and seconds
|
||||
hours, remainder = divmod(total_seconds, 3600)
|
||||
minutes, seconds = divmod(remainder, 60)
|
||||
|
||||
# Build the string
|
||||
parts = []
|
||||
if hours > 0:
|
||||
parts.append(f"{hours} hour{'s' if hours > 1 else ''}")
|
||||
if minutes > 0:
|
||||
parts.append(f"{minutes} minute{'s' if minutes > 1 else ''}")
|
||||
if seconds > 0 or not parts: # Include seconds even if zero if no other parts exist
|
||||
parts.append(f"{seconds} second{'s' if seconds > 1 else ''}")
|
||||
|
||||
return " and ".join(parts)
|
||||
|
||||
@staticmethod
|
||||
def _save_train_config(path: Path, training_spec: TrainingSpec) -> None:
|
||||
if training_spec.config_path is not None:
|
||||
with open(Path(training_spec.config_path), "r") as f:
|
||||
data = json.load(f)
|
||||
# Since the zip file can be moved to an arbitrary location, we do a slight
|
||||
# modification and override the image location with the absolute path of the image files
|
||||
data["examples"]["path"] = str(Path(training_spec.config_path).parent / data["examples"]["path"])
|
||||
else:
|
||||
checkpoint = ZipUtil.unzip(
|
||||
training_spec.checkpoint_path, "checkpoint.json", lambda x: json.load(open(x, "r"))
|
||||
)
|
||||
data = ZipUtil.unzip(
|
||||
training_spec.checkpoint_path, checkpoint["files"]["config"], lambda x: json.load(open(x, "r"))
|
||||
)
|
||||
|
||||
with open(path, "w", encoding="utf-8") as file:
|
||||
json.dump(data, file, indent=4)
|
||||
56
src/mflux/dreambooth/state/zip_util.py
Normal file
56
src/mflux/dreambooth/state/zip_util.py
Normal file
@ -0,0 +1,56 @@
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from zipfile import ZipFile
|
||||
|
||||
|
||||
class ZipUtil:
|
||||
@staticmethod
|
||||
def unzip(zip_path: str | Path, filename: str, loader: callable):
|
||||
zip_path = Path(zip_path)
|
||||
if not zip_path.exists():
|
||||
raise FileNotFoundError(f"ZIP file not found at: {zip_path}")
|
||||
|
||||
if not callable(loader):
|
||||
raise ValueError("The file_loader must be a callable (e.g., a function or lambda).")
|
||||
|
||||
with ZipFile(zip_path, "r") as zipf:
|
||||
# Normalize files names in the archive by stripping folders
|
||||
namelist = {os.path.basename(name): name for name in zipf.namelist()}
|
||||
|
||||
if filename not in namelist:
|
||||
raise FileNotFoundError(f"File '{filename}' not found in the ZIP archive.")
|
||||
|
||||
# Use the full path to extract the correct file
|
||||
archive_filename = namelist[filename]
|
||||
file_data = zipf.read(archive_filename)
|
||||
|
||||
# Create a temporary file to store the extracted content
|
||||
temp_file = tempfile.NamedTemporaryFile(suffix=f".{filename.split('.')[-1]}", delete=False)
|
||||
try:
|
||||
temp_file.write(file_data)
|
||||
temp_file.flush()
|
||||
temp_file.close()
|
||||
|
||||
# Call the file_loader with the temporary file path
|
||||
return loader(temp_file.name)
|
||||
finally:
|
||||
# Cleanup the temporary file
|
||||
if os.path.exists(temp_file.name):
|
||||
os.remove(temp_file.name)
|
||||
|
||||
@staticmethod
|
||||
def extract_all(zip_path: str | Path, output_dir: str | Path):
|
||||
zip_path = Path(zip_path)
|
||||
output_dir = Path(output_dir)
|
||||
|
||||
if not zip_path.exists():
|
||||
raise FileNotFoundError(f"ZIP file not found at: {zip_path}")
|
||||
|
||||
if not output_dir.exists():
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with ZipFile(zip_path, "r") as zipf:
|
||||
zipf.extractall(output_dir)
|
||||
|
||||
print(f"All files have been extracted to: {output_dir}")
|
||||
0
src/mflux/dreambooth/statistics/__init__.py
Normal file
0
src/mflux/dreambooth/statistics/__init__.py
Normal file
53
src/mflux/dreambooth/statistics/plotter.py
Normal file
53
src/mflux/dreambooth/statistics/plotter.py
Normal file
@ -0,0 +1,53 @@
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
from mflux.dreambooth.state.training_spec import TrainingSpec
|
||||
from mflux.dreambooth.state.training_state import TrainingState
|
||||
|
||||
|
||||
class Plotter:
|
||||
@staticmethod
|
||||
def update_loss_plot(training_spec: TrainingSpec, training_state: TrainingState) -> None:
|
||||
plt.style.use("bmh")
|
||||
|
||||
# Create figure with 16:9 aspect ratio
|
||||
plt.figure(figsize=(16, 9), dpi=300)
|
||||
|
||||
# Plot both lines and points
|
||||
stats = training_state.statistics
|
||||
plt.plot(stats.steps, stats.losses, "b-", linewidth=2, label="Validation Loss", zorder=1) # Line
|
||||
plt.plot(stats.steps, stats.losses, "bo", markersize=6, label="_nolegend_", zorder=2) # Points
|
||||
|
||||
# Customize the plot
|
||||
plt.title("Validation Loss Over Time", fontsize=16, pad=20)
|
||||
plt.xlabel("Steps", fontsize=12)
|
||||
plt.ylabel("Loss", fontsize=12)
|
||||
|
||||
# Set integer grid
|
||||
plt.grid(True, linestyle="--", alpha=0.7)
|
||||
ax = plt.gca()
|
||||
ax.xaxis.set_major_locator(plt.MaxNLocator(integer=True))
|
||||
|
||||
# Dynamic x-axis limit with 20% padding
|
||||
max_x = max(stats.steps)
|
||||
padding = max_x * 0.2
|
||||
plt.xlim(0, max_x + padding)
|
||||
|
||||
# Dynamic y-axis limit with 20% padding
|
||||
max_y = float(max(stats.losses))
|
||||
padding = max_y * 0.2
|
||||
plt.ylim(0, max_y + padding)
|
||||
|
||||
plt.legend(fontsize=12)
|
||||
|
||||
# Add margins for better visibility
|
||||
plt.margins(x=0.02)
|
||||
|
||||
# Tight layout to prevent label cutoff
|
||||
plt.tight_layout()
|
||||
|
||||
# Save to desktop with high PPI
|
||||
path = training_state.get_current_loss_plot_path(training_spec)
|
||||
plt.savefig(path, format="pdf", dpi=300, bbox_inches="tight")
|
||||
|
||||
# Close the figure to free memory
|
||||
plt.close()
|
||||
55
src/mflux/dreambooth/statistics/statistics.py
Normal file
55
src/mflux/dreambooth/statistics/statistics.py
Normal file
@ -0,0 +1,55 @@
|
||||
import datetime
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from mflux.dreambooth.state.zip_util import ZipUtil
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mflux.dreambooth.state.training_spec import TrainingSpec
|
||||
|
||||
|
||||
class Statistics:
|
||||
def __init__(self):
|
||||
self.steps: list[int] = []
|
||||
self.losses: list[float] = []
|
||||
self.times: list[datetime.datetime] = []
|
||||
|
||||
@staticmethod
|
||||
def from_spec(training_spec: "TrainingSpec") -> "Statistics":
|
||||
if training_spec.statistics is None:
|
||||
return Statistics()
|
||||
|
||||
if training_spec.statistics.state_path is None:
|
||||
return Statistics()
|
||||
|
||||
stats = Statistics()
|
||||
data = ZipUtil.unzip(
|
||||
zip_path=training_spec.checkpoint_path,
|
||||
filename=training_spec.statistics.state_path,
|
||||
loader=lambda x: json.load(open(x, "r")),
|
||||
)
|
||||
for entry in data:
|
||||
stats.steps.append(entry["step"])
|
||||
stats.losses.append(entry["loss"])
|
||||
stats.times.append(datetime.datetime.strptime(entry["time"], "%Y-%m-%d %H:%M:%S"))
|
||||
|
||||
return stats
|
||||
|
||||
def append_values(self, step: int, loss: float) -> None:
|
||||
self.steps.append(step)
|
||||
self.losses.append(loss)
|
||||
self.times.append(datetime.datetime.now())
|
||||
|
||||
def save(self, path: Path) -> None:
|
||||
loss_entries = [
|
||||
{
|
||||
"step": step,
|
||||
"loss": float(loss),
|
||||
"time": time.strftime("%Y-%m-%d %H:%M:%S")
|
||||
}
|
||||
for step, loss, time in zip(self.steps, self.losses, self.times)
|
||||
] # fmt: off
|
||||
|
||||
with open(path, "w", encoding="utf-8") as file:
|
||||
json.dump(loss_entries, file, indent=4)
|
||||
@ -16,3 +16,7 @@ class MFluxUserException(MFluxException):
|
||||
|
||||
class StopImageGenerationException(MFluxUserException):
|
||||
"""user has requested to stop a image generation in progress."""
|
||||
|
||||
|
||||
class StopTrainingException(MFluxUserException):
|
||||
"""user has requested to stop the training process in progress."""
|
||||
|
||||
@ -22,9 +22,11 @@ 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
|
||||
from mflux.weights.weight_handler_lora import WeightHandlerLoRA
|
||||
from mflux.weights.weight_util import WeightUtil
|
||||
|
||||
|
||||
class Flux1:
|
||||
class Flux1(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
model_config: ModelConfig,
|
||||
@ -33,6 +35,7 @@ class Flux1:
|
||||
lora_paths: list[str] | None = None,
|
||||
lora_scales: list[float] | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.lora_paths = lora_paths
|
||||
self.lora_scales = lora_scales
|
||||
self.model_config = model_config
|
||||
@ -48,32 +51,20 @@ class Flux1:
|
||||
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
|
||||
) # fmt: off
|
||||
# Set the weights and quantize the model
|
||||
weights = WeightHandler.load_regular_weights(repo_id=model_config.model_name, local_path=local_path)
|
||||
self.bits = WeightUtil.set_weights_and_quantize(
|
||||
quantize_arg=quantize,
|
||||
weights=weights,
|
||||
vae=self.vae,
|
||||
transformer=self.transformer,
|
||||
t5_text_encoder=self.t5_text_encoder,
|
||||
clip_text_encoder=self.clip_text_encoder,
|
||||
)
|
||||
|
||||
# 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
|
||||
# fmt: off
|
||||
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)
|
||||
# fmt: on
|
||||
|
||||
# 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)
|
||||
# Set LoRA weights
|
||||
lora_weights = WeightHandlerLoRA.load_lora_weights(transformer=self.transformer, lora_files=lora_paths, lora_scales=lora_scales) # fmt:off
|
||||
WeightHandlerLoRA.set_lora_weights(transformer=self.transformer, loras=lora_weights)
|
||||
|
||||
def generate_image(
|
||||
self,
|
||||
@ -100,8 +91,8 @@ class Flux1:
|
||||
# 2. Embed 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)
|
||||
prompt_embeds = self.t5_text_encoder(t5_tokens)
|
||||
pooled_prompt_embeds = self.clip_text_encoder(clip_tokens)
|
||||
|
||||
for gen_step, t in enumerate(time_steps, 1):
|
||||
try:
|
||||
@ -151,11 +142,11 @@ class Flux1:
|
||||
quantize=quantize,
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
def freeze(self, **kwargs):
|
||||
self.vae.freeze()
|
||||
self.transformer.freeze()
|
||||
self.t5_text_encoder.freeze()
|
||||
self.clip_text_encoder.freeze()
|
||||
|
||||
@ -10,7 +10,7 @@ class CLIPEmbeddings(nn.Module):
|
||||
self.position_embedding = nn.Embedding(num_embeddings=TokenizerCLIP.MAX_TOKEN_LENGTH, dims=dims)
|
||||
self.token_embedding = nn.Embedding(num_embeddings=49408, dims=dims)
|
||||
|
||||
def forward(self, tokens: mx.array) -> mx.array:
|
||||
def __call__(self, tokens: mx.array) -> mx.array:
|
||||
seq_length = tokens.shape[-1]
|
||||
position_ids = mx.arange(0, seq_length).reshape(1, seq_length)
|
||||
inputs_embeds = self.token_embedding(tokens)
|
||||
|
||||
@ -9,6 +9,6 @@ class CLIPEncoder(nn.Module):
|
||||
super().__init__()
|
||||
self.text_model = CLIPTextModel(dims=768, num_encoder_layers=12)
|
||||
|
||||
def forward(self, tokens: mx.array) -> mx.array:
|
||||
pooled_output = self.text_model.forward(tokens)
|
||||
def __call__(self, tokens: mx.array) -> mx.array:
|
||||
pooled_output = self.text_model(tokens)
|
||||
return pooled_output
|
||||
|
||||
@ -13,13 +13,13 @@ class CLIPEncoderLayer(nn.Module):
|
||||
self.mlp = CLIPMLP()
|
||||
self.layer_norm2 = nn.LayerNorm(dims=768)
|
||||
|
||||
def forward(self, hidden_states: mx.array, causal_attention_mask: mx.array) -> mx.array:
|
||||
def __call__(self, hidden_states: mx.array, causal_attention_mask: mx.array) -> mx.array:
|
||||
residual = hidden_states
|
||||
hidden_states = self.layer_norm1(hidden_states)
|
||||
hidden_states = self.self_attn.forward(hidden_states, causal_attention_mask)
|
||||
hidden_states = self.self_attn(hidden_states, causal_attention_mask)
|
||||
hidden_states = residual + hidden_states
|
||||
residual = hidden_states
|
||||
hidden_states = self.layer_norm2(hidden_states)
|
||||
hidden_states = self.mlp.forward(hidden_states)
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
return hidden_states
|
||||
|
||||
@ -8,7 +8,7 @@ class CLIPMLP(nn.Module):
|
||||
self.fc1 = nn.Linear(input_dims=768, output_dims=3072)
|
||||
self.fc2 = nn.Linear(input_dims=3072, output_dims=768)
|
||||
|
||||
def forward(self, hidden_states: mx.array) -> mx.array:
|
||||
def __call__(self, hidden_states: mx.array) -> mx.array:
|
||||
hidden_states = self.fc1(hidden_states)
|
||||
hidden_states = CLIPMLP.quick_gelu(hidden_states)
|
||||
hidden_states = self.fc2(hidden_states)
|
||||
|
||||
@ -14,7 +14,7 @@ class CLIPSdpaAttention(nn.Module):
|
||||
self.v_proj = nn.Linear(input_dims=768, output_dims=768)
|
||||
self.out_proj = nn.Linear(input_dims=768, output_dims=768)
|
||||
|
||||
def forward(self, hidden_states: mx.array, causal_attention_mask: mx.array) -> mx.array:
|
||||
def __call__(self, hidden_states: mx.array, causal_attention_mask: mx.array) -> mx.array:
|
||||
query = self.q_proj(hidden_states)
|
||||
key = self.k_proj(hidden_states)
|
||||
value = self.v_proj(hidden_states)
|
||||
|
||||
@ -12,10 +12,10 @@ class CLIPTextModel(nn.Module):
|
||||
self.embeddings = CLIPEmbeddings(dims)
|
||||
self.final_layer_norm = nn.LayerNorm(dims=768)
|
||||
|
||||
def forward(self, tokens: mx.array) -> (mx.array, mx.array):
|
||||
hidden_states = self.embeddings.forward(tokens)
|
||||
def __call__(self, tokens: mx.array) -> (mx.array, mx.array):
|
||||
hidden_states = self.embeddings(tokens)
|
||||
causal_attention_mask = CLIPTextModel.create_causal_attention_mask(hidden_states.shape)
|
||||
encoder_outputs = self.encoder.forward(hidden_states, causal_attention_mask)
|
||||
encoder_outputs = self.encoder(hidden_states, causal_attention_mask)
|
||||
last_hidden_state = self.final_layer_norm(encoder_outputs)
|
||||
pooled_output = last_hidden_state[0, mx.argmax(tokens, axis=-1)]
|
||||
return pooled_output
|
||||
|
||||
@ -11,8 +11,8 @@ class EncoderCLIP(nn.Module):
|
||||
super().__init__()
|
||||
self.layers = [CLIPEncoderLayer(i) for i in range(num_encoder_layers)]
|
||||
|
||||
def forward(self, tokens: mx.array, causal_attention_mask: mx.array) -> mx.array:
|
||||
def __call__(self, tokens: mx.array, causal_attention_mask: mx.array) -> mx.array:
|
||||
hidden_states = tokens
|
||||
for encoder_layer in self.layers:
|
||||
hidden_states = encoder_layer.forward(hidden_states, causal_attention_mask)
|
||||
hidden_states = encoder_layer(hidden_states, causal_attention_mask)
|
||||
return hidden_states
|
||||
|
||||
@ -13,8 +13,8 @@ class T5Attention(nn.Module):
|
||||
self.SelfAttention = T5SelfAttention()
|
||||
self.layer_norm = T5LayerNorm()
|
||||
|
||||
def forward(self, hidden_states: mx.array) -> mx.array:
|
||||
normed_hidden_states = self.layer_norm.forward(hidden_states)
|
||||
attention_output = self.SelfAttention.forward(normed_hidden_states)
|
||||
def __call__(self, hidden_states: mx.array) -> mx.array:
|
||||
normed_hidden_states = self.layer_norm(hidden_states)
|
||||
attention_output = self.SelfAttention(normed_hidden_states)
|
||||
hidden_states = hidden_states + attention_output
|
||||
return hidden_states
|
||||
|
||||
@ -11,8 +11,8 @@ class T5Block(nn.Module):
|
||||
self.attention = T5Attention()
|
||||
self.ff = T5FeedForward()
|
||||
|
||||
def forward(self, hidden_states: mx.array) -> mx.array:
|
||||
hidden_states = self.attention.forward(hidden_states)
|
||||
hidden_states = self.ff.forward(hidden_states)
|
||||
def __call__(self, hidden_states: mx.array) -> mx.array:
|
||||
hidden_states = self.attention(hidden_states)
|
||||
hidden_states = self.ff(hidden_states)
|
||||
outputs = hidden_states
|
||||
return outputs
|
||||
|
||||
@ -11,7 +11,7 @@ class T5DenseReluDense(nn.Module):
|
||||
self.wi_1 = nn.Linear(4096, 10240, bias=False)
|
||||
self.wo = nn.Linear(10240, 4096, bias=False)
|
||||
|
||||
def forward(self, hidden_states: mx.array) -> mx.array:
|
||||
def __call__(self, hidden_states: mx.array) -> mx.array:
|
||||
hidden_gelu = T5DenseReluDense.new_gelu(self.wi_0(hidden_states))
|
||||
hidden_linear = self.wi_1(hidden_states)
|
||||
hidden_states = hidden_gelu * hidden_linear
|
||||
|
||||
@ -12,9 +12,9 @@ class T5Encoder(nn.Module):
|
||||
self.t5_blocks = [T5Block(i) for i in range(24)]
|
||||
self.final_layer_norm = T5LayerNorm()
|
||||
|
||||
def forward(self, tokens: mx.array):
|
||||
def __call__(self, tokens: mx.array):
|
||||
hidden_states = self.shared(tokens)
|
||||
for block in self.t5_blocks:
|
||||
hidden_states = block.forward(hidden_states)
|
||||
hidden_states = self.final_layer_norm.forward(hidden_states)
|
||||
hidden_states = block(hidden_states)
|
||||
hidden_states = self.final_layer_norm(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
@ -13,7 +13,7 @@ class T5FeedForward(nn.Module):
|
||||
self.layer_norm = T5LayerNorm()
|
||||
self.DenseReluDense = T5DenseReluDense()
|
||||
|
||||
def forward(self, hidden_states: mx.array) -> mx.array:
|
||||
forwarded_states = self.layer_norm.forward(hidden_states)
|
||||
forwarded_states = self.DenseReluDense.forward(forwarded_states)
|
||||
def __call__(self, hidden_states: mx.array) -> mx.array:
|
||||
forwarded_states = self.layer_norm(hidden_states)
|
||||
forwarded_states = self.DenseReluDense(forwarded_states)
|
||||
return hidden_states + forwarded_states
|
||||
|
||||
@ -8,7 +8,7 @@ class T5LayerNorm(nn.Module):
|
||||
self.weight = mx.ones((4096,))
|
||||
self.variance_epsilon = 1e-06
|
||||
|
||||
def forward(self, hidden_states: mx.array) -> mx.array:
|
||||
def __call__(self, hidden_states: mx.array) -> mx.array:
|
||||
variance = mx.mean(
|
||||
mx.power(hidden_states.astype(mx.float32), 2),
|
||||
axis=-1,
|
||||
|
||||
@ -13,7 +13,7 @@ class T5SelfAttention(nn.Module):
|
||||
self.relative_attention_bias = nn.Embedding(32, 64)
|
||||
self.o = nn.Linear(4096, 4096, bias=False)
|
||||
|
||||
def forward(self, hidden_states: mx.array) -> mx.array:
|
||||
def __call__(self, hidden_states: mx.array) -> mx.array:
|
||||
query_states = T5SelfAttention.shape(self.q(hidden_states))
|
||||
key_states = T5SelfAttention.shape(self.k(hidden_states))
|
||||
value_states = T5SelfAttention.shape(self.v(hidden_states))
|
||||
|
||||
@ -11,7 +11,7 @@ class AdaLayerNormContinuous(nn.Module):
|
||||
self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2)
|
||||
self.norm = nn.LayerNorm(dims=embedding_dim, eps=1e-6, affine=False)
|
||||
|
||||
def forward(self, x: mx.array, text_embeddings: mx.array) -> mx.array:
|
||||
def __call__(self, x: mx.array, text_embeddings: mx.array) -> mx.array:
|
||||
text_embeddings = self.linear(nn.silu(text_embeddings).astype(Config.precision))
|
||||
chunk_size = self.embedding_dim
|
||||
scale = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
|
||||
|
||||
@ -8,7 +8,7 @@ class AdaLayerNormZero(nn.Module):
|
||||
self.linear = nn.Linear(3072, 18432)
|
||||
self.norm = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
|
||||
|
||||
def forward(self, x: mx.array, text_embeddings: mx.array):
|
||||
def __call__(self, x: mx.array, text_embeddings: mx.array):
|
||||
text_embeddings = self.linear(nn.silu(text_embeddings))
|
||||
chunk_size = 18432 // 6
|
||||
shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
|
||||
|
||||
@ -8,7 +8,7 @@ class AdaLayerNormZeroSingle(nn.Module):
|
||||
self.linear = nn.Linear(3072, 3 * 3072)
|
||||
self.norm = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
|
||||
|
||||
def forward(self, x: mx.array, text_embeddings: mx.array):
|
||||
def __call__(self, x: mx.array, text_embeddings: mx.array):
|
||||
text_embeddings = self.linear(nn.silu(text_embeddings))
|
||||
chunk_size = 9216 // 3
|
||||
shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
|
||||
|
||||
@ -9,7 +9,7 @@ class EmbedND(nn.Module):
|
||||
self.theta = 10000
|
||||
self.axes_dim = [16, 56, 56]
|
||||
|
||||
def forward(self, ids: mx.array) -> mx.array:
|
||||
def __call__(self, ids: mx.array) -> mx.array:
|
||||
emb = mx.concatenate(
|
||||
[EmbedND.rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(3)],
|
||||
axis=-3,
|
||||
|
||||
@ -9,7 +9,7 @@ class FeedForward(nn.Module):
|
||||
self.linear2 = nn.Linear(12288, 3072)
|
||||
self.activation_function = activation_function
|
||||
|
||||
def forward(self, hidden_states: mx.array) -> mx.array:
|
||||
def __call__(self, hidden_states: mx.array) -> mx.array:
|
||||
hidden_states = self.linear1(hidden_states)
|
||||
hidden_states = self.activation_function(hidden_states)
|
||||
hidden_states = self.linear2(hidden_states)
|
||||
|
||||
@ -8,7 +8,7 @@ class GuidanceEmbedder(nn.Module):
|
||||
self.linear_1 = nn.Linear(256, 3072)
|
||||
self.linear_2 = nn.Linear(3072, 3072)
|
||||
|
||||
def forward(self, sample: mx.array) -> mx.array:
|
||||
def __call__(self, sample: mx.array) -> mx.array:
|
||||
sample = self.linear_1(sample)
|
||||
sample = nn.silu(sample)
|
||||
sample = self.linear_2(sample)
|
||||
|
||||
@ -22,7 +22,7 @@ class JointAttention(nn.Module):
|
||||
self.norm_added_q = nn.RMSNorm(128)
|
||||
self.norm_added_k = nn.RMSNorm(128)
|
||||
|
||||
def forward(
|
||||
def __call__(
|
||||
self,
|
||||
hidden_states: mx.array,
|
||||
encoder_hidden_states: mx.array,
|
||||
|
||||
@ -18,22 +18,20 @@ class JointTransformerBlock(nn.Module):
|
||||
self.ff_context = FeedForward(activation_function=nn.gelu_approx)
|
||||
self.norm2_context = nn.LayerNorm(dims=1536, eps=1e-6, affine=False)
|
||||
|
||||
def forward(
|
||||
def __call__(
|
||||
self,
|
||||
hidden_states: mx.array,
|
||||
encoder_hidden_states: mx.array,
|
||||
text_embeddings: mx.array,
|
||||
rotary_embeddings: mx.array,
|
||||
) -> (mx.array, mx.array):
|
||||
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1.forward(
|
||||
hidden_states, text_embeddings
|
||||
)
|
||||
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, text_embeddings)
|
||||
|
||||
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context.forward(
|
||||
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
|
||||
x=encoder_hidden_states, text_embeddings=text_embeddings
|
||||
)
|
||||
|
||||
attn_output, context_attn_output = self.attn.forward(
|
||||
attn_output, context_attn_output = self.attn(
|
||||
hidden_states=norm_hidden_states,
|
||||
encoder_hidden_states=norm_encoder_hidden_states,
|
||||
image_rotary_emb=rotary_embeddings,
|
||||
@ -43,7 +41,7 @@ class JointTransformerBlock(nn.Module):
|
||||
hidden_states = hidden_states + attn_output
|
||||
norm_hidden_states = self.norm2(hidden_states)
|
||||
norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
|
||||
ff_output = self.ff.forward(norm_hidden_states)
|
||||
ff_output = self.ff(norm_hidden_states)
|
||||
ff_output = mx.expand_dims(gate_mlp, axis=1) * ff_output
|
||||
hidden_states = hidden_states + ff_output
|
||||
|
||||
@ -51,6 +49,6 @@ class JointTransformerBlock(nn.Module):
|
||||
encoder_hidden_states = encoder_hidden_states + context_attn_output
|
||||
norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
|
||||
norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
|
||||
context_ff_output = self.ff_context.forward(norm_encoder_hidden_states)
|
||||
context_ff_output = self.ff_context(norm_encoder_hidden_states)
|
||||
encoder_hidden_states = encoder_hidden_states + mx.expand_dims(c_gate_mlp, axis=1) * context_ff_output
|
||||
return encoder_hidden_states, hidden_states
|
||||
|
||||
@ -15,7 +15,7 @@ class SingleBlockAttention(nn.Module):
|
||||
self.norm_q = nn.RMSNorm(128)
|
||||
self.norm_k = nn.RMSNorm(128)
|
||||
|
||||
def forward(self, hidden_states: mx.array, image_rotary_emb: mx.array) -> (mx.array, mx.array):
|
||||
def __call__(self, hidden_states: mx.array, image_rotary_emb: mx.array) -> (mx.array, mx.array):
|
||||
query = self.to_q(hidden_states)
|
||||
key = self.to_k(hidden_states)
|
||||
value = self.to_v(hidden_states)
|
||||
|
||||
@ -16,16 +16,16 @@ class SingleTransformerBlock(nn.Module):
|
||||
self.attn = SingleBlockAttention()
|
||||
self.proj_out = nn.Linear(3072 + 4 * 3072, 3072)
|
||||
|
||||
def forward(
|
||||
def __call__(
|
||||
self,
|
||||
hidden_states: mx.array,
|
||||
text_embeddings: mx.array,
|
||||
rotary_embeddings: mx.array,
|
||||
) -> (mx.array, mx.array):
|
||||
residual = hidden_states
|
||||
norm_hidden_states, gate = self.norm.forward(x=hidden_states, text_embeddings=text_embeddings)
|
||||
norm_hidden_states, gate = self.norm(x=hidden_states, text_embeddings=text_embeddings)
|
||||
mlp_hidden_states = nn.gelu_approx(self.proj_mlp(norm_hidden_states))
|
||||
attn_output = self.attn.forward(
|
||||
attn_output = self.attn(
|
||||
hidden_states=norm_hidden_states,
|
||||
image_rotary_emb=rotary_embeddings,
|
||||
)
|
||||
|
||||
@ -8,7 +8,7 @@ class TextEmbedder(nn.Module):
|
||||
self.linear_1 = nn.Linear(768, 3072)
|
||||
self.linear_2 = nn.Linear(3072, 3072)
|
||||
|
||||
def forward(self, caption: mx.array) -> mx.array:
|
||||
def __call__(self, caption: mx.array) -> mx.array:
|
||||
hidden_states = self.linear_1(caption)
|
||||
hidden_states = nn.silu(hidden_states)
|
||||
hidden_states = self.linear_2(hidden_states)
|
||||
|
||||
@ -17,17 +17,17 @@ class TimeTextEmbed(nn.Module):
|
||||
self.guidance_embedder = GuidanceEmbedder() if model_config == ModelConfig.FLUX1_DEV else None
|
||||
self.timestep_embedder = TimestepEmbedder()
|
||||
|
||||
def forward(
|
||||
def __call__(
|
||||
self,
|
||||
time_step: mx.array,
|
||||
pooled_projection: mx.array,
|
||||
guidance: mx.array,
|
||||
) -> mx.array:
|
||||
time_steps_proj = self._time_proj(time_step)
|
||||
time_steps_emb = self.timestep_embedder.forward(time_steps_proj)
|
||||
time_steps_emb = self.timestep_embedder(time_steps_proj)
|
||||
if self.guidance_embedder is not None:
|
||||
time_steps_emb += self.guidance_embedder.forward(self._time_proj(guidance))
|
||||
pooled_projections = self.text_embedder.forward(pooled_projection)
|
||||
time_steps_emb += self.guidance_embedder(self._time_proj(guidance))
|
||||
pooled_projections = self.text_embedder(pooled_projection)
|
||||
conditioning = time_steps_emb + pooled_projections
|
||||
return conditioning.astype(Config.precision)
|
||||
|
||||
|
||||
@ -8,7 +8,7 @@ class TimestepEmbedder(nn.Module):
|
||||
self.linear_1 = nn.Linear(256, 3072)
|
||||
self.linear_2 = nn.Linear(3072, 3072)
|
||||
|
||||
def forward(self, sample: mx.array) -> mx.array:
|
||||
def __call__(self, sample: mx.array) -> mx.array:
|
||||
sample = self.linear_1(sample)
|
||||
sample = nn.silu(sample)
|
||||
sample = self.linear_2(sample)
|
||||
|
||||
@ -44,15 +44,15 @@ class Transformer(nn.Module):
|
||||
time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision)
|
||||
hidden_states = self.x_embedder(hidden_states)
|
||||
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)
|
||||
text_embeddings = self.time_text_embed(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)
|
||||
image_rotary_emb = self.pos_embed(ids)
|
||||
|
||||
for idx, block in enumerate(self.transformer_blocks):
|
||||
encoder_hidden_states, hidden_states = block.forward(
|
||||
encoder_hidden_states, hidden_states = block(
|
||||
hidden_states=hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
text_embeddings=text_embeddings,
|
||||
@ -66,7 +66,7 @@ class Transformer(nn.Module):
|
||||
hidden_states = mx.concatenate([encoder_hidden_states, hidden_states], axis=1)
|
||||
|
||||
for idx, block in enumerate(self.single_transformer_blocks):
|
||||
hidden_states = block.forward(
|
||||
hidden_states = block(
|
||||
hidden_states=hidden_states,
|
||||
text_embeddings=text_embeddings,
|
||||
rotary_embeddings=image_rotary_emb,
|
||||
@ -80,7 +80,7 @@ class Transformer(nn.Module):
|
||||
)
|
||||
|
||||
hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
|
||||
hidden_states = self.norm_out.forward(hidden_states, text_embeddings)
|
||||
hidden_states = self.norm_out(hidden_states, text_embeddings)
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
noise = hidden_states
|
||||
return noise
|
||||
|
||||
@ -13,7 +13,7 @@ class Attention(nn.Module):
|
||||
self.to_v = nn.Linear(512, 512)
|
||||
self.to_out = [nn.Linear(512, 512)]
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
input_array = mx.transpose(input_array, (0, 2, 3, 1))
|
||||
|
||||
B, H, W, C = input_array.shape
|
||||
|
||||
@ -55,7 +55,7 @@ class ResnetBlock2D(nn.Module):
|
||||
stride=(1, 1),
|
||||
)
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
input_array = mx.transpose(input_array, (0, 2, 3, 1))
|
||||
hidden_states = self.norm1(input_array.astype(mx.float32)).astype(Config.precision)
|
||||
hidden_states = nn.silu(hidden_states)
|
||||
|
||||
@ -14,8 +14,8 @@ class UnetMidBlock(nn.Module):
|
||||
ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512),
|
||||
]
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0].forward(input_array)
|
||||
hidden_states = self.attentions[0].forward(hidden_states)
|
||||
hidden_states = self.resnets[1].forward(hidden_states)
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0](input_array)
|
||||
hidden_states = self.attentions[0](hidden_states)
|
||||
hidden_states = self.resnets[1](hidden_states)
|
||||
return hidden_states
|
||||
|
||||
@ -13,6 +13,6 @@ class ConvIn(nn.Module):
|
||||
padding=1,
|
||||
)
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
input_array = mx.transpose(input_array, (0, 2, 3, 1))
|
||||
return mx.transpose(self.conv2d(input_array), (0, 3, 1, 2))
|
||||
|
||||
@ -15,7 +15,7 @@ class ConvNormOut(nn.Module):
|
||||
pytorch_compatible=True,
|
||||
)
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
input_array = mx.transpose(input_array, (0, 2, 3, 1))
|
||||
hidden_states = self.norm(input_array.astype(mx.float32)).astype(Config.precision)
|
||||
return mx.transpose(hidden_states, (0, 3, 1, 2))
|
||||
|
||||
@ -13,6 +13,6 @@ class ConvOut(nn.Module):
|
||||
padding=(1, 1),
|
||||
)
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
input_array = mx.transpose(input_array, (0, 2, 3, 1))
|
||||
return mx.transpose(self.conv2d(input_array), (0, 3, 1, 2))
|
||||
|
||||
@ -24,12 +24,12 @@ class Decoder(nn.Module):
|
||||
self.conv_norm_out = ConvNormOut()
|
||||
self.conv_out = ConvOut()
|
||||
|
||||
def decode(self, latents: mx.array) -> mx.array:
|
||||
latents = self.conv_in.forward(latents)
|
||||
latents = self.mid_block.forward(latents)
|
||||
def __call__(self, latents: mx.array) -> mx.array:
|
||||
latents = self.conv_in(latents)
|
||||
latents = self.mid_block(latents)
|
||||
for up_block in self.up_blocks:
|
||||
latents = up_block.forward(latents)
|
||||
latents = self.conv_norm_out.forward(latents)
|
||||
latents = up_block(latents)
|
||||
latents = self.conv_norm_out(latents)
|
||||
latents = nn.silu(latents)
|
||||
latents = self.conv_out.forward(latents)
|
||||
latents = self.conv_out(latents)
|
||||
return latents
|
||||
|
||||
@ -15,12 +15,12 @@ class UpBlock1Or2(nn.Module):
|
||||
]
|
||||
self.upsamplers = [UpSampler(conv_in=512, conv_out=512)]
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0].forward(input_array)
|
||||
hidden_states = self.resnets[1].forward(hidden_states)
|
||||
hidden_states = self.resnets[2].forward(hidden_states)
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0](input_array)
|
||||
hidden_states = self.resnets[1](hidden_states)
|
||||
hidden_states = self.resnets[2](hidden_states)
|
||||
|
||||
if self.upsamplers is not None:
|
||||
hidden_states = self.upsamplers[0].forward(hidden_states)
|
||||
hidden_states = self.upsamplers[0](hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
@ -17,12 +17,12 @@ class UpBlock3(nn.Module):
|
||||
# fmt: on
|
||||
self.upsamplers = [UpSampler(conv_in=256, conv_out=256)]
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0].forward(input_array)
|
||||
hidden_states = self.resnets[1].forward(hidden_states)
|
||||
hidden_states = self.resnets[2].forward(hidden_states)
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0](input_array)
|
||||
hidden_states = self.resnets[1](hidden_states)
|
||||
hidden_states = self.resnets[2](hidden_states)
|
||||
|
||||
if self.upsamplers is not None:
|
||||
hidden_states = self.upsamplers[0].forward(hidden_states)
|
||||
hidden_states = self.upsamplers[0](hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
@ -15,8 +15,8 @@ class UpBlock4(nn.Module):
|
||||
]
|
||||
# fmt: off
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0].forward(input_array)
|
||||
hidden_states = self.resnets[1].forward(hidden_states)
|
||||
hidden_states = self.resnets[2].forward(hidden_states)
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0](input_array)
|
||||
hidden_states = self.resnets[1](hidden_states)
|
||||
hidden_states = self.resnets[2](hidden_states)
|
||||
return hidden_states
|
||||
|
||||
@ -13,7 +13,7 @@ class UpSampler(nn.Module):
|
||||
padding=(1, 1),
|
||||
)
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
input_array = mx.transpose(input_array, (0, 2, 3, 1))
|
||||
hidden_states = UpSampler.up_sample_nearest(input_array)
|
||||
hidden_state = self.conv(hidden_states)
|
||||
|
||||
@ -13,6 +13,6 @@ class ConvIn(nn.Module):
|
||||
padding=1,
|
||||
)
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
input_array = mx.transpose(input_array, (0, 2, 3, 1))
|
||||
return mx.transpose(self.conv2d(input_array), (0, 3, 1, 2))
|
||||
|
||||
@ -13,7 +13,7 @@ class ConvNormOut(nn.Module):
|
||||
pytorch_compatible=True,
|
||||
)
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
input_array = mx.transpose(input_array, (0, 2, 3, 1))
|
||||
hidden_states = self.norm(input_array.astype(mx.float32)).astype(mx.float32)
|
||||
return mx.transpose(hidden_states, (0, 3, 1, 2))
|
||||
|
||||
@ -13,6 +13,6 @@ class ConvOut(nn.Module):
|
||||
padding=(1, 1),
|
||||
)
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
input_array = mx.transpose(input_array, (0, 2, 3, 1))
|
||||
return mx.transpose(self.conv2d(input_array), (0, 3, 1, 2))
|
||||
|
||||
@ -16,11 +16,11 @@ class DownBlock1(nn.Module):
|
||||
]
|
||||
self.downsamplers = [DownSampler(conv_in=128, conv_out=128)]
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0].forward(input_array)
|
||||
hidden_states = self.resnets[1].forward(hidden_states)
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0](input_array)
|
||||
hidden_states = self.resnets[1](hidden_states)
|
||||
|
||||
if self.downsamplers is not None:
|
||||
hidden_states = self.downsamplers[0].forward(hidden_states)
|
||||
hidden_states = self.downsamplers[0](hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
@ -16,11 +16,11 @@ class DownBlock2(nn.Module):
|
||||
# fmt: on
|
||||
self.downsamplers = [DownSampler(conv_in=256, conv_out=256)]
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0].forward(input_array)
|
||||
hidden_states = self.resnets[1].forward(hidden_states)
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0](input_array)
|
||||
hidden_states = self.resnets[1](hidden_states)
|
||||
|
||||
if self.downsamplers is not None:
|
||||
hidden_states = self.downsamplers[0].forward(hidden_states)
|
||||
hidden_states = self.downsamplers[0](hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
@ -16,11 +16,11 @@ class DownBlock3(nn.Module):
|
||||
# fmt: on
|
||||
self.downsamplers = [DownSampler(conv_in=512, conv_out=512)]
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0].forward(input_array)
|
||||
hidden_states = self.resnets[1].forward(hidden_states)
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0](input_array)
|
||||
hidden_states = self.resnets[1](hidden_states)
|
||||
|
||||
if self.downsamplers is not None:
|
||||
hidden_states = self.downsamplers[0].forward(hidden_states)
|
||||
hidden_states = self.downsamplers[0](hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
@ -14,7 +14,7 @@ class DownBlock4(nn.Module):
|
||||
]
|
||||
# fmt: on
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0].forward(input_array)
|
||||
hidden_states = self.resnets[1].forward(hidden_states)
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = self.resnets[0](input_array)
|
||||
hidden_states = self.resnets[1](hidden_states)
|
||||
return hidden_states
|
||||
|
||||
@ -12,7 +12,7 @@ class DownSampler(nn.Module):
|
||||
stride=(2, 2),
|
||||
)
|
||||
|
||||
def forward(self, input_array: mx.array) -> mx.array:
|
||||
def __call__(self, input_array: mx.array) -> mx.array:
|
||||
hidden_states = mx.pad(input_array, ((0, 0), (0, 0), (0, 1), (0, 1)))
|
||||
hidden_states = mx.transpose(hidden_states, (0, 2, 3, 1))
|
||||
hidden_state = self.conv(hidden_states)
|
||||
|
||||
@ -25,12 +25,12 @@ class Encoder(nn.Module):
|
||||
self.conv_norm_out = ConvNormOut()
|
||||
self.conv_out = ConvOut()
|
||||
|
||||
def encode(self, latents: mx.array) -> mx.array:
|
||||
latents = self.conv_in.forward(latents)
|
||||
def __call__(self, latents: mx.array) -> mx.array:
|
||||
latents = self.conv_in(latents)
|
||||
for down_block in self.down_blocks:
|
||||
latents = down_block.forward(latents)
|
||||
latents = self.mid_block.forward(latents)
|
||||
latents = self.conv_norm_out.forward(latents)
|
||||
latents = down_block(latents)
|
||||
latents = self.mid_block(latents)
|
||||
latents = self.conv_norm_out(latents)
|
||||
latents = nn.silu(latents)
|
||||
latents = self.conv_out.forward(latents)
|
||||
latents = self.conv_out(latents)
|
||||
return latents
|
||||
|
||||
@ -16,9 +16,9 @@ class VAE(nn.Module):
|
||||
|
||||
def decode(self, latents: mx.array) -> mx.array:
|
||||
scaled_latents = (latents / self.scaling_factor) + self.shift_factor
|
||||
return self.decoder.decode(scaled_latents)
|
||||
return self.decoder(scaled_latents)
|
||||
|
||||
def encode(self, latents: mx.array) -> mx.array:
|
||||
latents = self.encoder.encode(latents)
|
||||
latents = self.encoder(latents)
|
||||
mean, _ = mx.split(latents, 2, axis=1)
|
||||
return (mean - self.shift_factor) * self.scaling_factor
|
||||
|
||||
@ -4,6 +4,7 @@ import typing as t
|
||||
|
||||
import mlx.core as mx
|
||||
import PIL.Image
|
||||
import toml
|
||||
|
||||
from mflux.config.model_config import ModelConfig
|
||||
|
||||
@ -55,7 +56,7 @@ class GeneratedImage:
|
||||
return {
|
||||
# mflux_version is used by future metadata readers
|
||||
# to determine supportability of metadata-derived workflows
|
||||
"mflux_version": str(GeneratedImage.get_version()),
|
||||
"mflux_version": GeneratedImage.get_version(),
|
||||
"model": str(self.model_config.alias),
|
||||
"seed": self.seed,
|
||||
"steps": self.steps,
|
||||
@ -73,8 +74,27 @@ class GeneratedImage:
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_version():
|
||||
def get_version() -> str:
|
||||
version = GeneratedImage._get_version_from_toml()
|
||||
if version:
|
||||
return version
|
||||
|
||||
# Fallback to installed package version
|
||||
try:
|
||||
return importlib.metadata.version("mflux")
|
||||
return str(importlib.metadata.version("mflux"))
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
return "unknown"
|
||||
|
||||
@staticmethod
|
||||
def _get_version_from_toml() -> str | None:
|
||||
# Search for pyproject.toml by traversing up from the current working directory
|
||||
current_dir = pathlib.Path(__file__).resolve().parent
|
||||
for parent in current_dir.parents:
|
||||
pyproject_path = parent / "pyproject.toml"
|
||||
if pyproject_path.exists():
|
||||
try:
|
||||
pyproject_data = toml.load(pyproject_path)
|
||||
return pyproject_data.get("project", {}).get("version")
|
||||
except (toml.TomlDecodeError, KeyError, TypeError):
|
||||
return None
|
||||
return None
|
||||
|
||||
31
src/mflux/train.py
Normal file
31
src/mflux/train.py
Normal file
@ -0,0 +1,31 @@
|
||||
from mflux.dreambooth.dreambooth import DreamBooth
|
||||
from mflux.dreambooth.dreambooth_initializer import DreamBoothInitializer
|
||||
from mflux.error.exceptions import StopTrainingException
|
||||
from mflux.ui.cli.parsers import CommandLineParser
|
||||
|
||||
|
||||
def main():
|
||||
parser = CommandLineParser(description="Finetune a LoRA adapter")
|
||||
parser.add_model_arguments(require_model_arg=False)
|
||||
parser.add_training_arguments()
|
||||
args = parser.parse_args()
|
||||
|
||||
flux, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize(
|
||||
config_path=args.train_config,
|
||||
checkpoint_path=args.train_checkpoint
|
||||
) # fmt: off
|
||||
|
||||
try:
|
||||
DreamBooth.train(
|
||||
flux=flux,
|
||||
runtime_config=runtime_config,
|
||||
training_spec=training_spec,
|
||||
training_state=training_state
|
||||
) # fmt: off
|
||||
except StopTrainingException as stop_exc:
|
||||
training_state.save(training_spec)
|
||||
print(stop_exc)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -71,9 +71,19 @@ class CommandLineParser(argparse.ArgumentParser):
|
||||
self.supports_metadata_config = True
|
||||
self.add_argument("--config-from-metadata", "-C", type=Path, required=False, default=argparse.SUPPRESS, help="Re-use the parameters from prior metadata. Params from metadata are secondary to other args you provide.")
|
||||
|
||||
def add_training_arguments(self) -> None:
|
||||
self.add_argument("--train-config", type=str, required=False, help="Local path of the training configuration file")
|
||||
self.add_argument("--train-checkpoint", type=str, required=False, help="Local path of the checkpoint file which specifies how to continue the training process")
|
||||
|
||||
def parse_args(self, **kwargs) -> argparse.Namespace:
|
||||
namespace = super().parse_args()
|
||||
if hasattr(namespace, "path") and namespace.path is not None and namespace.model is None:
|
||||
|
||||
# Check if either training arguments are provided
|
||||
has_training_args = (hasattr(namespace, "train_config") and namespace.train_config is not None) or \
|
||||
(hasattr(namespace, "train_checkpoint") and namespace.train_checkpoint is not None)
|
||||
|
||||
# Only enforce model requirement for path if we're not in training mode
|
||||
if hasattr(namespace, "path") and namespace.path is not None and namespace.model is None and not has_training_args:
|
||||
self.error("--model must be specified when using --path")
|
||||
|
||||
if getattr(namespace, "config_from_metadata", False):
|
||||
@ -125,7 +135,8 @@ class CommandLineParser(argparse.ArgumentParser):
|
||||
if namespace.controlnet_save_canny == self.get_default("controlnet_save_canny") and (cnet_canny_from_metadata := prior_gen_metadata.get("controlnet_save_canny", None)):
|
||||
namespace.controlnet_save_canny = cnet_canny_from_metadata
|
||||
|
||||
if namespace.model is None:
|
||||
# Only require model if we're not in training mode
|
||||
if namespace.model is None and not has_training_args:
|
||||
self.error("--model / -m must be provided, or 'model' must be specified in the config file.")
|
||||
|
||||
if self.supports_image_generation and namespace.prompt is None:
|
||||
|
||||
@ -1,88 +0,0 @@
|
||||
import logging
|
||||
|
||||
from mlx.utils import tree_flatten
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LoraUtil:
|
||||
@staticmethod
|
||||
def apply_loras(
|
||||
transformer: dict,
|
||||
lora_files: list[str],
|
||||
lora_scales: list[float] | None = None,
|
||||
) -> None:
|
||||
lora_scales = LoraUtil._validate_lora_scales(lora_files, lora_scales)
|
||||
|
||||
for lora_file, lora_scale in zip(lora_files, lora_scales):
|
||||
LoraUtil._apply_lora(transformer, lora_file, lora_scale)
|
||||
|
||||
@staticmethod
|
||||
def _validate_lora_scales(lora_files: list[str], lora_scales: list[float]) -> list[float]:
|
||||
if len(lora_files) == 1:
|
||||
if not lora_scales:
|
||||
lora_scales = [1.0]
|
||||
if len(lora_scales) > 1:
|
||||
raise ValueError("Please provide a single scale for the LoRA, or skip it to default to 1")
|
||||
elif len(lora_files) > 1:
|
||||
if len(lora_files) != len(lora_scales):
|
||||
raise ValueError("When providing multiple LoRAs, be sure to specify a scale for each one respectively")
|
||||
return lora_scales
|
||||
|
||||
@staticmethod
|
||||
def _apply_lora(transformer: dict, lora_file: str, lora_scale: float) -> None:
|
||||
from mflux.weights.weight_handler import WeightHandler
|
||||
|
||||
lora_transformer, _ = WeightHandler.load_transformer(lora_path=lora_file)
|
||||
LoraUtil._apply_transformer(transformer, lora_transformer, lora_scale)
|
||||
|
||||
@staticmethod
|
||||
def _apply_transformer(transformer: dict, lora_transformer: dict, lora_scale: float) -> None:
|
||||
lora_weights = tree_flatten(lora_transformer)
|
||||
visited = {}
|
||||
|
||||
for key, weight in lora_weights:
|
||||
splits = key.split(".")
|
||||
target = transformer
|
||||
visiting = []
|
||||
for splitKey in splits:
|
||||
if isinstance(target, dict) and splitKey in target:
|
||||
target = target[splitKey]
|
||||
visiting.append(splitKey)
|
||||
elif isinstance(target, list) and len(target) > 0:
|
||||
if len(target) < int(splitKey):
|
||||
for _ in range(int(splitKey) - len(target) + 1):
|
||||
target.append({})
|
||||
|
||||
target = target[int(splitKey)]
|
||||
visiting.append(splitKey)
|
||||
else:
|
||||
parentKey = ".".join(visiting)
|
||||
if parentKey in visited and "lora_A" in visited[parentKey] and "lora_B" in visited[parentKey]:
|
||||
continue
|
||||
if not splitKey.startswith("lora_"):
|
||||
visiting.append(splitKey)
|
||||
parentKey = ".".join(visiting)
|
||||
if splitKey == "net":
|
||||
target["net"] = list({})
|
||||
target = target["net"]
|
||||
elif splitKey == "0":
|
||||
target.append({})
|
||||
target = target[0]
|
||||
continue
|
||||
elif splitKey == "proj":
|
||||
target[splitKey] = weight
|
||||
if parentKey not in visited:
|
||||
visited[parentKey] = {}
|
||||
continue
|
||||
if parentKey not in visited:
|
||||
visited[parentKey] = {}
|
||||
visited[parentKey][splitKey] = weight
|
||||
if "weight" not in target:
|
||||
raise ValueError(f"LoRA weights for layer {parentKey} cannot be loaded into the model.")
|
||||
if "lora_A" in visited[parentKey] and "lora_B" in visited[parentKey]:
|
||||
lora_a = visited[parentKey]["lora_A"]
|
||||
lora_b = visited[parentKey]["lora_B"]
|
||||
transWeight = target["weight"]
|
||||
weight = transWeight + lora_scale * (lora_b @ lora_a)
|
||||
target["weight"] = weight
|
||||
40
src/mflux/weights/quantization_util.py
Normal file
40
src/mflux/weights/quantization_util.py
Normal file
@ -0,0 +1,40 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import mlx.nn as nn
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet
|
||||
from mflux.weights.weight_handler import WeightHandler
|
||||
|
||||
|
||||
class QuantizationUtil:
|
||||
@staticmethod
|
||||
def quantize_model(
|
||||
vae: nn.Module,
|
||||
transformer: nn.Module,
|
||||
t5_text_encoder: nn.Module,
|
||||
clip_text_encoder: nn.Module,
|
||||
quantize: int,
|
||||
weights: "WeightHandler",
|
||||
) -> None:
|
||||
q_level = weights.meta_data.quantization_level
|
||||
|
||||
if quantize is not None or q_level is not None:
|
||||
bits = q_level if q_level is not None else quantize
|
||||
# fmt: off
|
||||
nn.quantize(vae, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=bits)
|
||||
nn.quantize(transformer, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 64, group_size=64, bits=bits)
|
||||
nn.quantize(t5_text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=bits)
|
||||
nn.quantize(clip_text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=bits)
|
||||
# fmt: on
|
||||
|
||||
@staticmethod
|
||||
def quantize_controlnet(
|
||||
quantize: int,
|
||||
weights: "WeightHandlerControlnet",
|
||||
transformer_controlnet: nn.Module,
|
||||
):
|
||||
q_level = weights.meta_data.quantization_level
|
||||
if quantize is not None or q_level is not None:
|
||||
bits = q_level if q_level is not None else quantize
|
||||
nn.quantize(transformer_controlnet, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 128, group_size=128, bits=bits) # fmt: off
|
||||
@ -1,3 +1,4 @@
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
@ -5,36 +6,65 @@ from huggingface_hub import snapshot_download
|
||||
from mlx.utils import tree_unflatten
|
||||
|
||||
from mflux.weights.lora_converter import LoRAConverter
|
||||
from mflux.weights.lora_util import LoraUtil
|
||||
from mflux.weights.weight_util import WeightUtil
|
||||
|
||||
|
||||
@dataclass
|
||||
class MetaData:
|
||||
quantization_level: int | None = None
|
||||
scale: float | None = None
|
||||
is_lora: bool = False
|
||||
is_mflux: bool = False
|
||||
|
||||
|
||||
class WeightHandler:
|
||||
def __init__(
|
||||
self,
|
||||
repo_id: str | None = None,
|
||||
local_path: str | None = None,
|
||||
lora_paths: list[str] | None = None,
|
||||
lora_scales: list[float] | None = None,
|
||||
meta_data: MetaData,
|
||||
clip_encoder: dict | None = None,
|
||||
t5_encoder: dict | None = None,
|
||||
vae: dict | None = None,
|
||||
transformer: dict | None = None,
|
||||
):
|
||||
root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id)
|
||||
|
||||
self.clip_encoder, _ = WeightHandler.load_clip_encoder(root_path=root_path)
|
||||
self.t5_encoder, _ = WeightHandler.load_t5_encoder(root_path=root_path)
|
||||
self.vae, _ = WeightHandler.load_vae(root_path=root_path)
|
||||
self.transformer, self.quantization_level = WeightHandler.load_transformer(root_path=root_path)
|
||||
|
||||
if lora_paths:
|
||||
LoraUtil.apply_loras(self.transformer, lora_paths, lora_scales)
|
||||
self.clip_encoder = clip_encoder
|
||||
self.t5_encoder = t5_encoder
|
||||
self.vae = vae
|
||||
self.transformer = transformer
|
||||
self.meta_data = meta_data
|
||||
|
||||
@staticmethod
|
||||
def load_clip_encoder(root_path: Path) -> (dict, int):
|
||||
weights, quantization_level = WeightHandler._get_weights("text_encoder", root_path)
|
||||
def load_regular_weights(
|
||||
repo_id: str | None = None,
|
||||
local_path: str | None = None,
|
||||
) -> "WeightHandler":
|
||||
# Load the weights from disk, huggingface cache, or download from huggingface
|
||||
root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id)
|
||||
|
||||
clip_encoder, _ = WeightHandler._load_clip_encoder(root_path=root_path)
|
||||
t5_encoder, _ = WeightHandler._load_t5_encoder(root_path=root_path)
|
||||
vae, _ = WeightHandler._load_vae(root_path=root_path)
|
||||
transformer, quantization_level, _ = WeightHandler.load_transformer(root_path=root_path)
|
||||
return WeightHandler(
|
||||
clip_encoder=clip_encoder,
|
||||
t5_encoder=t5_encoder,
|
||||
vae=vae,
|
||||
transformer=transformer,
|
||||
meta_data=MetaData(
|
||||
quantization_level=quantization_level,
|
||||
scale=None,
|
||||
is_lora=False,
|
||||
is_mflux=False,
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _load_clip_encoder(root_path: Path) -> (dict, int):
|
||||
weights, quantization_level, _ = WeightHandler._get_weights("text_encoder", root_path)
|
||||
return weights, quantization_level
|
||||
|
||||
@staticmethod
|
||||
def load_t5_encoder(root_path: Path) -> (dict, int):
|
||||
weights, quantization_level = WeightHandler._get_weights("text_encoder_2", root_path)
|
||||
def _load_t5_encoder(root_path: Path) -> (dict, int):
|
||||
weights, quantization_level, _ = WeightHandler._get_weights("text_encoder_2", root_path)
|
||||
|
||||
# Quantized weights (i.e. ones exported from this project) don't need any post-processing.
|
||||
if quantization_level is not None:
|
||||
@ -60,8 +90,8 @@ class WeightHandler:
|
||||
return weights, quantization_level
|
||||
|
||||
@staticmethod
|
||||
def load_transformer(root_path: Path | None = None, lora_path: str | None = None) -> (dict, int):
|
||||
weights, quantization_level = WeightHandler._get_weights("transformer", root_path, lora_path)
|
||||
def load_transformer(root_path: Path | None = None, lora_path: str | None = None) -> (dict, int, str | None):
|
||||
weights, quantization_level, mflux_version = WeightHandler._get_weights("transformer", root_path, lora_path)
|
||||
|
||||
if lora_path:
|
||||
if "transformer" not in weights:
|
||||
@ -69,8 +99,8 @@ class WeightHandler:
|
||||
weights = weights["transformer"]
|
||||
|
||||
# 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
|
||||
if quantization_level or mflux_version:
|
||||
return weights, quantization_level, mflux_version
|
||||
|
||||
# Reshape and process the huggingface weights
|
||||
if "transformer_blocks" in weights:
|
||||
@ -85,11 +115,11 @@ class WeightHandler:
|
||||
"linear1": block["ff_context"]["net"][0]["proj"],
|
||||
"linear2": block["ff_context"]["net"][2],
|
||||
}
|
||||
return weights, quantization_level
|
||||
return weights, quantization_level, mflux_version
|
||||
|
||||
@staticmethod
|
||||
def load_vae(root_path: Path) -> (dict, int):
|
||||
weights, quantization_level = WeightHandler._get_weights("vae", root_path)
|
||||
def _load_vae(root_path: Path) -> (dict, int):
|
||||
weights, quantization_level, _ = WeightHandler._get_weights("vae", root_path)
|
||||
|
||||
# Quantized weights (i.e. ones exported from this project) don't need any post-processing.
|
||||
if quantization_level is not None:
|
||||
@ -109,29 +139,35 @@ class WeightHandler:
|
||||
model_name: str,
|
||||
root_path: Path | None = None,
|
||||
lora_path: str | None = None,
|
||||
) -> (dict, int):
|
||||
) -> (dict, int, str | None):
|
||||
weights = []
|
||||
quantization_level = None
|
||||
mflux_version = None
|
||||
|
||||
if root_path is not None:
|
||||
for file in sorted(root_path.glob(model_name + "/*.safetensors")):
|
||||
quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level")
|
||||
weight = list(mx.load(str(file)).items())
|
||||
data = mx.load(str(file), return_metadata=True)
|
||||
weight = list(data[0].items())
|
||||
if len(data) > 1:
|
||||
quantization_level = data[1].get("quantization_level")
|
||||
weights.extend(weight)
|
||||
|
||||
if lora_path and root_path is None:
|
||||
weight = list(mx.load(lora_path).items())
|
||||
data = mx.load(lora_path, return_metadata=True)
|
||||
weight = list(data[0].items())
|
||||
if len(data) > 1:
|
||||
mflux_version = data[1].get("mflux_version", None)
|
||||
weights.extend(weight)
|
||||
|
||||
# Non huggingface weights (i.e. ones exported from this project) don't need any reshaping.
|
||||
if quantization_level is not None:
|
||||
return tree_unflatten(weights), quantization_level
|
||||
if quantization_level is not None or mflux_version is not None:
|
||||
return tree_unflatten(weights), quantization_level, mflux_version
|
||||
|
||||
# Huggingface weights needs to be reshaped
|
||||
weights = [WeightUtil.reshape_weights(k, v) for k, v in weights]
|
||||
weights = WeightUtil.flatten(weights)
|
||||
unflatten = tree_unflatten(weights)
|
||||
return unflatten, quantization_level
|
||||
return unflatten, quantization_level, mflux_version
|
||||
|
||||
@staticmethod
|
||||
def _download_or_get_cached_weights(repo_id: str) -> Path:
|
||||
|
||||
155
src/mflux/weights/weight_handler_lora.py
Normal file
155
src/mflux/weights/weight_handler_lora.py
Normal file
@ -0,0 +1,155 @@
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
from mlx.utils import tree_flatten, tree_unflatten
|
||||
|
||||
from mflux.dreambooth.lora_layers.fused_linear_lora_layer import FusedLoRALinear
|
||||
from mflux.dreambooth.lora_layers.linear_lora_layer import LoRALinear
|
||||
from mflux.dreambooth.lora_layers.lora_layers import LoRALayers
|
||||
from mflux.weights.weight_handler import MetaData, WeightHandler
|
||||
|
||||
|
||||
class WeightHandlerLoRA:
|
||||
def __init__(self, weight_handlers: list[WeightHandler]):
|
||||
self.weight_handlers = weight_handlers
|
||||
|
||||
@staticmethod
|
||||
def load_lora_weights(
|
||||
transformer: nn.Module,
|
||||
lora_files: list[str],
|
||||
lora_scales: list[float] | None = None,
|
||||
) -> list["WeightHandler"]:
|
||||
lora_weights = []
|
||||
if lora_files:
|
||||
lora_scales = WeightHandlerLoRA._validate_lora_scales(lora_files, lora_scales)
|
||||
for lora_file, lora_scale in zip(lora_files, lora_scales):
|
||||
weights, _, mflux_version = WeightHandler.load_transformer(lora_path=lora_file)
|
||||
weights = dict(tree_flatten(weights))
|
||||
weights = {key.removesuffix(".weight"): value for key, value in weights.items()}
|
||||
weights = {f"transformer.{key}": value for key, value in weights.items()}
|
||||
weights = {key: mx.transpose(value) for key, value in weights.items()}
|
||||
lora_transformer_dict = LoRALayers.transformer_dict_from_template(weights, transformer, lora_scale)
|
||||
transformer_weights = tree_unflatten(list(lora_transformer_dict.items()))["transformer"]
|
||||
weights = WeightHandler(
|
||||
clip_encoder=None,
|
||||
t5_encoder=None,
|
||||
vae=None,
|
||||
transformer=transformer_weights,
|
||||
meta_data=MetaData(
|
||||
quantization_level=None,
|
||||
scale=lora_scale,
|
||||
is_lora=True,
|
||||
is_mflux=True if mflux_version is not None else False,
|
||||
),
|
||||
)
|
||||
lora_weights.append(weights)
|
||||
|
||||
return lora_weights
|
||||
|
||||
@staticmethod
|
||||
def set_lora_weights(transformer: nn.Module, loras: list["WeightHandler"]) -> None:
|
||||
if loras:
|
||||
lora_transformer_weights = [lora.transformer for lora in loras]
|
||||
fused_weights = WeightHandlerLoRA._fuse_multiple_lora_dicts(lora_transformer_weights)
|
||||
fused_weights = WeightHandler(
|
||||
meta_data=MetaData(),
|
||||
clip_encoder=None,
|
||||
t5_encoder=None,
|
||||
vae=None,
|
||||
transformer=fused_weights,
|
||||
)
|
||||
|
||||
WeightHandlerLoRA.set_lora_layers(
|
||||
transformer_module=transformer,
|
||||
lora_layers=LoRALayers(weights=fused_weights)
|
||||
) # fmt:off
|
||||
|
||||
@staticmethod
|
||||
def _fuse_multiple_lora_dicts(dicts: list[dict]) -> dict:
|
||||
if not dicts:
|
||||
raise ValueError("No dictionaries provided for fusion.")
|
||||
if len(dicts) == 1:
|
||||
return dicts[0]
|
||||
|
||||
# Collect all unique keys across all dictionaries
|
||||
all_keys = set().union(*dicts)
|
||||
fused_dict = {}
|
||||
|
||||
for key in all_keys:
|
||||
# Get all values for this key, filtering out dictionaries that don't have it
|
||||
values = [d[key] for d in dicts if key in d]
|
||||
|
||||
# Skip if no values found (shouldn't happen due to how we collect keys)
|
||||
if not values:
|
||||
continue
|
||||
|
||||
first_value = values[0]
|
||||
|
||||
# Handle nested dictionaries
|
||||
if all(isinstance(v, dict) and not isinstance(v, LoRALinear) for v in values):
|
||||
fused_dict[key] = WeightHandlerLoRA._fuse_multiple_lora_dicts(values)
|
||||
|
||||
# Handle LoRALinear layers
|
||||
elif all(isinstance(v, LoRALinear) for v in values):
|
||||
fused_dict[key] = FusedLoRALinear(base_linear=first_value.linear, loras=values)
|
||||
|
||||
# Handle lists
|
||||
elif all(isinstance(v, list) for v in values):
|
||||
# Get the maximum length of all lists
|
||||
max_length = max(len(v) for v in values)
|
||||
|
||||
# Initialize the fused list
|
||||
fused_dict[key] = []
|
||||
|
||||
# Process each index up to the maximum length
|
||||
for idx in range(max_length):
|
||||
# Get elements at current index from lists that are long enough
|
||||
elements = [v[idx] for v in values if idx < len(v)]
|
||||
|
||||
if not elements:
|
||||
continue
|
||||
|
||||
# If elements are dicts or LoRALinear, recursively fuse them
|
||||
if all(isinstance(e, (dict, LoRALinear)) for e in elements):
|
||||
fused_element = WeightHandlerLoRA._fuse_multiple_lora_dicts([{str(idx): e} for e in elements])[str(idx)] # fmt:off
|
||||
fused_dict[key].append(fused_element)
|
||||
else:
|
||||
# For non-LoRALinear types, keep the first element
|
||||
fused_dict[key].append(elements[0])
|
||||
|
||||
else:
|
||||
types_str = ", ".join(type(v).__name__ for v in values)
|
||||
raise ValueError(f"Incompatible types for key {key}: {types_str}")
|
||||
|
||||
return fused_dict
|
||||
|
||||
@staticmethod
|
||||
def _validate_lora_scales(lora_files: list[str], lora_scales: list[float]) -> list[float]:
|
||||
if len(lora_files) == 1:
|
||||
if not lora_scales:
|
||||
lora_scales = [1.0]
|
||||
if len(lora_scales) > 1:
|
||||
raise ValueError("Please provide a single scale for the LoRA, or skip it to default to 1")
|
||||
elif len(lora_files) > 1:
|
||||
if len(lora_files) != len(lora_scales):
|
||||
raise ValueError("When providing multiple LoRAs, be sure to specify a scale for each one respectively")
|
||||
return lora_scales
|
||||
|
||||
@staticmethod
|
||||
def set_lora_layers(transformer_module: nn.Module, lora_layers: LoRALayers) -> None:
|
||||
transformer = lora_layers.layers.transformer
|
||||
|
||||
# Handle transformer_blocks
|
||||
transformer_blocks = transformer.get("transformer_blocks", [])
|
||||
for i, weights in enumerate(transformer_blocks):
|
||||
LoRALayers.set_transformer_block(
|
||||
transformer_block=transformer_module.transformer_blocks[i],
|
||||
dictionary=weights
|
||||
) # fmt:off
|
||||
|
||||
# Handle single_transformer_blocks
|
||||
single_transformer_blocks = transformer.get("single_transformer_blocks", [])
|
||||
for i, weights in enumerate(single_transformer_blocks):
|
||||
LoRALayers.set_single_transformer_block(
|
||||
single_transformer_block=transformer_module.single_transformer_blocks[i],
|
||||
dictionary=weights
|
||||
) # fmt:off
|
||||
@ -1,4 +1,13 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import mlx.nn as nn
|
||||
|
||||
from mflux.config.config import Config
|
||||
from mflux.weights.quantization_util import QuantizationUtil
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet
|
||||
from mflux.weights.weight_handler import WeightHandler
|
||||
|
||||
|
||||
class WeightUtil:
|
||||
@ -12,3 +21,65 @@ class WeightUtil:
|
||||
value = value.transpose(0, 2, 3, 1)
|
||||
value = value.reshape(-1).reshape(value.shape).astype(Config.precision)
|
||||
return [(key, value)]
|
||||
|
||||
@staticmethod
|
||||
def set_weights_and_quantize(
|
||||
quantize_arg: int | None,
|
||||
weights: "WeightHandler",
|
||||
vae: nn.Module,
|
||||
transformer: nn.Module,
|
||||
t5_text_encoder: nn.Module,
|
||||
clip_text_encoder: nn.Module,
|
||||
) -> int | None:
|
||||
if weights.meta_data.quantization_level is None and quantize_arg is None:
|
||||
WeightUtil._set_model_weights(weights, vae, transformer, t5_text_encoder, clip_text_encoder)
|
||||
return None
|
||||
|
||||
if weights.meta_data.quantization_level is None and quantize_arg is not None:
|
||||
bits = quantize_arg
|
||||
WeightUtil._set_model_weights(weights, vae, transformer, t5_text_encoder, clip_text_encoder)
|
||||
QuantizationUtil.quantize_model(vae, transformer, t5_text_encoder, clip_text_encoder, bits, weights) # fmt:off
|
||||
return bits
|
||||
|
||||
if weights.meta_data.quantization_level is not None:
|
||||
bits = weights.meta_data.quantization_level
|
||||
QuantizationUtil.quantize_model(vae, transformer, t5_text_encoder, clip_text_encoder, bits, weights) # fmt:off
|
||||
WeightUtil._set_model_weights(weights, vae, transformer, t5_text_encoder, clip_text_encoder)
|
||||
return bits
|
||||
|
||||
raise Exception("Error setting weights")
|
||||
|
||||
@staticmethod
|
||||
def set_controlnet_weights_and_quantize(
|
||||
quantize_arg: int | None,
|
||||
weights: "WeightHandlerControlnet",
|
||||
transformer_controlnet: nn.Module,
|
||||
) -> int | None:
|
||||
if weights.meta_data.quantization_level is None and quantize_arg is None:
|
||||
transformer_controlnet.update(weights.controlnet_transformer)
|
||||
return None
|
||||
|
||||
if weights.meta_data.quantization_level is None and quantize_arg is not None:
|
||||
bits = quantize_arg
|
||||
transformer_controlnet.update(weights.controlnet_transformer)
|
||||
QuantizationUtil.quantize_controlnet(bits, weights, transformer_controlnet)
|
||||
return bits
|
||||
|
||||
if weights.meta_data.quantization_level is not None:
|
||||
bits = weights.meta_data.quantization_level
|
||||
QuantizationUtil.quantize_controlnet(bits, weights, transformer_controlnet)
|
||||
transformer_controlnet.update(weights.controlnet_transformer)
|
||||
return bits
|
||||
|
||||
@staticmethod
|
||||
def _set_model_weights(
|
||||
weights: "WeightHandler",
|
||||
vae: nn.Module,
|
||||
transformer: nn.Module,
|
||||
t5_text_encoder: nn.Module,
|
||||
clip_text_encoder: nn.Module,
|
||||
):
|
||||
vae.update(weights.vae)
|
||||
transformer.update(weights.transformer)
|
||||
t5_text_encoder.update(weights.t5_encoder)
|
||||
clip_text_encoder.update(weights.clip_encoder)
|
||||
|
||||
0
tests/arg_parser/__init__.py
Normal file
0
tests/arg_parser/__init__.py
Normal file
0
tests/dreambooth/__init__.py
Normal file
0
tests/dreambooth/__init__.py
Normal file
88
tests/dreambooth/config/train.json
Normal file
88
tests/dreambooth/config/train.json
Normal file
@ -0,0 +1,88 @@
|
||||
{
|
||||
"model": "dev",
|
||||
"seed": 42,
|
||||
"steps": 2,
|
||||
"guidance": 3.0,
|
||||
"quantize": 4,
|
||||
"width": 128,
|
||||
"height": 128,
|
||||
"training_loop": {
|
||||
"num_epochs": 1,
|
||||
"batch_size": 1
|
||||
},
|
||||
"optimizer": {
|
||||
"name": "AdamW",
|
||||
"learning_rate": 1e-4
|
||||
},
|
||||
"save": {
|
||||
"output_path": "tests/dreambooth/tmp",
|
||||
"checkpoint_frequency": 1
|
||||
},
|
||||
"instrumentation": {
|
||||
"plot_frequency": 1,
|
||||
"generate_image_frequency": 50000,
|
||||
"validation_prompt": "photo of sks dog"
|
||||
},
|
||||
"lora_layers": {
|
||||
"transformer_blocks": {
|
||||
"block_range": {
|
||||
"start": 12,
|
||||
"end": 13
|
||||
},
|
||||
"layer_types": [
|
||||
"attn.to_q",
|
||||
"attn.to_k",
|
||||
"attn.to_v",
|
||||
"attn.add_q_proj",
|
||||
"attn.add_k_proj",
|
||||
"attn.add_v_proj",
|
||||
"attn.to_out",
|
||||
"attn.to_add_out",
|
||||
"ff.linear1",
|
||||
"ff.linear2",
|
||||
"ff_context.linear1",
|
||||
"ff_context.linear2"
|
||||
],
|
||||
"lora_rank": 4
|
||||
},
|
||||
"single_transformer_blocks": {
|
||||
"block_range": {
|
||||
"start": 17,
|
||||
"end": 18
|
||||
},
|
||||
"layer_types": [
|
||||
"proj_out",
|
||||
"proj_mlp",
|
||||
"attn.to_q",
|
||||
"attn.to_k",
|
||||
"attn.to_v"
|
||||
],
|
||||
"lora_rank": 4
|
||||
}
|
||||
},
|
||||
"examples": {
|
||||
"path": "../../../src/mflux/dreambooth/_example/images/",
|
||||
"images": [
|
||||
{
|
||||
"image": "01.jpeg",
|
||||
"prompt": "photo of sks dog"
|
||||
},
|
||||
{
|
||||
"image": "02.jpeg",
|
||||
"prompt": "photo of sks dog"
|
||||
},
|
||||
{
|
||||
"image": "03.jpeg",
|
||||
"prompt": "photo of sks dog"
|
||||
},
|
||||
{
|
||||
"image": "04.jpeg",
|
||||
"prompt": "photo of sks dog"
|
||||
},
|
||||
{
|
||||
"image": "05.jpeg",
|
||||
"prompt": "photo of sks dog"
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue
Block a user