From 9b95e8b0b630c8dfa97761e9d320e23496db3c44 Mon Sep 17 00:00:00 2001 From: Bryce Date: Sat, 9 Dec 2023 16:33:39 -0800 Subject: [PATCH] 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 | 3 +- docs/examples/immortal_pearl_earring.py | 3 +- imaginairy/__init__.py | 19 ------ imaginairy/api_refiners.py | 3 +- imaginairy/cli/clickshell_mod.py | 3 + imaginairy/cli/shared.py | 3 +- imaginairy/cli/unslow_the_cli.py | 79 ++++++++++++++++++++++ imaginairy/cli/upscale.py | 2 +- imaginairy/colorize.py | 7 +- imaginairy/http_app/utils.py | 2 +- imaginairy/prompt_schedules.py | 2 +- imaginairy/surprise_me.py | 10 +-- imaginairy/training_tools/image_prep.py | 3 +- imaginairy/video_sample.py | 3 +- scripts/asses_memory_usage.py | 8 +-- tests/__init__.py | 1 + tests/conftest.py | 4 +- tests/enhancers/test_facecrop.py | 2 +- tests/img_processors/test_control_modes.py | 2 +- tests/modules/diffusion/test_model.py | 18 +---- tests/modules/test_autoencoders.py | 2 +- tests/test_api.py | 3 +- tests/test_cli/__init__.py | 0 tests/{ => test_cli}/test_cmds.py | 26 ++++++- tests/test_enhancers.py | 3 +- tests/test_feather_tile.py | 2 +- tests/test_outpaint.py | 3 +- tests/test_schema/test_controlnetinput.py | 3 +- tests/test_schema/test_imagineprompt.py | 9 ++- tests/test_schema/test_lazy_load_image.py | 3 +- tests/test_utils/test_model_cache.py | 7 +- tests/utils.py | 20 ++++++ 32 files changed, 177 insertions(+), 81 deletions(-) create mode 100644 imaginairy/cli/unslow_the_cli.py create mode 100644 tests/test_cli/__init__.py rename tests/{ => test_cli}/test_cmds.py (79%) diff --git a/docs/examples/generate_doc_examples.py b/docs/examples/generate_doc_examples.py index 23f92bc..f0e05fd 100644 --- a/docs/examples/generate_doc_examples.py +++ b/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(): diff --git a/docs/examples/immortal_pearl_earring.py b/docs/examples/immortal_pearl_earring.py index aaaa5d9..ebd678a 100644 --- a/docs/examples/immortal_pearl_earring.py +++ b/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(): diff --git a/imaginairy/__init__.py b/imaginairy/__init__.py index 5203ae8..aede5c9 100644 --- a/imaginairy/__init__.py +++ b/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) diff --git a/imaginairy/api_refiners.py b/imaginairy/api_refiners.py index 76124bb..5db0f84 100644 --- a/imaginairy/api_refiners.py +++ b/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__) diff --git a/imaginairy/cli/clickshell_mod.py b/imaginairy/cli/clickshell_mod.py index 4ef3c0e..aa30cf9 100644 --- a/imaginairy/cli/clickshell_mod.py +++ b/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 diff --git a/imaginairy/cli/shared.py b/imaginairy/cli/shared.py index ad0a559..e2ed18f 100644 --- a/imaginairy/cli/shared.py +++ b/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: diff --git a/imaginairy/cli/unslow_the_cli.py b/imaginairy/cli/unslow_the_cli.py new file mode 100644 index 0000000..83892d7 --- /dev/null +++ b/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 diff --git a/imaginairy/cli/upscale.py b/imaginairy/cli/upscale.py index 9684239..6494513 100644 --- a/imaginairy/cli/upscale.py +++ b/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) diff --git a/imaginairy/colorize.py b/imaginairy/colorize.py index ca0ce33..f50f805 100644 --- a/imaginairy/colorize.py +++ b/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, ) diff --git a/imaginairy/http_app/utils.py b/imaginairy/http_app/utils.py index 6855196..2eb7618 100644 --- a/imaginairy/http_app/utils.py +++ b/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): diff --git a/imaginairy/prompt_schedules.py b/imaginairy/prompt_schedules.py index 3072180..31ba094 100644 --- a/imaginairy/prompt_schedules.py +++ b/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 diff --git a/imaginairy/surprise_me.py b/imaginairy/surprise_me.py index 01752bc..95a5d52 100644 --- a/imaginairy/surprise_me.py +++ b/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", diff --git a/imaginairy/training_tools/image_prep.py b/imaginairy/training_tools/image_prep.py index 88d07f1..5a1a331 100644 --- a/imaginairy/training_tools/image_prep.py +++ b/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__) diff --git a/imaginairy/video_sample.py b/imaginairy/video_sample.py index 7fe6c74..2dac33a 100644 --- a/imaginairy/video_sample.py +++ b/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, diff --git a/scripts/asses_memory_usage.py b/scripts/asses_memory_usage.py index afe68e2..9359a09 100644 --- a/scripts/asses_memory_usage.py +++ b/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, ) diff --git a/tests/__init__.py b/tests/__init__.py index ce4152c..635455c 100644 --- a/tests/__init__.py +++ b/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, "..")) diff --git a/tests/conftest.py b/tests/conftest.py index 18d8ad5..582b89f 100644 --- a/tests/conftest.py +++ b/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, diff --git a/tests/enhancers/test_facecrop.py b/tests/enhancers/test_facecrop.py index 63f73ce..b27b1cb 100644 --- a/tests/enhancers/test_facecrop.py +++ b/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__) diff --git a/tests/img_processors/test_control_modes.py b/tests/img_processors/test_control_modes.py index 90f31e8..1d02d59 100644 --- a/tests/img_processors/test_control_modes.py +++ b/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 diff --git a/tests/modules/diffusion/test_model.py b/tests/modules/diffusion/test_model.py index 0303a48..33e86a6 100644 --- a/tests/modules/diffusion/test_model.py +++ b/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(): diff --git a/tests/modules/test_autoencoders.py b/tests/modules/test_autoencoders.py index f68fe47..7523d10 100644 --- a/tests/modules/test_autoencoders.py +++ b/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 diff --git a/tests/test_api.py b/tests/test_api.py index 38abef4..6a0435f 100644 --- a/tests/test_api.py +++ b/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 diff --git a/tests/test_cli/__init__.py b/tests/test_cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_cmds.py b/tests/test_cli/test_cmds.py similarity index 79% rename from tests/test_cmds.py rename to tests/test_cli/test_cmds.py index 5c401db..4327386 100644 --- a/tests/test_cmds.py +++ b/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): diff --git a/tests/test_enhancers.py b/tests/test_enhancers.py index 3bd4d74..2642661 100644 --- a/tests/test_enhancers.py +++ b/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 diff --git a/tests/test_feather_tile.py b/tests/test_feather_tile.py index 748f52e..0b2d804 100644 --- a/tests/test_feather_tile.py +++ b/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] diff --git a/tests/test_outpaint.py b/tests/test_outpaint.py index 5027fb7..4edb2b1 100644 --- a/tests/test_outpaint.py +++ b/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 diff --git a/tests/test_schema/test_controlnetinput.py b/tests/test_schema/test_controlnetinput.py index fa257fa..3452bb1 100644 --- a/tests/test_schema/test_controlnetinput.py +++ b/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 diff --git a/tests/test_schema/test_imagineprompt.py b/tests/test_schema/test_imagineprompt.py index 4ac25b2..6b2aa8c 100644 --- a/tests/test_schema/test_imagineprompt.py +++ b/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 diff --git a/tests/test_schema/test_lazy_load_image.py b/tests/test_schema/test_lazy_load_image.py index 57bf9d6..b200401 100644 --- a/tests/test_schema/test_lazy_load_image.py +++ b/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 diff --git a/tests/test_utils/test_model_cache.py b/tests/test_utils/test_model_cache.py index d75edd6..21a0621 100644 --- a/tests/test_utils/test_model_cache.py +++ b/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" diff --git a/tests/utils.py b/tests/utils.py index 770e271..1765866 100644 --- a/tests/utils.py +++ b/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")