From d478771cc0a0598bec894ebf3ee4eedaeb457b29 Mon Sep 17 00:00:00 2001 From: Bryce Date: Fri, 15 Dec 2023 13:40:10 -0800 Subject: [PATCH] refactor: move a bunch of stuff to utils --- docs/examples/immortal_pearl_earring.py | 2 +- imaginairy/api.py | 33 +++++++++++---------- imaginairy/api_refiners.py | 25 +++++++++------- imaginairy/cli/colorize.py | 2 +- imaginairy/cli/edit_demo.py | 2 +- imaginairy/cli/shared.py | 8 ++--- imaginairy/cli/videogen.py | 2 +- imaginairy/enhancers/blur_detect.py | 2 +- imaginairy/enhancers/clip_masking.py | 6 ++-- imaginairy/enhancers/describe_image_blip.py | 2 +- imaginairy/enhancers/facecrop.py | 2 +- imaginairy/enhancers/prompt_expansion.py | 2 +- imaginairy/enhancers/upscale_riverwing.py | 2 +- imaginairy/img_processors/openpose.py | 2 +- imaginairy/model_manager.py | 2 +- imaginairy/modules/autoencoder.py | 2 +- imaginairy/modules/diffusion/ddpm.py | 2 +- imaginairy/modules/refiners_sd.py | 2 +- imaginairy/paths.py | 5 ---- imaginairy/samplers/base.py | 2 +- imaginairy/samplers/ddim.py | 2 +- imaginairy/samplers/kdiff.py | 2 +- imaginairy/samplers/plms.py | 2 +- imaginairy/schema.py | 6 ++-- imaginairy/{ => utils}/animations.py | 4 +-- imaginairy/{ => utils}/feather_tile.py | 0 imaginairy/{ => utils}/img_utils.py | 2 +- imaginairy/{ => utils}/log_utils.py | 6 ++-- imaginairy/{ => utils}/outpaint.py | 2 +- imaginairy/utils/paths.py | 4 +++ imaginairy/{ => utils}/prompt_schedules.py | 0 imaginairy/{ => utils}/roi_utils.py | 0 imaginairy/{ => utils}/surprise_me.py | 4 +-- imaginairy/vendored/k_diffusion/sampling.py | 2 +- imaginairy/video_sample.py | 2 +- scripts/controlnet_convert.py | 2 +- tests/conftest.py | 5 +++- tests/data/cuda-tests.csv | 2 ++ tests/img_processors/test_control_modes.py | 5 +++- tests/modules/test_autoencoders.py | 8 ++--- tests/test_api.py | 2 +- tests/test_cli/test_cmds.py | 2 +- tests/test_feather_tile.py | 7 +++-- tests/test_outpaint.py | 2 +- tests/test_prompt_schedules.py | 2 +- tests/test_roi_utils.py | 2 +- 46 files changed, 100 insertions(+), 84 deletions(-) delete mode 100644 imaginairy/paths.py rename imaginairy/{ => utils}/animations.py (98%) rename imaginairy/{ => utils}/feather_tile.py (100%) rename imaginairy/{ => utils}/img_utils.py (99%) rename imaginairy/{ => utils}/log_utils.py (97%) rename imaginairy/{ => utils}/outpaint.py (99%) create mode 100644 imaginairy/utils/paths.py rename imaginairy/{ => utils}/prompt_schedules.py (100%) rename imaginairy/{ => utils}/roi_utils.py (100%) rename imaginairy/{ => utils}/surprise_me.py (98%) diff --git a/docs/examples/immortal_pearl_earring.py b/docs/examples/immortal_pearl_earring.py index ebd678a..8af21b7 100644 --- a/docs/examples/immortal_pearl_earring.py +++ b/docs/examples/immortal_pearl_earring.py @@ -5,8 +5,8 @@ from PIL import ImageDraw, ImageFont from tqdm import tqdm from imaginairy.api import imagine -from imaginairy.log_utils import configure_logging from imaginairy.schema import ImaginePrompt, LazyLoadingImage, WeightedPrompt +from imaginairy.utils.log_utils import configure_logging def generate_image_morph_video(): diff --git a/imaginairy/api.py b/imaginairy/api.py index e7c5f80..a9e65ee 100755 --- a/imaginairy/api.py +++ b/imaginairy/api.py @@ -39,9 +39,9 @@ def imagine_image_files( ): from PIL import ImageDraw - from imaginairy.animations import make_bounce_animation - from imaginairy.img_utils import pillow_fit_image_within from imaginairy.utils import get_next_filenumber + from imaginairy.utils.animations import make_bounce_animation + from imaginairy.utils.img_utils import pillow_fit_image_within from imaginairy.video_sample import generate_video generated_imgs_path = os.path.join(outdir, "generated") @@ -237,30 +237,33 @@ def _generate_single_image_compvis( 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 ( + from imaginairy.model_manager import ( + get_diffusion_model, + get_model_default_image_size, + ) + from imaginairy.modules.midas.api import torch_image_to_depth_map + from imaginairy.safety import create_safety_score + from imaginairy.samplers import SOLVER_LOOKUP + from imaginairy.samplers.editing import CFGEditingDenoiser + from imaginairy.schema import ControlInput, ImagineResult, MaskMode + from imaginairy.utils import get_device, randn_seeded + from imaginairy.utils.img_utils import ( add_caption_to_image, pillow_fit_image_within, pillow_img_to_torch_image, pillow_mask_to_latent_mask, torch_img_to_pillow_img, ) - from imaginairy.log_utils import ( + from imaginairy.utils.log_utils import ( ImageLoggingContext, log_conditioning, log_img, log_latent, ) - from imaginairy.model_manager import ( - get_diffusion_model, - get_model_default_image_size, + from imaginairy.utils.outpaint import ( + outpaint_arg_str_parse, + prepare_image_for_outpaint, ) - from imaginairy.modules.midas.api import torch_image_to_depth_map - from imaginairy.outpaint import outpaint_arg_str_parse, prepare_image_for_outpaint - from imaginairy.safety import create_safety_score - from imaginairy.samplers import SOLVER_LOOKUP - from imaginairy.samplers.editing import CFGEditingDenoiser - from imaginairy.schema import ControlInput, ImagineResult, MaskMode - from imaginairy.utils import get_device, randn_seeded latent_channels = 4 downsampling_factor = 8 @@ -785,7 +788,7 @@ def combine_image(original_img, generated_img, mask_img): """Combine the generated image with the original image using the mask image.""" from PIL import Image - from imaginairy.log_utils import log_img + from imaginairy.utils.log_utils import log_img generated_img = generated_img.resize( original_img.size, diff --git a/imaginairy/api_refiners.py b/imaginairy/api_refiners.py index 38aba26..82fa889 100644 --- a/imaginairy/api_refiners.py +++ b/imaginairy/api_refiners.py @@ -35,26 +35,29 @@ def _generate_single_image( 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 ( + from imaginairy.model_manager import ( + get_diffusion_model_refiners, + get_model_default_image_size, + ) + from imaginairy.safety import create_safety_score + from imaginairy.samplers import SolverName + from imaginairy.schema import ImagineResult + from imaginairy.utils import get_device, randn_seeded + from imaginairy.utils.img_utils import ( add_caption_to_image, pillow_fit_image_within, pillow_img_to_torch_image, pillow_mask_to_latent_mask, ) - from imaginairy.log_utils import ( + from imaginairy.utils.log_utils import ( ImageLoggingContext, log_img, log_latent, ) - from imaginairy.model_manager import ( - get_diffusion_model_refiners, - get_model_default_image_size, + from imaginairy.utils.outpaint import ( + outpaint_arg_str_parse, + prepare_image_for_outpaint, ) - from imaginairy.outpaint import outpaint_arg_str_parse, prepare_image_for_outpaint - from imaginairy.safety import create_safety_score - from imaginairy.samplers import SolverName - from imaginairy.schema import ImagineResult - from imaginairy.utils import get_device, randn_seeded if dtype is None: dtype = torch.float16 if half_mode else torch.float32 @@ -451,7 +454,7 @@ def _calc_conditioning( ): import torch - from imaginairy.log_utils import log_conditioning + from imaginairy.utils.log_utils import log_conditioning # need to expand if doing batches neutral_conditioning = _prompts_to_embeddings(negative_prompts, text_encoder) diff --git a/imaginairy/cli/colorize.py b/imaginairy/cli/colorize.py index 23d6ba9..63651fb 100644 --- a/imaginairy/cli/colorize.py +++ b/imaginairy/cli/colorize.py @@ -39,8 +39,8 @@ def colorize_cmd(image_filepaths, outdir, repeats, caption): from tqdm import tqdm from imaginairy.colorize import colorize_img - from imaginairy.log_utils import configure_logging from imaginairy.schema import LazyLoadingImage + from imaginairy.utils.log_utils import configure_logging configure_logging() diff --git a/imaginairy/cli/edit_demo.py b/imaginairy/cli/edit_demo.py index aa32fd1..0be5e39 100644 --- a/imaginairy/cli/edit_demo.py +++ b/imaginairy/cli/edit_demo.py @@ -31,7 +31,7 @@ import click def edit_demo_cmd(image_paths, outdir, height, width): """Make some fun pre-set edits to input photos.""" from imaginairy.cli.shared import imaginairy_click_context - from imaginairy.surprise_me import create_surprise_me_images + from imaginairy.utils.surprise_me import create_surprise_me_images with imaginairy_click_context(): for image_path in image_paths: diff --git a/imaginairy/cli/shared.py b/imaginairy/cli/shared.py index ff7f2e3..be0413a 100644 --- a/imaginairy/cli/shared.py +++ b/imaginairy/cli/shared.py @@ -15,7 +15,7 @@ logger = logging.getLogger(__name__) def imaginairy_click_context(log_level="INFO"): from pydantic import ValidationError - from imaginairy.log_utils import configure_logging + from imaginairy.utils.log_utils import configure_logging errors_to_catch = (FileNotFoundError, ValidationError, ValueError) configure_logging(level=log_level) @@ -91,7 +91,7 @@ def _imagine_cmd( ) print(msg) - from imaginairy.log_utils import configure_logging + from imaginairy.utils.log_utils import configure_logging configure_logging(log_level) @@ -187,7 +187,7 @@ def _imagine_cmd( model_weights=model_weights_path, caption_text=caption_text, ) - from imaginairy.prompt_schedules import ( + from imaginairy.utils.prompt_schedules import ( parse_schedule_strs, prompt_mutator, ) @@ -224,7 +224,7 @@ def _imagine_cmd( comp_imgs = [LazyLoadingImage(filepath=f) for f in filenames] comp_imgs.reverse() - from imaginairy.animations import make_slideshow_animation + from imaginairy.utils.animations import make_slideshow_animation make_slideshow_animation( outpath=new_filename, diff --git a/imaginairy/cli/videogen.py b/imaginairy/cli/videogen.py index d8e19ea..6259d3d 100644 --- a/imaginairy/cli/videogen.py +++ b/imaginairy/cli/videogen.py @@ -70,7 +70,7 @@ def videogen_cmd( aimg videogen --start-image assets/rocket-wide.png """ - from imaginairy.log_utils import configure_logging + from imaginairy.utils.log_utils import configure_logging from imaginairy.video_sample import generate_video configure_logging() diff --git a/imaginairy/enhancers/blur_detect.py b/imaginairy/enhancers/blur_detect.py index 1647db8..c8f2a98 100644 --- a/imaginairy/enhancers/blur_detect.py +++ b/imaginairy/enhancers/blur_detect.py @@ -2,7 +2,7 @@ import cv2 -from imaginairy.img_utils import pillow_img_to_opencv_img +from imaginairy.utils.img_utils import pillow_img_to_opencv_img def calculate_blurriness_level(img): diff --git a/imaginairy/enhancers/clip_masking.py b/imaginairy/enhancers/clip_masking.py index 7ea94ac..49ca892 100644 --- a/imaginairy/enhancers/clip_masking.py +++ b/imaginairy/enhancers/clip_masking.py @@ -9,9 +9,9 @@ import PIL.Image import torch from torchvision import transforms -from imaginairy.img_utils import pillow_fit_image_within -from imaginairy.log_utils import log_img from imaginairy.schema import LazyLoadingImage +from imaginairy.utils.img_utils import pillow_fit_image_within +from imaginairy.utils.log_utils import log_img from imaginairy.vendored.clipseg import CLIPDensePredT weights_url = "https://github.com/timojl/clipseg/raw/master/weights/rd64-uni.pth" @@ -19,7 +19,7 @@ weights_url = "https://github.com/timojl/clipseg/raw/master/weights/rd64-uni.pth @lru_cache def clip_mask_model(): - from imaginairy.paths import PKG_ROOT + from imaginairy.utils.paths import PKG_ROOT model = CLIPDensePredT(version="ViT-B/16", reduce_dim=64, complex_trans_conv=True) model.eval() diff --git a/imaginairy/enhancers/describe_image_blip.py b/imaginairy/enhancers/describe_image_blip.py index 55f948b..411e98f 100644 --- a/imaginairy/enhancers/describe_image_blip.py +++ b/imaginairy/enhancers/describe_image_blip.py @@ -21,7 +21,7 @@ BLIP_EVAL_SIZE = 384 @lru_cache def blip_model(): - from imaginairy.paths import PKG_ROOT + from imaginairy.utils.paths import PKG_ROOT config_path = os.path.join( PKG_ROOT, "vendored", "blip", "configs", "med_config.json" diff --git a/imaginairy/enhancers/facecrop.py b/imaginairy/enhancers/facecrop.py index cbf295b..67009f3 100644 --- a/imaginairy/enhancers/facecrop.py +++ b/imaginairy/enhancers/facecrop.py @@ -3,7 +3,7 @@ import numpy as np from imaginairy.enhancers.face_restoration_codeformer import face_restore_helper -from imaginairy.roi_utils import resize_roi_coordinates, square_roi_coordinate +from imaginairy.utils.roi_utils import resize_roi_coordinates, square_roi_coordinate def detect_faces(img): diff --git a/imaginairy/enhancers/prompt_expansion.py b/imaginairy/enhancers/prompt_expansion.py index 0741e47..29eb2c0 100644 --- a/imaginairy/enhancers/prompt_expansion.py +++ b/imaginairy/enhancers/prompt_expansion.py @@ -7,7 +7,7 @@ import re from functools import lru_cache from string import Formatter -from imaginairy.paths import PKG_ROOT +from imaginairy.utils.paths import PKG_ROOT DEFAULT_PROMPT_LIBRARY_PATHS = [ os.path.join(PKG_ROOT, "vendored", "noodle_soup_prompts"), diff --git a/imaginairy/enhancers/upscale_riverwing.py b/imaginairy/enhancers/upscale_riverwing.py index d0a975f..b6d6821 100644 --- a/imaginairy/enhancers/upscale_riverwing.py +++ b/imaginairy/enhancers/upscale_riverwing.py @@ -7,9 +7,9 @@ import torch import torch.nn.functional as F from torch import nn -from imaginairy.log_utils import log_latent from imaginairy.model_manager import hf_hub_download from imaginairy.utils import get_device, platform_appropriate_autocast +from imaginairy.utils.log_utils import log_latent from imaginairy.vendored import k_diffusion as K from imaginairy.vendored.k_diffusion import layers from imaginairy.vendored.k_diffusion.models.image_v1 import ImageDenoiserModelV1 diff --git a/imaginairy/img_processors/openpose.py b/imaginairy/img_processors/openpose.py index 81e1444..13edcf7 100644 --- a/imaginairy/img_processors/openpose.py +++ b/imaginairy/img_processors/openpose.py @@ -11,9 +11,9 @@ import torch from scipy.ndimage.filters import gaussian_filter from torch import nn -from imaginairy.img_utils import torch_image_to_openvcv_img from imaginairy.model_manager import get_cached_url_path from imaginairy.utils import get_device +from imaginairy.utils.img_utils import torch_image_to_openvcv_img def pad_right_down_corner(img, stride, padValue): diff --git a/imaginairy/model_manager.py b/imaginairy/model_manager.py index 505c0d6..81e3d85 100644 --- a/imaginairy/model_manager.py +++ b/imaginairy/model_manager.py @@ -22,10 +22,10 @@ from safetensors.torch import load_file from imaginairy import config as iconfig from imaginairy.config import IMAGE_WEIGHTS_SHORT_NAMES, ModelArchitecture from imaginairy.modules import attention -from imaginairy.paths import PKG_ROOT from imaginairy.utils import get_device, instantiate_from_config from imaginairy.utils.model_cache import memory_managed_model from imaginairy.utils.named_resolutions import normalize_image_size +from imaginairy.utils.paths import PKG_ROOT logger = logging.getLogger(__name__) diff --git a/imaginairy/modules/autoencoder.py b/imaginairy/modules/autoencoder.py index ac42202..edb0a36 100644 --- a/imaginairy/modules/autoencoder.py +++ b/imaginairy/modules/autoencoder.py @@ -9,11 +9,11 @@ import pytorch_lightning as pl import torch from torch.cuda import OutOfMemoryError -from imaginairy.feather_tile import rebuild_image, tile_image from imaginairy.modules.diffusion.model import Decoder, Encoder from imaginairy.modules.distributions import DiagonalGaussianDistribution from imaginairy.modules.ema import LitEma from imaginairy.utils import instantiate_from_config +from imaginairy.utils.feather_tile import rebuild_image, tile_image logger = logging.getLogger(__name__) diff --git a/imaginairy/modules/diffusion/ddpm.py b/imaginairy/modules/diffusion/ddpm.py index 61e5de7..982e922 100644 --- a/imaginairy/modules/diffusion/ddpm.py +++ b/imaginairy/modules/diffusion/ddpm.py @@ -35,9 +35,9 @@ from imaginairy.modules.diffusion.util import ( ) from imaginairy.modules.distributions import DiagonalGaussianDistribution from imaginairy.modules.ema import LitEma -from imaginairy.paths import PKG_ROOT from imaginairy.samplers.kdiff import DPMPP2MSampler from imaginairy.utils import instantiate_from_config +from imaginairy.utils.paths import PKG_ROOT logger = logging.getLogger(__name__) __conditioning_keys__ = {"concat": "c_concat", "crossattn": "c_crossattn", "adm": "y"} diff --git a/imaginairy/modules/refiners_sd.py b/imaginairy/modules/refiners_sd.py index b7ed317..5e7f17c 100644 --- a/imaginairy/modules/refiners_sd.py +++ b/imaginairy/modules/refiners_sd.py @@ -24,7 +24,7 @@ from torch import Tensor, nn from torch.nn import functional as F from torch.nn.modules.utils import _pair -from imaginairy.feather_tile import rebuild_image, tile_image +from imaginairy.utils.feather_tile import rebuild_image, tile_image from imaginairy.weight_management.conversion import cast_weights logger = logging.getLogger(__name__) diff --git a/imaginairy/paths.py b/imaginairy/paths.py deleted file mode 100644 index 6e13eae..0000000 --- a/imaginairy/paths.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Code for defining package root path""" - -import os.path - -PKG_ROOT = os.path.dirname(__file__) diff --git a/imaginairy/samplers/base.py b/imaginairy/samplers/base.py index b64d844..faafd4f 100644 --- a/imaginairy/samplers/base.py +++ b/imaginairy/samplers/base.py @@ -7,13 +7,13 @@ from abc import ABC import numpy as np import torch -from imaginairy.log_utils import log_latent from imaginairy.modules.diffusion.util import ( extract_into_tensor, make_ddim_sampling_parameters, make_ddim_timesteps, ) from imaginairy.utils import get_device +from imaginairy.utils.log_utils import log_latent logger = logging.getLogger(__name__) diff --git a/imaginairy/samplers/ddim.py b/imaginairy/samplers/ddim.py index ecd28e4..c8813c6 100644 --- a/imaginairy/samplers/ddim.py +++ b/imaginairy/samplers/ddim.py @@ -7,7 +7,6 @@ import numpy as np import torch from tqdm import tqdm -from imaginairy.log_utils import increment_step, log_latent from imaginairy.modules.diffusion.util import extract_into_tensor, noise_like from imaginairy.samplers.base import ( ImageSolver, @@ -17,6 +16,7 @@ from imaginairy.samplers.base import ( mask_blend, ) from imaginairy.utils import get_device +from imaginairy.utils.log_utils import increment_step, log_latent logger = logging.getLogger(__name__) diff --git a/imaginairy/samplers/kdiff.py b/imaginairy/samplers/kdiff.py index 715f67f..a643e0a 100644 --- a/imaginairy/samplers/kdiff.py +++ b/imaginairy/samplers/kdiff.py @@ -7,7 +7,6 @@ from typing import Callable import torch from torch import nn -from imaginairy.log_utils import increment_step, log_latent from imaginairy.samplers.base import ( ImageSolver, SolverName, @@ -15,6 +14,7 @@ from imaginairy.samplers.base import ( mask_blend, ) from imaginairy.utils import get_device +from imaginairy.utils.log_utils import increment_step, log_latent from imaginairy.vendored.k_diffusion import sampling as k_sampling from imaginairy.vendored.k_diffusion.external import CompVisDenoiser, CompVisVDenoiser diff --git a/imaginairy/samplers/plms.py b/imaginairy/samplers/plms.py index 5e4148e..1b54ba7 100644 --- a/imaginairy/samplers/plms.py +++ b/imaginairy/samplers/plms.py @@ -7,7 +7,6 @@ import numpy as np import torch from tqdm import tqdm -from imaginairy.log_utils import increment_step, log_latent from imaginairy.modules.diffusion.util import extract_into_tensor, noise_like from imaginairy.samplers.base import ( ImageSolver, @@ -17,6 +16,7 @@ from imaginairy.samplers.base import ( mask_blend, ) from imaginairy.utils import get_device +from imaginairy.utils.log_utils import increment_step, log_latent logger = logging.getLogger(__name__) diff --git a/imaginairy/schema.py b/imaginairy/schema.py index 3406700..49bb6b2 100644 --- a/imaginairy/schema.py +++ b/imaginairy/schema.py @@ -435,7 +435,7 @@ class ImaginePrompt(BaseModel, protected_namespaces=()): @field_validator("outpaint", mode="after") def validate_outpaint(cls, v): - from imaginairy.outpaint import outpaint_arg_str_parse + from imaginairy.utils.outpaint import outpaint_arg_str_parse outpaint_arg_str_parse(v) return v @@ -712,11 +712,11 @@ class ImagineResult: ): import torch - from imaginairy.img_utils import ( + from imaginairy.utils import get_device, get_hardware_description + from imaginairy.utils.img_utils import ( model_latent_to_pillow_img, torch_img_to_pillow_img, ) - from imaginairy.utils import get_device, get_hardware_description self.prompt = prompt diff --git a/imaginairy/animations.py b/imaginairy/utils/animations.py similarity index 98% rename from imaginairy/animations.py rename to imaginairy/utils/animations.py index 17c4b79..f4fefc2 100644 --- a/imaginairy/animations.py +++ b/imaginairy/utils/animations.py @@ -4,13 +4,13 @@ import os.path import cv2 import torch -from imaginairy.img_utils import ( +from imaginairy.utils import shrink_list +from imaginairy.utils.img_utils import ( add_caption_to_image, imgpaths_to_imgs, model_latents_to_pillow_imgs, pillow_img_to_opencv_img, ) -from imaginairy.utils import shrink_list def make_bounce_animation( diff --git a/imaginairy/feather_tile.py b/imaginairy/utils/feather_tile.py similarity index 100% rename from imaginairy/feather_tile.py rename to imaginairy/utils/feather_tile.py diff --git a/imaginairy/img_utils.py b/imaginairy/utils/img_utils.py similarity index 99% rename from imaginairy/img_utils.py rename to imaginairy/utils/img_utils.py index 961c0ca..2318e2c 100644 --- a/imaginairy/img_utils.py +++ b/imaginairy/utils/img_utils.py @@ -17,9 +17,9 @@ import torch from einops import rearrange, repeat from PIL import Image, ImageDraw, ImageFont -from imaginairy.paths import PKG_ROOT from imaginairy.schema import LazyLoadingImage from imaginairy.utils import get_device +from imaginairy.utils.paths import PKG_ROOT def pillow_fit_image_within( diff --git a/imaginairy/log_utils.py b/imaginairy/utils/log_utils.py similarity index 97% rename from imaginairy/log_utils.py rename to imaginairy/utils/log_utils.py index 1a46f90..6d66f82 100644 --- a/imaginairy/log_utils.py +++ b/imaginairy/utils/log_utils.py @@ -122,7 +122,7 @@ class ImageLoggingContext: ) def log_latents(self, latents, description): - from imaginairy.img_utils import model_latents_to_pillow_imgs + from imaginairy.utils.img_utils import model_latents_to_pillow_imgs if "predicted_latent" in description: if self.progress_latent_callback is not None: @@ -169,7 +169,7 @@ class ImageLoggingContext: ) def log_progress_latent(self, latent): - from imaginairy.img_utils import model_latents_to_pillow_imgs + from imaginairy.utils.img_utils import model_latents_to_pillow_imgs if not self.progress_img_callback: return @@ -244,7 +244,7 @@ def configure_logging(level="INFO"): "formatters": { "standard": { "format": fmt, - "class": "imaginairy.log_utils.ColorIndentingFormatter", + "class": "imaginairy.utils.log_utils.ColorIndentingFormatter", }, }, "handlers": { diff --git a/imaginairy/outpaint.py b/imaginairy/utils/outpaint.py similarity index 99% rename from imaginairy/outpaint.py rename to imaginairy/utils/outpaint.py index 1b4ba3d..ad0e59f 100644 --- a/imaginairy/outpaint.py +++ b/imaginairy/utils/outpaint.py @@ -6,7 +6,7 @@ import torch from PIL import Image, ImageDraw from torch import nn -from imaginairy.img_utils import torch_img_to_pillow_img +from imaginairy.utils.img_utils import torch_img_to_pillow_img def outpaint_calculations( diff --git a/imaginairy/utils/paths.py b/imaginairy/utils/paths.py new file mode 100644 index 0000000..f4ff3f8 --- /dev/null +++ b/imaginairy/utils/paths.py @@ -0,0 +1,4 @@ +"""Code for defining package root path""" +import os + +PKG_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) diff --git a/imaginairy/prompt_schedules.py b/imaginairy/utils/prompt_schedules.py similarity index 100% rename from imaginairy/prompt_schedules.py rename to imaginairy/utils/prompt_schedules.py diff --git a/imaginairy/roi_utils.py b/imaginairy/utils/roi_utils.py similarity index 100% rename from imaginairy/roi_utils.py rename to imaginairy/utils/roi_utils.py diff --git a/imaginairy/surprise_me.py b/imaginairy/utils/surprise_me.py similarity index 98% rename from imaginairy/surprise_me.py rename to imaginairy/utils/surprise_me.py index 74457d5..471bfe0 100644 --- a/imaginairy/surprise_me.py +++ b/imaginairy/utils/surprise_me.py @@ -3,11 +3,11 @@ import logging import os.path -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, ImaginePrompt, LazyLoadingImage +from imaginairy.utils.animations import make_gif_animation +from imaginairy.utils.img_utils import add_caption_to_image, pillow_fit_image_within logger = logging.getLogger(__name__) diff --git a/imaginairy/vendored/k_diffusion/sampling.py b/imaginairy/vendored/k_diffusion/sampling.py index ab7b811..7a98386 100644 --- a/imaginairy/vendored/k_diffusion/sampling.py +++ b/imaginairy/vendored/k_diffusion/sampling.py @@ -6,7 +6,7 @@ from torch import nn from torchdiffeq import odeint from tqdm.auto import tqdm, trange -from imaginairy.log_utils import log_latent +from imaginairy.utils.log_utils import log_latent def append_zero(x): diff --git a/imaginairy/video_sample.py b/imaginairy/video_sample.py index 5169932..f50036f 100644 --- a/imaginairy/video_sample.py +++ b/imaginairy/video_sample.py @@ -20,7 +20,6 @@ from torchvision.transforms import ToTensor 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, @@ -28,6 +27,7 @@ from imaginairy.utils import ( instantiate_from_config, platform_appropriate_autocast, ) +from imaginairy.utils.paths import PKG_ROOT logger = logging.getLogger(__name__) diff --git a/scripts/controlnet_convert.py b/scripts/controlnet_convert.py index 4723003..33defd9 100644 --- a/scripts/controlnet_convert.py +++ b/scripts/controlnet_convert.py @@ -4,7 +4,7 @@ import torch from safetensors.torch import load_file, save_file from imaginairy.model_manager import get_cached_url_path -from imaginairy.paths import PKG_ROOT +from imaginairy.utils.paths import PKG_ROOT sd15_url = "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/889b629140e71758e1e0006e355c331a5744b4bf/v1-5-pruned-emaonly.ckpt" diff --git a/tests/conftest.py b/tests/conftest.py index 6f8ffa6..69a7f0e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,7 +15,6 @@ from urllib3 import HTTPConnectionPool 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, @@ -23,6 +22,10 @@ from imaginairy.utils import ( get_device, platform_appropriate_autocast, ) +from imaginairy.utils.log_utils import ( + configure_logging, + suppress_annoying_logs_and_warnings, +) from tests import TESTS_FOLDER if "pytest" in str(sys.argv): diff --git a/tests/data/cuda-tests.csv b/tests/data/cuda-tests.csv index 71e42fe..8076b83 100644 --- a/tests/data/cuda-tests.csv +++ b/tests/data/cuda-tests.csv @@ -7,6 +7,7 @@ test_control_images[depth-create_depth_map] test_control_images[hed-create_hed_edges] test_control_images[normal-create_normal_map] test_control_images[openpose-create_pose_map] +test_control_images[qrcode-adaptive_threshold_binarize] test_controlnet[canny] test_controlnet[colorize] test_controlnet[depth] @@ -15,6 +16,7 @@ test_controlnet[hed] test_controlnet[inpaint] test_controlnet[normal] test_controlnet[openpose] +test_controlnet[qrcode] test_controlnet[shuffle] test_describe_cmd test_describe_picture diff --git a/tests/img_processors/test_control_modes.py b/tests/img_processors/test_control_modes.py index 1d02d59..028bc0c 100644 --- a/tests/img_processors/test_control_modes.py +++ b/tests/img_processors/test_control_modes.py @@ -2,8 +2,11 @@ import pytest from lightning_fabric import seed_everything 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 imaginairy.utils.img_utils import ( + pillow_img_to_torch_image, + torch_img_to_pillow_img, +) from tests import TESTS_FOLDER from tests.utils import assert_image_similar_to_expectation diff --git a/tests/modules/test_autoencoders.py b/tests/modules/test_autoencoders.py index 7523d10..f40fc8d 100644 --- a/tests/modules/test_autoencoders.py +++ b/tests/modules/test_autoencoders.py @@ -4,14 +4,14 @@ from PIL import Image from torch.nn.functional import interpolate from imaginairy.enhancers.upscale_riverwing import upscale_latent -from imaginairy.img_utils import ( +from imaginairy.model_manager import get_diffusion_model +from imaginairy.schema import LazyLoadingImage +from imaginairy.utils import get_device +from imaginairy.utils.img_utils import ( pillow_fit_image_within, pillow_img_to_torch_image, 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 strat_combos = [ diff --git a/tests/test_api.py b/tests/test_api.py index 16f9b09..c84bbe5 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -4,9 +4,9 @@ import pytest 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, LazyLoadingImage, MaskMode from imaginairy.utils import get_device +from imaginairy.utils.img_utils import pillow_fit_image_within from . import TESTS_FOLDER from .utils import assert_image_similar_to_expectation diff --git a/tests/test_cli/test_cmds.py b/tests/test_cli/test_cmds.py index a46e890..ffcab05 100644 --- a/tests/test_cli/test_cmds.py +++ b/tests/test_cli/test_cmds.py @@ -4,13 +4,13 @@ from unittest import mock import pytest from click.testing import CliRunner -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 import surprise_me from imaginairy.utils.model_cache import GPUModelCache from tests import PROJECT_FOLDER, TESTS_FOLDER from tests.utils import Timer diff --git a/tests/test_feather_tile.py b/tests/test_feather_tile.py index 0b2d804..9541b53 100644 --- a/tests/test_feather_tile.py +++ b/tests/test_feather_tile.py @@ -2,9 +2,12 @@ import itertools import pytest -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 imaginairy.utils.feather_tile import rebuild_image, tile_image, tile_setup +from imaginairy.utils.img_utils import ( + pillow_img_to_torch_image, + torch_img_to_pillow_img, +) 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 4edb2b1..447d0b9 100644 --- a/tests/test_outpaint.py +++ b/tests/test_outpaint.py @@ -1,9 +1,9 @@ import pytest 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 imaginairy.utils.outpaint import outpaint_arg_str_parse from tests import TESTS_FOLDER from tests.utils import assert_image_similar_to_expectation diff --git a/tests/test_prompt_schedules.py b/tests/test_prompt_schedules.py index 5266ac9..ea401fd 100644 --- a/tests/test_prompt_schedules.py +++ b/tests/test_prompt_schedules.py @@ -1,7 +1,7 @@ import pytest -from imaginairy.prompt_schedules import parse_schedule_str from imaginairy.utils import frange +from imaginairy.utils.prompt_schedules import parse_schedule_str @pytest.mark.parametrize( diff --git a/tests/test_roi_utils.py b/tests/test_roi_utils.py index a44110d..cc919e0 100644 --- a/tests/test_roi_utils.py +++ b/tests/test_roi_utils.py @@ -1,7 +1,7 @@ import itertools import random -from imaginairy.roi_utils import ( +from imaginairy.utils.roi_utils import ( RoiNotInBoundsError, resize_roi_coordinates, square_roi_coordinate,