File size: 7,174 Bytes
d661b19 |
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 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 |
import torch
import numpy as np
from einops import rearrange
from torch import autocast
from contextlib import nullcontext
from math import sqrt
from adapt import ScoreAdapter
from cldm.model import create_model, load_state_dict
from lora_util import *
import warnings
from transformers import logging
warnings.filterwarnings("ignore", category=DeprecationWarning)
logging.set_verbosity_error()
device = torch.device("cuda")
def _sqrt(x):
if isinstance(x, float):
return sqrt(x)
else:
assert isinstance(x, torch.Tensor)
return torch.sqrt(x)
def load_embedding(model,embedding):
length=len(embedding['string_to_param']['*'])
voc=[]
for i in range(length):
voc.append(f'<{str(i)}>')
print(f"Added Token: {voc}")
model.cond_stage_model.tokenizer._add_tokens(voc)
x=torch.nn.Embedding(model.cond_stage_model.tokenizer.__len__(),768)
for params in x.parameters():
params.requires_grad=False
x.weight[:-length]=model.cond_stage_model.transformer.text_model.embeddings.token_embedding.weight
x.weight[-length:]=embedding['string_to_param']['*']
model.cond_stage_model.transformer.text_model.embeddings.token_embedding=x
def load_3DFuse(control,dir,alpha):
######################LOADCONTROL###########################
model = create_model(control['control_yaml']).cpu()
model.load_state_dict(load_state_dict(control['control_weight'], location='cuda'))
state_dict, l = merge("runwayml/stable-diffusion-v1-5",dir,alpha)
#######################OVERRIDE#############################
model.load_state_dict(state_dict,strict=False)
#######################ADDEMBBEDDING########################
load_embedding(model,l)
###############################################################
return model
class StableDiffusion(ScoreAdapter):
def __init__(self, variant, v2_highres, prompt, scale, precision, dir, alpha=1.0):
model=load_3DFuse(self.checkpoint_root(),dir,alpha)
self.model = model.cuda()
H , W = (512, 512)
ae_resolution_f = 8
self._device = self.model._device
self.prompt = prompt
self.scale = scale
self.precision = precision
self.precision_scope = autocast if self.precision == "autocast" else nullcontext
self._data_shape = (4, H // ae_resolution_f, W // ae_resolution_f)
self.cond_func = self.model.get_learned_conditioning
self.M = 1000
noise_schedule = "linear"
self.noise_schedule = noise_schedule
self.us = self.linear_us(self.M)
def data_shape(self):
return self._data_shape
@property
def σ_max(self):
return self.us[0]
@property
def σ_min(self):
return self.us[-1]
@torch.no_grad()
def denoise(self, xs, σ,control, **model_kwargs):
with self.precision_scope("cuda"):
with self.model.ema_scope():
N = xs.shape[0]
c = model_kwargs.pop('c')
uc = model_kwargs.pop('uc')
conditional_conditioning = {"c_concat": [control], "c_crossattn": [c]}
unconditional_conditioning = {"c_concat": [control], "c_crossattn": [uc]}
cond_t, σ = self.time_cond_vec(N, σ)
unscaled_xs = xs
xs = xs / _sqrt(1 + σ**2)
if uc is None or self.scale == 1.:
output = self.model.apply_model(xs, cond_t, c)
else:
x_in = torch.cat([xs] * 2)
t_in = torch.cat([cond_t] * 2)
c_in = dict()
for k in conditional_conditioning:
if isinstance(conditional_conditioning[k], list):
c_in[k] = [torch.cat([
unconditional_conditioning[k][i],
conditional_conditioning[k][i]]) for i in range(len(conditional_conditioning[k]))]
else:
c_in[k] = torch.cat([
unconditional_conditioning[k],
conditional_conditioning[k]])
e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
output = e_t_uncond + self.scale * (e_t - e_t_uncond)
if self.model.parameterization == "v":
output = self.model.predict_eps_from_z_and_v(xs, cond_t, output)
else:
output = output
Ds = unscaled_xs - σ * output
return Ds
def cond_info(self, batch_size):
prompts = batch_size * [self.prompt]
return self.prompts_emb(prompts)
@torch.no_grad()
def prompts_emb(self, prompts):
assert isinstance(prompts, list)
batch_size = len(prompts)
with self.precision_scope("cuda"):
with self.model.ema_scope():
cond = {}
c = self.cond_func(prompts)
cond['c'] = c
uc = None
if self.scale != 1.0:
uc = self.cond_func(batch_size * [""])
cond['uc'] = uc
return cond
def unet_is_cond(self):
return True
def use_cls_guidance(self):
return False
def snap_t_to_nearest_tick(self, t):
j = np.abs(t - self.us).argmin()
return self.us[j], j
def time_cond_vec(self, N, σ):
if isinstance(σ, float):
σ, j = self.snap_t_to_nearest_tick(σ) # σ might change due to snapping
cond_t = (self.M - 1) - j
cond_t = torch.tensor([cond_t] * N, device=self.device)
return cond_t, σ
else:
assert isinstance(σ, torch.Tensor)
σ = σ.reshape(-1).cpu().numpy()
σs = []
js = []
for elem in σ:
_σ, _j = self.snap_t_to_nearest_tick(elem)
σs.append(_σ)
js.append((self.M - 1) - _j)
cond_t = torch.tensor(js, device=self.device)
σs = torch.tensor(σs, device=self.device, dtype=torch.float32).reshape(-1, 1, 1, 1)
return cond_t, σs
@staticmethod
def linear_us(M=1000):
assert M == 1000
β_start = 0.00085
β_end = 0.0120
βs = np.linspace(β_start**0.5, β_end**0.5, M, dtype=np.float64)**2
αs = np.cumprod(1 - βs)
us = np.sqrt((1 - αs) / αs)
us = us[::-1]
return us
@torch.no_grad()
def encode(self, xs):
model = self.model
with self.precision_scope("cuda"):
with self.model.ema_scope():
zs = model.get_first_stage_encoding(
model.encode_first_stage(xs)
)
return zs
@torch.no_grad()
def decode(self, xs):
with self.precision_scope("cuda"):
with self.model.ema_scope():
xs = self.model.decode_first_stage(xs)
return xs
|