Apply no-implicit-optional fixes

src/mflux/models/depth_pro/feature_fusion_block_2d.py

@@ -14,7 +14,7 @@ class FeatureFusionBlock2d(nn.Module):
         self.deconv = nn.ConvTranspose2d(in_channels=num_features, out_channels=num_features, kernel_size=2, stride=2, padding=0, bias=False)  # fmt: off
         self.out_conv = nn.Conv2d(in_channels=num_features, out_channels=num_features, kernel_size=1, stride=1, padding=0, bias=True)  # fmt: off
 
-    def __call__(self, x0: mx.array, x1: mx.array = None) -> mx.array:
+    def __call__(self, x0: mx.array, x1: mx.array | None = None) -> mx.array:
         x = x0
         if x1 is not None:
             res = self.resnet1(x1)

src/mflux/models/depth_pro/upsample_block.py

@@ -21,7 +21,7 @@ class UpSampleBlock(nn.Module):
         )  # fmt: off
 
     @staticmethod
-    def _create_layers(dim_in: int, dim_out: int, upsample_layers: int, dim_int: int = None) -> list[nn.Module]:
+    def _create_layers(dim_in: int, dim_out: int, upsample_layers: int, dim_int: int | None = None) -> list[nn.Module]:
         if dim_int is None:
             dim_int = dim_out
 

src/mflux/post_processing/generated_image.py

@@ -1,5 +1,4 @@
 import importlib
-import typing as t
 from pathlib import Path
 
 import mlx.core as mx
@@ -78,7 +77,7 @@ class GeneratedImage:
 
     def save(
         self,
-        path: t.Union[str, Path],
+        path: str | Path,
         export_json_metadata: bool = False,
         overwrite: bool = False,
     ) -> None:

src/mflux/post_processing/image_util.py

@@ -1,6 +1,5 @@
 import json
 import logging
-import typing as t
 from pathlib import Path
 
 import mlx.core as mx
@@ -59,7 +58,7 @@ class ImageUtil:
         )
 
     @staticmethod
-    def to_composite_image(generated_images: t.List[GeneratedImage]) -> PIL.Image.Image:
+    def to_composite_image(generated_images: list[GeneratedImage]) -> PIL.Image.Image:
         # stitch horizontally
         total_width = sum(gen_img.image.width for gen_img in generated_images)
         max_height = max(gen_img.image.height for gen_img in generated_images)
@@ -119,7 +118,7 @@ class ImageUtil:
     @staticmethod
     def expand_image(
         image: PIL.Image.Image,
-        box_values: AbsoluteBoxValues = None,
+        box_values: AbsoluteBoxValues | None = None,
         top: int | str = 0,
         right: int | str = 0,
         bottom: int | str = 0,
@@ -161,7 +160,7 @@ class ImageUtil:
         orig_height: int,
         border_color: tuple,
         content_color: tuple,
-        box_values: AbsoluteBoxValues = None,
+        box_values: AbsoluteBoxValues | None = None,
         top: int | str = 0,
         right: int | str = 0,
         bottom: int | str = 0,
@@ -204,7 +203,7 @@ class ImageUtil:
     @staticmethod
     def save_image(
         image: PIL.Image.Image,
-        path: t.Union[str, Path],
+        path: str | Path,
         metadata: dict | None = None,
         export_json_metadata: bool = False,
         overwrite: bool = False,

src/mflux/weights/weight_handler_lora_huggingface.py

@@ -7,9 +7,9 @@ from huggingface_hub import snapshot_download
 class WeightHandlerLoRAHuggingFace:
     @staticmethod
     def download_loras(
-        lora_names: list[str] = None,
-        repo_id: str = None,
-        cache_dir: str = None,
+        lora_names: list[str] | None = None,
+        repo_id: str | None = None,
+        cache_dir: str | None = None,
     ) -> list[str]:
         if repo_id is None:
             return []
@@ -30,7 +30,7 @@ class WeightHandlerLoRAHuggingFace:
     def _download_lora(
         repo_id: str,
         lora_name: str,
-        cache_dir: str = None,
+        cache_dir: str | None = None,
     ) -> str:
         # Create cache directory if it doesn't exist
         if cache_dir is None:

tests/image_generation/helpers/image_generation_depth_test_helper.py

@@ -17,8 +17,8 @@ class ImageGeneratorDepthTestHelper:
         prompt: str,
         steps: int,
         seed: int,
-        height: int = None,
-        width: int = None,
+        height: int | None = None,
+        width: int | None = None,
         image_path: str | None = None,
         depth_image_path: str | None = None,
     ):

tests/image_generation/helpers/image_generation_in_context_test_helper.py

@@ -18,8 +18,8 @@ class ImageGeneratorInContextTestHelper:
         prompt: str,
         steps: int,
         seed: int,
-        height: int = None,
-        width: int = None,
+        height: int | None = None,
+        width: int | None = None,
         image_path: str | None = None,
         lora_style: str | None = None,
         lora_paths: list[str] | None = None,

tests/image_generation/helpers/image_generation_test_helper.py

@@ -16,8 +16,8 @@ class ImageGeneratorTestHelper:
         prompt: str,
         steps: int,
         seed: int,
-        height: int = None,
-        width: int = None,
+        height: int | None = None,
+        width: int | None = None,
         image_path: str | None = None,
         image_strength: float | None = None,
         lora_paths: list[str] | None = None,