Spaces:
Running
on
Zero
Running
on
Zero
| # Plug&Play Feature Injection | |
| import torch | |
| from typing import Any, Callable, Dict, List, Optional, Tuple, Union | |
| from random import randrange | |
| import PIL | |
| import numpy as np | |
| from tqdm import tqdm | |
| from torch.cuda.amp import custom_bwd, custom_fwd | |
| import torch.nn.functional as F | |
| from diffusers import ( | |
| StableDiffusionPipeline, | |
| StableDiffusionImg2ImgPipeline, | |
| DDIMScheduler, | |
| ) | |
| from diffusers.utils.torch_utils import randn_tensor | |
| from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import ( | |
| StableDiffusionPipelineOutput, | |
| retrieve_timesteps, | |
| PipelineImageInput | |
| ) | |
| from src.eunms import Scheduler_Type, Gradient_Averaging_Type, Epsilon_Update_Type | |
| def _backward_ddim(x_tm1, alpha_t, alpha_tm1, eps_xt): | |
| """ | |
| let a = alpha_t, b = alpha_{t - 1} | |
| We have a > b, | |
| x_{t} - x_{t - 1} = sqrt(a) ((sqrt(1/b) - sqrt(1/a)) * x_{t-1} + (sqrt(1/a - 1) - sqrt(1/b - 1)) * eps_{t-1}) | |
| From https://arxiv.org/pdf/2105.05233.pdf, section F. | |
| """ | |
| a, b = alpha_t, alpha_tm1 | |
| sa = a**0.5 | |
| sb = b**0.5 | |
| return sa * ((1 / sb) * x_tm1 + ((1 / a - 1) ** 0.5 - (1 / b - 1) ** 0.5) * eps_xt) | |
| class SDDDIMPipeline(StableDiffusionImg2ImgPipeline): | |
| # @torch.no_grad() | |
| def __call__( | |
| self, | |
| prompt: Union[str, List[str]] = None, | |
| image: PipelineImageInput = None, | |
| strength: float = 1.0, | |
| num_inversion_steps: Optional[int] = 50, | |
| timesteps: List[int] = None, | |
| guidance_scale: Optional[float] = 7.5, | |
| negative_prompt: Optional[Union[str, List[str]]] = None, | |
| num_images_per_prompt: Optional[int] = 1, | |
| eta: Optional[float] = 0.0, | |
| generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, | |
| prompt_embeds: Optional[torch.FloatTensor] = None, | |
| negative_prompt_embeds: Optional[torch.FloatTensor] = None, | |
| ip_adapter_image: Optional[PipelineImageInput] = None, | |
| output_type: Optional[str] = "pil", | |
| return_dict: bool = True, | |
| cross_attention_kwargs: Optional[Dict[str, Any]] = None, | |
| clip_skip: int = None, | |
| callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, | |
| callback_on_step_end_tensor_inputs: List[str] = ["latents"], | |
| opt_lr: float = 0.001, | |
| opt_iters: int = 1, | |
| opt_none_inference_steps: bool = False, | |
| opt_loss_kl_lambda: float = 10.0, | |
| num_inference_steps: int = 50, | |
| num_aprox_steps: int = 100, | |
| **kwargs, | |
| ): | |
| callback = kwargs.pop("callback", None) | |
| callback_steps = kwargs.pop("callback_steps", None) | |
| if callback is not None: | |
| deprecate( | |
| "callback", | |
| "1.0.0", | |
| "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", | |
| ) | |
| if callback_steps is not None: | |
| deprecate( | |
| "callback_steps", | |
| "1.0.0", | |
| "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", | |
| ) | |
| # 1. Check inputs. Raise error if not correct | |
| self.check_inputs( | |
| prompt, | |
| strength, | |
| callback_steps, | |
| negative_prompt, | |
| prompt_embeds, | |
| negative_prompt_embeds, | |
| callback_on_step_end_tensor_inputs, | |
| ) | |
| self._guidance_scale = guidance_scale | |
| self._clip_skip = clip_skip | |
| self._cross_attention_kwargs = cross_attention_kwargs | |
| # 2. Define call parameters | |
| if prompt is not None and isinstance(prompt, str): | |
| batch_size = 1 | |
| elif prompt is not None and isinstance(prompt, list): | |
| batch_size = len(prompt) | |
| else: | |
| batch_size = prompt_embeds.shape[0] | |
| device = self._execution_device | |
| # 3. Encode input prompt | |
| text_encoder_lora_scale = ( | |
| self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None | |
| ) | |
| prompt_embeds, negative_prompt_embeds = self.encode_prompt( | |
| prompt, | |
| device, | |
| num_images_per_prompt, | |
| self.do_classifier_free_guidance, | |
| negative_prompt, | |
| prompt_embeds=prompt_embeds, | |
| negative_prompt_embeds=negative_prompt_embeds, | |
| lora_scale=text_encoder_lora_scale, | |
| clip_skip=self.clip_skip, | |
| ) | |
| # For classifier free guidance, we need to do two forward passes. | |
| # Here we concatenate the unconditional and text embeddings into a single batch | |
| # to avoid doing two forward passes | |
| if self.do_classifier_free_guidance: | |
| prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) | |
| if ip_adapter_image is not None: | |
| image_embeds, negative_image_embeds = self.encode_image(ip_adapter_image, device, num_images_per_prompt) | |
| if self.do_classifier_free_guidance: | |
| image_embeds = torch.cat([negative_image_embeds, image_embeds]) | |
| # 4. Preprocess image | |
| image = self.image_processor.preprocess(image) | |
| # 5. set timesteps | |
| timesteps, num_inversion_steps = retrieve_timesteps(self.scheduler, num_inversion_steps, device, timesteps) | |
| timesteps, num_inversion_steps = self.get_timesteps(num_inversion_steps, strength, device) | |
| latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) | |
| _, num_inference_steps = retrieve_timesteps(self.scheduler_inference, num_inference_steps, device, None) | |
| # 6. Prepare latent variables | |
| with torch.no_grad(): | |
| latents = self.prepare_latents( | |
| image, | |
| latent_timestep, | |
| batch_size, | |
| num_images_per_prompt, | |
| prompt_embeds.dtype, | |
| device, | |
| generator, | |
| ) | |
| # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline | |
| extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) | |
| # 7.1 Add image embeds for IP-Adapter | |
| added_cond_kwargs = {"image_embeds": image_embeds} if ip_adapter_image is not None else None | |
| # 7.2 Optionally get Guidance Scale Embedding | |
| timestep_cond = None | |
| if self.unet.config.time_cond_proj_dim is not None: | |
| guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt) | |
| timestep_cond = self.get_guidance_scale_embedding( | |
| guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim | |
| ).to(device=device, dtype=latents.dtype) | |
| # 8. Denoising loop | |
| num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order | |
| self._num_timesteps = len(timesteps) | |
| prev_timestep = None | |
| self.prev_z = torch.clone(latents) | |
| self.prev_z4 = torch.clone(latents) | |
| self.z_0 = torch.clone(latents) | |
| g_cpu = torch.Generator().manual_seed(7865) | |
| self.noise = randn_tensor(self.z_0.shape, generator=g_cpu, device=self.z_0.device, dtype=self.z_0.dtype) | |
| all_latents = [latents.clone()] | |
| with self.progress_bar(total=num_inversion_steps) as progress_bar: | |
| for i, t in enumerate(reversed(timesteps)): | |
| z_tp1 = self.inversion_step(latents, | |
| t, | |
| prompt_embeds, | |
| added_cond_kwargs, | |
| prev_timestep=prev_timestep, | |
| num_aprox_steps=num_aprox_steps) | |
| if t in self.scheduler_inference.timesteps: | |
| z_tp1 = self.optimize_z_tp1(z_tp1, | |
| latents, | |
| t, | |
| prompt_embeds, | |
| added_cond_kwargs, | |
| nom_opt_iters=opt_iters, | |
| lr=opt_lr, | |
| opt_loss_kl_lambda=opt_loss_kl_lambda) | |
| prev_timestep = t | |
| latents = z_tp1 | |
| all_latents.append(latents.clone()) | |
| if callback_on_step_end is not None: | |
| callback_kwargs = {} | |
| for k in callback_on_step_end_tensor_inputs: | |
| callback_kwargs[k] = locals()[k] | |
| callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) | |
| latents = callback_outputs.pop("latents", latents) | |
| prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) | |
| negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) | |
| # call the callback, if provided | |
| if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): | |
| progress_bar.update() | |
| if callback is not None and i % callback_steps == 0: | |
| step_idx = i // getattr(self.scheduler, "order", 1) | |
| callback(step_idx, t, latents) | |
| image = latents | |
| # Offload all models | |
| self.maybe_free_model_hooks() | |
| return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=None), all_latents | |
| def noise_regularization(self, e_t, noise_pred_optimal): | |
| for _outer in range(self.cfg.num_reg_steps): | |
| if self.cfg.lambda_kl>0: | |
| _var = torch.autograd.Variable(e_t.detach().clone(), requires_grad=True) | |
| # l_kld = self.kl_divergence(_var) | |
| l_kld = self.patchify_latents_kl_divergence(_var, noise_pred_optimal) | |
| l_kld.backward() | |
| _grad = _var.grad.detach() | |
| _grad = torch.clip(_grad, -100, 100) | |
| e_t = e_t - self.cfg.lambda_kl*_grad | |
| if self.cfg.lambda_ac>0: | |
| for _inner in range(self.cfg.num_ac_rolls): | |
| _var = torch.autograd.Variable(e_t.detach().clone(), requires_grad=True) | |
| l_ac = self.auto_corr_loss(_var) | |
| l_ac.backward() | |
| _grad = _var.grad.detach()/self.cfg.num_ac_rolls | |
| e_t = e_t - self.cfg.lambda_ac*_grad | |
| e_t = e_t.detach() | |
| return e_t | |
| def auto_corr_loss(self, x, random_shift=True): | |
| B,C,H,W = x.shape | |
| assert B==1 | |
| x = x.squeeze(0) | |
| # x must be shape [C,H,W] now | |
| reg_loss = 0.0 | |
| for ch_idx in range(x.shape[0]): | |
| noise = x[ch_idx][None, None,:,:] | |
| while True: | |
| if random_shift: roll_amount = randrange(noise.shape[2]//2) | |
| else: roll_amount = 1 | |
| reg_loss += (noise*torch.roll(noise, shifts=roll_amount, dims=2)).mean()**2 | |
| reg_loss += (noise*torch.roll(noise, shifts=roll_amount, dims=3)).mean()**2 | |
| if noise.shape[2] <= 8: | |
| break | |
| noise = F.avg_pool2d(noise, kernel_size=2) | |
| return reg_loss | |
| def kl_divergence(self, x): | |
| _mu = x.mean() | |
| _var = x.var() | |
| return _var + _mu**2 - 1 - torch.log(_var+1e-7) | |
| # @torch.no_grad() | |
| def inversion_step( | |
| self, | |
| z_t: torch.tensor, | |
| t: torch.tensor, | |
| prompt_embeds, | |
| added_cond_kwargs, | |
| prev_timestep: Optional[torch.tensor] = None, | |
| num_aprox_steps: int = 100 | |
| ) -> torch.tensor: | |
| extra_step_kwargs = {} | |
| avg_range = self.cfg.gradient_averaging_first_step_range if t.item() < 250 else self.cfg.gradient_averaging_step_range | |
| # When doing more then one approximation step in the first step it adds artifacts | |
| if t.item() < 250: | |
| num_aprox_steps = min(self.cfg.max_num_aprox_steps_first_step, num_aprox_steps) | |
| approximated_z_tp1 = z_t.clone() | |
| nosie_pred_avg = None | |
| if self.cfg.num_reg_steps > 0: | |
| z_tp1_forward = self.scheduler.add_noise(self.z_0, self.noise, t.view((1))).detach() | |
| latent_model_input = torch.cat([z_tp1_forward] * 2) if self.do_classifier_free_guidance else z_tp1_forward | |
| latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) | |
| with torch.no_grad(): | |
| # predict the noise residual | |
| noise_pred_optimal = self.unet( | |
| latent_model_input, | |
| t, | |
| encoder_hidden_states=prompt_embeds, | |
| timestep_cond=None, | |
| cross_attention_kwargs=self.cross_attention_kwargs, | |
| added_cond_kwargs=added_cond_kwargs, | |
| return_dict=False, | |
| )[0].detach() | |
| else: | |
| noise_pred_optimal = None | |
| for i in range(num_aprox_steps + 1): | |
| latent_model_input = torch.cat([approximated_z_tp1] * 2) if self.do_classifier_free_guidance else approximated_z_tp1 | |
| latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) | |
| with torch.no_grad(): | |
| # predict the noise residual | |
| noise_pred = self.unet( | |
| latent_model_input, | |
| t, | |
| encoder_hidden_states=prompt_embeds, | |
| timestep_cond=None, | |
| cross_attention_kwargs=self.cross_attention_kwargs, | |
| added_cond_kwargs=added_cond_kwargs, | |
| return_dict=False, | |
| )[0] | |
| if i >= avg_range[0] and i < avg_range[1]: | |
| j = i - avg_range[0] | |
| if nosie_pred_avg is None: | |
| nosie_pred_avg = noise_pred.clone() | |
| else: | |
| nosie_pred_avg = j * nosie_pred_avg / (j + 1) + noise_pred / (j + 1) | |
| if self.cfg.gradient_averaging_type == Gradient_Averaging_Type.EACH_ITER: | |
| noise_pred = nosie_pred_avg.clone() | |
| # perform guidance | |
| if self.do_classifier_free_guidance: | |
| noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) | |
| noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) | |
| if i >= avg_range[0] or (self.cfg.gradient_averaging_type == Gradient_Averaging_Type.NONE and i > 0): | |
| noise_pred = self.noise_regularization(noise_pred, noise_pred_optimal) | |
| if self.cfg.scheduler_type == Scheduler_Type.EULER: | |
| approximated_z_tp1 = self.scheduler.inv_step(noise_pred, t, z_t, **extra_step_kwargs, return_dict=False)[0].detach() | |
| else: | |
| alpha_prod_t = self.scheduler.alphas_cumprod[t] | |
| alpha_prod_t_prev = ( | |
| self.scheduler.alphas_cumprod[prev_timestep] | |
| if prev_timestep is not None | |
| else self.scheduler.final_alpha_cumprod | |
| ) | |
| approximated_z_tp1 = _backward_ddim( | |
| x_tm1=z_t, | |
| alpha_t=alpha_prod_t, | |
| alpha_tm1=alpha_prod_t_prev, | |
| eps_xt=noise_pred, | |
| ) | |
| if self.cfg.gradient_averaging_type == Gradient_Averaging_Type.ON_END and nosie_pred_avg is not None: | |
| nosie_pred_avg = self.noise_regularization(nosie_pred_avg, noise_pred_optimal) | |
| if self.cfg.scheduler_type == Scheduler_Type.EULER: | |
| approximated_z_tp1 = self.scheduler.inv_step(nosie_pred_avg, t, z_t, **extra_step_kwargs, return_dict=False)[0].detach() | |
| else: | |
| alpha_prod_t = self.scheduler.alphas_cumprod[t] | |
| alpha_prod_t_prev = ( | |
| self.scheduler.alphas_cumprod[prev_timestep] | |
| if prev_timestep is not None | |
| else self.scheduler.final_alpha_cumprod | |
| ) | |
| approximated_z_tp1 = _backward_ddim( | |
| x_tm1=z_t, | |
| alpha_t=alpha_prod_t, | |
| alpha_tm1=alpha_prod_t_prev, | |
| eps_xt=nosie_pred_avg, | |
| ) | |
| if self.cfg.update_epsilon_type != Epsilon_Update_Type.NONE: | |
| latent_model_input = torch.cat([approximated_z_tp1] * 2) if self.do_classifier_free_guidance else approximated_z_tp1 | |
| latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) | |
| with torch.no_grad(): | |
| # predict the noise residual | |
| noise_pred = self.unet( | |
| latent_model_input, | |
| t, | |
| encoder_hidden_states=prompt_embeds, | |
| timestep_cond=None, | |
| cross_attention_kwargs=self.cross_attention_kwargs, | |
| added_cond_kwargs=added_cond_kwargs, | |
| return_dict=False, | |
| )[0] | |
| # perform guidance | |
| if self.do_classifier_free_guidance: | |
| noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) | |
| noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) | |
| self.scheduler.step_and_update_noise(noise_pred, t, approximated_z_tp1, z_t, return_dict=False, update_epsilon_type=self.cfg.update_epsilon_type) | |
| return approximated_z_tp1 | |
| def detach_before_opt(self, z_tp1, t, prompt_embeds, added_cond_kwargs): | |
| z_tp1 = z_tp1.detach() | |
| t = t.detach() | |
| prompt_embeds = prompt_embeds.detach() | |
| return z_tp1, t, prompt_embeds, added_cond_kwargs | |
| def opt_z_tp1_single_step( | |
| self, | |
| z_tp1, | |
| z_t, | |
| t, | |
| prompt_embeds, | |
| added_cond_kwargs, | |
| lr=0.001, | |
| opt_loss_kl_lambda=10.0, | |
| ): | |
| l1_loss = torch.nn.L1Loss(reduction='sum') | |
| mse = torch.nn.MSELoss(reduction='sum') | |
| extra_step_kwargs = {} | |
| self.unet.requires_grad_(False) | |
| z_tp1, t, prompt_embeds, added_cond_kwargs = self.detach_before_opt(z_tp1, t, prompt_embeds, added_cond_kwargs) | |
| z_tp1 = torch.nn.Parameter(z_tp1, requires_grad=True) | |
| optimizer = torch.optim.SGD([z_tp1], lr=lr, momentum=0.9) | |
| optimizer.zero_grad() | |
| self.unet.zero_grad() | |
| latent_model_input = torch.cat([z_tp1] * 2) if self.do_classifier_free_guidance else z_tp1 | |
| latent_model_input = self.scheduler_inference.scale_model_input(latent_model_input, t) | |
| noise_pred = self.unet( | |
| latent_model_input, | |
| t, | |
| encoder_hidden_states=prompt_embeds, | |
| timestep_cond=None, | |
| cross_attention_kwargs=self.cross_attention_kwargs, | |
| added_cond_kwargs=added_cond_kwargs, | |
| return_dict=False, | |
| )[0] | |
| # perform guidance | |
| if self.do_classifier_free_guidance: | |
| noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) | |
| noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) | |
| # # compute the previous noisy sample x_t -> x_t-1 | |
| z_t_hat = self.scheduler_inference.step(noise_pred, t, z_tp1, **extra_step_kwargs, return_dict=False)[0] | |
| direct_loss = 0.5 * mse(z_t_hat, z_t.detach()) + 0.5 * l1_loss(z_t_hat, z_t.detach()) | |
| kl_loss = torch.tensor([0]).to(z_t.device) | |
| loss = 1.0 * direct_loss + opt_loss_kl_lambda * kl_loss | |
| loss.backward() | |
| optimizer.step() | |
| print(f't: {t}\t total_loss: {format(loss.item(), ".3f")}\t\t direct_loss: {format(direct_loss.item(), ".3f")}\t\t kl_loss: {format(kl_loss.item(), ".3f")}') | |
| return z_tp1.detach() | |
| def optimize_z_tp1( | |
| self, | |
| z_tp1, | |
| z_t, | |
| t, | |
| prompt_embeds, | |
| added_cond_kwargs, | |
| nom_opt_iters=1, | |
| lr=0.001, | |
| opt_loss_kl_lambda=10.0, | |
| ): | |
| l1_loss = torch.nn.L1Loss(reduction='sum') | |
| mse = torch.nn.MSELoss(reduction='sum') | |
| extra_step_kwargs = {} | |
| self.unet.requires_grad_(False) | |
| z_tp1, t, prompt_embeds, added_cond_kwargs = self.detach_before_opt(z_tp1, t, prompt_embeds, added_cond_kwargs) | |
| z_tp1 = torch.nn.Parameter(z_tp1, requires_grad=True) | |
| optimizer = torch.optim.SGD([z_tp1], lr=lr, momentum=0.9) | |
| lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor = 0.5, verbose=True, patience=5, cooldown=3) | |
| max_loss = 99999999999999 | |
| z_tp1_forward = self.scheduler.add_noise(self.z_0, self.noise, t.view((1))).detach() | |
| z_tp1_best = None | |
| for i in range(nom_opt_iters): | |
| optimizer.zero_grad() | |
| self.unet.zero_grad() | |
| latent_model_input = torch.cat([z_tp1] * 2) if self.do_classifier_free_guidance else z_tp1 | |
| latent_model_input = self.scheduler_inference.scale_model_input(latent_model_input, t) | |
| noise_pred = self.unet( | |
| latent_model_input, | |
| t, | |
| encoder_hidden_states=prompt_embeds, | |
| timestep_cond=None, | |
| cross_attention_kwargs=self.cross_attention_kwargs, | |
| added_cond_kwargs=added_cond_kwargs, | |
| return_dict=False, | |
| )[0] | |
| # perform guidance | |
| if self.do_classifier_free_guidance: | |
| noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) | |
| noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) | |
| # # compute the previous noisy sample x_t -> x_t-1 | |
| z_t_hat = self.scheduler_inference.step(noise_pred, t, z_tp1, **extra_step_kwargs, return_dict=False)[0] | |
| direct_loss = 0.5 * mse(z_t_hat, z_t.detach()) + 0.5 * l1_loss(z_t_hat, z_t.detach()) | |
| kl_loss = self.patchify_latents_kl_divergence(z_tp1, z_tp1_forward) | |
| loss = 1.0 * direct_loss + opt_loss_kl_lambda * kl_loss | |
| loss.backward() | |
| best = False | |
| if loss < max_loss: | |
| max_loss = loss | |
| z_tp1_best = torch.clone(z_tp1) | |
| best = True | |
| lr_scheduler.step(loss) | |
| if optimizer.param_groups[0]['lr'] < 9e-06: | |
| break | |
| optimizer.step() | |
| print(f't: {t}\t\t iter: {i}\t total_loss: {format(loss.item(), ".3f")}\t\t direct_loss: {format(direct_loss.item(), ".3f")}\t\t kl_loss: {format(kl_loss.item(), ".3f")}\t\t best: {best}') | |
| if z_tp1_best is not None: | |
| z_tp1 = z_tp1_best | |
| self.prev_z4 = torch.clone(z_tp1) | |
| return z_tp1.detach() | |
| def opt_inv(self, | |
| z_t, | |
| t, | |
| prompt_embeds, | |
| added_cond_kwargs, | |
| prev_timestep, | |
| nom_opt_iters=1, | |
| lr=0.001, | |
| opt_none_inference_steps=False, | |
| opt_loss_kl_lambda=10.0, | |
| num_aprox_steps=100): | |
| z_tp1 = self.inversion_step(z_t, t, prompt_embeds, added_cond_kwargs, num_aprox_steps=num_aprox_steps) | |
| if t in self.scheduler_inference.timesteps: | |
| z_tp1 = self.optimize_z_tp1(z_tp1, z_t, t, prompt_embeds, added_cond_kwargs, nom_opt_iters=nom_opt_iters, lr=lr, opt_loss_kl_lambda=opt_loss_kl_lambda) | |
| return z_tp1 | |
| def latent2image(self, latents): | |
| needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast | |
| if needs_upcasting: | |
| self.upcast_vae() | |
| latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype) | |
| image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] | |
| # cast back to fp16 if needed | |
| # if needs_upcasting: | |
| # self.vae.to(dtype=torch.float16) | |
| return image | |
| def patchify_latents_kl_divergence(self, x0, x1): | |
| # devide x0 and x1 into patches (4x64x64) -> (4x4x4) | |
| PATCH_SIZE = 4 | |
| NUM_CHANNELS = 4 | |
| def patchify_tensor(input_tensor): | |
| patches = input_tensor.unfold(1, PATCH_SIZE, PATCH_SIZE).unfold(2, PATCH_SIZE, PATCH_SIZE).unfold(3, PATCH_SIZE, PATCH_SIZE) | |
| patches = patches.contiguous().view(-1, NUM_CHANNELS, PATCH_SIZE, PATCH_SIZE) | |
| return patches | |
| x0 = patchify_tensor(x0) | |
| x1 = patchify_tensor(x1) | |
| kl = self.latents_kl_divergence(x0, x1).sum() | |
| # for i in range(x0.shape[0]): | |
| # kl += self.latents_kl_divergence(x0[i], x1[i]) | |
| return kl | |
| def latents_kl_divergence(self, x0, x1): | |
| EPSILON = 1e-6 | |
| #{\displaystyle D_{\text{KL}}\left({\mathcal {N}}_{0}\parallel {\mathcal {N}}_{1}\right)={\frac {1}{2}}\left(\operatorname {tr} \left(\Sigma _{1}^{-1}\Sigma _{0}\right)-k+\left(\mu _{1}-\mu _{0}\right)^{\mathsf {T}}\Sigma _{1}^{-1}\left(\mu _{1}-\mu _{0}\right)+\ln \left({\frac {\det \Sigma _{1}}{\det \Sigma _{0}}}\right)\right).} | |
| x0 = x0.view(x0.shape[0], x0.shape[1], -1) | |
| x1 = x1.view(x1.shape[0], x1.shape[1], -1) | |
| mu0 = x0.mean(dim=-1) | |
| mu1 = x1.mean(dim=-1) | |
| var0 = x0.var(dim=-1) | |
| var1 = x1.var(dim=-1) | |
| kl = torch.log((var1 + EPSILON) / (var0 + EPSILON)) + (var0 + (mu0 - mu1)**2) / (var1 + EPSILON) - 1 | |
| kl = torch.abs(kl).sum(dim=-1) | |
| # kl = torch.linalg.norm(mu0 - mu1) + torch.linalg.norm(var0 - var1) | |
| # kl *= 1000 | |
| # sigma0 = torch.cov(x0) | |
| # sigma1 = torch.cov(x1) | |
| # inv_sigma1 = torch.inverse(sigma1.to(dtype=torch.float64)).to(dtype=x0.dtype) | |
| # k = x0.shape[1] | |
| # kl = 0.5 * (torch.trace(inv_sigma1 @ sigma0) - k + (mu1 - mu0).T @ inv_sigma1 @ (mu1 - mu0) + torch.log(torch.det(sigma1) / torch.det(sigma0))) | |
| return kl | |
| class SpecifyGradient(torch.autograd.Function): | |
| def forward(ctx, input_tensor, gt_grad): | |
| ctx.save_for_backward(gt_grad) | |
| # dummy loss value | |
| return torch.zeros([1], device=input_tensor.device, dtype=input_tensor.dtype) | |
| def backward(ctx, grad): | |
| gt_grad, = ctx.saved_tensors | |
| batch_size = len(gt_grad) | |
| return gt_grad / batch_size, None | |