From 77e5a2ef39d97d4c8382247223db79f97677c203 Mon Sep 17 00:00:00 2001 From: Anthony Wu <462072+anthonywu@users.noreply.github.com> Date: Fri, 20 Sep 2024 10:26:10 -0700 Subject: [PATCH 1/5] intro ruff config --- pyproject.toml | 47 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 6e9e1d6..5d79359 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,3 +46,50 @@ mflux-generate-controlnet = "mflux.generate_controlnet:main" [tool.setuptools.packages.find] where = ["src"] include = ["mflux*"] + +[tool.ruff] +line-length = 88 # same as 'black' default +indent-width = 4 +target-version = "py310" +respect-gitignore = true + +[tool.ruff.lint] +# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`) codes by default. +# Unlike Flake8, Ruff doesn't enable pycodestyle warnings (`W`) or +# McCabe complexity (`C901`) by default. +select = ["E4", "E7", "E9", "F"] +ignore = [] + +# Allow fix for all enabled rules (when `--fix`) is provided. +fixable = ["ALL"] +unfixable = [] + +# Allow unused variables when underscore-prefixed. +dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$" + +[tool.ruff.format] +# Like Black, use double quotes for strings. +quote-style = "double" + +# Like Black, indent with spaces, rather than tabs. +indent-style = "space" + +# Like Black, respect magic trailing commas. +skip-magic-trailing-comma = false + +# Like Black, automatically detect the appropriate line ending. +line-ending = "auto" + +# Enable auto-formatting of code examples in docstrings. Markdown, +# reStructuredText code/literal blocks and doctests are all supported. +# +# This is currently disabled by default, but it is planned for this +# to be opt-out in the future. +docstring-code-format = false + +# Set the line length limit used when formatting code snippets in +# docstrings. +# +# This only has an effect when the `docstring-code-format` setting is +# enabled. +docstring-code-line-length = "dynamic" From cb785e49f182e437060688089a0858e6944a77f9 Mon Sep 17 00:00:00 2001 From: Anthony Wu <462072+anthonywu@users.noreply.github.com> Date: Fri, 20 Sep 2024 19:06:49 -0700 Subject: [PATCH 2/5] introduce but not enforce .pre-commit-config.yaml --- .pre-commit-config.yaml | 12 ++++++++++++ 1 file changed, 12 insertions(+) create mode 100644 .pre-commit-config.yaml diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..df981c5 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,12 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + # Ruff version. + rev: v0.6.6 + hooks: + # Run the linter. + - id: ruff + types_or: [python, pyi] + args: [--fix] + # Run the formatter. + - id: ruff-format + types_or: [python, pyi] From 9e6ace4cc385aeb4610cb7b82621e918732dae94 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 21 Sep 2024 15:39:52 +0200 Subject: [PATCH 3/5] Apply ruff formating (minor manual modifications to some files which are subsequently excluded) --- src/mflux/__init__.py | 9 ++- src/mflux/config/config.py | 22 +++--- src/mflux/config/model_config.py | 10 +-- src/mflux/config/runtime_config.py | 3 +- src/mflux/controlnet/controlnet_util.py | 1 - src/mflux/controlnet/flux_controlnet.py | 20 ++--- .../controlnet/transformer_controlnet.py | 37 +++++----- src/mflux/flux/flux.py | 15 ++-- src/mflux/generate.py | 34 ++++----- src/mflux/generate_controlnet.py | 40 +++++----- .../clip_encoder/clip_embeddings.py | 1 - .../text_encoder/clip_encoder/clip_encoder.py | 1 - .../clip_encoder/clip_encoder_layer.py | 1 - .../text_encoder/clip_encoder/clip_mlp.py | 1 - .../clip_encoder/clip_sdpa_attention.py | 2 +- .../clip_encoder/clip_text_model.py | 1 - .../text_encoder/clip_encoder/encoder_clip.py | 10 +-- .../text_encoder/t5_encoder/t5_attention.py | 5 +- .../text_encoder/t5_encoder/t5_block.py | 1 - .../t5_encoder/t5_dense_relu_dense.py | 7 +- .../text_encoder/t5_encoder/t5_encoder.py | 1 - .../t5_encoder/t5_feed_forward.py | 9 +-- .../text_encoder/t5_encoder/t5_layer_norm.py | 7 +- .../t5_encoder/t5_self_attention.py | 3 +- .../transformer/ada_layer_norm_continous.py | 6 +- .../models/transformer/ada_layer_norm_zero.py | 13 ++-- .../transformer/ada_layer_norm_zero_single.py | 9 +-- src/mflux/models/transformer/embed_nd.py | 3 +- src/mflux/models/transformer/feed_forward.py | 1 - .../models/transformer/guidance_embedder.py | 1 - .../models/transformer/joint_attention.py | 32 +++++--- .../transformer/joint_transformer_block.py | 11 ++- .../transformer/single_block_attention.py | 13 ++-- .../transformer/single_transformer_block.py | 17 +++-- src/mflux/models/transformer/text_embedder.py | 1 - .../models/transformer/time_text_embed.py | 9 ++- .../models/transformer/timestep_embedder.py | 1 - src/mflux/models/transformer/transformer.py | 35 +++++---- src/mflux/models/vae/common/attention.py | 1 - .../models/vae/common/resnet_block_2d.py | 25 +++---- src/mflux/models/vae/common/unet_mid_block.py | 3 +- src/mflux/models/vae/decoder/conv_in.py | 1 - src/mflux/models/vae/decoder/conv_norm_out.py | 2 +- src/mflux/models/vae/decoder/decoder.py | 1 - .../models/vae/decoder/up_block_1_or_2.py | 3 +- src/mflux/models/vae/decoder/up_block_3.py | 1 - src/mflux/models/vae/decoder/up_block_4.py | 1 - src/mflux/models/vae/decoder/up_sampler.py | 1 - src/mflux/models/vae/encoder/conv_in.py | 1 - src/mflux/models/vae/encoder/conv_norm_out.py | 2 +- src/mflux/models/vae/encoder/down_block_1.py | 1 - src/mflux/models/vae/encoder/down_block_2.py | 1 - src/mflux/models/vae/encoder/down_block_3.py | 1 - src/mflux/models/vae/encoder/down_block_4.py | 1 - src/mflux/models/vae/encoder/down_sampler.py | 1 - src/mflux/models/vae/encoder/encoder.py | 1 - src/mflux/post_processing/generated_image.py | 73 +++++++++---------- src/mflux/post_processing/image_util.py | 19 +++-- src/mflux/save.py | 10 +-- src/mflux/tokenizer/t5_tokenizer.py | 1 - src/mflux/tokenizer/tokenizer_handler.py | 17 +++-- src/mflux/weights/lora_converter.py | 64 +++++++++------- src/mflux/weights/lora_util.py | 26 ++++--- src/mflux/weights/model_saver.py | 7 +- src/mflux/weights/weight_handler.py | 50 +++++++------ src/mflux/weights/weight_util.py | 1 - 66 files changed, 362 insertions(+), 347 deletions(-) diff --git a/src/mflux/__init__.py b/src/mflux/__init__.py index 9547ab8..3893264 100644 --- a/src/mflux/__init__.py +++ b/src/mflux/__init__.py @@ -5,4 +5,11 @@ from mflux.controlnet.flux_controlnet import Flux1Controlnet from mflux.flux.flux import Flux1 from mflux.post_processing.image_util import ImageUtil -__all__ = ["Flux1", "Flux1Controlnet", "Config", "ConfigControlnet", "ModelConfig", "ImageUtil"] +__all__ = [ + "Flux1", + "Flux1Controlnet", + "Config", + "ConfigControlnet", + "ModelConfig", + "ImageUtil", +] diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py index 1cc86b4..19f8974 100644 --- a/src/mflux/config/config.py +++ b/src/mflux/config/config.py @@ -9,11 +9,11 @@ class Config: precision: mx.Dtype = mx.bfloat16 def __init__( - self, - num_inference_steps: int = 4, - width: int = 1024, - height: int = 1024, - guidance: float = 4.0, + self, + num_inference_steps: int = 4, + width: int = 1024, + height: int = 1024, + guidance: float = 4.0, ): if width % 16 != 0 or height % 16 != 0: log.warning("Width and height should be multiples of 16. Rounding down.") @@ -25,12 +25,12 @@ class Config: class ConfigControlnet(Config): def __init__( - self, - num_inference_steps: int = 4, - width: int = 1024, - height: int = 1024, - guidance: float = 4.0, - controlnet_strength: float = 1.0, + self, + num_inference_steps: int = 4, + width: int = 1024, + height: int = 1024, + guidance: float = 4.0, + controlnet_strength: float = 1.0, ): super().__init__(num_inference_steps, width, height, guidance) self.controlnet_strength = controlnet_strength diff --git a/src/mflux/config/model_config.py b/src/mflux/config/model_config.py index 4eb0836..e67a581 100644 --- a/src/mflux/config/model_config.py +++ b/src/mflux/config/model_config.py @@ -6,11 +6,11 @@ class ModelConfig(Enum): FLUX1_SCHNELL = ("black-forest-labs/FLUX.1-schnell", "schnell", 1000, 256) def __init__( - self, - model_name: str, - alias: str, - num_train_steps: int, - max_sequence_length: int, + self, + model_name: str, + alias: str, + num_train_steps: int, + max_sequence_length: int, ): self.alias = alias self.model_name = model_name diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py index efb33cf..6b770a1 100644 --- a/src/mflux/config/runtime_config.py +++ b/src/mflux/config/runtime_config.py @@ -6,7 +6,6 @@ from mflux.config.model_config import ModelConfig class RuntimeConfig: - def __init__(self, config: Config | ConfigControlnet, model_config: ModelConfig): self.config = config self.model_config = model_config @@ -35,7 +34,7 @@ class RuntimeConfig: @property def num_train_steps(self) -> int: return self.model_config.num_train_steps - + @property def controlnet_strength(self) -> float: if isinstance(self.config, ConfigControlnet): diff --git a/src/mflux/controlnet/controlnet_util.py b/src/mflux/controlnet/controlnet_util.py index ab6a9c5..6e38c75 100644 --- a/src/mflux/controlnet/controlnet_util.py +++ b/src/mflux/controlnet/controlnet_util.py @@ -4,7 +4,6 @@ import PIL class ControlnetUtil: - @staticmethod def preprocess_canny(img: PIL.Image) -> PIL.Image: image_to_canny = np.array(img) diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index 4edf228..3c342d1 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -29,13 +29,13 @@ CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny" class Flux1Controlnet: def __init__( - self, - model_config: ModelConfig, - quantize: int | None = None, - local_path: str | None = None, - lora_paths: list[str] | None = None, - lora_scales: list[float] | None = None, - controlnet_path: str | None = None, + self, + model_config: ModelConfig, + quantize: int | None = None, + local_path: str | None = None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, + controlnet_path: str | None = None, ): self.lora_paths = lora_paths self.lora_scales = lora_scales @@ -112,7 +112,7 @@ class Flux1Controlnet: control_image = ControlnetUtil.preprocess_canny(control_image) controlnet_cond = ImageUtil.to_array(control_image) controlnet_cong = self.vae.encode(controlnet_cond) - # the rescaling in the next line is not in the huggingface code, but without it the images from + # the rescaling in the next line is not in the huggingface code, but without it the images from # the chosen controlnet model are very bad controlnet_cond = (controlnet_cong / self.vae.scaling_factor) + self.vae.shift_factor controlnet_cond = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width) @@ -158,7 +158,7 @@ class Flux1Controlnet: seed=seed, prompt=prompt, quantization=self.bits, - generation_time=time_steps.format_dict['elapsed'], + generation_time=time_steps.format_dict["elapsed"], lora_paths=self.lora_paths, lora_scales=self.lora_scales, config=config, @@ -170,7 +170,7 @@ class Flux1Controlnet: latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5)) latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2)) return latents - + @staticmethod def _pack_latents(latents: mx.array, height: int, width: int) -> mx.array: latents = mx.reshape(latents, (1, 16, height // 16, 2, width // 16, 2)) diff --git a/src/mflux/controlnet/transformer_controlnet.py b/src/mflux/controlnet/transformer_controlnet.py index 0406d8b..232fcfe 100644 --- a/src/mflux/controlnet/transformer_controlnet.py +++ b/src/mflux/controlnet/transformer_controlnet.py @@ -4,19 +4,22 @@ from mlx import nn from mflux.config.model_config import ModelConfig from mflux.config.runtime_config import RuntimeConfig from mflux.models.transformer.embed_nd import EmbedND -from mflux.models.transformer.joint_transformer_block import JointTransformerBlock -from mflux.models.transformer.single_transformer_block import SingleTransformerBlock +from mflux.models.transformer.joint_transformer_block import ( + JointTransformerBlock, +) +from mflux.models.transformer.single_transformer_block import ( + SingleTransformerBlock, +) from mflux.models.transformer.time_text_embed import TimeTextEmbed from mflux.models.transformer.transformer import Transformer class TransformerControlnet(nn.Module): - def __init__( - self, - model_config: ModelConfig, - num_blocks: int, - num_single_blocks: int, + self, + model_config: ModelConfig, + num_blocks: int, + num_single_blocks: int, ): super().__init__() self.pos_embed = EmbedND() @@ -33,13 +36,13 @@ class TransformerControlnet(nn.Module): self.controlnet_single_blocks = [nn.Linear(3072, 3072) for _ in range(num_single_blocks)] def forward( - self, - t: int, - prompt_embeds: mx.array, - pooled_prompt_embeds: mx.array, - hidden_states: mx.array, - controlnet_cond: mx.array, - config: RuntimeConfig, + self, + t: int, + prompt_embeds: mx.array, + pooled_prompt_embeds: mx.array, + hidden_states: mx.array, + controlnet_cond: mx.array, + config: RuntimeConfig, ) -> (list[mx.array], list[mx.array]): time_step = config.sigmas[t] * config.num_train_steps time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision) @@ -61,7 +64,7 @@ class TransformerControlnet(nn.Module): hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, text_embeddings=text_embeddings, - rotary_embeddings=image_rotary_emb + rotary_embeddings=image_rotary_emb, ) block_samples = block_samples + (hidden_states,) @@ -78,9 +81,9 @@ class TransformerControlnet(nn.Module): hidden_states = block.forward( hidden_states=hidden_states, text_embeddings=text_embeddings, - rotary_embeddings=image_rotary_emb + rotary_embeddings=image_rotary_emb, ) - single_block_samples = single_block_samples + (hidden_states[:, encoder_hidden_states.shape[1]:],) + single_block_samples = single_block_samples + (hidden_states[:, encoder_hidden_states.shape[1] :],) controlnet_single_block_samples = () for single_block_sample, controlnet_block in zip(single_block_samples, self.controlnet_single_blocks): diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index fd45ee9..55afdea 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -19,14 +19,13 @@ from mflux.weights.weight_handler import WeightHandler class Flux1: - def __init__( - self, - model_config: ModelConfig, - quantize: int | None = None, - local_path: str | None = None, - lora_paths: list[str] | None = None, - lora_scales: list[float] | None = None, + self, + model_config: ModelConfig, + quantize: int | None = None, + local_path: str | None = None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, ): self.lora_paths = lora_paths self.lora_scales = lora_scales @@ -110,7 +109,7 @@ class Flux1: seed=seed, prompt=prompt, quantization=self.bits, - generation_time=time_steps.format_dict['elapsed'], + generation_time=time_steps.format_dict["elapsed"], lora_paths=self.lora_paths, lora_scales=self.lora_scales, config=config, diff --git a/src/mflux/generate.py b/src/mflux/generate.py index fdcf8bd..6043a40 100644 --- a/src/mflux/generate.py +++ b/src/mflux/generate.py @@ -5,20 +5,20 @@ from mflux import Flux1, Config, ModelConfig def main(): - parser = argparse.ArgumentParser(description='Generate an image based on a prompt.') - parser.add_argument('--prompt', type=str, required=True, help='The textual description of the image to generate.') - parser.add_argument('--output', type=str, default="image.png", help='The filename for the output image. Default is "image.png".') - parser.add_argument('--model', "-m", type=str, required=True, choices=["dev", "schnell"], help='The model to use ("schnell" or "dev").') - parser.add_argument('--seed', type=int, default=None, help='Entropy Seed (Default is time-based random-seed)') - parser.add_argument('--height', type=int, default=1024, help='Image height (Default is 1024)') - parser.add_argument('--width', type=int, default=1024, help='Image width (Default is 1024)') - parser.add_argument('--steps', type=int, default=None, help='Inference Steps') - parser.add_argument('--guidance', type=float, default=3.5, help='Guidance Scale (Default is 3.5)') - parser.add_argument('--quantize', "-q", type=int, choices=[4, 8], default=None, help='Quantize the model (4 or 8, Default is None)') - parser.add_argument('--path', type=str, default=None, help='Local path for loading a model from disk') - parser.add_argument('--lora-paths', type=str, nargs='*', default=None, help='Local safetensors for applying LORA from disk') - parser.add_argument('--lora-scales', type=float, nargs='*', default=None, help='Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.') - parser.add_argument('--metadata', action='store_true', help='Export image metadata as a JSON file.') + parser = argparse.ArgumentParser(description="Generate an image based on a prompt.") + parser.add_argument("--prompt", type=str, required=True, help="The textual description of the image to generate.") + parser.add_argument("--output", type=str, default="image.png", help="The filename for the output image. Default is \"image.png\".") + parser.add_argument("--model", "-m", type=str, required=True, choices=["dev", "schnell"], help="The model to use (\"schnell\" or \"dev\").") + parser.add_argument("--seed", type=int, default=None, help="Entropy Seed (Default is time-based random-seed)") + parser.add_argument("--height", type=int, default=1024, help="Image height (Default is 1024)") + parser.add_argument("--width", type=int, default=1024, help="Image width (Default is 1024)") + parser.add_argument("--steps", type=int, default=None, help="Inference Steps") + parser.add_argument("--guidance", type=float, default=3.5, help="Guidance Scale (Default is 3.5)") + parser.add_argument("--quantize", "-q", type=int, choices=[4, 8], default=None, help="Quantize the model (4 or 8, Default is None)") + parser.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk") + parser.add_argument("--lora-paths", type=str, nargs="*", default=None, help="Local safetensors for applying LORA from disk") + parser.add_argument("--lora-scales", type=float, nargs="*", default=None, help="Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.") + parser.add_argument("--metadata", action="store_true", help="Export image metadata as a JSON file.") args = parser.parse_args() @@ -34,7 +34,7 @@ def main(): quantize=args.quantize, local_path=args.path, lora_paths=args.lora_paths, - lora_scales=args.lora_scales + lora_scales=args.lora_scales, ) # Generate an image @@ -46,12 +46,12 @@ def main(): height=args.height, width=args.width, guidance=args.guidance, - ) + ), ) # Save the image image.save(path=args.output, export_json_metadata=args.metadata) -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index 2e0784b..9f2ab64 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/generate_controlnet.py @@ -5,22 +5,22 @@ from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig, ImageUtil def main(): - parser = argparse.ArgumentParser(description='Generate an image based on a prompt.') - parser.add_argument('--prompt', type=str, required=True, help='The textual description of the image to generate.') - parser.add_argument('--control-image-path', type=str, required=True, help='Local path of the image to use as input for controlnet.') - parser.add_argument('--output', type=str, default="image.png", help='The filename for the output image. Default is "image.png".') - parser.add_argument('--model', "-m", type=str, required=True, choices=["dev", "schnell"], help='The model to use ("schnell" or "dev").') - parser.add_argument('--seed', type=int, default=None, help='Entropy Seed (Default is time-based random-seed)') - parser.add_argument('--height', type=int, default=1024, help='Image height (Default is 1024)') - parser.add_argument('--width', type=int, default=1024, help='Image width (Default is 1024)') - parser.add_argument('--steps', type=int, default=None, help='Inference Steps') - parser.add_argument('--guidance', type=float, default=3.5, help='Guidance Scale (Default is 3.5)') - parser.add_argument('--controlnet-strength', type=float, default=0.7, help='Controls how strongly the control image influences the output image. A value of 0.0 means no influence. (Default is 0.7)') - parser.add_argument('--quantize', "-q", type=int, choices=[4, 8], default=None, help='Quantize the model (4 or 8, Default is None)') - parser.add_argument('--path', type=str, default=None, help='Local path for loading a model from disk') - parser.add_argument('--lora-paths', type=str, nargs='*', default=None, help='Local safetensors for applying LORA from disk') - parser.add_argument('--lora-scales', type=float, nargs='*', default=None, help='Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.') - parser.add_argument('--metadata', action='store_true', help='Export image metadata as a JSON file.') + parser = argparse.ArgumentParser(description="Generate an image based on a prompt.") + parser.add_argument("--prompt", type=str, required=True, help="The textual description of the image to generate.") + parser.add_argument("--control-image-path", type=str, required=True, help="Local path of the image to use as input for controlnet.") + parser.add_argument("--output", type=str, default="image.png", help="The filename for the output image. Default is \"image.png\".") + parser.add_argument("--model", "-m", type=str, required=True, choices=["dev", "schnell"], help="The model to use (\"schnell\" or \"dev\").") + parser.add_argument("--seed", type=int, default=None, help="Entropy Seed (Default is time-based random-seed)") + parser.add_argument("--height", type=int, default=1024, help="Image height (Default is 1024)") + parser.add_argument("--width", type=int, default=1024, help="Image width (Default is 1024)") + parser.add_argument("--steps", type=int, default=None, help="Inference Steps") + parser.add_argument("--guidance", type=float, default=3.5, help="Guidance Scale (Default is 3.5)") + parser.add_argument("--controlnet-strength", type=float, default=0.7, help="Controls how strongly the control image influences the output image. A value of 0.0 means no influence. (Default is 0.7)") + parser.add_argument("--quantize", "-q", type=int, choices=[4, 8], default=None, help="Quantize the model (4 or 8, Default is None)") + parser.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk") + parser.add_argument("--lora-paths", type=str, nargs="*", default=None, help="Local safetensors for applying LORA from disk") + parser.add_argument("--lora-scales", type=float, nargs="*", default=None, help="Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.") + parser.add_argument("--metadata", action="store_true", help="Export image metadata as a JSON file.") args = parser.parse_args() @@ -36,7 +36,7 @@ def main(): quantize=args.quantize, local_path=args.path, lora_paths=args.lora_paths, - lora_scales=args.lora_scales + lora_scales=args.lora_scales, ) # Generate an image @@ -49,13 +49,13 @@ def main(): height=args.height, width=args.width, guidance=args.guidance, - controlnet_strength=args.controlnet_strength - ) + controlnet_strength=args.controlnet_strength, + ), ) # Save the image image.save(path=args.output, export_json_metadata=args.metadata) -if __name__ == '__main__': +if __name__ == "__main__": main() 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 9c4fd83..988c5d3 100644 --- a/src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py +++ b/src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py @@ -5,7 +5,6 @@ from mflux.tokenizer.clip_tokenizer import TokenizerCLIP class CLIPEmbeddings(nn.Module): - def __init__(self, dims: int): super().__init__() self.position_embedding = nn.Embedding(num_embeddings=TokenizerCLIP.MAX_TOKEN_LENGTH, dims=dims) 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 a8ab552..f1ac39b 100644 --- a/src/mflux/models/text_encoder/clip_encoder/clip_encoder.py +++ b/src/mflux/models/text_encoder/clip_encoder/clip_encoder.py @@ -5,7 +5,6 @@ from mflux.models.text_encoder.clip_encoder.clip_text_model import CLIPTextModel class CLIPEncoder(nn.Module): - def __init__(self): super().__init__() self.text_model = CLIPTextModel(dims=768, num_encoder_layers=12) 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 3d20300..972b39c 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 @@ -6,7 +6,6 @@ from mflux.models.text_encoder.clip_encoder.clip_sdpa_attention import CLIPSdpaA class CLIPEncoderLayer(nn.Module): - def __init__(self, layer: int): super().__init__() self.self_attn = CLIPSdpaAttention() 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 a1306f8..7ae9dd1 100644 --- a/src/mflux/models/text_encoder/clip_encoder/clip_mlp.py +++ b/src/mflux/models/text_encoder/clip_encoder/clip_mlp.py @@ -3,7 +3,6 @@ from mlx import nn class CLIPMLP(nn.Module): - def __init__(self): super().__init__() self.fc1 = nn.Linear(input_dims=768, output_dims=3072) 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 c970152..d9eded6 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 @@ -36,7 +36,7 @@ class CLIPSdpaAttention(nn.Module): scores = (query * scale) @ key.transpose(0, 1, 3, 2) scores = scores + mask attn = mx.softmax(scores, axis=-1) - hidden_states = (attn @ value) + hidden_states = attn @ value return hidden_states @staticmethod 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 1f25fa6..fc9eef3 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 @@ -6,7 +6,6 @@ from mflux.models.text_encoder.clip_encoder.encoder_clip import EncoderCLIP class CLIPTextModel(nn.Module): - def __init__(self, dims: int, num_encoder_layers: int): super().__init__() self.encoder = EncoderCLIP(num_encoder_layers) 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 7ec2b12..2ab9078 100644 --- a/src/mflux/models/text_encoder/clip_encoder/encoder_clip.py +++ b/src/mflux/models/text_encoder/clip_encoder/encoder_clip.py @@ -1,11 +1,12 @@ import mlx.core as mx from mlx import nn -from mflux.models.text_encoder.clip_encoder.clip_encoder_layer import CLIPEncoderLayer +from mflux.models.text_encoder.clip_encoder.clip_encoder_layer import ( + CLIPEncoderLayer, +) class EncoderCLIP(nn.Module): - def __init__(self, num_encoder_layers: int): super().__init__() self.layers = [CLIPEncoderLayer(i) for i in range(num_encoder_layers)] @@ -13,8 +14,5 @@ class EncoderCLIP(nn.Module): def forward(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.forward(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 f7ba4aa..6258716 100644 --- a/src/mflux/models/text_encoder/t5_encoder/t5_attention.py +++ b/src/mflux/models/text_encoder/t5_encoder/t5_attention.py @@ -2,11 +2,12 @@ import mlx.core as mx from mlx import nn from mflux.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm -from mflux.models.text_encoder.t5_encoder.t5_self_attention import T5SelfAttention +from mflux.models.text_encoder.t5_encoder.t5_self_attention import ( + T5SelfAttention, +) class T5Attention(nn.Module): - def __init__(self): super().__init__() self.SelfAttention = T5SelfAttention() 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 48f0fdc..8787f5b 100644 --- a/src/mflux/models/text_encoder/t5_encoder/t5_block.py +++ b/src/mflux/models/text_encoder/t5_encoder/t5_block.py @@ -6,7 +6,6 @@ from mflux.models.text_encoder.t5_encoder.t5_feed_forward import T5FeedForward class T5Block(nn.Module): - def __init__(self, layer: int): super().__init__() self.attention = T5Attention() 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 0da79e5..0318fe7 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 @@ -5,7 +5,6 @@ import mlx.core as mx class T5DenseReluDense(nn.Module): - def __init__(self): super().__init__() self.wi_0 = nn.Linear(4096, 10240, bias=False) @@ -21,4 +20,8 @@ class T5DenseReluDense(nn.Module): @staticmethod def new_gelu(input_array: mx.array) -> mx.array: - return 0.5 * input_array * (1.0 + mx.tanh(math.sqrt(2.0 / math.pi) * (input_array + 0.044715 * mx.power(input_array, 3.0)))) + return ( + 0.5 + * input_array + * (1.0 + mx.tanh(math.sqrt(2.0 / math.pi) * (input_array + 0.044715 * mx.power(input_array, 3.0)))) + ) 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 b9d4e3e..92f7741 100644 --- a/src/mflux/models/text_encoder/t5_encoder/t5_encoder.py +++ b/src/mflux/models/text_encoder/t5_encoder/t5_encoder.py @@ -6,7 +6,6 @@ from mflux.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm class T5Encoder(nn.Module): - def __init__(self): super().__init__() self.shared = nn.Embedding(num_embeddings=32128, dims=4096) 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 024e345..e5c52df 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 @@ -1,14 +1,13 @@ -import math - -from mlx import nn import mlx.core as mx +from mlx import nn -from mflux.models.text_encoder.t5_encoder.t5_dense_relu_dense import T5DenseReluDense +from mflux.models.text_encoder.t5_encoder.t5_dense_relu_dense import ( + T5DenseReluDense, +) from mflux.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm class T5FeedForward(nn.Module): - def __init__(self): super().__init__() self.layer_norm = T5LayerNorm() 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 987566a..b1bef54 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 @@ -3,13 +3,16 @@ from mlx import nn class T5LayerNorm(nn.Module): - def __init__(self): super().__init__() self.weight = mx.ones((4096,)) self.variance_epsilon = 1e-06 def forward(self, hidden_states: mx.array) -> mx.array: - variance = mx.mean(mx.power(hidden_states.astype(mx.float32), 2), axis=-1, keepdims=True) + variance = mx.mean( + mx.power(hidden_states.astype(mx.float32), 2), + axis=-1, + keepdims=True, + ) hidden_states = hidden_states * mx.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states 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 a628d85..d51708c 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 @@ -5,7 +5,6 @@ from mlx import nn class T5SelfAttention(nn.Module): - def __init__(self): super().__init__() self.q = nn.Linear(4096, 4096, bias=False) @@ -64,7 +63,7 @@ class T5SelfAttention(nn.Module): ).astype(mx.int32) relative_position_if_large = mx.minimum( relative_position_if_large, - mx.full(relative_position_if_large.shape, num_buckets - 1) + mx.full(relative_position_if_large.shape, num_buckets - 1), ) relative_buckets += mx.where(is_small, relative_position, relative_position_if_large) diff --git a/src/mflux/models/transformer/ada_layer_norm_continous.py b/src/mflux/models/transformer/ada_layer_norm_continous.py index becb08d..9538fd0 100644 --- a/src/mflux/models/transformer/ada_layer_norm_continous.py +++ b/src/mflux/models/transformer/ada_layer_norm_continous.py @@ -5,7 +5,6 @@ from mflux.config.config import Config class AdaLayerNormContinuous(nn.Module): - def __init__(self, embedding_dim: int, conditioning_embedding_dim: int): super().__init__() self.embedding_dim = embedding_dim @@ -15,8 +14,7 @@ class AdaLayerNormContinuous(nn.Module): def forward(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] - shift = text_embeddings[:, 1*chunk_size:2*chunk_size] + scale = text_embeddings[:, 0 * chunk_size : 1 * chunk_size] + shift = text_embeddings[:, 1 * chunk_size : 2 * chunk_size] x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] return x - diff --git a/src/mflux/models/transformer/ada_layer_norm_zero.py b/src/mflux/models/transformer/ada_layer_norm_zero.py index cb158ad..17d0956 100644 --- a/src/mflux/models/transformer/ada_layer_norm_zero.py +++ b/src/mflux/models/transformer/ada_layer_norm_zero.py @@ -3,7 +3,6 @@ import mlx.core as mx class AdaLayerNormZero(nn.Module): - def __init__(self): super().__init__() self.linear = nn.Linear(3072, 18432) @@ -12,11 +11,11 @@ class AdaLayerNormZero(nn.Module): def forward(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] - scale_msa = text_embeddings[:, 1*chunk_size:2*chunk_size] - gate_msa = text_embeddings[:, 2*chunk_size:3*chunk_size] - shift_mlp = text_embeddings[:, 3*chunk_size:4*chunk_size] - scale_mlp = text_embeddings[:, 4*chunk_size:5*chunk_size] - gate_mlp = text_embeddings[:, 5*chunk_size:6*chunk_size] + shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size] + scale_msa = text_embeddings[:, 1 * chunk_size : 2 * chunk_size] + gate_msa = text_embeddings[:, 2 * chunk_size : 3 * chunk_size] + shift_mlp = text_embeddings[:, 3 * chunk_size : 4 * chunk_size] + scale_mlp = text_embeddings[:, 4 * chunk_size : 5 * chunk_size] + gate_mlp = text_embeddings[:, 5 * chunk_size : 6 * chunk_size] x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] return x, gate_msa, shift_mlp, scale_mlp, gate_mlp 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 174e5ab..9299d67 100644 --- a/src/mflux/models/transformer/ada_layer_norm_zero_single.py +++ b/src/mflux/models/transformer/ada_layer_norm_zero_single.py @@ -3,17 +3,16 @@ import mlx.core as mx class AdaLayerNormZeroSingle(nn.Module): - def __init__(self): super().__init__() - self.linear = nn.Linear(3072, 3*3072) + 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): text_embeddings = self.linear(nn.silu(text_embeddings)) chunk_size = 9216 // 3 - shift_msa = text_embeddings[:, 0*chunk_size:1*chunk_size] - scale_msa = text_embeddings[:, 1*chunk_size:2*chunk_size] - gate_msa = text_embeddings[:, 2*chunk_size:3*chunk_size] + shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size] + scale_msa = text_embeddings[:, 1 * chunk_size : 2 * chunk_size] + gate_msa = text_embeddings[:, 2 * chunk_size : 3 * chunk_size] x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] return x, gate_msa diff --git a/src/mflux/models/transformer/embed_nd.py b/src/mflux/models/transformer/embed_nd.py index 03721dc..4d92dce 100644 --- a/src/mflux/models/transformer/embed_nd.py +++ b/src/mflux/models/transformer/embed_nd.py @@ -3,7 +3,6 @@ from mlx import nn class EmbedND(nn.Module): - def __init__(self): super().__init__() self.dim = 3072 @@ -20,7 +19,7 @@ class EmbedND(nn.Module): @staticmethod def rope(pos: mx.array, dim: int, theta: float) -> mx.array: scale = mx.arange(0, dim, 2, dtype=mx.float32) / dim - omega = 1.0 / (theta ** scale) + omega = 1.0 / (theta**scale) batch_size, seq_length = pos.shape pos_expanded = mx.expand_dims(pos, axis=-1) omega_expanded = mx.expand_dims(omega, axis=0) diff --git a/src/mflux/models/transformer/feed_forward.py b/src/mflux/models/transformer/feed_forward.py index cb01d8e..3a13d4b 100644 --- a/src/mflux/models/transformer/feed_forward.py +++ b/src/mflux/models/transformer/feed_forward.py @@ -3,7 +3,6 @@ import mlx.core as mx class FeedForward(nn.Module): - def __init__(self, activation_function): super().__init__() self.linear1 = nn.Linear(3072, 12288) diff --git a/src/mflux/models/transformer/guidance_embedder.py b/src/mflux/models/transformer/guidance_embedder.py index ca21128..91ba297 100644 --- a/src/mflux/models/transformer/guidance_embedder.py +++ b/src/mflux/models/transformer/guidance_embedder.py @@ -3,7 +3,6 @@ import mlx.core as mx class GuidanceEmbedder(nn.Module): - def __init__(self): super().__init__() self.linear_1 = nn.Linear(256, 3072) diff --git a/src/mflux/models/transformer/joint_attention.py b/src/mflux/models/transformer/joint_attention.py index 806b7a2..7c7cd80 100644 --- a/src/mflux/models/transformer/joint_attention.py +++ b/src/mflux/models/transformer/joint_attention.py @@ -23,10 +23,10 @@ class JointAttention(nn.Module): self.norm_added_k = nn.RMSNorm(128) def forward( - self, - hidden_states: mx.array, - encoder_hidden_states: mx.array, - image_rotary_emb: mx.array + self, + hidden_states: mx.array, + encoder_hidden_states: mx.array, + image_rotary_emb: mx.array, ) -> (mx.array, mx.array): residual = hidden_states @@ -45,9 +45,18 @@ class JointAttention(nn.Module): encoder_hidden_states_key_proj = self.add_k_proj(encoder_hidden_states) encoder_hidden_states_value_proj = self.add_v_proj(encoder_hidden_states) - encoder_hidden_states_query_proj = mx.transpose(mx.reshape(encoder_hidden_states_query_proj, (1, -1, 24, 128)), (0, 2, 1, 3)) - encoder_hidden_states_key_proj = mx.transpose(mx.reshape(encoder_hidden_states_key_proj, (1, -1, 24, 128)), (0, 2, 1, 3)) - encoder_hidden_states_value_proj = mx.transpose(mx.reshape(encoder_hidden_states_value_proj, (1, -1, 24, 128)), (0, 2, 1, 3)) + encoder_hidden_states_query_proj = mx.transpose( + mx.reshape(encoder_hidden_states_query_proj, (1, -1, 24, 128)), + (0, 2, 1, 3), + ) + encoder_hidden_states_key_proj = mx.transpose( + mx.reshape(encoder_hidden_states_key_proj, (1, -1, 24, 128)), + (0, 2, 1, 3), + ) + encoder_hidden_states_value_proj = mx.transpose( + mx.reshape(encoder_hidden_states_value_proj, (1, -1, 24, 128)), + (0, 2, 1, 3), + ) encoder_hidden_states_query_proj = self.norm_added_q(encoder_hidden_states_query_proj) encoder_hidden_states_key_proj = self.norm_added_k(encoder_hidden_states_key_proj) @@ -60,10 +69,13 @@ class JointAttention(nn.Module): hidden_states = JointAttention.attention(query, key, value) hidden_states = mx.transpose(hidden_states, (0, 2, 1, 3)) - hidden_states = mx.reshape(hidden_states, (self.batch_size, -1, self.num_heads * self.head_dimension)) + hidden_states = mx.reshape( + hidden_states, + (self.batch_size, -1, self.num_heads * self.head_dimension), + ) encoder_hidden_states, hidden_states = ( hidden_states[:, : encoder_hidden_states.shape[1]], - hidden_states[:, encoder_hidden_states.shape[1]:], + hidden_states[:, encoder_hidden_states.shape[1] :], ) hidden_states = self.to_out[0](hidden_states) @@ -76,7 +88,7 @@ class JointAttention(nn.Module): scale = 1 / mx.sqrt(query.shape[-1]) scores = (query * scale) @ key.transpose(0, 1, 3, 2) attn = mx.softmax(scores, axis=-1) - hidden_states = (attn @ value) + hidden_states = attn @ value return hidden_states @staticmethod diff --git a/src/mflux/models/transformer/joint_transformer_block.py b/src/mflux/models/transformer/joint_transformer_block.py index 3083d53..d6f2349 100644 --- a/src/mflux/models/transformer/joint_transformer_block.py +++ b/src/mflux/models/transformer/joint_transformer_block.py @@ -7,7 +7,6 @@ from mflux.models.transformer.joint_attention import JointAttention class JointTransformerBlock(nn.Module): - def __init__(self, layer): super().__init__() self.layer = layer @@ -20,11 +19,11 @@ class JointTransformerBlock(nn.Module): self.norm2_context = nn.LayerNorm(dims=1536, eps=1e-6, affine=False) def forward( - self, - hidden_states: mx.array, - encoder_hidden_states: mx.array, - text_embeddings: mx.array, - rotary_embeddings: mx.array + 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) diff --git a/src/mflux/models/transformer/single_block_attention.py b/src/mflux/models/transformer/single_block_attention.py index 8b38fe7..7f806a3 100644 --- a/src/mflux/models/transformer/single_block_attention.py +++ b/src/mflux/models/transformer/single_block_attention.py @@ -15,11 +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 forward(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) @@ -35,7 +31,10 @@ class SingleBlockAttention(nn.Module): hidden_states = SingleBlockAttention.attention(query, key, value) hidden_states = mx.transpose(hidden_states, (0, 2, 1, 3)) - hidden_states = mx.reshape(hidden_states, (self.batch_size, -1, self.num_heads * self.head_dimension)) + hidden_states = mx.reshape( + hidden_states, + (self.batch_size, -1, self.num_heads * self.head_dimension), + ) return hidden_states @@ -44,7 +43,7 @@ class SingleBlockAttention(nn.Module): scale = 1 / mx.sqrt(query.shape[-1]) scores = (query * scale) @ key.transpose(0, 1, 3, 2) attn = mx.softmax(scores, axis=-1) - hidden_states = (attn @ value) + hidden_states = attn @ value return hidden_states @staticmethod diff --git a/src/mflux/models/transformer/single_transformer_block.py b/src/mflux/models/transformer/single_transformer_block.py index 4b51cf9..31cdef6 100644 --- a/src/mflux/models/transformer/single_transformer_block.py +++ b/src/mflux/models/transformer/single_transformer_block.py @@ -1,25 +1,26 @@ import mlx.core as mx from mlx import nn -from mflux.models.transformer.ada_layer_norm_zero_single import AdaLayerNormZeroSingle +from mflux.models.transformer.ada_layer_norm_zero_single import ( + AdaLayerNormZeroSingle, +) from mflux.models.transformer.single_block_attention import SingleBlockAttention class SingleTransformerBlock(nn.Module): - def __init__(self, layer): super().__init__() self.layer = layer self.norm = AdaLayerNormZeroSingle() - self.proj_mlp = nn.Linear(3072, 4*3072) + self.proj_mlp = nn.Linear(3072, 4 * 3072) self.attn = SingleBlockAttention() - self.proj_out = nn.Linear(3072 + 4*3072, 3072) + self.proj_out = nn.Linear(3072 + 4 * 3072, 3072) def forward( - self, - hidden_states: mx.array, - text_embeddings: mx.array, - rotary_embeddings: mx.array + 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) diff --git a/src/mflux/models/transformer/text_embedder.py b/src/mflux/models/transformer/text_embedder.py index 120ef20..6fe09e2 100644 --- a/src/mflux/models/transformer/text_embedder.py +++ b/src/mflux/models/transformer/text_embedder.py @@ -3,7 +3,6 @@ import mlx.core as mx class TextEmbedder(nn.Module): - def __init__(self): super().__init__() self.linear_1 = nn.Linear(768, 3072) diff --git a/src/mflux/models/transformer/time_text_embed.py b/src/mflux/models/transformer/time_text_embed.py index 7333e4c..a3db91a 100644 --- a/src/mflux/models/transformer/time_text_embed.py +++ b/src/mflux/models/transformer/time_text_embed.py @@ -10,14 +10,18 @@ from mflux.models.transformer.guidance_embedder import GuidanceEmbedder class TimeTextEmbed(nn.Module): - def __init__(self, model_config: ModelConfig): super().__init__() self.text_embedder = TextEmbedder() self.guidance_embedder = GuidanceEmbedder() if model_config == ModelConfig.FLUX1_DEV else None self.timestep_embedder = TimestepEmbedder() - def forward(self, time_step: mx.array, pooled_projection: mx.array, guidance: mx.array) -> mx.array: + def forward( + 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) if self.guidance_embedder is not None: @@ -37,4 +41,3 @@ class TimeTextEmbed(nn.Module): emb = mx.concatenate([mx.sin(emb), mx.cos(emb)], axis=-1) emb = mx.concatenate([emb[:, half_dim:], emb[:, :half_dim]], axis=-1) return emb - diff --git a/src/mflux/models/transformer/timestep_embedder.py b/src/mflux/models/transformer/timestep_embedder.py index 8d947a1..dfe50fd 100644 --- a/src/mflux/models/transformer/timestep_embedder.py +++ b/src/mflux/models/transformer/timestep_embedder.py @@ -3,7 +3,6 @@ import mlx.core as mx class TimestepEmbedder(nn.Module): - def __init__(self): super().__init__() self.linear_1 = nn.Linear(256, 3072) diff --git a/src/mflux/models/transformer/transformer.py b/src/mflux/models/transformer/transformer.py index cd32d83..387304a 100644 --- a/src/mflux/models/transformer/transformer.py +++ b/src/mflux/models/transformer/transformer.py @@ -5,15 +5,20 @@ from mlx import nn from mflux.config.model_config import ModelConfig from mflux.config.runtime_config import RuntimeConfig -from mflux.models.transformer.ada_layer_norm_continous import AdaLayerNormContinuous +from mflux.models.transformer.ada_layer_norm_continous import ( + AdaLayerNormContinuous, +) from mflux.models.transformer.embed_nd import EmbedND -from mflux.models.transformer.joint_transformer_block import JointTransformerBlock -from mflux.models.transformer.single_transformer_block import SingleTransformerBlock +from mflux.models.transformer.joint_transformer_block import ( + JointTransformerBlock, +) +from mflux.models.transformer.single_transformer_block import ( + SingleTransformerBlock, +) from mflux.models.transformer.time_text_embed import TimeTextEmbed class Transformer(nn.Module): - def __init__(self, model_config: ModelConfig): super().__init__() self.pos_embed = EmbedND() @@ -26,14 +31,14 @@ class Transformer(nn.Module): self.proj_out = nn.Linear(3072, 64) def predict( - self, - t: int, - prompt_embeds: mx.array, - pooled_prompt_embeds: mx.array, - hidden_states: mx.array, - config: RuntimeConfig, - controlnet_block_samples: list[mx.array] | None = None, - controlnet_single_block_samples: list[mx.array] | None = None, + self, + t: int, + prompt_embeds: mx.array, + pooled_prompt_embeds: mx.array, + hidden_states: mx.array, + config: RuntimeConfig, + controlnet_block_samples: list[mx.array] | None = None, + controlnet_single_block_samples: list[mx.array] | None = None, ) -> mx.array: time_step = config.sigmas[t] * config.num_train_steps time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision) @@ -51,7 +56,7 @@ class Transformer(nn.Module): hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, text_embeddings=text_embeddings, - rotary_embeddings=image_rotary_emb + rotary_embeddings=image_rotary_emb, ) if controlnet_block_samples is not None and len(controlnet_block_samples) > 0: interval_control = len(self.transformer_blocks) / len(controlnet_block_samples) @@ -64,7 +69,7 @@ class Transformer(nn.Module): hidden_states = block.forward( hidden_states=hidden_states, text_embeddings=text_embeddings, - rotary_embeddings=image_rotary_emb + rotary_embeddings=image_rotary_emb, ) if controlnet_single_block_samples is not None and len(controlnet_single_block_samples) > 0: interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples) @@ -74,7 +79,7 @@ class Transformer(nn.Module): + controlnet_single_block_samples[idx // interval_control] ) - hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...] + hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...] hidden_states = self.norm_out.forward(hidden_states, text_embeddings) hidden_states = self.proj_out(hidden_states) noise = hidden_states diff --git a/src/mflux/models/vae/common/attention.py b/src/mflux/models/vae/common/attention.py index 16e8546..f1ba8b3 100644 --- a/src/mflux/models/vae/common/attention.py +++ b/src/mflux/models/vae/common/attention.py @@ -5,7 +5,6 @@ from mflux.config.config import Config class Attention(nn.Module): - def __init__(self): super().__init__() self.group_norm = nn.GroupNorm(32, 512, pytorch_compatible=True) diff --git a/src/mflux/models/vae/common/resnet_block_2d.py b/src/mflux/models/vae/common/resnet_block_2d.py index 0bf73c5..067c2ff 100644 --- a/src/mflux/models/vae/common/resnet_block_2d.py +++ b/src/mflux/models/vae/common/resnet_block_2d.py @@ -5,18 +5,17 @@ from mflux.config.config import Config class ResnetBlock2D(nn.Module): - def __init__( - self, - norm1: int, - conv1_in: int, - conv1_out: int, - norm2: int, - conv2_in: int, - conv2_out: int, - conv_shortcut_in: int | None = None, - conv_shortcut_out: int | None = None, - is_conv_shortcut: bool = False + self, + norm1: int, + conv1_in: int, + conv1_out: int, + norm2: int, + conv2_in: int, + conv2_out: int, + conv_shortcut_in: int | None = None, + conv_shortcut_out: int | None = None, + is_conv_shortcut: bool = False, ): super().__init__() self.norm1 = nn.GroupNorm( @@ -24,14 +23,14 @@ class ResnetBlock2D(nn.Module): dims=norm1, eps=1e-6, affine=True, - pytorch_compatible=True + pytorch_compatible=True, ) self.norm2 = nn.GroupNorm( num_groups=32, dims=norm2, eps=1e-6, affine=True, - pytorch_compatible=True + pytorch_compatible=True, ) self.conv1 = nn.Conv2d( in_channels=conv1_in, diff --git a/src/mflux/models/vae/common/unet_mid_block.py b/src/mflux/models/vae/common/unet_mid_block.py index c9bda8a..bb1a423 100644 --- a/src/mflux/models/vae/common/unet_mid_block.py +++ b/src/mflux/models/vae/common/unet_mid_block.py @@ -6,13 +6,12 @@ from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D class UnetMidBlock(nn.Module): - def __init__(self): super().__init__() self.attentions = [Attention()] self.resnets = [ ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512), - ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512) + 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: diff --git a/src/mflux/models/vae/decoder/conv_in.py b/src/mflux/models/vae/decoder/conv_in.py index 9e1edd4..5c0eda9 100644 --- a/src/mflux/models/vae/decoder/conv_in.py +++ b/src/mflux/models/vae/decoder/conv_in.py @@ -3,7 +3,6 @@ import mlx.nn as nn class ConvIn(nn.Module): - def __init__(self): super().__init__() self.conv2d = nn.Conv2d( diff --git a/src/mflux/models/vae/decoder/conv_norm_out.py b/src/mflux/models/vae/decoder/conv_norm_out.py index 9194e93..bfdcb17 100644 --- a/src/mflux/models/vae/decoder/conv_norm_out.py +++ b/src/mflux/models/vae/decoder/conv_norm_out.py @@ -12,7 +12,7 @@ class ConvNormOut(nn.Module): dims=128, eps=1e-6, affine=True, - pytorch_compatible=True + pytorch_compatible=True, ) def forward(self, input_array: mx.array) -> mx.array: diff --git a/src/mflux/models/vae/decoder/decoder.py b/src/mflux/models/vae/decoder/decoder.py index c7c43f2..e61481a 100644 --- a/src/mflux/models/vae/decoder/decoder.py +++ b/src/mflux/models/vae/decoder/decoder.py @@ -11,7 +11,6 @@ from mflux.models.vae.decoder.up_block_4 import UpBlock4 class Decoder(nn.Module): - def __init__(self): super().__init__() self.conv_in = ConvIn() 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 69c646b..5763e48 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 @@ -6,13 +6,12 @@ from mflux.models.vae.decoder.up_sampler import UpSampler class UpBlock1Or2(nn.Module): - def __init__(self): super().__init__() self.resnets = [ ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512), ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512), - ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512) + ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512), ] self.upsamplers = [UpSampler(conv_in=512, conv_out=512)] diff --git a/src/mflux/models/vae/decoder/up_block_3.py b/src/mflux/models/vae/decoder/up_block_3.py index f7662bb..fd51799 100644 --- a/src/mflux/models/vae/decoder/up_block_3.py +++ b/src/mflux/models/vae/decoder/up_block_3.py @@ -6,7 +6,6 @@ from mflux.models.vae.decoder.up_sampler import UpSampler class UpBlock3(nn.Module): - def __init__(self): super().__init__() self.resnets = [ diff --git a/src/mflux/models/vae/decoder/up_block_4.py b/src/mflux/models/vae/decoder/up_block_4.py index 9fe9d4f..4b59341 100644 --- a/src/mflux/models/vae/decoder/up_block_4.py +++ b/src/mflux/models/vae/decoder/up_block_4.py @@ -5,7 +5,6 @@ from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D class UpBlock4(nn.Module): - def __init__(self): super().__init__() self.resnets = [ diff --git a/src/mflux/models/vae/decoder/up_sampler.py b/src/mflux/models/vae/decoder/up_sampler.py index e348489..46f147f 100644 --- a/src/mflux/models/vae/decoder/up_sampler.py +++ b/src/mflux/models/vae/decoder/up_sampler.py @@ -4,7 +4,6 @@ from mlx import nn class UpSampler(nn.Module): - def __init__(self, conv_in: int, conv_out: int): super().__init__() self.conv = nn.Conv2d( diff --git a/src/mflux/models/vae/encoder/conv_in.py b/src/mflux/models/vae/encoder/conv_in.py index 6384ed0..a999821 100644 --- a/src/mflux/models/vae/encoder/conv_in.py +++ b/src/mflux/models/vae/encoder/conv_in.py @@ -3,7 +3,6 @@ import mlx.nn as nn class ConvIn(nn.Module): - def __init__(self): super().__init__() self.conv2d = nn.Conv2d( diff --git a/src/mflux/models/vae/encoder/conv_norm_out.py b/src/mflux/models/vae/encoder/conv_norm_out.py index 83b211a..693a8b8 100644 --- a/src/mflux/models/vae/encoder/conv_norm_out.py +++ b/src/mflux/models/vae/encoder/conv_norm_out.py @@ -10,7 +10,7 @@ class ConvNormOut(nn.Module): dims=512, eps=1e-6, affine=True, - pytorch_compatible=True + pytorch_compatible=True, ) def forward(self, input_array: mx.array) -> mx.array: diff --git a/src/mflux/models/vae/encoder/down_block_1.py b/src/mflux/models/vae/encoder/down_block_1.py index 04b1b8d..482d8a4 100644 --- a/src/mflux/models/vae/encoder/down_block_1.py +++ b/src/mflux/models/vae/encoder/down_block_1.py @@ -6,7 +6,6 @@ from mflux.models.vae.encoder.down_sampler import DownSampler class DownBlock1(nn.Module): - def __init__(self): super().__init__() self.resnets = [ diff --git a/src/mflux/models/vae/encoder/down_block_2.py b/src/mflux/models/vae/encoder/down_block_2.py index 81542c9..b382409 100644 --- a/src/mflux/models/vae/encoder/down_block_2.py +++ b/src/mflux/models/vae/encoder/down_block_2.py @@ -6,7 +6,6 @@ from mflux.models.vae.encoder.down_sampler import DownSampler class DownBlock2(nn.Module): - def __init__(self): super().__init__() self.resnets = [ diff --git a/src/mflux/models/vae/encoder/down_block_3.py b/src/mflux/models/vae/encoder/down_block_3.py index 6526e46..6a66146 100644 --- a/src/mflux/models/vae/encoder/down_block_3.py +++ b/src/mflux/models/vae/encoder/down_block_3.py @@ -6,7 +6,6 @@ from mflux.models.vae.encoder.down_sampler import DownSampler class DownBlock3(nn.Module): - def __init__(self): super().__init__() self.resnets = [ diff --git a/src/mflux/models/vae/encoder/down_block_4.py b/src/mflux/models/vae/encoder/down_block_4.py index fbed7fe..c6926f1 100644 --- a/src/mflux/models/vae/encoder/down_block_4.py +++ b/src/mflux/models/vae/encoder/down_block_4.py @@ -5,7 +5,6 @@ from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D class DownBlock4(nn.Module): - def __init__(self): super().__init__() self.resnets = [ diff --git a/src/mflux/models/vae/encoder/down_sampler.py b/src/mflux/models/vae/encoder/down_sampler.py index 073a08e..b06cd6d 100644 --- a/src/mflux/models/vae/encoder/down_sampler.py +++ b/src/mflux/models/vae/encoder/down_sampler.py @@ -4,7 +4,6 @@ from mlx import nn class DownSampler(nn.Module): - def __init__(self, conv_in: int, conv_out: int): super().__init__() self.conv = nn.Conv2d( diff --git a/src/mflux/models/vae/encoder/encoder.py b/src/mflux/models/vae/encoder/encoder.py index 594a760..919dac7 100644 --- a/src/mflux/models/vae/encoder/encoder.py +++ b/src/mflux/models/vae/encoder/encoder.py @@ -13,7 +13,6 @@ from mflux.models.vae.encoder.down_block_4 import DownBlock4 class Encoder(nn.Module): - def __init__(self): super().__init__() self.conv_in = ConvIn() diff --git a/src/mflux/post_processing/generated_image.py b/src/mflux/post_processing/generated_image.py index 91aab40..1ba24b9 100644 --- a/src/mflux/post_processing/generated_image.py +++ b/src/mflux/post_processing/generated_image.py @@ -12,21 +12,20 @@ log = logging.getLogger(__name__) class GeneratedImage: - def __init__( - self, - image: PIL.Image.Image, - model_config: ModelConfig, - seed: int, - prompt: str, - steps: int, - guidance: float | None, - precision: mx.Dtype, - quantization: int, - generation_time: float, - lora_paths: list[str], - lora_scales: list[float], - controlnet_strength: float | None = None, + self, + image: PIL.Image.Image, + model_config: ModelConfig, + seed: int, + prompt: str, + steps: int, + guidance: float | None, + precision: mx.Dtype, + quantization: int, + generation_time: float, + lora_paths: list[str], + lora_scales: list[float], + controlnet_strength: float | None = None, ): self.image = image self.model_config = model_config @@ -61,7 +60,7 @@ class GeneratedImage: # Optionally save json metadata file if export_json_metadata: - with open(f"{file_path.with_suffix('.json')}", 'w') as json_file: + with open(f"{file_path.with_suffix('.json')}", "w") as json_file: json.dump(self._get_metadata(), json_file, indent=4) # Embed metadata @@ -79,33 +78,27 @@ class GeneratedImage: metadata_str = str(metadata) # Convert the string to bytes (using UTF-8 encoding) - user_comment_bytes = metadata_str.encode('utf-8') + user_comment_bytes = metadata_str.encode("utf-8") # Define the UserComment tag ID USER_COMMENT_TAG_ID = 0x9286 # Create an EXIF dictionary exif_dict = { - '0th': {}, - 'Exif': { - USER_COMMENT_TAG_ID: user_comment_bytes - }, - 'GPS': {}, - '1st': {}, - 'thumbnail': None + "0th": {}, + "Exif": {USER_COMMENT_TAG_ID: user_comment_bytes}, + "GPS": {}, + "1st": {}, + "thumbnail": None, } # Create a piexif-compatible dictionary structure - exif_piexif_dict = { - 'Exif': { - USER_COMMENT_TAG_ID: user_comment_bytes - } - } + exif_piexif_dict = {"Exif": {USER_COMMENT_TAG_ID: user_comment_bytes}} # Load the image and embed the EXIF data image = PIL.Image.open(path) exif_bytes = piexif.dump(exif_piexif_dict) - image.info['exif'] = exif_bytes + image.info["exif"] = exif_bytes # Save the image with metadata image.save(path, exif=exif_bytes) @@ -115,15 +108,15 @@ class GeneratedImage: def _get_metadata(self) -> dict: return { - 'model': str(self.model_config.alias), - 'seed': str(self.seed), - 'steps': str(self.steps), - 'guidance': "None" if self.model_config == ModelConfig.FLUX1_SCHNELL else str(self.guidance), - 'precision': f"{self.precision}", - 'quantization': "None" if self.quantization is None else f"{self.quantization} bit", - 'generation_time': f"{self.generation_time:.2f} seconds", - 'lora_paths': ', '.join(self.lora_paths) if self.lora_paths else '', - 'lora_scales': ', '.join([f"{scale:.2f}" for scale in self.lora_scales]) if self.lora_scales else '', - 'prompt': self.prompt, - 'controlnet_strength': "None" if self.controlnet_strength is None else f"{self.controlnet_strength:.2f}", + "model": str(self.model_config.alias), + "seed": str(self.seed), + "steps": str(self.steps), + "guidance": "None" if self.model_config == ModelConfig.FLUX1_SCHNELL else str(self.guidance), + "precision": f"{self.precision}", + "quantization": "None" if self.quantization is None else f"{self.quantization} bit", + "generation_time": f"{self.generation_time:.2f} seconds", + "lora_paths": ", ".join(self.lora_paths) if self.lora_paths else "", + "lora_scales": ", ".join([f"{scale:.2f}" for scale in self.lora_scales]) if self.lora_scales else "", + "prompt": self.prompt, + "controlnet_strength": "None" if self.controlnet_strength is None else f"{self.controlnet_strength:.2f}", } diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index 9764bc2..8f83356 100644 --- a/src/mflux/post_processing/image_util.py +++ b/src/mflux/post_processing/image_util.py @@ -9,17 +9,16 @@ from mflux.post_processing.generated_image import GeneratedImage class ImageUtil: - @staticmethod def to_image( - decoded_latents: mx.array, - seed: int, - prompt: str, - quantization: int, - generation_time: float, - lora_paths: list[str], - lora_scales: list[float], - config: RuntimeConfig, + decoded_latents: mx.array, + seed: int, + prompt: str, + quantization: int, + generation_time: float, + lora_paths: list[str], + lora_scales: list[float], + config: RuntimeConfig, ) -> GeneratedImage: normalized = ImageUtil._denormalize(decoded_latents) normalized_numpy = ImageUtil._to_numpy(normalized) @@ -73,7 +72,7 @@ class ImageUtil: array = mx.transpose(array, (0, 3, 1, 2)) array = ImageUtil._normalize(array) return array - + @staticmethod def load_image(path: str) -> Image.Image: return Image.open(path) diff --git a/src/mflux/save.py b/src/mflux/save.py index 49409ca..bbbaf83 100644 --- a/src/mflux/save.py +++ b/src/mflux/save.py @@ -4,10 +4,10 @@ from mflux import Flux1, ModelConfig def main(): - parser = argparse.ArgumentParser(description='Save a quantized version of Flux.1 to disk.') - parser.add_argument('--path', type=str, required=True, help='Local path for loading a model from disk') - parser.add_argument('--model', "-m", type=str, required=True, choices=["dev", "schnell"], help='The model to use ("schnell" or "dev").') - parser.add_argument('--quantize', "-q", type=int, choices=[4, 8], default=8, help='Quantize the model (4 or 8, Default is 8)') + parser = argparse.ArgumentParser(description="Save a quantized version of Flux.1 to disk.") + parser.add_argument("--path", type=str, required=True, help="Local path for loading a model from disk") + parser.add_argument("--model", "-m", type=str, required=True, choices=["dev", "schnell"], help="The model to use (\"schnell\" or \"dev\").") + parser.add_argument("--quantize", "-q", type=int, choices=[4, 8], default=8, help="Quantize the model (4 or 8, Default is 8)") args = parser.parse_args() @@ -23,5 +23,5 @@ def main(): print(f"Model saved at {args.path}\n") -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/src/mflux/tokenizer/t5_tokenizer.py b/src/mflux/tokenizer/t5_tokenizer.py index 2ce659f..b59c1f0 100644 --- a/src/mflux/tokenizer/t5_tokenizer.py +++ b/src/mflux/tokenizer/t5_tokenizer.py @@ -3,7 +3,6 @@ from transformers import T5Tokenizer class TokenizerT5: - def __init__(self, tokenizer: T5Tokenizer, max_length: int = 256): self.tokenizer = tokenizer self.max_length = max_length diff --git a/src/mflux/tokenizer/tokenizer_handler.py b/src/mflux/tokenizer/tokenizer_handler.py index 7640e6c..68ae3b6 100644 --- a/src/mflux/tokenizer/tokenizer_handler.py +++ b/src/mflux/tokenizer/tokenizer_handler.py @@ -7,19 +7,23 @@ from mflux.tokenizer.clip_tokenizer import TokenizerCLIP class TokenizerHandler: - - def __init__(self, repo_id: str, max_t5_length: int = 256, local_path: str | None = None): + def __init__( + self, + repo_id: str, + max_t5_length: int = 256, + local_path: str | None = None, + ): root_path = Path(local_path) if local_path else TokenizerHandler._download_or_get_cached_tokenizers(repo_id) self.clip = transformers.CLIPTokenizer.from_pretrained( pretrained_model_name_or_path=root_path / "tokenizer", local_files_only=True, - max_length=TokenizerCLIP.MAX_TOKEN_LENGTH + max_length=TokenizerCLIP.MAX_TOKEN_LENGTH, ) self.t5 = transformers.T5Tokenizer.from_pretrained( pretrained_model_name_or_path=root_path / "tokenizer_2", local_files_only=True, - max_length=max_t5_length + max_length=max_t5_length, ) @staticmethod @@ -27,9 +31,6 @@ class TokenizerHandler: return Path( snapshot_download( repo_id=repo_id, - allow_patterns=[ - "tokenizer/**", - "tokenizer_2/**" - ] + allow_patterns=["tokenizer/**", "tokenizer_2/**"], ) ) diff --git a/src/mflux/weights/lora_converter.py b/src/mflux/weights/lora_converter.py index 8e8a7d2..67ab6b3 100644 --- a/src/mflux/weights/lora_converter.py +++ b/src/mflux/weights/lora_converter.py @@ -10,8 +10,8 @@ logger = logging.getLogger(__name__) # This script is based on `convert_flux_lora.py` from `kohya-ss/sd-scripts`. # For more info, see: https://github.com/kohya-ss/sd-scripts/blob/sd3/networks/convert_flux_lora.py -class LoRAConverter: +class LoRAConverter: @staticmethod def load_weights(lora_path: str) -> dict: state_dict = LoRAConverter._load_pytorch_weights(lora_path) @@ -38,7 +38,7 @@ class LoRAConverter: source, target, f"lora_unet_double_blocks_{i}_img_attn_proj", - f"transformer.transformer_blocks.{i}.attn.to_out.0" + f"transformer.transformer_blocks.{i}.attn.to_out.0", ) LoRAConverter._convert_to_diffusers_cat( source, @@ -54,25 +54,25 @@ class LoRAConverter: source, target, f"lora_unet_double_blocks_{i}_img_mlp_0", - f"transformer.transformer_blocks.{i}.ff.net.0.proj" + f"transformer.transformer_blocks.{i}.ff.net.0.proj", ) LoRAConverter._convert_to_diffusers( source, target, f"lora_unet_double_blocks_{i}_img_mlp_2", - f"transformer.transformer_blocks.{i}.ff.net.2" + f"transformer.transformer_blocks.{i}.ff.net.2", ) LoRAConverter._convert_to_diffusers( source, target, f"lora_unet_double_blocks_{i}_img_mod_lin", - f"transformer.transformer_blocks.{i}.norm1.linear" + f"transformer.transformer_blocks.{i}.norm1.linear", ) LoRAConverter._convert_to_diffusers( source, target, f"lora_unet_double_blocks_{i}_txt_attn_proj", - f"transformer.transformer_blocks.{i}.attn.to_add_out" + f"transformer.transformer_blocks.{i}.attn.to_add_out", ) LoRAConverter._convert_to_diffusers_cat( source, @@ -88,19 +88,19 @@ class LoRAConverter: source, target, f"lora_unet_double_blocks_{i}_txt_mlp_0", - f"transformer.transformer_blocks.{i}.ff_context.net.0.proj" + f"transformer.transformer_blocks.{i}.ff_context.net.0.proj", ) LoRAConverter._convert_to_diffusers( source, target, f"lora_unet_double_blocks_{i}_txt_mlp_2", - f"transformer.transformer_blocks.{i}.ff_context.net.2" + f"transformer.transformer_blocks.{i}.ff_context.net.2", ) LoRAConverter._convert_to_diffusers( source, target, f"lora_unet_double_blocks_{i}_txt_mod_lin", - f"transformer.transformer_blocks.{i}.norm1_context.linear" + f"transformer.transformer_blocks.{i}.norm1_context.linear", ) for i in range(38): @@ -120,12 +120,13 @@ class LoRAConverter: source, target, f"lora_unet_single_blocks_{i}_linear2", - f"transformer.single_transformer_blocks.{i}.proj_out" + f"transformer.single_transformer_blocks.{i}.proj_out", ) LoRAConverter._convert_to_diffusers( source, - target, f"lora_unet_single_blocks_{i}_modulation_lin", - f"transformer.single_transformer_blocks.{i}.norm.linear" + target, + f"lora_unet_single_blocks_{i}_modulation_lin", + f"transformer.single_transformer_blocks.{i}.norm.linear", ) if len(source) > 0: @@ -133,12 +134,7 @@ class LoRAConverter: return target @staticmethod - def _convert_to_diffusers( - source: dict, - target: dict, - source_key: str, - target_key: str - ): + def _convert_to_diffusers(source: dict, target: dict, source_key: str, target_key: str): if source_key + ".lora_down.weight" not in source: return down_weight = source.pop(source_key + ".lora_down.weight") @@ -160,11 +156,11 @@ class LoRAConverter: @staticmethod def _convert_to_diffusers_cat( - source: dict, - target: dict, - source_key: str, - target_keys: list[str], - dims=None + source: dict, + target: dict, + source_key: str, + target_keys: list[str], + dims=None, ): if source_key + ".lora_down.weight" not in source: return @@ -204,7 +200,12 @@ class LoRAConverter: if j == k: continue is_sparse = is_sparse and torch.all( - up_weight[i: i + dims[j], k * diffusers_rank: (k + 1) * diffusers_rank] == 0) + up_weight[ + i : i + dims[j], + k * diffusers_rank : (k + 1) * diffusers_rank, + ] + == 0 + ) i += dims[j] if is_sparse: logger.info(f"weight is sparse: {source_key}") @@ -220,12 +221,23 @@ class LoRAConverter: target.update({k: v for k, v in zip(diffusers_up_keys, torch.split(up_weight, dims, dim=0))}) else: # down_weight is chunked to each split - target.update({k: v for k, v in zip(diffusers_down_keys, torch.chunk(down_weight, num_splits, dim=0))}) + target.update( + { + k: v + for k, v in zip( + diffusers_down_keys, + torch.chunk(down_weight, num_splits, dim=0), + ) + } + ) # up_weight is sparse: only non-zero values are copied to each split i = 0 for j in range(len(dims)): - target[diffusers_up_keys[j]] = up_weight[i: i + dims[j], j * diffusers_rank: (j + 1) * diffusers_rank].contiguous() + target[diffusers_up_keys[j]] = up_weight[ + i : i + dims[j], + j * diffusers_rank : (j + 1) * diffusers_rank, + ].contiguous() i += dims[j] @staticmethod diff --git a/src/mflux/weights/lora_util.py b/src/mflux/weights/lora_util.py index 11bf3d2..dcd9f47 100644 --- a/src/mflux/weights/lora_util.py +++ b/src/mflux/weights/lora_util.py @@ -6,9 +6,12 @@ log = logging.getLogger(__name__) class LoraUtil: - @staticmethod - def apply_loras(transformer: dict, lora_files: list[str], lora_scales: list[float] | None = None) -> None: + 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): @@ -29,6 +32,7 @@ class LoraUtil: @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) @@ -54,14 +58,14 @@ class LoraUtil: visiting.append(splitKey) else: parentKey = ".".join(visiting) - if parentKey in visited and 'lora_A' in visited[parentKey] and 'lora_B' in visited[parentKey]: + 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'] + target["net"] = list({}) + target = target["net"] elif splitKey == "0": target.append({}) target = target[0] @@ -74,11 +78,11 @@ class LoraUtil: if parentKey not in visited: visited[parentKey] = {} visited[parentKey][splitKey] = weight - if not 'weight' in target: + if not "weight" 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'] + 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 + target["weight"] = weight diff --git a/src/mflux/weights/model_saver.py b/src/mflux/weights/model_saver.py index 0e7d372..da133e9 100644 --- a/src/mflux/weights/model_saver.py +++ b/src/mflux/weights/model_saver.py @@ -7,7 +7,6 @@ from transformers import CLIPTokenizer, T5Tokenizer class ModelSaver: - @staticmethod def save_model(model, bits: int, base_path: str): # Save the tokenizers @@ -32,7 +31,11 @@ class ModelSaver: path.mkdir(parents=True, exist_ok=True) weights = ModelSaver._split_weights(base_path, dict(tree_flatten(model.parameters()))) for i, weight in enumerate(weights): - mx.save_safetensors(str(path / f"{i}.safetensors"), weight, {"quantization_level": str(bits)}) + mx.save_safetensors( + str(path / f"{i}.safetensors"), + weight, + {"quantization_level": str(bits)}, + ) @staticmethod def _split_weights(base_path: str, weights: dict, max_file_size_gb: int = 2) -> list: diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py index dbd25f1..56d25cc 100644 --- a/src/mflux/weights/weight_handler.py +++ b/src/mflux/weights/weight_handler.py @@ -11,13 +11,12 @@ from mflux.weights.weight_util import WeightUtil 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, + self, + repo_id: str | None = None, + local_path: str | None = None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, ): root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id) @@ -66,7 +65,7 @@ class WeightHandler: weights, quantization_level = WeightHandler._get_weights("transformer", root_path, lora_path) if lora_path: - if 'transformer' not in weights: + if "transformer" not in weights: weights = LoRAConverter.load_weights(lora_path) weights = weights["transformer"] @@ -80,18 +79,23 @@ class WeightHandler: if block.get("ff") is not None: block["ff"] = { "linear1": block["ff"]["net"][0]["proj"], - "linear2": block["ff"]["net"][2] + "linear2": block["ff"]["net"][2], } if block.get("ff_context") is not None: block["ff_context"] = { "linear1": block["ff_context"]["net"][0]["proj"], - "linear2": block["ff_context"]["net"][2] + "linear2": block["ff_context"]["net"][2], } return weights, quantization_level - + @staticmethod def load_controlnet_transformer(controlnet_id: str) -> (dict, int): - controlnet_path = Path(snapshot_download(repo_id=controlnet_id,allow_patterns=["*.safetensors","config.json"])) + controlnet_path = Path( + snapshot_download( + repo_id=controlnet_id, + allow_patterns=["*.safetensors", "config.json"], + ) + ) file = next(controlnet_path.glob("diffusion_pytorch_model.safetensors")) quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level") weights = list(mx.load(str(file)).items()) @@ -112,12 +116,12 @@ class WeightHandler: for block in weights["transformer_blocks"]: block["ff"] = { "linear1": block["ff"]["net"][0]["proj"], - "linear2": block["ff"]["net"][2] + "linear2": block["ff"]["net"][2], } if block.get("ff_context") is not None: block["ff_context"] = { "linear1": block["ff_context"]["net"][0]["proj"], - "linear2": block["ff_context"]["net"][2] + "linear2": block["ff_context"]["net"][2], } config = json.load(open(controlnet_path / "config.json")) return weights, quantization_level, config @@ -131,16 +135,20 @@ class WeightHandler: return weights, quantization_level # Reshape and process the huggingface weights - weights['decoder']['conv_in'] = {'conv2d': weights['decoder']['conv_in']} - weights['decoder']['conv_out'] = {'conv2d': weights['decoder']['conv_out']} - weights['decoder']['conv_norm_out'] = {'norm': weights['decoder']['conv_norm_out']} - weights['encoder']['conv_in'] = {'conv2d': weights['encoder']['conv_in']} - weights['encoder']['conv_out'] = {'conv2d': weights['encoder']['conv_out']} - weights['encoder']['conv_norm_out'] = {'norm': weights['encoder']['conv_norm_out']} + weights["decoder"]["conv_in"] = {"conv2d": weights["decoder"]["conv_in"]} + weights["decoder"]["conv_out"] = {"conv2d": weights["decoder"]["conv_out"]} + weights["decoder"]["conv_norm_out"] = {"norm": weights["decoder"]["conv_norm_out"]} + weights["encoder"]["conv_in"] = {"conv2d": weights["encoder"]["conv_in"]} + weights["encoder"]["conv_out"] = {"conv2d": weights["encoder"]["conv_out"]} + weights["encoder"]["conv_norm_out"] = {"norm": weights["encoder"]["conv_norm_out"]} return weights, quantization_level @staticmethod - def _get_weights(model_name: str, root_path: Path | None = None, lora_path: str | None = None) -> (dict, int): + def _get_weights( + model_name: str, + root_path: Path | None = None, + lora_path: str | None = None, + ) -> (dict, int): weights = [] quantization_level = None @@ -174,6 +182,6 @@ class WeightHandler: "text_encoder_2/*.safetensors", "transformer/*.safetensors", "vae/*.safetensors", - ] + ], ) ) diff --git a/src/mflux/weights/weight_util.py b/src/mflux/weights/weight_util.py index dd16ad3..3cd0625 100644 --- a/src/mflux/weights/weight_util.py +++ b/src/mflux/weights/weight_util.py @@ -2,7 +2,6 @@ from mflux.config.config import Config class WeightUtil: - @staticmethod def flatten(params): return [(k, v) for p in params for (k, v) in p] From fd0add32dd0e17b7ffb46e96267cbdf3ed2e479c Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 21 Sep 2024 15:41:46 +0200 Subject: [PATCH 4/5] Update ruff setup: 120 character line-length and exclude files --- pyproject.toml | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5d79359..5cca839 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,13 +48,27 @@ where = ["src"] include = ["mflux*"] [tool.ruff] -line-length = 88 # same as 'black' default +line-length = 120 indent-width = 4 target-version = "py310" respect-gitignore = true +exclude = [ + "src/mflux/generate.py", + "src/mflux/generate_controlnet.py", + "src/mflux/save.py", + "src/mflux/flux/flux.py", + "src/mflux/controlnet/flux_controlnet.py", + "src/mflux/models/transformer/joint_transformer_block.py", + "src/mflux/models/vae/common/resnet_block_2d.py", + "src/mflux/models/vae/common/unet_mid_block.py", + "src/mflux/models/vae/encoder/down_block_2.py", + "src/mflux/models/vae/encoder/down_block_3.py", + "src/mflux/models/vae/decoder/up_block_3.py", + "src/mflux/models/vae/decoder/up_block_4.py" + ] [tool.ruff.lint] -# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`) codes by default. +# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`) codes by default. # Unlike Flake8, Ruff doesn't enable pycodestyle warnings (`W`) or # McCabe complexity (`C901`) by default. select = ["E4", "E7", "E9", "F"] From 6ef9ceb259579048d9a109c3a4ca8c2d1557523a Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 21 Sep 2024 19:51:22 +0200 Subject: [PATCH 5/5] Fix linting issues --- src/mflux/models/transformer/joint_attention.py | 2 -- src/mflux/post_processing/generated_image.py | 9 --------- src/mflux/weights/lora_converter.py | 1 - src/mflux/weights/lora_util.py | 2 +- 4 files changed, 1 insertion(+), 13 deletions(-) diff --git a/src/mflux/models/transformer/joint_attention.py b/src/mflux/models/transformer/joint_attention.py index 7c7cd80..6ea3f62 100644 --- a/src/mflux/models/transformer/joint_attention.py +++ b/src/mflux/models/transformer/joint_attention.py @@ -28,8 +28,6 @@ class JointAttention(nn.Module): encoder_hidden_states: mx.array, image_rotary_emb: mx.array, ) -> (mx.array, mx.array): - residual = hidden_states - query = self.to_q(hidden_states) key = self.to_k(hidden_states) value = self.to_v(hidden_states) diff --git a/src/mflux/post_processing/generated_image.py b/src/mflux/post_processing/generated_image.py index 1ba24b9..ad6c694 100644 --- a/src/mflux/post_processing/generated_image.py +++ b/src/mflux/post_processing/generated_image.py @@ -83,15 +83,6 @@ class GeneratedImage: # Define the UserComment tag ID USER_COMMENT_TAG_ID = 0x9286 - # Create an EXIF dictionary - exif_dict = { - "0th": {}, - "Exif": {USER_COMMENT_TAG_ID: user_comment_bytes}, - "GPS": {}, - "1st": {}, - "thumbnail": None, - } - # Create a piexif-compatible dictionary structure exif_piexif_dict = {"Exif": {USER_COMMENT_TAG_ID: user_comment_bytes}} diff --git a/src/mflux/weights/lora_converter.py b/src/mflux/weights/lora_converter.py index 67ab6b3..9570a04 100644 --- a/src/mflux/weights/lora_converter.py +++ b/src/mflux/weights/lora_converter.py @@ -25,7 +25,6 @@ class LoRAConverter: def _load_pytorch_weights(lora_path: str) -> dict: state_dict = {} with safe_open(lora_path, framework="pt") as f: - metadata = f.metadata() for k in f.keys(): state_dict[k] = f.get_tensor(k) return state_dict diff --git a/src/mflux/weights/lora_util.py b/src/mflux/weights/lora_util.py index dcd9f47..9145914 100644 --- a/src/mflux/weights/lora_util.py +++ b/src/mflux/weights/lora_util.py @@ -78,7 +78,7 @@ class LoraUtil: if parentKey not in visited: visited[parentKey] = {} visited[parentKey][splitKey] = weight - if not "weight" in target: + 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"]