Fix str/Path related typing

src/mflux/config/runtime_config.py

@@ -1,4 +1,5 @@
 import logging
+from pathlib import Path
 
 import mlx.core as mx
 import numpy as np
@@ -48,7 +49,7 @@ class RuntimeConfig:
         return self.model_config.num_train_steps
 
     @property
-    def image_path(self) -> str:
+    def image_path(self) -> Path | None:
         return self.config.image_path
 
     @property
@@ -56,15 +57,15 @@ class RuntimeConfig:
         return self.config.image_strength
 
     @property
-    def depth_image_path(self) -> str | None:
+    def depth_image_path(self) -> Path | None:
         return self.config.depth_image_path
 
     @property
-    def redux_image_paths(self) -> str | None:
+    def redux_image_paths(self) -> list[Path] | None:
         return self.config.redux_image_paths
 
     @property
-    def masked_image_path(self) -> str | None:
+    def masked_image_path(self) -> Path | None:
         return self.config.masked_image_path
 
     @property

src/mflux/flux_tools/depth/depth_util.py

@@ -1,5 +1,6 @@
 import logging
 import os
+from pathlib import Path
 
 import mlx.core as mx
 import PIL.Image
@@ -19,8 +20,8 @@ class DepthUtil:
         vae: VAE,
         depth_pro: DepthPro,
         config: RuntimeConfig,
-        image_path: str | None = None,
-        depth_image_path: str | None = None,
+        image_path: str | Path | None = None,
+        depth_image_path: str | Path | None = None,
     ) -> (mx.array, PIL.Image.Image):
         # 1. Create the depth map or use existing one
         depth_image_path, depth_image = DepthUtil.get_or_create_depth_map(
@@ -44,9 +45,9 @@ class DepthUtil:
     @staticmethod
     def get_or_create_depth_map(
         depth_pro: DepthPro,
-        image_path: str | None = None,
-        depth_map_path: str | None = None,
-    ) -> tuple[str, PIL.Image.Image | None]:
+        image_path: str | Path | None = None,
+        depth_map_path: str | Path | None = None,
+    ) -> tuple[str | Path, PIL.Image.Image | None]:
         # 1. If a depth map path is provided, use it directly
         if depth_map_path:
             if not os.path.exists(depth_map_path):

src/mflux/flux_tools/fill/mask_util.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import mlx.core as mx
 
 from mflux.config.runtime_config import RuntimeConfig
@@ -12,8 +14,8 @@ class MaskUtil:
         vae: VAE,
         config: RuntimeConfig,
         latents: mx.array,
-        img_path: str,
-        mask_path: str | None,
+        img_path: str | Path,
+        mask_path: str | Path | None,
     ) -> mx.array:
         if not img_path or not mask_path:
             # Return empty latents if no image or mask is provided

src/mflux/flux_tools/redux/flux_redux.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import mlx.core as mx
 from mlx import nn
 from tqdm import tqdm
@@ -158,7 +160,7 @@ class Flux1Redux(nn.Module):
         clip_tokenizer: TokenizerCLIP,
         t5_text_encoder: T5Encoder,
         clip_text_encoder: CLIPEncoder,
-        image_paths: list[str],
+        image_paths: list[str] | list[Path],
         image_encoder: SiglipVisionTransformer,
         image_embedder: ReduxEncoder,
     ) -> (mx.array, mx.array):

src/mflux/flux_tools/redux/redux_util.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import mlx.core as mx
 
 from mflux import ImageUtil
@@ -8,7 +10,7 @@ from mflux.models.siglip_vision_transformer.siglip_vision_transformer import Sig
 class ReduxUtil:
     @staticmethod
     def embed_images(
-        image_paths: list[str],
+        image_paths: list[str] | list[Path],
         image_encoder: SiglipVisionTransformer,
         image_embedder: ReduxEncoder,
     ) -> list[mx.array]:  # fmt:off
@@ -24,7 +26,7 @@ class ReduxUtil:
 
     @staticmethod
     def _embed_single_image(
-        image_path: str,
+        image_path: str | Path,
         image_encoder: SiglipVisionTransformer,
         image_embedder: ReduxEncoder,
     ) -> mx.array:  # fmt:off

src/mflux/post_processing/image_util.py

@@ -27,12 +27,12 @@ class ImageUtil:
         generation_time: float,
         lora_paths: list[str],
         lora_scales: list[float],
-        controlnet_image_path: str | None = None,
-        image_path: str | None = None,
-        redux_image_paths: list[str] | None = None,
+        controlnet_image_path: str | Path | None = None,
+        image_path: str | Path | None = None,
+        redux_image_paths: list[str] | list[Path] | None = None,
         image_strength: float | None = None,
-        masked_image_path: str | None = None,
-        depth_image_path: str | None = None,
+        masked_image_path: str | Path | None = None,
+        depth_image_path: str | Path | None = None,
     ) -> GeneratedImage:
         normalized = ImageUtil._denormalize(decoded_latents)
         normalized_numpy = ImageUtil._to_numpy(normalized)
@@ -240,7 +240,7 @@ class ImageUtil:
             log.error(f"Error saving image: {e}")
 
     @staticmethod
-    def _embed_metadata(metadata: dict, path: str) -> None:
+    def _embed_metadata(metadata: dict, path: str | Path) -> None:
         try:
             # Convert metadata dictionary to a string
             metadata_str = json.dumps(metadata)