Merge pull request #173 from akx/path-types

Fix up some path-related types

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/dreambooth/state/zip_util.py

@@ -6,7 +6,9 @@ from zipfile import ZipFile
 
 class ZipUtil:
     @staticmethod
-    def unzip(zip_path: str | Path, 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)
         if not zip_path.exists():
             raise FileNotFoundError(f"ZIP file not found at: {zip_path}")

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,
+        image_path: str | Path | None = None,
-        depth_image_path: str | 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,
+        image_path: str | Path | None = None,
-        depth_map_path: str | None = None,
+        depth_map_path: str | Path | None = None,
-    ) -> tuple[str, PIL.Image.Image | 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,
+        img_path: str | Path,
-        mask_path: str | None,
+        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/latent_creator/latent_creator.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import mlx.core as mx
 
 from mflux.models.vae.vae import VAE
@@ -11,7 +13,7 @@ class Img2Img:
         vae: VAE,
         sigmas: mx.array,
         init_time_step: int,
-        image_path: int,
+        image_path: str | Path | None,
     ):
         self.vae = vae
         self.sigmas = sigmas
@@ -39,9 +41,7 @@ class LatentCreator:
         img2img: Img2Img,
     ) -> mx.array:
         # 0. Determine type of image generation
-        is_text2img = img2img.image_path is None
+        if img2img.image_path is None:
-
-        if is_text2img:
             # 1. Create the pure noise
             return LatentCreator.create(
                 seed=seed,
@@ -71,7 +71,7 @@ class LatentCreator:
             )
 
     @staticmethod
-    def encode_image(vae: VAE, image_path: str, height: int, width: int):
+    def encode_image(vae: VAE, image_path: str | Path, height: int, width: int):
         scaled_user_image = ImageUtil.scale_to_dimensions(
             image=ImageUtil.load_image(image_path).convert("RGB"),
             target_width=width,

src/mflux/post_processing/generated_image.py

@@ -1,6 +1,6 @@
 import importlib
-import pathlib
 import typing as t
+from pathlib import Path
 
 import mlx.core as mx
 import PIL.Image
@@ -23,13 +23,13 @@ class GeneratedImage:
         generation_time: float,
         lora_paths: list[str],
         lora_scales: list[float],
-        controlnet_image_path: str | pathlib.Path | None = None,
+        controlnet_image_path: str | Path | None = None,
         controlnet_strength: float | None = None,
-        image_path: str | pathlib.Path | None = None,
+        image_path: str | Path | None = None,
         image_strength: float | None = None,
-        masked_image_path: str | pathlib.Path | None = None,
+        masked_image_path: str | Path | None = None,
-        depth_image_path: str | pathlib.Path | None = None,
+        depth_image_path: str | Path | None = None,
-        redux_image_paths: list[str] | list[pathlib.Path] | None = None,
+        redux_image_paths: list[str] | list[Path] | None = None,
     ):
         self.image = image
         self.model_config = model_config
@@ -78,7 +78,7 @@ class GeneratedImage:
 
     def save(
         self,
-        path: t.Union[str, pathlib.Path],
+        path: t.Union[str, Path],
         export_json_metadata: bool = False,
         overwrite: bool = False,
     ) -> None:
@@ -129,7 +129,7 @@ class GeneratedImage:
     @staticmethod
     def _get_version_from_toml() -> str | None:
         # Search for pyproject.toml by traversing up from the current working directory
-        current_dir = pathlib.Path(__file__).resolve().parent
+        current_dir = Path(__file__).resolve().parent
         for parent in current_dir.parents:
             pyproject_path = parent / "pyproject.toml"
             if pyproject_path.exists():

src/mflux/post_processing/image_util.py

@@ -1,7 +1,7 @@
 import json
 import logging
-import pathlib
 import typing as t
+from pathlib import Path
 
 import mlx.core as mx
 import numpy as np
@@ -27,12 +27,12 @@ class ImageUtil:
         generation_time: float,
         lora_paths: list[str],
         lora_scales: list[float],
-        controlnet_image_path: str | None = None,
+        controlnet_image_path: str | Path | None = None,
-        image_path: str | None = None,
+        image_path: str | Path | None = None,
-        redux_image_paths: list[str] | None = None,
+        redux_image_paths: list[str] | list[Path] | None = None,
         image_strength: float | None = None,
-        masked_image_path: str | None = None,
+        masked_image_path: str | Path | None = None,
-        depth_image_path: str | None = None,
+        depth_image_path: str | Path | None = None,
     ) -> GeneratedImage:
         normalized = ImageUtil._denormalize(decoded_latents)
         normalized_numpy = ImageUtil._to_numpy(normalized)
@@ -113,7 +113,7 @@ class ImageUtil:
         return array
 
     @staticmethod
-    def load_image(path: str | pathlib.Path) -> PIL.Image.Image:
+    def load_image(path: str | Path) -> PIL.Image.Image:
         return PIL.Image.open(path)
 
     @staticmethod
@@ -204,12 +204,12 @@ class ImageUtil:
     @staticmethod
     def save_image(
         image: PIL.Image.Image,
-        path: t.Union[str, pathlib.Path],
+        path: t.Union[str, Path],
         metadata: dict | None = None,
         export_json_metadata: bool = False,
         overwrite: bool = False,
     ) -> None:
-        file_path = pathlib.Path(path)
+        file_path = Path(path)
         file_path.parent.mkdir(parents=True, exist_ok=True)
         file_name = file_path.stem
         file_extension = file_path.suffix
@@ -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)

tools/create_outpaint_image_canvas_and_mask.py

@@ -1,5 +1,5 @@
-import pathlib
 import sys
+from pathlib import Path
 
 from mflux.post_processing.image_util import ImageUtil
 from mflux.ui.box_values import AbsoluteBoxValues, BoxValues
@@ -8,7 +8,7 @@ from mflux.ui.cli.parsers import CommandLineParser
 
 def main():
     parser = CommandLineParser(description="Create expanded canvas and mask for outpainting")
-    parser.add_argument("image_path", type=pathlib.Path, help="Path to the input image file")
+    parser.add_argument("image_path", type=Path, help="Path to the input image file")
     parser.add_image_outpaint_arguments(required=True)
     args = parser.parse_args()