refactor: move a bunch of stuff to utils

This commit is contained in:
Bryce 2023-12-15 13:40:10 -08:00 committed by Bryce Drennan
parent 987af23abe
commit d478771cc0
46 changed files with 100 additions and 84 deletions

View File

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

View File

@ -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,

View File

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

View File

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

View File

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

View File

@ -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,

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -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

View File

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

View File

@ -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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -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

1 test_cache_ordering
7 test_control_images[hed-create_hed_edges]
8 test_control_images[normal-create_normal_map]
9 test_control_images[openpose-create_pose_map]
10 test_control_images[qrcode-adaptive_threshold_binarize]
11 test_controlnet[canny]
12 test_controlnet[colorize]
13 test_controlnet[depth]
16 test_controlnet[inpaint]
17 test_controlnet[normal]
18 test_controlnet[openpose]
19 test_controlnet[qrcode]
20 test_controlnet[shuffle]
21 test_describe_cmd
22 test_describe_picture

View File

@ -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

View File

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

View File

@ -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

View File

@ -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

View File

@ -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]

View File

@ -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

View File

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

View File

@ -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,