Spaces:
Paused
Paused
File size: 8,014 Bytes
a121792 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 |
import argparse, os, sys, glob
import torch
import PIL
import numpy as np
from omegaconf import OmegaConf
from PIL import Image
from tqdm import tqdm, trange
from itertools import islice
from einops import rearrange, repeat
from torchvision.utils import make_grid
import time
from pytorch_lightning import seed_everything
from torch import autocast
from contextlib import contextmanager, nullcontext
from ldm.util import instantiate_from_config
from ldm.models.diffusion.ddim import DDIMSampler
from ldm.models.diffusion.plms import PLMSSampler
def chunk(it, size):
it = iter(it)
return iter(lambda: tuple(islice(it, size)), ())
def torch_gc():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
def load_model_from_config(config, ckpt, verbose=False):
print(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']}")
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)
if len(u) > 0 and verbose:
print("unexpected keys:")
print(u)
model.cuda()
model.half()
model.eval()
return model
def load_img(image, W, H):
w, h = image.size
print(f"loaded input image of size ({w}, {h})")
image = image.resize((int(W), int(H)), resample=PIL.Image.LANCZOS)
print(f"resize input image to size ({W}, {H})")
image = np.array(image).astype(np.float32) / 255.0
image = image[None].transpose(0, 3, 1, 2)
image = torch.from_numpy(image)
return 2.*image - 1.
class AppModel():
def __init__(self,):
self.config = OmegaConf.load("configs/stable-diffusion/v1-inference.yaml")
self.model = load_model_from_config(self.config, "models/ldm/stable-diffusion-v1/model.ckpt")
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
self.device = device
self.model = self.model.to(device)
self.sampler = PLMSSampler(self.model)
self.img_sampler = DDIMSampler(self.model)
self.C = 4 # latent channels
self.f = 8 # downsampling factors
def run_with_prompt(self, seed, prompt, n_samples, W, H, scale, ddim_steps, strength=0., init_img=None):
torch_gc()
seed_everything(seed)
ddim_eta=0.0
assert prompt is not None
print(f"Prompt: {prompt}")
batch_size = n_samples
data = [batch_size * [prompt]]
start_code = None
n_rows = int(n_samples**0.5)
precision_scope = autocast
if init_img is None:
with torch.no_grad():
with precision_scope(device_type='cuda', dtype=torch.float16):
with self.model.ema_scope():
all_samples = list()
for prompts in tqdm(data, desc="data"):
torch_gc()
uc = None
if scale != 1.0:
uc = self.model.get_learned_conditioning(batch_size * [""])
if isinstance(prompts, tuple):
prompts = list(prompts)
c = self.model.get_learned_conditioning(prompts)
shape = [self.C, H // self.f, W // self.f]
samples_ddim, _ = self.sampler.sample(S=ddim_steps,
conditioning=c,
batch_size=n_samples,
shape=shape,
verbose=False,
unconditional_guidance_scale=scale,
unconditional_conditioning=uc,
eta=ddim_eta,
x_T=start_code)
x_samples_ddim = self.model.decode_first_stage(samples_ddim)
x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
for x_sample in x_samples_ddim:
x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
image = Image.fromarray(x_sample.astype(np.uint8))
all_samples.append(image)
# additionally, grid image
grid = torch.stack([x_samples_ddim], 0)
grid = rearrange(grid, 'n b c h w -> (n b) c h w')
grid = make_grid(grid, nrow=n_rows)
# to image
grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
grid = grid.astype(np.uint8)
torch_gc()
return grid, all_samples
else:
init_image = load_img(init_img, W, H).to(self.device)
init_image = repeat(init_image, '1 ... -> b ...', b=batch_size)
torch_gc()
with precision_scope(device_type='cuda', dtype=torch.float16):
init_latent = self.model.get_first_stage_encoding(self.model.encode_first_stage(init_image)) # move to latent space
torch_gc()
sampler = self.img_sampler
sampler.make_schedule(ddim_num_steps=ddim_steps, ddim_eta=ddim_eta, verbose=False)
assert 0. <= strength < 1., 'can only work with strength in [0.0, 1.0)'
t_enc = int(strength * ddim_steps)
print(f"target t_enc is {t_enc} steps")
with torch.no_grad():
with precision_scope(device_type='cuda', dtype=torch.float16):
with self.model.ema_scope():
all_samples = list()
for prompts in tqdm(data, desc="data"):
uc = None
if scale != 1.0:
uc = self.model.get_learned_conditioning(batch_size * [""])
if isinstance(prompts, tuple):
prompts = list(prompts)
c = self.model.get_learned_conditioning(prompts)
# encode (scaled latent)
z_enc = sampler.stochastic_encode(init_latent, torch.tensor([t_enc]*batch_size).to(self.device))
# decode it
samples = sampler.decode(z_enc, c, t_enc, unconditional_guidance_scale=scale,
unconditional_conditioning=uc,)
x_samples = self.model.decode_first_stage(samples)
x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
for x_sample in x_samples:
x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
image = Image.fromarray(x_sample.astype(np.uint8))
all_samples.append(image)
# additionally, save as grid
grid = torch.stack([x_samples], 0)
grid = rearrange(grid, 'n b c h w -> (n b) c h w')
grid = make_grid(grid, nrow=n_rows)
# to image
grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
grid = grid.astype(np.uint8)
torch_gc()
return grid, all_samples
if __name__ == "__main__":
main()
|