perf: improve cli startup time

- do not provide automatically imported api functions and objects in `imaginairy` root module
- horrible hack to overcome horrible design choices by easy_install/setuptools

The hack modifies the installed script to remove the __import__ pkg_resources line

If we don't do this then the scripts will be slow to start up because of
pkg_resources.require() which is called by setuptools to ensure the
"correct" version of the package is installed.

before modification example:
```
__requires__ = 'imaginAIry==14.0.0b5'
__import__('pkg_resources').require('imaginAIry==14.0.0b5')
__file__ = '/home/user/projects/imaginairy/imaginairy/bin/aimg'
with open(__file__) as f:
    exec(compile(f.read(), __file__, 'exec'))
```

docs/examples/generate_doc_examples.py

@@ -1,4 +1,5 @@
-from imaginairy import ImaginePrompt, LazyLoadingImage, imagine_image_files
+from imaginairy.api import imagine_image_files
+from imaginairy.schema import ImaginePrompt, LazyLoadingImage
 
 
 def main():

docs/examples/immortal_pearl_earring.py

@@ -4,8 +4,9 @@ import cv2
 from PIL import ImageDraw, ImageFont
 from tqdm import tqdm
 
-from imaginairy import ImaginePrompt, LazyLoadingImage, WeightedPrompt, imagine
+from imaginairy.api import imagine
 from imaginairy.log_utils import configure_logging
+from imaginairy.schema import ImaginePrompt, LazyLoadingImage, WeightedPrompt
 
 
 def generate_image_morph_video():

imaginairy/__init__.py

@@ -4,22 +4,3 @@ import os
 os.putenv("PYTORCH_ENABLE_MPS_FALLBACK", "1")
 # use more memory than we should
 os.putenv("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.0")
-
-import sys  # noqa
-
-from .api import imagine, imagine_image_files  # noqa
-from .schema import (  # noqa
-    ImaginePrompt,
-    ImagineResult,
-    LazyLoadingImage,
-    WeightedPrompt,
-)
-
-# if python version is 3.11 or higher, throw an exception
-if sys.version_info >= (3, 11):
-    msg = (
-        "Imaginairy is not compatible with Python 3.11 or higher. Please use Python 3.8 - 3.10.\n"
-        "This is due to torch 1.13 not supporting Python 3.11 and this library not having yet switched "
-        "to torch 2.0"
-    )
-    raise RuntimeError(msg)

imaginairy/api_refiners.py

@@ -1,10 +1,9 @@
 import logging
 from typing import List, Optional
 
-from imaginairy import ImaginePrompt, WeightedPrompt
 from imaginairy.config import CONTROL_CONFIG_SHORTCUTS
 from imaginairy.model_manager import load_controlnet_adapter
-from imaginairy.schema import MaskMode
+from imaginairy.schema import ImaginePrompt, MaskMode, WeightedPrompt
 
 logger = logging.getLogger(__name__)
 

imaginairy/cli/clickshell_mod.py

@@ -88,6 +88,9 @@ class ImagineColorsCommand(HelpColorsCommand):
         super().__init__(*args, **kwargs)
         self.help_headers_color = "yellow"
         self.help_options_color = "green"
+        from imaginairy.cli.unslow_the_cli import unslowify_scripts_safe
+
+        unslowify_scripts_safe()
 
     def parse_args(self, ctx, args):
         # run the parser for ourselves to preserve the passed order

imaginairy/cli/shared.py

@@ -110,7 +110,8 @@ def _imagine_cmd(
         f"Received {len(prompt_texts)} prompt(s) and {len(init_images)} input image(s). Will repeat the generations {repeats} times to create {total_image_count} images."
     )
 
-    from imaginairy import ImaginePrompt, LazyLoadingImage, imagine_image_files
+    from imaginairy.api import imagine_image_files
+    from imaginairy.schema import ImaginePrompt, LazyLoadingImage
 
     new_init_images = []
     for _init_image in init_images:

imaginairy/cli/unslow_the_cli.py

@@ -0,0 +1,79 @@
+"""
+horrible hack to overcome horrible design choices by easy_install/setuptools
+
+If we don't do this then the scripts will be slow to start up because of
+pkg_resources.require() which is called by setuptools to ensure the
+"correct" version of the package is installed.
+"""
+import os
+
+
+def log(text):
+    # for debugging
+    pass
+    # print(text)
+
+
+def find_script_path(script_name):
+    for path in os.environ["PATH"].split(os.pathsep):
+        script_path = os.path.join(path, script_name)
+        if os.path.isfile(script_path):
+            return script_path
+    return None
+
+
+def is_already_modified():
+    return bool(os.environ.get("IMAGINAIRY_SCRIPT_MODIFIED"))
+
+
+def remove_pkg_resources_requirement(script_path):
+    import shutil
+    import tempfile
+
+    with open(script_path) as file:
+        lines = file.readlines()
+
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file:
+        for line in lines:
+            if "__import__('pkg_resources').require" not in line:
+                temp_file.write(line)
+            else:
+                temp_file.write(
+                    '\nimport os\nos.environ["IMAGINAIRY_SCRIPT_MODIFIED"] = "1"\n'
+                )
+                log(f"Writing to {temp_file.name}")
+
+    # Preserve the original file permissions
+    original_permissions = os.stat(script_path).st_mode
+    os.chmod(temp_file.name, original_permissions)
+
+    # Replace the original file with the modified one
+    shutil.move(temp_file.name, script_path)
+    log(f"Replaced {script_path}")
+
+
+has_run = False
+
+
+def unslowify_scripts():
+    global has_run
+
+    if has_run or is_already_modified():
+        return
+
+    has_run = True
+    script_names = ["aimg", "imagine"]
+
+    for script_name in script_names:
+        script_path = find_script_path(script_name)
+        log(f"Found script {script_name} at {script_path}")
+
+        if script_path:
+            remove_pkg_resources_requirement(script_path)
+
+
+def unslowify_scripts_safe():
+    try:  # noqa
+        unslowify_scripts()
+    except Exception:  # noqa
+        pass

imaginairy/cli/upscale.py

@@ -29,9 +29,9 @@ def upscale_cmd(image_filepaths, outdir, fix_faces, fix_faces_fidelity):
 
     from tqdm import tqdm
 
-    from imaginairy import LazyLoadingImage
     from imaginairy.enhancers.face_restoration_codeformer import enhance_faces
     from imaginairy.enhancers.upscale_realesrgan import upscale_image
+    from imaginairy.schema import LazyLoadingImage
     from imaginairy.utils import glob_expand_paths
 
     os.makedirs(outdir, exist_ok=True)

imaginairy/colorize.py

@@ -2,9 +2,9 @@ import logging
 
 from PIL import Image, ImageEnhance, ImageStat
 
-from imaginairy import ImaginePrompt, imagine
+from imaginairy.api import imagine
 from imaginairy.enhancers.describe_image_blip import generate_caption
-from imaginairy.schema import ControlInput
+from imaginairy.schema import ControlInput, ImaginePrompt
 
 logger = logging.getLogger(__name__)
 
@@ -31,8 +31,7 @@ def colorize_img(img, max_width=1024, max_height=1024, caption=None):
         init_image=img,
         init_image_strength=0.0,
         control_inputs=control_inputs,
-        width=min(img.width, max_width),
-        height=min(img.height, max_height),
+        size=(min(img.width, max_width), min(img.height, max_height)),
         steps=30,
         prompt_strength=12,
     )

imaginairy/http_app/utils.py

@@ -1,7 +1,7 @@
 import base64
 from io import BytesIO
 
-from imaginairy import imagine
+from imaginairy.api import imagine
 
 
 def generate_image(prompt):

imaginairy/prompt_schedules.py

@@ -2,7 +2,7 @@ import csv
 import re
 from copy import copy
 
-from imaginairy import ImaginePrompt
+from imaginairy.schema import ImaginePrompt
 from imaginairy.utils import frange
 
 

imaginairy/surprise_me.py

@@ -1,16 +1,10 @@
-"""
-
-
-aimg.
-"""
-
 import os.path
 
-from imaginairy import ImaginePrompt, LazyLoadingImage, imagine_image_files
 from imaginairy.animations import make_gif_animation
+from imaginairy.api import imagine_image_files
 from imaginairy.enhancers.facecrop import detect_faces
 from imaginairy.img_utils import add_caption_to_image, pillow_fit_image_within
-from imaginairy.schema import ControlInput
+from imaginairy.schema import ControlInput, ImaginePrompt, LazyLoadingImage
 
 preserve_head_kwargs = {
     "mask_prompt": "head|face",

imaginairy/training_tools/image_prep.py

@@ -6,10 +6,11 @@ import re
 from PIL import Image
 from tqdm import tqdm
 
-from imaginairy import ImaginePrompt, LazyLoadingImage, imagine
+from imaginairy.api import imagine
 from imaginairy.enhancers.face_restoration_codeformer import enhance_faces
 from imaginairy.enhancers.facecrop import detect_faces, generate_face_crops
 from imaginairy.enhancers.upscale_realesrgan import upscale_image
+from imaginairy.schema import ImaginePrompt, LazyLoadingImage
 from imaginairy.vendored.smart_crop import SmartCrop
 
 logger = logging.getLogger(__name__)

imaginairy/video_sample.py

@@ -16,9 +16,10 @@ from omegaconf import OmegaConf
 from PIL import Image
 from torchvision.transforms import ToTensor
 
-from imaginairy import LazyLoadingImage, config
+from imaginairy import config
 from imaginairy.model_manager import get_cached_url_path
 from imaginairy.paths import PKG_ROOT
+from imaginairy.schema import LazyLoadingImage
 from imaginairy.utils import (
     default,
     get_device,

scripts/asses_memory_usage.py

@@ -1,22 +1,22 @@
 import torch
 from torch.cuda import OutOfMemoryError
 
-from imaginairy import ImaginePrompt, imagine_image_files
+from imaginairy.api import imagine_image_files
+from imaginairy.schema import ImaginePrompt
 from imaginairy.utils import get_device
 
 
 def assess_memory_usage():
     assert get_device() == "cuda"
     img_size = 3048
-    prompt = ImaginePrompt("strawberries", width=64, height=64, seed=1)
+    prompt = ImaginePrompt("strawberries", size=64, seed=1)
     imagine_image_files([prompt], outdir="outputs")
     datalog = []
     while True:
         torch.cuda.reset_peak_memory_stats()
         prompt = ImaginePrompt(
             "beautiful landscape, Unreal Engine 5, RTX, AAA Game, Detailed 3D Render, Cinema4D",
-            width=img_size,
-            height=img_size,
+            size=img_size,
             seed=1,
             steps=2,
         )

tests/__init__.py

@@ -1,3 +1,4 @@
 import os.path
 
 TESTS_FOLDER = os.path.abspath(os.path.dirname(__file__))
+PROJECT_FOLDER = os.path.abspath(os.path.join(TESTS_FOLDER, ".."))

tests/conftest.py

@@ -10,8 +10,10 @@ import responses
 from tqdm import tqdm
 from urllib3 import HTTPConnectionPool
 
-from imaginairy import ImaginePrompt, api, imagine
+from imaginairy import api
+from imaginairy.api import imagine
 from imaginairy.log_utils import configure_logging, suppress_annoying_logs_and_warnings
+from imaginairy.schema import ImaginePrompt
 from imaginairy.utils import (
     fix_torch_group_norm,
     fix_torch_nn_layer_norm,

tests/enhancers/test_facecrop.py

@@ -1,7 +1,7 @@
 import logging
 
-from imaginairy import LazyLoadingImage
 from imaginairy.enhancers.facecrop import generate_face_crops
+from imaginairy.schema import LazyLoadingImage
 from tests import TESTS_FOLDER
 
 logger = logging.getLogger(__name__)

tests/img_processors/test_control_modes.py

@@ -1,9 +1,9 @@
 import pytest
 from lightning_fabric import seed_everything
 
-from imaginairy import LazyLoadingImage
 from imaginairy.img_processors.control_modes import CONTROL_MODES
 from imaginairy.img_utils import pillow_img_to_torch_image, torch_img_to_pillow_img
+from imaginairy.schema import LazyLoadingImage
 from tests import TESTS_FOLDER
 from tests.utils import assert_image_similar_to_expectation
 

tests/modules/diffusion/test_model.py

@@ -1,23 +1,7 @@
-import time
-
 import torch
 
 from imaginairy.utils import get_device
-
-
-class Timer:
-    def __init__(self, name):
-        self.name = name
-        self.start = None
-
-    def __enter__(self):
-        self.start = time.perf_counter()
-        return self
-
-    def __exit__(self, *args):
-        elapsed = time.perf_counter() - self.start
-
-        print(f"{self.name} took {elapsed*1000:.2f} ms")
+from tests.utils import Timer
 
 
 def test_nonlinearity():

tests/modules/test_autoencoders.py

@@ -3,7 +3,6 @@ import pytest
 from PIL import Image
 from torch.nn.functional import interpolate
 
-from imaginairy import LazyLoadingImage
 from imaginairy.enhancers.upscale_riverwing import upscale_latent
 from imaginairy.img_utils import (
     pillow_fit_image_within,
@@ -11,6 +10,7 @@ from imaginairy.img_utils import (
     torch_img_to_pillow_img,
 )
 from imaginairy.model_manager import get_diffusion_model
+from imaginairy.schema import LazyLoadingImage
 from imaginairy.utils import get_device
 from tests import TESTS_FOLDER
 

tests/test_api.py

@@ -2,11 +2,10 @@ import os.path
 
 import pytest
 
-from imaginairy import LazyLoadingImage
 from imaginairy.api import imagine, imagine_image_files
 from imaginairy.img_processors.control_modes import CONTROL_MODES
 from imaginairy.img_utils import pillow_fit_image_within
-from imaginairy.schema import ControlInput, ImaginePrompt, MaskMode
+from imaginairy.schema import ControlInput, ImaginePrompt, LazyLoadingImage, MaskMode
 from imaginairy.utils import get_device
 
 from . import TESTS_FOLDER

tests/test_cli/test_cmds.py

@@ -1,15 +1,37 @@
+import subprocess
 from unittest import mock
 
+import pytest
 from click.testing import CliRunner
 
-from imaginairy import ImaginePrompt, LazyLoadingImage, surprise_me
+from imaginairy import surprise_me
 from imaginairy.cli.edit import edit_cmd
 from imaginairy.cli.edit_demo import edit_demo_cmd
 from imaginairy.cli.imagine import imagine_cmd
 from imaginairy.cli.main import aimg
 from imaginairy.cli.upscale import upscale_cmd
+from imaginairy.schema import ImaginePrompt, LazyLoadingImage
 from imaginairy.utils.model_cache import GPUModelCache
-from tests import TESTS_FOLDER
+from tests import PROJECT_FOLDER, TESTS_FOLDER
+from tests.utils import Timer
+
+
+@pytest.mark.parametrize("subcommand_name", aimg.commands.keys())
+def test_cmd_help_time(subcommand_name):
+    cmd_parts = [
+        "python",
+        "-X",
+        "importtime",
+        "imaginairy/cli/main.py",
+        subcommand_name,
+        "--help",
+    ]
+    with Timer(f"{subcommand_name} --help") as t:
+        result = subprocess.run(
+            cmd_parts, check=False, capture_output=True, cwd=PROJECT_FOLDER
+        )
+    assert result.returncode == 0, result.stderr
+    assert t.elapsed < 1.0, f"{t.elapsed} > 1.0"
 
 
 def test_imagine_cmd(monkeypatch):

tests/test_enhancers.py

@@ -2,12 +2,13 @@ import pytest
 from PIL import Image
 from pytorch_lightning import seed_everything
 
-from imaginairy import ImaginePrompt, imagine
+from imaginairy.api import imagine
 from imaginairy.enhancers.bool_masker import MASK_PROMPT
 from imaginairy.enhancers.clip_masking import get_img_mask
 from imaginairy.enhancers.describe_image_blip import generate_caption
 from imaginairy.enhancers.describe_image_clip import find_img_text_similarity
 from imaginairy.enhancers.face_restoration_codeformer import enhance_faces
+from imaginairy.schema import ImaginePrompt
 from imaginairy.utils import get_device
 from tests import TESTS_FOLDER
 from tests.utils import assert_image_similar_to_expectation

tests/test_feather_tile.py

@@ -2,9 +2,9 @@ import itertools
 
 import pytest
 
-from imaginairy import LazyLoadingImage
 from imaginairy.feather_tile import rebuild_image, tile_image, tile_setup
 from imaginairy.img_utils import pillow_img_to_torch_image, torch_img_to_pillow_img
+from imaginairy.schema import LazyLoadingImage
 from tests import TESTS_FOLDER
 
 img_ratios = [0.2, 0.242, 0.3, 0.33333333, 0.5, 0.75, 1, 4 / 3.0, 16 / 9.0, 2, 21 / 9.0]

tests/test_outpaint.py

@@ -1,7 +1,8 @@
 import pytest
 
-from imaginairy import ImaginePrompt, LazyLoadingImage, imagine
+from imaginairy.api import imagine
 from imaginairy.outpaint import outpaint_arg_str_parse
+from imaginairy.schema import ImaginePrompt, LazyLoadingImage
 from imaginairy.utils import get_device
 from tests import TESTS_FOLDER
 from tests.utils import assert_image_similar_to_expectation

tests/test_schema/test_controlnetinput.py

@@ -1,8 +1,7 @@
 import pytest
 from pydantic import ValidationError
 
-from imaginairy import LazyLoadingImage
-from imaginairy.schema import ControlInput
+from imaginairy.schema import ControlInput, LazyLoadingImage
 from tests import TESTS_FOLDER
 
 

tests/test_schema/test_imagineprompt.py

@@ -1,8 +1,13 @@
 import pytest
 from pydantic import ValidationError
 
-from imaginairy import LazyLoadingImage, config
-from imaginairy.schema import ControlInput, ImaginePrompt, WeightedPrompt
+from imaginairy import config
+from imaginairy.schema import (
+    ControlInput,
+    ImaginePrompt,
+    LazyLoadingImage,
+    WeightedPrompt,
+)
 from imaginairy.utils.data_distorter import DataDistorter
 from tests import TESTS_FOLDER
 

tests/test_schema/test_lazy_load_image.py

@@ -5,8 +5,7 @@ import pytest
 from PIL import Image
 from pydantic import BaseModel
 
-from imaginairy import LazyLoadingImage
-from imaginairy.schema import InvalidUrlError
+from imaginairy.schema import InvalidUrlError, LazyLoadingImage
 from tests import TESTS_FOLDER
 
 

tests/test_utils/test_model_cache.py

@@ -1,7 +1,8 @@
 import pytest
 from torch import nn
 
-from imaginairy import ImaginePrompt, imagine
+from imaginairy.api import imagine
+from imaginairy.schema import ImaginePrompt
 from imaginairy.utils import get_device
 from imaginairy.utils.model_cache import GPUModelCache
 
@@ -40,7 +41,9 @@ def create_model_of_n_bytes(n):
 def test_memory_usage(filename_base_for_orig_outputs, model_version):
     """Test that we can switch between model versions."""
     prompt_text = "valley, fairytale treehouse village covered, , matte painting, highly detailed, dynamic lighting, cinematic, realism, realistic, photo real, sunset, detailed, high contrast, denoised, centered, michael whelan"
-    prompts = [ImaginePrompt(prompt_text, model=model_version, seed=1, steps=30)]
+    prompts = [
+        ImaginePrompt(prompt_text, model_weights=model_version, seed=1, steps=30)
+    ]
 
     for i, result in enumerate(imagine(prompts)):
         img_path = f"{filename_base_for_orig_outputs}_{result.prompt.prompt_text}_{result.prompt.model}.png"

tests/utils.py

@@ -1,3 +1,5 @@
+import time
+
 import numpy as np
 from PIL import Image
 
@@ -23,3 +25,21 @@ def calc_norm_sum_sq_diff(img, img2):
     )
     norm_sum_sq_diff = sum_sq_diff / np.sqrt(sum_sq_diff)
     return norm_sum_sq_diff
+
+
+class Timer:
+    def __init__(self, name):
+        self.name = name
+        self.start = None
+        self.elapsed = None
+        self.end = None
+
+    def __enter__(self):
+        self.start = time.perf_counter()
+        return self
+
+    def __exit__(self, *args):
+        self.end = time.perf_counter()
+        self.elapsed = self.end - self.start
+
+        print(f"{self.name} took {self.elapsed*1000:.2f} ms")