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'))
```
This commit is contained in:
Bryce 2023-12-09 16:33:39 -08:00 committed by Bryce Drennan
parent 2bd6cb264b
commit 9b95e8b0b6
32 changed files with 177 additions and 81 deletions

View File

@ -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(): def main():

View File

@ -4,8 +4,9 @@ import cv2
from PIL import ImageDraw, ImageFont from PIL import ImageDraw, ImageFont
from tqdm import tqdm 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.log_utils import configure_logging
from imaginairy.schema import ImaginePrompt, LazyLoadingImage, WeightedPrompt
def generate_image_morph_video(): def generate_image_morph_video():

View File

@ -4,22 +4,3 @@ import os
os.putenv("PYTORCH_ENABLE_MPS_FALLBACK", "1") os.putenv("PYTORCH_ENABLE_MPS_FALLBACK", "1")
# use more memory than we should # use more memory than we should
os.putenv("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.0") 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)

View File

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

View File

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

View File

@ -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." 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 = [] new_init_images = []
for _init_image in init_images: for _init_image in init_images:

View File

@ -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

View File

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

View File

@ -2,9 +2,9 @@ import logging
from PIL import Image, ImageEnhance, ImageStat 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.enhancers.describe_image_blip import generate_caption
from imaginairy.schema import ControlInput from imaginairy.schema import ControlInput, ImaginePrompt
logger = logging.getLogger(__name__) 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=img,
init_image_strength=0.0, init_image_strength=0.0,
control_inputs=control_inputs, control_inputs=control_inputs,
width=min(img.width, max_width), size=(min(img.width, max_width), min(img.height, max_height)),
height=min(img.height, max_height),
steps=30, steps=30,
prompt_strength=12, prompt_strength=12,
) )

View File

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

View File

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

View File

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

View File

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

View File

@ -16,9 +16,10 @@ from omegaconf import OmegaConf
from PIL import Image from PIL import Image
from torchvision.transforms import ToTensor 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.model_manager import get_cached_url_path
from imaginairy.paths import PKG_ROOT from imaginairy.paths import PKG_ROOT
from imaginairy.schema import LazyLoadingImage
from imaginairy.utils import ( from imaginairy.utils import (
default, default,
get_device, get_device,

View File

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

View File

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

View File

@ -10,8 +10,10 @@ import responses
from tqdm import tqdm from tqdm import tqdm
from urllib3 import HTTPConnectionPool 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.log_utils import configure_logging, suppress_annoying_logs_and_warnings
from imaginairy.schema import ImaginePrompt
from imaginairy.utils import ( from imaginairy.utils import (
fix_torch_group_norm, fix_torch_group_norm,
fix_torch_nn_layer_norm, fix_torch_nn_layer_norm,

View File

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

View File

@ -1,9 +1,9 @@
import pytest import pytest
from lightning_fabric import seed_everything from lightning_fabric import seed_everything
from imaginairy import LazyLoadingImage
from imaginairy.img_processors.control_modes import CONTROL_MODES 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.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 import TESTS_FOLDER
from tests.utils import assert_image_similar_to_expectation from tests.utils import assert_image_similar_to_expectation

View File

@ -1,23 +1,7 @@
import time
import torch import torch
from imaginairy.utils import get_device from imaginairy.utils import get_device
from tests.utils import Timer
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")
def test_nonlinearity(): def test_nonlinearity():

View File

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

View File

@ -2,11 +2,10 @@ import os.path
import pytest import pytest
from imaginairy import LazyLoadingImage
from imaginairy.api import imagine, imagine_image_files from imaginairy.api import imagine, imagine_image_files
from imaginairy.img_processors.control_modes import CONTROL_MODES from imaginairy.img_processors.control_modes import CONTROL_MODES
from imaginairy.img_utils import pillow_fit_image_within 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 imaginairy.utils import get_device
from . import TESTS_FOLDER from . import TESTS_FOLDER

View File

View File

@ -1,15 +1,37 @@
import subprocess
from unittest import mock from unittest import mock
import pytest
from click.testing import CliRunner 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 import edit_cmd
from imaginairy.cli.edit_demo import edit_demo_cmd from imaginairy.cli.edit_demo import edit_demo_cmd
from imaginairy.cli.imagine import imagine_cmd from imaginairy.cli.imagine import imagine_cmd
from imaginairy.cli.main import aimg from imaginairy.cli.main import aimg
from imaginairy.cli.upscale import upscale_cmd from imaginairy.cli.upscale import upscale_cmd
from imaginairy.schema import ImaginePrompt, LazyLoadingImage
from imaginairy.utils.model_cache import GPUModelCache 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): def test_imagine_cmd(monkeypatch):

View File

@ -2,12 +2,13 @@ import pytest
from PIL import Image from PIL import Image
from pytorch_lightning import seed_everything 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.bool_masker import MASK_PROMPT
from imaginairy.enhancers.clip_masking import get_img_mask from imaginairy.enhancers.clip_masking import get_img_mask
from imaginairy.enhancers.describe_image_blip import generate_caption from imaginairy.enhancers.describe_image_blip import generate_caption
from imaginairy.enhancers.describe_image_clip import find_img_text_similarity from imaginairy.enhancers.describe_image_clip import find_img_text_similarity
from imaginairy.enhancers.face_restoration_codeformer import enhance_faces from imaginairy.enhancers.face_restoration_codeformer import enhance_faces
from imaginairy.schema import ImaginePrompt
from imaginairy.utils import get_device from imaginairy.utils import get_device
from tests import TESTS_FOLDER from tests import TESTS_FOLDER
from tests.utils import assert_image_similar_to_expectation from tests.utils import assert_image_similar_to_expectation

View File

@ -2,9 +2,9 @@ import itertools
import pytest import pytest
from imaginairy import LazyLoadingImage
from imaginairy.feather_tile import rebuild_image, tile_image, tile_setup 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.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 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] 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]

View File

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

View File

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

View File

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

View File

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

View File

@ -1,7 +1,8 @@
import pytest import pytest
from torch import nn 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 import get_device
from imaginairy.utils.model_cache import GPUModelCache 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): def test_memory_usage(filename_base_for_orig_outputs, model_version):
"""Test that we can switch between model versions.""" """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" 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)): for i, result in enumerate(imagine(prompts)):
img_path = f"{filename_base_for_orig_outputs}_{result.prompt.prompt_text}_{result.prompt.model}.png" img_path = f"{filename_base_for_orig_outputs}_{result.prompt.prompt_text}_{result.prompt.model}.png"

View File

@ -1,3 +1,5 @@
import time
import numpy as np import numpy as np
from PIL import Image 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) norm_sum_sq_diff = sum_sq_diff / np.sqrt(sum_sq_diff)
return norm_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")