From 9b6cd01171218728462cddb94e1d625334309f27 Mon Sep 17 00:00:00 2001
From: Anthony Wu <462072+anthonywu@users.noreply.github.com>
Date: Thu, 10 Oct 2024 11:56:14 -0700
Subject: [PATCH] introduce ruff isort

---
 pyproject.toml                                | 20 +++++++++++++++++--
 src/mflux/__init__.py                         |  5 ++---
 src/mflux/controlnet/controlnet_util.py       |  1 -
 src/mflux/controlnet/flux_controlnet.py       |  8 +++++---
 src/mflux/flux/flux.py                        |  4 ++--
 src/mflux/generate.py                         |  2 +-
 src/mflux/generate_controlnet.py              |  2 +-
 .../clip_encoder/clip_embeddings.py           |  2 +-
 .../clip_encoder/clip_encoder_layer.py        |  2 +-
 .../t5_encoder/t5_dense_relu_dense.py         |  2 +-
 .../transformer/ada_layer_norm_continuous.py  |  2 +-
 .../models/transformer/ada_layer_norm_zero.py |  2 +-
 .../transformer/ada_layer_norm_zero_single.py |  2 +-
 src/mflux/models/transformer/feed_forward.py  |  2 +-
 .../models/transformer/guidance_embedder.py   |  2 +-
 src/mflux/models/transformer/text_embedder.py |  2 +-
 .../models/transformer/time_text_embed.py     |  5 +++--
 .../models/transformer/timestep_embedder.py   |  2 +-
 src/mflux/models/vae/decoder/decoder.py       |  2 +-
 src/mflux/models/vae/decoder/up_sampler.py    |  1 -
 src/mflux/models/vae/encoder/down_sampler.py  |  1 -
 src/mflux/models/vae/encoder/encoder.py       |  1 -
 src/mflux/post_processing/generated_image.py  |  3 ++-
 src/mflux/post_processing/image_util.py       | 10 ++++++----
 src/mflux/weights/lora_converter.py           |  1 +
 ...image_generation_controlnet_test_helper.py |  2 +-
 tests/helpers/image_generation_test_helper.py |  2 +-
 27 files changed, 54 insertions(+), 36 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 90a514b..be18f64 100644
--- a/pyproject.toml
+++ b/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"
\ No newline at end of file
+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",
+]
diff --git a/src/mflux/__init__.py b/src/mflux/__init__.py
index e53c33f..553f634 100644
--- a/src/mflux/__init__.py
+++ b/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__ = [
diff --git a/src/mflux/controlnet/controlnet_util.py b/src/mflux/controlnet/controlnet_util.py
index 7f083cc..64cb27c 100644
--- a/src/mflux/controlnet/controlnet_util.py
+++ b/src/mflux/controlnet/controlnet_util.py
@@ -5,7 +5,6 @@ import cv2
 import numpy as np
 import PIL.Image
 
-
 log = logging.getLogger(__name__)
 
 
diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py
index 73f5523..5f4687e 100644
--- a/src/mflux/controlnet/flux_controlnet.py
+++ b/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)}")
 
diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py
index 634bf32..be0b367 100644
--- a/src/mflux/flux/flux.py
+++ b/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)}")
 
diff --git a/src/mflux/generate.py b/src/mflux/generate.py
index d6726f0..2184ef7 100644
--- a/src/mflux/generate.py
+++ b/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():
diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py
index 470236d..28c1a09 100644
--- a/src/mflux/generate_controlnet.py
+++ b/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():
diff --git a/src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py b/src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py
index 988c5d3..d8b1b17 100644
--- a/src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py
+++ b/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
 
diff --git a/src/mflux/models/text_encoder/clip_encoder/clip_encoder_layer.py b/src/mflux/models/text_encoder/clip_encoder/clip_encoder_layer.py
index 972b39c..8df7028 100644
--- a/src/mflux/models/text_encoder/clip_encoder/clip_encoder_layer.py
+++ b/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
diff --git a/src/mflux/models/text_encoder/t5_encoder/t5_dense_relu_dense.py b/src/mflux/models/text_encoder/t5_encoder/t5_dense_relu_dense.py
index 0318fe7..a6af0ea 100644
--- a/src/mflux/models/text_encoder/t5_encoder/t5_dense_relu_dense.py
+++ b/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):
diff --git a/src/mflux/models/transformer/ada_layer_norm_continuous.py b/src/mflux/models/transformer/ada_layer_norm_continuous.py
index 9538fd0..94a07cc 100644
--- a/src/mflux/models/transformer/ada_layer_norm_continuous.py
+++ b/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
 
diff --git a/src/mflux/models/transformer/ada_layer_norm_zero.py b/src/mflux/models/transformer/ada_layer_norm_zero.py
index 17d0956..deeba1c 100644
--- a/src/mflux/models/transformer/ada_layer_norm_zero.py
+++ b/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):
diff --git a/src/mflux/models/transformer/ada_layer_norm_zero_single.py b/src/mflux/models/transformer/ada_layer_norm_zero_single.py
index 9299d67..655a324 100644
--- a/src/mflux/models/transformer/ada_layer_norm_zero_single.py
+++ b/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):
diff --git a/src/mflux/models/transformer/feed_forward.py b/src/mflux/models/transformer/feed_forward.py
index 3a13d4b..6124fe7 100644
--- a/src/mflux/models/transformer/feed_forward.py
+++ b/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):
diff --git a/src/mflux/models/transformer/guidance_embedder.py b/src/mflux/models/transformer/guidance_embedder.py
index 91ba297..ed9d1c4 100644
--- a/src/mflux/models/transformer/guidance_embedder.py
+++ b/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):
diff --git a/src/mflux/models/transformer/text_embedder.py b/src/mflux/models/transformer/text_embedder.py
index 6fe09e2..6034b74 100644
--- a/src/mflux/models/transformer/text_embedder.py
+++ b/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):
diff --git a/src/mflux/models/transformer/time_text_embed.py b/src/mflux/models/transformer/time_text_embed.py
index a3db91a..5dbd84e 100644
--- a/src/mflux/models/transformer/time_text_embed.py
+++ b/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):
diff --git a/src/mflux/models/transformer/timestep_embedder.py b/src/mflux/models/transformer/timestep_embedder.py
index dfe50fd..b72aea2 100644
--- a/src/mflux/models/transformer/timestep_embedder.py
+++ b/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):
diff --git a/src/mflux/models/vae/decoder/decoder.py b/src/mflux/models/vae/decoder/decoder.py
index e61481a..86a50cd 100644
--- a/src/mflux/models/vae/decoder/decoder.py
+++ b/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
diff --git a/src/mflux/models/vae/decoder/up_sampler.py b/src/mflux/models/vae/decoder/up_sampler.py
index 46f147f..2312f45 100644
--- a/src/mflux/models/vae/decoder/up_sampler.py
+++ b/src/mflux/models/vae/decoder/up_sampler.py
@@ -1,5 +1,4 @@
 import mlx.core as mx
-
 from mlx import nn
 
 
diff --git a/src/mflux/models/vae/encoder/down_sampler.py b/src/mflux/models/vae/encoder/down_sampler.py
index b06cd6d..856ecc5 100644
--- a/src/mflux/models/vae/encoder/down_sampler.py
+++ b/src/mflux/models/vae/encoder/down_sampler.py
@@ -1,5 +1,4 @@
 import mlx.core as mx
-
 from mlx import nn
 
 
diff --git a/src/mflux/models/vae/encoder/encoder.py b/src/mflux/models/vae/encoder/encoder.py
index 919dac7..e73d490 100644
--- a/src/mflux/models/vae/encoder/encoder.py
+++ b/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
diff --git a/src/mflux/post_processing/generated_image.py b/src/mflux/post_processing/generated_image.py
index f2320d8..5d85c4f 100644
--- a/src/mflux/post_processing/generated_image.py
+++ b/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
 
diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py
index ba7c07f..61a1d40 100644
--- a/src/mflux/post_processing/image_util.py
+++ b/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
diff --git a/src/mflux/weights/lora_converter.py b/src/mflux/weights/lora_converter.py
index 9570a04..6552036 100644
--- a/src/mflux/weights/lora_converter.py
+++ b/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
diff --git a/tests/helpers/image_generation_controlnet_test_helper.py b/tests/helpers/image_generation_controlnet_test_helper.py
index 5ed4f6a..c5f3b45 100644
--- a/tests/helpers/image_generation_controlnet_test_helper.py
+++ b/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
 
 
diff --git a/tests/helpers/image_generation_test_helper.py b/tests/helpers/image_generation_test_helper.py
index 72eb3f7..8f35a11 100644
--- a/tests/helpers/image_generation_test_helper.py
+++ b/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: