Add option to override saved images

This commit is contained in:
filipstrand 2025-03-01 21:32:14 +01:00
parent 8e5e09f7d8
commit d80b6e520f
2 changed files with 19 additions and 12 deletions

View File

@ -44,10 +44,15 @@ class GeneratedImage:
self.init_image_path = init_image_path
self.init_image_strength = init_image_strength
def save(self, path: t.Union[str, pathlib.Path], export_json_metadata: bool = False) -> None:
def save(
self,
path: t.Union[str, pathlib.Path],
export_json_metadata: bool = False,
overwrite: bool = False,
) -> None:
from mflux import ImageUtil
ImageUtil.save_image(self.image, path, self._get_metadata(), export_json_metadata)
ImageUtil.save_image(self.image, path, self._get_metadata(), export_json_metadata, overwrite)
def _get_metadata(self) -> dict:
"""Generate metadata for reference as well as input data for

View File

@ -114,22 +114,24 @@ class ImageUtil:
@staticmethod
def save_image(
image: PIL.Image.Image,
path: t.Union[str, pathlib.Path],
metadata: dict | None = None,
export_json_metadata: bool = False
image: PIL.Image.Image,
path: t.Union[str, pathlib.Path],
metadata: dict | None = None,
export_json_metadata: bool = False,
overwrite: bool = False
) -> None: # fmt: off
file_path = pathlib.Path(path)
file_path.parent.mkdir(parents=True, exist_ok=True)
file_name = file_path.stem
file_extension = file_path.suffix
# If a file already exists, create a new name with a counter
counter = 1
while file_path.exists():
new_name = f"{file_name}_{counter}{file_extension}"
file_path = file_path.with_name(new_name)
counter += 1
# If a file already exists and overwrite is False, create a new name with a counter
if not overwrite:
counter = 1
while file_path.exists():
new_name = f"{file_name}_{counter}{file_extension}"
file_path = file_path.with_name(new_name)
counter += 1
try:
# Save image without metadata first