Spaces:
Runtime error
Runtime error
| from transformers import CLIPTextModel, CLIPTokenizer, logging | |
| from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler | |
| from diffusers.utils.torch_utils import randn_tensor | |
| # suppress partial model loading warning | |
| logging.set_verbosity_error() | |
| import os | |
| from tqdm import tqdm, trange | |
| import torch | |
| import torch.nn as nn | |
| import argparse | |
| from torchvision.io import write_video | |
| from pathlib import Path | |
| from utils import * | |
| import torchvision.transforms as T | |
| import cv2 | |
| import numpy as np | |
| def get_timesteps(scheduler, num_inference_steps, strength, device): | |
| # get the original timestep using init_timestep | |
| init_timestep = min(int(num_inference_steps * strength), num_inference_steps) | |
| t_start = max(num_inference_steps - init_timestep, 0) | |
| timesteps = scheduler.timesteps[t_start:] | |
| return timesteps, num_inference_steps - t_start | |
| class Preprocess(nn.Module): | |
| def __init__(self, device, opt, vae, tokenizer, text_encoder, unet,scheduler, hf_key=None): | |
| super().__init__() | |
| self.device = device | |
| self.sd_version = opt["sd_version"] | |
| self.use_depth = False | |
| self.config = opt | |
| print(f'[INFO] loading stable diffusion...') | |
| if hf_key is not None: | |
| print(f'[INFO] using hugging face custom model key: {hf_key}') | |
| model_key = hf_key | |
| elif self.sd_version == '2.1': | |
| model_key = "stabilityai/stable-diffusion-2-1-base" | |
| elif self.sd_version == '2.0': | |
| model_key = "stabilityai/stable-diffusion-2-base" | |
| elif self.sd_version == '1.5' or self.sd_version == 'ControlNet': | |
| model_key = "runwayml/stable-diffusion-v1-5" | |
| elif self.sd_version == 'depth': | |
| model_key = "stabilityai/stable-diffusion-2-depth" | |
| else: | |
| raise ValueError(f'Stable-diffusion version {self.sd_version} not supported.') | |
| self.model_key = model_key | |
| # Create model | |
| # self.vae = AutoencoderKL.from_pretrained(model_key, subfolder="vae", revision="fp16", | |
| # torch_dtype=torch.float16).to(self.device) | |
| # self.tokenizer = CLIPTokenizer.from_pretrained(model_key, subfolder="tokenizer") | |
| # self.text_encoder = CLIPTextModel.from_pretrained(model_key, subfolder="text_encoder", revision="fp16", | |
| # torch_dtype=torch.float16).to(self.device) | |
| # self.unet = UNet2DConditionModel.from_pretrained(model_key, subfolder="unet", revision="fp16", | |
| # torch_dtype=torch.float16).to(self.device) | |
| self.vae = vae | |
| self.tokenizer = tokenizer | |
| self.text_encoder = text_encoder | |
| self.unet = unet | |
| self.scheduler=scheduler | |
| self.total_inverted_latents = {} | |
| self.noise_total = None # will contain all zs if inversion == 'ddpm', var name chosen to match the save path of zs used in pr https://github.com/omerbt/TokenFlow/pull/24/files# | |
| self.paths, self.frames, self.latents = self.get_data(self.config["data_path"], self.config["n_frames"]) | |
| print("self.frames", self.frames.shape) | |
| print("self.latents", self.latents.shape) | |
| if self.sd_version == 'ControlNet': | |
| from diffusers import ControlNetModel, StableDiffusionControlNetPipeline | |
| controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16).to(self.device) | |
| control_pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16 | |
| ).to(self.device) | |
| self.unet = control_pipe.unet | |
| self.controlnet = control_pipe.controlnet | |
| self.canny_cond = self.get_canny_cond() | |
| elif self.sd_version == 'depth': | |
| self.depth_maps = self.prepare_depth_maps() | |
| self.scheduler = scheduler | |
| self.unet.enable_xformers_memory_efficient_attention() | |
| print(f'[INFO] loaded stable diffusion!') | |
| def prepare_depth_maps(self, model_type='DPT_Large', device='cuda'): | |
| depth_maps = [] | |
| midas = torch.hub.load("intel-isl/MiDaS", model_type) | |
| midas.to(device) | |
| midas.eval() | |
| midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms") | |
| if model_type == "DPT_Large" or model_type == "DPT_Hybrid": | |
| transform = midas_transforms.dpt_transform | |
| else: | |
| transform = midas_transforms.small_transform | |
| for i in range(len(self.paths)): | |
| img = cv2.imread(self.paths[i]) | |
| img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
| latent_h = img.shape[0] // 8 | |
| latent_w = img.shape[1] // 8 | |
| input_batch = transform(img).to(device) | |
| prediction = midas(input_batch) | |
| depth_map = torch.nn.functional.interpolate( | |
| prediction.unsqueeze(1), | |
| size=(latent_h, latent_w), | |
| mode="bicubic", | |
| align_corners=False, | |
| ) | |
| depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True) | |
| depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True) | |
| depth_map = 2.0 * (depth_map - depth_min) / (depth_max - depth_min) - 1.0 | |
| depth_maps.append(depth_map) | |
| return torch.cat(depth_maps).to(self.device).to(torch.float16) | |
| def get_canny_cond(self): | |
| canny_cond = [] | |
| for image in self.frames.cpu().permute(0, 2, 3, 1): | |
| image = np.uint8(np.array(255 * image)) | |
| low_threshold = 100 | |
| high_threshold = 200 | |
| image = cv2.Canny(image, low_threshold, high_threshold) | |
| image = image[:, :, None] | |
| image = np.concatenate([image, image, image], axis=2) | |
| image = torch.from_numpy((image.astype(np.float32) / 255.0)) | |
| canny_cond.append(image) | |
| canny_cond = torch.stack(canny_cond).permute(0, 3, 1, 2).to(self.device).to(torch.float16) | |
| return canny_cond | |
| 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, | |
| ) | |
| # 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 | |
| def encode_text(self, prompts, device=None): | |
| if device is None: | |
| device = self.device | |
| text_inputs = self.tokenizer( | |
| prompts, | |
| padding="max_length", | |
| max_length=self.tokenizer.model_max_length, | |
| return_tensors="pt", | |
| ) | |
| text_input_ids = text_inputs.input_ids | |
| if text_input_ids.shape[-1] > self.tokenizer.model_max_length: | |
| removed_text = self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length:]) | |
| print( | |
| "The following part of your input was truncated because CLIP can only handle sequences up to" | |
| f" {self.tokenizer.model_max_length} tokens: {removed_text}" | |
| ) | |
| text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length] | |
| text_embeddings = self.text_encoder(text_input_ids.to(device))[0] | |
| return text_embeddings | |
| def get_text_embeds(self, prompt, negative_prompt, device="cuda"): | |
| text_embeddings = self.encode_text(prompt, device=device) | |
| uncond_embeddings = self.encode_text(negative_prompt, device=device) | |
| text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) | |
| return text_embeddings | |
| def decode_latents(self, latents): | |
| decoded = [] | |
| batch_size = 8 | |
| for b in range(0, latents.shape[0], batch_size): | |
| latents_batch = 1 / 0.18215 * latents[b:b + batch_size] | |
| imgs = self.vae.decode(latents_batch).sample | |
| imgs = (imgs / 2 + 0.5).clamp(0, 1) | |
| decoded.append(imgs) | |
| return torch.cat(decoded) | |
| def encode_imgs(self, imgs, batch_size=10, deterministic=True): | |
| imgs = 2 * imgs - 1 | |
| latents = [] | |
| for i in range(0, len(imgs), batch_size): | |
| posterior = self.vae.encode(imgs[i:i + batch_size]).latent_dist | |
| latent = posterior.mean if deterministic else posterior.sample() | |
| latents.append(latent * self.vae.config.scaling_factor) | |
| latents = torch.cat(latents) | |
| return latents | |
| def get_data(self, frames_path, n_frames): | |
| # load frames | |
| if not self.config["frames"]: | |
| paths = [f"{frames_path}/%05d.png" % i for i in range(n_frames)] | |
| print(paths) | |
| if not os.path.exists(paths[0]): | |
| paths = [f"{frames_path}/%05d.jpg" % i for i in range(n_frames)] | |
| self.paths = paths | |
| frames = [Image.open(path).convert('RGB') for path in paths] | |
| if frames[0].size[0] == frames[0].size[1]: | |
| frames = [frame.resize((512, 512), resample=Image.Resampling.LANCZOS) for frame in frames] | |
| else: | |
| frames = self.config["frames"][:n_frames] | |
| frames = torch.stack([T.ToTensor()(frame) for frame in frames]).to(torch.float16).to(self.device) | |
| # encode to latents | |
| latents = self.encode_imgs(frames, deterministic=True).to(torch.float16).to(self.device) | |
| print("frames", frames.shape) | |
| print("latents", latents.shape) | |
| if not self.config["frames"]: | |
| return paths, frames, latents | |
| else: | |
| return None, frames, latents | |
| def ddim_inversion(self, cond, latent_frames, save_path, batch_size, save_latents=True, timesteps_to_save=None): | |
| timesteps = reversed(self.scheduler.timesteps) | |
| timesteps_to_save = timesteps_to_save if timesteps_to_save is not None else timesteps | |
| return_inverted_latents = self.config["frames"] is not None | |
| for i, t in enumerate(tqdm(timesteps)): | |
| for b in range(0, latent_frames.shape[0], int(batch_size)): | |
| x_batch = latent_frames[b:b + batch_size] | |
| model_input = x_batch | |
| cond_batch = cond.repeat(x_batch.shape[0], 1, 1) | |
| if self.sd_version == 'depth': | |
| depth_maps = torch.cat([self.depth_maps[b: b + batch_size]]) | |
| model_input = torch.cat([x_batch, depth_maps],dim=1) | |
| alpha_prod_t = self.scheduler.alphas_cumprod[t] | |
| alpha_prod_t_prev = ( | |
| self.scheduler.alphas_cumprod[timesteps[i - 1]] | |
| if i > 0 else self.scheduler.final_alpha_cumprod | |
| ) | |
| 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 | |
| eps = self.unet(model_input, t, encoder_hidden_states=cond_batch).sample if self.sd_version != 'ControlNet' \ | |
| else self.controlnet_pred(x_batch, t, cond_batch, torch.cat([self.canny_cond[b: b + batch_size]])) | |
| pred_x0 = (x_batch - sigma_prev * eps) / mu_prev | |
| latent_frames[b:b + batch_size] = mu * pred_x0 + sigma * eps | |
| if return_inverted_latents and t in timesteps_to_save: | |
| self.total_inverted_latents[f'noisy_latents_{t}'] = latent_frames.clone() | |
| if save_latents and t in timesteps_to_save: | |
| torch.save(latent_frames, os.path.join(save_path, 'latents', f'noisy_latents_{t}.pt')) | |
| if save_latents: | |
| torch.save(latent_frames, os.path.join(save_path, 'latents', f'noisy_latents_{t}.pt')) | |
| if return_inverted_latents: | |
| self.total_inverted_latents[f'noisy_latents_{t}'] = latent_frames.clone() | |
| return latent_frames | |
| def ddpm_inversion(self, cond, | |
| latent_frames, | |
| batch_size, | |
| num_inversion_steps, | |
| save_path=None, | |
| save_latents=True, | |
| eta: float = 1.0, | |
| skip_steps=20): | |
| timesteps = self.scheduler.timesteps | |
| return_inverted_latents = self.config["frames"] is not None | |
| variance_noise_shape = ( | |
| num_inversion_steps, | |
| *latent_frames.shape) | |
| x0 = latent_frames | |
| t_to_idx = {int(v): k for k, v in enumerate(timesteps)} | |
| xts = torch.zeros(size=variance_noise_shape, device=self.device, dtype=cond.dtype) | |
| for t in reversed(timesteps): | |
| idx = t_to_idx[int(t)] | |
| for b in range(0, x0.shape[0], batch_size): | |
| x_batch = x0[b:b + batch_size] | |
| noise = randn_tensor(shape=x_batch.shape, device=self.device, dtype=x0.dtype) | |
| xts[idx, b:b + batch_size] = self.scheduler.add_noise(x_batch, noise, t) | |
| xts = torch.cat([xts, x0.unsqueeze(0)], dim=0) | |
| zs = torch.zeros(size=variance_noise_shape, device=self.device, dtype=cond.dtype) | |
| for t in tqdm(timesteps): | |
| idx = t_to_idx[int(t)] | |
| # 1. predict noise residual | |
| for b in range(0, x0.shape[0], batch_size): | |
| xt = xts[idx, b:b + batch_size] | |
| cond_batch = cond.repeat(xt.shape[0], 1, 1) | |
| noise_pred = self.unet(xt, timestep=t, encoder_hidden_states=cond_batch).sample | |
| xtm1 = xts[idx + 1, b:b + batch_size] | |
| z, xtm1_corrected = compute_noise(self.scheduler, xtm1, xt, t, noise_pred, eta) | |
| zs[idx, b:b + batch_size] = z | |
| # correction to avoid error accumulation | |
| xts[idx + 1, b:b + batch_size] = xtm1_corrected | |
| if save_latents: | |
| torch.save(xts[idx], os.path.join(save_path, 'latents', f'noisy_latents_{t}.pt')) | |
| if return_inverted_latents: | |
| self.total_inverted_latents[f'noisy_latents_{t}'] = xts[idx].clone() | |
| if save_path: | |
| torch.save(xts[idx], os.path.join(save_path, 'latents', f'noisy_latents_{t}.pt')) | |
| torch.save(zs, os.path.join(save_path, 'latents', f'noise_total.pt')) | |
| if return_inverted_latents: | |
| self.total_inverted_latents[f'noisy_latents_{t}'] = xts[idx].clone() | |
| self.noise_total = zs.clone() | |
| return xts[skip_steps].expand(latent_frames.shape[0], -1, -1, -1), zs | |
| def prepare_extra_step_kwargs(self, eta): | |
| # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature | |
| # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. | |
| # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 | |
| # and should be between [0, 1] | |
| accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) | |
| extra_step_kwargs = {} | |
| if accepts_eta: | |
| extra_step_kwargs["eta"] = eta | |
| # check if the scheduler accepts generator | |
| accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) | |
| return extra_step_kwargs | |
| def ddpm_sample(self, init_latents, cond, batch_size, num_inversion_steps, skip_steps, eta, zs_all, | |
| guidance_scale=0): | |
| use_ddpm = True | |
| do_classifier_free_guidance = guidance_scale > 1.0 | |
| total_latents = init_latents | |
| self.scheduler.set_timesteps(num_inversion_steps, device=device) | |
| timesteps = self.scheduler.timesteps | |
| zs_total = zs_all[skip_steps:] | |
| if use_ddpm: | |
| t_to_idx = {int(v): k for k, v in enumerate(timesteps[-zs_total.shape[0]:])} | |
| timesteps = timesteps[-zs_total.shape[0]:] | |
| num_warmup_steps = len(timesteps) - num_inversion_steps * self.scheduler.order | |
| extra_step_kwargs = self.prepare_extra_step_kwargs(eta) | |
| for i, t in enumerate(tqdm(timesteps)): | |
| for b in range(0, total_latents.shape[0], batch_size): | |
| latents = total_latents[b:b + batch_size] | |
| if do_classifier_free_guidance: | |
| latent_model_input = torch.cat([latents] * 2) | |
| else: | |
| latent_model_input = latents | |
| cond_batch = cond.repeat(latents.shape[0], 1, 1) | |
| latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) | |
| noise_pred = self.unet( | |
| latent_model_input, | |
| t, | |
| encoder_hidden_states=cond_batch, | |
| return_dict=False, | |
| )[0] | |
| if do_classifier_free_guidance: | |
| noise_pred_out = noise_pred.chunk(2) # [b,4, 64, 64] | |
| noise_pred_uncond, noise_pred_text = noise_pred_out[0], noise_pred_out[1] | |
| # default text guidance | |
| noise_guidance = guidance_scale * (noise_pred_text - noise_pred_uncond) | |
| noise_pred = noise_pred_uncond + noise_guidance | |
| idx = t_to_idx[int(t)] | |
| zs = zs_total[idx, b:b + batch_size] | |
| latents = self.scheduler.step(noise_pred, t, latents, variance_noise=zs, | |
| **extra_step_kwargs).prev_sample | |
| total_latents[b:b + batch_size] = latents | |
| return total_latents | |
| def ddim_sample(self, x, cond, batch_size): | |
| timesteps = self.scheduler.timesteps | |
| for i, t in enumerate(tqdm(timesteps)): | |
| for b in range(0, x.shape[0], batch_size): | |
| x_batch = x[b:b + batch_size] | |
| model_input = x_batch | |
| cond_batch = cond.repeat(x_batch.shape[0], 1, 1) | |
| if self.sd_version == 'depth': | |
| depth_maps = torch.cat([self.depth_maps[b: b + batch_size]]) | |
| model_input = torch.cat([x_batch, depth_maps],dim=1) | |
| alpha_prod_t = self.scheduler.alphas_cumprod[t] | |
| alpha_prod_t_prev = ( | |
| self.scheduler.alphas_cumprod[timesteps[i + 1]] | |
| if i < len(timesteps) - 1 | |
| else self.scheduler.final_alpha_cumprod | |
| ) | |
| mu = alpha_prod_t ** 0.5 | |
| sigma = (1 - alpha_prod_t) ** 0.5 | |
| mu_prev = alpha_prod_t_prev ** 0.5 | |
| sigma_prev = (1 - alpha_prod_t_prev) ** 0.5 | |
| eps = self.unet(model_input, t, encoder_hidden_states=cond_batch).sample if self.sd_version != 'ControlNet' \ | |
| else self.controlnet_pred(x_batch, t, cond_batch, torch.cat([self.canny_cond[b: b + batch_size]])) | |
| pred_x0 = (x_batch - sigma * eps) / mu | |
| x[b:b + batch_size] = mu_prev * pred_x0 + sigma_prev * eps | |
| return x | |
| def extract_latents(self, | |
| num_steps, | |
| save_path, | |
| batch_size, | |
| timesteps_to_save, | |
| inversion_prompt='', | |
| skip_steps=20, | |
| inversion_type='ddim', | |
| eta=1.0, | |
| reconstruction=False): | |
| self.scheduler.set_timesteps(num_steps) | |
| cond = self.get_text_embeds(inversion_prompt, "")[1].unsqueeze(0) | |
| latent_frames = self.latents | |
| if inversion_type == 'ddim': | |
| inverted_x= self.ddim_inversion(cond, | |
| latent_frames, | |
| save_path, | |
| batch_size=batch_size, | |
| save_latents=True if save_path else False, | |
| timesteps_to_save=timesteps_to_save) | |
| if reconstruction: | |
| latent_reconstruction = self.ddim_sample(inverted_x, cond, batch_size=batch_size) | |
| rgb_reconstruction = self.decode_latents(latent_reconstruction) | |
| return (self.frames, self.latents, self.total_inverted_latents), rgb_reconstruction | |
| else: | |
| return (self.frames, self.latents, self.total_inverted_latents), None | |
| elif inversion_type == 'ddpm': | |
| inverted_x, zs = self.ddpm_inversion(cond, | |
| latent_frames, | |
| save_path= save_path, | |
| batch_size=batch_size, | |
| save_latents=True if save_path else False, | |
| num_inversion_steps=num_steps, | |
| eta=eta, | |
| skip_steps=skip_steps) | |
| cond = self.encode_text(inversion_prompt) | |
| if reconstruction: | |
| latent_reconstruction = self.ddpm_sample(init_latents=inverted_x, | |
| cond=cond, batch_size=batch_size, | |
| num_inversion_steps=num_steps, skip_steps=skip_steps, | |
| eta=eta, zs_all=zs) | |
| rgb_reconstruction = self.decode_latents(latent_reconstruction) | |
| return (self.frames, self.latents, self.total_inverted_latents, self.noise_total), rgb_reconstruction | |
| else: | |
| return (self.frames, self.latents, self.total_inverted_latents, self.noise_total), None | |
| else: | |
| raise NotImplementedError() | |
| def compute_noise(scheduler, prev_latents, latents, timestep, noise_pred, eta): | |
| # 1. get previous step value (=t-1) | |
| prev_timestep = timestep - scheduler.config.num_train_timesteps // scheduler.num_inference_steps | |
| # 2. compute alphas, betas | |
| alpha_prod_t = scheduler.alphas_cumprod[timestep] | |
| alpha_prod_t_prev = ( | |
| scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else scheduler.final_alpha_cumprod | |
| ) | |
| beta_prod_t = 1 - alpha_prod_t | |
| # 3. compute predicted original sample from predicted noise also called | |
| # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf | |
| pred_original_sample = (latents - beta_prod_t ** (0.5) * noise_pred) / alpha_prod_t ** (0.5) | |
| # 4. Clip "predicted x_0" | |
| if scheduler.config.clip_sample: | |
| pred_original_sample = torch.clamp(pred_original_sample, -1, 1) | |
| # 5. compute variance: "sigma_t(η)" -> see formula (16) | |
| # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1) | |
| variance = scheduler._get_variance(timestep, prev_timestep) | |
| std_dev_t = eta * variance ** (0.5) | |
| # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf | |
| pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t ** 2) ** (0.5) * noise_pred | |
| # modifed so that updated xtm1 is returned as well (to avoid error accumulation) | |
| mu_xt = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction | |
| noise = (prev_latents - mu_xt) / (variance ** (0.5) * eta) | |
| return noise, mu_xt + (eta * variance ** 0.5) * noise | |
| def prep(opt): | |
| # timesteps to save | |
| if opt["sd_version"] == '2.1': | |
| model_key = "stabilityai/stable-diffusion-2-1-base" | |
| elif opt["sd_version"] == '2.0': | |
| model_key = "stabilityai/stable-diffusion-2-base" | |
| elif opt["sd_version"] == '1.5' or opt["sd_version"] == 'ControlNet': | |
| model_key = "runwayml/stable-diffusion-v1-5" | |
| elif opt["sd_version"] == 'depth': | |
| model_key = "stabilityai/stable-diffusion-2-depth" | |
| toy_scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler") | |
| toy_scheduler.set_timesteps(opt["save_steps"]) | |
| timesteps_to_save, num_inference_steps = get_timesteps(toy_scheduler, num_inference_steps=opt["save_steps"], | |
| strength=1.0, | |
| device=device) | |
| seed_everything(opt["seed"]) | |
| if not opt["frames"]: # original non demo setting | |
| save_path = os.path.join(opt["save_dir"], | |
| f'inversion_{opt[inversion]}', | |
| f'sd_{opt["sd_version"]}', | |
| Path(opt["data_path"]).stem, | |
| f'steps_{opt["steps"]}', | |
| f'nframes_{opt["n_frames"]}') | |
| os.makedirs(os.path.join(save_path, f'latents'), exist_ok=True) | |
| if opt[inversion] == 'ddpm': | |
| os.makedirs(os.path.join(save_path, f'latents'), exist_ok=True) | |
| add_dict_to_yaml_file(os.path.join(opt["save_dir"], 'inversion_prompts.yaml'), Path(opt["data_path"]).stem, opt["inversion_prompt"]) | |
| # save inversion prompt in a txt file | |
| with open(os.path.join(save_path, 'inversion_prompt.txt'), 'w') as f: | |
| f.write(opt["inversion_prompt"]) | |
| else: | |
| save_path = None | |
| model = Preprocess(device, | |
| config, | |
| vae=vae, | |
| text_encoder=text_encoder, | |
| scheduler=scheduler, | |
| tokenizer=tokenizer, | |
| unet=unet) | |
| frames_and_latents, rgb_reconstruction = model.extract_latents( | |
| num_steps=model.config["steps"], | |
| save_path=save_path, | |
| batch_size=model.config["batch_size"], | |
| timesteps_to_save=timesteps_to_save, | |
| inversion_prompt=model.config["inversion_prompt"], | |
| inversion_type=model.config[inversion], | |
| skip_steps=model.config[skip_steps], | |
| reconstruction=model.config[reconstruct] | |
| ) | |
| if model.config[inversion] == 'ddpm': | |
| frames, latents, total_inverted_latents, zs = frames_and_latents | |
| return frames, latents, total_inverted_latents, zs, rgb_reconstruction | |
| else: | |
| frames, latents, total_inverted_latents = frames_and_latents | |
| return frames, latents, total_inverted_latents, rgb_reconstruction | |