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.constants as const 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_MultiControlNet(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_1 = ControlNetModel.from_pretrained(hf_cn_path[0], torch_dtype=self.dtype).to(self.device, self.dtype) controlnet_2 = ControlNetModel.from_pretrained(hf_cn_path[1], torch_dtype=self.dtype).to(self.device, self.dtype) pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(hf_path, controlnet=[controlnet_1, controlnet_2], 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_1, self.preprocess_name_2 = preprocess_name.split('-') 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): # prompt, negative_prompt: [str] # Tokenize text and get embeddings 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] # Do the same for unconditional embeddings 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] # Cat for final embeddings 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_1, control_image_2, text_embeddings, t, guidance_scale, current_sampling_percent): # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes. latent_model_input = torch.cat([latents] * 2) control_image_1 = torch.cat([control_image_1] * 2) control_image_2 = torch.cat([control_image_2] * 2) latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) # compute the percentage of total steps we are at noise_pred = self.pred_controlnet_sampling(current_sampling_percent, latent_model_input, t, text_embeddings, [control_image_1, control_image_2]) # perform guidance noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 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_1, list_of_pils_2 = [], [] for frame_pil in list_of_image_pils: frame_pil_1 = pu.pixel_perfect_process(np.array(frame_pil, dtype='uint8'), self.preprocess_name_1) frame_pil_2 = pu.pixel_perfect_process(np.array(frame_pil, dtype='uint8'), self.preprocess_name_2) list_of_pils_1.append(frame_pil_1) list_of_pils_2.append(frame_pil_2) control_images_1 = np.array(list_of_pils_1) control_images_2 = np.array(list_of_pils_2) control_img_1 = ipu.create_grid_from_numpy(control_images_1, grid_size=self.grid) control_img_1 = PIL.Image.fromarray(control_img_1).convert("L") control_img_2 = ipu.create_grid_from_numpy(control_images_2, grid_size=self.grid) control_img_2 = PIL.Image.fromarray(control_img_2).convert("L") return control_img_1, control_img_2 @torch.no_grad() def shuffle_latents(self, latents, control_image_1, control_image_2, indices): rand_i = torch.randperm(self.total_frame_number).tolist() # latents, _ = fu.prepare_key_grid_latents(latents, self.grid, self.grid, rand_i) # control_image, _ = fu.prepare_key_grid_latents(control_image, self.grid, self.grid, rand_i) latents_l, controls_l_1, controls_l_2, 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_1, _ = fu.prepare_key_grid_latents(control_image_1, self.grid, self.grid, rand_indices) control_keyframe_2, _ = fu.prepare_key_grid_latents(control_image_2, self.grid, self.grid, rand_indices) latents_l.append(latents_keyframe) controls_l_1.append(control_keyframe_1) controls_l_2.append(control_keyframe_2) randx.extend(rand_indices) rand_i = randx.copy() latents = torch.cat(latents_l, dim=0) control_image_1 = torch.cat(controls_l_1, dim=0) control_image_2 = torch.cat(controls_l_2, dim=0) indices = [indices[i] for i in rand_i] return latents, indices, control_image_1, control_image_2 @torch.no_grad() def batch_denoise(self, latents, control_image_1, control_image_2, indices, t, guidance_scale, current_sampling_percent): latents_l, controls_l_1, controls_l_2 = [], [], [] control_split_1 = control_image_1.split(self.batch_size, dim=0) control_split_2 = control_image_2.split(self.batch_size, dim=0) latents_split = latents.split(self.batch_size, dim=0) for idx in range(len(control_split_1)): 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_1[idx], control_split_2[idx], txt_embed, t, guidance_scale, current_sampling_percent) latents_l.append(latents) controls_l_1.append(control_split_1[idx]) controls_l_2.append(control_split_2[idx]) latents = torch.cat(latents_l, dim=0) controls_1 = torch.cat(controls_l_1, dim=0) controls_2 = torch.cat(controls_l_2, dim=0) return latents, indices, controls_1, controls_2 @torch.no_grad() def reverse_diffusion(self, latents=None, control_image_1=None, control_image_2=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_1, control_image_2 = self.shuffle_latents(latents, control_image_1, control_image_2, indices) if self.cond_step_start < current_sampling_percent: latents, indices, control_image_1, control_image_2 = self.batch_denoise(latents, control_image_1, control_image_2, indices, t, guidance_scale, current_sampling_percent) else: latents, indices, control_image_1, control_image_2 = self.batch_denoise(latents, control_image_1, control_image_2, indices, t, 0.0, current_sampling_percent) return latents, indices, control_image_1, control_image_2 @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): 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=self.controlnet_conditioning_scale, return_dict=False, ) # apply the denoising network 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_1, control_batch_2, 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_1, control_batch_2 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) # latent, indices, control_batch = self.shuffle_latents(latent, control_batch, indices) latents_l = [] latents_split = latents.split(self.batch_size, dim=0) control_batch_split_1 = control_batch_1.split(self.batch_size, dim=0) control_batch_split_2 = control_batch_2.split(self.batch_size, dim=0) for idx in range(len(latents_split)): cond_batch = inv_cond.repeat(latents_split[idx].shape[0], 1, 1) # print(cond_batch.shape, latents_split[idx].shape, control_batch_split_1[idx].shape, control_batch_split_2[idx].shape) # input() latents = self.ddim_step(latents_split[idx], t, cond_batch, alpha_prod_t, alpha_prod_t_prev, control_batch_split_1[idx], control_batch_split_2[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_1, control_batch_2 def ddim_step(self, latent_frames, t, cond_batch, alpha_prod_t, alpha_prod_t_prev, control_batch_1, control_batch_2): 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_1, control_batch_2]) 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_1 = torch.load(os.path.join(self.controls_path, 'control_1.pt')).to(self.device) control_torch_2 = torch.load(os.path.join(self.controls_path, 'control_2.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_1, control_torch_list_2 = [], [] for image_pil in image_pil_list: width, height = image_pil.size # control_pil = PIL.Image.fromarray(pu.pixel_perfect_process(np.array(image_pil, dtype='uint8'), self.preprocess_name)) control_pil_1, control_pil_2 = self.preprocess_control_grid(image_pil) control_image_1 = self.prepare_control_image(control_pil_1, width, height) control_image_2 = self.prepare_control_image(control_pil_2, width, height) control_torch_list_1.append(control_image_1) control_torch_list_2.append(control_image_2) image_torch_list.append(ipu.pil_img_to_torch_tensor(image_pil)) control_torch_1 = torch.cat(control_torch_list_1, dim=0).to(self.device) control_torch_2 = torch.cat(control_torch_list_2, dim=0).to(self.device) img_torch = torch.cat(image_torch_list, dim=0).to(self.device) torch.save(control_torch_1, os.path.join(self.controls_path, 'control_1.pt')) torch.save(control_torch_2, os.path.join(self.controls_path, 'control_2.pt')) torch.save(img_torch, os.path.join(self.controls_path, 'img.pt')) return img_torch, control_torch_1, control_torch_2 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.no_grad() def __preprocess_inversion_input(self, init_latents, control_batch_1, control_batch_2): 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_1 = [fu.flatten_grid(el.unsqueeze(0), self.grid) for el in control_batch_1] control_batch_1 = torch.cat(control_batch_flattens_1, dim=-1) control_batch_1 = torch.cat(torch.chunk(control_batch_1, self.total_frame_number, dim=-1), dim=0) control_batch_flattens_2 = [fu.flatten_grid(el.unsqueeze(0), self.grid) for el in control_batch_2] control_batch_2 = torch.cat(control_batch_flattens_2, dim=-1) control_batch_2 = torch.cat(torch.chunk(control_batch_2, self.total_frame_number, dim=-1), dim=0) return init_latents, control_batch_1, control_batch_2 @torch.no_grad() def __postprocess_inversion_input(self, latents_inverted, control_batch_1, control_batch_2): 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_1 = torch.cat([fu.unflatten_grid(torch.cat([a for a in control_batch_1[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_2 = torch.cat([fu.unflatten_grid(torch.cat([a for a in control_batch_2[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_1, control_batch_2 @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 = [float(x) for x in input_dict['controlnet_conditioning_scale'].split('-')] 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.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_1, control_batch_2 = 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_1, control_batch_2 = self.__preprocess_inversion_input(init_latents_pre, control_batch_1, control_batch_2) latents_inverted, indices, control_batch_1, control_batch_2 = self.ddim_inversion(init_latents, control_batch_1, control_batch_2, indices) latents_inverted, control_batch_1, control_batch_2 = self.__postprocess_inversion_input(latents_inverted, control_batch_1, control_batch_2) 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_1, controls_2 = self.reverse_diffusion(latents_inverted, control_batch_1, control_batch_2, 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_1 = self.order_grids(ipu.torch_to_pil_img_batch(controls_1), indices) ordered_control_frames_2 = self.order_grids(ipu.torch_to_pil_img_batch(controls_2), indices) return ordered_img_frames, ordered_control_frames_1, ordered_control_frames_2