refactor: move a bunch of stuff to utils

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():

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 (
+    from imaginairy.utils.outpaint import (
-        get_diffusion_model,
+        outpaint_arg_str_parse,
-        get_model_default_image_size,
+        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,

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 (
+    from imaginairy.utils.outpaint import (
-        get_diffusion_model_refiners,
+        outpaint_arg_str_parse,
-        get_model_default_image_size,
+        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)

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()
 

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:

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,

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()

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):

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()

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"

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):

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"),

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

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):

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__)
 

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__)
 

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"}

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__)

imaginairy/paths.py

@@ -1,5 +0,0 @@
-"""Code for defining package root path"""
-
-import os.path
-
-PKG_ROOT = os.path.dirname(__file__)

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__)
 

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__)
 

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
 

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__)
 

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
 

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(

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(

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": {

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(

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__), ".."))

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__)
 

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):

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__)
 

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"
 

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):

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

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
 

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 = [

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

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

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]

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
 

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(

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,