feature: adds --format flag for cli of imagine, upscale, and edit.

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}")

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,

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,

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",

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

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, "

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
 

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

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