mirror of
https://github.com/brycedrennan/imaginAIry
synced 2024-10-31 03:20:40 +00:00
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:
parent
2bd6cb264b
commit
9b95e8b0b6
@ -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():
|
||||||
|
@ -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():
|
||||||
|
@ -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)
|
|
||||||
|
@ -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__)
|
||||||
|
|
||||||
|
@ -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
|
||||||
|
@ -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:
|
||||||
|
79
imaginairy/cli/unslow_the_cli.py
Normal file
79
imaginairy/cli/unslow_the_cli.py
Normal 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
|
@ -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)
|
||||||
|
@ -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,
|
||||||
)
|
)
|
||||||
|
@ -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):
|
||||||
|
@ -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
|
||||||
|
|
||||||
|
|
||||||
|
@ -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",
|
||||||
|
@ -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__)
|
||||||
|
@ -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,
|
||||||
|
@ -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,
|
||||||
)
|
)
|
||||||
|
@ -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, ".."))
|
||||||
|
@ -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,
|
||||||
|
@ -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__)
|
||||||
|
@ -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
|
||||||
|
|
||||||
|
@ -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():
|
||||||
|
@ -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
|
||||||
|
|
||||||
|
@ -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
|
||||||
|
0
tests/test_cli/__init__.py
Normal file
0
tests/test_cli/__init__.py
Normal 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):
|
@ -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
|
||||||
|
@ -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]
|
||||||
|
@ -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
|
||||||
|
@ -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
|
||||||
|
|
||||||
|
|
||||||
|
@ -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
|
||||||
|
|
||||||
|
@ -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
|
||||||
|
|
||||||
|
|
||||||
|
@ -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"
|
||||||
|
@ -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")
|
||||||
|
Loading…
Reference in New Issue
Block a user