diff --git a/pyproject.toml b/pyproject.toml index 5cca839..7e4fd83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,17 +12,17 @@ authors = [{ name = "Filip Strand", email = "strand.filip@gmail.com" }] maintainers = [{ name = "Filip Strand", email = "strand.filip@gmail.com" }] requires-python = ">=3.10" dependencies = [ - "huggingface-hub>=0.24.5", - "mlx>=0.16.0", - "numpy>=2.0.1", - "opencv-python>=4.10.0", - "piexif>=1.1.3", - "pillow>=10.4.0", - "safetensors>=0.4.4", - "sentencepiece>=0.2.0", - "torch>=2.3.1", - "tqdm>=4.66.5", - "transformers>=4.44.0", + "huggingface-hub>=0.24.5,<1.0", + "mlx>=0.16.0,<1.0", + "numpy>=2.0.1,<3.0", + "opencv-python>=4.10.0,<5.0", + "piexif>=1.1.3,<2.0", + "pillow>=10.4.0,<11.0", + "safetensors>=0.4.4,<1.0", + "sentencepiece>=0.2.0,<1.0", + "torch>=2.3.1,<3.0", + "tqdm>=4.66.5,<5.0", + "transformers>=4.44.0,<5.0", ] classifiers = [ "Framework :: MLX", @@ -52,20 +52,6 @@ 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. diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index 3c342d1..c9e5b1c 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -58,7 +58,7 @@ class Flux1Controlnet: 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: @@ -68,16 +68,20 @@ class Flux1Controlnet: 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, ctrlnet_quantization_level, controlnet_config = WeightHandler.load_controlnet_transformer(controlnet_id=CONTROLNET_ID) + weights_controlnet, ctrlnet_quantization_level, controlnet_config = WeightHandler.load_controlnet_transformer( + controlnet_id=CONTROLNET_ID + ) self.transformer_controlnet = TransformerControlnet( model_config=model_config, num_blocks=controlnet_config["num_layers"], @@ -90,25 +94,29 @@ class Flux1Controlnet: self.bits = None if quantize is not None or ctrlnet_quantization_level is not None: self.bits = ctrlnet_quantization_level if ctrlnet_quantization_level is not None else quantize + # fmt: off nn.quantize(self.transformer_controlnet, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 128, group_size=128, bits=self.bits) + # fmt: on if ctrlnet_quantization_level is not None: self.transformer_controlnet.update(weights_controlnet) - def generate_image(self, seed: int, prompt: str, control_image: PIL.Image.Image, config: ConfigControlnet = ConfigControlnet()) -> GeneratedImage: + def generate_image(self, seed: int, prompt: str, control_image: PIL.Image.Image, config: ConfigControlnet = ConfigControlnet()) -> GeneratedImage: # fmt: off # Create a new runtime config based on the model type and input parameters config = RuntimeConfig(config, self.model_config) time_steps = tqdm(range(config.num_inference_steps)) if config.height != control_image.height or config.width != control_image.width: - log.warning(f"Control image has different dimensions than the model. Resizing to {config.width}x{config.height}") + log.warning( + f"Control image has different dimensions than the model. Resizing to {config.width}x{config.height}" + ) control_image = control_image.resize((config.width, config.height), PIL.Image.LANCZOS) # 1. Create the initial latents latents = mx.random.normal( shape=[1, (config.height // 16) * (config.width // 16), 64], key=mx.random.key(seed) - ) + ) # fmt: off control_image = ControlnetUtil.preprocess_canny(control_image) controlnet_cond = ImageUtil.to_array(control_image) controlnet_cong = self.vae.encode(controlnet_cond) diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index 55afdea..2b05e26 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -48,7 +48,7 @@ class Flux1: 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: @@ -58,10 +58,12 @@ class Flux1: 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: @@ -76,7 +78,7 @@ class Flux1: latents = mx.random.normal( shape=[1, (config.height // 16) * (config.width // 16), 64], key=mx.random.key(seed) - ) + ) # fmt: off # 2. Embedd the prompt t5_tokens = self.t5_tokenizer.tokenize(prompt) diff --git a/src/mflux/generate.py b/src/mflux/generate.py index 6043a40..ded7371 100644 --- a/src/mflux/generate.py +++ b/src/mflux/generate.py @@ -5,6 +5,7 @@ from mflux import Flux1, Config, ModelConfig def main(): + # fmt: off parser = argparse.ArgumentParser(description="Generate an image based on a prompt.") parser.add_argument("--prompt", type=str, required=True, help="The textual description of the image to generate.") parser.add_argument("--output", type=str, default="image.png", help="The filename for the output image. Default is \"image.png\".") @@ -19,6 +20,7 @@ def main(): 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.") + # fmt: on args = parser.parse_args() diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index 9f2ab64..257b806 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/generate_controlnet.py @@ -5,6 +5,7 @@ from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig, ImageUtil def main(): + # fmt: off parser = argparse.ArgumentParser(description="Generate an image based on a prompt.") parser.add_argument("--prompt", type=str, required=True, help="The textual description of the image to generate.") parser.add_argument("--control-image-path", type=str, required=True, help="Local path of the image to use as input for controlnet.") @@ -21,6 +22,7 @@ def main(): 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.") + # fmt: on args = parser.parse_args() diff --git a/src/mflux/models/transformer/joint_transformer_block.py b/src/mflux/models/transformer/joint_transformer_block.py index d6f2349..e347761 100644 --- a/src/mflux/models/transformer/joint_transformer_block.py +++ b/src/mflux/models/transformer/joint_transformer_block.py @@ -25,11 +25,12 @@ class JointTransformerBlock(nn.Module): 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.forward( + hidden_states, text_embeddings + ) norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context.forward( - x=encoder_hidden_states, - text_embeddings=text_embeddings + x=encoder_hidden_states, text_embeddings=text_embeddings ) attn_output, context_attn_output = self.attn.forward( diff --git a/src/mflux/models/vae/common/resnet_block_2d.py b/src/mflux/models/vae/common/resnet_block_2d.py index 067c2ff..2c60335 100644 --- a/src/mflux/models/vae/common/resnet_block_2d.py +++ b/src/mflux/models/vae/common/resnet_block_2d.py @@ -4,6 +4,7 @@ from mlx import nn from mflux.config.config import Config +# fmt: off class ResnetBlock2D(nn.Module): def __init__( self, diff --git a/src/mflux/models/vae/decoder/up_block_3.py b/src/mflux/models/vae/decoder/up_block_3.py index fd51799..5959b94 100644 --- a/src/mflux/models/vae/decoder/up_block_3.py +++ b/src/mflux/models/vae/decoder/up_block_3.py @@ -8,11 +8,13 @@ from mflux.models.vae.decoder.up_sampler import UpSampler class UpBlock3(nn.Module): def __init__(self): super().__init__() + # fmt: off self.resnets = [ ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256, is_conv_shortcut=True, conv_shortcut_in=512, conv_shortcut_out=256), ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256), ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256), ] + # fmt: on self.upsamplers = [UpSampler(conv_in=256, conv_out=256)] def forward(self, input_array: mx.array) -> mx.array: diff --git a/src/mflux/models/vae/decoder/up_block_4.py b/src/mflux/models/vae/decoder/up_block_4.py index 4b59341..dd78356 100644 --- a/src/mflux/models/vae/decoder/up_block_4.py +++ b/src/mflux/models/vae/decoder/up_block_4.py @@ -7,11 +7,13 @@ from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D class UpBlock4(nn.Module): def __init__(self): super().__init__() + # fmt: off self.resnets = [ ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128, is_conv_shortcut=True, conv_shortcut_in=256, conv_shortcut_out=128), ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128), ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128), ] + # fmt: off def forward(self, input_array: mx.array) -> mx.array: hidden_states = self.resnets[0].forward(input_array) diff --git a/src/mflux/models/vae/encoder/down_block_1.py b/src/mflux/models/vae/encoder/down_block_1.py index 482d8a4..ce5116b 100644 --- a/src/mflux/models/vae/encoder/down_block_1.py +++ b/src/mflux/models/vae/encoder/down_block_1.py @@ -9,8 +9,10 @@ class DownBlock1(nn.Module): def __init__(self): super().__init__() self.resnets = [ + # fmt: off ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128), ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128), + # fmt: on ] self.downsamplers = [DownSampler(conv_in=128, conv_out=128)] diff --git a/src/mflux/models/vae/encoder/down_block_2.py b/src/mflux/models/vae/encoder/down_block_2.py index b382409..afa05a5 100644 --- a/src/mflux/models/vae/encoder/down_block_2.py +++ b/src/mflux/models/vae/encoder/down_block_2.py @@ -8,10 +8,12 @@ from mflux.models.vae.encoder.down_sampler import DownSampler class DownBlock2(nn.Module): def __init__(self): super().__init__() + # fmt: off self.resnets = [ ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256, is_conv_shortcut=True, conv_shortcut_in=128, conv_shortcut_out=256), ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256), ] + # fmt: on self.downsamplers = [DownSampler(conv_in=256, conv_out=256)] def forward(self, input_array: mx.array) -> mx.array: diff --git a/src/mflux/models/vae/encoder/down_block_3.py b/src/mflux/models/vae/encoder/down_block_3.py index 6a66146..c1cac99 100644 --- a/src/mflux/models/vae/encoder/down_block_3.py +++ b/src/mflux/models/vae/encoder/down_block_3.py @@ -8,10 +8,12 @@ from mflux.models.vae.encoder.down_sampler import DownSampler class DownBlock3(nn.Module): def __init__(self): super().__init__() + # fmt: off self.resnets = [ ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512, is_conv_shortcut=True, conv_shortcut_in=256, conv_shortcut_out=512), ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512), ] + # fmt: on self.downsamplers = [DownSampler(conv_in=512, conv_out=512)] def forward(self, input_array: mx.array) -> mx.array: diff --git a/src/mflux/models/vae/encoder/down_block_4.py b/src/mflux/models/vae/encoder/down_block_4.py index c6926f1..45c917a 100644 --- a/src/mflux/models/vae/encoder/down_block_4.py +++ b/src/mflux/models/vae/encoder/down_block_4.py @@ -7,10 +7,12 @@ from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D class DownBlock4(nn.Module): def __init__(self): super().__init__() + # fmt: off 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), ] + # fmt: on def forward(self, input_array: mx.array) -> mx.array: hidden_states = self.resnets[0].forward(input_array) diff --git a/src/mflux/save.py b/src/mflux/save.py index bbbaf83..8bd4a6f 100644 --- a/src/mflux/save.py +++ b/src/mflux/save.py @@ -4,10 +4,12 @@ from mflux import Flux1, ModelConfig def main(): + # fmt: off 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)") + # fmt: on args = parser.parse_args()