diff --git a/.gitignore b/.gitignore index c518e73..64907ac 100644 --- a/.gitignore +++ b/.gitignore @@ -18,3 +18,8 @@ # build/ generated from 'pip install -e .' in developer mode build/ +.gz +.fdb_latexmk +.fls +.aux +.log diff --git a/README.md b/README.md index 533e340..462f93c 100644 --- a/README.md +++ b/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. + +![image](src/mflux/assets/dreambooth.jpg) + +*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 diff --git a/pyproject.toml b/pyproject.toml index ef61272..cb32f94 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"] diff --git a/src/mflux/assets/dreambooth.jpg b/src/mflux/assets/dreambooth.jpg new file mode 100644 index 0000000..310d827 Binary files /dev/null and b/src/mflux/assets/dreambooth.jpg differ diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index 1fc62b8..0ff1a47 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -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") diff --git a/src/mflux/controlnet/transformer_controlnet.py b/src/mflux/controlnet/transformer_controlnet.py index 50c2db3..6d9d4ad 100644 --- a/src/mflux/controlnet/transformer_controlnet.py +++ b/src/mflux/controlnet/transformer_controlnet.py @@ -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, diff --git a/src/mflux/controlnet/weight_handler_controlnet.py b/src/mflux/controlnet/weight_handler_controlnet.py index 00654c9..0cdfd70 100644 --- a/src/mflux/controlnet/weight_handler_controlnet.py +++ b/src/mflux/controlnet/weight_handler_controlnet.py @@ -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 diff --git a/tests/helpers/__init__.py b/src/mflux/dreambooth/__init__.py similarity index 100% rename from tests/helpers/__init__.py rename to src/mflux/dreambooth/__init__.py diff --git a/src/mflux/dreambooth/_example/images/01.jpeg b/src/mflux/dreambooth/_example/images/01.jpeg new file mode 100644 index 0000000..e839e5e Binary files /dev/null and b/src/mflux/dreambooth/_example/images/01.jpeg differ diff --git a/src/mflux/dreambooth/_example/images/02.jpeg b/src/mflux/dreambooth/_example/images/02.jpeg new file mode 100644 index 0000000..35dd2ab Binary files /dev/null and b/src/mflux/dreambooth/_example/images/02.jpeg differ diff --git a/src/mflux/dreambooth/_example/images/03.jpeg b/src/mflux/dreambooth/_example/images/03.jpeg new file mode 100644 index 0000000..cd760c6 Binary files /dev/null and b/src/mflux/dreambooth/_example/images/03.jpeg differ diff --git a/src/mflux/dreambooth/_example/images/04.jpeg b/src/mflux/dreambooth/_example/images/04.jpeg new file mode 100644 index 0000000..16ab824 Binary files /dev/null and b/src/mflux/dreambooth/_example/images/04.jpeg differ diff --git a/src/mflux/dreambooth/_example/images/05.jpeg b/src/mflux/dreambooth/_example/images/05.jpeg new file mode 100644 index 0000000..44d5967 Binary files /dev/null and b/src/mflux/dreambooth/_example/images/05.jpeg differ diff --git a/src/mflux/dreambooth/_example/train.json b/src/mflux/dreambooth/_example/train.json new file mode 100644 index 0000000..3259082 --- /dev/null +++ b/src/mflux/dreambooth/_example/train.json @@ -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" + } + ] + } +} \ No newline at end of file diff --git a/src/mflux/dreambooth/dataset/__init__.py b/src/mflux/dreambooth/dataset/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/dreambooth/dataset/batch.py b/src/mflux/dreambooth/dataset/batch.py new file mode 100644 index 0000000..39987ca --- /dev/null +++ b/src/mflux/dreambooth/dataset/batch.py @@ -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 diff --git a/src/mflux/dreambooth/dataset/dataset.py b/src/mflux/dreambooth/dataset/dataset.py new file mode 100644 index 0000000..5ef3f7e --- /dev/null +++ b/src/mflux/dreambooth/dataset/dataset.py @@ -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 diff --git a/src/mflux/dreambooth/dataset/dreambooth_preprocessing.py b/src/mflux/dreambooth/dataset/dreambooth_preprocessing.py new file mode 100644 index 0000000..30f62bd --- /dev/null +++ b/src/mflux/dreambooth/dataset/dreambooth_preprocessing.py @@ -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] diff --git a/src/mflux/dreambooth/dataset/iterator.py b/src/mflux/dreambooth/dataset/iterator.py new file mode 100644 index 0000000..772b045 --- /dev/null +++ b/src/mflux/dreambooth/dataset/iterator.py @@ -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)) diff --git a/src/mflux/dreambooth/dreambooth.py b/src/mflux/dreambooth/dreambooth.py new file mode 100644 index 0000000..cd45f3b --- /dev/null +++ b/src/mflux/dreambooth/dreambooth.py @@ -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) diff --git a/src/mflux/dreambooth/dreambooth_initializer.py b/src/mflux/dreambooth/dreambooth_initializer.py new file mode 100644 index 0000000..ab7ce10 --- /dev/null +++ b/src/mflux/dreambooth/dreambooth_initializer.py @@ -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 diff --git a/src/mflux/dreambooth/lora_layers/__init__.py b/src/mflux/dreambooth/lora_layers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/dreambooth/lora_layers/fused_linear_lora_layer.py b/src/mflux/dreambooth/lora_layers/fused_linear_lora_layer.py new file mode 100644 index 0000000..9a70e90 --- /dev/null +++ b/src/mflux/dreambooth/lora_layers/fused_linear_lora_layer.py @@ -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 diff --git a/src/mflux/dreambooth/lora_layers/linear_lora_layer.py b/src/mflux/dreambooth/lora_layers/linear_lora_layer.py new file mode 100644 index 0000000..05bcbb8 --- /dev/null +++ b/src/mflux/dreambooth/lora_layers/linear_lora_layer.py @@ -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 diff --git a/src/mflux/dreambooth/lora_layers/lora_layers.py b/src/mflux/dreambooth/lora_layers/lora_layers.py new file mode 100644 index 0000000..2923f77 --- /dev/null +++ b/src/mflux/dreambooth/lora_layers/lora_layers.py @@ -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), + }, + ) diff --git a/src/mflux/dreambooth/optimization/__init__.py b/src/mflux/dreambooth/optimization/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/dreambooth/optimization/_loss_derivation/dreambooth_loss.pdf b/src/mflux/dreambooth/optimization/_loss_derivation/dreambooth_loss.pdf new file mode 100644 index 0000000..adeaeab Binary files /dev/null and b/src/mflux/dreambooth/optimization/_loss_derivation/dreambooth_loss.pdf differ diff --git a/src/mflux/dreambooth/optimization/_loss_derivation/dreambooth_loss.tex b/src/mflux/dreambooth/optimization/_loss_derivation/dreambooth_loss.tex new file mode 100644 index 0000000..4d71fca --- /dev/null +++ b/src/mflux/dreambooth/optimization/_loss_derivation/dreambooth_loss.tex @@ -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} diff --git a/src/mflux/dreambooth/optimization/dreambooth_loss.py b/src/mflux/dreambooth/optimization/dreambooth_loss.py new file mode 100644 index 0000000..c985833 --- /dev/null +++ b/src/mflux/dreambooth/optimization/dreambooth_loss.py @@ -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() diff --git a/src/mflux/dreambooth/optimization/optimizer.py b/src/mflux/dreambooth/optimization/optimizer.py new file mode 100644 index 0000000..5cddeea --- /dev/null +++ b/src/mflux/dreambooth/optimization/optimizer.py @@ -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) diff --git a/src/mflux/dreambooth/state/__init__.py b/src/mflux/dreambooth/state/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/dreambooth/state/training_spec.py b/src/mflux/dreambooth/state/training_spec.py new file mode 100644 index 0000000..7438427 --- /dev/null +++ b/src/mflux/dreambooth/state/training_spec.py @@ -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 diff --git a/src/mflux/dreambooth/state/training_state.py b/src/mflux/dreambooth/state/training_state.py new file mode 100644 index 0000000..da2aca7 --- /dev/null +++ b/src/mflux/dreambooth/state/training_state.py @@ -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) diff --git a/src/mflux/dreambooth/state/zip_util.py b/src/mflux/dreambooth/state/zip_util.py new file mode 100644 index 0000000..9f5767a --- /dev/null +++ b/src/mflux/dreambooth/state/zip_util.py @@ -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}") diff --git a/src/mflux/dreambooth/statistics/__init__.py b/src/mflux/dreambooth/statistics/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/dreambooth/statistics/plotter.py b/src/mflux/dreambooth/statistics/plotter.py new file mode 100644 index 0000000..111d3c3 --- /dev/null +++ b/src/mflux/dreambooth/statistics/plotter.py @@ -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() diff --git a/src/mflux/dreambooth/statistics/statistics.py b/src/mflux/dreambooth/statistics/statistics.py new file mode 100644 index 0000000..9cfbce2 --- /dev/null +++ b/src/mflux/dreambooth/statistics/statistics.py @@ -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) diff --git a/src/mflux/error/exceptions.py b/src/mflux/error/exceptions.py index 7f28df6..e2693ac 100644 --- a/src/mflux/error/exceptions.py +++ b/src/mflux/error/exceptions.py @@ -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.""" diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index 0410e17..08e0de5 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -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() diff --git a/src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py b/src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py index d8b1b17..0410e13 100644 --- a/src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py +++ b/src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py @@ -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) diff --git a/src/mflux/models/text_encoder/clip_encoder/clip_encoder.py b/src/mflux/models/text_encoder/clip_encoder/clip_encoder.py index f1ac39b..9b641e3 100644 --- a/src/mflux/models/text_encoder/clip_encoder/clip_encoder.py +++ b/src/mflux/models/text_encoder/clip_encoder/clip_encoder.py @@ -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 diff --git a/src/mflux/models/text_encoder/clip_encoder/clip_encoder_layer.py b/src/mflux/models/text_encoder/clip_encoder/clip_encoder_layer.py index 8df7028..55e2a3f 100644 --- a/src/mflux/models/text_encoder/clip_encoder/clip_encoder_layer.py +++ b/src/mflux/models/text_encoder/clip_encoder/clip_encoder_layer.py @@ -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 diff --git a/src/mflux/models/text_encoder/clip_encoder/clip_mlp.py b/src/mflux/models/text_encoder/clip_encoder/clip_mlp.py index 7ae9dd1..fc40bdb 100644 --- a/src/mflux/models/text_encoder/clip_encoder/clip_mlp.py +++ b/src/mflux/models/text_encoder/clip_encoder/clip_mlp.py @@ -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) diff --git a/src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py b/src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py index d9eded6..e56a0b2 100644 --- a/src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py +++ b/src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py @@ -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) diff --git a/src/mflux/models/text_encoder/clip_encoder/clip_text_model.py b/src/mflux/models/text_encoder/clip_encoder/clip_text_model.py index fc9eef3..66bfe08 100644 --- a/src/mflux/models/text_encoder/clip_encoder/clip_text_model.py +++ b/src/mflux/models/text_encoder/clip_encoder/clip_text_model.py @@ -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 diff --git a/src/mflux/models/text_encoder/clip_encoder/encoder_clip.py b/src/mflux/models/text_encoder/clip_encoder/encoder_clip.py index 2ab9078..2199732 100644 --- a/src/mflux/models/text_encoder/clip_encoder/encoder_clip.py +++ b/src/mflux/models/text_encoder/clip_encoder/encoder_clip.py @@ -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 diff --git a/src/mflux/models/text_encoder/t5_encoder/t5_attention.py b/src/mflux/models/text_encoder/t5_encoder/t5_attention.py index 6258716..9280da5 100644 --- a/src/mflux/models/text_encoder/t5_encoder/t5_attention.py +++ b/src/mflux/models/text_encoder/t5_encoder/t5_attention.py @@ -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 diff --git a/src/mflux/models/text_encoder/t5_encoder/t5_block.py b/src/mflux/models/text_encoder/t5_encoder/t5_block.py index 8787f5b..64c789a 100644 --- a/src/mflux/models/text_encoder/t5_encoder/t5_block.py +++ b/src/mflux/models/text_encoder/t5_encoder/t5_block.py @@ -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 diff --git a/src/mflux/models/text_encoder/t5_encoder/t5_dense_relu_dense.py b/src/mflux/models/text_encoder/t5_encoder/t5_dense_relu_dense.py index a6af0ea..bc5d549 100644 --- a/src/mflux/models/text_encoder/t5_encoder/t5_dense_relu_dense.py +++ b/src/mflux/models/text_encoder/t5_encoder/t5_dense_relu_dense.py @@ -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 diff --git a/src/mflux/models/text_encoder/t5_encoder/t5_encoder.py b/src/mflux/models/text_encoder/t5_encoder/t5_encoder.py index 92f7741..b053fc1 100644 --- a/src/mflux/models/text_encoder/t5_encoder/t5_encoder.py +++ b/src/mflux/models/text_encoder/t5_encoder/t5_encoder.py @@ -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 diff --git a/src/mflux/models/text_encoder/t5_encoder/t5_feed_forward.py b/src/mflux/models/text_encoder/t5_encoder/t5_feed_forward.py index e5c52df..03d6f5b 100644 --- a/src/mflux/models/text_encoder/t5_encoder/t5_feed_forward.py +++ b/src/mflux/models/text_encoder/t5_encoder/t5_feed_forward.py @@ -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 diff --git a/src/mflux/models/text_encoder/t5_encoder/t5_layer_norm.py b/src/mflux/models/text_encoder/t5_encoder/t5_layer_norm.py index b1bef54..9dcc849 100644 --- a/src/mflux/models/text_encoder/t5_encoder/t5_layer_norm.py +++ b/src/mflux/models/text_encoder/t5_encoder/t5_layer_norm.py @@ -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, diff --git a/src/mflux/models/text_encoder/t5_encoder/t5_self_attention.py b/src/mflux/models/text_encoder/t5_encoder/t5_self_attention.py index d51708c..08fb66b 100644 --- a/src/mflux/models/text_encoder/t5_encoder/t5_self_attention.py +++ b/src/mflux/models/text_encoder/t5_encoder/t5_self_attention.py @@ -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)) diff --git a/src/mflux/models/transformer/ada_layer_norm_continuous.py b/src/mflux/models/transformer/ada_layer_norm_continuous.py index 94a07cc..801e54c 100644 --- a/src/mflux/models/transformer/ada_layer_norm_continuous.py +++ b/src/mflux/models/transformer/ada_layer_norm_continuous.py @@ -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] diff --git a/src/mflux/models/transformer/ada_layer_norm_zero.py b/src/mflux/models/transformer/ada_layer_norm_zero.py index deeba1c..98f6173 100644 --- a/src/mflux/models/transformer/ada_layer_norm_zero.py +++ b/src/mflux/models/transformer/ada_layer_norm_zero.py @@ -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] diff --git a/src/mflux/models/transformer/ada_layer_norm_zero_single.py b/src/mflux/models/transformer/ada_layer_norm_zero_single.py index 655a324..425543f 100644 --- a/src/mflux/models/transformer/ada_layer_norm_zero_single.py +++ b/src/mflux/models/transformer/ada_layer_norm_zero_single.py @@ -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] diff --git a/src/mflux/models/transformer/embed_nd.py b/src/mflux/models/transformer/embed_nd.py index 4d92dce..71605d0 100644 --- a/src/mflux/models/transformer/embed_nd.py +++ b/src/mflux/models/transformer/embed_nd.py @@ -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, diff --git a/src/mflux/models/transformer/feed_forward.py b/src/mflux/models/transformer/feed_forward.py index 6124fe7..5b962ea 100644 --- a/src/mflux/models/transformer/feed_forward.py +++ b/src/mflux/models/transformer/feed_forward.py @@ -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) diff --git a/src/mflux/models/transformer/guidance_embedder.py b/src/mflux/models/transformer/guidance_embedder.py index ed9d1c4..76ab6b9 100644 --- a/src/mflux/models/transformer/guidance_embedder.py +++ b/src/mflux/models/transformer/guidance_embedder.py @@ -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) diff --git a/src/mflux/models/transformer/joint_attention.py b/src/mflux/models/transformer/joint_attention.py index 6ea3f62..2dd991a 100644 --- a/src/mflux/models/transformer/joint_attention.py +++ b/src/mflux/models/transformer/joint_attention.py @@ -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, diff --git a/src/mflux/models/transformer/joint_transformer_block.py b/src/mflux/models/transformer/joint_transformer_block.py index e347761..b89b4e6 100644 --- a/src/mflux/models/transformer/joint_transformer_block.py +++ b/src/mflux/models/transformer/joint_transformer_block.py @@ -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 diff --git a/src/mflux/models/transformer/single_block_attention.py b/src/mflux/models/transformer/single_block_attention.py index 7f806a3..8b9aec4 100644 --- a/src/mflux/models/transformer/single_block_attention.py +++ b/src/mflux/models/transformer/single_block_attention.py @@ -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) diff --git a/src/mflux/models/transformer/single_transformer_block.py b/src/mflux/models/transformer/single_transformer_block.py index 31cdef6..1bab2a4 100644 --- a/src/mflux/models/transformer/single_transformer_block.py +++ b/src/mflux/models/transformer/single_transformer_block.py @@ -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, ) diff --git a/src/mflux/models/transformer/text_embedder.py b/src/mflux/models/transformer/text_embedder.py index 6034b74..d5696be 100644 --- a/src/mflux/models/transformer/text_embedder.py +++ b/src/mflux/models/transformer/text_embedder.py @@ -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) diff --git a/src/mflux/models/transformer/time_text_embed.py b/src/mflux/models/transformer/time_text_embed.py index 5dbd84e..3492a83 100644 --- a/src/mflux/models/transformer/time_text_embed.py +++ b/src/mflux/models/transformer/time_text_embed.py @@ -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) diff --git a/src/mflux/models/transformer/timestep_embedder.py b/src/mflux/models/transformer/timestep_embedder.py index b72aea2..ccfa0e1 100644 --- a/src/mflux/models/transformer/timestep_embedder.py +++ b/src/mflux/models/transformer/timestep_embedder.py @@ -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) diff --git a/src/mflux/models/transformer/transformer.py b/src/mflux/models/transformer/transformer.py index 9922cce..dbabac8 100644 --- a/src/mflux/models/transformer/transformer.py +++ b/src/mflux/models/transformer/transformer.py @@ -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 diff --git a/src/mflux/models/vae/common/attention.py b/src/mflux/models/vae/common/attention.py index f1ba8b3..a23059c 100644 --- a/src/mflux/models/vae/common/attention.py +++ b/src/mflux/models/vae/common/attention.py @@ -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 diff --git a/src/mflux/models/vae/common/resnet_block_2d.py b/src/mflux/models/vae/common/resnet_block_2d.py index 2c60335..f235944 100644 --- a/src/mflux/models/vae/common/resnet_block_2d.py +++ b/src/mflux/models/vae/common/resnet_block_2d.py @@ -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) diff --git a/src/mflux/models/vae/common/unet_mid_block.py b/src/mflux/models/vae/common/unet_mid_block.py index bb1a423..f6af721 100644 --- a/src/mflux/models/vae/common/unet_mid_block.py +++ b/src/mflux/models/vae/common/unet_mid_block.py @@ -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 diff --git a/src/mflux/models/vae/decoder/conv_in.py b/src/mflux/models/vae/decoder/conv_in.py index 5c0eda9..f481bdc 100644 --- a/src/mflux/models/vae/decoder/conv_in.py +++ b/src/mflux/models/vae/decoder/conv_in.py @@ -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)) diff --git a/src/mflux/models/vae/decoder/conv_norm_out.py b/src/mflux/models/vae/decoder/conv_norm_out.py index bfdcb17..05e5fe8 100644 --- a/src/mflux/models/vae/decoder/conv_norm_out.py +++ b/src/mflux/models/vae/decoder/conv_norm_out.py @@ -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)) diff --git a/src/mflux/models/vae/decoder/conv_out.py b/src/mflux/models/vae/decoder/conv_out.py index ad6a694..5e75f1d 100644 --- a/src/mflux/models/vae/decoder/conv_out.py +++ b/src/mflux/models/vae/decoder/conv_out.py @@ -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)) diff --git a/src/mflux/models/vae/decoder/decoder.py b/src/mflux/models/vae/decoder/decoder.py index 86a50cd..1196dd0 100644 --- a/src/mflux/models/vae/decoder/decoder.py +++ b/src/mflux/models/vae/decoder/decoder.py @@ -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 diff --git a/src/mflux/models/vae/decoder/up_block_1_or_2.py b/src/mflux/models/vae/decoder/up_block_1_or_2.py index 5763e48..9de27af 100644 --- a/src/mflux/models/vae/decoder/up_block_1_or_2.py +++ b/src/mflux/models/vae/decoder/up_block_1_or_2.py @@ -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 diff --git a/src/mflux/models/vae/decoder/up_block_3.py b/src/mflux/models/vae/decoder/up_block_3.py index 5959b94..c3f8a46 100644 --- a/src/mflux/models/vae/decoder/up_block_3.py +++ b/src/mflux/models/vae/decoder/up_block_3.py @@ -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 diff --git a/src/mflux/models/vae/decoder/up_block_4.py b/src/mflux/models/vae/decoder/up_block_4.py index dd78356..b69235c 100644 --- a/src/mflux/models/vae/decoder/up_block_4.py +++ b/src/mflux/models/vae/decoder/up_block_4.py @@ -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 diff --git a/src/mflux/models/vae/decoder/up_sampler.py b/src/mflux/models/vae/decoder/up_sampler.py index 2312f45..5e5be36 100644 --- a/src/mflux/models/vae/decoder/up_sampler.py +++ b/src/mflux/models/vae/decoder/up_sampler.py @@ -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) diff --git a/src/mflux/models/vae/encoder/conv_in.py b/src/mflux/models/vae/encoder/conv_in.py index a999821..56119c5 100644 --- a/src/mflux/models/vae/encoder/conv_in.py +++ b/src/mflux/models/vae/encoder/conv_in.py @@ -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)) diff --git a/src/mflux/models/vae/encoder/conv_norm_out.py b/src/mflux/models/vae/encoder/conv_norm_out.py index 693a8b8..64269a8 100644 --- a/src/mflux/models/vae/encoder/conv_norm_out.py +++ b/src/mflux/models/vae/encoder/conv_norm_out.py @@ -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)) diff --git a/src/mflux/models/vae/encoder/conv_out.py b/src/mflux/models/vae/encoder/conv_out.py index f008208..63a7e4a 100644 --- a/src/mflux/models/vae/encoder/conv_out.py +++ b/src/mflux/models/vae/encoder/conv_out.py @@ -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)) diff --git a/src/mflux/models/vae/encoder/down_block_1.py b/src/mflux/models/vae/encoder/down_block_1.py index ce5116b..96b1ee8 100644 --- a/src/mflux/models/vae/encoder/down_block_1.py +++ b/src/mflux/models/vae/encoder/down_block_1.py @@ -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 diff --git a/src/mflux/models/vae/encoder/down_block_2.py b/src/mflux/models/vae/encoder/down_block_2.py index afa05a5..cf258e1 100644 --- a/src/mflux/models/vae/encoder/down_block_2.py +++ b/src/mflux/models/vae/encoder/down_block_2.py @@ -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 diff --git a/src/mflux/models/vae/encoder/down_block_3.py b/src/mflux/models/vae/encoder/down_block_3.py index c1cac99..f94a5f9 100644 --- a/src/mflux/models/vae/encoder/down_block_3.py +++ b/src/mflux/models/vae/encoder/down_block_3.py @@ -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 diff --git a/src/mflux/models/vae/encoder/down_block_4.py b/src/mflux/models/vae/encoder/down_block_4.py index 45c917a..7ffd825 100644 --- a/src/mflux/models/vae/encoder/down_block_4.py +++ b/src/mflux/models/vae/encoder/down_block_4.py @@ -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 diff --git a/src/mflux/models/vae/encoder/down_sampler.py b/src/mflux/models/vae/encoder/down_sampler.py index 856ecc5..53f364d 100644 --- a/src/mflux/models/vae/encoder/down_sampler.py +++ b/src/mflux/models/vae/encoder/down_sampler.py @@ -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) diff --git a/src/mflux/models/vae/encoder/encoder.py b/src/mflux/models/vae/encoder/encoder.py index e73d490..c3b3859 100644 --- a/src/mflux/models/vae/encoder/encoder.py +++ b/src/mflux/models/vae/encoder/encoder.py @@ -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 diff --git a/src/mflux/models/vae/vae.py b/src/mflux/models/vae/vae.py index 280afbd..9ac7af5 100644 --- a/src/mflux/models/vae/vae.py +++ b/src/mflux/models/vae/vae.py @@ -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 diff --git a/src/mflux/post_processing/generated_image.py b/src/mflux/post_processing/generated_image.py index ca8c21e..893da42 100644 --- a/src/mflux/post_processing/generated_image.py +++ b/src/mflux/post_processing/generated_image.py @@ -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 diff --git a/src/mflux/train.py b/src/mflux/train.py new file mode 100644 index 0000000..1d035c1 --- /dev/null +++ b/src/mflux/train.py @@ -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() diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index c6da341..215a4ed 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/src/mflux/ui/cli/parsers.py @@ -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: diff --git a/src/mflux/weights/lora_util.py b/src/mflux/weights/lora_util.py deleted file mode 100644 index 9145914..0000000 --- a/src/mflux/weights/lora_util.py +++ /dev/null @@ -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 diff --git a/src/mflux/weights/quantization_util.py b/src/mflux/weights/quantization_util.py new file mode 100644 index 0000000..3283e30 --- /dev/null +++ b/src/mflux/weights/quantization_util.py @@ -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 diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py index 1db8f95..7fb0470 100644 --- a/src/mflux/weights/weight_handler.py +++ b/src/mflux/weights/weight_handler.py @@ -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: diff --git a/src/mflux/weights/weight_handler_lora.py b/src/mflux/weights/weight_handler_lora.py new file mode 100644 index 0000000..22edb7d --- /dev/null +++ b/src/mflux/weights/weight_handler_lora.py @@ -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 diff --git a/src/mflux/weights/weight_util.py b/src/mflux/weights/weight_util.py index 3cd0625..8d603a7 100644 --- a/src/mflux/weights/weight_util.py +++ b/src/mflux/weights/weight_util.py @@ -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) diff --git a/tests/arg_parser/__init__.py b/tests/arg_parser/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_cli_argparser.py b/tests/arg_parser/test_cli_argparser.py similarity index 100% rename from tests/test_cli_argparser.py rename to tests/arg_parser/test_cli_argparser.py diff --git a/tests/dreambooth/__init__.py b/tests/dreambooth/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/dreambooth/config/train.json b/tests/dreambooth/config/train.json new file mode 100644 index 0000000..1dcbbaf --- /dev/null +++ b/tests/dreambooth/config/train.json @@ -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" + } + ] + } +} \ No newline at end of file diff --git a/tests/dreambooth/test_dreambooth_dataset_iterator.py b/tests/dreambooth/test_dreambooth_dataset_iterator.py new file mode 100644 index 0000000..82b6fa0 --- /dev/null +++ b/tests/dreambooth/test_dreambooth_dataset_iterator.py @@ -0,0 +1,255 @@ +from pathlib import Path +from typing import List, Set + +import mlx.core as mx +import pytest + +from mflux.dreambooth.dataset.batch import Example +from mflux.dreambooth.dataset.dataset import Dataset +from mflux.dreambooth.dataset.iterator import Iterator + + +def get_dataset(num_examples: int) -> Dataset: + examples = [ + Example( + example_id=i, + prompt=f"prompt_{i}", + image_path=Path(f"image_{i}.jpg"), + encoded_image=mx.zeros((1,)), + prompt_embeds=mx.zeros((1,)), + pooled_prompt_embeds=mx.zeros((1,)), + ) + for i in range(num_examples) + ] + return Dataset(examples=examples) + + +@pytest.fixture(params=[4, 5, 8, 25, 103]) +def dataset(request): + return get_dataset(request.param) + + +@pytest.fixture(params=[1, 2, 3]) +def batch_size(request): + return request.param + + +def test_batch_size_consistency(dataset, batch_size): + """Test that batches are of the specified size except for the last one.""" + iterator = Iterator(dataset, batch_size) + + batch_sizes = [] + example_ids = set() + + for batch in iterator: + batch_sizes.append(len(batch.examples)) + for example in batch.examples: + example_ids.add(example.example_id) + + if len(example_ids) == dataset.size(): + break + + # Check all batches except the last are of specified size + assert all(size == batch_size for size in batch_sizes[:-1]) + # Check last batch handles remainder + assert batch_sizes[-1] == dataset.size() % batch_size or batch_size + + +def test_complete_coverage(dataset, batch_size): + """Test that all examples are seen exactly once before reset.""" + iterator = Iterator(dataset, batch_size) + seen_examples: Set[int] = set() + + for batch in iterator: + for example in batch.examples: + seen_examples.add(example.example_id) + if len(seen_examples) == dataset.size(): + break + + expected_examples = set(range(dataset.size())) + assert seen_examples == expected_examples + + +def test_state_restoration(dataset, batch_size): + """Test that iterator can be saved and restored to the same state.""" + iterator1 = Iterator(dataset, batch_size) + + first_batches = [] + num_batches_to_compare = 3 + for _ in range(num_batches_to_compare): + batch = next(iterator1) + first_batches.append([e.example_id for e in batch.examples]) + + state_dict = iterator1.to_dict() + iterator2 = Iterator.from_dict(state_dict, dataset) + + for _ in range(num_batches_to_compare): + batch1 = next(iterator1) + batch2 = next(iterator2) + ids1 = [e.example_id for e in batch1.examples] + ids2 = [e.example_id for e in batch2.examples] + assert ids1 == ids2 + + +def test_state_restoration_across_epochs(dataset, batch_size): + """Test that state restoration works when crossing epoch boundaries.""" + iterator1 = Iterator(dataset, batch_size) + + num_iterations = (dataset.size() + batch_size - 1) // batch_size + 1 + for _ in range(num_iterations): + next(iterator1) + + state_dict = iterator1.to_dict() + iterator2 = Iterator.from_dict(state_dict, dataset) + + batch1 = next(iterator1) + batch2 = next(iterator2) + ids1 = [e.example_id for e in batch1.examples] + ids2 = [e.example_id for e in batch2.examples] + assert ids1 == ids2 + + +def test_randomization(dataset, batch_size): + """Test that different seeds produce different sequences.""" + iterator1 = Iterator(dataset, batch_size=2, seed=1) + iterator2 = Iterator(dataset, batch_size=2, seed=2) + + batch1 = next(iterator1) + batch2 = next(iterator2) + + ids1 = [e.example_id for e in batch1.examples] + ids2 = [e.example_id for e in batch2.examples] + + assert ids1 != ids2 + + +def test_epoch_transition(dataset, batch_size): + """Test that epoch transition creates a new permutation.""" + iterator = Iterator(dataset, batch_size) + + num_batches_in_epoch = (dataset.size() + batch_size - 1) // batch_size + first_epoch_examples: List[List[int]] = [] + for _ in range(num_batches_in_epoch): + batch = next(iterator) + first_epoch_examples.append([e.example_id for e in batch.examples]) + + second_epoch_batch = next(iterator) + second_epoch_ids = [e.example_id for e in second_epoch_batch.examples] + + assert all(0 <= id < dataset.size() for id in second_epoch_ids) + assert len(second_epoch_ids) == min(batch_size, dataset.size()) + + +def test_state_consistency(dataset, batch_size): + """Test that state remains consistent across iterations within the same epoch.""" + iterator = Iterator(dataset, batch_size) + + initial_state = iterator.to_dict() + next(iterator) + new_state = iterator.to_dict() + + assert initial_state != new_state + + iterator2 = Iterator.from_dict(initial_state, dataset) + next(iterator2) + restored_state = iterator2.to_dict() + + assert restored_state["position"] == new_state["position"] + assert restored_state["epoch"] == new_state["epoch"] + assert restored_state["num_iterations"] == new_state["num_iterations"] + assert restored_state["current_permutation"] == new_state["current_permutation"] + + +def test_fixed_num_epochs(dataset, batch_size): + """Test that iterator stops after specified number of epochs.""" + num_epochs = 100 + iterator = Iterator(dataset, batch_size, num_epochs=num_epochs) + + total_examples_seen = 0 + epochs_completed = 0 + examples_per_epoch = set() + + try: + while True: + batch = next(iterator) + batch_examples = {e.example_id for e in batch.examples} + total_examples_seen += len(batch.examples) + + examples_per_epoch.update(batch_examples) + + if len(examples_per_epoch) == dataset.size(): + epochs_completed += 1 + examples_per_epoch = set() + + except StopIteration: + pass + + assert epochs_completed == num_epochs + assert total_examples_seen == dataset.size() * num_epochs + + +def test_iteration_counting(dataset, batch_size): + """Test that iteration counting works correctly.""" + iterator = Iterator(dataset, batch_size) + + num_iterations = 5 + for _ in range(num_iterations): + next(iterator) + + assert iterator.num_iterations == num_iterations + + +def test_iteration_counting_across_epochs(dataset, batch_size): + """Test that iteration counting works correctly across epoch boundaries.""" + iterator = Iterator(dataset, batch_size) + + batches_per_epoch = (dataset.size() + batch_size - 1) // batch_size + total_iterations = int(batches_per_epoch * 1.5) + + for _ in range(total_iterations): + next(iterator) + + assert iterator.num_iterations == total_iterations + + +def test_random_seed_consistency(dataset, batch_size): + """Test that the same seed produces the same sequence of batches.""" + seed = 42 + iterator1 = Iterator(dataset, batch_size, seed=seed) + iterator2 = Iterator(dataset, batch_size, seed=seed) + + for _ in range(5): + batch1 = next(iterator1) + batch2 = next(iterator2) + + ids1 = [e.example_id for e in batch1.examples] + ids2 = [e.example_id for e in batch2.examples] + assert ids1 == ids2 + + +def test_state_serialization(dataset, batch_size): + """Test that the iterator state can be serialized and deserialized.""" + iterator = Iterator(dataset, batch_size) + next(iterator) + + state_dict = iterator.to_dict() + restored_iterator = Iterator.from_dict(state_dict, dataset) + + batch = next(restored_iterator) + assert len(batch.examples) > 0 + + +def test_state_restoration_after_exhaustion(dataset, batch_size): + """Test that state restoration works correctly after iterator exhaustion.""" + iterator = Iterator(dataset, batch_size, num_epochs=10) + try: + while True: + next(iterator) + except StopIteration: + pass + + state_dict = iterator.to_dict() + restored_iterator = Iterator.from_dict(state_dict, dataset) + + with pytest.raises(StopIteration): + next(restored_iterator) diff --git a/tests/dreambooth/test_resume_training.py b/tests/dreambooth/test_resume_training.py new file mode 100644 index 0000000..423119b --- /dev/null +++ b/tests/dreambooth/test_resume_training.py @@ -0,0 +1,87 @@ +import os +import shutil + +import mlx.core as mx + +from mflux.dreambooth.dreambooth import DreamBooth +from mflux.dreambooth.dreambooth_initializer import DreamBoothInitializer +from mflux.dreambooth.state.zip_util import ZipUtil + +CHECKPOINT_3 = "tests/dreambooth/tmp/_checkpoints/0000003_checkpoint.zip" +CHECKPOINT_4 = "tests/dreambooth/tmp/_checkpoints/0000004_checkpoint.zip" +CHECKPOINT_5 = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint.zip" + + +class TestResumeTraining: + def test_resume_training(self): + try: + # Given: A small training run from scratch for 5 steps (as described in the config)... + fluxA, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize( + config_path="tests/dreambooth/config/train.json", + checkpoint_path=None + ) # fmt:off + DreamBooth.train( + flux=fluxA, + runtime_config=runtime_config, + training_spec=training_spec, + training_state=training_state + ) # fmt: off + del fluxA, runtime_config, training_spec, training_state + # ...where we can inspect the training state after 5 runs... + adapter_after_5_steps = ZipUtil.unzip( + zip_path=CHECKPOINT_5, + filename="0000005_adapter.safetensors", + loader=lambda x: dict(mx.load(x).items()), + ) + # ...and deleting the outputs past step 3 to be sure we regenerate them after this... + TestResumeTraining.delete_file(CHECKPOINT_4) + TestResumeTraining.delete_file(CHECKPOINT_5) + + # When: Resuming the training from step 3... + fluxB, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize( + config_path=None, + checkpoint_path=CHECKPOINT_3 + ) # fmt:off + DreamBooth.train( + flux=fluxB, + runtime_config=runtime_config, + training_spec=training_spec, + training_state=training_state + ) # fmt: off + del fluxB, runtime_config, training_spec, training_state + # ...where we can inspect the training state after 2 additional runs... + adapter_after_5_steps_resumed = ZipUtil.unzip( + zip_path=CHECKPOINT_5, + filename="0000005_adapter.safetensors", + loader=lambda x: dict(mx.load(x).items()), + ) + + # Then: We want to confirm that the weights *exactly* match + assert len(adapter_after_5_steps.keys()) == len(adapter_after_5_steps_resumed.keys()) + for key in adapter_after_5_steps.keys(): + if key in adapter_after_5_steps_resumed: + array1 = adapter_after_5_steps[key] + array2 = adapter_after_5_steps_resumed[key] + assert mx.array_equal(array1, array2) + finally: + # cleanup + TestResumeTraining.delete_folder("tests/dreambooth/tmp") + + @staticmethod + def delete_file(path: str) -> None: + try: + if os.path.exists(path): + os.remove(path) + print(f"Deleted ZIP file: {path}") + else: + print("The specified ZIP file does not exist.") + except FileNotFoundError: + print(f"File not found: {path}") + except PermissionError: + print(f"Permission denied: {path}") + except OSError as e: + print(f"OS error while deleting file: {e}") + + @staticmethod + def delete_folder(path: str) -> None: + return shutil.rmtree(path) diff --git a/tests/dreambooth/test_train_and_load_weights.py b/tests/dreambooth/test_train_and_load_weights.py new file mode 100644 index 0000000..10928eb --- /dev/null +++ b/tests/dreambooth/test_train_and_load_weights.py @@ -0,0 +1,74 @@ +import shutil + +import numpy as np + +from mflux import Config, Flux1, ModelConfig +from mflux.dreambooth.dreambooth import DreamBooth +from mflux.dreambooth.dreambooth_initializer import DreamBoothInitializer +from mflux.dreambooth.state.zip_util import ZipUtil + +CHECKPOINT = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint.zip" +OUTPUT_DIR = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint" +LORA_FILE = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint/0000005_adapter.safetensors" + + +class TestTrainAndLoadWeights: + def test_train_and_load_weights(self): + try: + # Given: A small training run from scratch for 5 steps (as described in the config)... + fluxA, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize( + config_path="tests/dreambooth/config/train.json", + checkpoint_path=None + ) # fmt:off + DreamBooth.train( + flux=fluxA, + runtime_config=runtime_config, + training_spec=training_spec, + training_state=training_state + ) # fmt: off + # ...we generate an image with the flux instance with the trained weights + image1 = fluxA.generate_image( + seed=42, + prompt="test", + config=Config( + num_inference_steps=20, + height=128, + width=128, + ), + ) + del fluxA, runtime_config, training_spec, training_state + # unzip so that LoRA adapter can be read later... + ZipUtil.extract_all(zip_path=CHECKPOINT, output_dir=OUTPUT_DIR) + + # When: Loading a new Flux instance with the trained LoRA... + fluxB = Flux1( + model_config=ModelConfig.FLUX1_DEV, + quantize=4, + lora_paths=[LORA_FILE], + lora_scales=[1.0], + ) # fmt: off + + # ...and generating the same image from that + image2 = fluxB.generate_image( + seed=42, + prompt="test", + config=Config( + num_inference_steps=20, + height=128, + width=128, + ), + ) + + # Then: We want to confirm that the images *exactly* match + np.testing.assert_array_equal( + np.array(image1.image), + np.array(image2.image), + err_msg="Generated image doesn't match reference image.", + ) + finally: + # cleanup + TestTrainAndLoadWeights.delete_folder("tests/dreambooth/tmp") + + @staticmethod + def delete_folder(path: str) -> None: + return shutil.rmtree(path) diff --git a/tests/image_generation/__init__.py b/tests/image_generation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/image_generation/helpers/__init__.py b/tests/image_generation/helpers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/helpers/image_generation_controlnet_test_helper.py b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py similarity index 95% rename from tests/helpers/image_generation_controlnet_test_helper.py rename to tests/image_generation/helpers/image_generation_controlnet_test_helper.py index c5f3b45..9a47cd3 100644 --- a/tests/helpers/image_generation_controlnet_test_helper.py +++ b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py @@ -4,7 +4,7 @@ import numpy as np from PIL import Image from mflux import ConfigControlnet, Flux1Controlnet, ModelConfig -from tests.helpers.image_generation_test_helper import ImageGeneratorTestHelper +from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper class ImageGeneratorControlnetTestHelper: diff --git a/tests/helpers/image_generation_test_helper.py b/tests/image_generation/helpers/image_generation_test_helper.py similarity index 96% rename from tests/helpers/image_generation_test_helper.py rename to tests/image_generation/helpers/image_generation_test_helper.py index d1d8929..0ef548c 100644 --- a/tests/helpers/image_generation_test_helper.py +++ b/tests/image_generation/helpers/image_generation_test_helper.py @@ -66,4 +66,4 @@ class ImageGeneratorTestHelper: def resolve_path(path) -> Path | None: if path is None: return None - return Path(__file__).parent.parent / "resources" / path + return Path(__file__).parent.parent.parent / "resources" / path diff --git a/tests/test_generate_image.py b/tests/image_generation/test_generate_image.py similarity index 73% rename from tests/test_generate_image.py rename to tests/image_generation/test_generate_image.py index d234361..9f18b06 100644 --- a/tests/test_generate_image.py +++ b/tests/image_generation/test_generate_image.py @@ -1,5 +1,5 @@ from mflux import ModelConfig -from tests.helpers.image_generation_test_helper import ImageGeneratorTestHelper +from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper class TestImageGenerator: @@ -43,6 +43,20 @@ class TestImageGenerator: lora_scales=[1.0], ) + def test_image_generation_dev_multiple_loras(self): + ImageGeneratorTestHelper.assert_matches_reference_image( + reference_image_path="reference_dev_lora_multiple.png", + output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME, + model_config=ModelConfig.FLUX1_DEV, + steps=15, + seed=42, + height=341, + width=768, + prompt="Renaissance painting, mkym this is made of wool, burger", + lora_paths=["FLUX-dev-lora-MiaoKa-Yarn-World.safetensors", "Flux_-_Renaissance_art_style.safetensors"], + lora_scales=[0.4, 0.6], + ) + def test_image_generation_dev_image_to_image(self): ImageGeneratorTestHelper.assert_matches_reference_image( reference_image_path="reference_dev_image_to_image_result.png", diff --git a/tests/test_generate_image_controlnet.py b/tests/image_generation/test_generate_image_controlnet.py similarity index 94% rename from tests/test_generate_image_controlnet.py rename to tests/image_generation/test_generate_image_controlnet.py index 35b1f84..cfc52d9 100644 --- a/tests/test_generate_image_controlnet.py +++ b/tests/image_generation/test_generate_image_controlnet.py @@ -1,5 +1,5 @@ from mflux import ModelConfig -from tests.helpers.image_generation_controlnet_test_helper import ImageGeneratorControlnetTestHelper +from tests.image_generation.helpers.image_generation_controlnet_test_helper import ImageGeneratorControlnetTestHelper class TestImageGeneratorControlnet: diff --git a/tests/model_saving/__init__.py b/tests/model_saving/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/model_saving/test_model_saving.py b/tests/model_saving/test_model_saving.py new file mode 100644 index 0000000..9d2fc08 --- /dev/null +++ b/tests/model_saving/test_model_saving.py @@ -0,0 +1,58 @@ +import shutil + +import numpy as np + +from mflux import Config, Flux1, ModelConfig + +PATH = "tests/4bit/" + + +class TestModelSaving: + def test_save_and_load_4bit_model(self): + try: + # given a saved quantized model (and an image from that model) + fluxA = Flux1( + model_config=ModelConfig.FLUX1_DEV, + quantize=4, + ) + image1 = fluxA.generate_image( + seed=42, + prompt="Luxury food photograph", + config=Config( + num_inference_steps=15, + height=341, + width=768, + ), + ) + fluxA.save_model(PATH) + del fluxA + + # when loading the quantized model (also without specifying bits) + fluxB = Flux1( + model_config=ModelConfig.FLUX1_DEV, + local_path=PATH, + ) + + # then we can load the model and generate the identical image + image2 = fluxB.generate_image( + seed=42, + prompt="Luxury food photograph", + config=Config( + num_inference_steps=15, + height=341, + width=768, + ), + ) + np.testing.assert_array_equal( + np.array(image1.image), + np.array(image2.image), + err_msg="image2 doesn't match image1.", + ) + + finally: + # cleanup + TestModelSaving.delete_folder(PATH) + + @staticmethod + def delete_folder(path: str) -> None: + return shutil.rmtree(path) diff --git a/tests/resources/Flux_-_Renaissance_art_style.safetensors b/tests/resources/Flux_-_Renaissance_art_style.safetensors new file mode 100644 index 0000000..af8a24b Binary files /dev/null and b/tests/resources/Flux_-_Renaissance_art_style.safetensors differ diff --git a/tests/resources/reference_controlnet_dev_lora.png b/tests/resources/reference_controlnet_dev_lora.png index 086c4dc..d6c2ffb 100644 Binary files a/tests/resources/reference_controlnet_dev_lora.png and b/tests/resources/reference_controlnet_dev_lora.png differ diff --git a/tests/resources/reference_dev_image_to_image_result.png b/tests/resources/reference_dev_image_to_image_result.png index 0e9bde9..5adfe09 100644 Binary files a/tests/resources/reference_dev_image_to_image_result.png and b/tests/resources/reference_dev_image_to_image_result.png differ diff --git a/tests/resources/reference_dev_lora.png b/tests/resources/reference_dev_lora.png index e9255bd..f52cb8f 100644 Binary files a/tests/resources/reference_dev_lora.png and b/tests/resources/reference_dev_lora.png differ diff --git a/tests/resources/reference_dev_lora_multiple.png b/tests/resources/reference_dev_lora_multiple.png new file mode 100644 index 0000000..9229ae5 Binary files /dev/null and b/tests/resources/reference_dev_lora_multiple.png differ