feature: print version

- fix: default config wasn't being selected for custom models

Author: brycedrennan

.gitignore

@@ -25,4 +25,5 @@ tests/vastai_cli.py
 *.kgrind
 *.pyprof
 **/.polyscope.ini
-**/imgui.ini
+**/imgui.ini
+**/.eggs

README.md

@@ -253,9 +253,9 @@ docker run -it --gpus all -v $HOME/.cache/huggingface:/root/.cache/huggingface -
 ## ChangeLog
 
 **7.6.0**
+- fix: default model config was broken
+- feature: print version with `--version`
 - feature: ability to load safetensors
-
-**7.5.0**
 - feature:  🎉 outpainting. Examples: `--outpaint up10,down300,left50,right50` or `--outpaint all100` or `--outpaint u100,d200,l300,r400`
 
 **7.4.3**

imaginairy/__init__.py

@@ -12,6 +12,7 @@
     LazyLoadingImage,
     WeightedPrompt,
 )
+from .version import __version__  # noqa
 
 # https://stackoverflow.com/questions/71738218/module-pil-has-not-attribute-resampling
 if not hasattr(PIL.Image, "Resampling"):  # Pillow<9.0

imaginairy/cmds.py

@@ -5,9 +5,8 @@
 import click
 from click_shell import shell
 
-from imaginairy import LazyLoadingImage, config, generate_caption
+from imaginairy import LazyLoadingImage, __version__, config, generate_caption
 from imaginairy.api import imagine_image_files
-from imaginairy.config import MODEL_SHORT_NAMES
 from imaginairy.enhancers.prompt_expansion import expand_prompts
 from imaginairy.log_utils import configure_logging
 from imaginairy.samplers import SAMPLER_TYPE_OPTIONS
@@ -198,7 +197,7 @@
 @click.option(
     "--model-weights-path",
     "--model",
-    help=f"Model to use. Should be one of {', '.join(MODEL_SHORT_NAMES)}, or a path to custom weights.",
+    help=f"Model to use. Should be one of {', '.join(config.MODEL_SHORT_NAMES)}, or a path to custom weights.",
     show_default=True,
     default=config.DEFAULT_MODEL,
 )
@@ -215,6 +214,12 @@
     default=None,
     multiple=True,
 )
+@click.option(
+    "--version",
+    default=False,
+    is_flag=True,
+    help="Print the version and exit.",
+)
 @click.pass_context
 def imagine_cmd(
     ctx,
@@ -249,11 +254,16 @@ def imagine_cmd(
     model_weights_path,
     model_config_path,
     prompt_library_path,
+    version,  # noqa
 ):
     """Have the AI generate images. alias:imagine."""
     if ctx.invoked_subcommand is not None:
         return
 
+    if version:
+        print(__version__)
+        return
+
     if quiet:
         log_level = "ERROR"
     configure_logging(log_level)
@@ -329,6 +339,12 @@ def aimg():
     pass
 
 
+@aimg.command()
+def version():
+    """Print the version."""
+    print(__version__)
+
+
 @click.argument("image_filepaths", nargs=-1)
 @aimg.command()
 def describe(image_filepaths):
@@ -381,7 +397,7 @@ def describe(image_filepaths):
     "--model-weights-path",
     "--model",
     "model",
-    help=f"Model to use. Should be one of {', '.join(MODEL_SHORT_NAMES)}, or a path to custom weights.",
+    help=f"Model to use. Should be one of {', '.join(config.MODEL_SHORT_NAMES)}, or a path to custom weights.",
     show_default=True,
     default=config.DEFAULT_MODEL,
 )

imaginairy/model_manager.py

@@ -241,6 +241,9 @@ def resolve_model_paths(
 
     if model_metadata_c:
         config_path = model_metadata_c.config_path
+
+    if config_path is None:
+        config_path = iconfig.MODEL_CONFIG_SHORTCUTS[iconfig.DEFAULT_MODEL].config_path
     model_metadata = model_metadata_w or model_metadata_c
     logger.debug(f"Loading model weights from: {weights_path}")
     logger.debug(f"Loading model config from:  {config_path}")

imaginairy/version.py

@@ -0,0 +1,6 @@
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("imaginairy")
+except PackageNotFoundError:
+    __version__ = None

tests/test_model_manager.py

@@ -0,0 +1,16 @@
+from imaginairy import config
+from imaginairy.model_manager import resolve_model_paths
+
+
+def test_resolved_paths():
+    """Test that the resolved model path is correct."""
+    model_metadata, weights_path, config_path = resolve_model_paths()
+    assert model_metadata.short_name == config.DEFAULT_MODEL
+    assert model_metadata.config_path == config_path
+    default_config_path = config_path
+
+    model_metadata, weights_path, config_path = resolve_model_paths(
+        weights_path="foo.ckpt"
+    )
+    assert weights_path == "foo.ckpt"
+    assert config_path == default_config_path

tests/test_outpaint.py

@@ -2,10 +2,12 @@
 
 from imaginairy import ImaginePrompt, LazyLoadingImage, imagine
 from imaginairy.outpaint import outpaint_arg_str_parse
+from imaginairy.utils import get_device
 from tests import TESTS_FOLDER
 from tests.utils import assert_image_similar_to_expectation
 
 
+@pytest.mark.skipif(get_device() == "cpu", reason="Too slow to run on CPU")
 def test_outpainting_outpaint(filename_base_for_outputs):
     img = LazyLoadingImage(
         filepath=f"{TESTS_FOLDER}/data/girl_with_a_pearl_earring.jpg"

tox.ini

@@ -8,7 +8,7 @@ filterwarnings =
 
 [pylama]
 format = pylint
-skip = */.tox/*,*/.env/*,build/*,*/downloads/*,other/*,prolly_delete/*,downloads/*,imaginairy/vendored/*,testing_support/vastai_cli_official.py
+skip = */.tox/*,*/.env/*,build/*,*/downloads/*,other/*,prolly_delete/*,downloads/*,imaginairy/vendored/*,testing_support/vastai_cli_official.py,.eggs/*
 linters = pylint,pycodestyle,pyflakes,mypy
 ignore =
     Z999,C0103,C0301,C0302,C0114,C0115,C0116,