🧹 Cleanup/bump pre-commits; intro mypy config, quick type ignores/fixes (#221)

Co-authored-by: Anthony Wu <pls-file-gh-issue@users.noreply.github.com>

.pre-commit-config.yaml

@@ -1,16 +1,21 @@
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.11.13
+    rev: v0.12.1
     hooks:
       # Run the linter.
-      - id: ruff
+      - id: ruff-check
         types_or: [python, pyi]
         args: [--fix]
       # Run the formatter.
       - id: ruff-format
         types_or: [python, pyi]
   - repo: https://github.com/crate-ci/typos
-    rev: v1.31.0
+    rev: v1.33.1
     hooks:
       - id: typos
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.16.1
+    hooks:
+      - id: mypy
+        files: ^src/

pyproject.toml

@@ -145,3 +145,23 @@ section-order = [
     "first-party",
     "local-folder",
 ]
+
+[tool.mypy]
+disable_error_code = [
+    # TODO: for each error code - clean them up in a single PR then remove entry
+    "annotation-unchecked",
+    "arg-type",
+    "assignment",
+    "attr-defined",
+    "import-untyped",
+    "index",
+    "union-attr",
+    "return",
+    "return-value",
+    "var-annotated"
+]
+error_summary = true
+ignore_missing_imports = true
+implicit_optional = true
+# we should feel free to use new hinting features even if we back support to earlier versions
+python_version = 3.12

src/mflux/callbacks/instances/battery_saver.py

@@ -56,7 +56,7 @@ class BatterySaver(BeforeLoopCallback):
     def __init__(self, battery_percentage_stop_limit=10):
         self.limit = battery_percentage_stop_limit
 
-    def call_before_loop(self, **kwargs) -> None:
+    def call_before_loop(self, **kwargs) -> None:  # type: ignore
         current_pct: int | None = get_battery_percentage()
         if current_pct is not None and current_pct <= self.limit:
             raise StopImageGenerationException(f"Battery below {self.limit}% threshold: {current_pct}%")

src/mflux/config/runtime_config.py

@@ -33,7 +33,7 @@ class RuntimeConfig:
         self.config.width = value
 
     @property
-    def guidance(self) -> float | None:
+    def guidance(self) -> float:
         return self.config.guidance
 
     @property
@@ -82,10 +82,10 @@ class RuntimeConfig:
 
         if is_img2img:
             # 1. Clamp strength to [0, 1]
-            strength = max(0.0, min(1.0, self.config.image_strength))
+            strength = max(0.0, min(1.0, self.config.image_strength))  # type: ignore
 
             # 2. Return start time in [1, floor(num_steps * strength)]
-            return max(1, int(self.num_inference_steps * strength))
+            return max(1, int(self.num_inference_steps * strength))  # type: ignore
         else:
             return 0
 

src/mflux/dreambooth/dataset/dataset.py

@@ -27,10 +27,9 @@ class Dataset:
         examples = Dataset._create_examples(flux, raw_data, width=width, height=height)
 
         # Expend the original dataset to get more training data with variations
-        augmented_examples = []
+        augmented_examples = [
-        for example in examples:
+            variation for example in examples for variation in DreamBoothPreProcessing.augment(example)
-            [augmented_examples.append(variation) for variation in DreamBoothPreProcessing.augment(example)]
+        ]
-
         # Dataset is now prepared
         return Dataset(augmented_examples)
 

src/mflux/dreambooth/dataset/iterator.py

@@ -156,7 +156,7 @@ class Iterator:
         return Batch(examples=examples, rng=self.rng)
 
     def total_number_of_steps(self) -> int:
-        return self.total_examples * self.num_epochs
+        return self.total_examples * self.num_epochs  # type: ignore
 
     def save(self, path: Path) -> None:
         self.to_json(str(path))

src/mflux/dreambooth/state/training_state.py

@@ -74,7 +74,7 @@ class TrainingState:
                     if file_path.exists():
                         zipf.write(file_path, file_path.name)
 
-    def _create_checkpoint_data(self, training_spec: TrainingSpec, start_date_time: datetime) -> dict:
+    def _create_checkpoint_data(self, training_spec: TrainingSpec, start_date_time: datetime.datetime) -> dict:
         now = datetime.datetime.now()
         return {
             "metadata": {
@@ -122,7 +122,7 @@ class TrainingState:
         return None if path is None else str(Path(path).resolve())
 
     @staticmethod
-    def _format_duration(start: datetime, end: datetime) -> str:
+    def _format_duration(start: datetime.datetime, end: datetime.datetime) -> str:
         # Calculate the duration
         duration = end - start
         total_seconds = int(duration.total_seconds())

src/mflux/dreambooth/state/zip_util.py

@@ -1,12 +1,13 @@
 import os
 import tempfile
 from pathlib import Path
+from typing import Callable
 from zipfile import ZipFile
 
 
 class ZipUtil:
     @staticmethod
-    def unzip(zip_path: str | Path | None, filename: str, loader: callable):
+    def unzip(zip_path: str | Path | None, filename: str, loader: Callable):
         if not zip_path:  # Would be nicer to do this in typing, but that's more effort on the callers' side
             raise ValueError("zip_path cannot be None")
         zip_path = Path(zip_path)

src/mflux/flux_tools/redux/flux_redux.py

@@ -58,14 +58,14 @@ class Flux1Redux(nn.Module):
         config: Config,
     ) -> GeneratedImage:
         # 0. Create a new runtime config based on the model type and input parameters
-        config = RuntimeConfig(config, self.model_config)
+        runtime_config = RuntimeConfig(config, self.model_config)
-        time_steps = tqdm(range(config.init_time_step, config.num_inference_steps))
+        time_steps = tqdm(range(runtime_config.init_time_step, runtime_config.num_inference_steps))
 
         # 1. Create the initial latents
         latents = LatentCreator.create(
             seed=seed,
-            height=config.height,
+            height=runtime_config.height,
-            width=config.width,
+            width=runtime_config.width,
         )
 
         # 2. Get prompt embeddings by fusing the prompt and image embeddings
@@ -76,10 +76,10 @@ class Flux1Redux(nn.Module):
             clip_tokenizer=self.clip_tokenizer,
             t5_text_encoder=self.t5_text_encoder,
             clip_text_encoder=self.clip_text_encoder,
-            image_paths=config.redux_image_paths,
+            image_paths=runtime_config.redux_image_paths,
             image_encoder=self.image_encoder,
             image_embedder=self.image_embedder,
-            image_strengths=config.redux_image_strengths,
+            image_strengths=runtime_config.redux_image_strengths,
         )  # fmt: off
 
         # (Optional) Call subscribers for beginning of loop
@@ -87,7 +87,7 @@ class Flux1Redux(nn.Module):
             seed=seed,
             prompt=prompt,
             latents=latents,
-            config=config,
+            config=runtime_config,
         )  # fmt: off
 
         for t in time_steps:
@@ -95,14 +95,14 @@ class Flux1Redux(nn.Module):
                 # 3.t Predict the noise
                 noise = self.transformer(
                     t=t,
-                    config=config,
+                    config=runtime_config,
                     hidden_states=latents,
                     prompt_embeds=prompt_embeds,
                     pooled_prompt_embeds=pooled_prompt_embeds,
                 )
 
                 # 4.t Take one denoise step
-                dt = config.sigmas[t + 1] - config.sigmas[t]
+                dt = runtime_config.sigmas[t + 1] - runtime_config.sigmas[t]
                 latents += noise * dt
 
                 # (Optional) Call subscribers in-loop
@@ -111,7 +111,7 @@ class Flux1Redux(nn.Module):
                     seed=seed,
                     prompt=prompt,
                     latents=latents,
-                    config=config,
+                    config=runtime_config,
                     time_steps=time_steps,
                 )  # fmt: off
 
@@ -124,7 +124,7 @@ class Flux1Redux(nn.Module):
                     seed=seed,
                     prompt=prompt,
                     latents=latents,
-                    config=config,
+                    config=runtime_config,
                     time_steps=time_steps,
                 )
                 raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
@@ -134,23 +134,23 @@ class Flux1Redux(nn.Module):
             seed=seed,
             prompt=prompt,
             latents=latents,
-            config=config,
+            config=runtime_config,
         )  # fmt: off
 
         # 7. Decode the latent array and return the image
-        latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
+        latents = ArrayUtil.unpack_latents(latents=latents, height=runtime_config.height, width=runtime_config.width)
         decoded = self.vae.decode(latents)
         return ImageUtil.to_image(
             decoded_latents=decoded,
-            config=config,
+            config=runtime_config,
             seed=seed,
             prompt=prompt,
             quantization=self.bits,
             lora_paths=self.lora_paths,
             lora_scales=self.lora_scales,
-            redux_image_paths=config.redux_image_paths,
+            redux_image_paths=runtime_config.redux_image_paths,
-            redux_image_strengths=config.redux_image_strengths,
+            redux_image_strengths=runtime_config.redux_image_strengths,
-            image_strength=config.image_strength,
+            image_strength=runtime_config.image_strength,
             generation_time=time_steps.format_dict["elapsed"],
         )
 

src/mflux/models/vae/common/resnet_block_2d.py

@@ -66,7 +66,7 @@ class ResnetBlock2D(nn.Module):
         hidden_states = self.norm2(hidden_states.astype(mx.float32)).astype(Config.precision)
         hidden_states = nn.silu(hidden_states)
         hidden_states = self.conv2(hidden_states)
-        if self.is_conv_shortcut:
+        if self.conv_shortcut is not None:
             input_array = self.conv_shortcut(input_array)
         output_tensor = input_array + hidden_states
         return mx.transpose(output_tensor, (0, 3, 1, 2))

src/mflux/ui/cli/parsers.py

@@ -190,7 +190,7 @@ class CommandLineParser(argparse.ArgumentParser):
         self.add_argument("--train-config", type=str, required=False, help="Local path of the training configuration file")
         self.add_argument("--train-checkpoint", type=str, required=False, help="Local path of the checkpoint file which specifies how to continue the training process")
 
-    def parse_args(self, **kwargs) -> argparse.Namespace:
+    def parse_args(self) -> argparse.Namespace:  # type: ignore
         namespace = super().parse_args()
 
         # Check if either training arguments are provided

src/mflux/upscale.py

@@ -68,14 +68,16 @@ def _calculate_output_dimensions(args) -> tuple[int, int]:
     output_width, output_height = orig_image.size
 
     if isinstance(args.height, ScaleFactor):
-        output_height: int = args.height.get_scaled_value(orig_image.height)
+        output_height: int = args.height.get_scaled_value(orig_image.height)  # type: ignore
+
     else:
-        output_height = args.height
+        output_height = args.height  # type: ignore
 
     if isinstance(args.width, ScaleFactor):
-        output_width: int = args.width.get_scaled_value(orig_image.width)
+        output_width: int = args.width.get_scaled_value(orig_image.width)  # type: ignore
+
     else:
-        output_width = args.width
+        output_width = args.width  # type: ignore
 
     # Check if dimensions exceed safe limits
     total_pixels = output_height * output_width