diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0d5f5fd..7964059 100644 --- a/.pre-commit-config.yaml +++ b/.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/ diff --git a/pyproject.toml b/pyproject.toml index d8fa229..3844aff 100644 --- a/pyproject.toml +++ b/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 diff --git a/src/mflux/callbacks/instances/battery_saver.py b/src/mflux/callbacks/instances/battery_saver.py index 4e51258..267ccfc 100644 --- a/src/mflux/callbacks/instances/battery_saver.py +++ b/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}%") diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py index 6ee11bd..809373c 100644 --- a/src/mflux/config/runtime_config.py +++ b/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 diff --git a/src/mflux/dreambooth/dataset/dataset.py b/src/mflux/dreambooth/dataset/dataset.py index 5ef3f7e..0c386a1 100644 --- a/src/mflux/dreambooth/dataset/dataset.py +++ b/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 = [] - for example in examples: - [augmented_examples.append(variation) for variation in DreamBoothPreProcessing.augment(example)] - + augmented_examples = [ + variation for example in examples for variation in DreamBoothPreProcessing.augment(example) + ] # Dataset is now prepared return Dataset(augmented_examples) diff --git a/src/mflux/dreambooth/dataset/iterator.py b/src/mflux/dreambooth/dataset/iterator.py index 3eff4a8..c058d5b 100644 --- a/src/mflux/dreambooth/dataset/iterator.py +++ b/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)) diff --git a/src/mflux/dreambooth/state/training_state.py b/src/mflux/dreambooth/state/training_state.py index 33d61ef..2d3a0f5 100644 --- a/src/mflux/dreambooth/state/training_state.py +++ b/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()) diff --git a/src/mflux/dreambooth/state/zip_util.py b/src/mflux/dreambooth/state/zip_util.py index 4bca683..569fcac 100644 --- a/src/mflux/dreambooth/state/zip_util.py +++ b/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) diff --git a/src/mflux/flux_tools/redux/flux_redux.py b/src/mflux/flux_tools/redux/flux_redux.py index 3be8623..ded162b 100644 --- a/src/mflux/flux_tools/redux/flux_redux.py +++ b/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) - time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) + runtime_config = RuntimeConfig(config, self.model_config) + 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, - width=config.width, + height=runtime_config.height, + 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_strengths=config.redux_image_strengths, - image_strength=config.image_strength, + redux_image_paths=runtime_config.redux_image_paths, + redux_image_strengths=runtime_config.redux_image_strengths, + image_strength=runtime_config.image_strength, generation_time=time_steps.format_dict["elapsed"], ) diff --git a/src/mflux/models/vae/common/resnet_block_2d.py b/src/mflux/models/vae/common/resnet_block_2d.py index 356a627..fe1e65a 100644 --- a/src/mflux/models/vae/common/resnet_block_2d.py +++ b/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)) diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index 032d803..4131b4a 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/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 diff --git a/src/mflux/upscale.py b/src/mflux/upscale.py index 8c98033..226177c 100644 --- a/src/mflux/upscale.py +++ b/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