Merge pull request #61 from anthonywu/qol-followup

quality of life maintenance for ruff/pre-commit

Makefile

@@ -23,6 +23,14 @@ expect-uv:
 		echo "Please install uv to continue:\n    https://github.com/astral-sh/uv?tab=readme-ov-file#installation"; \
 	fi
 
+# assume reasonably pre-commit is a safe dependency given its wide support (e.g. GitHub Actions integration)
+.PHONY: ensure-pre-commit
+	@if ! /usr/bin/which -s pre-commit; then \
+		echo "pre-commit required for submitting commits before pull requests. Using uv tool to install pre-commit."; \
+		uv tool install pre-commit; \
+	fi
+
+
 # assume reasonably that if user has installed uv, they trust ruff from the same team
 .PHONY: ensure-ruff
 ensure-ruff:
@@ -41,16 +49,18 @@ venv-init: expect-arm64 expect-uv
 
 # Install dependencies
 .PHONY: install
-install: venv-init
-	# 🏗️ Installing dependencies...
+install: venv-init ensure-pre-commit
+	# 🏗️ Installing dependencies and pre-commit hooks...
 	uv pip install -e .
 	# ✅ Dependencies installed.
+	pre-commit install
+	# ✅ Pre-commit hooks installed.
 
 # Run linters
 .PHONY: lint
 lint: ensure-ruff
 	# 🏗️ Running linters, your files will not be mutated.
-	# Use 'make autofix' to auto-apply fixes."
+	# Use 'make check' to auto-apply fixes."
 	ruff check
 	# ✅ Linting complete."
 
@@ -63,15 +73,11 @@ format: ensure-ruff
 	git diff --stat
 	# ✅ Formatting complete. Please review your git diffs, if any.
 
-# Run ruff auto lint and format
-# use || true to force possibly failing format/lint checks to allow rest of make helper to continue
-.PHONY: autofix
-autofix: ensure-ruff
-	# 🏗️ Running linter and formatters on files...
-	@(ruff format --respect-gitignore --target-version py310 || true)
-	# ✅ auto formatting complete
-	@(ruff check --fix || true)
-	# ✅ auto lint fixes complete, check for issues ruff cannot auto fix
+# Run ruff auto lint and format via pre-commit hook
+.PHONY: check
+check: ensure-ruff
+	# 🏗️ Running pre-commit linter and formatters on files...
+	@(pre-commit run --all-files)
 
 # Run tests
 .PHONY: test
@@ -96,7 +102,7 @@ help:
 	@echo "  make install     - Install project dev dependencies"
 	@echo "  make lint        - Run ruff python linter"
 	@echo "  make format      - Run ruff code formatter"
-	@echo "  make autofix     - Run linters auto fixes *and* style formatter"
+	@echo "  make check       - Run linters auto fixes *and* style formatter via pre-commit hook"
 	@echo "  make test        - Run tests"
 	@echo "  make clean       - Remove the virtual environment"
 	@echo "  make help        - Show this help message"

README.md

@@ -53,7 +53,7 @@ pip install -U mflux
 3. Follow format and lint checks prior to submitting Pull Requests. The recommended `make lint` and `make format` installs and uses [`ruff`](https://github.com/astral-sh/ruff). You can setup your editor/IDE to lint/format automatically, or use our provided `make` helpers:
   - `make format` - formats your code
   - `make lint` - shows your lint errors and warnings, but does not auto fix
-  - `make autofix` - formats your code **and** attempts to auto fix lint errors
+  - `make check` - via `pre-commit` hooks, formats your code **and** attempts to auto fix lint errors
   - consult official [`ruff` documentation](https://docs.astral.sh/ruff/) on advanced usages
 
 </details>

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.

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)

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)

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()
 

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()
 

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(

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,

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:

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)

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)]
 

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:

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:

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)

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()