content: """
  wild mixture of
  https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
  https://github.com/openai/improved-diffusion/blob/e94489283bb876ac1477d5dd7709bbbd2d9902ce/improved_diffusion/gaussian_diffusion.py
  https://github.com/CompVis/taming-transformers
  -- merci
  """
  
  import torch
  import torch.nn as nn
  import numpy as np
  import pytorch_lightning as pl
  from torch.optim.lr_scheduler import LambdaLR
  from einops import rearrange, repeat
  from contextlib import contextmanager, nullcontext
  from functools import partial
  import itertools
  from tqdm import tqdm
  from torchvision.utils import make_grid
  from pytorch_lightning.utilities.rank_zero import rank_zero_only
  from omegaconf import ListConfig
  
  from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config
  from ldm.modules.ema import LitEma
  from ldm.modules.distributions.distributions import normal_kl, DiagonalGaussianDistribution
  from ldm.models.autoencoder import VQModelInterface, IdentityFirstStage, AutoencoderKL
  from ldm.modules.diffusionmodules.util import make_beta_schedule, extract_into_tensor, noise_like
  from ldm.models.diffusion.ddim import DDIMSampler
  from ldm.modules.attention import CrossAttention
  
  
  __conditioning_keys__ = {'concat': 'c_concat',
                           'crossattn': 'c_crossattn',
                           'adm': 'y'}
  
  
  def disabled_train(self, mode=True):
      """Overwrite model.train with this function to make sure train/eval mode
      does not change anymore."""
      return self
  
  
  def uniform_on_device(r1, r2, shape, device):
      return (r1 - r2) * torch.rand(*shape, device=device) + r2
  
  
  class DDPM(pl.LightningModule):
      # classic DDPM with Gaussian diffusion, in image space
      def __init__(self,
                   unet_config,
                   timesteps=1000,
                   beta_schedule="linear",
                   loss_type="l2",
                   ckpt_path=None,
                   ignore_keys=[],
                   load_only_unet=False,
                   monitor="val/loss",
                   use_ema=True,
                   first_stage_key="image",
                   image_size=256,
                   channels=3,
                   log_every_t=100,
                   clip_denoised=True,
                   linear_start=1e-4,
                   linear_end=2e-2,
                   cosine_s=8e-3,
                   given_betas=None,
                   original_elbo_weight=0.,
                   v_posterior=0.,  # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta
                   l_simple_weight=1.,
                   conditioning_key=None,
                   parameterization="eps",  # all assuming fixed variance schedules
                   scheduler_config=None,
                   use_positional_encodings=False,
                   learn_logvar=False,
                   logvar_init=0.,
                   make_it_fit=False,
                   ucg_training=None,
                   ):
          super().__init__()
          assert parameterization in ["eps", "x0"], 'currently only supporting "eps" and "x0"'
          self.parameterization = parameterization
          print(f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode")
          self.cond_stage_model = None
          self.clip_denoised = clip_denoised
          self.log_every_t = log_every_t
          self.first_stage_key = first_stage_key
          self.image_size = image_size  # try conv?
          self.channels = channels
          self.use_positional_encodings = use_positional_encodings
          self.model = DiffusionWrapper(unet_config, conditioning_key)
          count_params(self.model, verbose=True)
          self.use_ema = use_ema
          if self.use_ema:
              self.model_ema = LitEma(self.model)
              print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
  
          self.use_scheduler = scheduler_config is not None
          if self.use_scheduler:
              self.scheduler_config = scheduler_config
  
          self.v_posterior = v_posterior
          self.original_elbo_weight = original_elbo_weight
          self.l_simple_weight = l_simple_weight
  
          if monitor is not None:
              self.monitor = monitor
          self.make_it_fit = make_it_fit
          if ckpt_path is not None:
              self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys, only_model=load_only_unet)
  
          self.register_schedule(given_betas=given_betas, beta_schedule=beta_schedule, timesteps=timesteps,
                                 linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s)
  
          self.loss_type = loss_type
  
          self.learn_logvar = learn_logvar
          self.logvar = torch.full(fill_value=logvar_init, size=(self.num_timesteps,))
          if self.learn_logvar:
              self.logvar = nn.Parameter(self.logvar, requires_grad=True)
  
          self.ucg_training = ucg_training or dict()
          if self.ucg_training:
              self.ucg_prng = np.random.RandomState()
  
      def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
                            linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
          if exists(given_betas):
              betas = given_betas
          else:
              betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end,
                                         cosine_s=cosine_s)
          alphas = 1. - betas
          alphas_cumprod = np.cumprod(alphas, axis=0)
          alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
  
          timesteps, = betas.shape
          self.num_timesteps = int(timesteps)
          self.linear_start = linear_start
          self.linear_end = linear_end
          assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep'
  
          to_torch = partial(torch.tensor, dtype=torch.float32)
  
          self.register_buffer('betas', to_torch(betas))
          self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
          self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
  
          # calculations for diffusion q(x_t | x_{t-1}) and others
          self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
          self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
          self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
          self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
          self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
  
          # calculations for posterior q(x_{t-1} | x_t, x_0)
          posterior_variance = (1 - self.v_posterior) * betas * (1. - alphas_cumprod_prev) / (
                      1. - alphas_cumprod) + self.v_posterior * betas
          # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
          self.register_buffer('posterior_variance', to_torch(posterior_variance))
          # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
          self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
          self.register_buffer('posterior_mean_coef1', to_torch(
              betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
          self.register_buffer('posterior_mean_coef2', to_torch(
              (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
  
          if self.parameterization == "eps":
              lvlb_weights = self.betas ** 2 / (
                          2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod))
          elif self.parameterization == "x0":
              lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * 1 - torch.Tensor(alphas_cumprod))
          else:
              raise NotImplementedError("mu not supported")
          # TODO how to choose this term
          lvlb_weights[0] = lvlb_weights[1]
          self.register_buffer('lvlb_weights', lvlb_weights, persistent=False)
          assert not torch.isnan(self.lvlb_weights).all()
  
      @contextmanager
      def ema_scope(self, context=None):
          if self.use_ema:
              self.model_ema.store(self.model.parameters())
              self.model_ema.copy_to(self.model)
              if context is not None:
                  print(f"{context}: Switched to EMA weights")
          try:
              yield None
          finally:
              if self.use_ema:
                  self.model_ema.restore(self.model.parameters())
                  if context is not None:
                      print(f"{context}: Restored training weights")
  
      @torch.no_grad()
      def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
          sd = torch.load(path, map_location="cpu")
          if "state_dict" in list(sd.keys()):
              sd = sd["state_dict"]
          keys = list(sd.keys())
  
          if self.make_it_fit:
              n_params = len([name for name, _ in
                              itertools.chain(self.named_parameters(),
                                              self.named_buffers())])
              for name, param in tqdm(
                      itertools.chain(self.named_parameters(),
                                      self.named_buffers()),
                      desc="Fitting old weights to new weights",
                      total=n_params
              ):
                  if not name in sd:
                      continue
                  old_shape = sd[name].shape
                  new_shape = param.shape
                  assert len(old_shape)==len(new_shape)
                  if len(new_shape) > 2:
                      # we only modify first two axes
                      assert new_shape[2:] == old_shape[2:]
                  # assumes first axis corresponds to output dim
                  if not new_shape == old_shape:
                      new_param = param.clone()
                      old_param = sd[name]
                      if len(new_shape) == 1:
                          for i in range(new_param.shape[0]):
                              new_param[i] = old_param[i % old_shape[0]]
                      elif len(new_shape) >= 2:
                          for i in range(new_param.shape[0]):
                              for j in range(new_param.shape[1]):
                                  new_param[i, j] = old_param[i % old_shape[0], j % old_shape[1]]
  
                          n_used_old = torch.ones(old_shape[1])
                          for j in range(new_param.shape[1]):
                              n_used_old[j % old_shape[1]] += 1
                          n_used_new = torch.zeros(new_shape[1])
                          for j in range(new_param.shape[1]):
                              n_used_new[j] = n_used_old[j % old_shape[1]]
  
                          n_used_new = n_used_new[None, :]
                          while len(n_used_new.shape) < len(new_shape):
                              n_used_new = n_used_new.unsqueeze(-1)
                          new_param /= n_used_new
  
                      sd[name] = new_param
  
          missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
              sd, strict=False)
          print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
          if len(missing) > 0:
              print(f"Missing Keys: {missing}")
          if len(unexpected) > 0:
              print(f"Unexpected Keys: {unexpected}")
  
      def q_mean_variance(self, x_start, t):
          """
          Get the distribution q(x_t | x_0).
          :param x_start: the [N x C x ...] tensor of noiseless inputs.
          :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
          :return: A tuple (mean, variance, log_variance), all of x_start's shape.
          """
          mean = (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start)
          variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
          log_variance = extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape)
          return mean, variance, log_variance
  
      def predict_start_from_noise(self, x_t, t, noise):
          return (
                  extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
                  extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
          )
  
      def q_posterior(self, x_start, x_t, t):
          posterior_mean = (
                  extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start +
                  extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
          )
          posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape)
          posterior_log_variance_clipped = extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape)
          return posterior_mean, posterior_variance, posterior_log_variance_clipped
  
      def p_mean_variance(self, x, t, clip_denoised: bool):
          model_out = self.model(x, t)
          if self.parameterization == "eps":
              x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
          elif self.parameterization == "x0":
              x_recon = model_out
          if clip_denoised:
              x_recon.clamp_(-1., 1.)
  
          model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
          return model_mean, posterior_variance, posterior_log_variance
  
      @torch.no_grad()
      def p_sample(self, x, t, clip_denoised=True, repeat_noise=False):
          b, *_, device = *x.shape, x.device
          model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised)
          noise = noise_like(x.shape, device, repeat_noise)
          # no noise when t == 0
          nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
          return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
  
      @torch.no_grad()
      def p_sample_loop(self, shape, return_intermediates=False):
          device = self.betas.device
          b = shape[0]
          img = torch.randn(shape, device=device)
          intermediates = [img]
          for i in tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps):
              img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long),
                                  clip_denoised=self.clip_denoised)
              if i % self.log_every_t == 0 or i == self.num_timesteps - 1:
                  intermediates.append(img)
          if return_intermediates:
              return img, intermediates
          return img
  
      @torch.no_grad()
      def sample(self, batch_size=16, return_intermediates=False):
          image_size = self.image_size
          channels = self.channels
          return self.p_sample_loop((batch_size, channels, image_size, image_size),
                                    return_intermediates=return_intermediates)
  
      def q_sample(self, x_start, t, noise=None):
          noise = default(noise, lambda: torch.randn_like(x_start))
          return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
                  extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
  
      def get_loss(self, pred, target, mean=True):
          if self.loss_type == 'l1':
              loss = (target - pred).abs()
              if mean:
                  loss = loss.mean()
          elif self.loss_type == 'l2':
              if mean:
                  loss = torch.nn.functional.mse_loss(target, pred)
              else:
                  loss = torch.nn.functional.mse_loss(target, pred, reduction='none')
          else:
              raise NotImplementedError("unknown loss type '{loss_type}'")
  
          return loss
  
      def p_losses(self, x_start, t, noise=None):
          noise = default(noise, lambda: torch.randn_like(x_start))
          x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
          model_out = self.model(x_noisy, t)
  
          loss_dict = {}
          if self.parameterization == "eps":
              target = noise
          elif self.parameterization == "x0":
              target = x_start
          else:
              raise NotImplementedError(f"Paramterization {self.parameterization} not yet supported")
  
          loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3])
  
          log_prefix = 'train' if self.training else 'val'
  
          loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()})
          loss_simple = loss.mean() * self.l_simple_weight
  
          loss_vlb = (self.lvlb_weights[t] * loss).mean()
          loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb})
  
          loss = loss_simple + self.original_elbo_weight * loss_vlb
  
          loss_dict.update({f'{log_prefix}/loss': loss})
  
          return loss, loss_dict
  
      def forward(self, x, *args, **kwargs):
          # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size
          # assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
          t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
          return self.p_losses(x, t, *args, **kwargs)
  
      def get_input(self, batch, k):
          x = batch[k]
          if len(x.shape) == 3:
              x = x[..., None]
          x = rearrange(x, 'b h w c -> b c h w')
          x = x.to(memory_format=torch.contiguous_format).float()
          return x
  
      def shared_step(self, batch):
          x = self.get_input(batch, self.first_stage_key)
          loss, loss_dict = self(x)
          return loss, loss_dict
  
      def training_step(self, batch, batch_idx):
          for k in self.ucg_training:
              p = self.ucg_training[k]["p"]
              val = self.ucg_training[k]["val"]
              if val is None:
                  val = ""
              for i in range(len(batch[k])):
                  if self.ucg_prng.choice(2, p=[1-p, p]):
                      batch[k][i] = val
  
          loss, loss_dict = self.shared_step(batch)
  
          self.log_dict(loss_dict, prog_bar=True,
                        logger=True, on_step=True, on_epoch=True)
  
          self.log("global_step", self.global_step,
                   prog_bar=True, logger=True, on_step=True, on_epoch=False)
  
          if self.use_scheduler:
              lr = self.optimizers().param_groups[0]['lr']
              self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)
  
          return loss
  
      @torch.no_grad()
      def validation_step(self, batch, batch_idx):
          _, loss_dict_no_ema = self.shared_step(batch)
          with self.ema_scope():
              _, loss_dict_ema = self.shared_step(batch)
              loss_dict_ema = {key + '_ema': loss_dict_ema[key] for key in loss_dict_ema}
          self.log_dict(loss_dict_no_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
          self.log_dict(loss_dict_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
  
      def on_train_batch_end(self, *args, **kwargs):
          if self.use_ema:
              self.model_ema(self.model)
  
      def _get_rows_from_list(self, samples):
          n_imgs_per_row = len(samples)
          denoise_grid = rearrange(samples, 'n b c h w -> b n c h w')
          denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
          denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
          return denoise_grid
  
      @torch.no_grad()
      def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs):
          log = dict()
          x = self.get_input(batch, self.first_stage_key)
          N = min(x.shape[0], N)
          n_row = min(x.shape[0], n_row)
          x = x.to(self.device)[:N]
          log["inputs"] = x
  
          # get diffusion row
          diffusion_row = list()
          x_start = x[:n_row]
  
          for t in range(self.num_timesteps):
              if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
                  t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
                  t = t.to(self.device).long()
                  noise = torch.randn_like(x_start)
                  x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
                  diffusion_row.append(x_noisy)
  
          log["diffusion_row"] = self._get_rows_from_list(diffusion_row)
  
          if sample:
              # get denoise row
              with self.ema_scope("Plotting"):
                  samples, denoise_row = self.sample(batch_size=N, return_intermediates=True)
  
              log["samples"] = samples
              log["denoise_row"] = self._get_rows_from_list(denoise_row)
  
          if return_keys:
              if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
                  return log
              else:
                  return {key: log[key] for key in return_keys}
          return log
  
      def configure_optimizers(self):
          lr = self.learning_rate
          params = list(self.model.parameters())
          if self.learn_logvar:
              params = params + [self.logvar]
          opt = torch.optim.AdamW(params, lr=lr)
          return opt
  
  
  class LatentDiffusion(DDPM):
      """main class"""
      def __init__(self,
                   first_stage_config,
                   cond_stage_config,
                   num_timesteps_cond=None,
                   cond_stage_key="image",
                   cond_stage_trainable=False,
                   concat_mode=True,
                   cond_stage_forward=None,
                   conditioning_key=None,
                   scale_factor=1.0,
                   scale_by_std=False,
                   unet_trainable=True,
                   *args, **kwargs):
          self.num_timesteps_cond = default(num_timesteps_cond, 1)
          self.scale_by_std = scale_by_std
          assert self.num_timesteps_cond <= kwargs['timesteps']
          # for backwards compatibility after implementation of DiffusionWrapper
          if conditioning_key is None:
              conditioning_key = 'concat' if concat_mode else 'crossattn'
          if cond_stage_config == '__is_unconditional__':
              conditioning_key = None
          ckpt_path = kwargs.pop("ckpt_path", None)
          ignore_keys = kwargs.pop("ignore_keys", [])
          super().__init__(conditioning_key=conditioning_key, *args, **kwargs)
          self.concat_mode = concat_mode
          self.cond_stage_trainable = cond_stage_trainable
          self.unet_trainable = unet_trainable
          self.cond_stage_key = cond_stage_key
          try:
              self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1
          except:
              self.num_downs = 0
          if not scale_by_std:
              self.scale_factor = scale_factor
          else:
              self.register_buffer('scale_factor', torch.tensor(scale_factor))
          self.instantiate_first_stage(first_stage_config)
          self.instantiate_cond_stage(cond_stage_config)
          self.cond_stage_forward = cond_stage_forward
  
          # construct linear projection layer for concatenating image CLIP embedding and RT
          self.cc_projection = nn.Linear(772, 768)
          nn.init.eye_(list(self.cc_projection.parameters())[0][:768, :768])
          nn.init.zeros_(list(self.cc_projection.parameters())[1])
          self.cc_projection.requires_grad_(True)
          
          self.clip_denoised = False
          self.bbox_tokenizer = None
  
          self.restarted_from_ckpt = False
          if ckpt_path is not None:
              self.init_from_ckpt(ckpt_path, ignore_keys)
              self.restarted_from_ckpt = True
  
      def make_cond_schedule(self, ):
          self.cond_ids = torch.full(size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long)
          ids = torch.round(torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)).long()
          self.cond_ids[:self.num_timesteps_cond] = ids
  
      @rank_zero_only
      @torch.no_grad()
      def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
          # only for very first batch
          if self.scale_by_std and self.current_epoch == 0 and self.global_step == 0 and batch_idx == 0 and not self.restarted_from_ckpt:
              assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously'
              # set rescale weight to 1./std of encodings
              print("### USING STD-RESCALING ###")
              x = super().get_input(batch, self.first_stage_key)
              x = x.to(self.device)
              encoder_posterior = self.encode_first_stage(x)
              z = self.get_first_stage_encoding(encoder_posterior).detach()
              del self.scale_factor
              self.register_buffer('scale_factor', 1. / z.flatten().std())
              print(f"setting self.scale_factor to {self.scale_factor}")
              print("### USING STD-RESCALING ###")
  
      def register_schedule(self,
                            given_betas=None, beta_schedule="linear", timesteps=1000,
                            linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
          super().register_schedule(given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s)
  
          self.shorten_cond_schedule = self.num_timesteps_cond > 1
          if self.shorten_cond_schedule:
              self.make_cond_schedule()
  
      def instantiate_first_stage(self, config):
          model = instantiate_from_config(config)
          self.first_stage_model = model.eval()
          self.first_stage_model.train = disabled_train
          for param in self.first_stage_model.parameters():
              param.requires_grad = False
  
      def instantiate_cond_stage(self, config):
          if not self.cond_stage_trainable:
              if config == "__is_first_stage__":
                  print("Using first stage also as cond stage.")
                  self.cond_stage_model = self.first_stage_model
              elif config == "__is_unconditional__":
                  print(f"Training {self.__class__.__name__} as an unconditional model.")
                  self.cond_stage_model = None
                  # self.be_unconditional = True
              else:
                  model = instantiate_from_config(config)
                  self.cond_stage_model = model.eval()
                  self.cond_stage_model.train = disabled_train
                  for param in self.cond_stage_model.parameters():
                      param.requires_grad = False
          else:
              assert config != '__is_first_stage__'
              assert config != '__is_unconditional__'
              model = instantiate_from_config(config)
              self.cond_stage_model = model
  
      def _get_denoise_row_from_list(self, samples, desc='', force_no_decoder_quantization=False):
          denoise_row = []
          for zd in tqdm(samples, desc=desc):
              denoise_row.append(self.decode_first_stage(zd.to(self.device),
                                                              force_not_quantize=force_no_decoder_quantization))
          n_imgs_per_row = len(denoise_row)
          denoise_row = torch.stack(denoise_row)  # n_log_step, n_row, C, H, W
          denoise_grid = rearrange(denoise_row, 'n b c h w -> b n c h w')
          denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
          denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
          return denoise_grid
  
      def get_first_stage_encoding(self, encoder_posterior):
          if isinstance(encoder_posterior, DiagonalGaussianDistribution):
              z = encoder_posterior.sample()
          elif isinstance(encoder_posterior, torch.Tensor):
              z = encoder_posterior
          else:
              raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented")
          return self.scale_factor * z
  
      def get_learned_conditioning(self, c):
          if self.cond_stage_forward is None:
              if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode):
                  c = self.cond_stage_model.encode(c)
                  if isinstance(c, DiagonalGaussianDistribution):
                      c = c.mode()
              else:
                  c = self.cond_stage_model(c)
          else:
              assert hasattr(self.cond_stage_model, self.cond_stage_forward)
              c = getattr(self.cond_stage_model, self.cond_stage_forward)(c)
          return c
  
      def meshgrid(self, h, w):
          y = torch.arange(0, h).view(h, 1, 1).repeat(1, w, 1)
          x = torch.arange(0, w).view(1, w, 1).repeat(h, 1, 1)
  
          arr = torch.cat([y, x], dim=-1)
          return arr
  
      def delta_border(self, h, w):
          """
          :param h: height
          :param w: width
          :return: normalized distance to image border,
           wtith min distance = 0 at border and max dist = 0.5 at image center
          """
          lower_right_corner = torch.tensor([h - 1, w - 1]).view(1, 1, 2)
          arr = self.meshgrid(h, w) / lower_right_corner
          dist_left_up = torch.min(arr, dim=-1, keepdims=True)[0]
          dist_right_down = torch.min(1 - arr, dim=-1, keepdims=True)[0]
          edge_dist = torch.min(torch.cat([dist_left_up, dist_right_down], dim=-1), dim=-1)[0]
          return edge_dist
  
      def get_weighting(self, h, w, Ly, Lx, device):
          weighting = self.delta_border(h, w)
          weighting = torch.clip(weighting, self.split_input_params["clip_min_weight"],
                                 self.split_input_params["clip_max_weight"], )
          weighting = weighting.view(1, h * w, 1).repeat(1, 1, Ly * Lx).to(device)
  
          if self.split_input_params["tie_braker"]:
              L_weighting = self.delta_border(Ly, Lx)
              L_weighting = torch.clip(L_weighting,
                                       self.split_input_params["clip_min_tie_weight"],
                                       self.split_input_params["clip_max_tie_weight"])
  
              L_weighting = L_weighting.view(1, 1, Ly * Lx).to(device)
              weighting = weighting * L_weighting
          return weighting
  
      def get_fold_unfold(self, x, kernel_size, stride, uf=1, df=1):  # todo load once not every time, shorten code
          """
          :param x: img of size (bs, c, h, w)
          :return: n img crops of size (n, bs, c, kernel_size[0], kernel_size[1])
          """
          bs, nc, h, w = x.shape
  
          # number of crops in image
          Ly = (h - kernel_size[0]) // stride[0] + 1
          Lx = (w - kernel_size[1]) // stride[1] + 1
  
          if uf == 1 and df == 1:
              fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
              unfold = torch.nn.Unfold(**fold_params)
  
              fold = torch.nn.Fold(output_size=x.shape[2:], **fold_params)
  
              weighting = self.get_weighting(kernel_size[0], kernel_size[1], Ly, Lx, x.device).to(x.dtype)
              normalization = fold(weighting).view(1, 1, h, w)  # normalizes the overlap
              weighting = weighting.view((1, 1, kernel_size[0], kernel_size[1], Ly * Lx))
  
          elif uf > 1 and df == 1:
              fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
              unfold = torch.nn.Unfold(**fold_params)
  
              fold_params2 = dict(kernel_size=(kernel_size[0] * uf, kernel_size[0] * uf),
                                  dilation=1, padding=0,
                                  stride=(stride[0] * uf, stride[1] * uf))
              fold = torch.nn.Fold(output_size=(x.shape[2] * uf, x.shape[3] * uf), **fold_params2)
  
              weighting = self.get_weighting(kernel_size[0] * uf, kernel_size[1] * uf, Ly, Lx, x.device).to(x.dtype)
              normalization = fold(weighting).view(1, 1, h * uf, w * uf)  # normalizes the overlap
              weighting = weighting.view((1, 1, kernel_size[0] * uf, kernel_size[1] * uf, Ly * Lx))
  
          elif df > 1 and uf == 1:
              fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
              unfold = torch.nn.Unfold(**fold_params)
  
              fold_params2 = dict(kernel_size=(kernel_size[0] // df, kernel_size[0] // df),
                                  dilation=1, padding=0,
                                  stride=(stride[0] // df, stride[1] // df))
              fold = torch.nn.Fold(output_size=(x.shape[2] // df, x.shape[3] // df), **fold_params2)
  
              weighting = self.get_weighting(kernel_size[0] // df, kernel_size[1] // df, Ly, Lx, x.device).to(x.dtype)
              normalization = fold(weighting).view(1, 1, h // df, w // df)  # normalizes the overlap
              weighting = weighting.view((1, 1, kernel_size[0] // df, kernel_size[1] // df, Ly * Lx))
  
          else:
              raise NotImplementedError
  
          return fold, unfold, normalization, weighting
  
      
      @torch.no_grad()
  