Add files via upload
add
This commit is contained in:
smthemex
@@ -0,0 +1,163 @@
|
||||
CONFIG_2_0 = {
|
||||
"model_config": {
|
||||
"image_size": 64,
|
||||
"num_channels": 384,
|
||||
"num_res_blocks": 3,
|
||||
"channel_mult": "",
|
||||
"num_heads": 1,
|
||||
"num_head_channels": 64,
|
||||
"num_heads_upsample": -1,
|
||||
"attention_resolutions": "32,16,8",
|
||||
"dropout": 0,
|
||||
"model_dim": 768,
|
||||
"use_scale_shift_norm": True,
|
||||
"resblock_updown": True,
|
||||
"use_fp16": False,
|
||||
"cache_text_emb": True,
|
||||
"text_encoder_in_dim1": 1024,
|
||||
"text_encoder_in_dim2": 640,
|
||||
"pooling_type": "from_model",
|
||||
"in_channels": 4,
|
||||
"out_channels": 8,
|
||||
"up": False,
|
||||
"inpainting": False,
|
||||
},
|
||||
"diffusion_config": {
|
||||
"learn_sigma": True,
|
||||
"sigma_small": False,
|
||||
"steps": 1000,
|
||||
"noise_schedule": "linear",
|
||||
"timestep_respacing": "",
|
||||
"use_kl": False,
|
||||
"predict_xstart": False,
|
||||
"rescale_timesteps": True,
|
||||
"rescale_learned_sigmas": True,
|
||||
"linear_start": 0.0001,
|
||||
"linear_end": 0.02,
|
||||
},
|
||||
"image_enc_params": {
|
||||
"name": "AutoencoderKL",
|
||||
"scale": 0.0512,
|
||||
"params": {
|
||||
"ckpt_path": "",
|
||||
"embed_dim": 4,
|
||||
"ddconfig": {
|
||||
"double_z": True,
|
||||
"z_channels": 4,
|
||||
"resolution": 256,
|
||||
"in_channels": 3,
|
||||
"out_ch": 3,
|
||||
"ch": 128,
|
||||
"ch_mult": [1, 2, 4, 4],
|
||||
"num_res_blocks": 2,
|
||||
"attn_resolutions": [],
|
||||
"dropout": 0.0,
|
||||
},
|
||||
},
|
||||
},
|
||||
"text_enc_params1": {"model_path": "", "model_name": "multiclip"},
|
||||
"text_enc_params2": {"model_path": "", "model_name": "MT5EncoderModel"},
|
||||
"tokenizer_name1": "",
|
||||
"tokenizer_name2": "",
|
||||
}
|
||||
|
||||
CONFIG_2_1 = {
|
||||
"clip_name": "ViT-L/14",
|
||||
"clip_image_size": 224,
|
||||
"tokenizer_name": "",
|
||||
"image_enc_params": {
|
||||
"name": "MOVQ",
|
||||
"scale": 1,
|
||||
"ckpt_path": "",
|
||||
"params": {
|
||||
"embed_dim": 4,
|
||||
"n_embed": 16384,
|
||||
"ddconfig": {
|
||||
"double_z": False,
|
||||
"z_channels": 4,
|
||||
"resolution": 256,
|
||||
"in_channels": 3,
|
||||
"out_ch": 3,
|
||||
"ch": 128,
|
||||
"ch_mult": [1, 2, 2, 4],
|
||||
"num_res_blocks": 2,
|
||||
"attn_resolutions": [32],
|
||||
"dropout": 0.0,
|
||||
},
|
||||
},
|
||||
},
|
||||
"text_enc_params": {
|
||||
"model_path": "",
|
||||
"model_name": "multiclip",
|
||||
"in_features": 1024,
|
||||
"out_features": 768,
|
||||
},
|
||||
"prior": {
|
||||
"clip_mean_std_path": "ViT-L-14_stats.th",
|
||||
"params": {
|
||||
"model": {
|
||||
"type": "prior",
|
||||
"diffusion_sampler": "uniform",
|
||||
"hparams": {
|
||||
"text_ctx": 77,
|
||||
"xf_width": 2048,
|
||||
"xf_layers": 20,
|
||||
"xf_heads": 32,
|
||||
"xf_final_ln": True,
|
||||
"xf_padding": False,
|
||||
"text_drop": 0.2,
|
||||
"clip_dim": 768,
|
||||
"clip_xf_width": 768,
|
||||
},
|
||||
},
|
||||
"diffusion": {
|
||||
"steps": 1000,
|
||||
"learn_sigma": False,
|
||||
"sigma_small": True,
|
||||
"noise_schedule": "cosine",
|
||||
"use_kl": False,
|
||||
"predict_xstart": True,
|
||||
"rescale_learned_sigmas": False,
|
||||
"timestep_respacing": "",
|
||||
},
|
||||
},
|
||||
},
|
||||
"model_config": {
|
||||
"version": "2.1",
|
||||
"image_size": 64,
|
||||
"num_channels": 384,
|
||||
"num_res_blocks": 3,
|
||||
"channel_mult": "",
|
||||
"num_heads": 1,
|
||||
"num_head_channels": 64,
|
||||
"num_heads_upsample": -1,
|
||||
"attention_resolutions": "32,16,8",
|
||||
"dropout": 0,
|
||||
"model_dim": 768,
|
||||
"use_scale_shift_norm": True,
|
||||
"resblock_updown": True,
|
||||
"use_fp16": True,
|
||||
"cache_text_emb": True,
|
||||
"text_encoder_in_dim1": 1024,
|
||||
"text_encoder_in_dim2": 768,
|
||||
"image_encoder_in_dim": 768,
|
||||
"num_image_embs": 10,
|
||||
"pooling_type": "from_model",
|
||||
"in_channels": 4,
|
||||
"out_channels": 8,
|
||||
"use_flash_attention": False,
|
||||
},
|
||||
"diffusion_config": {
|
||||
"learn_sigma": True,
|
||||
"sigma_small": False,
|
||||
"steps": 1000,
|
||||
"noise_schedule": "linear",
|
||||
"timestep_respacing": "",
|
||||
"use_kl": False,
|
||||
"predict_xstart": False,
|
||||
"rescale_timesteps": True,
|
||||
"rescale_learned_sigmas": True,
|
||||
"linear_start": 0.00085,
|
||||
"linear_end": 0.012,
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,552 @@
|
||||
from transformers import AutoTokenizer
|
||||
from PIL import Image
|
||||
import cv2
|
||||
import torch
|
||||
from omegaconf import OmegaConf
|
||||
import math
|
||||
from copy import deepcopy
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import clip
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from .model.text_encoders import TextEncoder
|
||||
from .vqgan.autoencoder import VQModelInterface, AutoencoderKL, MOVQ
|
||||
from .model.samplers import DDIMSampler, PLMSSampler
|
||||
from .model.model_creation import create_model, create_gaussian_diffusion
|
||||
from .model.prior import PriorDiffusionModel, CustomizedTokenizer
|
||||
from .utils import prepare_image, q_sample, process_images, prepare_mask
|
||||
|
||||
|
||||
class Kandinsky2_1:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config,
|
||||
model_path,
|
||||
prior_path,
|
||||
device,
|
||||
task_type="text2img"
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.use_fp16 = self.config["model_config"]["use_fp16"]
|
||||
self.task_type = task_type
|
||||
self.clip_image_size = config["clip_image_size"]
|
||||
if task_type == "text2img":
|
||||
self.config["model_config"]["up"] = False
|
||||
self.config["model_config"]["inpainting"] = False
|
||||
elif task_type == "inpainting":
|
||||
self.config["model_config"]["up"] = False
|
||||
self.config["model_config"]["inpainting"] = True
|
||||
else:
|
||||
raise ValueError("Only text2img and inpainting is available")
|
||||
|
||||
self.tokenizer1 = AutoTokenizer.from_pretrained(self.config["tokenizer_name"])
|
||||
self.tokenizer2 = CustomizedTokenizer()
|
||||
clip_mean, clip_std = torch.load(
|
||||
config["prior"]["clip_mean_std_path"], map_location="cpu"
|
||||
)
|
||||
|
||||
self.prior = PriorDiffusionModel(
|
||||
config["prior"]["params"],
|
||||
self.tokenizer2,
|
||||
clip_mean,
|
||||
clip_std,
|
||||
)
|
||||
self.prior.load_state_dict(torch.load(prior_path), strict=False)
|
||||
if self.use_fp16:
|
||||
self.prior = self.prior.half()
|
||||
self.text_encoder = TextEncoder(**self.config["text_enc_params"])
|
||||
if self.use_fp16:
|
||||
self.text_encoder = self.text_encoder.half()
|
||||
|
||||
self.clip_model, self.preprocess = clip.load(
|
||||
config["clip_name"], device=self.device, jit=False
|
||||
)
|
||||
self.clip_model.eval()
|
||||
|
||||
if self.config["image_enc_params"] is not None:
|
||||
self.use_image_enc = True
|
||||
self.scale = self.config["image_enc_params"]["scale"]
|
||||
if self.config["image_enc_params"]["name"] == "AutoencoderKL":
|
||||
self.image_encoder = AutoencoderKL(
|
||||
**self.config["image_enc_params"]["params"]
|
||||
)
|
||||
elif self.config["image_enc_params"]["name"] == "VQModelInterface":
|
||||
self.image_encoder = VQModelInterface(
|
||||
**self.config["image_enc_params"]["params"]
|
||||
)
|
||||
elif self.config["image_enc_params"]["name"] == "MOVQ":
|
||||
self.image_encoder = MOVQ(**self.config["image_enc_params"]["params"])
|
||||
self.image_encoder.load_state_dict(
|
||||
torch.load(self.config["image_enc_params"]["ckpt_path"])
|
||||
)
|
||||
self.image_encoder.eval()
|
||||
else:
|
||||
self.use_image_enc = False
|
||||
|
||||
self.config["model_config"]["cache_text_emb"] = True
|
||||
self.model = create_model(**self.config["model_config"])
|
||||
self.model.load_state_dict(torch.load(model_path))
|
||||
if self.use_fp16:
|
||||
self.model.convert_to_fp16()
|
||||
self.image_encoder = self.image_encoder.half()
|
||||
|
||||
self.model_dtype = torch.float16
|
||||
else:
|
||||
self.model_dtype = torch.float32
|
||||
|
||||
self.image_encoder = self.image_encoder.to(self.device).eval()
|
||||
self.text_encoder = self.text_encoder.to(self.device).eval()
|
||||
self.prior = self.prior.to(self.device).eval()
|
||||
self.model.eval()
|
||||
self.model.to(self.device)
|
||||
|
||||
def get_new_h_w(self, h, w):
|
||||
new_h = h // 64
|
||||
if h % 64 != 0:
|
||||
new_h += 1
|
||||
new_w = w // 64
|
||||
if w % 64 != 0:
|
||||
new_w += 1
|
||||
return new_h * 8, new_w * 8
|
||||
|
||||
@torch.no_grad()
|
||||
def encode_text(self, text_encoder, tokenizer, prompt, batch_size):
|
||||
text_encoding = tokenizer(
|
||||
[prompt] * batch_size + [""] * batch_size,
|
||||
max_length=77,
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
return_attention_mask=True,
|
||||
add_special_tokens=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
tokens = text_encoding["input_ids"].to(self.device)
|
||||
mask = text_encoding["attention_mask"].to(self.device)
|
||||
|
||||
full_emb, pooled_emb = text_encoder(tokens=tokens, mask=mask)
|
||||
return full_emb, pooled_emb
|
||||
|
||||
@torch.no_grad()
|
||||
def generate_clip_emb(
|
||||
self,
|
||||
prompt,
|
||||
batch_size=1,
|
||||
prior_cf_scale=4,
|
||||
prior_steps="25",
|
||||
negative_prior_prompt="",
|
||||
seed=None,
|
||||
):
|
||||
prompts_batch = [prompt for _ in range(batch_size)]
|
||||
prior_cf_scales_batch = [prior_cf_scale] * len(prompts_batch)
|
||||
prior_cf_scales_batch = torch.tensor(prior_cf_scales_batch, device=self.device)
|
||||
max_txt_length = self.prior.model.text_ctx
|
||||
tok, mask = self.tokenizer2.padded_tokens_and_mask(
|
||||
prompts_batch, max_txt_length
|
||||
)
|
||||
cf_token, cf_mask = self.tokenizer2.padded_tokens_and_mask(
|
||||
[negative_prior_prompt], max_txt_length
|
||||
)
|
||||
if not (cf_token.shape == tok.shape):
|
||||
cf_token = cf_token.expand(tok.shape[0], -1)
|
||||
cf_mask = cf_mask.expand(tok.shape[0], -1)
|
||||
tok = torch.cat([tok, cf_token], dim=0)
|
||||
mask = torch.cat([mask, cf_mask], dim=0)
|
||||
tok, mask = tok.to(device=self.device), mask.to(device=self.device)
|
||||
|
||||
x = self.clip_model.token_embedding(tok).type(self.clip_model.dtype)
|
||||
x = x + self.clip_model.positional_embedding.type(self.clip_model.dtype)
|
||||
x = x.permute(1, 0, 2) # NLD -> LND|
|
||||
x = self.clip_model.transformer(x)
|
||||
x = x.permute(1, 0, 2) # LND -> NLD
|
||||
x = self.clip_model.ln_final(x).type(self.clip_model.dtype)
|
||||
txt_feat_seq = x
|
||||
txt_feat = (x[torch.arange(x.shape[0]), tok.argmax(dim=-1)] @ self.clip_model.text_projection)
|
||||
txt_feat, txt_feat_seq = txt_feat.float().to(self.device), txt_feat_seq.float().to(self.device)
|
||||
img_feat = self.prior(
|
||||
txt_feat,
|
||||
txt_feat_seq,
|
||||
mask,
|
||||
prior_cf_scales_batch,
|
||||
timestep_respacing=prior_steps,
|
||||
seed=seed,
|
||||
)
|
||||
return img_feat.to(self.model_dtype)
|
||||
|
||||
@torch.no_grad()
|
||||
def encode_images(self, image, is_pil=False):
|
||||
if is_pil:
|
||||
image = self.preprocess(image).unsqueeze(0).to(self.device)
|
||||
return self.clip_model.encode_image(image).to(self.model_dtype)
|
||||
|
||||
@torch.no_grad()
|
||||
def generate_img(
|
||||
self,
|
||||
prompt,
|
||||
img_prompt,
|
||||
batch_size=1,
|
||||
diffusion=None,
|
||||
guidance_scale=7,
|
||||
init_step=None,
|
||||
noise=None,
|
||||
init_img=None,
|
||||
img_mask=None,
|
||||
h=512,
|
||||
w=512,
|
||||
sampler="ddim_sampler",
|
||||
num_steps=50,
|
||||
seed=None
|
||||
):
|
||||
new_h, new_w = self.get_new_h_w(h, w)
|
||||
full_batch_size = batch_size * 2
|
||||
model_kwargs = {}
|
||||
generator = torch.Generator(device='cuda').manual_seed(seed) if seed is not None else None
|
||||
if init_img is not None and self.use_fp16:
|
||||
init_img = init_img.half()
|
||||
if img_mask is not None and self.use_fp16:
|
||||
img_mask = img_mask.half()
|
||||
model_kwargs["full_emb"], model_kwargs["pooled_emb"] = self.encode_text(
|
||||
text_encoder=self.text_encoder,
|
||||
tokenizer=self.tokenizer1,
|
||||
prompt=prompt,
|
||||
batch_size=batch_size,
|
||||
)
|
||||
model_kwargs["image_emb"] = img_prompt
|
||||
|
||||
if self.task_type == "inpainting":
|
||||
init_img = init_img.to(self.device)
|
||||
img_mask = img_mask.to(self.device)
|
||||
model_kwargs["inpaint_image"] = init_img * img_mask
|
||||
model_kwargs["inpaint_mask"] = img_mask
|
||||
|
||||
def model_fn(x_t, ts, **kwargs):
|
||||
half = x_t[: len(x_t) // 2]
|
||||
combined = torch.cat([half, half], dim=0)
|
||||
model_out = self.model(combined, ts, **kwargs)
|
||||
eps, rest = model_out[:, :4], model_out[:, 4:]
|
||||
cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
|
||||
half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps)
|
||||
eps = torch.cat([half_eps, half_eps], dim=0)
|
||||
if sampler == "p_sampler":
|
||||
return torch.cat([eps, rest], dim=1)
|
||||
else:
|
||||
return eps
|
||||
|
||||
if noise is not None:
|
||||
noise = noise.float()
|
||||
if self.task_type == "inpainting":
|
||||
def denoised_fun(x_start):
|
||||
x_start = x_start.clamp(-2, 2)
|
||||
return x_start * (1 - img_mask) + init_img * img_mask
|
||||
else:
|
||||
def denoised_fun(x):
|
||||
return x.clamp(-2, 2)
|
||||
|
||||
if sampler == "p_sampler":
|
||||
self.model.del_cache()
|
||||
samples = diffusion.p_sample_loop(
|
||||
model_fn,
|
||||
(full_batch_size, 4, new_h, new_w),
|
||||
device=self.device,
|
||||
noise=noise,
|
||||
progress=True,
|
||||
model_kwargs=model_kwargs,
|
||||
init_step=init_step,
|
||||
denoised_fn=denoised_fun,
|
||||
generator=generator
|
||||
)[:batch_size]
|
||||
self.model.del_cache()
|
||||
else:
|
||||
if sampler == "ddim_sampler":
|
||||
sampler = DDIMSampler(
|
||||
model=model_fn,
|
||||
old_diffusion=diffusion,
|
||||
schedule="linear",
|
||||
)
|
||||
elif sampler == "plms_sampler":
|
||||
sampler = PLMSSampler(
|
||||
model=model_fn,
|
||||
old_diffusion=diffusion,
|
||||
schedule="linear",
|
||||
)
|
||||
else:
|
||||
raise ValueError("Only ddim_sampler and plms_sampler is available")
|
||||
|
||||
self.model.del_cache()
|
||||
samples, _ = sampler.sample(
|
||||
num_steps,
|
||||
batch_size * 2,
|
||||
(4, new_h, new_w),
|
||||
conditioning=model_kwargs,
|
||||
x_T=noise,
|
||||
init_step=init_step,
|
||||
)
|
||||
self.model.del_cache()
|
||||
samples = samples[:batch_size]
|
||||
|
||||
if self.use_image_enc:
|
||||
if self.use_fp16:
|
||||
samples = samples.half()
|
||||
samples = self.image_encoder.decode(samples / self.scale)
|
||||
|
||||
samples = samples[:, :, :h, :w]
|
||||
return process_images(samples)
|
||||
|
||||
@torch.no_grad()
|
||||
def create_zero_img_emb(self, batch_size):
|
||||
img = torch.zeros(1, 3, self.clip_image_size, self.clip_image_size).to(self.device)
|
||||
return self.encode_images(img, is_pil=False).repeat(batch_size, 1)
|
||||
|
||||
@torch.no_grad()
|
||||
def generate_text2img(
|
||||
self,
|
||||
prompt,
|
||||
num_steps=100,
|
||||
batch_size=1,
|
||||
guidance_scale=7,
|
||||
h=512,
|
||||
w=512,
|
||||
sampler="ddim_sampler",
|
||||
prior_cf_scale=4,
|
||||
prior_steps="25",
|
||||
negative_prior_prompt="",
|
||||
negative_decoder_prompt="",
|
||||
):
|
||||
# generate clip embeddings
|
||||
image_emb = self.generate_clip_emb(
|
||||
prompt,
|
||||
batch_size=batch_size,
|
||||
prior_cf_scale=prior_cf_scale,
|
||||
prior_steps=prior_steps,
|
||||
negative_prior_prompt=negative_prior_prompt,
|
||||
)
|
||||
if negative_decoder_prompt == "":
|
||||
zero_image_emb = self.create_zero_img_emb(batch_size=batch_size)
|
||||
else:
|
||||
zero_image_emb = self.generate_clip_emb(
|
||||
negative_decoder_prompt,
|
||||
batch_size=batch_size,
|
||||
prior_cf_scale=prior_cf_scale,
|
||||
prior_steps=prior_steps,
|
||||
negative_prior_prompt=negative_prior_prompt,
|
||||
)
|
||||
|
||||
image_emb = torch.cat([image_emb, zero_image_emb], dim=0).to(self.device)
|
||||
|
||||
# load diffusion
|
||||
config = deepcopy(self.config)
|
||||
if sampler == "p_sampler":
|
||||
config["diffusion_config"]["timestep_respacing"] = str(num_steps)
|
||||
diffusion = create_gaussian_diffusion(**config["diffusion_config"])
|
||||
|
||||
return self.generate_img(
|
||||
prompt=prompt,
|
||||
img_prompt=image_emb,
|
||||
batch_size=batch_size,
|
||||
guidance_scale=guidance_scale,
|
||||
h=h,
|
||||
w=w,
|
||||
sampler=sampler,
|
||||
num_steps=num_steps,
|
||||
diffusion=diffusion,
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def mix_images(
|
||||
self,
|
||||
images_texts,
|
||||
weights,
|
||||
num_steps=100,
|
||||
batch_size=1,
|
||||
guidance_scale=7,
|
||||
h=512,
|
||||
w=512,
|
||||
sampler="ddim_sampler",
|
||||
prior_cf_scale=4,
|
||||
prior_steps="25",
|
||||
negative_prior_prompt="",
|
||||
negative_decoder_prompt="",
|
||||
):
|
||||
assert len(images_texts) == len(weights) and len(images_texts) > 0
|
||||
|
||||
# generate clip embeddings
|
||||
image_emb = None
|
||||
for i in range(len(images_texts)):
|
||||
if image_emb is None:
|
||||
if type(images_texts[i]) == str:
|
||||
image_emb = weights[i] * self.generate_clip_emb(
|
||||
images_texts[i],
|
||||
batch_size=1,
|
||||
prior_cf_scale=prior_cf_scale,
|
||||
prior_steps=prior_steps,
|
||||
negative_prior_prompt=negative_prior_prompt,
|
||||
)
|
||||
else:
|
||||
image_emb = self.encode_images(images_texts[i], is_pil=True) * weights[i]
|
||||
else:
|
||||
if type(images_texts[i]) == str:
|
||||
image_emb = image_emb + weights[i] * self.generate_clip_emb(
|
||||
images_texts[i],
|
||||
batch_size=1,
|
||||
prior_cf_scale=prior_cf_scale,
|
||||
prior_steps=prior_steps,
|
||||
negative_prior_prompt=negative_prior_prompt,
|
||||
)
|
||||
else:
|
||||
image_emb = image_emb + self.encode_images(images_texts[i], is_pil=True) * weights[i]
|
||||
|
||||
image_emb = image_emb.repeat(batch_size, 1)
|
||||
if negative_decoder_prompt == "":
|
||||
zero_image_emb = self.create_zero_img_emb(batch_size=batch_size)
|
||||
else:
|
||||
zero_image_emb = self.generate_clip_emb(
|
||||
negative_decoder_prompt,
|
||||
batch_size=batch_size,
|
||||
prior_cf_scale=prior_cf_scale,
|
||||
prior_steps=prior_steps,
|
||||
negative_prior_prompt=negative_prior_prompt,
|
||||
)
|
||||
image_emb = torch.cat([image_emb, zero_image_emb], dim=0).to(self.device)
|
||||
|
||||
# load diffusion
|
||||
config = deepcopy(self.config)
|
||||
if sampler == "p_sampler":
|
||||
config["diffusion_config"]["timestep_respacing"] = str(num_steps)
|
||||
diffusion = create_gaussian_diffusion(**config["diffusion_config"])
|
||||
return self.generate_img(
|
||||
prompt="",
|
||||
img_prompt=image_emb,
|
||||
batch_size=batch_size,
|
||||
guidance_scale=guidance_scale,
|
||||
h=h,
|
||||
w=w,
|
||||
sampler=sampler,
|
||||
num_steps=num_steps,
|
||||
diffusion=diffusion,
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def generate_img2img(
|
||||
self,
|
||||
prompt,
|
||||
pil_img,
|
||||
strength=0.7,
|
||||
num_steps=100,
|
||||
batch_size=1,
|
||||
guidance_scale=7,
|
||||
h=512,
|
||||
w=512,
|
||||
sampler="ddim_sampler",
|
||||
prior_cf_scale=4,
|
||||
prior_steps="25",
|
||||
):
|
||||
# generate clip embeddings
|
||||
image_emb = self.generate_clip_emb(
|
||||
prompt,
|
||||
batch_size=batch_size,
|
||||
prior_cf_scale=prior_cf_scale,
|
||||
prior_steps=prior_steps,
|
||||
)
|
||||
zero_image_emb = self.create_zero_img_emb(batch_size=batch_size)
|
||||
image_emb = torch.cat([image_emb, zero_image_emb], dim=0).to(self.device)
|
||||
|
||||
# load diffusion
|
||||
config = deepcopy(self.config)
|
||||
if sampler == "p_sampler":
|
||||
config["diffusion_config"]["timestep_respacing"] = str(num_steps)
|
||||
diffusion = create_gaussian_diffusion(**config["diffusion_config"])
|
||||
|
||||
image = prepare_image(pil_img, h=h, w=w).to(self.device)
|
||||
if self.use_fp16:
|
||||
image = image.half()
|
||||
image = self.image_encoder.encode(image) * self.scale
|
||||
|
||||
start_step = int(diffusion.num_timesteps * (1 - strength))
|
||||
image = q_sample(
|
||||
image,
|
||||
torch.tensor(diffusion.timestep_map[start_step - 1]).to(self.device),
|
||||
schedule_name=config["diffusion_config"]["noise_schedule"],
|
||||
num_steps=config["diffusion_config"]["steps"],
|
||||
)
|
||||
|
||||
image = image.repeat(2, 1, 1, 1)
|
||||
return self.generate_img(
|
||||
prompt=prompt,
|
||||
img_prompt=image_emb,
|
||||
batch_size=batch_size,
|
||||
guidance_scale=guidance_scale,
|
||||
h=h,
|
||||
w=w,
|
||||
sampler=sampler,
|
||||
num_steps=num_steps,
|
||||
diffusion=diffusion,
|
||||
noise=image,
|
||||
init_step=start_step,
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def generate_inpainting(
|
||||
self,
|
||||
prompt,
|
||||
pil_img,
|
||||
img_mask,
|
||||
num_steps=100,
|
||||
batch_size=1,
|
||||
guidance_scale=7,
|
||||
h=512,
|
||||
w=512,
|
||||
sampler="ddim_sampler",
|
||||
prior_cf_scale=4,
|
||||
prior_steps="25",
|
||||
negative_prior_prompt="",
|
||||
negative_decoder_prompt="",
|
||||
):
|
||||
# generate clip embeddings
|
||||
image_emb = self.generate_clip_emb(
|
||||
prompt,
|
||||
batch_size=batch_size,
|
||||
prior_cf_scale=prior_cf_scale,
|
||||
prior_steps=prior_steps,
|
||||
negative_prior_prompt=negative_prior_prompt,
|
||||
)
|
||||
zero_image_emb = self.create_zero_img_emb(batch_size=batch_size)
|
||||
image_emb = torch.cat([image_emb, zero_image_emb], dim=0).to(self.device)
|
||||
|
||||
# load diffusion
|
||||
config = deepcopy(self.config)
|
||||
if sampler == "p_sampler":
|
||||
config["diffusion_config"]["timestep_respacing"] = str(num_steps)
|
||||
diffusion = create_gaussian_diffusion(**config["diffusion_config"])
|
||||
image = prepare_image(pil_img, w, h).to(self.device)
|
||||
if self.use_fp16:
|
||||
image = image.half()
|
||||
image = self.image_encoder.encode(image) * self.scale
|
||||
image_shape = tuple(image.shape[-2:])
|
||||
img_mask = torch.from_numpy(img_mask).unsqueeze(0).unsqueeze(0)
|
||||
img_mask = F.interpolate(
|
||||
img_mask,
|
||||
image_shape,
|
||||
mode="nearest",
|
||||
)
|
||||
img_mask = prepare_mask(img_mask).to(self.device)
|
||||
if self.use_fp16:
|
||||
img_mask = img_mask.half()
|
||||
image = image.repeat(2, 1, 1, 1)
|
||||
img_mask = img_mask.repeat(2, 1, 1, 1)
|
||||
|
||||
return self.generate_img(
|
||||
prompt=prompt,
|
||||
img_prompt=image_emb,
|
||||
batch_size=batch_size,
|
||||
guidance_scale=guidance_scale,
|
||||
h=h,
|
||||
w=w,
|
||||
sampler=sampler,
|
||||
num_steps=num_steps,
|
||||
diffusion=diffusion,
|
||||
init_img=image,
|
||||
img_mask=img_mask,
|
||||
)
|
||||
@@ -0,0 +1,189 @@
|
||||
from PIL import Image
|
||||
import cv2
|
||||
import torch
|
||||
import math
|
||||
from copy import deepcopy
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline, KandinskyV22Img2ImgPipeline, \
|
||||
KandinskyV22InpaintPipeline
|
||||
from transformers import CLIPVisionModelWithProjection
|
||||
from diffusers.models import UNet2DConditionModel
|
||||
from transformers import CLIPVisionModelWithProjection
|
||||
from diffusers.models import UNet2DConditionModel
|
||||
from typing import Optional, Any
|
||||
|
||||
|
||||
class Kandinsky2_2:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
device,
|
||||
task_type="text2img",
|
||||
):
|
||||
self.device = device
|
||||
self.task_type = task_type
|
||||
self.image_encoder = CLIPVisionModelWithProjection.from_pretrained('kandinsky-community/kandinsky-2-2-prior',
|
||||
subfolder='image_encoder').to(
|
||||
torch.float16).to(self.device)
|
||||
if task_type == "text2img":
|
||||
self.unet = UNet2DConditionModel.from_pretrained('kandinsky-community/kandinsky-2-2-decoder',
|
||||
subfolder='unet').to(torch.float16).to(self.device)
|
||||
self.prior = KandinskyV22PriorPipeline.from_pretrained('kandinsky-community/kandinsky-2-2-prior',
|
||||
image_encoder=self.image_encoder,
|
||||
torch_dtype=torch.float16)
|
||||
self.prior = self.prior.to(self.device)
|
||||
self.decoder = KandinskyV22Pipeline.from_pretrained('kandinsky-community/kandinsky-2-2-decoder',
|
||||
unet=self.unet, torch_dtype=torch.float16)
|
||||
self.decoder = self.decoder.to(self.device)
|
||||
elif task_type == "inpainting":
|
||||
self.unet = UNet2DConditionModel.from_pretrained('kandinsky-community/kandinsky-2-2-decoder-inpaint',
|
||||
subfolder='unet').to(torch.float16).to(self.device)
|
||||
self.prior = KandinskyV22PriorPipeline.from_pretrained('kandinsky-community/kandinsky-2-2-prior',
|
||||
image_encoder=self.image_encoder,
|
||||
torch_dtype=torch.float16)
|
||||
self.prior = self.prior.to(self.device)
|
||||
self.decoder = KandinskyV22InpaintPipeline.from_pretrained(
|
||||
'kandinsky-community/kandinsky-2-2-decoder-inpaint', unet=self.unet, torch_dtype=torch.float16)
|
||||
self.decoder = self.decoder.to(self.device)
|
||||
elif task_type == "img2img":
|
||||
self.unet = UNet2DConditionModel.from_pretrained('kandinsky-community/kandinsky-2-2-decoder',
|
||||
subfolder='unet').to(torch.float16).to(self.device)
|
||||
self.prior = KandinskyV22PriorPipeline.from_pretrained('kandinsky-community/kandinsky-2-2-prior',
|
||||
image_encoder=self.image_encoder,
|
||||
torch_dtype=torch.float16)
|
||||
self.prior = self.prior.to(self.device)
|
||||
self.decoder = KandinskyV22Img2ImgPipeline.from_pretrained('kandinsky-community/kandinsky-2-2-decoder',
|
||||
unet=self.unet, torch_dtype=torch.float16)
|
||||
self.decoder = self.decoder.to(self.device)
|
||||
else:
|
||||
raise ValueError("Only text2img, img2img, inpainting is available")
|
||||
|
||||
def get_new_h_w(self, h, w):
|
||||
new_h = h // 64
|
||||
if h % 64 != 0:
|
||||
new_h += 1
|
||||
new_w = w // 64
|
||||
if w % 64 != 0:
|
||||
new_w += 1
|
||||
return new_h * 64, new_w * 64
|
||||
|
||||
def generate_text2img(
|
||||
self,
|
||||
prompt,
|
||||
batch_size=1,
|
||||
decoder_steps=50,
|
||||
prior_steps=25,
|
||||
decoder_guidance_scale=4,
|
||||
prior_guidance_scale=4,
|
||||
h=512,
|
||||
w=512,
|
||||
negative_prior_prompt="",
|
||||
negative_decoder_prompt="",
|
||||
):
|
||||
h, w = self.get_new_h_w(h, w)
|
||||
img_emb = self.prior(prompt=prompt, num_inference_steps=prior_steps,
|
||||
num_images_per_prompt=batch_size, guidance_scale=prior_guidance_scale,
|
||||
negative_prompt=negative_prior_prompt)
|
||||
negative_emb = self.prior(prompt=negative_decoder_prompt, num_inference_steps=prior_steps,
|
||||
num_images_per_prompt=batch_size, guidance_scale=prior_guidance_scale)
|
||||
if negative_decoder_prompt == "":
|
||||
negative_emb = negative_emb.negative_image_embeds
|
||||
else:
|
||||
negative_emb = negative_emb.image_embeds
|
||||
images = self.decoder(image_embeds=img_emb.image_embeds, negative_image_embeds=negative_emb,
|
||||
num_inference_steps=decoder_steps, height=h,
|
||||
width=w, guidance_scale=decoder_guidance_scale).images
|
||||
return images
|
||||
|
||||
def generate_img2img(
|
||||
self,
|
||||
prompt,
|
||||
image,
|
||||
strength=0.4,
|
||||
batch_size=1,
|
||||
decoder_steps=100,
|
||||
prior_steps=25,
|
||||
decoder_guidance_scale=4,
|
||||
prior_guidance_scale=4,
|
||||
h=512,
|
||||
w=512,
|
||||
negative_prior_prompt="",
|
||||
negative_decoder_prompt="",
|
||||
):
|
||||
h, w = self.get_new_h_w(h, w)
|
||||
img_emb = self.prior(prompt=prompt, num_inference_steps=prior_steps,
|
||||
num_images_per_prompt=batch_size, guidance_scale=prior_guidance_scale,
|
||||
negative_prompt=negative_prior_prompt)
|
||||
negative_emb = self.prior(prompt=negative_prior_prompt, num_inference_steps=prior_steps,
|
||||
num_images_per_prompt=batch_size, guidance_scale=prior_guidance_scale)
|
||||
if negative_decoder_prompt == "":
|
||||
negative_emb = negative_emb.negative_image_embeds
|
||||
else:
|
||||
negative_emb = negative_emb.image_embeds
|
||||
images = self.decoder(image_embeds=img_emb.image_embeds, negative_image_embeds=negative_emb,
|
||||
num_inference_steps=decoder_steps, height=h,
|
||||
width=w, guidance_scale=decoder_guidance_scale,
|
||||
strength=strength, image=image).images
|
||||
return images
|
||||
|
||||
def mix_images(
|
||||
self,
|
||||
images_texts,
|
||||
weights,
|
||||
batch_size=1,
|
||||
decoder_steps=50,
|
||||
prior_steps=25,
|
||||
decoder_guidance_scale=4,
|
||||
prior_guidance_scale=4,
|
||||
h=512,
|
||||
w=512,
|
||||
negative_prior_prompt="",
|
||||
negative_decoder_prompt="",
|
||||
):
|
||||
assert len(images_texts) == len(weights) and len(images_texts) > 0
|
||||
|
||||
img_emb = self.prior.interpolate(images_and_prompts=images_texts, weights=weights,
|
||||
num_inference_steps=prior_steps, num_images_per_prompt=batch_size,
|
||||
guidance_scale=prior_guidance_scale, negative_prompt=negative_prior_prompt)
|
||||
negative_emb = self.prior(prompt=negative_prior_prompt, num_inference_steps=prior_steps,
|
||||
num_images_per_prompt=batch_size, guidance_scale=prior_guidance_scale)
|
||||
if negative_decoder_prompt == "":
|
||||
negative_emb = negative_emb.negative_image_embeds
|
||||
else:
|
||||
negative_emb = negative_emb.image_embeds
|
||||
images = self.decoder(image_embeds=img_emb.image_embeds, negative_image_embeds=negative_emb,
|
||||
num_inference_steps=decoder_steps, height=h,
|
||||
width=w, guidance_scale=decoder_guidance_scale).images
|
||||
return images
|
||||
|
||||
def generate_inpainting(
|
||||
self,
|
||||
prompt,
|
||||
pil_img,
|
||||
img_mask,
|
||||
batch_size=1,
|
||||
decoder_steps=50,
|
||||
prior_steps=25,
|
||||
decoder_guidance_scale=4,
|
||||
prior_guidance_scale=4,
|
||||
h=512,
|
||||
w=512,
|
||||
negative_prior_prompt="",
|
||||
negative_decoder_prompt="",
|
||||
):
|
||||
|
||||
img_emb = self.prior(prompt=prompt, num_inference_steps=prior_steps,
|
||||
num_images_per_prompt=batch_size, guidance_scale=prior_guidance_scale,
|
||||
negative_prompt=negative_prior_prompt)
|
||||
negative_emb = self.prior(prompt=negative_prior_prompt, num_inference_steps=prior_steps,
|
||||
num_images_per_prompt=batch_size, guidance_scale=prior_guidance_scale)
|
||||
if negative_decoder_prompt == "":
|
||||
negative_emb = negative_emb.negative_image_embeds
|
||||
else:
|
||||
negative_emb = negative_emb.image_embeds
|
||||
images = self.decoder(image_embeds=img_emb.image_embeds, negative_image_embeds=negative_emb,
|
||||
num_inference_steps=decoder_steps, height=h,
|
||||
width=w, guidance_scale=decoder_guidance_scale,
|
||||
image=pil_img, mask_image=img_mask).images
|
||||
return images
|
||||
@@ -0,0 +1,381 @@
|
||||
from PIL import Image
|
||||
import cv2
|
||||
import torch
|
||||
from omegaconf import OmegaConf
|
||||
import math
|
||||
from copy import deepcopy
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from .model.samplers import DDIMSampler, PLMSSampler
|
||||
from .model.model_creation import create_model, create_gaussian_diffusion
|
||||
from .model.text_encoders import TextEncoder
|
||||
from .vqgan.autoencoder import VQModelInterface, AutoencoderKL
|
||||
from .utils import prepare_image, q_sample, process_images, prepare_mask
|
||||
|
||||
|
||||
class Kandinsky2:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config,
|
||||
model_path,
|
||||
device,
|
||||
task_type="text2img"
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.task_type = task_type
|
||||
if task_type == "text2img" or task_type == "img2img":
|
||||
self.config["model_config"]["up"] = False
|
||||
self.config["model_config"]["inpainting"] = False
|
||||
elif task_type == "inpainting":
|
||||
self.config["model_config"]["up"] = False
|
||||
self.config["model_config"]["inpainting"] = True
|
||||
else:
|
||||
raise ValueError("Only text2img, img2img and inpainting is available")
|
||||
|
||||
self.tokenizer1 = AutoTokenizer.from_pretrained(self.config["tokenizer_name1"])
|
||||
self.tokenizer2 = AutoTokenizer.from_pretrained(self.config["tokenizer_name2"])
|
||||
|
||||
self.text_encoder1 = (
|
||||
TextEncoder(**self.config["text_enc_params1"]).to(self.device).eval()
|
||||
)
|
||||
self.text_encoder2 = (
|
||||
TextEncoder(**self.config["text_enc_params2"]).to(self.device).eval()
|
||||
)
|
||||
|
||||
self.use_fp16 = self.config["model_config"]["use_fp16"]
|
||||
|
||||
if self.config["image_enc_params"] is not None:
|
||||
self.use_image_enc = True
|
||||
self.scale = self.config["image_enc_params"]["scale"]
|
||||
if self.config["image_enc_params"]["name"] == "AutoencoderKL":
|
||||
self.image_encoder = AutoencoderKL(
|
||||
**self.config["image_enc_params"]["params"]
|
||||
).to(self.device)
|
||||
elif self.config["image_enc_params"]["name"] == "VQModelInterface":
|
||||
self.image_encoder = VQModelInterface(
|
||||
**self.config["image_enc_params"]["params"]
|
||||
).to(self.device)
|
||||
self.image_encoder.eval()
|
||||
else:
|
||||
self.use_image_enc = False
|
||||
|
||||
self.config["model_config"]["cache_text_emb"] = True
|
||||
self.model = create_model(**self.config["model_config"])
|
||||
self.model.load_state_dict(torch.load(model_path), strict=False)
|
||||
if self.use_fp16:
|
||||
self.model.convert_to_fp16()
|
||||
self.text_encoder1 = self.text_encoder1.half()
|
||||
self.text_encoder2 = self.text_encoder2.half()
|
||||
self.image_encoder = self.image_encoder.half()
|
||||
self.model.eval()
|
||||
self.model.to(self.device)
|
||||
|
||||
def get_new_h_w(self, h, w):
|
||||
new_h = h // 64
|
||||
if h % 64 != 0:
|
||||
new_h += 1
|
||||
new_w = w // 64
|
||||
if w % 64 != 0:
|
||||
new_w += 1
|
||||
return new_h * 8, new_w * 8
|
||||
|
||||
@torch.no_grad()
|
||||
def encode_text(self, text_encoder, tokenizer, prompt, batch_size):
|
||||
text_encoding = tokenizer(
|
||||
[prompt] * batch_size + [""] * batch_size,
|
||||
max_length=77,
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
return_attention_mask=True,
|
||||
add_special_tokens=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
tokens = text_encoding["input_ids"].to(self.device)
|
||||
mask = text_encoding["attention_mask"].to(self.device)
|
||||
|
||||
full_emb, pooled_emb = text_encoder(tokens=tokens, mask=mask)
|
||||
return full_emb, pooled_emb
|
||||
|
||||
@torch.no_grad()
|
||||
def generate_img(
|
||||
self,
|
||||
prompt,
|
||||
batch_size=1,
|
||||
diffusion=None,
|
||||
num_steps=50,
|
||||
guidance_scale=7,
|
||||
progress=True,
|
||||
dynamic_threshold_v=99.5,
|
||||
denoised_type="dynamic_threshold",
|
||||
init_step=None,
|
||||
noise=None,
|
||||
init_img=None,
|
||||
img_mask=None,
|
||||
h=512,
|
||||
w=512,
|
||||
sampler="ddim_sampler",
|
||||
ddim_eta=0.8,
|
||||
):
|
||||
new_h, new_w = self.get_new_h_w(h, w)
|
||||
full_batch_size = batch_size * 2
|
||||
model_kwargs = {}
|
||||
if noise is not None and self.use_fp16:
|
||||
noise = noise.half()
|
||||
if init_img is not None and self.use_fp16:
|
||||
init_img = init_img.half()
|
||||
if img_mask is not None and self.use_fp16:
|
||||
img_mask = img_mask.half()
|
||||
model_kwargs["full_emb1"], model_kwargs["pooled_emb1"] = self.encode_text(
|
||||
text_encoder=self.text_encoder1,
|
||||
tokenizer=self.tokenizer1,
|
||||
prompt=prompt,
|
||||
batch_size=batch_size,
|
||||
)
|
||||
model_kwargs["full_emb2"], model_kwargs["pooled_emb2"] = self.encode_text(
|
||||
text_encoder=self.text_encoder2,
|
||||
tokenizer=self.tokenizer2,
|
||||
prompt=prompt,
|
||||
batch_size=batch_size,
|
||||
)
|
||||
if self.task_type == "inpainting":
|
||||
init_img = init_img.to(self.device)
|
||||
img_mask = img_mask.to(self.device)
|
||||
model_kwargs["inpaint_image"] = init_img * img_mask
|
||||
model_kwargs["inpaint_mask"] = img_mask
|
||||
|
||||
def model_fn(x_t, ts, **kwargs):
|
||||
half = x_t[: len(x_t) // 2]
|
||||
combined = torch.cat([half, half], dim=0)
|
||||
model_out = self.model(combined, ts, **kwargs)
|
||||
eps, rest = model_out[:, :4], model_out[:, 4:]
|
||||
cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
|
||||
half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps)
|
||||
eps = torch.cat([half_eps, half_eps], dim=0)
|
||||
if sampler == "p_sampler":
|
||||
return torch.cat([eps, rest], dim=1)
|
||||
else:
|
||||
return eps
|
||||
|
||||
if self.task_type == "inpainting":
|
||||
def denoised_fn(x_start):
|
||||
if denoised_type == "dynamic_threshold":
|
||||
x2 = torch.clone(x_start).cpu().detach().numpy()
|
||||
p = dynamic_threshold_v
|
||||
s = np.percentile(np.abs(x2), p, axis=tuple(range(1, x2.ndim)))[0]
|
||||
s = max(s, 1.0)
|
||||
x_start = torch.clip(x_start, -s, s) / s
|
||||
elif denoised_type == "clip_denoised":
|
||||
x_start = x_start.clamp(-1, 1)
|
||||
return x_start * (1 - img_mask) + init_img * img_mask
|
||||
|
||||
denoised_function = denoised_fn
|
||||
else:
|
||||
def denoised_fn(x):
|
||||
if denoised_type == "dynamic_threshold":
|
||||
x2 = torch.clone(x).cpu().detach().numpy()
|
||||
p = dynamic_threshold_v
|
||||
s = np.percentile(np.abs(x2), p, axis=tuple(range(1, x2.ndim)))[0]
|
||||
s = max(s, 1.0)
|
||||
x = torch.clip(x, -s, s) / s
|
||||
return x
|
||||
elif denoised_type == "clip_denoised":
|
||||
return x.clamp(-1, 1)
|
||||
return x
|
||||
|
||||
denoised_function = None
|
||||
|
||||
if sampler == "p_sampler":
|
||||
self.model.del_cache()
|
||||
samples = diffusion.p_sample_loop(
|
||||
model_fn,
|
||||
(full_batch_size, 4, new_h, new_w),
|
||||
device=self.device,
|
||||
noise=noise,
|
||||
progress=progress,
|
||||
model_kwargs=model_kwargs,
|
||||
init_step=init_step,
|
||||
denoised_fn=denoised_function,
|
||||
)[:batch_size]
|
||||
self.model.del_cache()
|
||||
else:
|
||||
if sampler == "ddim_sampler":
|
||||
sampler = DDIMSampler(
|
||||
model=model_fn,
|
||||
old_diffusion=diffusion,
|
||||
schedule="linear",
|
||||
)
|
||||
elif sampler == "plms_sampler":
|
||||
sampler = PLMSSampler(
|
||||
model=model_fn,
|
||||
old_diffusion=diffusion,
|
||||
schedule="linear",
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"Only p_sampler, ddim_sampler and plms_sampler is available"
|
||||
)
|
||||
self.model.del_cache()
|
||||
|
||||
if sampler == "ddim_sampler":
|
||||
samples, _ = sampler.sample(
|
||||
num_steps,
|
||||
batch_size * 2,
|
||||
(4, new_h, new_w),
|
||||
conditioning=model_kwargs,
|
||||
x_T=noise,
|
||||
init_step=init_step,
|
||||
eta=ddim_eta,
|
||||
)
|
||||
else:
|
||||
samples, _ = sampler.sample(
|
||||
num_steps,
|
||||
batch_size * 2,
|
||||
(4, new_h, new_w),
|
||||
conditioning=model_kwargs,
|
||||
x_T=noise,
|
||||
init_step=init_step,
|
||||
)
|
||||
self.model.del_cache()
|
||||
samples = samples[:batch_size]
|
||||
if self.use_image_enc:
|
||||
if self.use_fp16:
|
||||
samples = samples.half()
|
||||
samples = self.image_encoder.decode(samples / self.scale)
|
||||
samples = samples[:, :, :h, :w]
|
||||
return process_images(samples)
|
||||
|
||||
@torch.no_grad()
|
||||
def generate_text2img(
|
||||
self,
|
||||
prompt,
|
||||
num_steps=100,
|
||||
batch_size=1,
|
||||
guidance_scale=7,
|
||||
progress=True,
|
||||
dynamic_threshold_v=99.5,
|
||||
denoised_type="dynamic_threshold",
|
||||
h=512,
|
||||
w=512,
|
||||
sampler="ddim_sampler",
|
||||
ddim_eta=0.05,
|
||||
):
|
||||
config = deepcopy(self.config)
|
||||
if sampler == "p_sampler":
|
||||
config["diffusion_config"]["timestep_respacing"] = str(num_steps)
|
||||
diffusion = create_gaussian_diffusion(**config["diffusion_config"])
|
||||
|
||||
return self.generate_img(
|
||||
prompt=prompt,
|
||||
batch_size=batch_size,
|
||||
diffusion=diffusion,
|
||||
num_steps=num_steps,
|
||||
guidance_scale=guidance_scale,
|
||||
progress=progress,
|
||||
dynamic_threshold_v=dynamic_threshold_v,
|
||||
denoised_type=denoised_type,
|
||||
h=h,
|
||||
w=w,
|
||||
sampler=sampler,
|
||||
ddim_eta=ddim_eta,
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def generate_img2img(
|
||||
self,
|
||||
prompt,
|
||||
pil_img,
|
||||
strength=0.7,
|
||||
num_steps=100,
|
||||
guidance_scale=7,
|
||||
progress=True,
|
||||
dynamic_threshold_v=99.5,
|
||||
denoised_type="dynamic_threshold",
|
||||
sampler="ddim_sampler",
|
||||
ddim_eta=0.05,
|
||||
):
|
||||
config = deepcopy(self.config)
|
||||
if sampler == "p_sampler":
|
||||
config["diffusion_config"]["timestep_respacing"] = str(num_steps)
|
||||
diffusion = create_gaussian_diffusion(**config["diffusion_config"])
|
||||
|
||||
image = prepare_image(pil_img).to(self.device)
|
||||
if self.use_fp16:
|
||||
image = image.half()
|
||||
image = self.image_encoder.encode(image).sample() * self.scale
|
||||
start_step = int(diffusion.num_timesteps * (1 - strength))
|
||||
image = q_sample(
|
||||
image,
|
||||
torch.tensor(diffusion.timestep_map[start_step - 1]).to(self.device),
|
||||
schedule_name=config["diffusion_config"]["noise_schedule"],
|
||||
num_steps=config["diffusion_config"]["steps"],
|
||||
)
|
||||
image = image.repeat(2, 1, 1, 1)
|
||||
return self.generate_img(
|
||||
prompt=prompt,
|
||||
batch_size=1,
|
||||
diffusion=diffusion,
|
||||
noise=image,
|
||||
num_steps=num_steps,
|
||||
guidance_scale=guidance_scale,
|
||||
progress=progress,
|
||||
dynamic_threshold_v=dynamic_threshold_v,
|
||||
denoised_type=denoised_type,
|
||||
init_step=start_step,
|
||||
sampler=sampler,
|
||||
ddim_eta=ddim_eta,
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def generate_inpainting(
|
||||
self,
|
||||
prompt,
|
||||
pil_img,
|
||||
img_mask,
|
||||
num_steps=100,
|
||||
guidance_scale=7,
|
||||
progress=True,
|
||||
dynamic_threshold_v=99.5,
|
||||
denoised_type="dynamic_threshold",
|
||||
sampler="ddim_sampler",
|
||||
ddim_eta=0.05,
|
||||
):
|
||||
config = deepcopy(self.config)
|
||||
if sampler == "p_sampler":
|
||||
config["diffusion_config"]["timestep_respacing"] = str(num_steps)
|
||||
diffusion = create_gaussian_diffusion(**config["diffusion_config"])
|
||||
|
||||
image = prepare_image(pil_img).to(self.device)
|
||||
if self.use_fp16:
|
||||
image = image.half()
|
||||
image = self.image_encoder.encode(image).sample() * self.scale
|
||||
image_shape = tuple(image.shape[-2:])
|
||||
img_mask = torch.from_numpy(img_mask).unsqueeze(0).unsqueeze(0)
|
||||
img_mask = F.interpolate(
|
||||
img_mask,
|
||||
image_shape,
|
||||
mode="nearest",
|
||||
)
|
||||
img_mask = prepare_mask(img_mask).to(self.device)
|
||||
if self.use_fp16:
|
||||
img_mask = img_mask.half()
|
||||
image = image.repeat(2, 1, 1, 1)
|
||||
img_mask = img_mask.repeat(2, 1, 1, 1)
|
||||
return self.generate_img(
|
||||
prompt=prompt,
|
||||
batch_size=1,
|
||||
diffusion=diffusion,
|
||||
num_steps=num_steps,
|
||||
guidance_scale=guidance_scale,
|
||||
progress=progress,
|
||||
dynamic_threshold_v=dynamic_threshold_v,
|
||||
denoised_type=denoised_type,
|
||||
init_img=image,
|
||||
img_mask=img_mask,
|
||||
sampler=sampler,
|
||||
ddim_eta=ddim_eta,
|
||||
)
|
||||
@@ -0,0 +1,72 @@
|
||||
import math
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import importlib
|
||||
from .model.utils import get_named_beta_schedule, _extract_into_tensor
|
||||
from copy import deepcopy
|
||||
|
||||
|
||||
|
||||
def prepare_mask(mask):
|
||||
mask = mask.float()[0]
|
||||
old_mask = deepcopy(mask)
|
||||
for i in range(mask.shape[1]):
|
||||
for j in range(mask.shape[2]):
|
||||
if old_mask[0][i][j] == 1:
|
||||
continue
|
||||
if i != 0:
|
||||
mask[:, i - 1, j] = 0
|
||||
if j != 0:
|
||||
mask[:, i, j - 1] = 0
|
||||
if i != 0 and j != 0:
|
||||
mask[:, i - 1, j - 1] = 0
|
||||
if i != mask.shape[1] - 1:
|
||||
mask[:, i + 1, j] = 0
|
||||
if j != mask.shape[2] - 1:
|
||||
mask[:, i, j + 1] = 0
|
||||
if i != mask.shape[1] - 1 and j != mask.shape[2] - 1:
|
||||
mask[:, i + 1, j + 1] = 0
|
||||
return mask.unsqueeze(0)
|
||||
|
||||
|
||||
def prepare_image(pil_image, w=512, h=512):
|
||||
pil_image = pil_image.resize((w, h), resample=Image.BICUBIC, reducing_gap=1)
|
||||
arr = np.array(pil_image.convert("RGB"))
|
||||
arr = arr.astype(np.float32) / 127.5 - 1
|
||||
arr = np.transpose(arr, [2, 0, 1])
|
||||
image = torch.from_numpy(arr).unsqueeze(0)
|
||||
return image
|
||||
|
||||
|
||||
def q_sample(x_start, t, schedule_name="linear", num_steps=1000, noise=None):
|
||||
betas = get_named_beta_schedule(schedule_name, num_steps)
|
||||
alphas = 1.0 - betas
|
||||
alphas_cumprod = np.cumprod(alphas, axis=0)
|
||||
sqrt_alphas_cumprod = np.sqrt(alphas_cumprod)
|
||||
sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - alphas_cumprod)
|
||||
if noise is None:
|
||||
noise = torch.randn_like(x_start)
|
||||
assert noise.shape == x_start.shape
|
||||
return (
|
||||
_extract_into_tensor(sqrt_alphas_cumprod, t, x_start.shape) * x_start
|
||||
+ _extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
|
||||
)
|
||||
|
||||
|
||||
def process_images(batch):
|
||||
scaled = (
|
||||
((batch + 1) * 127.5)
|
||||
.round()
|
||||
.clamp(0, 255)
|
||||
.to(torch.uint8)
|
||||
.to("cpu")
|
||||
.permute(0, 2, 3, 1)
|
||||
.numpy()
|
||||
)
|
||||
images = []
|
||||
for i in range(scaled.shape[0]):
|
||||
images.append(Image.fromarray(scaled[i]))
|
||||
return images
|
||||
|
||||
Reference in New Issue
Block a user