# 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):
@staticmethod
@custom_fwd
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)
@staticmethod
@custom_bwd
def backward(ctx, grad):
gt_grad, = ctx.saved_tensors
batch_size = len(gt_grad)
return gt_grad / batch_size, None