import pdb from diffusion import gaussian_diffusion as gd from diffusion.respace import SpacedDiffusion, space_timesteps def create_gaussian_diffusion(DiffusionClass=SpacedDiffusion,use_ddim=False): noise_schedule = 'cosine' sigma_small = True lambda_vel = 0.0 lambda_rcxyz = 0.0 lambda_fc = 0.0 # default params predict_xstart = True # we always predict x_start (a.k.a. x0), that's our deal! steps = 1000 scale_beta = 1. # no scaling timestep_respacing =None if use_ddim: timestep_respacing = 'ddim50' # can be used for ddim sampling, we don't use it. learn_sigma = False rescale_timesteps = False betas = gd.get_named_beta_schedule(noise_schedule, steps, scale_beta) loss_type = gd.LossType.MSE if not timestep_respacing: timestep_respacing = [steps] return DiffusionClass( use_timesteps=space_timesteps(steps, timestep_respacing), betas=betas, model_mean_type=( gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X ), model_var_type=( ( gd.ModelVarType.FIXED_LARGE if not sigma_small else gd.ModelVarType.FIXED_SMALL ) if not learn_sigma else gd.ModelVarType.LEARNED_RANGE ), loss_type=loss_type, rescale_timesteps=rescale_timesteps, lambda_vel=lambda_vel, lambda_rcxyz=lambda_rcxyz, lambda_fc=lambda_fc, )