PR review

This commit is contained in:
Anthony Wu 2024-10-06 14:59:15 -10:00
parent fcc9b21d19
commit 5de208d1ba
3 changed files with 9 additions and 7 deletions

View File

@ -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):

View File

@ -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)

View File

@ -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