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/dreambooth/state/zip_util.py b/src/mflux/dreambooth/state/zip_util.py index 9f5767a..4bca683 100644 --- a/src/mflux/dreambooth/state/zip_util.py +++ b/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}") 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/latent_creator/latent_creator.py b/src/mflux/latent_creator/latent_creator.py index 0bd9a8b..b2d277b 100644 --- a/src/mflux/latent_creator/latent_creator.py +++ b/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 is_text2img: + if img2img.image_path is None: # 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, diff --git a/src/mflux/post_processing/generated_image.py b/src/mflux/post_processing/generated_image.py index 4250475..1eea9cb 100644 --- a/src/mflux/post_processing/generated_image.py +++ b/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, - depth_image_path: str | pathlib.Path | None = None, - redux_image_paths: list[str] | list[pathlib.Path] | None = None, + masked_image_path: str | Path | None = None, + depth_image_path: str | 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(): diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index 599c9c8..997579d 100644 --- a/src/mflux/post_processing/image_util.py +++ b/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, - 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) @@ -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) diff --git a/tools/create_outpaint_image_canvas_and_mask.py b/tools/create_outpaint_image_canvas_and_mask.py index 4b59f02..76d8722 100755 --- a/tools/create_outpaint_image_canvas_and_mask.py +++ b/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()