Archives
/
imaginAIry
mirror of https://github.com/brycedrennan/imaginAIry
2
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

refactor: use logging instead of prints

Browse Source
pull/1/head
Bryce 2 years ago
parent f782fac570
commit a26e002a86
10 changed files with 98 additions and 95 deletions
Show Stats Download Patch File Download Diff File

40
imaginairy/imagine.py
Unescape Escape View File

@ -1,9 +1,11 @@
import argparse
import logging
import os
import random
import re
import subprocess
from contextlib import nullcontext
from functools import lru_cache
import PIL
import numpy as np
@ -19,22 +21,21 @@ from imaginairy.models.diffusion.plms import PLMSSampler
from imaginairy.utils import get_device, instantiate_from_config
LIB_PATH = os.path.dirname(__file__)
logger = logging.getLogger(__name__)
def load_model_from_config(config, ckpt, verbose=False):
print(f"Loading model from {ckpt}")
logger.info(f"Loading model from {ckpt}")
pl_sd = torch.load(ckpt, map_location="cpu")
if "global_step" in pl_sd:
print(f"Global Step: {pl_sd['global_step']}")
logger.info(f"Global Step: {pl_sd['global_step']}")
sd = pl_sd["state_dict"]
model = instantiate_from_config(config.model)
m, u = model.load_state_dict(sd, strict=False)
if len(m) > 0 and verbose:
print("missing keys:")
print(m)
logger.info(f"missing keys: {m}")
if len(u) > 0 and verbose:
print("unexpected keys:")
print(u)
logger.info(f"unexpected keys: {u}")
model.cuda()
model.eval()
@ -222,7 +223,7 @@ class ImaginePrompt:
def load_img(path, max_height=512, max_width=512):
image = Image.open(path).convert("RGB")
w, h = image.size
print(f"loaded input image of size ({w}, {h}) from {path}")
logger.info(f"loaded input image of size ({w}, {h}) from {path}")
resize_ratio = min(max_width / w, max_height / h)
w, h = int(w * resize_ratio), int(h * resize_ratio)
w, h = map(lambda x: x - x % 64, (w, h)) # resize to integer multiple of 32
@ -233,23 +234,27 @@ def load_img(path, max_height=512, max_width=512):
return 2.0 * image - 1.0, w, h
@lru_cache()
def load_model():
config = ("data/stable-diffusion-v1.yaml",)
ckpt = ("data/stable-diffusion-v1-4.ckpt",)
config = OmegaConf.load(f"{LIB_PATH}/{config}")
model = load_model_from_config(config, f"{LIB_PATH}/{ckpt}")
model = model.to(get_device())
return model
def imagine(
prompts,
config="data/stable-diffusion-v1.yaml",
ckpt="data/stable-diffusion-v1-4.ckpt",
outdir="outputs/txt2img-samples",
fixed_code=None,
latent_channels=4,
downsampling_factor=8,
precision="autocast",
skip_save=False,
ddim_eta=0.0,
):
config = OmegaConf.load(f"{LIB_PATH}/{config}")
model = load_model_from_config(config, f"{LIB_PATH}/{ckpt}")
model = model.to(get_device())
model = load_model()
os.makedirs(outdir, exist_ok=True)
outpath = outdir
@ -296,11 +301,6 @@ def imagine(
)
start_code = None
# if fixed_code:
# start_code = torch.randn(
# [1, latent_channels, prompt.height, prompt.width],
# device=get_device(),
# )
sampler = get_sampler(prompt.sampler_type, model)
if prompt.init_image:
generation_strength = 1 - prompt.init_image_strength

15
imaginairy/models/autoencoder.py
Unescape Escape View File

@ -1,14 +1,15 @@
import torch
import pytorch_lightning as pl
import torch.nn.functional as F
import logging
import pytorch_lightning as pl
import torch
from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
from imaginairy.modules.diffusionmodules.model import Encoder, Decoder
from imaginairy.modules.distributions import DiagonalGaussianDistribution
from imaginairy.utils import instantiate_from_config
logger = logging.getLogger(__name__)
class VQModel(pl.LightningModule):
def __init__(
@ -51,7 +52,7 @@ class VQModel(pl.LightningModule):
self.monitor = monitor
self.batch_resize_range = batch_resize_range
if self.batch_resize_range is not None:
print(
logger.info(
f"{self.__class__.__name__}: Using per-batch resizing in range {batch_resize_range}."
)
@ -117,10 +118,10 @@ class AutoencoderKL(pl.LightningModule):
for k in keys:
for ik in ignore_keys:
if k.startswith(ik):
print("Deleting key {} from state_dict.".format(k))
logger.info("Deleting key {} from state_dict.".format(k))
del sd[k]
self.load_state_dict(sd, strict=False)
print(f"Restored from {path}")
logger.info(f"Restored from {path}")
def encode(self, x):
h = self.encoder(x)

15
imaginairy/models/diffusion/ddim.py
Unescape Escape View File

@ -1,7 +1,8 @@
"""SAMPLING ONLY."""
import logging
import torch
import numpy as np
import torch
from tqdm import tqdm
from imaginairy.modules.diffusionmodules.util import (
@ -12,6 +13,8 @@ from imaginairy.modules.diffusionmodules.util import (
)
from imaginairy.utils import get_device
logger = logging.getLogger(__name__)
class DDIMSampler:
def __init__(self, model, schedule="linear", **kwargs):
@ -117,12 +120,12 @@ class DDIMSampler:
if isinstance(conditioning, dict):
cbs = conditioning[list(conditioning.keys())[0]].shape[0]
if cbs != batch_size:
print(
logger.warning(
f"Warning: Got {cbs} conditionings but batch-size is {batch_size}"
)
else:
if conditioning.shape[0] != batch_size:
print(
logger.warning(
f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}"
)
@ -130,7 +133,7 @@ class DDIMSampler:
# sampling
C, H, W = shape
size = (batch_size, C, H, W)
print(f"Data shape for DDIM sampling is {size}, eta {eta}")
logger.info(f"Data shape for DDIM sampling is {size}, eta {eta}")
samples, intermediates = self.ddim_sampling(
conditioning,
@ -203,7 +206,7 @@ class DDIMSampler:
else np.flip(timesteps)
)
total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
print(f"Running DDIM Sampling with {total_steps} timesteps")
logger.info(f"Running DDIM Sampling with {total_steps} timesteps")
iterator = tqdm(time_range, desc="DDIM Sampler", total=total_steps)
@ -352,7 +355,7 @@ class DDIMSampler:
time_range = np.flip(timesteps)
total_steps = timesteps.shape[0]
print(f"Running DDIM Sampling with {total_steps} timesteps")
logger.info(f"Running DDIM Sampling with {total_steps} timesteps")
iterator = tqdm(time_range, desc="Decoding image", total=total_steps)
x_dec = x_latent

34
imaginairy/models/diffusion/ddpm.py
Unescape Escape View File

@ -5,14 +5,14 @@ https://github.com/openai/improved-diffusion/blob/e94489283bb876ac1477d5dd7709bb
https://github.com/CompVis/taming-transformers
-- merci
"""
import logging
from functools import partial
import numpy as np
import pytorch_lightning as pl
import torch
import torch.nn as nn
from einops import rearrange, repeat
from einops import rearrange
from torchvision.utils import make_grid
from tqdm import tqdm
@ -21,12 +21,12 @@ from imaginairy.models.autoencoder import (
)
from imaginairy.modules.diffusionmodules.util import (
make_beta_schedule,
extract_into_tensor,
noise_like,
)
from imaginairy.modules.distributions import DiagonalGaussianDistribution
from imaginairy.utils import print_params, instantiate_from_config
from imaginairy.utils import log_params, instantiate_from_config
logger = logging.getLogger(__name__)
__conditioning_keys__ = {"concat": "c_concat", "crossattn": "c_crossattn", "adm": "y"}
@ -77,7 +77,7 @@ class DDPM(pl.LightningModule):
"x0",
], 'currently only supporting "eps" and "x0"'
self.parameterization = parameterization
print(
logger.info(
f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode"
)
self.cond_stage_model = None
@ -88,7 +88,7 @@ class DDPM(pl.LightningModule):
self.channels = channels
self.use_positional_encodings = use_positional_encodings
self.model = DiffusionWrapper(unet_config, conditioning_key)
print_params(self.model)
log_params(self.model)
self.use_scheduler = scheduler_config is not None
if self.use_scheduler:
@ -309,10 +309,12 @@ class LatentDiffusion(DDPM):
def instantiate_cond_stage(self, config):
if not self.cond_stage_trainable:
if config == "__is_first_stage__":
print("Using first stage also as cond stage.")
logger.info("Using first stage also as cond stage.")
self.cond_stage_model = self.first_stage_model
elif config == "__is_unconditional__":
print(f"Training {self.__class__.__name__} as an unconditional model.")
logger.info(
f"Training {self.__class__.__name__} as an unconditional model."
)
self.cond_stage_model = None
# self.be_unconditional = True
else:
@ -576,11 +578,11 @@ class LatentDiffusion(DDPM):
bs, nc, h, w = z.shape
if ks[0] > h or ks[1] > w:
ks = (min(ks[0], h), min(ks[1], w))
print("reducing Kernel")
logger.info("reducing Kernel")
if stride[0] > h or stride[1] > w:
stride = (min(stride[0], h), min(stride[1], w))
print("reducing stride")
logger.info("reducing stride")
fold, unfold, normalization, weighting = self.get_fold_unfold(
z, ks, stride, uf=uf
@ -643,11 +645,11 @@ class LatentDiffusion(DDPM):
bs, nc, h, w = x.shape
if ks[0] > h or ks[1] > w:
ks = (min(ks[0], h), min(ks[1], w))
print("reducing Kernel")
logger.info("reducing Kernel")
if stride[0] > h or stride[1] > w:
stride = (min(stride[0], h), min(stride[1], w))
print("reducing stride")
logger.info("reducing stride")
fold, unfold, normalization, weighting = self.get_fold_unfold(
x, ks, stride, df=df
@ -774,23 +776,21 @@ class LatentDiffusion(DDPM):
)
for bbox in patch_limits
] # list of length l with tensors of shape (1, 2)
print(patch_limits_tknzd[0].shape)
# cut tknzd crop position from conditioning
assert isinstance(cond, dict), "cond must be dict to be fed into model"
cut_cond = cond["c_crossattn"][0][..., :-2].to(self.device)
print(cut_cond.shape)
adapted_cond = torch.stack(
[torch.cat([cut_cond, p], dim=1) for p in patch_limits_tknzd]
)
adapted_cond = rearrange(adapted_cond, "l b n -> (l b) n")
print(adapted_cond.shape)
adapted_cond = self.get_learned_conditioning(adapted_cond)
print(adapted_cond.shape)
adapted_cond = rearrange(
adapted_cond, "(l b) n d -> l b n d", l=z.shape[-1]
)
print(adapted_cond.shape)
cond_list = [{"c_crossattn": [e]} for e in adapted_cond]

13
imaginairy/models/diffusion/plms.py
Unescape Escape View File

@ -1,7 +1,8 @@
"""SAMPLING ONLY."""
import logging
import torch
import numpy as np
import torch
from tqdm import tqdm
from imaginairy.modules.diffusionmodules.util import (
@ -11,6 +12,8 @@ from imaginairy.modules.diffusionmodules.util import (
)
from imaginairy.utils import get_device
logger = logging.getLogger(__name__)
class PLMSSampler(object):
def __init__(self, model, schedule="linear", **kwargs):
@ -118,12 +121,12 @@ class PLMSSampler(object):
if isinstance(conditioning, dict):
cbs = conditioning[list(conditioning.keys())[0]].shape[0]
if cbs != batch_size:
print(
logger.warning(
f"Warning: Got {cbs} conditionings but batch-size is {batch_size}"
)
else:
if conditioning.shape[0] != batch_size:
print(
logger.warning(
f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}"
)
@ -131,7 +134,7 @@ class PLMSSampler(object):
# sampling
C, H, W = shape
size = (batch_size, C, H, W)
print(f"Data shape for PLMS sampling is {size}")
logger.info(f"Data shape for PLMS sampling is {size}")
samples, intermediates = self.plms_sampling(
conditioning,
@ -204,7 +207,7 @@ class PLMSSampler(object):
else np.flip(timesteps)
)
total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
print(f"Running PLMS Sampling with {total_steps} timesteps")
logger.info(f"Running PLMS Sampling with {total_steps} timesteps")
iterator = tqdm(time_range, desc="PLMS Sampler", total=total_steps)
old_eps = []

22
imaginairy/modules/clip_embedders.py
Unescape Escape View File

@ -1,23 +1,14 @@
import clip
import kornia
import torch
import torch.nn as nn
import clip
from einops import repeat
from transformers import CLIPTokenizer, CLIPTextModel
import kornia
from imaginairy.utils import get_device, print_params
from imaginairy.utils import get_device
class AbstractEncoder(nn.Module):
def __init__(self):
super().__init__()
def encode(self, *args, **kwargs):
raise NotImplementedError
class FrozenCLIPEmbedder(AbstractEncoder):
class FrozenCLIPEmbedder:
"""Uses the CLIP transformer encoder for text (from Hugging Face)"""
def __init__(
@ -139,8 +130,3 @@ class FrozenClipImageEmbedder(nn.Module):
def forward(self, x):
# x is assumed to be in range [-1,1]
return self.model.encode_image(self.preprocess(x))
if __name__ == "__main__":
model = FrozenCLIPEmbedder()
print_params(model, verbose=True)

17
imaginairy/modules/diffusionmodules/model.py
Unescape Escape View File

@ -1,4 +1,5 @@
# pytorch_diffusion + derived encoder decoder
import logging
import math
import numpy as np
@ -9,6 +10,8 @@ from einops import rearrange
from imaginairy.modules.attention import LinearAttention
from imaginairy.utils import instantiate_from_config
logger = logging.getLogger(__name__)
def get_timestep_embedding(timesteps, embedding_dim):
"""
@ -193,7 +196,9 @@ class AttnBlock(nn.Module):
def make_attn(in_channels, attn_type="vanilla"):
assert attn_type in ["vanilla", "linear", "none"], f"attn_type {attn_type} unknown"
print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
logger.info(
f"making attention of type '{attn_type}' with {in_channels} in_channels"
)
if attn_type == "vanilla":
return AttnBlock(in_channels)
elif attn_type == "none":
@ -356,10 +361,8 @@ class Decoder(nn.Module):
block_in = ch * ch_mult[self.num_resolutions - 1]
curr_res = resolution // 2 ** (self.num_resolutions - 1)
self.z_shape = (1, z_channels, curr_res, curr_res)
print(
"Working with z of shape {} = {} dimensions.".format(
self.z_shape, np.prod(self.z_shape)
)
logger.info(
f"Working with z of shape {self.z_shape} = {np.prod(self.z_shape)} dimensions."
)
# z to block_in
@ -513,7 +516,7 @@ class Upsampler(nn.Module):
assert out_size >= in_size
num_blocks = int(np.log2(out_size // in_size)) + 1
factor_up = 1.0 + (out_size % in_size)
print(
logger.info(
f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}"
)
self.rescaler = LatentRescaler(
@ -545,7 +548,7 @@ class Resize(nn.Module):
self.with_conv = learned
self.mode = mode
if self.with_conv:
print(
logger.info(
f"Note: {self.__class__.__name} uses learned downsampling and will ignore the fixed {mode} mode"
)
raise NotImplementedError()

29
imaginairy/modules/diffusionmodules/util.py
Unescape Escape View File

@ -8,6 +8,7 @@
# thanks!
import logging
import math
import numpy as np
@ -17,6 +18,8 @@ from einops import repeat
from imaginairy.utils import instantiate_from_config
logger = logging.getLogger(__name__)
def make_beta_schedule(
schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3
@ -64,7 +67,9 @@ def frange(start, stop, step):
def make_ddim_timesteps(
ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True
ddim_discr_method,
num_ddim_timesteps,
num_ddpm_timesteps,
):
if ddim_discr_method == "uniform":
c = num_ddpm_timesteps / num_ddim_timesteps
@ -82,28 +87,26 @@ def make_ddim_timesteps(
# assert ddim_timesteps.shape[0] == num_ddim_timesteps
# add one to get the final alpha values right (the ones from first scale to data during sampling)
steps_out = ddim_timesteps + 1
if verbose:
print(f"Selected timesteps for ddim sampler: {steps_out}")
logger.debug(f"Selected timesteps for ddim sampler: {steps_out}")
return steps_out
def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True):
def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta):
# select alphas for computing the variance schedule
alphas = alphacums[ddim_timesteps]
alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
# according the the formula provided in https://arxiv.org/abs/2010.02502
# according to the formula provided in https://arxiv.org/abs/2010.02502
sigmas = eta * np.sqrt(
(1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev)
)
if verbose:
print(
f"Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}"
)
print(
f"For the chosen value of eta, which is {eta}, "
f"this results in the following sigma_t schedule for ddim sampler {sigmas}"
)
logger.debug(
f"Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}"
)
logger.debug(
f"For the chosen value of eta, which is {eta}, "
f"this results in the following sigma_t schedule for ddim sampler {sigmas}"
)
return sigmas, alphas, alphas_prev

7
imaginairy/utils.py
Unescape Escape View File

@ -1,8 +1,11 @@
import importlib
import logging
from functools import lru_cache
import torch
logger = logging.getLogger(__name__)
@lru_cache()
def get_device():
@ -14,9 +17,9 @@ def get_device():
return "cpu"
def print_params(model):
def log_params(model):
total_params = sum(p.numel() for p in model.parameters())
print(f"{model.__class__.__name__} has {total_params * 1.e-6:.2f} M params.")
logger.info(f"{model.__class__.__name__} has {total_params * 1.e-6:.2f} M params.")
def instantiate_from_config(config):

1
setup.py
Unescape Escape View File

@ -6,6 +6,7 @@ setup(
description='AI imagined images.',
packages=find_packages("imaginairy"),
install_requires=[
'click',
'torch',
'numpy',
'tqdm',

Write Preview
Loading…
Cancel
Save