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

Author: jaydrennan

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 @@
 
 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 _record_step(img, description, image_count, step_count, prompt):
         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 @@
 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 @@ def temp_f():
         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 @@
 
 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 @@
 @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_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 @@
 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