diff --git a/imaginairy/api/generate.py b/imaginairy/api/generate.py index 3806b34..11f71bc 100755 --- a/imaginairy/api/generate.py +++ b/imaginairy/api/generate.py @@ -2,8 +2,11 @@ import logging import os +from datetime import datetime, timezone from typing import TYPE_CHECKING, Callable +from imaginairy.config import DEFAULT_SHARED_FILE_FORMAT_TEMPLATE, DEFAULT_SHARED_OUTDIR + if TYPE_CHECKING: from imaginairy.schema import ImaginePrompt @@ -24,10 +27,10 @@ _most_recent_result = None def imagine_image_files( prompts: "list[ImaginePrompt] | ImaginePrompt", - outdir: str, + outdir: str = DEFAULT_SHARED_OUTDIR, + format_template: str = DEFAULT_SHARED_FILE_FORMAT_TEMPLATE, precision: str = "autocast", record_step_images: bool = False, - output_file_extension: str = "jpg", print_caption: bool = False, make_gif: bool = False, make_compare_gif: bool = False, @@ -60,14 +63,18 @@ def imagine_image_files( from imaginairy.api.video_sample import generate_video from imaginairy.utils import get_next_filenumber, prompt_normalized from imaginairy.utils.animations import make_bounce_animation + from imaginairy.utils.format_file_name import format_filename from imaginairy.utils.img_utils import pillow_fit_image_within generated_imgs_path = os.path.join(outdir, "generated") os.makedirs(generated_imgs_path, exist_ok=True) base_count = get_next_filenumber(generated_imgs_path) - output_file_extension = output_file_extension.lower() - if output_file_extension not in {"jpg", "png"}: + + format_template, output_file_extension = os.path.splitext(format_template) + if not output_file_extension: + output_file_extension = ".jpg" + elif output_file_extension not in {".jpg", ".png"}: raise ValueError("Must output a png or jpg") if not isinstance(prompts, list): @@ -101,15 +108,29 @@ def imagine_image_files( if prompt.init_image: img_str = f"_img2img-{prompt.init_image_strength}" - basefilename = ( - f"{base_count:06}_{prompt.seed}_{prompt.solver_type.replace('_', '')}{prompt.steps}_" - f"PS{prompt.prompt_strength}{img_str}_{prompt_normalized(prompt.prompt_text)}" + now = datetime.now(timezone.utc) + + format_data = { + "file_sequence_number": f"{base_count:06}", + "seed": f"{prompt.seed}", + "steps": f"{prompt.steps}", + "solver_type": f"{prompt.solver_type.replace('_', '')}", + "prompt_strength": f"PS{prompt.prompt_strength}", + "prompt_text": f"{prompt_normalized(prompt.prompt_text)}", + "img_str": f"{img_str}", + "file_extension": f"{output_file_extension}", + "now": now, + } + + basefilename = format_filename( + format_template=format_template, data=format_data ) + for image_type in result.images: subpath = os.path.join(outdir, image_type) os.makedirs(subpath, exist_ok=True) filepath = os.path.join( - subpath, f"{basefilename}_[{image_type}].{output_file_extension}" + subpath, f"{basefilename}_[{image_type}]{output_file_extension}" ) result.save(filepath, image_type=image_type) logger.info(f" {image_type:<22} {filepath}") diff --git a/imaginairy/cli/edit.py b/imaginairy/cli/edit.py index cdce40e..f6f4a4b 100644 --- a/imaginairy/cli/edit.py +++ b/imaginairy/cli/edit.py @@ -53,7 +53,7 @@ def edit_cmd( negative_prompt, prompt_strength, outdir, - output_file_extension, + format_template, repeats, size, steps, @@ -109,7 +109,7 @@ def edit_cmd( init_image=image_paths, init_image_strength=image_strength, outdir=outdir, - output_file_extension=output_file_extension, + format_template=format_template, repeats=repeats, size=size, steps=steps, diff --git a/imaginairy/cli/imagine.py b/imaginairy/cli/imagine.py index c166366..27716d6 100644 --- a/imaginairy/cli/imagine.py +++ b/imaginairy/cli/imagine.py @@ -84,7 +84,6 @@ def imagine_cmd( init_image, init_image_strength, outdir, - output_file_extension, repeats, size, steps, @@ -122,6 +121,7 @@ def imagine_cmd( control_strength, control_mode, videogen, + format_template, ): """ Generate images via AI. @@ -192,7 +192,7 @@ def imagine_cmd( init_image=init_image, init_image_strength=init_image_strength, outdir=outdir, - output_file_extension=output_file_extension, + format_template=format_template, repeats=repeats, size=size, steps=steps, diff --git a/imaginairy/cli/shared.py b/imaginairy/cli/shared.py index 220ddb4..9dd0fda 100644 --- a/imaginairy/cli/shared.py +++ b/imaginairy/cli/shared.py @@ -7,6 +7,7 @@ from contextlib import contextmanager import click from imaginairy import config +from imaginairy.config import DEFAULT_SHARED_FILE_FORMAT_TEMPLATE logger = logging.getLogger(__name__) @@ -36,7 +37,7 @@ def _imagine_cmd( init_image, init_image_strength, outdir, - output_file_extension, + format_template, repeats, size, steps, @@ -220,8 +221,8 @@ def _imagine_cmd( filenames = imagine_image_files( prompts, outdir=outdir, + format_template=format_template, record_step_images=show_work, - output_file_extension=output_file_extension, print_caption=caption, precision=precision, make_gif=make_gif, @@ -320,11 +321,15 @@ common_options = [ help="Where to write results to.", ), click.option( - "--output-file-extension", - default="jpg", - show_default=True, - type=click.Choice(["jpg", "png"]), - help="Where to write results to.", + "--format", + "format_template", + default=DEFAULT_SHARED_FILE_FORMAT_TEMPLATE, + type=str, + help="Formats the file name. Default value will save '{file_sequence_number:06}_{seed}_{solver_type}{steps}_PS{prompt_strength}{img_str}_{prompt_text}' to the default or specified directory." + " {original_filename}: original name without the extension;" + "{file_sequence_number:pad}: sequence number in directory, can make zero-padded (e.g., 06 for six digits).;" + " {seed}: seed used in generation. {steps}: number of steps used in generation. {prompt_strength}: strength of the prompt. {img_str}: the init image name. {prompt_text}: the prompt text. {solver_type}: the solver used.;" + "{now:%Y-%m-%d:%H-%M-%S}: current date and time, customizable using standard strftime format codes.", ), click.option( "-r", diff --git a/imaginairy/cli/upscale.py b/imaginairy/cli/upscale.py index aa3d72a..44abb0e 100644 --- a/imaginairy/cli/upscale.py +++ b/imaginairy/cli/upscale.py @@ -10,16 +10,19 @@ from imaginairy.config import DEFAULT_UPSCALE_MODEL logger = logging.getLogger(__name__) -DEFAULT_FORMAT_TEMPLATE = "{original_filename}.upscaled{file_extension}" +DEFAULT_UPSCALE_FORMAT_TEMPLATE = "{original_filename}.upscaled{file_extension}" + +DEV_UPSCALE_FORMAT_TEMPLATE = ( + "{file_sequence_number:06}_{algorithm}_{original_filename}.upscaled{file_extension}" +) +DEV_DEFAULT_OUTDIR = "./outputs/upscaled" @click.argument("image_filepaths", nargs=-1, required=False) @click.option( "--outdir", - default="./outputs/upscaled", - show_default=True, type=click.Path(), - help="Where to write results to. Default will be where the directory of the original file.", + help="Where to write results to. Default will be where the directory of the original file directory.", ) @click.option("--fix-faces", is_flag=True) @click.option( @@ -40,7 +43,7 @@ DEFAULT_FORMAT_TEMPLATE = "{original_filename}.upscaled{file_extension}" @click.option( "--format", "format_template", - default="{original_filename}.upscaled{file_extension}", + default="DEFAULT", type=str, help="Formats the file name. Default value will save '{original_filename}.upscaled{file_extension}' to the original directory." " {original_filename}: original name without the extension;" @@ -78,7 +81,12 @@ def upscale_cmd( click.echo(f"{model_name}") return - os.makedirs(outdir, exist_ok=True) + if outdir or format_template == "DEV": + if format_template == "DEV" and outdir is None: + format_template = DEV_UPSCALE_FORMAT_TEMPLATE + outdir = DEV_DEFAULT_OUTDIR + os.makedirs(outdir, exist_ok=True) + image_filepaths = glob_expand_paths(image_filepaths) if not image_filepaths: @@ -88,11 +96,13 @@ def upscale_cmd( return if format_template == "DEV": - format_template = "{file_sequence_number:06}_{algorithm}_{original_filename}.upscaled{file_extension}" + format_template = DEV_UPSCALE_FORMAT_TEMPLATE elif format_template == "DEFAULT": - format_template = DEFAULT_FORMAT_TEMPLATE + format_template = DEFAULT_UPSCALE_FORMAT_TEMPLATE for p in tqdm(image_filepaths): + if outdir is None: + outdir = os.path.dirname(p) savepath = os.path.join(outdir, os.path.basename(p)) if p.startswith("http"): img = LazyLoadingImage(url=p) @@ -107,9 +117,6 @@ def upscale_cmd( if fix_faces: img = enhance_faces(img, fidelity=fix_faces_fidelity) - if format_template == DEFAULT_FORMAT_TEMPLATE: - outdir = os.path.dirname(p) + "/" - file_base_name, extension = os.path.splitext(os.path.basename(p)) base_count = len(os.listdir(outdir)) diff --git a/imaginairy/config.py b/imaginairy/config.py index 458b749..2497961 100644 --- a/imaginairy/config.py +++ b/imaginairy/config.py @@ -7,6 +7,9 @@ DEFAULT_MODEL_WEIGHTS = "sd15" DEFAULT_SOLVER = "ddim" DEFAULT_UPSCALE_MODEL = "realesrgan-x2-plus" +DEFAULT_SHARED_FILE_FORMAT_TEMPLATE = "{file_sequence_number:06}_{seed}_{solver_type}{steps}_PS{prompt_strength}{img_str}_{prompt_text}" +DEFAULT_SHARED_OUTDIR = "./outputs" + DEFAULT_NEGATIVE_PROMPT = ( "Ugly, duplication, duplicates, mutilation, deformed, mutilated, mutation, twisted body, disfigured, bad anatomy, " "out of frame, extra fingers, mutated hands, " diff --git a/imaginairy/utils/format_file_name.py b/imaginairy/utils/format_file_name.py index 33d90d6..596a7c0 100644 --- a/imaginairy/utils/format_file_name.py +++ b/imaginairy/utils/format_file_name.py @@ -2,13 +2,21 @@ import os from urllib.parse import urlparse +class FileFormat: + def __init__(self, format_template: str, directory: dict): + self.format_template = format_template + self.directory = directory + + def __str__(self): + return format_filename(self.format_template, self.data) + + def format_filename(format_template: str, data: dict) -> str: """ Formats the filename based on the provided template and variables. """ if not isinstance(format_template, str): raise TypeError("format argument must be a string") - filename = format_template.format(**data) return filename diff --git a/tests/test_api/test_generate.py b/tests/test_api/test_generate.py index 32b7a6c..3c310a6 100644 --- a/tests/test_api/test_generate.py +++ b/tests/test_api/test_generate.py @@ -1,8 +1,13 @@ +import os import os.path +from unittest.mock import patch import pytest -from imaginairy.api import imagine, imagine_image_files +from imaginairy.api import imagine +from imaginairy.api.generate import ( + imagine_image_files, +) from imaginairy.img_processors.control_modes import CONTROL_MODES from imaginairy.schema import ControlInput, ImaginePrompt, LazyLoadingImage, MaskMode from imaginairy.utils import get_device @@ -370,3 +375,163 @@ def test_large_image(filename_base_for_outputs): img_path = f"{filename_base_for_outputs}.png" assert_image_similar_to_expectation(result.img, img_path=img_path, threshold=35000) + + +class MockPrompt: + def __init__(self, seed, steps): + self.seed = seed + self.steps = steps + self.is_intermediate = False + self.init_image = None + self.init_image_strength = 0.5 + self.solver_type = "k_dpm_2_a" + self.prompt_strength = 1.0 + self.prompt_text = "a dog" + + +class MockResult: + def __init__(self, prompt, images): + self.prompt = prompt + self.images = images + + def save(self, filepath, image_type): + pass + + +@pytest.mark.parametrize( + ( + "prompts", + "outdir", + "format_template", + "precision", + "record_step_images", + "print_caption", + "make_gif", + "make_compare_gif", + "return_filename_type", + "videogen", + "expected_exception", + "expected_files", + ), + [ + # default test case, no outdir, no format_template + ( + [MockPrompt(123, 10)], + None, + None, + None, + False, + False, + False, + False, + None, + False, + None, + ["_kdpm2a10_PSPS1.0_a_dog_[generated].jpg"], + ), + # Custom outdir + ( + [MockPrompt(123, 10)], + "./hello/world/", + None, + None, + False, + False, + False, + False, + None, + False, + None, + ["./hello/world/generated/"], + ), + # Custom format template + ( + [MockPrompt(123, 10)], + None, + "{solver_type}_custom", + None, + False, + False, + False, + False, + None, + False, + None, + ["kdpm2a_custom"], + ), + # using .png in format + ( + [MockPrompt(123, 10)], + None, + "{solver_type}_custom.png", + None, + False, + False, + False, + False, + None, + False, + None, + ["[generated].png"], + ), + # Combination of custom outdir and format template + ( + [MockPrompt(123, 10)], + "./hello/world/", + "{prompt_text}_custom_{solver_type}", + "autocast", + False, + False, + False, + False, + None, + False, + None, + ["./hello/world/generated/a_dog_custom_kdpm2a_[generated].jpg"], + ), + ], +) +def test_imagine_image_file_formatting( + prompts, + outdir, + format_template, + precision, + record_step_images, + print_caption, + make_gif, + make_compare_gif, + return_filename_type, + videogen, + expected_exception, + expected_files, +): + with patch("imaginairy.api.generate.imagine") as mock_imagine: + mock_imagine.return_value = [ + MockResult(MockPrompt(123, 10), {"generated": "path/to/image.jpg"}) + ] + + params = locals() + assert len(params) == 13 + test_params = {} + + skip_params = [ + "expected_exception", + "mock_imagine", + "expected_files", + ] + + for key, value in params.items(): + if key in skip_params: + pass + elif value: + test_params[key] = value + + if expected_exception: + with pytest.raises(expected_exception): + imagine_image_files(**test_params) + else: + filenames = imagine_image_files(**test_params) + assert isinstance(filenames, list) + assert len(filenames) == len(expected_files) + for expected_file, actual_file in zip(expected_files, filenames): + assert expected_file in actual_file diff --git a/tests/test_enhancers/test_upscale.py b/tests/test_enhancers/test_upscale.py index 3fe50c1..65f58b8 100644 --- a/tests/test_enhancers/test_upscale.py +++ b/tests/test_enhancers/test_upscale.py @@ -2,7 +2,6 @@ from unittest.mock import Mock, patch import pytest from click.testing import CliRunner -from PIL import Image from imaginairy.cli.upscale import ( upscale_cmd, @@ -10,33 +9,99 @@ from imaginairy.cli.upscale import ( from tests import TESTS_FOLDER -@pytest.fixture() -def mock_pil_save(): - with patch.object(Image, "save", autospec=True) as mock_save: - yield mock_save - - -def test_upscale_cmd_format_option(): +@pytest.mark.parametrize( + ("format_option", "outdir_option", "expected_directory", "expected_filename"), + [ + # Test no given format with no outdir specified + (None, None, "/tests/data/", "sand_upscale_difficult.upscaled.jpg"), + # Test no given format with outdir specified + ( + None, + "tests/data/temp/", + "tests/data/temp/", + "sand_upscale_difficult.upscaled.jpg", + ), + # Test given format with no outdir specified + ( + "{original_filename}{original_filename}.upscaled{file_extension}", + None, + "/tests/data/", + "sand_upscale_difficultsand_upscale_difficult.upscaled.jpg", + ), + # Test given format and given directory + ( + "{original_filename}{original_filename}.upscaled{file_extension}", + "tests/data/temp/", + "tests/data/temp/", + "sand_upscale_difficultsand_upscale_difficult.upscaled.jpg", + ), + # Test default config with 'DEFAULT' keyword and no outdir specified + ("DEFAULT", None, "/tests/data/", ".upscaled"), + # Test 'DEV' config with no outdir specified + ( + "DEV", + None, + "./outputs/upscaled", + "000000_realesrgan-x2-plus_sand_upscale_difficult.upscaled.jpg", + ), + # Test 'DEFAULT' config with outdir specified + ( + "DEFAULT", + "tests/data/temp/", + "tests/data/temp/", + "tests/data/temp/sand_upscale_difficult.upscaled.jpg", + ), + # Test 'DEV' config with outdir specified + ( + "DEV", + "tests/data/temp/", + "tests/data/temp/", + "tests/data/temp/000000_realesrgan-x2-plus_sand_upscale_difficult.upscaled.jpg", + ), + # save directory specified in both format and outdir + ( + "tests/data/temp/{original_filename}.upscaled{file_extension}", + "tests/data/temp/", + "tests/data/temp/", + "tests/data/temp/sand_upscale_difficult.upscaled.jpg", + ), + # save directory specified in format but not outdir + ( + "tests/data/temp/{original_filename}.upscaled{file_extension}", + None, + "/tests/data/temp/", + "tests/data/temp/sand_upscale_difficult.upscaled.jpg", + ), + ], +) +def test_upscale_cmd_format_option( + format_option, outdir_option, expected_directory, expected_filename +): runner = CliRunner() mock_img = Mock() mock_img.save = Mock() + command_args = ["tests/data/sand_upscale_difficult.jpg"] + if format_option: + command_args.extend(["--format", format_option]) + if outdir_option: + command_args.extend(["--outdir", outdir_option]) + with patch.multiple( "imaginairy.enhancers.upscale", upscale_image=Mock(return_value=mock_img) ), patch( "imaginairy.utils.glob_expand_paths", new=Mock(return_value=[f"{TESTS_FOLDER}/data/sand_upscale_difficult.jpg"]), ): - result = runner.invoke( - upscale_cmd, - [ - "tests/data/sand_upscale_difficult.jpg", - "--format", - "{original_filename}_upscaled_{file_sequence_number}_{algorithm}_{now}", - ], - ) + result = runner.invoke(upscale_cmd, command_args) assert result.exit_code == 0 assert "Saved to " in result.output mock_img.save.assert_called() # Check if save method was called + saved_path = mock_img.save.call_args[0][ + 0 + ] # Get the path where the image was saved + + assert expected_directory in saved_path + assert expected_filename in saved_path