fix: fixes typing and updates test to work with upscale cli template

brycedrennan · Jan 25, 2024 · f3c17a9 · f3c17a9
1 parent 633ff84
commit f3c17a9
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/imaginairy/cli/upscale.py b/imaginairy/cli/upscale.py
@@ -39,7 +39,7 @@
 @click.option(
     "--format",
     "format_template",
-    default="{file_sequence_number:06}_{algorithm}_{original_filename}_UPSCALED{file_extension}",
+    default="{original_filename}.UPSCALED{file_extension}",
     type=str,
     help="Formats the file name. {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).; {algorithm}: upscaling algorithm; "

diff --git a/imaginairy/enhancers/upscale.py b/imaginairy/enhancers/upscale.py
@@ -77,9 +77,11 @@ def upscale_image(
     return image
 
 
-def load_image(img: LazyLoadingImage):
+def load_image(img: LazyLoadingImage | Image.Image):
+    if isinstance(img, LazyLoadingImage):
+        img = img.as_pillow()
     transform = transforms.ToTensor()
-    image_tensor = transform(img.as_pillow())
+    image_tensor = transform(img)
 
     image_tensor = image_tensor.unsqueeze(0)
     return image_tensor.to(get_device())
diff --git a/tests/test_enhancers/test_upscale.py b/tests/test_enhancers/test_upscale.py
@@ -33,7 +33,7 @@ def test_upscale_cmd_format_option():
             [
                 "tests/data/sand_upscale_difficult.jpg",
                 "--format",
-                "{original}_upscaled_{number}_{algorithm}_{date}_{time}",
+                "{original_filename}_upscaled_{file_sequence_number}_{algorithm}_{now}",
             ],
         )