Merge pull request #72 from anthonywu/intro-isort-linter

introduce ruff isort

pyproject.toml

@@ -62,7 +62,7 @@ respect-gitignore = true
 # 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 = ["BLE", "E4", "E7", "E9", "F", "ICN", "LOG", "PERF", "W"]
+select = ["BLE", "E4", "E7", "E9", "F", "I", "ICN", "LOG", "PERF", "W"]
 ignore = []
 
 # Allow fix for all enabled rules (when `--fix`) is provided.
@@ -102,4 +102,20 @@ docstring-code-line-length = "dynamic"
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 python_files = "test_*.py"
-addopts = "-v"
+addopts = "-v"
+
+# https://docs.astral.sh/ruff/settings/#lintisort
+[tool.ruff.lint.isort]
+case-sensitive = false
+combine-as-imports = true
+force-wrap-aliases = true
+from-first = false
+known-first-party = ["src"]
+relative-imports-order = "furthest-to-closest"
+section-order = [
+    "future",
+    "standard-library",
+    "third-party",
+    "first-party",
+    "local-folder",
+]

src/mflux/__init__.py

@@ -1,9 +1,8 @@
-from mflux.config.config import Config
-from mflux.config.config import ConfigControlnet
+from mflux.config.config import Config, ConfigControlnet
 from mflux.config.model_config import ModelConfig
 from mflux.controlnet.flux_controlnet import Flux1Controlnet
-from mflux.flux.flux import Flux1
 from mflux.error.exceptions import StopImageGenerationException
+from mflux.flux.flux import Flux1
 from mflux.post_processing.image_util import ImageUtil
 
 __all__ = [

src/mflux/controlnet/controlnet_util.py

@@ -5,7 +5,6 @@ import cv2
 import numpy as np
 import PIL.Image
 
-
 log = logging.getLogger(__name__)
 
 

src/mflux/controlnet/flux_controlnet.py

@@ -1,9 +1,11 @@
 import logging
+from pathlib import Path
+from typing import TYPE_CHECKING
+
 import mlx.core as mx
 from mlx import nn
-from pathlib import Path
 from tqdm import tqdm
-from typing import TYPE_CHECKING
+
 from mflux.config.config import ConfigControlnet
 from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
@@ -183,7 +185,7 @@ class Flux1Controlnet:
                 # Evaluate to enable progress tracking
                 mx.eval(latents)
 
-            except KeyboardInterrupt:
+            except KeyboardInterrupt:  # noqa: PERF203
                 stepwise_handler.handle_interruption()
                 raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
 

src/mflux/flux/flux.py

@@ -8,7 +8,6 @@ from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.error.exceptions import StopImageGenerationException
-from mflux.post_processing.stepwise_handler import StepwiseHandler
 from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
 from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
 from mflux.models.transformer.transformer import Transformer
@@ -16,6 +15,7 @@ from mflux.models.vae.vae import VAE
 from mflux.post_processing.array_util import ArrayUtil
 from mflux.post_processing.generated_image import GeneratedImage
 from mflux.post_processing.image_util import ImageUtil
+from mflux.post_processing.stepwise_handler import StepwiseHandler
 from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
 from mflux.tokenizer.t5_tokenizer import TokenizerT5
 from mflux.tokenizer.tokenizer_handler import TokenizerHandler
@@ -126,7 +126,7 @@ class Flux1:
                 # Evaluate to enable progress tracking
                 mx.eval(latents)
 
-            except KeyboardInterrupt:
+            except KeyboardInterrupt:  # noqa: PERF203
                 stepwise_handler.handle_interruption()
                 raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
 

src/mflux/generate.py

@@ -2,7 +2,7 @@ import argparse
 import time
 from pathlib import Path
 
-from mflux import Flux1, Config, ModelConfig, StopImageGenerationException
+from mflux import Config, Flux1, ModelConfig, StopImageGenerationException
 
 
 def main():

src/mflux/generate_controlnet.py

@@ -2,7 +2,7 @@ import argparse
 import time
 from pathlib import Path
 
-from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig, StopImageGenerationException
+from mflux import ConfigControlnet, Flux1Controlnet, ModelConfig, StopImageGenerationException
 
 
 def main():

src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py

@@ -1,5 +1,5 @@
-from mlx import nn
 import mlx.core as mx
+from mlx import nn
 
 from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
 

src/mflux/models/text_encoder/clip_encoder/clip_encoder_layer.py

@@ -1,5 +1,5 @@
-from mlx import nn
 import mlx.core as mx
+from mlx import nn
 
 from mflux.models.text_encoder.clip_encoder.clip_mlp import CLIPMLP
 from mflux.models.text_encoder.clip_encoder.clip_sdpa_attention import CLIPSdpaAttention

src/mflux/models/text_encoder/t5_encoder/t5_dense_relu_dense.py

@@ -1,7 +1,7 @@
 import math
 
-from mlx import nn
 import mlx.core as mx
+from mlx import nn
 
 
 class T5DenseReluDense(nn.Module):

src/mflux/models/transformer/ada_layer_norm_continuous.py

@@ -1,5 +1,5 @@
-from mlx import nn
 import mlx.core as mx
+from mlx import nn
 
 from mflux.config.config import Config
 

src/mflux/models/transformer/ada_layer_norm_zero.py

@@ -1,5 +1,5 @@
-from mlx import nn
 import mlx.core as mx
+from mlx import nn
 
 
 class AdaLayerNormZero(nn.Module):

src/mflux/models/transformer/ada_layer_norm_zero_single.py

@@ -1,5 +1,5 @@
-from mlx import nn
 import mlx.core as mx
+from mlx import nn
 
 
 class AdaLayerNormZeroSingle(nn.Module):

src/mflux/models/transformer/feed_forward.py

@@ -1,5 +1,5 @@
-from mlx import nn
 import mlx.core as mx
+from mlx import nn
 
 
 class FeedForward(nn.Module):

src/mflux/models/transformer/guidance_embedder.py

@@ -1,5 +1,5 @@
-from mlx import nn
 import mlx.core as mx
+from mlx import nn
 
 
 class GuidanceEmbedder(nn.Module):

src/mflux/models/transformer/text_embedder.py

@@ -1,5 +1,5 @@
-from mlx import nn
 import mlx.core as mx
+from mlx import nn
 
 
 class TextEmbedder(nn.Module):

src/mflux/models/transformer/time_text_embed.py

@@ -1,12 +1,13 @@
 import math
-from mlx import nn
+
 import mlx.core as mx
+from mlx import nn
 
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
+from mflux.models.transformer.guidance_embedder import GuidanceEmbedder
 from mflux.models.transformer.text_embedder import TextEmbedder
 from mflux.models.transformer.timestep_embedder import TimestepEmbedder
-from mflux.models.transformer.guidance_embedder import GuidanceEmbedder
 
 
 class TimeTextEmbed(nn.Module):

src/mflux/models/transformer/timestep_embedder.py

@@ -1,5 +1,5 @@
-from mlx import nn
 import mlx.core as mx
+from mlx import nn
 
 
 class TimestepEmbedder(nn.Module):

src/mflux/models/vae/decoder/decoder.py

@@ -1,10 +1,10 @@
 import mlx.core as mx
 from mlx import nn
 
+from mflux.models.vae.common.unet_mid_block import UnetMidBlock
 from mflux.models.vae.decoder.conv_in import ConvIn
 from mflux.models.vae.decoder.conv_norm_out import ConvNormOut
 from mflux.models.vae.decoder.conv_out import ConvOut
-from mflux.models.vae.common.unet_mid_block import UnetMidBlock
 from mflux.models.vae.decoder.up_block_1_or_2 import UpBlock1Or2
 from mflux.models.vae.decoder.up_block_3 import UpBlock3
 from mflux.models.vae.decoder.up_block_4 import UpBlock4

src/mflux/models/vae/decoder/up_sampler.py

@@ -1,5 +1,4 @@
 import mlx.core as mx
-
 from mlx import nn
 
 

src/mflux/models/vae/encoder/down_sampler.py

@@ -1,5 +1,4 @@
 import mlx.core as mx
-
 from mlx import nn
 
 

src/mflux/models/vae/encoder/encoder.py

@@ -2,7 +2,6 @@ import mlx.core as mx
 from mlx import nn
 
 from mflux.models.vae.common.unet_mid_block import UnetMidBlock
-
 from mflux.models.vae.encoder.conv_in import ConvIn
 from mflux.models.vae.encoder.conv_norm_out import ConvNormOut
 from mflux.models.vae.encoder.conv_out import ConvOut

src/mflux/post_processing/generated_image.py

@@ -1,8 +1,9 @@
 import importlib
 import pathlib
 import typing as t
-import PIL.Image
+
 import mlx.core as mx
+import PIL.Image
 
 from mflux.config.model_config import ModelConfig
 

src/mflux/post_processing/image_util.py

@@ -1,13 +1,15 @@
-import typing as t
 import json
 import logging
 import pathlib
-import PIL
-import PIL.Image
+import typing as t
+
 import mlx.core as mx
 import numpy as np
-from PIL import Image
 import piexif
+import PIL
+import PIL.Image
+from PIL import Image
+
 from mflux.config.config import ConfigControlnet
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.post_processing.generated_image import GeneratedImage

src/mflux/weights/lora_converter.py

@@ -1,4 +1,5 @@
 import logging
+
 import mlx.core as mx
 import torch
 from mlx.utils import tree_unflatten

tests/helpers/image_generation_controlnet_test_helper.py

@@ -3,7 +3,7 @@ import os
 import numpy as np
 from PIL import Image
 
-from mflux import ModelConfig, Flux1Controlnet, ConfigControlnet
+from mflux import ConfigControlnet, Flux1Controlnet, ModelConfig
 from tests.helpers.image_generation_test_helper import ImageGeneratorTestHelper
 
 

tests/helpers/image_generation_test_helper.py

@@ -4,7 +4,7 @@ from pathlib import Path
 import numpy as np
 from PIL import Image
 
-from mflux import Flux1, Config, ModelConfig
+from mflux import Config, Flux1, ModelConfig
 
 
 class ImageGeneratorTestHelper: