diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py index 85f5eb3..6cba76c 100644 --- a/src/mflux/config/runtime_config.py +++ b/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 diff --git a/src/mflux/flux_tools/depth/depth_util.py b/src/mflux/flux_tools/depth/depth_util.py index cbac56b..0e6595f 100644 --- a/src/mflux/flux_tools/depth/depth_util.py +++ b/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): diff --git a/src/mflux/flux_tools/fill/mask_util.py b/src/mflux/flux_tools/fill/mask_util.py index d2ef628..b681dd1 100644 --- a/src/mflux/flux_tools/fill/mask_util.py +++ b/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 diff --git a/src/mflux/flux_tools/redux/flux_redux.py b/src/mflux/flux_tools/redux/flux_redux.py index 6756a3b..ea86f61 100644 --- a/src/mflux/flux_tools/redux/flux_redux.py +++ b/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): diff --git a/src/mflux/flux_tools/redux/redux_util.py b/src/mflux/flux_tools/redux/redux_util.py index f5ef347..dfc5692 100644 --- a/src/mflux/flux_tools/redux/redux_util.py +++ b/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 diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index 4c6d696..997579d 100644 --- a/src/mflux/post_processing/image_util.py +++ b/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)