From 20909a23e452efafcedb125fc0f7138de6a8689c Mon Sep 17 00:00:00 2001 From: smthemex <138738845+smthemex@users.noreply.github.com> Date: Wed, 12 Jun 2024 19:48:53 +0800 Subject: [PATCH] Add files via upload add --- kandinsky2/configs.py | 163 +++++++++ kandinsky2/kandinsky2_1_model.py | 552 +++++++++++++++++++++++++++++++ kandinsky2/kandinsky2_2_model.py | 189 +++++++++++ kandinsky2/kandinsky2_model.py | 381 +++++++++++++++++++++ kandinsky2/utils.py | 72 ++++ 5 files changed, 1357 insertions(+) create mode 100644 kandinsky2/configs.py create mode 100644 kandinsky2/kandinsky2_1_model.py create mode 100644 kandinsky2/kandinsky2_2_model.py create mode 100644 kandinsky2/kandinsky2_model.py create mode 100644 kandinsky2/utils.py diff --git a/kandinsky2/configs.py b/kandinsky2/configs.py new file mode 100644 index 0000000..165277e --- /dev/null +++ b/kandinsky2/configs.py @@ -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, + }, +} \ No newline at end of file diff --git a/kandinsky2/kandinsky2_1_model.py b/kandinsky2/kandinsky2_1_model.py new file mode 100644 index 0000000..343cdad --- /dev/null +++ b/kandinsky2/kandinsky2_1_model.py @@ -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, + ) diff --git a/kandinsky2/kandinsky2_2_model.py b/kandinsky2/kandinsky2_2_model.py new file mode 100644 index 0000000..2af93dc --- /dev/null +++ b/kandinsky2/kandinsky2_2_model.py @@ -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 diff --git a/kandinsky2/kandinsky2_model.py b/kandinsky2/kandinsky2_model.py new file mode 100644 index 0000000..b372d3a --- /dev/null +++ b/kandinsky2/kandinsky2_model.py @@ -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, + ) diff --git a/kandinsky2/utils.py b/kandinsky2/utils.py new file mode 100644 index 0000000..fdb1ea7 --- /dev/null +++ b/kandinsky2/utils.py @@ -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 +