import random
import os
import PIL
import torch
import warnings
warnings.filterwarnings("ignore")
from transformers import set_seed
from tqdm import tqdm
from transformers import logging
from diffusers import ControlNetModel, StableDiffusionControlNetImg2ImgPipeline, DDIMScheduler
import torch.nn as nn
import numpy as np
import utils.feature_utils as fu
import utils.preprocesser_utils as pu
import utils.image_process_utils as ipu
logging.set_verbosity_error()
def set_seed_lib(seed):
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
random.seed(seed)
set_seed(seed)
@torch.no_grad()
class RAVE(nn.Module):
def __init__(self, device):
super().__init__()
self.device = device
self.dtype = torch.float
@torch.no_grad()
def __init_pipe(self, hf_cn_path, hf_path):
controlnet = ControlNetModel.from_pretrained(hf_cn_path, torch_dtype=self.dtype).to(self.device, self.dtype)
pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(hf_path, controlnet=controlnet, torch_dtype=self.dtype).to(self.device, self.dtype)
pipe.enable_model_cpu_offload()
pipe.enable_xformers_memory_efficient_attention()
return pipe
@torch.no_grad()
def init_models(self, hf_cn_path, hf_path, preprocess_name, model_id=None):
if model_id is None or model_id == "None":
pipe = self.__init_pipe(hf_cn_path, hf_path)
else:
pipe = self.__init_pipe(hf_cn_path, model_id)
self.preprocess_name = preprocess_name
self._prepare_control_image = pipe.prepare_control_image
self.run_safety_checker = pipe.run_safety_checker
self.tokenizer = pipe.tokenizer
self.text_encoder = pipe.text_encoder
self.vae = pipe.vae
self.unet = pipe.unet
self.controlnet = pipe.controlnet
self.scheduler_config = pipe.scheduler.config
del pipe
@torch.no_grad()
def get_text_embeds(self, prompt, negative_prompt):
cond_input = self.tokenizer(prompt, padding='max_length', max_length=self.tokenizer.model_max_length, truncation=True, return_tensors='pt')
cond_embeddings = self.text_encoder(cond_input.input_ids.to(self.device))[0]
uncond_input = self.tokenizer(negative_prompt, padding='max_length', max_length=self.tokenizer.model_max_length, return_tensors='pt')
uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]
return cond_embeddings, uncond_embeddings
@torch.no_grad()
def prepare_control_image(self, control_pil, width, height):
control_image = self._prepare_control_image(
image=control_pil,
width=width,
height=height,
device=self.device,
dtype=self.controlnet.dtype,
batch_size=1,
num_images_per_prompt=1
)
return control_image
@torch.no_grad()
def pred_controlnet_sampling(self, current_sampling_percent, latent_model_input, t, text_embeddings, control_image):
if (current_sampling_percent < self.controlnet_guidance_start or current_sampling_percent > self.controlnet_guidance_end):
down_block_res_samples = None
mid_block_res_sample = None
else:
down_block_res_samples, mid_block_res_sample = self.controlnet(
latent_model_input,
t,
conditioning_scale=self.controlnet_conditioning_scale,
encoder_hidden_states=text_embeddings,
controlnet_cond=control_image,
return_dict=False,
)
noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings,
down_block_additional_residuals=down_block_res_samples,
mid_block_additional_residual=mid_block_res_sample)['sample']
return noise_pred
@torch.no_grad()
def denoising_step(self, latents, control_image, text_embeddings, t, guidance_scale, current_sampling_percent):
latent_model_input = torch.cat([latents] * 2)
control_image = torch.cat([control_image] * 2)
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
noise_pred = self.pred_controlnet_sampling(current_sampling_percent, latent_model_input, t, text_embeddings, control_image)
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
latents = self.scheduler.step(noise_pred, t, latents)['prev_sample']
return latents
@torch.no_grad()
def preprocess_control_grid(self, image_pil):
list_of_image_pils = fu.pil_grid_to_frames(image_pil, grid_size=self.grid) # List[C, W, H] -> len = num_frames
list_of_pils = [pu.pixel_perfect_process(np.array(frame_pil, dtype='uint8'), self.preprocess_name) for frame_pil in list_of_image_pils]
control_images = np.array(list_of_pils)
control_img = ipu.create_grid_from_numpy(control_images, grid_size=self.grid)
control_img = PIL.Image.fromarray(control_img).convert("L")
return control_img
@torch.no_grad()
def shuffle_latents(self, latents, control_image, indices):
rand_i = torch.randperm(self.total_frame_number).tolist()
latents_l, controls_l, randx = [], [], []
for j in range(self.sample_size):
rand_indices = rand_i[j*self.grid_frame_number:(j+1)*self.grid_frame_number]
latents_keyframe, _ = fu.prepare_key_grid_latents(latents, self.grid, self.grid, rand_indices)
control_keyframe, _ = fu.prepare_key_grid_latents(control_image, self.grid, self.grid, rand_indices)
latents_l.append(latents_keyframe)
controls_l.append(control_keyframe)
randx.extend(rand_indices)
rand_i = randx.copy()
latents = torch.cat(latents_l, dim=0)
control_image = torch.cat(controls_l, dim=0)
indices = [indices[i] for i in rand_i]
return latents, indices, control_image
@torch.no_grad()
def batch_denoise(self, latents, control_image, indices, t, guidance_scale, current_sampling_percent):
latents_l, controls_l = [], []
control_split = control_image.split(self.batch_size, dim=0)
latents_split = latents.split(self.batch_size, dim=0)
for idx in range(len(control_split)):
txt_embed = torch.cat([self.uncond_embeddings] * len(latents_split[idx]) + [self.cond_embeddings] * len(latents_split[idx]))
latents = self.denoising_step(latents_split[idx], control_split[idx], txt_embed, t, guidance_scale, current_sampling_percent)
latents_l.append(latents)
controls_l.append(control_split[idx])
latents = torch.cat(latents_l, dim=0)
controls = torch.cat(controls_l, dim=0)
return latents, indices, controls
@torch.no_grad()
def reverse_diffusion(self, latents=None, control_image=None, guidance_scale=7.5, indices=None):
self.scheduler.set_timesteps(self.num_inference_steps, device=self.device)
with torch.autocast('cuda'):
for i, t in tqdm(enumerate(self.scheduler.timesteps), desc='reverse_diffusion'):
indices = list(indices)
current_sampling_percent = i / len(self.scheduler.timesteps)
if self.is_shuffle:
latents, indices, control_image = self.shuffle_latents(latents, control_image, indices)
if self.cond_step_start < current_sampling_percent:
latents, indices, controls = self.batch_denoise(latents, control_image, indices, t, guidance_scale, current_sampling_percent)
else:
latents, indices, controls = self.batch_denoise(latents, control_image, indices, t, 0.0, current_sampling_percent)
return latents, indices, controls
@torch.no_grad()
def encode_imgs(self, img_torch):
latents_l = []
splits = img_torch.split(self.batch_size_vae, dim=0)
for split in splits:
image = 2 * split - 1
posterior = self.vae.encode(image).latent_dist
latents = posterior.mean * self.vae.config.scaling_factor
latents_l.append(latents)
return torch.cat(latents_l, dim=0)
@torch.no_grad()
def decode_latents(self, latents: torch.Tensor):
image_l = []
splits = latents.split(self.batch_size_vae, dim=0)
for split in splits:
image = self.vae.decode(split / self.vae.config.scaling_factor, return_dict=False)[0]
image = (image / 2 + 0.5).clamp(0, 1)
image_l.append(image)
return torch.cat(image_l, dim=0)
@torch.no_grad()
def controlnet_pred(self, latent_model_input, t, text_embed_input, controlnet_cond):
down_block_res_samples, mid_block_res_sample = self.controlnet(
latent_model_input,
t,
encoder_hidden_states=text_embed_input,
controlnet_cond=controlnet_cond,
conditioning_scale=1,
return_dict=False,
)
noise_pred = self.unet(
latent_model_input,
t,
encoder_hidden_states=text_embed_input,
cross_attention_kwargs={},
down_block_additional_residuals=down_block_res_samples,
mid_block_additional_residual=mid_block_res_sample,
return_dict=False,
)[0]
return noise_pred
@torch.no_grad()
def ddim_inversion(self, latents, control_batch, indices):
k = None
els = os.listdir(self.inverse_path)
els = [el for el in els if el.endswith('.pt')]
for k,inv_path in enumerate(sorted(els, key=lambda x: int(x.split('.')[0]))):
latents[k] = torch.load(os.path.join(self.inverse_path, inv_path)).to(device=self.device)
self.inverse_scheduler = DDIMScheduler.from_config(self.scheduler_config)
self.inverse_scheduler.set_timesteps(self.num_inversion_step, device=self.device)
self.timesteps = reversed(self.inverse_scheduler.timesteps)
if k == (latents.shape[0]-1):
return latents, indices, control_batch
inv_cond = torch.cat([self.inv_uncond_embeddings] * 1 + [self.inv_cond_embeddings] * 1)[1].unsqueeze(0)
for i, t in enumerate(tqdm(self.timesteps)):
alpha_prod_t = self.inverse_scheduler.alphas_cumprod[t]
alpha_prod_t_prev = (self.inverse_scheduler.alphas_cumprod[self.timesteps[i - 1]] if i > 0 else self.inverse_scheduler.final_alpha_cumprod)
if k is not None:
if len(latents[:k+1].shape) == 3:
latents[:k+1] = latents[:k+1].unsqueeze(0)
latents_l = [] if k is None else [latents[:k+1]]
latents_split = latents.split(self.inv_batch_size, dim=0) if k is None else latents[k+1:].split(self.inv_batch_size, dim=0)
control_batch_split = control_batch.split(self.inv_batch_size, dim=0) if k is None else control_batch[k+1:].split(self.inv_batch_size, dim=0)
for idx in range(len(latents_split)):
cond_batch = inv_cond.repeat(latents_split[idx].shape[0], 1, 1)
latents = self.ddim_step(latents_split[idx], t, cond_batch, alpha_prod_t, alpha_prod_t_prev, control_batch_split[idx])
latents_l.append(latents)
latents = torch.cat(latents_l, dim=0)
for k,i in enumerate(latents):
torch.save(i.detach().cpu(), f'{self.inverse_path}/{str(k).zfill(5)}.pt')
return latents, indices, control_batch
def ddim_step(self, latent_frames, t, cond_batch, alpha_prod_t, alpha_prod_t_prev, control_batch):
mu = alpha_prod_t ** 0.5
mu_prev = alpha_prod_t_prev ** 0.5
sigma = (1 - alpha_prod_t) ** 0.5
sigma_prev = (1 - alpha_prod_t_prev) ** 0.5
if self.give_control_inversion:
eps = self.controlnet_pred(latent_frames, t, text_embed_input=cond_batch, controlnet_cond=control_batch)
else:
eps = self.unet(latent_frames, t, encoder_hidden_states=cond_batch, return_dict=False)[0]
pred_x0 = (latent_frames - sigma_prev * eps) / mu_prev
latent_frames = mu * pred_x0 + sigma * eps
return latent_frames
def process_image_batch(self, image_pil_list):
if len(os.listdir(self.controls_path)) > 0:
control_torch = torch.load(os.path.join(self.controls_path, 'control.pt')).to(self.device)
img_torch = torch.load(os.path.join(self.controls_path, 'img.pt')).to(self.device)
else:
image_torch_list = []
control_torch_list = []
for image_pil in image_pil_list:
width, height = image_pil.size
control_pil = self.preprocess_control_grid(image_pil)
control_image = self.prepare_control_image(control_pil, width, height)
control_torch_list.append(control_image)
image_torch_list.append(ipu.pil_img_to_torch_tensor(image_pil))
control_torch = torch.cat(control_torch_list, dim=0).to(self.device)
img_torch = torch.cat(image_torch_list, dim=0).to(self.device)
torch.save(control_torch, os.path.join(self.controls_path, 'control.pt'))
torch.save(img_torch, os.path.join(self.controls_path, 'img.pt'))
return img_torch, control_torch
def order_grids(self, list_of_pils, indices):
k = []
for i in range(len(list_of_pils)):
k.extend(fu.pil_grid_to_frames(list_of_pils[i], self.grid))
frames = [k[indices.index(i)] for i in np.arange(len(indices))]
return frames
@torch.autocast(dtype=torch.float16, device_type='cuda')
def batched_denoise_step(self, x, t, indices):
batch_size = self.config["batch_size"]
denoised_latents = []
pivotal_idx = torch.randint(batch_size, (len(x)//batch_size,)) + torch.arange(0,len(x),batch_size)
self.denoise_step(x[pivotal_idx], t, indices[pivotal_idx])
for i, b in enumerate(range(0, len(x), batch_size)):
denoised_latents.append(self.denoise_step(x[b:b + batch_size], t, indices[b:b + batch_size]))
denoised_latents = torch.cat(denoised_latents)
return denoised_latents
@torch.no_grad()
def __preprocess_inversion_input(self, init_latents, control_batch):
list_of_flattens = [fu.flatten_grid(el.unsqueeze(0), self.grid) for el in init_latents]
init_latents = torch.cat(list_of_flattens, dim=-1)
init_latents = torch.cat(torch.chunk(init_latents, self.total_frame_number, dim=-1), dim=0)
control_batch_flattens = [fu.flatten_grid(el.unsqueeze(0), self.grid) for el in control_batch]
control_batch = torch.cat(control_batch_flattens, dim=-1)
control_batch = torch.cat(torch.chunk(control_batch, self.total_frame_number, dim=-1), dim=0)
return init_latents, control_batch
@torch.no_grad()
def __postprocess_inversion_input(self, latents_inverted, control_batch):
latents_inverted = torch.cat([fu.unflatten_grid(torch.cat([a for a in latents_inverted[i*self.grid_frame_number:(i+1)*self.grid_frame_number]], dim=-1).unsqueeze(0), self.grid) for i in range(self.sample_size)] , dim=0)
control_batch = torch.cat([fu.unflatten_grid(torch.cat([a for a in control_batch[i*self.grid_frame_number:(i+1)*self.grid_frame_number]], dim=-1).unsqueeze(0), self.grid) for i in range(self.sample_size)] , dim=0)
return latents_inverted, control_batch
@torch.no_grad()
def __call__(self, input_dict):
set_seed_lib(input_dict['seed'])
self.grid_size = input_dict['grid_size']
self.sample_size = input_dict['sample_size']
self.grid_frame_number = self.grid_size * self.grid_size
self.total_frame_number = (self.grid_frame_number) * self.sample_size
self.grid = [self.grid_size, self.grid_size]
self.cond_step_start = input_dict['cond_step_start']
self.controlnet_guidance_start = input_dict['controlnet_guidance_start']
self.controlnet_guidance_end = input_dict['controlnet_guidance_end']
self.controlnet_conditioning_scale = input_dict['controlnet_conditioning_scale']
self.positive_prompts = input_dict['positive_prompts']
self.negative_prompts = input_dict['negative_prompts']
self.inversion_prompt = input_dict['inversion_prompt']
self.batch_size = input_dict['batch_size']
self.inv_batch_size = self.batch_size * self.grid_size * self.grid_size
self.batch_size_vae = input_dict['batch_size_vae']
self.num_inference_steps = input_dict['num_inference_steps']
self.num_inversion_step = input_dict['num_inversion_step']
self.inverse_path = input_dict['inverse_path']
self.controls_path = input_dict['control_path']
self.is_ddim_inversion = input_dict['is_ddim_inversion']
self.is_shuffle = input_dict['is_shuffle']
self.give_control_inversion = input_dict['give_control_inversion']
self.guidance_scale = input_dict['guidance_scale']
indices = list(np.arange(self.total_frame_number))
img_batch, control_batch = self.process_image_batch(input_dict['image_pil_list'])
init_latents_pre = self.encode_imgs(img_batch)
self.scheduler = DDIMScheduler.from_config(self.scheduler_config)
self.scheduler.set_timesteps(self.num_inference_steps, device=self.device)
self.inv_cond_embeddings, self.inv_uncond_embeddings = self.get_text_embeds(self.inversion_prompt, "")
if self.is_ddim_inversion:
init_latents, control_batch = self.__preprocess_inversion_input(init_latents_pre, control_batch)
latents_inverted, indices, control_batch = self.ddim_inversion(init_latents, control_batch, indices)
latents_inverted, control_batch = self.__postprocess_inversion_input(latents_inverted, control_batch)
else:
init_latents_pre = torch.cat([init_latents_pre], dim=0)
noise = torch.randn_like(init_latents_pre)
latents_inverted = self.scheduler.add_noise(init_latents_pre, noise, self.scheduler.timesteps[:1])
self.cond_embeddings, self.uncond_embeddings = self.get_text_embeds(self.positive_prompts, self.negative_prompts)
latents_denoised, indices, controls = self.reverse_diffusion(latents_inverted, control_batch, self.guidance_scale, indices=indices)
image_torch = self.decode_latents(latents_denoised)
ordered_img_frames = self.order_grids(ipu.torch_to_pil_img_batch(image_torch), indices)
ordered_control_frames = self.order_grids(ipu.torch_to_pil_img_batch(controls), indices)
return ordered_img_frames, ordered_control_frames