save intermediate state score_sde

run.py

+#!/usr/bin/env python3
+import numpy as np
+import PIL
+import functools
+
+import models
+from models import utils as mutils
+from models import ncsnv2
+from models import ncsnpp
+from models import ddpm as ddpm_model
+from models import layerspp
+from models import layers
+from models import normalization
+
+from utils import restore_checkpoint
+
+import sampling
+from sde_lib import VESDE, VPSDE, subVPSDE
+from sampling import (NoneCorrector, 
+                      ReverseDiffusionPredictor, 
+                      LangevinCorrector,
+                      EulerMaruyamaPredictor, 
+                      AncestralSamplingPredictor, 
+                      NonePredictor,
+                      AnnealedLangevinDynamics)
+import datasets
+import torch
+
+
+torch.manual_seed(0)
+
+
+#class NewVESDE(SDE):
+#  def __init__(self, sigma_min=0.01, sigma_max=50, N=1000):
+#    """Construct a Variance Exploding SDE.
+#
+#    Args:
+#      sigma_min: smallest sigma.
+#      sigma_max: largest sigma.
+#      N: number of discretization steps
+#    """
+#    super().__init__(N)
+#    self.sigma_min = sigma_min
+#    self.sigma_max = sigma_max
+#    self.discrete_sigmas = torch.exp(torch.linspace(np.log(self.sigma_min), np.log(self.sigma_max), N))
+#    self.N = N
+#
+#  @property
+#  def T(self):
+#    return 1
+#
+#  def sde(self, x, t):
+#    sigma = self.sigma_min * (self.sigma_max / self.sigma_min) ** t
+#    drift = torch.zeros_like(x)
+#    diffusion = sigma * torch.sqrt(torch.tensor(2 * (np.log(self.sigma_max) - np.log(self.sigma_min)),
+#                                                device=t.device))
+#    return drift, diffusion
+#
+#  def marginal_prob(self, x, t):
+#    std = self.sigma_min * (self.sigma_max / self.sigma_min) ** t
+#    mean = x
+#    return mean, std
+#
+#  def prior_sampling(self, shape):
+#    return torch.randn(*shape) * self.sigma_max
+#
+#  def prior_logp(self, z):
+#    shape = z.shape
+#    N = np.prod(shape[1:])
+#    return -N / 2. * np.log(2 * np.pi * self.sigma_max ** 2) - torch.sum(z ** 2, dim=(1, 2, 3)) / (2 * self.sigma_max ** 2)
+#
+#  def discretize(self, x, t):
+#    """SMLD(NCSN) discretization."""
+#    timestep = (t * (self.N - 1) / self.T).long()
+#    sigma = self.discrete_sigmas.to(t.device)[timestep]
+#    adjacent_sigma = torch.where(timestep == 0, torch.zeros_like(t),
+#                                 self.discrete_sigmas[timestep - 1].to(t.device))
+#    f = torch.zeros_like(x)
+#    G = torch.sqrt(sigma ** 2 - adjacent_sigma ** 2)
+#    return f, G
+
+
+class NewReverseDiffusionPredictor:
+
+  def __init__(self, sde, score_fn, probability_flow=False):
+    super().__init__()
+    self.sde = sde
+    self.probability_flow = probability_flow
+    self.score_fn = score_fn
+
+  def discretize(self, x, t):
+    timestep = (t * (self.sde.N - 1) / self.sde.T).long()
+    sigma = self.sde.discrete_sigmas.to(t.device)[timestep]
+    adjacent_sigma = torch.where(timestep == 0, torch.zeros_like(t),
+                                 self.sde.discrete_sigmas[timestep - 1].to(t.device))
+    f = torch.zeros_like(x)
+    G = torch.sqrt(sigma ** 2 - adjacent_sigma ** 2)
+
+    labels = self.sde.marginal_prob(torch.zeros_like(x), t)[1]
+    result = self.score_fn(x, labels)
+
+    rev_f = f - G[:, None, None, None] ** 2 * result * (0.5 if self.probability_flow else 1.)
+    rev_G = torch.zeros_like(G) if self.probability_flow else G
+    return rev_f, rev_G
+
+  def update_fn(self, x, t):
+    f, G = self.discretize(x, t)
+    z = torch.randn_like(x)
+    x_mean = x - f
+    x = x_mean + G[:, None, None, None] * z
+    return x, x_mean
+
+
+class NewLangevinCorrector:
+
+  def __init__(self, sde, score_fn, snr, n_steps):
+    super().__init__()
+    self.sde = sde
+    self.score_fn = score_fn
+    self.snr = snr
+    self.n_steps = n_steps
+
+  def update_fn(self, x, t):
+    sde = self.sde
+    score_fn = self.score_fn
+    n_steps = self.n_steps
+    target_snr = self.snr
+    if isinstance(sde, VPSDE) or isinstance(sde, subVPSDE):
+      timestep = (t * (sde.N - 1) / sde.T).long()
+      alpha = sde.alphas.to(t.device)[timestep]
+    else:
+      alpha = torch.ones_like(t)
+
+    for i in range(n_steps):
+      labels = sde.marginal_prob(torch.zeros_like(x), t)[1]
+      grad = score_fn(x, labels)
+      noise = torch.randn_like(x)
+      grad_norm = torch.norm(grad.reshape(grad.shape[0], -1), dim=-1).mean()
+      noise_norm = torch.norm(noise.reshape(noise.shape[0], -1), dim=-1).mean()
+      step_size = (target_snr * noise_norm / grad_norm) ** 2 * 2 * alpha
+      x_mean = x + step_size[:, None, None, None] * grad
+      x = x_mean + torch.sqrt(step_size * 2)[:, None, None, None] * noise
+
+    return x, x_mean
+
+
+
+def save_image(x):
+#    image_processed = x.cpu().permute(0, 2, 3, 1)
+#    image_processed = (image_processed + 1.0) * 127.5
+#    image_processed = image_processed.numpy().astype(np.uint8)
+    image_processed = np.clip(x.permute(0, 2, 3, 1).cpu().numpy() * 255, 0, 255).astype(np.uint8)
+    image_pil = PIL.Image.fromarray(image_processed[0])
+
+    # 6. save image
+    image_pil.save("../images/hey.png")
+
+
+#x = np.load("cifar10.npy")
+#
+#save_image(x)
+# @title Load the score-based model
+sde = 'VESDE' #@param ['VESDE', 'VPSDE', 'subVPSDE'] {"type": "string"}
+if sde.lower() == 'vesde':
+  from configs.ve import cifar10_ncsnpp_continuous as configs
+  ckpt_filename = "exp/ve/cifar10_ncsnpp_continuous/checkpoint_24.pth"
+#  from configs.ve import ffhq_ncsnpp_continuous as configs
+#  ckpt_filename = "exp/ve/ffhq_1024_ncsnpp_continuous/checkpoint_60.pth"
+  config = configs.get_config()  
+  config.model.num_scales = 1000
+  sde = VESDE(sigma_min=config.model.sigma_min, sigma_max=config.model.sigma_max, N=config.model.num_scales)
+  sampling_eps = 1e-5
+elif sde.lower() == 'vpsde':
+  from configs.vp import cifar10_ddpmpp_continuous as configs  
+  ckpt_filename = "exp/vp/cifar10_ddpmpp_continuous/checkpoint_8.pth"
+  config = configs.get_config()
+  sde = VPSDE(beta_min=config.model.beta_min, beta_max=config.model.beta_max, N=config.model.num_scales)
+  sampling_eps = 1e-3
+elif sde.lower() == 'subvpsde':
+  from configs.subvp import cifar10_ddpmpp_continuous as configs
+  ckpt_filename = "exp/subvp/cifar10_ddpmpp_continuous/checkpoint_26.pth"
+  config = configs.get_config()
+  sde = subVPSDE(beta_min=config.model.beta_min, beta_max=config.model.beta_max, N=config.model.num_scales)
+  sampling_eps = 1e-3
+
+batch_size = 1 #@param {"type":"integer"}
+config.training.batch_size = batch_size
+config.eval.batch_size = batch_size
+
+random_seed = 0 #@param {"type": "integer"}
+
+score_model = mutils.create_model(config)
+
+loaded_state = torch.load(ckpt_filename)
+score_model.load_state_dict(loaded_state["model"], strict=False)
+
+inverse_scaler = datasets.get_data_inverse_scaler(config)
+predictor = ReverseDiffusionPredictor #@param ["EulerMaruyamaPredictor", "AncestralSamplingPredictor", "ReverseDiffusionPredictor", "None"] {"type": "raw"}
+corrector = LangevinCorrector #@param ["LangevinCorrector", "AnnealedLangevinDynamics", "None"] {"type": "raw"}
+
+def image_grid(x):
+  size = config.data.image_size
+  channels = config.data.num_channels
+  img = x.reshape(-1, size, size, channels)
+  w = int(np.sqrt(img.shape[0]))
+  img = img.reshape((w, w, size, size, channels)).transpose((0, 2, 1, 3, 4)).reshape((w * size, w * size, channels))
+  return img
+
+#@title PC sampling
+img_size = config.data.image_size
+channels = config.data.num_channels
+shape = (batch_size, channels, img_size, img_size)
+probability_flow = False
+snr = 0.16 #@param {"type": "number"}
+n_steps =  1#@param {"type": "integer"}
+
+
+def shared_predictor_update_fn(x, t, sde, model, predictor, probability_flow, continuous):
+  """A wrapper that configures and returns the update function of predictors."""
+  score_fn = mutils.get_score_fn(sde, model, train=False, continuous=continuous)
+  if predictor is None:
+    # Corrector-only sampler
+    predictor_obj = NonePredictor(sde, score_fn, probability_flow)
+  else:
+    predictor_obj = predictor(sde, score_fn, probability_flow)
+  return predictor_obj.update_fn(x, t)
+
+
+def shared_corrector_update_fn(x, t, sde, model, corrector, continuous, snr, n_steps):
+  """A wrapper tha configures and returns the update function of correctors."""
+  score_fn = mutils.get_score_fn(sde, model, train=False, continuous=continuous)
+  if corrector is None:
+    # Predictor-only sampler
+    corrector_obj = NoneCorrector(sde, score_fn, snr, n_steps)
+  else:
+    corrector_obj = corrector(sde, score_fn, snr, n_steps)
+  return corrector_obj.update_fn(x, t)
+
+
+continuous = config.training.continuous
+
+
+predictor_update_fn = functools.partial(shared_predictor_update_fn,
+                                          sde=sde,
+                                          predictor=predictor,
+                                          probability_flow=probability_flow,
+                                          continuous=continuous)
+
+corrector_update_fn = functools.partial(shared_corrector_update_fn,
+                                          sde=sde,
+                                          corrector=corrector,
+                                          continuous=continuous,
+                                          snr=snr,
+                                          n_steps=n_steps)
+
+device = "cuda"
+model = score_model.to(device)
+denoise = False
+
+new_corrector = NewLangevinCorrector(sde=sde, score_fn=model, snr=snr, n_steps=n_steps)
+new_predictor = NewReverseDiffusionPredictor(sde=sde, score_fn=model)
+
+
+with torch.no_grad():
+    # Initial sample
+    x = sde.prior_sampling(shape).to(device)
+    timesteps = torch.linspace(sde.T, sampling_eps, sde.N, device=device)
+
+    for i in range(sde.N):
+        t = timesteps[i]
+        vec_t = torch.ones(shape[0], device=t.device) * t
+#        x, x_mean = corrector_update_fn(x, vec_t, model=model)
+#        x, x_mean = predictor_update_fn(x, vec_t, model=model)
+        x, x_mean = new_corrector.update_fn(x, vec_t)
+        x, x_mean = new_predictor.update_fn(x, vec_t)
+
+    x, n = inverse_scaler(x_mean if denoise else x), sde.N * (n_steps + 1)
+
+
+# for 5
+#assert x.abs().sum().cpu().item() - 106114.90625 < 1e-2, "sum wrong"
+#assert x.abs().mean().cpu().item() - 34.5426139831543 < 1e-4, "mean wrong"
+
+# for 1000
+assert x.abs().sum().cpu().item() - 436.5811 < 1e-2, "sum wrong"
+assert x.abs().mean().cpu().item() - 0.1421 < 1e-4, "mean wrong"
+
+save_image(x)