Merge pull request #58 from anthonywu/intro-ruff-config

introduce standardized `ruff check` and `ruff format` config

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

pyproject.toml

@@ -46,3 +46,64 @@ mflux-generate-controlnet = "mflux.generate_controlnet:main"
 [tool.setuptools.packages.find]
 where = ["src"]
 include = ["mflux*"]
+
+[tool.ruff]
+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.
+# 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"

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",
+]

src/mflux/config/config.py

@@ -9,11 +9,11 @@ class Config:
     precision: mx.Dtype = mx.bfloat16
 
     def __init__(
-            self,
+        self,
-            num_inference_steps: int = 4,
+        num_inference_steps: int = 4,
-            width: int = 1024,
+        width: int = 1024,
-            height: int = 1024,
+        height: int = 1024,
-            guidance: float = 4.0,
+        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,
+        self,
-            num_inference_steps: int = 4,
+        num_inference_steps: int = 4,
-            width: int = 1024,
+        width: int = 1024,
-            height: int = 1024,
+        height: int = 1024,
-            guidance: float = 4.0,
+        guidance: float = 4.0,
-            controlnet_strength: float = 1.0,
+        controlnet_strength: float = 1.0,
     ):
         super().__init__(num_inference_steps, width, height, guidance)
         self.controlnet_strength = controlnet_strength

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,
+        self,
-            model_name: str,
+        model_name: str,
-            alias: str,
+        alias: str,
-            num_train_steps: int,
+        num_train_steps: int,
-            max_sequence_length: int,
+        max_sequence_length: int,
     ):
         self.alias = alias
         self.model_name = model_name

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

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)

src/mflux/controlnet/flux_controlnet.py

@@ -29,13 +29,13 @@ CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny"
 
 class Flux1Controlnet:
     def __init__(
-            self,
+        self,
-            model_config: ModelConfig,
+        model_config: ModelConfig,
-            quantize: int | None = None,
+        quantize: int | None = None,
-            local_path: str | None = None,
+        local_path: str | None = None,
-            lora_paths: list[str] | None = None,
+        lora_paths: list[str] | None = None,
-            lora_scales: list[float] | None = None,
+        lora_scales: list[float] | None = None,
-            controlnet_path: str | None = None,
+        controlnet_path: str | None = None,
     ):
         self.lora_paths = lora_paths
         self.lora_scales = lora_scales
@@ -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,

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.joint_transformer_block import (
-from mflux.models.transformer.single_transformer_block import SingleTransformerBlock
+    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,
+        self,
-            model_config: ModelConfig,
+        model_config: ModelConfig,
-            num_blocks: int,
+        num_blocks: int,
-            num_single_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,
+        self,
-            t: int,
+        t: int,
-            prompt_embeds: mx.array,
+        prompt_embeds: mx.array,
-            pooled_prompt_embeds: mx.array,
+        pooled_prompt_embeds: mx.array,
-            hidden_states: mx.array,
+        hidden_states: mx.array,
-            controlnet_cond: mx.array,
+        controlnet_cond: mx.array,
-            config: RuntimeConfig,
+        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):

src/mflux/flux/flux.py

@@ -19,14 +19,13 @@ from mflux.weights.weight_handler import WeightHandler
 
 
 class Flux1:
-
     def __init__(
-            self,
+        self,
-            model_config: ModelConfig,
+        model_config: ModelConfig,
-            quantize: int | None = None,
+        quantize: int | None = None,
-            local_path: str | None = None,
+        local_path: str | None = None,
-            lora_paths: list[str] | None = None,
+        lora_paths: list[str] | None = None,
-            lora_scales: list[float] | 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,

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 = 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("--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("--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("--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("--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("--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("--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("--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("--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("--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("--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-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("--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.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()

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 = 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("--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("--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("--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("--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("--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("--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("--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("--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("--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("--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("--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("--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-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("--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.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()

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)

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)

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

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)

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

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)

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 = encoder_layer.forward(hidden_states, causal_attention_mask)
-                hidden_states,
-                causal_attention_mask
-            )
         return hidden_states

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

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

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

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)

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

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

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)

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]
+        scale = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
-        shift = text_embeddings[:, 1*chunk_size:2*chunk_size]
+        shift = text_embeddings[:, 1 * chunk_size : 2 * chunk_size]
         x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
         return x
-

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]
+        shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
-        scale_msa = text_embeddings[:, 1*chunk_size:2*chunk_size]
+        scale_msa = text_embeddings[:, 1 * chunk_size : 2 * chunk_size]
-        gate_msa = text_embeddings[:, 2*chunk_size:3*chunk_size]
+        gate_msa = text_embeddings[:, 2 * chunk_size : 3 * chunk_size]
-        shift_mlp = text_embeddings[:, 3*chunk_size:4*chunk_size]
+        shift_mlp = text_embeddings[:, 3 * chunk_size : 4 * chunk_size]
-        scale_mlp = text_embeddings[:, 4*chunk_size:5*chunk_size]
+        scale_mlp = text_embeddings[:, 4 * chunk_size : 5 * chunk_size]
-        gate_mlp = text_embeddings[:, 5*chunk_size:6*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

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]
+        shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
-        scale_msa = text_embeddings[:, 1*chunk_size:2*chunk_size]
+        scale_msa = text_embeddings[:, 1 * chunk_size : 2 * chunk_size]
-        gate_msa = text_embeddings[:, 2*chunk_size:3*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

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)

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)

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)

src/mflux/models/transformer/joint_attention.py

@@ -23,13 +23,11 @@ class JointAttention(nn.Module):
         self.norm_added_k = nn.RMSNorm(128)
 
     def forward(
-            self,
+        self,
-            hidden_states: mx.array,
+        hidden_states: mx.array,
-            encoder_hidden_states: mx.array,
+        encoder_hidden_states: mx.array,
-            image_rotary_emb: 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)
@@ -45,9 +43,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_query_proj = mx.transpose(
-        encoder_hidden_states_key_proj = mx.transpose(mx.reshape(encoder_hidden_states_key_proj, (1, -1, 24, 128)), (0, 2, 1, 3))
+            mx.reshape(encoder_hidden_states_query_proj, (1, -1, 24, 128)),
-        encoder_hidden_states_value_proj = mx.transpose(mx.reshape(encoder_hidden_states_value_proj, (1, -1, 24, 128)),   (0, 2, 1, 3))
+            (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 +67,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 +86,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

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,
+        self,
-            hidden_states: mx.array,
+        hidden_states: mx.array,
-            encoder_hidden_states: mx.array,
+        encoder_hidden_states: mx.array,
-            text_embeddings: mx.array,
+        text_embeddings: mx.array,
-            rotary_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)
 

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(
+    def forward(self, hidden_states: mx.array, image_rotary_emb: mx.array) -> (mx.array, mx.array):
-            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

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,
+        self,
-            hidden_states: mx.array,
+        hidden_states: mx.array,
-            text_embeddings: mx.array,
+        text_embeddings: mx.array,
-            rotary_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)

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)

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
-

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)

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.joint_transformer_block import (
-from mflux.models.transformer.single_transformer_block import SingleTransformerBlock
+    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,
+        self,
-            t: int,
+        t: int,
-            prompt_embeds: mx.array,
+        prompt_embeds: mx.array,
-            pooled_prompt_embeds: mx.array,
+        pooled_prompt_embeds: mx.array,
-            hidden_states: mx.array,
+        hidden_states: mx.array,
-            config: RuntimeConfig,
+        config: RuntimeConfig,
-            controlnet_block_samples: list[mx.array] | None = None,
+        controlnet_block_samples: list[mx.array] | None = None,
-            controlnet_single_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

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)

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,
+        self,
-            norm1: int,
+        norm1: int,
-            conv1_in: int,
+        conv1_in: int,
-            conv1_out: int,
+        conv1_out: int,
-            norm2: int,
+        norm2: int,
-            conv2_in: int,
+        conv2_in: int,
-            conv2_out: int,
+        conv2_out: int,
-            conv_shortcut_in: int | None = None,
+        conv_shortcut_in: int | None = None,
-            conv_shortcut_out: int | None = None,
+        conv_shortcut_out: int | None = None,
-            is_conv_shortcut: bool = False
+        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,

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:

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(

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:

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

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

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 = [

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 = [

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(

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(

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:

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 = [

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 = [

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 = [

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 = [

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(

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

src/mflux/post_processing/generated_image.py

@@ -12,21 +12,20 @@ log = logging.getLogger(__name__)
 
 
 class GeneratedImage:
-
     def __init__(
-            self,
+        self,
-            image: PIL.Image.Image,
+        image: PIL.Image.Image,
-            model_config: ModelConfig,
+        model_config: ModelConfig,
-            seed: int,
+        seed: int,
-            prompt: str,
+        prompt: str,
-            steps: int,
+        steps: int,
-            guidance: float | None,
+        guidance: float | None,
-            precision: mx.Dtype,
+        precision: mx.Dtype,
-            quantization: int,
+        quantization: int,
-            generation_time: float,
+        generation_time: float,
-            lora_paths: list[str],
+        lora_paths: list[str],
-            lora_scales: list[float],
+        lora_scales: list[float],
-            controlnet_strength: float | None = None,
+        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,18 @@ 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
-            }
-
             # Create a piexif-compatible dictionary structure
-            exif_piexif_dict = {
+            exif_piexif_dict = {"Exif": {USER_COMMENT_TAG_ID: user_comment_bytes}}
-                '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 +99,15 @@ class GeneratedImage:
 
     def _get_metadata(self) -> dict:
         return {
-            'model': str(self.model_config.alias),
+            "model": str(self.model_config.alias),
-            'seed': str(self.seed),
+            "seed": str(self.seed),
-            'steps': str(self.steps),
+            "steps": str(self.steps),
-            'guidance': "None" if self.model_config == ModelConfig.FLUX1_SCHNELL else str(self.guidance),
+            "guidance": "None" if self.model_config == ModelConfig.FLUX1_SCHNELL else str(self.guidance),
-            'precision': f"{self.precision}",
+            "precision": f"{self.precision}",
-            'quantization': "None" if self.quantization is None else f"{self.quantization} bit",
+            "quantization": "None" if self.quantization is None else f"{self.quantization} bit",
-            'generation_time': f"{self.generation_time:.2f} seconds",
+            "generation_time": f"{self.generation_time:.2f} seconds",
-            'lora_paths': ', '.join(self.lora_paths) if self.lora_paths else '',
+            "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 '',
+            "lora_scales": ", ".join([f"{scale:.2f}" for scale in self.lora_scales]) if self.lora_scales else "",
-            'prompt': self.prompt,
+            "prompt": self.prompt,
-            'controlnet_strength': "None" if self.controlnet_strength is None else f"{self.controlnet_strength:.2f}",
+            "controlnet_strength": "None" if self.controlnet_strength is None else f"{self.controlnet_strength:.2f}",
         }

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,
+        decoded_latents: mx.array,
-            seed: int,
+        seed: int,
-            prompt: str,
+        prompt: str,
-            quantization: int,
+        quantization: int,
-            generation_time: float,
+        generation_time: float,
-            lora_paths: list[str],
+        lora_paths: list[str],
-            lora_scales: list[float],
+        lora_scales: list[float],
-            config: RuntimeConfig,
+        config: RuntimeConfig,
     ) -> GeneratedImage:
         normalized = ImageUtil._denormalize(decoded_latents)
         normalized_numpy = ImageUtil._to_numpy(normalized)

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 = 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("--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("--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.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()

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

src/mflux/tokenizer/tokenizer_handler.py

@@ -7,19 +7,23 @@ from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
 
 
 class TokenizerHandler:
-
+    def __init__(
-    def __init__(self, repo_id: str, max_t5_length: int = 256, local_path: str | None = None):
+        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=[
+                allow_patterns=["tokenizer/**", "tokenizer_2/**"],
-                    "tokenizer/**",
-                    "tokenizer_2/**"
-                ]
             )
         )

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)
@@ -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
@@ -38,7 +37,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 +53,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 +87,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 +119,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",
+                target,
-                f"transformer.single_transformer_blocks.{i}.norm.linear"
+                f"lora_unet_single_blocks_{i}_modulation_lin",
+                f"transformer.single_transformer_blocks.{i}.norm.linear",
             )
 
         if len(source) > 0:
@@ -133,12 +133,7 @@ class LoRAConverter:
         return target
 
     @staticmethod
-    def _convert_to_diffusers(
+    def _convert_to_diffusers(source: dict, target: dict, source_key: str, target_key: str):
-            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 +155,11 @@ class LoRAConverter:
 
     @staticmethod
     def _convert_to_diffusers_cat(
-            source: dict,
+        source: dict,
-            target: dict,
+        target: dict,
-            source_key: str,
+        source_key: str,
-            target_keys: list[str],
+        target_keys: list[str],
-            dims=None
+        dims=None,
     ):
         if source_key + ".lora_down.weight" not in source:
             return
@@ -204,7 +199,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 +220,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

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["net"] = list({})
-                            target = target['net']
+                            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 "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]:
+                    if "lora_A" in visited[parentKey] and "lora_B" in visited[parentKey]:
-                        lora_a = visited[parentKey]['lora_A']
+                        lora_a = visited[parentKey]["lora_A"]
-                        lora_b = visited[parentKey]['lora_B']
+                        lora_b = visited[parentKey]["lora_B"]
-                        transWeight = target['weight']
+                        transWeight = target["weight"]
                         weight = transWeight + lora_scale * (lora_b @ lora_a)
-                        target['weight'] = weight
+                        target["weight"] = weight

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:

src/mflux/weights/weight_handler.py

@@ -11,13 +11,12 @@ from mflux.weights.weight_util import WeightUtil
 
 
 class WeightHandler:
-
     def __init__(
-            self,
+        self,
-            repo_id: str | None = None,
+        repo_id: str | None = None,
-            local_path: str | None = None,
+        local_path: str | None = None,
-            lora_paths: list[str] | None = None,
+        lora_paths: list[str] | None = None,
-            lora_scales: list[float] | 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_in"] = {"conv2d": weights["decoder"]["conv_in"]}
-        weights['decoder']['conv_out'] = {'conv2d': weights['decoder']['conv_out']}
+        weights["decoder"]["conv_out"] = {"conv2d": weights["decoder"]["conv_out"]}
-        weights['decoder']['conv_norm_out'] = {'norm': weights['decoder']['conv_norm_out']}
+        weights["decoder"]["conv_norm_out"] = {"norm": weights["decoder"]["conv_norm_out"]}
-        weights['encoder']['conv_in'] = {'conv2d': weights['encoder']['conv_in']}
+        weights["encoder"]["conv_in"] = {"conv2d": weights["encoder"]["conv_in"]}
-        weights['encoder']['conv_out'] = {'conv2d': weights['encoder']['conv_out']}
+        weights["encoder"]["conv_out"] = {"conv2d": weights["encoder"]["conv_out"]}
-        weights['encoder']['conv_norm_out'] = {'norm': weights['encoder']['conv_norm_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",
-                ]
+                ],
             )
         )

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]