tests: better coverage of cli commands

add "empty" model that just uses empty state for models

imaginairy/model_manager.py

@@ -83,6 +83,9 @@ class MemoryAwareModel:
 
 
 def load_tensors(tensorfile, map_location=None):
+    if tensorfile == "empty":
+        # used for testing
+        return {}
     if tensorfile.endswith(".ckpt"):
         return torch.load(tensorfile, map_location=map_location)
     if tensorfile.endswith(".safetensors"):

tests/test_cmds.py

@@ -1,12 +1,10 @@
-import pytest
 from click.testing import CliRunner
 
-from imaginairy.cmds import edit_image, imagine_cmd
-from imaginairy.utils import get_device
+from imaginairy import ImaginePrompt, LazyLoadingImage, surprise_me
+from imaginairy.cmds import aimg, edit_demo, edit_image, imagine_cmd, upscale_cmd
 from tests import TESTS_FOLDER
 
 
-@pytest.mark.skipif(get_device() == "cpu", reason="Too slow to run on CPU")
 def test_imagine_cmd():
     runner = CliRunner()
     result = runner.invoke(
@@ -14,17 +12,20 @@ def test_imagine_cmd():
         [
             "gold coins",
             "--steps",
-            "25",
+            "2",
             "--outdir",
             f"{TESTS_FOLDER}/test_output",
             "--seed",
             "703425280",
+            "--model",
+            "empty",
+            "--outdir",
+            f"{TESTS_FOLDER}/test_output",
         ],
     )
     assert result.exit_code == 0
 
 
-@pytest.mark.skipif(get_device() == "cpu", reason="Too slow to run on CPU")
 def test_edit_cmd():
     runner = CliRunner()
     result = runner.invoke(
@@ -35,6 +36,66 @@ def test_edit_cmd():
             "1",
             "-p",
             "turn the dog into a cat",
+            "--model",
+            "empty",
+            "--outdir",
+            f"{TESTS_FOLDER}/test_output",
+        ],
+    )
+    assert result.exit_code == 0
+
+
+def test_aimg_shell():
+    runner = CliRunner()
+    result = runner.invoke(
+        aimg,
+        [],
+    )
+    assert "Starting imaginAIry shell" in result.output
+    assert result.exit_code == 0
+
+
+def test_edit_demo(monkeypatch):
+    runner = CliRunner()
+
+    def mock_surprise_me_prompts(*args, **kwargs):
+        return [
+            ImaginePrompt(
+                "",
+                steps=1,
+                width=256,
+                height=256,
+                model="empty",
+            )
+        ]
+
+    monkeypatch.setattr(surprise_me, "surprise_me_prompts", mock_surprise_me_prompts)
+    surprise_me.generic_prompts = []
+    result = runner.invoke(
+        edit_demo,
+        [
+            f"{TESTS_FOLDER}/data/dog.jpg",
+            "--outdir",
+            f"{TESTS_FOLDER}/test_output",
+        ],
+    )
+    assert result.exit_code == 0
+
+
+def test_upscale(monkeypatch):
+    from imaginairy.enhancers import upscale_realesrgan
+
+    def mock_upscale_image(*args, **kwargs):
+        return LazyLoadingImage(filepath=f"{TESTS_FOLDER}/data/dog.jpg")
+
+    monkeypatch.setattr(upscale_realesrgan, "upscale_image", mock_upscale_image)
+    runner = CliRunner()
+    result = runner.invoke(
+        upscale_cmd,
+        [
+            f"{TESTS_FOLDER}/data/dog.jpg",
+            "--outdir",
+            f"{TESTS_FOLDER}/test_output",
         ],
     )
     assert result.exit_code == 0