Loss and Likelihood Based Membership Inference of Diffusion Models
import numpy as np

        
def get_metrics(label, score, fixed_fpr=0.001):
    """
    Compute TPR at FPR
    """
    from sklearn.metrics import roc_curve
    
    fpr, tpr, _ = roc_curve(label, score)

    
    tpr_at_low_fpr = tpr[np.where(fpr <= fixed_fpr)[0][-1]]
    
    
    return tpr_at_low_fpr


def plot_all_steps(prediction,save_path=None):
    '''
    For loss based attacks.
    TPR at low FPR for all diffusion steps.
    '''
    from sklearn.metrics import roc_curve
    import matplotlib.pyplot as plt
    import seaborn as sns
    sns.set_style("darkgrid")
    palette_tab = plt.get_cmap('tab10')
    color_alfa = 1
    
    x_idx = np.arange(0,1000,1)

    import matplotlib.pyplot as plt
    import matplotlib
    matplotlib.rc('font', size=15)        

    plt.plot(x_idx, prediction[:,0], linewidth=2, label='TPR@10%FPR',color=palette_tab(0))
    plt.plot(x_idx, prediction[:,1], linewidth=2, label='TPR@1%FPR',color=palette_tab(1))
    plt.plot(x_idx, prediction[:,2], linewidth=2, label='TPR@0.1%FPR',color=palette_tab(2))
    plt.plot(x_idx, prediction[:,3], linewidth=2, label='TPR@0.01%FPR',color=palette_tab(4))

    

    plt.ylim([0, 1.05])
    # axis labels
    plt.xlabel('Diffusion Steps',fontsize=20)
    plt.ylabel('True Positive Rate',fontsize=20)
    plt.xticks(fontsize=20)
    plt.yticks(fontsize=20)
    # show the legend
    plt.legend(fontsize=14)
    # show the plot
    if save_path != None:
        plt.savefig(save_path, bbox_inches = 'tight', pad_inches = 0.01,dpi=300) 
    plt.show()
    

def plot_one_step(labels, predictions,save_path=None):
    '''
    For loss based attacks.
    ROC curve for one step.
    '''
    from sklearn.metrics import roc_curve
    import seaborn as sns
    sns.set_style("darkgrid")
    
    fpr_0, tpr_0, _ = roc_curve(labels, predictions[0])
    fpr_1, tpr_1, _ = roc_curve(labels, predictions[1])
    fpr_2, tpr_2, _ = roc_curve(labels, predictions[2])
    fpr_3, tpr_3, _ = roc_curve(labels, predictions[3])
    fpr_4, tpr_4, _ = roc_curve(labels, predictions[4])
    fpr_5, tpr_5, _ = roc_curve(labels, predictions[5])
    

    # generate a no skill prediction (majority class)
    ns_probs = [0 for _ in range(len(labels))]
    ns_fpr, ns_tpr, _ = roc_curve(labels, ns_probs)

    import matplotlib.pyplot as plt
    import matplotlib
    matplotlib.rc('font', size=15)    
    plt.plot(ns_fpr, ns_tpr, linestyle='--', label='',c = 'grey')

    plt.plot(fpr_0, tpr_0, linewidth=2, label='t=0')
    plt.plot(fpr_1, tpr_1, linewidth=3, label='t=200')
    plt.plot(fpr_2, tpr_2, linewidth=2, label='t=500')
    plt.plot(fpr_3, tpr_3, linewidth=2, label='t=600')
    plt.plot(fpr_4, tpr_4, linewidth=2, label='t=800')
    plt.plot(fpr_5, tpr_5, linewidth=2, label='t=999')
    
    plt.yscale('log')
    plt.xscale('log')
    plt.xlim([1e-5, 1])
    plt.ylim([1e-5, 2.0])
    # axis labels
    plt.xlabel('False Positive Rate',fontsize=20)
    plt.ylabel('True Positive Rate',fontsize=20)
    plt.xticks(fontsize=20)
    plt.yticks(fontsize=20)
    # show the legend
    plt.legend(fontsize=14)
    if save_path!=None:
        plt.savefig(save_path,bbox_inches = 'tight',pad_inches = 0.03,dpi=300)          
    # show the plot
    plt.show()


def plot_TPR_FPR(labels, predictions,save_path=None):
    '''
    For likelihood based attacks.
    ROC curve.
    '''    
    from sklearn.metrics import roc_curve
    import seaborn as sns
    sns.set_style("darkgrid")
    
    fpr_0, tpr_0, _ = roc_curve(labels, predictions)

    

    # generate a no skill prediction (majority class)
    ns_probs = [0 for _ in range(len(labels))]
    ns_fpr, ns_tpr, _ = roc_curve(labels, ns_probs)

    import matplotlib.pyplot as plt

    plt.plot(ns_fpr, ns_tpr, linestyle='--', label='',c = 'grey')

    plt.plot(fpr_0, tpr_0, linewidth=2, label='Likelihood based attacks')

    
    plt.yscale('log')
    plt.xscale('log')
    plt.xlim([1e-5, 1])
    plt.ylim([1e-5, 1])
    # axis labels
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    # show the legend
    plt.legend()
    if save_path!=None:
        plt.savefig(save_path,bbox_inches = 'tight',pad_inches = 0)       
    # show the plot
    plt.show()

attacking_diffusion_models/compute_metrics.py
HailongHuPri/MIDM
import ml_collections
import torch


def get_default_configs():
  config = ml_collections.ConfigDict()
  # training
  config.training = training = ml_collections.ConfigDict()
  config.training.batch_size = 32
  training.n_iters = 500001
  training.snapshot_freq = 50000
  training.log_freq = 50
  training.eval_freq = 100
  ## store additional checkpoints for preemption in cloud computing environments
  training.snapshot_freq_for_preemption = 10000
  ## produce samples at each snapshot.
  training.snapshot_sampling = True
  training.likelihood_weighting = False
  training.continuous = True
  training.reduce_mean = False

  # sampling
  config.sampling = sampling = ml_collections.ConfigDict()
  sampling.n_steps_each = 1
  sampling.noise_removal = True
  sampling.probability_flow = False
  sampling.snr = 0.17

  # evaluation
  config.eval = evaluate = ml_collections.ConfigDict()
  evaluate.begin_ckpt = 10
  evaluate.end_ckpt = 10
  evaluate.batch_size = 10
  evaluate.enable_sampling = True
  evaluate.num_samples = 10
  evaluate.enable_loss = False
  evaluate.enable_bpd = False
  evaluate.bpd_dataset = 'train'

  # data
  config.data = data = ml_collections.ConfigDict()
  data.dataset = 'FFHQ'
  data.image_size = 64
  data.random_flip = True
  data.uniform_dequantization = False
  data.centered = False
  data.num_channels = 3

  # model
  config.model = model = ml_collections.ConfigDict()
  model.sigma_max = 90.
  model.sigma_min = 0.01
  model.num_scales = 1000
  model.beta_min = 0.1
  model.beta_max = 20.
  model.dropout = 0.1
  model.embedding_type = 'fourier'

  # optimization
  config.optim = optim = ml_collections.ConfigDict()
  optim.weight_decay = 0
  optim.optimizer = 'Adam'
  optim.lr = 2e-4
  optim.beta1 = 0.9
  optim.eps = 1e-8
  optim.warmup = 5000
  optim.grad_clip = 1.

  config.seed = 42
  config.device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

  return config
attacking_diffusion_models/configs/default_ffhq_configs.py
HailongHuPri/MIDM
import ml_collections
import torch


def get_default_configs():
  config = ml_collections.ConfigDict()
  # training
  config.training = training = ml_collections.ConfigDict()
  config.training.batch_size = 32
  training.n_iters = 500001
  training.snapshot_freq = 50000
  training.log_freq = 50
  training.eval_freq = 100
  ## store additional checkpoints for preemption in cloud computing environments
  training.snapshot_freq_for_preemption = 10000
  ## produce samples at each snapshot.
  training.snapshot_sampling = True
  training.likelihood_weighting = False
  training.continuous = True
  training.reduce_mean = False

  # sampling
  config.sampling = sampling = ml_collections.ConfigDict()
  sampling.n_steps_each = 1
  sampling.noise_removal = True
  sampling.probability_flow = False
  sampling.snr = 0.16

  # evaluation
  config.eval = evaluate = ml_collections.ConfigDict()
  evaluate.begin_ckpt = 9
  evaluate.end_ckpt = 26
  evaluate.batch_size = 1024
  evaluate.enable_sampling = False
  evaluate.num_samples = 50000
  evaluate.enable_loss = True
  evaluate.enable_bpd = False
  evaluate.bpd_dataset = 'test'

  # data
  config.data = data = ml_collections.ConfigDict()
  data.dataset = 'FFHQ'
  data.image_size = 64
  data.random_flip = True
  data.uniform_dequantization = False
  data.centered = False
  data.num_channels = 3


  # model
  config.model = model = ml_collections.ConfigDict()
  model.sigma_min = 0.01
  model.sigma_max = 50
  model.num_scales = 1000
  model.beta_min = 0.1
  model.beta_max = 20.
  model.dropout = 0.1
  model.embedding_type = 'fourier'

  # optimization
  config.optim = optim = ml_collections.ConfigDict()
  optim.weight_decay = 0
  optim.optimizer = 'Adam'
  optim.lr = 2e-4
  optim.beta1 = 0.9
  optim.eps = 1e-8
  optim.warmup = 5000
  optim.grad_clip = 1.

  config.seed = 42
  config.device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

  return config
attacking_diffusion_models/configs/default_ffhq_configs_vesde.py
HailongHuPri/MIDM
# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Training VE continuous on FFHQ"""

from configs.default_ffhq_configs_vesde import get_default_configs


def get_config():
  config = get_default_configs()
  # training
  training = config.training
  training.sde = 'vesde'
  training.continuous = True
  training.reduce_mean = True

  # sampling
  sampling = config.sampling
  sampling.method = 'pc'
  sampling.predictor = 'euler_maruyama'
  sampling.corrector = 'none'

  # data
  data = config.data
  data.dataset = 'FFHQ'
  data.centered = False
  data.tfrecords_path = 'Set the data path in your environment (*-r06.tfrecords).'
  data.image_size = 64


  # model
  model = config.model
  model.name = 'ncsnpp'
  model.scale_by_sigma = False
  model.ema_rate = 0.9999
  model.normalization = 'GroupNorm'
  model.nonlinearity = 'swish'
  model.nf = 128
  model.ch_mult = (1, 2, 2, 2)
  model.num_res_blocks = 4
  model.attn_resolutions = (16,)
  model.resamp_with_conv = True
  model.conditional = True
  model.fir = False
  model.fir_kernel = [1, 3, 3, 1]
  model.skip_rescale = True
  model.resblock_type = 'biggan'
  model.progressive = 'none'
  model.progressive_input = 'none'
  model.progressive_combine = 'sum'
  model.attention_type = 'ddpm'
  model.init_scale = 0.
  model.embedding_type = 'positional'
  model.fourier_scale = 16
  model.conv_size = 3
  



  return config

attacking_diffusion_models/configs/ve/ffhq_ve_cont.py
HailongHuPri/MIDM
# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Training SMLD (VE discrete) on FFHQ"""

from configs.default_ffhq_configs import get_default_configs


def get_config():
  config = get_default_configs()
  # training
  training = config.training
  training.sde = 'vesde'
  training.continuous = False

  # sampling
  sampling = config.sampling
  sampling.method = 'pc'
  sampling.predictor = 'reverse_diffusion'
  sampling.corrector = 'langevin'

  # model
  model = config.model
  model.name = 'ddpm'
  model.scale_by_sigma = True
  model.ema_rate = 0.999
  model.normalization = 'GroupNorm'
  model.nonlinearity = 'swish'
  model.nf = 128
  model.ch_mult = (1, 2, 2, 2)
  model.num_res_blocks = 2
  model.attn_resolutions = (16,)
  model.resamp_with_conv = True
  model.conditional = True
  model.conv_size = 3


  # data
  data = config.data
  data.dataset = 'FFHQ'
  data.image_size = 64
  data.centered = False
  data.tfrecords_path = 'Set the data path in your environment (*-r06.tfrecords).'


  
  return config



attacking_diffusion_models/configs/ve/ffhq_ve_discr.py
HailongHuPri/MIDM
# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Training DDPM (VP discrete) on FFHQ"""

from configs.default_ffhq_configs import get_default_configs


def get_config():
  config = get_default_configs()

  # training
  training = config.training
  training.sde = 'vpsde'
  training.continuous = False
  training.reduce_mean = True

  # sampling
  sampling = config.sampling
  sampling.method = 'pc'
  sampling.predictor = 'ancestral_sampling'
  sampling.corrector = 'none'

  # data
  data = config.data
  data.dataset = 'FFHQ'
  data.centered = True
  data.tfrecords_path = 'Set the data path in your environment (*-r06.tfrecords).'
  data.image_size = 64

  # model
  model = config.model
  model.name = 'ddpm'
  model.scale_by_sigma = False
  model.num_scales = 1000
  model.ema_rate = 0.9999
  model.normalization = 'GroupNorm'
  model.nonlinearity = 'swish'
  model.nf = 128
  model.ch_mult = (1, 2, 2, 2)
  model.num_res_blocks = 2
  model.attn_resolutions = (16,)
  model.resamp_with_conv = True
  model.conditional = True

  # optim
  optim = config.optim
  optim.lr = 2e-5
  



  return config

attacking_diffusion_models/configs/vp/ddpm/ffhq_ddpm.py
HailongHuPri/MIDM
# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Training VP continuous on FFHQ"""

from configs.default_ffhq_configs import get_default_configs


def get_config():
  config = get_default_configs()
  # training
  training = config.training
  training.sde = 'vpsde'
  training.continuous = True
  training.reduce_mean = True

  # sampling
  sampling = config.sampling
  sampling.method = 'pc'
  sampling.predictor = 'euler_maruyama'
  sampling.corrector = 'none'

  # data
  data = config.data
  data.dataset = 'FFHQ'
  data.centered = True
  data.tfrecords_path = 'Set the data path in your environment (*-r06.tfrecords).'
  data.image_size = 64


  # model
  model = config.model
  model.name = 'ncsnpp'
  model.scale_by_sigma = False
  model.ema_rate = 0.9999
  model.normalization = 'GroupNorm'
  model.nonlinearity = 'swish'
  model.nf = 128
  model.ch_mult = (1, 2, 2, 2)
  model.num_res_blocks = 4
  model.attn_resolutions = (16,)
  model.resamp_with_conv = True
  model.conditional = True
  model.fir = False
  model.fir_kernel = [1, 3, 3, 1]
  model.skip_rescale = True
  model.resblock_type = 'biggan'
  model.progressive = 'none'
  model.progressive_input = 'none'
  model.progressive_combine = 'sum'
  model.attention_type = 'ddpm'
  model.init_scale = 0.
  model.embedding_type = 'positional'
  model.fourier_scale = 16
  model.conv_size = 3
  



  return config

attacking_diffusion_models/configs/vp/ffhq_vp_cont.py
HailongHuPri/MIDM
# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# pylint: skip-file
# pytype: skip-file
"""Various sampling methods."""

import torch
import numpy as np
from scipy import integrate
from models import utils as mutils


def get_div_fn(fn):
  """Create the divergence function of `fn` using the Hutchinson-Skilling trace estimator."""

  def div_fn(x, t, eps):
    with torch.enable_grad():
      x.requires_grad_(True)
      fn_eps = torch.sum(fn(x, t) * eps)
      grad_fn_eps = torch.autograd.grad(fn_eps, x)[0]
    x.requires_grad_(False)
    return torch.sum(grad_fn_eps * eps, dim=tuple(range(1, len(x.shape))))

  return div_fn


def get_likelihood_fn(sde, inverse_scaler, hutchinson_type='Rademacher',
                      rtol=1e-5, atol=1e-5, method='RK45', eps=1e-5, continuous=True):
  """Create a function to compute the unbiased log-likelihood estimate of a given data point.

  Args:
    sde: A `sde_lib.SDE` object that represents the forward SDE.
    inverse_scaler: The inverse data normalizer.
    hutchinson_type: "Rademacher" or "Gaussian". The type of noise for Hutchinson-Skilling trace estimator.
    rtol: A `float` number. The relative tolerance level of the black-box ODE solver.
    atol: A `float` number. The absolute tolerance level of the black-box ODE solver.
    method: A `str`. The algorithm for the black-box ODE solver.
      See documentation for `scipy.integrate.solve_ivp`.
    eps: A `float` number. The probability flow ODE is integrated to `eps` for numerical stability.

  Returns:
    A function that a batch of data points and returns the log-likelihoods in bits/dim,
      the latent code, and the number of function evaluations cost by computation.
  """

  def drift_fn(model, x, t):
    """The drift function of the reverse-time SDE."""
    score_fn = mutils.get_score_fn(sde, model, train=False, continuous=continuous)
    # Probability flow ODE is a special case of Reverse SDE
    rsde = sde.reverse(score_fn, probability_flow=True)
    return rsde.sde(x, t)[0]

  def div_fn(model, x, t, noise):
    return get_div_fn(lambda xx, tt: drift_fn(model, xx, tt))(x, t, noise)

  def likelihood_fn(model, data):
    """Compute an unbiased estimate to the log-likelihood in bits/dim.

    Args:
      model: A score model.
      data: A PyTorch tensor.

    Returns:
      bpd: A PyTorch tensor of shape [batch size]. The log-likelihoods on `data` in bits/dim.
      z: A PyTorch tensor of the same shape as `data`. The latent representation of `data` under the
        probability flow ODE.
      nfe: An integer. The number of function evaluations used for running the black-box ODE solver.
    """
    with torch.no_grad():
      shape = data.shape
      if hutchinson_type == 'Gaussian':
        epsilon = torch.randn_like(data)
      elif hutchinson_type == 'Rademacher':
        epsilon = torch.randint_like(data, low=0, high=2).float() * 2 - 1.
      else:
        raise NotImplementedError(f"Hutchinson type {hutchinson_type} unknown.")

      def ode_func(t, x):
#         import ipdb;ipdb.set_trace() # debugging starts here
        sample = mutils.from_flattened_numpy(x[:-shape[0]], shape).to(data.device).type(torch.float32)
        vec_t = torch.ones(sample.shape[0], device=sample.device) * t
        drift = mutils.to_flattened_numpy(drift_fn(model, sample, vec_t))
        logp_grad = mutils.to_flattened_numpy(div_fn(model, sample, vec_t, epsilon))
        
        return np.concatenate([drift, logp_grad], axis=0)

      init = np.concatenate([mutils.to_flattened_numpy(data), np.zeros((shape[0],))], axis=0)
      solution = integrate.solve_ivp(ode_func, (eps, sde.T), init, rtol=rtol, atol=atol, method=method)
      nfe = solution.nfev
      zp = solution.y[:, -1]
      z = mutils.from_flattened_numpy(zp[:-shape[0]], shape).to(data.device).type(torch.float32)
      delta_logp = mutils.from_flattened_numpy(zp[-shape[0]:], (shape[0],)).to(data.device).type(torch.float32)
      prior_logp = sde.prior_logp(z)
      bpd = -(prior_logp + delta_logp) / np.log(2)
      N = np.prod(shape[1:])
      bpd = bpd / N
      # A hack to convert log-likelihoods to bits/dim
      offset = 7. - inverse_scaler(-1.)
      bpd = bpd + offset
      return bpd, z, nfe

  return likelihood_fn

attacking_diffusion_models/likelihood.py
HailongHuPri/MIDM
# -*- coding: utf-8 -*-



import os
import h5py
import random
import timeit
import argparse
import numpy as np

import torch
import datasets

# Keep the import below for registering all model definitions
import losses
import sde_lib
import likelihood
from models import utils as mutils
from models import ddpm, ncsnv2, ncsnpp
from models.ema import ExponentialMovingAverage


from compute_metrics import get_metrics






def load_checkpoint(ckpt_dir, state, device):

    loaded_state = torch.load(ckpt_dir, map_location=device)
    state['optimizer'].load_state_dict(loaded_state['optimizer'])
    state['model'].load_state_dict(loaded_state['model'], strict=False)
    state['ema'].load_state_dict(loaded_state['ema'])
    state['step'] = loaded_state['step']
    return state

def load_target_dataset(data_path, random_idx_path, memb_idx, num_images = 10000, seed = 1000):
    with h5py.File(data_path, "r") as f:
        images = f['images'][:]   
        print(images.shape)    
    
    shuffled_idx = np.load(random_idx_path)  # load training data as members
    all_member_images = images[shuffled_idx[:memb_idx]]
    all_nonmember_images = images[shuffled_idx[memb_idx:]]
    
    np.random.seed(seed)       
    shuffled_idx = np.arange(all_member_images.shape[0])
    np.random.shuffle(shuffled_idx)    
    member_images = all_member_images[shuffled_idx[:num_images]]
    print('member_images', member_images.shape)
    
    np.random.seed(seed)     
    shuffled_idx = np.arange(all_nonmember_images.shape[0])
    np.random.shuffle(shuffled_idx)
    nonmember_images = all_nonmember_images[shuffled_idx[:num_images]]  
    print('nonmember_images', nonmember_images.shape)    
    
    all_images=np.concatenate((member_images,nonmember_images),axis=0)
    return all_images


def main(args):

    seed=args.seed
    random.seed(seed)
    torch.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    np.random.seed(seed)
    
    
    # Initialize model.
    diff_types = args.diff_types
    if diff_types == 'ddpm':
        from configs.vp.ddpm import ffhq_ddpm as configs  
    elif diff_types == 'vpsde':
        from configs.vp import ffhq_vp_cont as configs  
    elif diff_types == 'smld':
        from configs.ve import ffhq_ve_discr as configs
    elif diff_types == 'vesde':  
        from configs.ve import ffhq_ve_cont as configs
    
    config = configs.get_config()
    score_model = mutils.create_model(config)
    ema = ExponentialMovingAverage(score_model.parameters(), decay=config.model.ema_rate)
    optimizer = losses.get_optimizer(config, score_model.parameters())
    state = dict(optimizer=optimizer, model=score_model, ema=ema, step=0)


    # load a trained model
    checkpoint_meta_dir = args.model_path
    state = load_checkpoint(checkpoint_meta_dir, state, config.device)
    ema.copy_to(score_model.parameters()) 
    

    
    # load a dataset (including member and nonmember data)
    memb_idx=args.memb_idx
    num_images=args.num_images
    random_idx_path=args.shuffled_idx_file
    data_path=args.data_path
    images = load_target_dataset(data_path, random_idx_path, memb_idx, num_images, seed = 1000)
    print(images.shape)
    images = np.clip(images/ 255.0, 0.0, 1.0).astype(np.float32) # [0,255] ==> [0,1]


    if config.training.sde.lower() == 'vpsde':
      sde = sde_lib.VPSDE(beta_min=config.model.beta_min, beta_max=config.model.beta_max, N=config.model.num_scales)
      sampling_eps = 1e-3
    elif config.training.sde.lower() == 'vesde':
      sde = sde_lib.VESDE(sigma_min=config.model.sigma_min, sigma_max=config.model.sigma_max, N=config.model.num_scales)
      sampling_eps = 1e-5
    else:
      raise NotImplementedError(f"SDE {config.training.sde} unknown.")   
    
    inverse_scaler = datasets.get_data_inverse_scaler(config)
    
    n_images = images.shape[0]
    im_batch_size=1
    all_likelihoods=[]
    
    if diff_types == 'ddpm' or diff_types == 'vpsde':   
        if diff_types == 'ddpm':
            likelihood_fn = likelihood.get_likelihood_fn(sde, inverse_scaler, continuous=False)
        elif diff_types == 'vpsde':
            likelihood_fn = likelihood.get_likelihood_fn(sde, inverse_scaler, continuous=True)            
        for i_batch in range(0, n_images, im_batch_size):
            if i_batch*im_batch_size%100==0:
                print(i_batch)
            batch = torch.from_numpy(images[i_batch : i_batch + im_batch_size, :]).to(config.device).float()
            batch = (batch * 2. - 1.)
            
            bpd = likelihood_fn(score_model, batch)[0]
            bpd = bpd.detach().cpu().numpy().reshape(-1)
            all_likelihoods.extend(bpd)
    
            
            if i_batch*im_batch_size%100==0:
                all_likelihoods_np = np.array(all_likelihoods)
                if args.save_dir!='None':
                    if not os.path.exists(args.save_dir):
                        os.makedirs(args.save_dir)
                    h5py_file = f'{args.save_dir}/likelihood_ffhq_{args.num_images}_{args.diff_types}.h5py'
                    with h5py.File(h5py_file, "w") as f:
                        f.create_dataset('likelihood', data = all_likelihoods_np, dtype='float32')   
                        
    elif diff_types == 'smld' or diff_types == 'vesde': 
        if diff_types == 'smld':
            likelihood_fn = likelihood.get_likelihood_fn(sde, inverse_scaler, continuous=False)
        elif diff_types == 'vesde':
            likelihood_fn = likelihood.get_likelihood_fn(sde, inverse_scaler, continuous=True)  
        for i_batch in range(0, n_images, im_batch_size):
            if i_batch*im_batch_size%100==0:
                print(i_batch)
            batch = torch.from_numpy(images[i_batch : i_batch + im_batch_size, :]).to(config.device).float()
            
            bpd = likelihood_fn(score_model, batch)[0]
            bpd = bpd.detach().cpu().numpy().reshape(-1)
            all_likelihoods.extend(bpd)
    

            if i_batch*im_batch_size%100==0:
                all_likelihoods_np = np.array(all_likelihoods)
                if args.save_dir!='None':
                    if not os.path.exists(args.save_dir):
                        os.makedirs(args.save_dir)
                    h5py_file = f'{args.save_dir}/likelihood_ffhq_{args.num_images}_{args.diff_types}.h5py'
                    with h5py.File(h5py_file, "w") as f:
                        f.create_dataset('likelihood', data = all_likelihoods_np, dtype='float32')   

        
    all_likelihoods_np = np.array(all_likelihoods)                
    if args.save_dir!='None':
        if not os.path.exists(args.save_dir):
            os.makedirs(args.save_dir)
        h5py_file = f'{args.save_dir}/likelihood_ffhq_{args.num_images}_{args.diff_types}.h5py'
        with h5py.File(h5py_file, "w") as f:
            f.create_dataset('likelihood', data = all_likelihoods_np, dtype='float32')   

    # attack performance: TFP@FPR
    likelihood_attack_performance(h5py_file,args.num_images)

def likelihood_attack_performance(data_path,num_images):

    with h5py.File(data_path, "r") as f:
        resluts = f['likelihood'][:]    # likelihood, larger is better
    print(resluts.shape)
        
    
    labels = np.concatenate((np.zeros(num_images), np.ones(num_images))) 


    all_steps_tpr_at_low_fpr=[]


    tpr_at_low_fpr_1 =get_metrics(labels, resluts, fixed_fpr=0.1)
    tpr_at_low_fpr_2 =get_metrics(labels, resluts, fixed_fpr=0.01)
    tpr_at_low_fpr_3 =get_metrics(labels, resluts, fixed_fpr=0.001)
    tpr_at_low_fpr_4 =get_metrics(labels, resluts, fixed_fpr=0.0001)

        

    all_steps_tpr_at_low_fpr.append([tpr_at_low_fpr_1,tpr_at_low_fpr_2,tpr_at_low_fpr_3,tpr_at_low_fpr_4])

    print('finished!')    

    all_steps_tpr_at_low_fpr_ = np.reshape(all_steps_tpr_at_low_fpr,(-1,4))
    

    print('[TPR@10%FPR,TPR@1%FPR,TPR@0.1%FPR,TPR@0.01%FPR]')
    print(all_steps_tpr_at_low_fpr_*100)
                
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Loss attack')
    
   
    parser.add_argument('--seed', default=2022, type=int, help='')
    parser.add_argument('--model_path', default='./', type=str, help='Path to the target model.')   
    parser.add_argument('--diff_types', default='ddpm', type=str, help='') 
    parser.add_argument('--shuffled_idx_file', default='./', type=str, help='The index file for training images (e.g. ffhq_1000_idx.npy).')   
    parser.add_argument('--data_path', default='./', type=str, help='Path to the file containing all images in the h5py format (e.g. ffhq_all.h5py).') 
    parser.add_argument('--memb_idx', default=1000, type=int, help='The index for member(training) samples in a dataset.')
    parser.add_argument('--num_images', default=1000, type=int, help='The number of images chosen for inference.')
    

    parser.add_argument('--save_dir', default='./', type=str, help='Path to save the results of the attack.')   
    

     
    args, other_args = parser.parse_known_args()
    
    print('Settings:')
    for arg in vars(args):
        print('\t{}: {}'.format(arg, getattr(args, arg)))
    ### RUN
    main(args)
attacking_diffusion_models/likelihood_attack.py
HailongHuPri/MIDM
# -*- coding: utf-8 -*-


import os
import h5py
import timeit
import random
import argparse
import numpy as np

import torch

# Keep the import below for registering all model definitions
import losses
import sde_lib
from models import utils as mutils
from models import ddpm, ncsnv2, ncsnpp
from models.ema import ExponentialMovingAverage

from compute_metrics import get_metrics


def load_checkpoint(ckpt_dir, state, device):

    loaded_state = torch.load(ckpt_dir, map_location=device)
    state['optimizer'].load_state_dict(loaded_state['optimizer'])
    state['model'].load_state_dict(loaded_state['model'], strict=False)
    state['ema'].load_state_dict(loaded_state['ema'])
    state['step'] = loaded_state['step']
    return state

def load_target_dataset(data_path, random_idx_path, memb_idx, num_images = 10000, seed = 1000):
    with h5py.File(data_path, "r") as f:
        images = f['images'][:]   
        print(images.shape)    
    
    shuffled_idx = np.load(random_idx_path)  # load training data as members
    all_member_images = images[shuffled_idx[:memb_idx]]
    all_nonmember_images = images[shuffled_idx[memb_idx:]]
    
    np.random.seed(seed)       
    shuffled_idx = np.arange(all_member_images.shape[0])
    np.random.shuffle(shuffled_idx)    
    member_images = all_member_images[shuffled_idx[:num_images]]
    print('member_images', member_images.shape)
    
    np.random.seed(seed)     
    shuffled_idx = np.arange(all_nonmember_images.shape[0])
    np.random.shuffle(shuffled_idx)
    nonmember_images = all_nonmember_images[shuffled_idx[:num_images]]  
    print('nonmember_images', nonmember_images.shape)    
    
    all_images=np.concatenate((member_images,nonmember_images),axis=0)
    return all_images
 
def get_loss_ddpm(state, batch, train=False,reduce_mean=True):
    model = state['model']
    #  training hyperparameters
    N = 1000
    beta_min=0.1
    beta_max=20
    discrete_betas = torch.linspace(beta_min / N, beta_max / N, N)
    alphas = 1. - discrete_betas
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
    sqrt_1m_alphas_cumprod = torch.sqrt(1. - alphas_cumprod)
    
    reduce_op = torch.mean if reduce_mean else lambda *args, **kwargs: 0.5 * torch.sum(*args, **kwargs)
    
    all_loss=[]
    im_batch_size=200
    with torch.no_grad():
        model_fn = mutils.get_model_fn(model, train=train)
            
        labels = torch.arange(0, N, device=batch.device).reshape(1,-1).repeat_interleave(batch.size(0),dim=0).reshape(-1)
        sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(batch.device)
        sqrt_1m_alphas_cumprod = sqrt_1m_alphas_cumprod.to(batch.device)
        
        batch = batch.repeat_interleave(N,dim=0)
        noise = torch.randn_like(batch)
        perturbed_data = sqrt_alphas_cumprod[labels, None, None, None] * batch + \
                         sqrt_1m_alphas_cumprod[labels, None, None, None] * noise
        
        for i_batch in range(0, N, im_batch_size): 
            
            score = model_fn(perturbed_data[i_batch : i_batch + im_batch_size, :], labels[i_batch : i_batch + im_batch_size])
            losses = torch.square(score - noise[i_batch : i_batch + im_batch_size, :])
            losses = reduce_op(losses.reshape(losses.shape[0], -1), dim=-1)
            all_loss += losses.cpu().numpy().tolist()
    
    return all_loss

def get_loss_vpsde(state, batch, sde, t, train=False,reduce_mean=True,eps=1e-5):
    model = state['model']
    
    #  training hyperparameters
    reduce_op = torch.mean if reduce_mean else lambda *args, **kwargs: 0.5 * torch.sum(*args, **kwargs)
    
    all_loss=[]
    im_batch_size=200
    N=1000 # here N is noise times
    with torch.no_grad():
        
        score_fn = mutils.get_score_fn(sde, model, train=train, continuous=True)
        t = t.to(batch.device).reshape(1,-1).repeat_interleave(batch.size(0),dim=0).reshape(-1)
        batch = batch.repeat_interleave(N,dim=0)
        z = torch.randn_like(batch)
        mean, std = sde.marginal_prob(batch, t)
        perturbed_data = mean + std[:, None, None, None] * z
        
        
        for i_batch in range(0, N, im_batch_size): 
            score = score_fn(perturbed_data[i_batch : i_batch + im_batch_size, :], t[i_batch : i_batch + im_batch_size])
            losses = torch.square(score * std[i_batch : i_batch + im_batch_size, None, None, None] + z[i_batch : i_batch + im_batch_size])
            losses = reduce_op(losses.reshape(losses.shape[0], -1), dim=-1)
            
            all_loss += losses.cpu().numpy().tolist()

    return all_loss

    
def get_loss_smld(state, batch, train=False,reduce_mean=True):
    model = state['model']
    #  training hyperparameters
    N = 1000
    sigma_min=0.001
    sigma_max=50
  
    
    discrete_sigmas = torch.exp(torch.linspace(np.log(sigma_min), np.log(sigma_max), N))
    smld_sigma_array = torch.flip(discrete_sigmas, dims=(0,))
    reduce_op = torch.mean if reduce_mean else lambda *args, **kwargs: 0.5 * torch.sum(*args, **kwargs)
    all_loss=[]
    im_batch_size=200
    with torch.no_grad():
        model_fn = mutils.get_model_fn(model, train=train)
            
        labels = torch.arange(0, N, device=batch.device).reshape(1,-1).repeat_interleave(batch.size(0),dim=0).reshape(-1)
        sigmas = smld_sigma_array.to(batch.device)[labels]
        noise = torch.randn_like(batch) * sigmas[:, None, None, None]
        perturbed_data = noise + batch
         
        for i_batch in range(0, N, im_batch_size): 
            
            score = model_fn(perturbed_data[i_batch : i_batch + im_batch_size, :], labels[i_batch : i_batch + im_batch_size])
            target = -noise[i_batch : i_batch + im_batch_size, :] / (sigmas[i_batch : i_batch + im_batch_size] ** 2)[:, None, None, None]
            losses = torch.square(score - target)
            losses = reduce_op(losses.reshape(losses.shape[0], -1), dim=-1) * sigmas[i_batch : i_batch + im_batch_size] ** 2
            
            all_loss += losses.cpu().numpy().tolist()
    
    return all_loss

def get_loss_vesde(state, batch, sde, t, train=False,reduce_mean=True,eps=1e-5):
    model = state['model']
    #  training hyperparameters
    reduce_op = torch.mean if reduce_mean else lambda *args, **kwargs: 0.5 * torch.sum(*args, **kwargs)
    
    all_loss=[]
    im_batch_size=200
    N=1000 # here N is noise times
    with torch.no_grad():
        
        score_fn = mutils.get_score_fn(sde, model, train=train, continuous=True)
        t = t.to(batch.device).reshape(1,-1).repeat_interleave(batch.size(0),dim=0).reshape(-1)
        batch = batch.repeat_interleave(N,dim=0)
        z = torch.randn_like(batch)
        mean, std = sde.marginal_prob(batch, t)
        perturbed_data = mean + std[:, None, None, None] * z
        
    
        for i_batch in range(0, N, im_batch_size): 
            score = score_fn(perturbed_data[i_batch : i_batch + im_batch_size, :], t[i_batch : i_batch + im_batch_size])
            losses = torch.square(score * std[i_batch : i_batch + im_batch_size, None, None, None] + z[i_batch : i_batch + im_batch_size])
            losses = reduce_op(losses.reshape(losses.shape[0], -1), dim=-1)
            
            all_loss += losses.cpu().numpy().tolist()

    return all_loss

def main(args):

    seed=args.seed
    random.seed(seed)
    torch.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    np.random.seed(seed)
    

    # Initialize model.
    diff_types = args.diff_types
    if diff_types == 'ddpm':
        from configs.vp.ddpm import ffhq_ddpm as configs  
    elif diff_types == 'vpsde':
        from configs.vp import ffhq_vp_cont as configs  
    elif diff_types == 'smld':
        from configs.ve import ffhq_ve_discr as configs
    elif diff_types == 'vesde':  
        from configs.ve import ffhq_ve_cont as configs
    
    config = configs.get_config()
    score_model = mutils.create_model(config)
    ema = ExponentialMovingAverage(score_model.parameters(), decay=config.model.ema_rate)
    optimizer = losses.get_optimizer(config, score_model.parameters())
    state = dict(optimizer=optimizer, model=score_model, ema=ema, step=0)


    # load a trained model
    checkpoint_meta_dir = args.model_path
    state = load_checkpoint(checkpoint_meta_dir, state, config.device)
    ema.copy_to(score_model.parameters()) 
    
    
    # load a dataset (including member and nonmember data)
    memb_idx=args.memb_idx
    num_images=args.num_images
    random_idx_path=args.shuffled_idx_file
    data_path=args.data_path
    images = load_target_dataset(data_path, random_idx_path, memb_idx, num_images, seed = 1000)
    print(images.shape)
    images = np.clip(images/ 255.0, 0.0, 1.0).astype(np.float32) # [0,255] ==> [0,1]

    
    if config.training.sde.lower() == 'vpsde':
      sde = sde_lib.VPSDE(beta_min=config.model.beta_min, beta_max=config.model.beta_max, N=config.model.num_scales)
      sampling_eps = 1e-3
    elif config.training.sde.lower() == 'vesde':
      sde = sde_lib.VESDE(sigma_min=config.model.sigma_min, sigma_max=config.model.sigma_max, N=config.model.num_scales)
      sampling_eps = 1e-5
    else:
      raise NotImplementedError(f"SDE {config.training.sde} unknown.")   
      
    # compute loss values
    n_images = images.shape[0]
    im_batch_size=1
    all_losses=[]
    if diff_types == 'ddpm':     
        for i_batch in range(0, n_images, im_batch_size):
            if i_batch*im_batch_size%100==0:
                print(i_batch)
            batch = torch.from_numpy(images[i_batch : i_batch + im_batch_size, :]).to(config.device).float()
            batch = (batch * 2. - 1.)
            loss = get_loss_ddpm(state, batch)
            all_losses.append(loss)

            
            if i_batch*im_batch_size%100==0:
                all_losses_np = np.array(all_losses)
                if args.save_dir!='None':
                    if not os.path.exists(args.save_dir):
                        os.makedirs(args.save_dir)
                    h5py_file = f'{args.save_dir}/loss_ffhq_{args.num_images}_{args.diff_types}.h5py'
                    with h5py.File(h5py_file, "w") as f:
                        f.create_dataset('losses', data = all_losses_np, dtype='float32')   
                        
    elif diff_types == 'vpsde':
        t = torch.rand(1000) * (1 - 1e-5) + 1e-5
        for i_batch in range(0, n_images, im_batch_size):
            if i_batch*im_batch_size%100==0:
                print(i_batch)
            batch = torch.from_numpy(images[i_batch : i_batch + im_batch_size, :]).to(config.device).float()
            batch = (batch * 2. - 1.)

            loss = get_loss_vpsde(state, batch, sde, t)
            all_losses.append(loss) 

            if i_batch*im_batch_size%100==0:
                all_losses_np = np.array(all_losses)
                if args.save_dir!='None':
                    if not os.path.exists(args.save_dir):
                        os.makedirs(args.save_dir)
                    h5py_file = f'{args.save_dir}/loss_ffhq_{args.num_images}_{args.diff_types}.h5py'
                    with h5py.File(h5py_file, "w") as f:
                        f.create_dataset('losses', data = all_losses_np, dtype='float32')  
                        f.create_dataset('t', data = t, dtype='float32')  
                        
                        
    elif diff_types == 'smld':
        for i_batch in range(0, n_images, im_batch_size):
            if i_batch*im_batch_size%100==0:
                print(i_batch)
            batch = torch.from_numpy(images[i_batch : i_batch + im_batch_size, :]).to(config.device).float()

            loss = get_loss_smld(state, batch)
            all_losses.append(loss)


            if i_batch*im_batch_size%100==0:
                all_losses_np = np.array(all_losses)
                if args.save_dir!='None':
                    if not os.path.exists(args.save_dir):
                        os.makedirs(args.save_dir)
                    h5py_file = f'{args.save_dir}/loss_ffhq_{args.num_images}_{args.diff_types}.h5py'
                    with h5py.File(h5py_file, "w") as f:
                        f.create_dataset('losses', data = all_losses_np, dtype='float32')   
                        
                        
    elif diff_types == 'vesde':      
        t = torch.rand(1000) * (1 - 1e-5) + 1e-5
        for i_batch in range(0, n_images, im_batch_size):
            if i_batch*im_batch_size%100==0:
                print(i_batch)
            batch = torch.from_numpy(images[i_batch : i_batch + im_batch_size, :]).to(config.device).float()
            loss = get_loss_vesde(state, batch, sde, t)
            all_losses.append(loss)  

            if i_batch*im_batch_size%100==0:
                all_losses_np = np.array(all_losses)
                if args.save_dir!='None':
                    if not os.path.exists(args.save_dir):
                        os.makedirs(args.save_dir)
                    h5py_file = f'{args.save_dir}/loss_ffhq_{args.num_images}_{args.diff_types}.h5py'
                    with h5py.File(h5py_file, "w") as f:
                        f.create_dataset('losses', data = all_losses_np, dtype='float32')   
                        f.create_dataset('t', data = t, dtype='float32')  
                        
                        
    
    all_losses_np = np.array(all_losses)                
    if args.save_dir!='None':
        if not os.path.exists(args.save_dir):
            os.makedirs(args.save_dir)
        h5py_file = f'{args.save_dir}/loss_ffhq_{args.num_images}_{args.diff_types}.h5py'
        if diff_types == 'ddpm' or diff_types == 'smld':     
            with h5py.File(h5py_file, "w") as f:
                f.create_dataset('losses', data = all_losses_np, dtype='float32')   
        elif diff_types == 'vpsde' or diff_types == 'vesde':     
            t = t.numpy()  
            with h5py.File(h5py_file, "w") as f:
                f.create_dataset('losses', data = all_losses_np, dtype='float32')       
                f.create_dataset('t', data = t, dtype='float32')  
                
    # attack performance: TFP@FPR
    attack_performance(h5py_file, diff_types,args.num_images)
    
def attack_performance(data_path, diff_types,num_images):
    if diff_types == 'ddpm' or diff_types == 'smld':   
        with h5py.File(data_path, "r") as f:
            resluts = f['losses'][:]    # losses, smaller is better
        print(resluts.shape)
        
    elif diff_types == 'vpsde' or diff_types == 'vesde':         
        with h5py.File(data_path, "r") as f:
            _resluts = f['losses'][:]    # losses, smaller is better
            _t=f['t'][:]
        print(_resluts.shape)

        sort_idx= np.argsort(_t)  
        t = _t[sort_idx]
        resluts = _resluts[:,sort_idx] # sorting by t
        print(resluts.shape)
        print(t.shape)
    
    labels = np.concatenate((np.zeros(num_images), np.ones(num_images))) 


    all_steps_tpr_at_low_fpr=[]

    print(resluts.shape[1])
    for i in range(resluts.shape[1]):
        tpr_at_low_fpr_1 =get_metrics(labels, resluts[:,i], fixed_fpr=0.1)
        tpr_at_low_fpr_2 =get_metrics(labels, resluts[:,i], fixed_fpr=0.01)
        tpr_at_low_fpr_3 =get_metrics(labels, resluts[:,i], fixed_fpr=0.001)
        tpr_at_low_fpr_4 =get_metrics(labels, resluts[:,i], fixed_fpr=0.0001)

        

        all_steps_tpr_at_low_fpr.append([tpr_at_low_fpr_1,tpr_at_low_fpr_2,tpr_at_low_fpr_3,tpr_at_low_fpr_4])

    print('finished!')    

    all_steps_tpr_at_low_fpr_ = np.reshape(all_steps_tpr_at_low_fpr,(-1,4))
    
    diffusion_steps = args.diffusion_steps.split(',')
    print('step, [TPR@10%FPR,TPR@1%FPR,TPR@0.1%FPR,TPR@0.01%FPR]')
    for i in diffusion_steps:
        print('t=',i,' ', all_steps_tpr_at_low_fpr_[int(i)]*100)



            
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Loss attack')
    
   
    parser.add_argument('--seed', default=2022, type=int, help='')
    parser.add_argument('--model_path', default='./', type=str, help='Path to the target model.')   
    parser.add_argument('--diff_types', default='ddpm', type=str, help='') 
    parser.add_argument('--shuffled_idx_file', default='./', type=str, help='The index file for training images (e.g. ffhq_1000_idx.npy).')   
    parser.add_argument('--data_path', default='./', type=str, help='Path to the file containing all images in the h5py format (e.g. ffhq_all.h5py).') 
    parser.add_argument('--memb_idx', default=1000, type=int, help='The index for member(training) samples in a dataset.')
    parser.add_argument('--num_images', default=1000, type=int, help='The number of images chosen for inference.')
    

    parser.add_argument('--save_dir', default='./', type=str, help='Path to save the results of the attack.')   
    
    parser.add_argument('--diffusion_steps', default='0,200,500', type=str, help='The diffusion steps, separate in comma [0-999].')
     
    args, other_args = parser.parse_known_args()
    
    print('Settings:')
    for arg in vars(args):
        print('\t{}: {}'.format(arg, getattr(args, arg)))
    ### RUN
    main(args)
attacking_diffusion_models/loss_attack.py
HailongHuPri/MIDM
# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# pylint: skip-file
# pytype: skip-file
"""Various sampling methods."""
import functools

import torch
import numpy as np
import abc

from models.utils import from_flattened_numpy, to_flattened_numpy, get_score_fn
from scipy import integrate
import sde_lib
from models import utils as mutils

_CORRECTORS = {}
_PREDICTORS = {}


def register_predictor(cls=None, *, name=None):
  """A decorator for registering predictor classes."""

  def _register(cls):
    if name is None:
      local_name = cls.__name__
    else:
      local_name = name
    if local_name in _PREDICTORS:
      raise ValueError(f'Already registered model with name: {local_name}')
    _PREDICTORS[local_name] = cls
    return cls

  if cls is None:
    return _register
  else:
    return _register(cls)


def register_corrector(cls=None, *, name=None):
  """A decorator for registering corrector classes."""

  def _register(cls):
    if name is None:
      local_name = cls.__name__
    else:
      local_name = name
    if local_name in _CORRECTORS:
      raise ValueError(f'Already registered model with name: {local_name}')
    _CORRECTORS[local_name] = cls
    return cls

  if cls is None:
    return _register
  else:
    return _register(cls)


def get_predictor(name):
  return _PREDICTORS[name]


def get_corrector(name):
  return _CORRECTORS[name]


def get_sampling_fn(config, sde, shape, inverse_scaler, eps):
  """Create a sampling function.

  Args:
    config: A `ml_collections.ConfigDict` object that contains all configuration information.
    sde: A `sde_lib.SDE` object that represents the forward SDE.
    shape: A sequence of integers representing the expected shape of a single sample.
    inverse_scaler: The inverse data normalizer function.
    eps: A `float` number. The reverse-time SDE is only integrated to `eps` for numerical stability.

  Returns:
    A function that takes random states and a replicated training state and outputs samples with the
      trailing dimensions matching `shape`.
  """

  sampler_name = config.sampling.method
  # Probability flow ODE sampling with black-box ODE solvers
  if sampler_name.lower() == 'ode':
    sampling_fn = get_ode_sampler(sde=sde,
                                  shape=shape,
                                  inverse_scaler=inverse_scaler,
                                  denoise=config.sampling.noise_removal,
                                  eps=eps,
                                  device=config.device)
  # Predictor-Corrector sampling. Predictor-only and Corrector-only samplers are special cases.
  elif sampler_name.lower() == 'pc':
    predictor = get_predictor(config.sampling.predictor.lower())
    corrector = get_corrector(config.sampling.corrector.lower())
    sampling_fn = get_pc_sampler(sde=sde,
                                 shape=shape,
                                 predictor=predictor,
                                 corrector=corrector,
                                 inverse_scaler=inverse_scaler,
                                 snr=config.sampling.snr,
                                 n_steps=config.sampling.n_steps_each,
                                 probability_flow=config.sampling.probability_flow,
                                 continuous=config.training.continuous,
                                 denoise=config.sampling.noise_removal,
                                 eps=eps,
                                 device=config.device)
  else:
    raise ValueError(f"Sampler name {sampler_name} unknown.")

  return sampling_fn


class Predictor(abc.ABC):
  """The abstract class for a predictor algorithm."""

  def __init__(self, sde, score_fn, probability_flow=False):
    super().__init__()
    self.sde = sde
    # Compute the reverse SDE/ODE
    self.rsde = sde.reverse(score_fn, probability_flow)
    self.score_fn = score_fn

  @abc.abstractmethod
  def update_fn(self, x, t):
    """One update of the predictor.

    Args:
      x: A PyTorch tensor representing the current state
      t: A Pytorch tensor representing the current time step.

    Returns:
      x: A PyTorch tensor of the next state.
      x_mean: A PyTorch tensor. The next state without random noise. Useful for denoising.
    """
    pass


class Corrector(abc.ABC):
  """The abstract class for a corrector algorithm."""

  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

  @abc.abstractmethod
  def update_fn(self, x, t):
    """One update of the corrector.

    Args:
      x: A PyTorch tensor representing the current state
      t: A PyTorch tensor representing the current time step.

    Returns:
      x: A PyTorch tensor of the next state.
      x_mean: A PyTorch tensor. The next state without random noise. Useful for denoising.
    """
    pass


@register_predictor(name='euler_maruyama')
class EulerMaruyamaPredictor(Predictor):
  def __init__(self, sde, score_fn, probability_flow=False):
    super().__init__(sde, score_fn, probability_flow)

  def update_fn(self, x, t):
    dt = -1. / self.rsde.N
    z = torch.randn_like(x)
    drift, diffusion = self.rsde.sde(x, t)
    x_mean = x + drift * dt
    x = x_mean + diffusion[:, None, None, None] * np.sqrt(-dt) * z
    return x, x_mean


@register_predictor(name='reverse_diffusion')
class ReverseDiffusionPredictor(Predictor):
  def __init__(self, sde, score_fn, probability_flow=False):
    super().__init__(sde, score_fn, probability_flow)

  def update_fn(self, x, t):
    f, G = self.rsde.discretize(x, t)
    z = torch.randn_like(x)
    x_mean = x - f
    x = x_mean + G[:, None, None, None] * z
    return x, x_mean


@register_predictor(name='ancestral_sampling')
class AncestralSamplingPredictor(Predictor):
  """The ancestral sampling predictor. Currently only supports VE/VP SDEs."""

  def __init__(self, sde, score_fn, probability_flow=False):
    super().__init__(sde, score_fn, probability_flow)
    if not isinstance(sde, sde_lib.VPSDE) and not isinstance(sde, sde_lib.VESDE):
      raise NotImplementedError(f"SDE class {sde.__class__.__name__} not yet supported.")
    assert not probability_flow, "Probability flow not supported by ancestral sampling"

  def vesde_update_fn(self, x, t):
    sde = self.sde
    timestep = (t * (sde.N - 1) / sde.T).long()
    sigma = sde.discrete_sigmas[timestep]
    adjacent_sigma = torch.where(timestep == 0, torch.zeros_like(t), sde.discrete_sigmas.to(t.device)[timestep - 1])
    score = self.score_fn(x, t)
    x_mean = x + score * (sigma ** 2 - adjacent_sigma ** 2)[:, None, None, None]
    std = torch.sqrt((adjacent_sigma ** 2 * (sigma ** 2 - adjacent_sigma ** 2)) / (sigma ** 2))
    noise = torch.randn_like(x)
    x = x_mean + std[:, None, None, None] * noise
    return x, x_mean

  def vpsde_update_fn(self, x, t):
    sde = self.sde
    timestep = (t * (sde.N - 1) / sde.T).long()
    beta = sde.discrete_betas.to(t.device)[timestep]
    score = self.score_fn(x, t)
    x_mean = (x + beta[:, None, None, None] * score) / torch.sqrt(1. - beta)[:, None, None, None]
    noise = torch.randn_like(x)
    x = x_mean + torch.sqrt(beta)[:, None, None, None] * noise
    return x, x_mean

  def update_fn(self, x, t):
    if isinstance(self.sde, sde_lib.VESDE):
      return self.vesde_update_fn(x, t)
    elif isinstance(self.sde, sde_lib.VPSDE):
      return self.vpsde_update_fn(x, t)


@register_predictor(name='none')
class NonePredictor(Predictor):
  """An empty predictor that does nothing."""

  def __init__(self, sde, score_fn, probability_flow=False):
    pass

  def update_fn(self, x, t):
    return x, x


@register_corrector(name='langevin')
class LangevinCorrector(Corrector):
  def __init__(self, sde, score_fn, snr, n_steps):
    super().__init__(sde, score_fn, snr, n_steps)
    if not isinstance(sde, sde_lib.VPSDE) \
        and not isinstance(sde, sde_lib.VESDE) \
        and not isinstance(sde, sde_lib.subVPSDE):
      raise NotImplementedError(f"SDE class {sde.__class__.__name__} not yet supported.")

  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, sde_lib.VPSDE) or isinstance(sde, sde_lib.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):
      grad = score_fn(x, t)
      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


@register_corrector(name='ald')
class AnnealedLangevinDynamics(Corrector):
  """The original annealed Langevin dynamics predictor in NCSN/NCSNv2.

  We include this corrector only for completeness. It was not directly used in our paper.
  """

  def __init__(self, sde, score_fn, snr, n_steps):
    super().__init__(sde, score_fn, snr, n_steps)
    if not isinstance(sde, sde_lib.VPSDE) \
        and not isinstance(sde, sde_lib.VESDE) \
        and not isinstance(sde, sde_lib.subVPSDE):
      raise NotImplementedError(f"SDE class {sde.__class__.__name__} not yet supported.")

  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, sde_lib.VPSDE) or isinstance(sde, sde_lib.subVPSDE):
      timestep = (t * (sde.N - 1) / sde.T).long()
      alpha = sde.alphas.to(t.device)[timestep]
    else:
      alpha = torch.ones_like(t)

    std = self.sde.marginal_prob(x, t)[1]

    for i in range(n_steps):
      grad = score_fn(x, t)
      noise = torch.randn_like(x)
      step_size = (target_snr * std) ** 2 * 2 * alpha
      x_mean = x + step_size[:, None, None, None] * grad
      x = x_mean + noise * torch.sqrt(step_size * 2)[:, None, None, None]

    return x, x_mean


@register_corrector(name='none')
class NoneCorrector(Corrector):
  """An empty corrector that does nothing."""

  def __init__(self, sde, score_fn, snr, n_steps):
    pass

  def update_fn(self, x, t):
    return x, x


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)


def get_pc_sampler(sde, shape, predictor, corrector, inverse_scaler, snr,
                   n_steps=1, probability_flow=False, continuous=False,
                   denoise=True, eps=1e-3, device='cuda'):
  """Create a Predictor-Corrector (PC) sampler.

  Args:
    sde: An `sde_lib.SDE` object representing the forward SDE.
    shape: A sequence of integers. The expected shape of a single sample.
    predictor: A subclass of `sampling.Predictor` representing the predictor algorithm.
    corrector: A subclass of `sampling.Corrector` representing the corrector algorithm.
    inverse_scaler: The inverse data normalizer.
    snr: A `float` number. The signal-to-noise ratio for configuring correctors.
    n_steps: An integer. The number of corrector steps per predictor update.
    probability_flow: If `True`, solve the reverse-time probability flow ODE when running the predictor.
    continuous: `True` indicates that the score model was continuously trained.
    denoise: If `True`, add one-step denoising to the final samples.
    eps: A `float` number. The reverse-time SDE and ODE are integrated to `epsilon` to avoid numerical issues.
    device: PyTorch device.

  Returns:
    A sampling function that returns samples and the number of function evaluations during sampling.
  """
  # Create predictor & corrector update functions
  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)

  def pc_sampler(model):
    """ The PC sampler funciton.

    Args:
      model: A score model.
    Returns:
      Samples, number of function evaluations.
    """
    with torch.no_grad():
      # Initial sample
      x = sde.prior_sampling(shape).to(device)
      timesteps = torch.linspace(sde.T, eps, sde.N, device=device)

      for i in range(sde.N):
        t = timesteps[i]
        vec_t = torch.ones(shape[0], device=t.device) * t
#         import ipdb;ipdb.set_trace() # debugging starts here
        x = x.type(torch.cuda.FloatTensor)
        x, x_mean = corrector_update_fn(x, vec_t, model=model)
        x = x.type(torch.cuda.FloatTensor)
        x, x_mean = predictor_update_fn(x, vec_t, model=model)
      return inverse_scaler(x_mean if denoise else x), sde.N * (n_steps + 1)

  return pc_sampler


def get_ode_sampler(sde, shape, inverse_scaler,
                    denoise=False, rtol=1e-5, atol=1e-5,
                    method='RK45', eps=1e-3, device='cuda'):
  """Probability flow ODE sampler with the black-box ODE solver.

  Args:
    sde: An `sde_lib.SDE` object that represents the forward SDE.
    shape: A sequence of integers. The expected shape of a single sample.
    inverse_scaler: The inverse data normalizer.
    denoise: If `True`, add one-step denoising to final samples.
    rtol: A `float` number. The relative tolerance level of the ODE solver.
    atol: A `float` number. The absolute tolerance level of the ODE solver.
    method: A `str`. The algorithm used for the black-box ODE solver.
      See the documentation of `scipy.integrate.solve_ivp`.
    eps: A `float` number. The reverse-time SDE/ODE will be integrated to `eps` for numerical stability.
    device: PyTorch device.

  Returns:
    A sampling function that returns samples and the number of function evaluations during sampling.
  """

  def denoise_update_fn(model, x):
    score_fn = get_score_fn(sde, model, train=False, continuous=True)
    # Reverse diffusion predictor for denoising
    predictor_obj = ReverseDiffusionPredictor(sde, score_fn, probability_flow=False)
    vec_eps = torch.ones(x.shape[0], device=x.device) * eps
    _, x = predictor_obj.update_fn(x, vec_eps)
    return x

  def drift_fn(model, x, t):
    """Get the drift function of the reverse-time SDE."""
    score_fn = get_score_fn(sde, model, train=False, continuous=True)
    rsde = sde.reverse(score_fn, probability_flow=True)
    return rsde.sde(x, t)[0]

  def ode_sampler(model, z=None):
    """The probability flow ODE sampler with black-box ODE solver.

    Args:
      model: A score model.
      z: If present, generate samples from latent code `z`.
    Returns:
      samples, number of function evaluations.
    """
    with torch.no_grad():
      # Initial sample
      if z is None:
        # If not represent, sample the latent code from the prior distibution of the SDE.
        x = sde.prior_sampling(shape).to(device)
      else:
        x = z

      def ode_func(t, x):
        x = from_flattened_numpy(x, shape).to(device).type(torch.float32)
        vec_t = torch.ones(shape[0], device=x.device) * t
        drift = drift_fn(model, x, vec_t)
        return to_flattened_numpy(drift)

      # Black-box ODE solver for the probability flow ODE
      solution = integrate.solve_ivp(ode_func, (sde.T, eps), to_flattened_numpy(x),
                                     rtol=rtol, atol=atol, method=method)
      nfe = solution.nfev
      x = torch.tensor(solution.y[:, -1]).reshape(shape).to(device).type(torch.float32)

      # Denoising is equivalent to running one predictor step without adding noise
      if denoise:
        x = denoise_update_fn(model, x)

      x = inverse_scaler(x)
      return x, nfe

  return ode_sampler

attacking_diffusion_models/sampling.py
HailongHuPri/MIDM
"""Abstract SDE classes, Reverse SDE, and VE/VP SDEs."""
import abc
import torch
import numpy as np


class SDE(abc.ABC):
  """SDE abstract class. Functions are designed for a mini-batch of inputs."""

  def __init__(self, N):
    """Construct an SDE.

    Args:
      N: number of discretization time steps.
    """
    super().__init__()
    self.N = N

  @property
  @abc.abstractmethod
  def T(self):
    """End time of the SDE."""
    pass

  @abc.abstractmethod
  def sde(self, x, t):
    pass

  @abc.abstractmethod
  def marginal_prob(self, x, t):
    """Parameters to determine the marginal distribution of the SDE, $p_t(x)$."""
    pass

  @abc.abstractmethod
  def prior_sampling(self, shape):
    """Generate one sample from the prior distribution, $p_T(x)$."""
    pass

  @abc.abstractmethod
  def prior_logp(self, z):
    """Compute log-density of the prior distribution.

    Useful for computing the log-likelihood via probability flow ODE.

    Args:
      z: latent code
    Returns:
      log probability density
    """
    pass

  def discretize(self, x, t):
    """Discretize the SDE in the form: x_{i+1} = x_i + f_i(x_i) + G_i z_i.

    Useful for reverse diffusion sampling and probabiliy flow sampling.
    Defaults to Euler-Maruyama discretization.

    Args:
      x: a torch tensor
      t: a torch float representing the time step (from 0 to `self.T`)

    Returns:
      f, G
    """
    dt = 1 / self.N
    drift, diffusion = self.sde(x, t)
    f = drift * dt
    G = diffusion * torch.sqrt(torch.tensor(dt, device=t.device))
    return f, G

  def reverse(self, score_fn, probability_flow=False):
    """Create the reverse-time SDE/ODE.

    Args:
      score_fn: A time-dependent score-based model that takes x and t and returns the score.
      probability_flow: If `True`, create the reverse-time ODE used for probability flow sampling.
    """
    N = self.N
    T = self.T
    sde_fn = self.sde
    discretize_fn = self.discretize

    # Build the class for reverse-time SDE.
    class RSDE(self.__class__):
      def __init__(self):
        self.N = N
        self.probability_flow = probability_flow

      @property
      def T(self):
        return T

      def sde(self, x, t):
        """Create the drift and diffusion functions for the reverse SDE/ODE."""
        drift, diffusion = sde_fn(x, t)
        score = score_fn(x, t)
        drift = drift - diffusion[:, None, None, None] ** 2 * score * (0.5 if self.probability_flow else 1.)
        # Set the diffusion function to zero for ODEs.
        diffusion = 0. if self.probability_flow else diffusion
        return drift, diffusion

      def discretize(self, x, t):
        """Create discretized iteration rules for the reverse diffusion sampler."""
        f, G = discretize_fn(x, t)
        rev_f = f - G[:, None, None, None] ** 2 * score_fn(x, t) * (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

    return RSDE()


class VPSDE(SDE):
  def __init__(self, beta_min=0.1, beta_max=20, N=1000):
    """Construct a Variance Preserving SDE.

    Args:
      beta_min: value of beta(0)
      beta_max: value of beta(1)
      N: number of discretization steps
    """
    super().__init__(N)
    self.beta_0 = beta_min
    self.beta_1 = beta_max
    self.N = N
    self.discrete_betas = torch.linspace(beta_min / N, beta_max / N, N)
    self.alphas = 1. - self.discrete_betas
    self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
    self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
    self.sqrt_1m_alphas_cumprod = torch.sqrt(1. - self.alphas_cumprod)

  @property
  def T(self):
    return 1

  def sde(self, x, t):
    beta_t = self.beta_0 + t * (self.beta_1 - self.beta_0)
    drift = -0.5 * beta_t[:, None, None, None] * x
    diffusion = torch.sqrt(beta_t)
    return drift, diffusion

  def marginal_prob(self, x, t):
    log_mean_coeff = -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
    mean = torch.exp(log_mean_coeff[:, None, None, None]) * x
    std = torch.sqrt(1. - torch.exp(2. * log_mean_coeff))
    return mean, std

  def prior_sampling(self, shape):
    return torch.randn(*shape)

  def prior_logp(self, z):
    shape = z.shape
    N = np.prod(shape[1:])
    logps = -N / 2. * np.log(2 * np.pi) - torch.sum(z ** 2, dim=(1, 2, 3)) / 2.
    return logps

  def discretize(self, x, t):
    """DDPM discretization."""
    timestep = (t * (self.N - 1) / self.T).long()
    beta = self.discrete_betas.to(x.device)[timestep]
    alpha = self.alphas.to(x.device)[timestep]
    sqrt_beta = torch.sqrt(beta)
    f = torch.sqrt(alpha)[:, None, None, None] * x - x
    G = sqrt_beta
    return f, G


class subVPSDE(SDE):
  def __init__(self, beta_min=0.1, beta_max=20, N=1000):
    """Construct the sub-VP SDE that excels at likelihoods.

    Args:
      beta_min: value of beta(0)
      beta_max: value of beta(1)
      N: number of discretization steps
    """
    super().__init__(N)
    self.beta_0 = beta_min
    self.beta_1 = beta_max
    self.N = N

  @property
  def T(self):
    return 1

  def sde(self, x, t):
    beta_t = self.beta_0 + t * (self.beta_1 - self.beta_0)
    drift = -0.5 * beta_t[:, None, None, None] * x
    discount = 1. - torch.exp(-2 * self.beta_0 * t - (self.beta_1 - self.beta_0) * t ** 2)
    diffusion = torch.sqrt(beta_t * discount)
    return drift, diffusion

  def marginal_prob(self, x, t):
    log_mean_coeff = -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
    mean = torch.exp(log_mean_coeff)[:, None, None, None] * x
    std = 1 - torch.exp(2. * log_mean_coeff)
    return mean, std

  def prior_sampling(self, shape):
    return torch.randn(*shape)

  def prior_logp(self, z):
    shape = z.shape
    N = np.prod(shape[1:])
    return -N / 2. * np.log(2 * np.pi) - torch.sum(z ** 2, dim=(1, 2, 3)) / 2.


class VESDE(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))
    adjacent_sigma = torch.where(timestep == 0, torch.zeros_like(t),
                                 self.discrete_sigmas.to(t.device)[timestep - 1])
    f = torch.zeros_like(x)
    G = torch.sqrt(sigma ** 2 - adjacent_sigma ** 2)
    return f, G
attacking_diffusion_models/sde_lib.py
HailongHuPri/MIDM
import torch
import tensorflow as tf
import os
import logging


def restore_checkpoint(ckpt_dir, state, device):
  if not tf.io.gfile.exists(ckpt_dir):
    tf.io.gfile.makedirs(os.path.dirname(ckpt_dir))
    logging.warning(f"No checkpoint found at {ckpt_dir}. "
                    f"Returned the same state as input")
    return state
  else:
    logging.warning(f"Restoring checkpoint from {ckpt_dir}. ")    
    loaded_state = torch.load(ckpt_dir, map_location=device)
    state['optimizer'].load_state_dict(loaded_state['optimizer'])
    state['model'].load_state_dict(loaded_state['model'], strict=False)
    state['ema'].load_state_dict(loaded_state['ema'])
    state['step'] = loaded_state['step']
    return state


def save_checkpoint(ckpt_dir, state):
  saved_state = {
    'optimizer': state['optimizer'].state_dict(),
    'model': state['model'].state_dict(),
    'ema': state['ema'].state_dict(),
    'step': state['step']
  }
  torch.save(saved_state, ckpt_dir)
attacking_diffusion_models/utils.py
HailongHuPri/MIDM
# -*- coding: utf-8 -*-


import warnings
warnings.filterwarnings('ignore')

import os
import h5py
import timeit
import argparse
import numpy as np

import PIL.Image
from tqdm import tqdm

import tensorflow as tf



class TFRecordExporter:
    def __init__(self, tfrecord_dir, expected_images, print_progress=True, progress_interval=10):
        self.tfrecord_dir       = tfrecord_dir
        self.tfr_prefix         = os.path.join(self.tfrecord_dir, os.path.basename(self.tfrecord_dir))
        self.expected_images    = expected_images
        self.cur_images         = 0
        self.shape              = None
        self.resolution_log2    = None
        self.tfr_writers        = []
        self.print_progress     = print_progress
        self.progress_interval  = progress_interval
        if self.print_progress:
            print('Creating dataset "%s"' % tfrecord_dir)
        if not os.path.isdir(self.tfrecord_dir):
            os.makedirs(self.tfrecord_dir)
        assert(os.path.isdir(self.tfrecord_dir))
        
    def close(self):
        if self.print_progress:
            print('%-40s\r' % 'Flushing data...', end='', flush=True)
        for tfr_writer in self.tfr_writers:
            tfr_writer.close()
        self.tfr_writers = []
        if self.print_progress:
            print('%-40s\r' % '', end='', flush=True)
            print('Added %d images.' % self.cur_images)

    def choose_shuffled_order(self): 
        order = np.arange(self.expected_images)
        np.random.RandomState(123).shuffle(order)
        return order

    def add_image(self, img):
        if self.print_progress and self.cur_images % self.progress_interval == 0:
            print('%d / %d\r' % (self.cur_images, self.expected_images), end='', flush=True)
        if self.shape is None:
            self.shape = img.shape
            self.resolution_log2 = int(np.log2(self.shape[1]))
            assert self.shape[0] in [1, 3]
            assert self.shape[1] == self.shape[2]
            assert self.shape[1] == 2**self.resolution_log2
            tfr_opt = tf.compat.v1.python_io.TFRecordOptions(tf.compat.v1.python_io.TFRecordCompressionType.NONE)
            
            for lod in range(self.resolution_log2 - 1):
                tfr_file = self.tfr_prefix + '-r%02d.tfrecords' % (self.resolution_log2 - lod)
                self.tfr_writers.append(tf.compat.v1.python_io.TFRecordWriter(tfr_file, tfr_opt))
        assert img.shape == self.shape
        for lod, tfr_writer in enumerate(self.tfr_writers):
            if lod:
                img = img.astype(np.float32)
                img = (img[:, 0::2, 0::2] + img[:, 0::2, 1::2] + img[:, 1::2, 0::2] + img[:, 1::2, 1::2]) * 0.25
            quant = np.rint(img).clip(0, 255).astype(np.uint8)
            ex = tf.train.Example(features=tf.train.Features(feature={
                'shape': tf.train.Feature(int64_list=tf.train.Int64List(value=quant.shape)),
                'data': tf.train.Feature(bytes_list=tf.train.BytesList(value=[quant.tostring()]))}))
            tfr_writer.write(ex.SerializeToString())
        self.cur_images += 1

    def add_labels(self, labels):
        if self.print_progress:
            print('%-40s\r' % 'Saving labels...', end='', flush=True)
        assert labels.shape[0] == self.cur_images
        with open(self.tfr_prefix + '-rxx.labels', 'wb') as f:
            np.save(f, labels.astype(np.float32))
            
    def __enter__(self):
        return self
    
    def __exit__(self, *args):
        self.close()



def create_from_h5py(tfrecord_dir, h5py_file, mode, shuffle = 0):
    print('Loading h5py file from "%s"' % h5py_file)
    with h5py.File(h5py_file, "r") as hf:
        images = hf[mode][:]  
    assert images[0].shape == (3, 64, 64)
    num_img = images.shape[0]
    
    with TFRecordExporter(tfrecord_dir, num_img) as tfr:
        order = tfr.choose_shuffled_order() if shuffle else np.arange(num_img)
        for idx in range(order.size):
            tfr.add_image(images[order[idx]])
            
            
def ffhq(h5py_file, img_dir, resolution=64, nc = 3):
    
    image_filenames = []
    for root,dirs,files in os.walk(img_dir):
        for file in files:
            image_filenames.append(os.path.join(root, file))
    img_files = sorted(image_filenames)
    # delete .txt file
    if '.txt' in img_files[-1]:
        _ = img_files.pop(-1)

    num_img = len(img_files)
    print('total images>>>', num_img)
    
    
    images = np.zeros((num_img, nc, resolution, resolution), dtype=np.uint8)
    for idx, img_file in tqdm(enumerate(img_files)):
        img = np.asarray(PIL.Image.open(img_file))
        assert img.shape == (128, 128, 3)
        img = PIL.Image.fromarray(img, 'RGB')
        img = img.resize((resolution, resolution), PIL.Image.ANTIALIAS)
        img = np.asarray(img)
        img = img.transpose(2, 0, 1) 
        images[idx] = img
    
    
    with h5py.File(h5py_file, "w") as f:
        f.create_dataset('images', data = images, dtype='uint8')

def make_sub_dataset(h5py_file, selected_imgs_file, shuffled_idx_file, num_images, random_seed):
    with h5py.File(h5py_file, "r") as f:
        images = f['images'][:]  
        print(images.shape)


    np.random.seed(random_seed)       
    shuffled_idx = np.arange(images.shape[0])
    print(shuffled_idx.shape)
    np.random.shuffle(shuffled_idx)    
    np.save(shuffled_idx_file,shuffled_idx) 
    print(shuffled_idx[:10])


    images_selected = images[shuffled_idx[:num_images]]

    with h5py.File(selected_imgs_file, "w") as f:
        f.create_dataset('images', data = images_selected, dtype='uint8')
    
def main(args):
    
    start = timeit.default_timer()
    # Set seed
    random_seed=args.seed
    np.random.seed(random_seed)
    if not os.path.exists(args.base_dir):
        os.mkdir(args.base_dir)
    # transform dataset into .h5py dataset    
    print('transform dataset into .h5py dataset...')
    h5py_file=args.base_dir + 'ffhq_all.h5py'
    ffhq(h5py_file, args.img_dir, args.resolution, nc = 3)
  
   # Transforming the dataset into a .h5py dataset, the default of the shuffled_idx_file is sequential.
    if args.sub_dataset != 'True' or args.tf_dataset !='True':
        shuffled_idx_file=args.base_dir + f'ffhq_{args.num_images}_idx.npy'
        shuffled_idx = np.arange(args.num_images)
        print(shuffled_idx.shape)
        np.save(shuffled_idx_file,shuffled_idx) 

        
    # Spliting a dataset into  samples and nonmember samples (random select)
    
    if args.sub_dataset == 'True':
        print('make sub dataset...')
        selected_imgs_file=args.base_dir + f'ffhq_{args.num_images}.h5py'
        shuffled_idx_file=args.base_dir + f'ffhq_{args.num_images}_idx.npy'
        make_sub_dataset(h5py_file, selected_imgs_file, 
                                  shuffled_idx_file, args.num_images, 
                                  random_seed)
    
    # Making tfrecords dataset for training target models
    
    if args.sub_dataset == 'True' and args.tf_dataset=='True':
        print('make tfrecords dataset...')
        tfrecord_dir = args.base_dir + f'tfrecord_dir_ffhq_{args.num_images}'
        if not os.path.exists(tfrecord_dir):
            os.mkdir(tfrecord_dir)
    
        create_from_h5py(tfrecord_dir, selected_imgs_file, mode = 'images')        

    stop = timeit.default_timer()
    print("Time elapses: {}s".format(stop - start))    
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Dataset Preparation')
    
   
    parser.add_argument('--seed', default=2022, type=int, help='')
    parser.add_argument('--resolution', default=64, type=int, help='')
    parser.add_argument('--num_images', default=1000, type=int, help='The number of member(training) samples.')
    

    parser.add_argument('--img_dir', default='./', type=str, help='Path for the raw image folder.')   
    parser.add_argument('--base_dir', default='./', type=str, help='Path for the base folder to save results.') 

    parser.add_argument('--sub_dataset', default='True', type=str, help='Spliting a dataset into member(training) samples and nonmember samples.')
    parser.add_argument('--tf_dataset', default='True', type=str, help='Making tfrecords dataset for training.')
       
    args, other_args = parser.parse_known_args()
    
    print('Settings:')
    for arg in vars(args):
        print('\t{}: {}'.format(arg, getattr(args, arg)))
    ### RUN
    main(args)
    
data_process.py
HailongHuPri/MIDM
import ml_collections
import torch


def get_default_configs():
  config = ml_collections.ConfigDict()
  # training
  config.training = training = ml_collections.ConfigDict()
  config.training.batch_size = 32
  training.n_iters = 500001
  training.snapshot_freq = 50000
  training.log_freq = 50
  training.eval_freq = 100
  ## store additional checkpoints for preemption in cloud computing environments
  training.snapshot_freq_for_preemption = 10000
  ## produce samples at each snapshot.
  training.snapshot_sampling = True
  training.likelihood_weighting = False
  training.continuous = True
  training.reduce_mean = False

  # sampling
  config.sampling = sampling = ml_collections.ConfigDict()
  sampling.n_steps_each = 1
  sampling.noise_removal = True
  sampling.probability_flow = False
  sampling.snr = 0.17

  # evaluation
  config.eval = evaluate = ml_collections.ConfigDict()
  evaluate.begin_ckpt = 10
  evaluate.end_ckpt = 10
  evaluate.batch_size = 10
  evaluate.enable_sampling = True
  evaluate.num_samples = 10
  evaluate.enable_loss = False
  evaluate.enable_bpd = False
  evaluate.bpd_dataset = 'train'

  # data
  config.data = data = ml_collections.ConfigDict()
  data.dataset = 'FFHQ'
  data.image_size = 64
  data.random_flip = True
  data.uniform_dequantization = False
  data.centered = False
  data.num_channels = 3

  # model
  config.model = model = ml_collections.ConfigDict()
  model.sigma_max = 90.
  model.sigma_min = 0.01
  model.num_scales = 1000
  model.beta_min = 0.1
  model.beta_max = 20.
  model.dropout = 0.1
  model.embedding_type = 'fourier'

  # optimization
  config.optim = optim = ml_collections.ConfigDict()
  optim.weight_decay = 0
  optim.optimizer = 'Adam'
  optim.lr = 2e-4
  optim.beta1 = 0.9
  optim.eps = 1e-8
  optim.warmup = 5000
  optim.grad_clip = 1.

  config.seed = 42
  config.device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

  return config
diffusion_models/configs/default_ffhq_configs.py
HailongHuPri/MIDM
import ml_collections
import torch


def get_default_configs():
  config = ml_collections.ConfigDict()
  # training
  config.training = training = ml_collections.ConfigDict()
  config.training.batch_size = 32
  training.n_iters = 500001
  training.snapshot_freq = 50000
  training.log_freq = 50
  training.eval_freq = 100
  ## store additional checkpoints for preemption in cloud computing environments
  training.snapshot_freq_for_preemption = 10000
  ## produce samples at each snapshot.
  training.snapshot_sampling = True
  training.likelihood_weighting = False
  training.continuous = True
  training.reduce_mean = False

  # sampling
  config.sampling = sampling = ml_collections.ConfigDict()
  sampling.n_steps_each = 1
  sampling.noise_removal = True
  sampling.probability_flow = False
  sampling.snr = 0.16

  # evaluation
  config.eval = evaluate = ml_collections.ConfigDict()
  evaluate.begin_ckpt = 9
  evaluate.end_ckpt = 26
  evaluate.batch_size = 1024
  evaluate.enable_sampling = False
  evaluate.num_samples = 50000
  evaluate.enable_loss = True
  evaluate.enable_bpd = False
  evaluate.bpd_dataset = 'test'

  # data
  config.data = data = ml_collections.ConfigDict()
  data.dataset = 'FFHQ'
  data.image_size = 64
  data.random_flip = True
  data.uniform_dequantization = False
  data.centered = False
  data.num_channels = 3


  # model
  config.model = model = ml_collections.ConfigDict()
  model.sigma_min = 0.01
  model.sigma_max = 50
  model.num_scales = 1000
  model.beta_min = 0.1
  model.beta_max = 20.
  model.dropout = 0.1
  model.embedding_type = 'fourier'

  # optimization
  config.optim = optim = ml_collections.ConfigDict()
  optim.weight_decay = 0
  optim.optimizer = 'Adam'
  optim.lr = 2e-4
  optim.beta1 = 0.9
  optim.eps = 1e-8
  optim.warmup = 5000
  optim.grad_clip = 1.

  config.seed = 42
  config.device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

  return config
diffusion_models/configs/default_ffhq_configs_vesde.py
HailongHuPri/MIDM
# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Training VE continuous on FFHQ"""

from configs.default_ffhq_configs_vesde import get_default_configs


def get_config():
  config = get_default_configs()
  # training
  training = config.training
  training.sde = 'vesde'
  training.continuous = True
  training.reduce_mean = True

  # sampling
  sampling = config.sampling
  sampling.method = 'pc'
  sampling.predictor = 'euler_maruyama'
  sampling.corrector = 'none'

  # data
  data = config.data
  data.dataset = 'FFHQ'
  data.centered = False
  data.tfrecords_path = 'Set the data path in your environment (*-r06.tfrecords).'
  data.image_size = 64


  # model
  model = config.model
  model.name = 'ncsnpp'
  model.scale_by_sigma = False
  model.ema_rate = 0.9999
  model.normalization = 'GroupNorm'
  model.nonlinearity = 'swish'
  model.nf = 128
  model.ch_mult = (1, 2, 2, 2)
  model.num_res_blocks = 4
  model.attn_resolutions = (16,)
  model.resamp_with_conv = True
  model.conditional = True
  model.fir = False
  model.fir_kernel = [1, 3, 3, 1]
  model.skip_rescale = True
  model.resblock_type = 'biggan'
  model.progressive = 'none'
  model.progressive_input = 'none'
  model.progressive_combine = 'sum'
  model.attention_type = 'ddpm'
  model.init_scale = 0.
  model.embedding_type = 'positional'
  model.fourier_scale = 16
  model.conv_size = 3
  



  return config

diffusion_models/configs/ve/ffhq_ve_cont.py
HailongHuPri/MIDM
# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Training SMLD (VE discrete) on FFHQ"""

from configs.default_ffhq_configs import get_default_configs


def get_config():
  config = get_default_configs()
  # training
  training = config.training
  training.sde = 'vesde'
  training.continuous = False

  # sampling
  sampling = config.sampling
  sampling.method = 'pc'
  sampling.predictor = 'reverse_diffusion'
  sampling.corrector = 'langevin'

  # model
  model = config.model
  model.name = 'ddpm'
  model.scale_by_sigma = True
  model.ema_rate = 0.999
  model.normalization = 'GroupNorm'
  model.nonlinearity = 'swish'
  model.nf = 128
  model.ch_mult = (1, 2, 2, 2)
  model.num_res_blocks = 2
  model.attn_resolutions = (16,)
  model.resamp_with_conv = True
  model.conditional = True
  model.conv_size = 3


  # data
  data = config.data
  data.dataset = 'FFHQ'
  data.image_size = 64
  data.centered = False
  data.tfrecords_path = 'Set the data path in your environment (*-r06.tfrecords).'


  
  return config



diffusion_models/configs/ve/ffhq_ve_discr.py
HailongHuPri/MIDM
# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Training DDPM (VP discrete) on FFHQ"""

from configs.default_ffhq_configs import get_default_configs


def get_config():
  config = get_default_configs()

  # training
  training = config.training
  training.sde = 'vpsde'
  training.continuous = False
  training.reduce_mean = True

  # sampling
  sampling = config.sampling
  sampling.method = 'pc'
  sampling.predictor = 'ancestral_sampling'
  sampling.corrector = 'none'

  # data
  data = config.data
  data.dataset = 'FFHQ'
  data.centered = True
  data.tfrecords_path = 'Set the data path in your environment (*-r06.tfrecords).'
  data.image_size = 64

  # model
  model = config.model
  model.name = 'ddpm'
  model.scale_by_sigma = False
  model.num_scales = 1000
  model.ema_rate = 0.9999
  model.normalization = 'GroupNorm'
  model.nonlinearity = 'swish'
  model.nf = 128
  model.ch_mult = (1, 2, 2, 2)
  model.num_res_blocks = 2
  model.attn_resolutions = (16,)
  model.resamp_with_conv = True
  model.conditional = True

  # optim
  optim = config.optim
  optim.lr = 2e-5
  



  return config

diffusion_models/configs/vp/ddpm/ffhq_ddpm.py
HailongHuPri/MIDM
# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Training VP continuous on FFHQ"""

from configs.default_ffhq_configs import get_default_configs


def get_config():
  config = get_default_configs()
  # training
  training = config.training
  training.sde = 'vpsde'
  training.continuous = True
  training.reduce_mean = True

  # sampling
  sampling = config.sampling
  sampling.method = 'pc'
  sampling.predictor = 'euler_maruyama'
  sampling.corrector = 'none'

  # data
  data = config.data
  data.dataset = 'FFHQ'
  data.centered = True
  data.tfrecords_path = 'Set the data path in your environment (*-r06.tfrecords).'
  data.image_size = 64


  # model
  model = config.model
  model.name = 'ncsnpp'
  model.scale_by_sigma = False
  model.ema_rate = 0.9999
  model.normalization = 'GroupNorm'
  model.nonlinearity = 'swish'
  model.nf = 128
  model.ch_mult = (1, 2, 2, 2)
  model.num_res_blocks = 4
  model.attn_resolutions = (16,)
  model.resamp_with_conv = True
  model.conditional = True
  model.fir = False
  model.fir_kernel = [1, 3, 3, 1]
  model.skip_rescale = True
  model.resblock_type = 'biggan'
  model.progressive = 'none'
  model.progressive_input = 'none'
  model.progressive_combine = 'sum'
  model.attention_type = 'ddpm'
  model.init_scale = 0.
  model.embedding_type = 'positional'
  model.fourier_scale = 16
  model.conv_size = 3
  



  return config

diffusion_models/configs/vp/ffhq_vp_cont.py
HailongHuPri/MIDM
import argparse,random
import h5py


import numpy as np

import os

import torch
from losses import get_optimizer
from models.ema import ExponentialMovingAverage


from utils import restore_checkpoint


import models
from models import utils as mutils
from models import ncsnpp
from models import ddpm as ddpm_model

import sampling
from sde_lib import VESDE, VPSDE
from sampling import ReverseDiffusionPredictor, LangevinCorrector
import datasets


def main(args):
    random_seed=args.seed
    random.seed(random_seed)
    torch.manual_seed(random_seed)
    torch.backends.cudnn.deterministic = True
    np.random.seed(random_seed)
    
    # @title Load the score-based model
    # sde = 'VPSDE' #@param ['VESDE', 'VPSDE', 'subVPSDE'] {"type": "string"}
    diff_types = args.diff_types
    if diff_types == 'ddpm':
        sde = 'VPSDE'
        if sde.lower() == 'vpsde':
          from configs.vp.ddpm import ffhq_ddpm as configs  
          ckpt_filename = args.model_path
          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 diff_types == 'vpsde':
        sde = 'VPSDE'
        if sde.lower() == 'vpsde':
          from configs.vp import ffhq_vp_cont as configs  
          ckpt_filename = args.model_path
          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 diff_types == 'smld':
        sde = 'VESDE'
        if sde.lower() == 'vesde':
          from configs.ve import ffhq_ve_discr as configs
          ckpt_filename = args.model_path
          config = configs.get_config()  
          sde = VESDE(sigma_min=config.model.sigma_min, sigma_max=config.model.sigma_max, N=config.model.num_scales)
          sampling_eps = 1e-5
    elif diff_types == 'vesde':  
        sde = 'VESDE'
        if sde.lower() == 'vesde':
          from configs.ve import ffhq_ve_cont as configs
          ckpt_filename = args.model_path
          config = configs.get_config()  
          sde = VESDE(sigma_min=config.model.sigma_min, sigma_max=config.model.sigma_max, N=config.model.num_scales)
          sampling_eps = 1e-5
    
    num_imgs=args.num_imgs
    batch_size =   args.batch_size  #@param {"type":"integer"}
    config.training.batch_size = batch_size
    config.eval.batch_size = batch_size
    # import ipdb;ipdb.set_trace() # debugging starts here
    # random_seed = 0 #@param {"type": "integer"}
    

    inverse_scaler = datasets.get_data_inverse_scaler(config)
    score_model = mutils.create_model(config)
    
    optimizer = get_optimizer(config, score_model.parameters())
    ema = ExponentialMovingAverage(score_model.parameters(),
                                   decay=config.model.ema_rate)
    state = dict(step=0, optimizer=optimizer,
                 model=score_model, ema=ema)
    
    state = restore_checkpoint(ckpt_filename, state, config.device)
    ema.copy_to(score_model.parameters())
    
    
    #@title PC sampling
    img_size = config.data.image_size
    channels = config.data.num_channels
    shape = (batch_size, channels, img_size, img_size)
    predictor = ReverseDiffusionPredictor #@param ["EulerMaruyamaPredictor", "AncestralSamplingPredictor", "ReverseDiffusionPredictor", "None"] {"type": "raw"}
    corrector = LangevinCorrector #@param ["LangevinCorrector", "AnnealedLangevinDynamics", "None"] {"type": "raw"}
    snr = 0.16 #@param {"type": "number"}
    n_steps =  1#@param {"type": "integer"}
    probability_flow = False #@param {"type": "boolean"}
    
    all_x = []
    
    
    sampling_fn = sampling.get_pc_sampler(sde, shape, predictor, corrector,
                                          inverse_scaler, snr, n_steps=n_steps,
                                          probability_flow=probability_flow,
                                          continuous=config.training.continuous,
                                          eps=sampling_eps, device=config.device)
    
    # import timeit
    
    for i in range(0, num_imgs, batch_size):
        # start = timeit.default_timer()
        print(i)
        x, n = sampling_fn(score_model)
        all_x.append(x.detach().cpu().numpy())
        # stop = timeit.default_timer()
        # print("Time elapses: {}s".format(stop - start))
    all_x = np.vstack(all_x)
    fake_images = np.clip(np.rint(all_x * 255.0), 0.0, 255.0).astype(np.uint8) # [-1,1] => [0,255]
    

    if args.save_dir!='None':
        if not os.path.exists(args.save_dir):
            os.makedirs(args.save_dir)
        h5py_file = f'{args.save_dir}/query_ffhq_1k_{args.diff_types}_{random_seed}.h5py'
        with h5py.File(h5py_file, "w") as f:
            f.create_dataset('images', data = fake_images, dtype='uint8')
        
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='VAE')
    
   
    parser.add_argument('--seed', default=0, type=int, help='')
    parser.add_argument('--batch_size', default=100, type=int, help='')
    parser.add_argument('--num_imgs', default=1000, type=int, help='')
    

    parser.add_argument('--save_dir', default='./', type=str, help='') 
    parser.add_argument('--model_path', default='./', type=str, help='')   
    
    parser.add_argument('--diff_types', default='ddpm', type=str, help='') 

       
    args, other_args = parser.parse_known_args()
    
    print('Settings:')
    for arg in vars(args):
        print('\t{}: {}'.format(arg, getattr(args, arg)))
    ### RUN
    main(args)
    
    
    
diffusion_models/generating.py
HailongHuPri/MIDM
# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# pylint: skip-file
# pytype: skip-file
"""Various sampling methods."""
import functools

import torch
import numpy as np
import abc

from models.utils import from_flattened_numpy, to_flattened_numpy, get_score_fn
from scipy import integrate
import sde_lib
from models import utils as mutils

_CORRECTORS = {}
_PREDICTORS = {}


def register_predictor(cls=None, *, name=None):
  """A decorator for registering predictor classes."""

  def _register(cls):
    if name is None:
      local_name = cls.__name__
    else:
      local_name = name
    if local_name in _PREDICTORS:
      raise ValueError(f'Already registered model with name: {local_name}')
    _PREDICTORS[local_name] = cls
    return cls

  if cls is None:
    return _register
  else:
    return _register(cls)


def register_corrector(cls=None, *, name=None):
  """A decorator for registering corrector classes."""

  def _register(cls):
    if name is None:
      local_name = cls.__name__
    else:
      local_name = name
    if local_name in _CORRECTORS:
      raise ValueError(f'Already registered model with name: {local_name}')
    _CORRECTORS[local_name] = cls
    return cls

  if cls is None:
    return _register
  else:
    return _register(cls)


def get_predictor(name):
  return _PREDICTORS[name]


def get_corrector(name):
  return _CORRECTORS[name]


def get_sampling_fn(config, sde, shape, inverse_scaler, eps):
  """Create a sampling function.

  Args:
    config: A `ml_collections.ConfigDict` object that contains all configuration information.
    sde: A `sde_lib.SDE` object that represents the forward SDE.
    shape: A sequence of integers representing the expected shape of a single sample.
    inverse_scaler: The inverse data normalizer function.
    eps: A `float` number. The reverse-time SDE is only integrated to `eps` for numerical stability.

  Returns:
    A function that takes random states and a replicated training state and outputs samples with the
      trailing dimensions matching `shape`.
  """

  sampler_name = config.sampling.method
  # Probability flow ODE sampling with black-box ODE solvers
  if sampler_name.lower() == 'ode':
    sampling_fn = get_ode_sampler(sde=sde,
                                  shape=shape,
                                  inverse_scaler=inverse_scaler,
                                  denoise=config.sampling.noise_removal,
                                  eps=eps,
                                  device=config.device)
  # Predictor-Corrector sampling. Predictor-only and Corrector-only samplers are special cases.
  elif sampler_name.lower() == 'pc':
    predictor = get_predictor(config.sampling.predictor.lower())
    corrector = get_corrector(config.sampling.corrector.lower())
    sampling_fn = get_pc_sampler(sde=sde,
                                 shape=shape,
                                 predictor=predictor,
                                 corrector=corrector,
                                 inverse_scaler=inverse_scaler,
                                 snr=config.sampling.snr,
                                 n_steps=config.sampling.n_steps_each,
                                 probability_flow=config.sampling.probability_flow,
                                 continuous=config.training.continuous,
                                 denoise=config.sampling.noise_removal,
                                 eps=eps,
                                 device=config.device)
  else:
    raise ValueError(f"Sampler name {sampler_name} unknown.")

  return sampling_fn


class Predictor(abc.ABC):
  """The abstract class for a predictor algorithm."""

  def __init__(self, sde, score_fn, probability_flow=False):
    super().__init__()
    self.sde = sde
    # Compute the reverse SDE/ODE
    self.rsde = sde.reverse(score_fn, probability_flow)
    self.score_fn = score_fn

  @abc.abstractmethod
  def update_fn(self, x, t):
    """One update of the predictor.

    Args:
      x: A PyTorch tensor representing the current state
      t: A Pytorch tensor representing the current time step.

    Returns:
      x: A PyTorch tensor of the next state.
      x_mean: A PyTorch tensor. The next state without random noise. Useful for denoising.
    """
    pass


class Corrector(abc.ABC):
  """The abstract class for a corrector algorithm."""

  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

  @abc.abstractmethod
  def update_fn(self, x, t):
    """One update of the corrector.

    Args:
      x: A PyTorch tensor representing the current state
      t: A PyTorch tensor representing the current time step.

    Returns:
      x: A PyTorch tensor of the next state.
      x_mean: A PyTorch tensor. The next state without random noise. Useful for denoising.
    """
    pass


@register_predictor(name='euler_maruyama')
class EulerMaruyamaPredictor(Predictor):
  def __init__(self, sde, score_fn, probability_flow=False):
    super().__init__(sde, score_fn, probability_flow)

  def update_fn(self, x, t):
    dt = -1. / self.rsde.N
    z = torch.randn_like(x)
    drift, diffusion = self.rsde.sde(x, t)
    x_mean = x + drift * dt
    x = x_mean + diffusion[:, None, None, None] * np.sqrt(-dt) * z
    return x, x_mean


@register_predictor(name='reverse_diffusion')
class ReverseDiffusionPredictor(Predictor):
  def __init__(self, sde, score_fn, probability_flow=False):
    super().__init__(sde, score_fn, probability_flow)

  def update_fn(self, x, t):
    f, G = self.rsde.discretize(x, t)
    z = torch.randn_like(x)
    x_mean = x - f
    x = x_mean + G[:, None, None, None] * z
    return x, x_mean


@register_predictor(name='ancestral_sampling')
class AncestralSamplingPredictor(Predictor):
  """The ancestral sampling predictor. Currently only supports VE/VP SDEs."""

  def __init__(self, sde, score_fn, probability_flow=False):
    super().__init__(sde, score_fn, probability_flow)
    if not isinstance(sde, sde_lib.VPSDE) and not isinstance(sde, sde_lib.VESDE):
      raise NotImplementedError(f"SDE class {sde.__class__.__name__} not yet supported.")
    assert not probability_flow, "Probability flow not supported by ancestral sampling"

  def vesde_update_fn(self, x, t):
    sde = self.sde
    timestep = (t * (sde.N - 1) / sde.T).long()
    sigma = sde.discrete_sigmas[timestep]
    adjacent_sigma = torch.where(timestep == 0, torch.zeros_like(t), sde.discrete_sigmas.to(t.device)[timestep - 1])
    score = self.score_fn(x, t)
    x_mean = x + score * (sigma ** 2 - adjacent_sigma ** 2)[:, None, None, None]
    std = torch.sqrt((adjacent_sigma ** 2 * (sigma ** 2 - adjacent_sigma ** 2)) / (sigma ** 2))
    noise = torch.randn_like(x)
    x = x_mean + std[:, None, None, None] * noise
    return x, x_mean

  def vpsde_update_fn(self, x, t):
    sde = self.sde
    timestep = (t * (sde.N - 1) / sde.T).long()
    beta = sde.discrete_betas.to(t.device)[timestep]
    score = self.score_fn(x, t)
    x_mean = (x + beta[:, None, None, None] * score) / torch.sqrt(1. - beta)[:, None, None, None]
    noise = torch.randn_like(x)
    x = x_mean + torch.sqrt(beta)[:, None, None, None] * noise
    return x, x_mean

  def update_fn(self, x, t):
    if isinstance(self.sde, sde_lib.VESDE):
      return self.vesde_update_fn(x, t)
    elif isinstance(self.sde, sde_lib.VPSDE):
      return self.vpsde_update_fn(x, t)


@register_predictor(name='none')
class NonePredictor(Predictor):
  """An empty predictor that does nothing."""

  def __init__(self, sde, score_fn, probability_flow=False):
    pass

  def update_fn(self, x, t):
    return x, x


@register_corrector(name='langevin')
class LangevinCorrector(Corrector):
  def __init__(self, sde, score_fn, snr, n_steps):
    super().__init__(sde, score_fn, snr, n_steps)
    if not isinstance(sde, sde_lib.VPSDE) \
        and not isinstance(sde, sde_lib.VESDE) \
        and not isinstance(sde, sde_lib.subVPSDE):
      raise NotImplementedError(f"SDE class {sde.__class__.__name__} not yet supported.")

  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, sde_lib.VPSDE) or isinstance(sde, sde_lib.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):
      grad = score_fn(x, t)
      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


@register_corrector(name='ald')
class AnnealedLangevinDynamics(Corrector):
  """The original annealed Langevin dynamics predictor in NCSN/NCSNv2.

  We include this corrector only for completeness. It was not directly used in our paper.
  """

  def __init__(self, sde, score_fn, snr, n_steps):
    super().__init__(sde, score_fn, snr, n_steps)
    if not isinstance(sde, sde_lib.VPSDE) \
        and not isinstance(sde, sde_lib.VESDE) \
        and not isinstance(sde, sde_lib.subVPSDE):
      raise NotImplementedError(f"SDE class {sde.__class__.__name__} not yet supported.")

  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, sde_lib.VPSDE) or isinstance(sde, sde_lib.subVPSDE):
      timestep = (t * (sde.N - 1) / sde.T).long()
      alpha = sde.alphas.to(t.device)[timestep]
    else:
      alpha = torch.ones_like(t)

    std = self.sde.marginal_prob(x, t)[1]

    for i in range(n_steps):
      grad = score_fn(x, t)
      noise = torch.randn_like(x)
      step_size = (target_snr * std) ** 2 * 2 * alpha
      x_mean = x + step_size[:, None, None, None] * grad
      x = x_mean + noise * torch.sqrt(step_size * 2)[:, None, None, None]

    return x, x_mean


@register_corrector(name='none')
class NoneCorrector(Corrector):
  """An empty corrector that does nothing."""

  def __init__(self, sde, score_fn, snr, n_steps):
    pass

  def update_fn(self, x, t):
    return x, x


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)


def get_pc_sampler(sde, shape, predictor, corrector, inverse_scaler, snr,
                   n_steps=1, probability_flow=False, continuous=False,
                   denoise=True, eps=1e-3, device='cuda'):
  """Create a Predictor-Corrector (PC) sampler.

  Args:
    sde: An `sde_lib.SDE` object representing the forward SDE.
    shape: A sequence of integers. The expected shape of a single sample.
    predictor: A subclass of `sampling.Predictor` representing the predictor algorithm.
    corrector: A subclass of `sampling.Corrector` representing the corrector algorithm.
    inverse_scaler: The inverse data normalizer.
    snr: A `float` number. The signal-to-noise ratio for configuring correctors.
    n_steps: An integer. The number of corrector steps per predictor update.
    probability_flow: If `True`, solve the reverse-time probability flow ODE when running the predictor.
    continuous: `True` indicates that the score model was continuously trained.
    denoise: If `True`, add one-step denoising to the final samples.
    eps: A `float` number. The reverse-time SDE and ODE are integrated to `epsilon` to avoid numerical issues.
    device: PyTorch device.

  Returns:
    A sampling function that returns samples and the number of function evaluations during sampling.
  """
  # Create predictor & corrector update functions
  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)

  def pc_sampler(model):
    """ The PC sampler funciton.

    Args:
      model: A score model.
    Returns:
      Samples, number of function evaluations.
    """
    with torch.no_grad():
      # Initial sample
      x = sde.prior_sampling(shape).to(device)
      timesteps = torch.linspace(sde.T, eps, sde.N, device=device)

      for i in range(sde.N):
        t = timesteps[i]
        vec_t = torch.ones(shape[0], device=t.device) * t
#         import ipdb;ipdb.set_trace() # debugging starts here
        x = x.type(torch.cuda.FloatTensor)
        x, x_mean = corrector_update_fn(x, vec_t, model=model)
        x = x.type(torch.cuda.FloatTensor)
        x, x_mean = predictor_update_fn(x, vec_t, model=model)
      return inverse_scaler(x_mean if denoise else x), sde.N * (n_steps + 1)

  return pc_sampler


def get_ode_sampler(sde, shape, inverse_scaler,
                    denoise=False, rtol=1e-5, atol=1e-5,
                    method='RK45', eps=1e-3, device='cuda'):
  """Probability flow ODE sampler with the black-box ODE solver.

  Args:
    sde: An `sde_lib.SDE` object that represents the forward SDE.
    shape: A sequence of integers. The expected shape of a single sample.
    inverse_scaler: The inverse data normalizer.
    denoise: If `True`, add one-step denoising to final samples.
    rtol: A `float` number. The relative tolerance level of the ODE solver.
    atol: A `float` number. The absolute tolerance level of the ODE solver.
    method: A `str`. The algorithm used for the black-box ODE solver.
      See the documentation of `scipy.integrate.solve_ivp`.
    eps: A `float` number. The reverse-time SDE/ODE will be integrated to `eps` for numerical stability.
    device: PyTorch device.

  Returns:
    A sampling function that returns samples and the number of function evaluations during sampling.
  """

  def denoise_update_fn(model, x):
    score_fn = get_score_fn(sde, model, train=False, continuous=True)
    # Reverse diffusion predictor for denoising
    predictor_obj = ReverseDiffusionPredictor(sde, score_fn, probability_flow=False)
    vec_eps = torch.ones(x.shape[0], device=x.device) * eps
    _, x = predictor_obj.update_fn(x, vec_eps)
    return x

  def drift_fn(model, x, t):
    """Get the drift function of the reverse-time SDE."""
    score_fn = get_score_fn(sde, model, train=False, continuous=True)
    rsde = sde.reverse(score_fn, probability_flow=True)
    return rsde.sde(x, t)[0]

  def ode_sampler(model, z=None):
    """The probability flow ODE sampler with black-box ODE solver.

    Args:
      model: A score model.
      z: If present, generate samples from latent code `z`.
    Returns:
      samples, number of function evaluations.
    """
    with torch.no_grad():
      # Initial sample
      if z is None:
        # If not represent, sample the latent code from the prior distibution of the SDE.
        x = sde.prior_sampling(shape).to(device)
      else:
        x = z

      def ode_func(t, x):
        x = from_flattened_numpy(x, shape).to(device).type(torch.float32)
        vec_t = torch.ones(shape[0], device=x.device) * t
        drift = drift_fn(model, x, vec_t)
        return to_flattened_numpy(drift)

      # Black-box ODE solver for the probability flow ODE
      solution = integrate.solve_ivp(ode_func, (sde.T, eps), to_flattened_numpy(x),
                                     rtol=rtol, atol=atol, method=method)
      nfe = solution.nfev
      x = torch.tensor(solution.y[:, -1]).reshape(shape).to(device).type(torch.float32)

      # Denoising is equivalent to running one predictor step without adding noise
      if denoise:
        x = denoise_update_fn(model, x)

      x = inverse_scaler(x)
      return x, nfe

  return ode_sampler

diffusion_models/sampling.py
HailongHuPri/MIDM
"""Abstract SDE classes, Reverse SDE, and VE/VP SDEs."""
import abc
import torch
import numpy as np


class SDE(abc.ABC):
  """SDE abstract class. Functions are designed for a mini-batch of inputs."""

  def __init__(self, N):
    """Construct an SDE.

    Args:
      N: number of discretization time steps.
    """
    super().__init__()
    self.N = N

  @property
  @abc.abstractmethod
  def T(self):
    """End time of the SDE."""
    pass

  @abc.abstractmethod
  def sde(self, x, t):
    pass

  @abc.abstractmethod
  def marginal_prob(self, x, t):
    """Parameters to determine the marginal distribution of the SDE, $p_t(x)$."""
    pass

  @abc.abstractmethod
  def prior_sampling(self, shape):
    """Generate one sample from the prior distribution, $p_T(x)$."""
    pass

  @abc.abstractmethod
  def prior_logp(self, z):
    """Compute log-density of the prior distribution.

    Useful for computing the log-likelihood via probability flow ODE.

    Args:
      z: latent code
    Returns:
      log probability density
    """
    pass

  def discretize(self, x, t):
    """Discretize the SDE in the form: x_{i+1} = x_i + f_i(x_i) + G_i z_i.

    Useful for reverse diffusion sampling and probabiliy flow sampling.
    Defaults to Euler-Maruyama discretization.

    Args:
      x: a torch tensor
      t: a torch float representing the time step (from 0 to `self.T`)

    Returns:
      f, G
    """
    dt = 1 / self.N
    drift, diffusion = self.sde(x, t)
    f = drift * dt
    G = diffusion * torch.sqrt(torch.tensor(dt, device=t.device))
    return f, G

  def reverse(self, score_fn, probability_flow=False):
    """Create the reverse-time SDE/ODE.

    Args:
      score_fn: A time-dependent score-based model that takes x and t and returns the score.
      probability_flow: If `True`, create the reverse-time ODE used for probability flow sampling.
    """
    N = self.N
    T = self.T
    sde_fn = self.sde
    discretize_fn = self.discretize

    # Build the class for reverse-time SDE.
    class RSDE(self.__class__):
      def __init__(self):
        self.N = N
        self.probability_flow = probability_flow

      @property
      def T(self):
        return T

      def sde(self, x, t):
        """Create the drift and diffusion functions for the reverse SDE/ODE."""
        drift, diffusion = sde_fn(x, t)
        score = score_fn(x, t)
        drift = drift - diffusion[:, None, None, None] ** 2 * score * (0.5 if self.probability_flow else 1.)
        # Set the diffusion function to zero for ODEs.
        diffusion = 0. if self.probability_flow else diffusion
        return drift, diffusion

      def discretize(self, x, t):
        """Create discretized iteration rules for the reverse diffusion sampler."""
        f, G = discretize_fn(x, t)
        rev_f = f - G[:, None, None, None] ** 2 * score_fn(x, t) * (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

    return RSDE()


class VPSDE(SDE):
  def __init__(self, beta_min=0.1, beta_max=20, N=1000):
    """Construct a Variance Preserving SDE.

    Args:
      beta_min: value of beta(0)
      beta_max: value of beta(1)
      N: number of discretization steps
    """
    super().__init__(N)
    self.beta_0 = beta_min
    self.beta_1 = beta_max
    self.N = N
    self.discrete_betas = torch.linspace(beta_min / N, beta_max / N, N)
    self.alphas = 1. - self.discrete_betas
    self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
    self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
    self.sqrt_1m_alphas_cumprod = torch.sqrt(1. - self.alphas_cumprod)

  @property
  def T(self):
    return 1

  def sde(self, x, t):
    beta_t = self.beta_0 + t * (self.beta_1 - self.beta_0)
    drift = -0.5 * beta_t[:, None, None, None] * x
    diffusion = torch.sqrt(beta_t)
    return drift, diffusion

  def marginal_prob(self, x, t):
    log_mean_coeff = -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
    mean = torch.exp(log_mean_coeff[:, None, None, None]) * x
    std = torch.sqrt(1. - torch.exp(2. * log_mean_coeff))
    return mean, std

  def prior_sampling(self, shape):
    return torch.randn(*shape)

  def prior_logp(self, z):
    shape = z.shape
    N = np.prod(shape[1:])
    logps = -N / 2. * np.log(2 * np.pi) - torch.sum(z ** 2, dim=(1, 2, 3)) / 2.
    return logps

  def discretize(self, x, t):
    """DDPM discretization."""
    timestep = (t * (self.N - 1) / self.T).long()
    beta = self.discrete_betas.to(x.device)[timestep]
    alpha = self.alphas.to(x.device)[timestep]
    sqrt_beta = torch.sqrt(beta)
    f = torch.sqrt(alpha)[:, None, None, None] * x - x
    G = sqrt_beta
    return f, G


class subVPSDE(SDE):
  def __init__(self, beta_min=0.1, beta_max=20, N=1000):
    """Construct the sub-VP SDE that excels at likelihoods.

    Args:
      beta_min: value of beta(0)
      beta_max: value of beta(1)
      N: number of discretization steps
    """
    super().__init__(N)
    self.beta_0 = beta_min
    self.beta_1 = beta_max
    self.N = N

  @property
  def T(self):
    return 1

  def sde(self, x, t):
    beta_t = self.beta_0 + t * (self.beta_1 - self.beta_0)
    drift = -0.5 * beta_t[:, None, None, None] * x
    discount = 1. - torch.exp(-2 * self.beta_0 * t - (self.beta_1 - self.beta_0) * t ** 2)
    diffusion = torch.sqrt(beta_t * discount)
    return drift, diffusion

  def marginal_prob(self, x, t):
    log_mean_coeff = -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
    mean = torch.exp(log_mean_coeff)[:, None, None, None] * x
    std = 1 - torch.exp(2. * log_mean_coeff)
    return mean, std

  def prior_sampling(self, shape):
    return torch.randn(*shape)

  def prior_logp(self, z):
    shape = z.shape
    N = np.prod(shape[1:])
    return -N / 2. * np.log(2 * np.pi) - torch.sum(z ** 2, dim=(1, 2, 3)) / 2.


class VESDE(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))
    adjacent_sigma = torch.where(timestep == 0, torch.zeros_like(t),
                                 self.discrete_sigmas.to(t.device)[timestep - 1])
    f = torch.zeros_like(x)
    G = torch.sqrt(sigma ** 2 - adjacent_sigma ** 2)
    return f, G
diffusion_models/sde_lib.py
HailongHuPri/MIDM
import torch
import tensorflow as tf
import os
import logging


def restore_checkpoint(ckpt_dir, state, device):
  if not tf.io.gfile.exists(ckpt_dir):
    tf.io.gfile.makedirs(os.path.dirname(ckpt_dir))
    logging.warning(f"No checkpoint found at {ckpt_dir}. "
                    f"Returned the same state as input")
    return state
  else:
    logging.warning(f"Restoring checkpoint from {ckpt_dir}. ")    
    loaded_state = torch.load(ckpt_dir, map_location=device)
    state['optimizer'].load_state_dict(loaded_state['optimizer'])
    state['model'].load_state_dict(loaded_state['model'], strict=False)
    state['ema'].load_state_dict(loaded_state['ema'])
    state['step'] = loaded_state['step']
    return state


def save_checkpoint(ckpt_dir, state):
  saved_state = {
    'optimizer': state['optimizer'].state_dict(),
    'model': state['model'].state_dict(),
    'ema': state['ema'].state_dict(),
    'step': state['step']
  }
  torch.save(saved_state, ckpt_dir)
diffusion_models/utils.py
HailongHuPri/MIDM