Formatting

Author: ZachNagengast

.flake8

@@ -1,5 +1,6 @@
 [flake8]
 max-line-length = 120
+extend-ignore = E203
 filename = *.py
 
 [isort]

.gitignore

@@ -27,6 +27,7 @@ __pycache__/
 # Distribution / packaging
 .Python
 build/
+.build/
 develop-eggs/
 dist/
 downloads/

.pre-commit-config.yaml

@@ -11,6 +11,7 @@ repos:
       - id: black
         name: black
         language: python
+        args: ["--config", ".flake8"]
 
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.5.0

python/src/diffusionkit/mlx/__init__.py

@@ -39,7 +39,7 @@
     "sd3-8b-unreleased": "models/sd3_8b_beta.safetensors",  # unreleased
     "argmaxinc/mlx-FLUX.1-schnell": "argmaxinc/mlx-FLUX.1-schnell",
     "argmaxinc/mlx-FLUX.1-schnell-4bit-quantized": "argmaxinc/mlx-FLUX.1-schnell-4bit-quantized",
-    "argmaxinc/mlx-FLUX.1-dev": "argmaxinc/mlx-FLUX.1-dev"
+    "argmaxinc/mlx-FLUX.1-dev": "argmaxinc/mlx-FLUX.1-dev",
 }
 
 T5_MAX_LENGTH = {
@@ -49,6 +49,7 @@
     "argmaxinc/mlx-FLUX.1-dev": 512,
 }
 
+
 class DiffusionKitInferenceContext(AppleSiliconContextMixin, InferenceContextSpec):
     def code_spec(self):
         return {}

python/src/diffusionkit/mlx/config.py

@@ -105,7 +105,7 @@ def hidden_size(self) -> int:
     use_qk_norm=True,
     float16_dtype=mx.bfloat16,
     guidance_embed=True,  # Add this line
-    dtype=mx.bfloat16
+    dtype=mx.bfloat16,
 )
 
 

python/src/diffusionkit/mlx/mmdit.py

@@ -29,7 +29,9 @@ def __init__(self, config: MMDiTConfig):
         self.config = config
 
         if config.guidance_embed:
-            self.guidance_in = MLPEmbedder(in_dim=config.frequency_embed_dim, hidden_dim=config.hidden_size)
+            self.guidance_in = MLPEmbedder(
+                in_dim=config.frequency_embed_dim, hidden_dim=config.hidden_size
+            )
         else:
             self.guidance_in = nn.Identity()
 
@@ -939,6 +941,7 @@ def apply(q_or_k: mx.array, rope: mx.array) -> mx.array:
             .flatten(-2)
         )
 
+
 class MLPEmbedder(nn.Module):
     def __init__(self, in_dim: int, hidden_dim: int):
         super().__init__()

python/src/diffusionkit/mlx/scripts/generate_images.py

@@ -111,7 +111,7 @@ def cli():
     args.a16 = True
 
     if "FLUX" in args.model_version and args.cfg > 0.0:
-        logger.warning("Disabling CFG for FLUX.1-schnell model.")
+        logger.warning(f"Disabling CFG for {args.model_version} model.")
         args.cfg = 0.0
 
     if args.benchmark_mode: