diff --git a/src/mflux/exceptions.py b/src/mflux/exceptions.py index a92b499..7f28df6 100644 --- a/src/mflux/exceptions.py +++ b/src/mflux/exceptions.py @@ -3,11 +3,11 @@ class MFluxException(Exception): class ImageSavingException(MFluxException): - """error ocurred while attempting to save image to storage.""" + """error occurred while attempting to save image to storage.""" class MetadataEmbedException(MFluxException): - """error ocurred while attempting to embed metadata in image""" + """error occurred while attempting to embed metadata in image""" class MFluxUserException(MFluxException): diff --git a/src/mflux/post_processing/generated_image.py b/src/mflux/post_processing/generated_image.py index 251e57b..f2320d8 100644 --- a/src/mflux/post_processing/generated_image.py +++ b/src/mflux/post_processing/generated_image.py @@ -1,5 +1,6 @@ import importlib - +import pathlib +import typing as t import PIL.Image import mlx.core as mx @@ -37,7 +38,7 @@ class GeneratedImage: self.controlnet_image = controlnet_image_path self.controlnet_strength = controlnet_strength - def save(self, path: str, export_json_metadata: bool = False) -> None: + def save(self, path: t.Union[str, pathlib.Path], export_json_metadata: bool = False) -> None: from mflux import ImageUtil ImageUtil.save_image(self.image, path, self._get_metadata(), export_json_metadata) diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index 726df64..ba7c07f 100644 --- a/src/mflux/post_processing/image_util.py +++ b/src/mflux/post_processing/image_util.py @@ -1,7 +1,7 @@ import typing as t import json import logging -from pathlib import Path +import pathlib import PIL import PIL.Image import mlx.core as mx @@ -47,6 +47,7 @@ class ImageUtil: controlnet_strength=config.controlnet_strength if isinstance(config.config, ConfigControlnet) else None, ) + @staticmethod def to_composite_image(generated_images: t.List[GeneratedImage]) -> Image: # stitch horizontally total_width = sum(gen_img.image.width for gen_img in generated_images) @@ -100,11 +101,11 @@ class ImageUtil: @staticmethod def save_image( image: PIL.Image.Image, - path: str, + path: t.Union[str, pathlib.Path], metadata: dict | None = None, export_json_metadata: bool = False ) -> None: # fmt: off - file_path = Path(path) + file_path = pathlib.Path(path) file_path.parent.mkdir(parents=True, exist_ok=True) file_name = file_path.stem file_extension = file_path.suffix