From 9eacf5e7ededcf742835ebc02a21f771263b0463 Mon Sep 17 00:00:00 2001 From: Bryce Drennan Date: Fri, 3 Feb 2023 00:43:04 -0500 Subject: [PATCH] perf: improve startup time by doing some imports lazily (#233) just running `aimg --help` or `aimg --version` was very slow due to all the imports being brought in eagerly Before changes `aimg --help` `2.24s user 4.05s system 184% cpu 3.416 total` After changes: `0.04s user 0.02s system 8% cpu 0.625 total` Used `PYTHONPROFILEIMPORTTIME=1 aimg --help` to find time consuming imports. Also switched to using `scripts` instead of `entrypoints` since the scripts are much faster. Made duplicate SAMPLER_TYPE_OPTIONS that can be loaded without loading all the samplers themselves. Likely a breaking change - not sure. --- .github/workflows/ci.yaml | 2 +- README.md | 2 + imaginairy/__init__.py | 9 +---- imaginairy/api.py | 79 ++++++++++++++++++++++----------------- imaginairy/bin/aimg | 6 +++ imaginairy/bin/imagine | 6 +++ imaginairy/cmds.py | 78 ++++++++++++++++++++++++-------------- imaginairy/config.py | 15 ++++++++ imaginairy/debug_info.py | 5 ++- imaginairy/log_utils.py | 16 +++++--- imaginairy/safety.py | 6 +-- imaginairy/schema.py | 29 ++++++++------ imaginairy/version.py | 11 +++--- setup.py | 8 +--- tests/test_config.py | 6 +++ tox.ini | 2 +- 16 files changed, 171 insertions(+), 109 deletions(-) create mode 100644 imaginairy/bin/aimg create mode 100644 imaginairy/bin/imagine create mode 100644 tests/test_config.py diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 82ba4e2..edf9730 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -33,7 +33,7 @@ jobs: python-version: 3.9 - name: Install dependencies run: | - python -m pip install --disable-pip-version-check black isort + python -m pip install --disable-pip-version-check black==22.12.0 isort==5.11.4 - name: Autoformatter run: | black --diff . diff --git a/README.md b/README.md index eb49c26..bc327e2 100644 --- a/README.md +++ b/README.md @@ -291,6 +291,8 @@ docker run -it --gpus all -v $HOME/.cache/huggingface:/root/.cache/huggingface - ## ChangeLog +- perf: cli now has minimal overhead such that `aimg --help` runs in ~650ms instead of ~3400ms + **8.3.1** - fix: init-image-strength type diff --git a/imaginairy/__init__.py b/imaginairy/__init__.py index d4f3eb9..f13573b 100644 --- a/imaginairy/__init__.py +++ b/imaginairy/__init__.py @@ -1,19 +1,12 @@ -import os.path +import os # tells pytorch to allow MPS usage (for Mac M1 compatibility) os.putenv("PYTORCH_ENABLE_MPS_FALLBACK", "1") -import PIL.Image # noqa from .api import imagine, imagine_image_files # noqa -from .enhancers.describe_image_blip import generate_caption # noqa from .schema import ( # noqa ImaginePrompt, ImagineResult, 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 - PIL.Image.Resampling = PIL.Image diff --git a/imaginairy/api.py b/imaginairy/api.py index 37cefd1..093efe3 100755 --- a/imaginairy/api.py +++ b/imaginairy/api.py @@ -2,41 +2,7 @@ import logging import os import re -import numpy as np -import torch -import torch.nn -from einops import rearrange, repeat -from PIL import Image, ImageDraw, ImageOps -from pytorch_lightning import seed_everything -from torch.cuda import OutOfMemoryError - -from imaginairy.animations import make_bounce_animation -from imaginairy.enhancers.clip_masking import get_img_mask -from imaginairy.enhancers.describe_image_blip import generate_caption -from imaginairy.enhancers.face_restoration_codeformer import enhance_faces -from imaginairy.enhancers.upscale_realesrgan import upscale_image -from imaginairy.img_utils import pillow_fit_image_within, pillow_img_to_torch_image -from imaginairy.log_utils import ( - ImageLoggingContext, - log_conditioning, - log_img, - log_latent, -) -from imaginairy.model_manager import get_diffusion_model -from imaginairy.modules.midas.utils import AddMiDaS -from imaginairy.outpaint import outpaint_arg_str_parse, prepare_image_for_outpaint -from imaginairy.safety import SafetyMode, create_safety_score -from imaginairy.samplers import SAMPLER_LOOKUP -from imaginairy.samplers.base import NoiseSchedule, noise_an_image -from imaginairy.samplers.editing import CFGEditingDenoiser -from imaginairy.schema import ImaginePrompt, ImagineResult -from imaginairy.utils import ( - fix_torch_group_norm, - fix_torch_nn_layer_norm, - get_device, - platform_appropriate_autocast, - randn_seeded, -) +from imaginairy.schema import SafetyMode logger = logging.getLogger(__name__) @@ -64,6 +30,11 @@ def imagine_image_files( make_compare_gif=False, return_filename_type="generated", ): + from PIL import ImageDraw + + from imaginairy.animations import make_bounce_animation + from imaginairy.img_utils import pillow_fit_image_within + generated_imgs_path = os.path.join(outdir, "generated") os.makedirs(generated_imgs_path, exist_ok=True) @@ -167,6 +138,16 @@ def imagine( add_caption=False, unsafe_retry_count=1, ): + import torch.nn + + from imaginairy.schema import ImaginePrompt + from imaginairy.utils import ( + fix_torch_group_norm, + fix_torch_nn_layer_norm, + get_device, + platform_appropriate_autocast, + ) + prompts = [ImaginePrompt(prompts)] if isinstance(prompts, str) else prompts prompts = [prompts] if isinstance(prompts, ImaginePrompt) else prompts @@ -214,6 +195,34 @@ def _generate_single_image( half_mode=None, add_caption=False, ): + import numpy as np + import torch.nn + from einops import rearrange, repeat + from PIL import Image, ImageOps + from pytorch_lightning import seed_everything + from torch.cuda import OutOfMemoryError + + from imaginairy.enhancers.clip_masking import get_img_mask + from imaginairy.enhancers.describe_image_blip import generate_caption + from imaginairy.enhancers.face_restoration_codeformer import enhance_faces + from imaginairy.enhancers.upscale_realesrgan import upscale_image + from imaginairy.img_utils import pillow_fit_image_within, pillow_img_to_torch_image + from imaginairy.log_utils import ( + ImageLoggingContext, + log_conditioning, + log_img, + log_latent, + ) + from imaginairy.model_manager import get_diffusion_model + from imaginairy.modules.midas.utils import AddMiDaS + from imaginairy.outpaint import outpaint_arg_str_parse, prepare_image_for_outpaint + from imaginairy.safety import create_safety_score + from imaginairy.samplers import SAMPLER_LOOKUP + from imaginairy.samplers.base import NoiseSchedule, noise_an_image + from imaginairy.samplers.editing import CFGEditingDenoiser + from imaginairy.schema import ImaginePrompt, ImagineResult + from imaginairy.utils import get_device, randn_seeded + latent_channels = 4 downsampling_factor = 8 batch_size = 1 diff --git a/imaginairy/bin/aimg b/imaginairy/bin/aimg new file mode 100644 index 0000000..e547adc --- /dev/null +++ b/imaginairy/bin/aimg @@ -0,0 +1,6 @@ +#!/bin/env python +# -*- coding: utf-8 -*- + +if __name__ == "__main__": + from imaginairy import cmds + cmds.aimg() \ No newline at end of file diff --git a/imaginairy/bin/imagine b/imaginairy/bin/imagine new file mode 100644 index 0000000..f24fe6b --- /dev/null +++ b/imaginairy/bin/imagine @@ -0,0 +1,6 @@ +#!/bin/env python +# -*- coding: utf-8 -*- + +if __name__ == "__main__": + from imaginairy import cmds + cmds.imagine_cmd() \ No newline at end of file diff --git a/imaginairy/cmds.py b/imaginairy/cmds.py index 3bce057..0d2e0f1 100644 --- a/imaginairy/cmds.py +++ b/imaginairy/cmds.py @@ -1,29 +1,10 @@ import logging import math -import os.path import click from click_shell import shell -from tqdm import tqdm -from imaginairy import LazyLoadingImage, __version__, config, generate_caption -from imaginairy.animations import make_bounce_animation -from imaginairy.api import imagine_image_files -from imaginairy.debug_info import get_debug_info -from imaginairy.enhancers.prompt_expansion import expand_prompts -from imaginairy.enhancers.upscale_realesrgan import upscale_image -from imaginairy.log_utils import configure_logging -from imaginairy.prompt_schedules import parse_schedule_strs, prompt_mutator -from imaginairy.samplers import SAMPLER_TYPE_OPTIONS -from imaginairy.schema import ImaginePrompt -from imaginairy.surprise_me import create_surprise_me_images -from imaginairy.train import train_diffusion_model -from imaginairy.training_tools.image_prep import ( - create_class_images, - get_image_filenames, - prep_images, -) -from imaginairy.training_tools.prune_model import prune_diffusion_ckpt +from imaginairy import config logger = logging.getLogger(__name__) @@ -111,7 +92,7 @@ logger = logging.getLogger(__name__) "--sampler", default=config.DEFAULT_SAMPLER, show_default=True, - type=click.Choice(SAMPLER_TYPE_OPTIONS), + type=click.Choice(config.SAMPLER_TYPE_OPTIONS), help="What sampling strategy to use.", ) @click.option( @@ -413,7 +394,7 @@ def imagine_cmd( "--sampler", default=config.DEFAULT_SAMPLER, show_default=True, - type=click.Choice(SAMPLER_TYPE_OPTIONS), + type=click.Choice(config.SAMPLER_TYPE_OPTIONS), help="What sampling strategy to use.", ) @click.option( @@ -603,6 +584,9 @@ def edit_image( # noqa arg_schedules, make_compilation_animation, ): + from imaginairy.log_utils import configure_logging + from imaginairy.surprise_me import create_surprise_me_images + init_image_strength = 1 if surprise_me and prompt_texts: raise ValueError("Cannot use surprise_me and prompt_texts together") @@ -698,11 +682,23 @@ def _imagine_cmd( make_compilation_animation=False, ): """Have the AI generate images. alias:imagine.""" + import os.path + + from imaginairy import LazyLoadingImage + from imaginairy.animations import make_bounce_animation + from imaginairy.api import imagine_image_files + from imaginairy.enhancers.prompt_expansion import expand_prompts + from imaginairy.log_utils import configure_logging + from imaginairy.prompt_schedules import parse_schedule_strs, prompt_mutator + from imaginairy.schema import ImaginePrompt + if ctx.invoked_subcommand is not None: return if version: - print(__version__) + from imaginairy.version import get_version + + print(get_version()) return if quiet: @@ -816,13 +812,14 @@ def aimg(): ✨ Run `aimg` to start a persistent shell session. This makes generation and editing much quicker since the model can stay loaded in memory. """ - configure_logging() @aimg.command() def version(): """Print the version.""" - print(__version__) + from imaginairy.version import get_version + + print(get_version()) @click.argument("image_filepaths", nargs=-1) @@ -838,6 +835,13 @@ def upscale_cmd(image_filepaths, outdir): """ Upscale an image 4x using AI. """ + import os.path + + from tqdm import tqdm + + from imaginairy import LazyLoadingImage + from imaginairy.enhancers.upscale_realesrgan import upscale_image + os.makedirs(outdir, exist_ok=True) for p in tqdm(image_filepaths): @@ -847,7 +851,7 @@ def upscale_cmd(image_filepaths, outdir): else: img = LazyLoadingImage(filepath=p) logger.info( - f"Upscaling {p} from {img.width}x{img.height } to {img.width * 4}x{img.height*4} and saving it to {savepath}" + f"Upscaling {p} from {img.width}x{img.height} to {img.width * 4}x{img.height * 4} and saving it to {savepath}" ) img = upscale_image(img) @@ -859,6 +863,10 @@ def upscale_cmd(image_filepaths, outdir): @aimg.command() def describe(image_filepaths): """Generate text descriptions of images.""" + + from imaginairy import LazyLoadingImage + from imaginairy.enhancers.describe_image_blip import generate_caption + imgs = [] for p in image_filepaths: if p.startswith("http"): @@ -991,6 +999,15 @@ def train_concept( You can find a lot of relevant instructions here: https://github.com/JoePenna/Dreambooth-Stable-Diffusion """ + import os.path + + from imaginairy.train import train_diffusion_model + from imaginairy.training_tools.image_prep import ( + create_class_images, + get_image_filenames, + prep_images, + ) + target_size = 512 # Step 1. Crop and enhance the training images prepped_images_path = os.path.join(concept_images_dir, "prepped-images") @@ -1081,7 +1098,9 @@ def prepare_images(images_dir, is_person, target_size): aimg prep-images --person ./images/selfies aimg prep-images ./images/toy-train """ - configure_logging() + + from imaginairy.training_tools.image_prep import prep_images + prep_images(images_dir=images_dir, is_person=is_person, target_size=target_size) @@ -1097,8 +1116,9 @@ def prune_ckpt(ckpt_paths): Example: aimg prune-ckpt ./path/to/checkpoint.ckpt """ + from imaginairy.training_tools.prune_model import prune_diffusion_ckpt + click.secho("Pruning checkpoint files...") - configure_logging() for p in ckpt_paths: prune_diffusion_ckpt(p) @@ -1108,6 +1128,8 @@ def system_info(): """ Display system information. Submit this when reporting bugs. """ + from imaginairy.debug_info import get_debug_info + for k, v in get_debug_info().items(): k += ":" click.secho(f"{k: <30} {v}") diff --git a/imaginairy/config.py b/imaginairy/config.py index f2296c5..19d862e 100644 --- a/imaginairy/config.py +++ b/imaginairy/config.py @@ -142,3 +142,18 @@ MODEL_CONFIG_SHORTCUTS["openjourney"] = MODEL_CONFIG_SHORTCUTS["openjourney-v2"] MODEL_CONFIG_SHORTCUTS["oj"] = MODEL_CONFIG_SHORTCUTS["openjourney-v2"] MODEL_SHORT_NAMES = sorted(MODEL_CONFIG_SHORTCUTS.keys()) + +SAMPLER_TYPE_OPTIONS = [ + "plms", + "ddim", + "k_dpm_fast", + "k_dpm_adaptive", + "k_lms", + "k_dpm_2", + "k_dpm_2_a", + "k_dpmpp_2m", + "k_dpmpp_2s_a", + "k_euler", + "k_euler_a", + "k_heun", +] diff --git a/imaginairy/debug_info.py b/imaginairy/debug_info.py index 8897fa5..e4cc24c 100644 --- a/imaginairy/debug_info.py +++ b/imaginairy/debug_info.py @@ -3,13 +3,14 @@ import sys import torch -from imaginairy import __version__ from imaginairy.utils import get_device, get_hardware_description +from imaginairy.version import get_version def get_debug_info(): + data = { - "imaginairy_version": __version__, + "imaginairy_version": get_version(), "imaginairy_path": os.path.dirname(__file__), "python_version": f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}", "python_installation_path": sys.executable, diff --git a/imaginairy/log_utils.py b/imaginairy/log_utils.py index e037a9b..0dfde7b 100644 --- a/imaginairy/log_utils.py +++ b/imaginairy/log_utils.py @@ -4,12 +4,6 @@ import re import time import warnings -import torch -from pytorch_lightning import _logger as pytorch_logger -from torchvision.transforms import ToPILImage -from transformers.modeling_utils import logger as modeling_logger -from transformers.utils.logging import _configure_library_root_logger - _CURRENT_LOGGING_CONTEXT = None logger = logging.getLogger(__name__) @@ -159,6 +153,9 @@ class ImageLoggingContext: def log_img(self, img, description): if not self.debug_img_callback: return + import torch + from torchvision.transforms import ToPILImage + self.image_count += 1 if isinstance(img, torch.Tensor): img = ToPILImage()(img.squeeze().cpu().detach()) @@ -200,6 +197,8 @@ def filesafe_text(t): def conditioning_to_img(conditioning): + from torchvision.transforms import ToPILImage + return ToPILImage()(conditioning) @@ -252,6 +251,9 @@ def configure_logging(level="INFO"): def disable_transformers_custom_logging(): + from transformers.modeling_utils import logger as modeling_logger + from transformers.utils.logging import _configure_library_root_logger + _configure_library_root_logger() _logger = modeling_logger.parent _logger.handlers = [] @@ -263,6 +265,8 @@ def disable_transformers_custom_logging(): def disable_pytorch_lighting_custom_logging(): + from pytorch_lightning import _logger as pytorch_logger + try: from pytorch_lightning.utilities.seed import log # noqa diff --git a/imaginairy/safety.py b/imaginairy/safety.py index 1aa9afa..3af6870 100644 --- a/imaginairy/safety.py +++ b/imaginairy/safety.py @@ -6,15 +6,11 @@ from diffusers.pipelines.stable_diffusion import safety_checker as safety_checke from transformers import AutoFeatureExtractor from imaginairy.enhancers.blur_detect import is_blurry +from imaginairy.schema import SafetyMode logger = logging.getLogger(__name__) -class SafetyMode: - STRICT = "strict" - RELAXED = "relaxed" - - class SafetyResult: # increase this value to create a stronger `nfsw` filter # at the cost of increasing the possibility of filtering benign images diff --git a/imaginairy/schema.py b/imaginairy/schema.py index 47f032d..3fc94c0 100644 --- a/imaginairy/schema.py +++ b/imaginairy/schema.py @@ -4,17 +4,8 @@ import logging import os.path import random from datetime import datetime, timezone -from functools import lru_cache - -import requests -from PIL import Image, ImageOps -from urllib3.exceptions import LocationParseError -from urllib3.util import parse_url from imaginairy import config -from imaginairy.model_manager import get_model_default_image_size -from imaginairy.samplers import SAMPLER_LOOKUP, SamplerName -from imaginairy.utils import get_device, get_hardware_description logger = logging.getLogger(__name__) @@ -36,6 +27,9 @@ class LazyLoadingImage: # validate url is valid url if url: + from urllib3.exceptions import LocationParseError + from urllib3.util import parse_url + try: parsed_url = parse_url(url) except LocationParseError: @@ -53,6 +47,7 @@ class LazyLoadingImage: raise AttributeError() if self._img: return getattr(self._img, key) + from PIL import Image, ImageOps if self._lazy_filepath: self._img = Image.open(self._lazy_filepath) @@ -60,6 +55,8 @@ class LazyLoadingImage: f"Loaded input 🖼 of size {self._img.size} from {self._lazy_filepath}" ) elif self._lazy_url: + import requests + self._img = Image.open( requests.get(self._lazy_url, stream=True, timeout=60).raw ) @@ -148,6 +145,8 @@ class ImaginePrompt: self.validate() def validate(self): + from imaginairy.samplers import SAMPLER_LOOKUP, SamplerName + self.prompts = self.process_prompt_input(self.prompts) if self.tile_mode is True: @@ -182,6 +181,8 @@ class ImaginePrompt: ) if self.height is None or self.width is None or self.steps is None: + from imaginairy.model_manager import get_model_default_image_size + SamplerCls = SAMPLER_LOOKUP[self.sampler_type] self.steps = self.steps or SamplerCls.default_steps self.width = self.width or get_model_default_image_size(self.model) @@ -281,6 +282,8 @@ class ImagineResult: timings=None, progress_latents=None, ): + from imaginairy.utils import get_device, get_hardware_description + self.prompt = prompt self.images = {"generated": img} @@ -327,6 +330,8 @@ class ImagineResult: return " ".join(f"{k}:{v:.2f}s" for k, v in self.timings.items()) def _exif(self): + from PIL import Image + exif = Image.Exif() exif[ExifCodes.ImageDescription] = self.prompt.prompt_description() exif[ExifCodes.UserComment] = json.dumps(self.metadata_dict()) @@ -349,6 +354,6 @@ class ImagineResult: img.convert("RGB").save(save_path, exif=self._exif()) -@lru_cache(maxsize=2) -def _get_briefly_cached_url(url): - return requests.get(url, timeout=60) +class SafetyMode: + STRICT = "strict" + RELAXED = "relaxed" diff --git a/imaginairy/version.py b/imaginairy/version.py index 845e42c..cc00fb6 100644 --- a/imaginairy/version.py +++ b/imaginairy/version.py @@ -1,6 +1,7 @@ -from importlib.metadata import PackageNotFoundError, version +def get_version(): + from importlib.metadata import PackageNotFoundError, version -try: - __version__ = version("imaginairy") -except PackageNotFoundError: - __version__ = None + try: + return version("imaginairy") + except PackageNotFoundError: + return None diff --git a/setup.py b/setup.py index b748c90..0d13b4c 100644 --- a/setup.py +++ b/setup.py @@ -16,16 +16,12 @@ setup( "Source": "https://github.com/brycedrennan/imaginAIry", }, packages=find_packages(include=("imaginairy", "imaginairy.*")), - entry_points={ - "console_scripts": [ - "imagine=imaginairy.cmds:imagine_cmd", - "aimg=imaginairy.cmds:aimg", - ], - }, + scripts=["imaginairy/bin/aimg", "imaginairy/bin/imagine"], package_data={ "imaginairy": [ "configs/*.yaml", "data/*.*", + "bin/*.*", "enhancers/phraselists/*.txt", "vendored/clip/*.txt.gz", "vendored/clipseg/*.pth", diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..5bc784c --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,6 @@ +from imaginairy import config +from imaginairy.samplers import SAMPLER_TYPE_OPTIONS + + +def test_sampler_options(): + assert set(config.SAMPLER_TYPE_OPTIONS) == set(SAMPLER_TYPE_OPTIONS) diff --git a/tox.ini b/tox.ini index a44b91b..c2d44f9 100644 --- a/tox.ini +++ b/tox.ini @@ -11,7 +11,7 @@ format = pylint 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, + Z999,C0103,C0301,C0302,C0114,C0115,C0116,C0415, Z999,D100,D101,D102,D103,D105,D106,D107,D200,D202,D203,D205,D212,D400,D401,D406,D407,D413,D415,D417, Z999,E203,E501,E1101,E1131,E1135,E1136, Z999,R0901,R0902,R0903,R0904,R0193,R0912,R0913,R0914,R0915,R1702,