Spaces:
Running
on
Zero
Running
on
Zero
| 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, | |
| ) |