Certain step counts for DDIM cause out of bounds error_add one to get the final alpha values right

报错记录： 

Traceback (most recent call last):
  File "/home/yiming/Projects/EvDiffusion/scripts/frame_voxel/test_sr_ver1.py", line 86, in <module>
    samples_ddim, _ = sampler.sample(S=opt.steps,
  File "/home/yiming/anaconda3/envs/ldm/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/yiming/Projects/EvDiffusion/ldm/models/diffusion/ddim.py", line 98, in sample
    self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
  File "/home/yiming/Projects/EvDiffusion/ldm/models/diffusion/ddim.py", line 44, in make_schedule
    ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
  File "/home/yiming/Projects/EvDiffusion/ldm/modules/diffusionmodules/util.py", line 65, in make_ddim_sampling_parameters
    alphas = alphacums[ddim_timesteps]
IndexError: index 1000 is out of bounds for dimension 0 with size 1000

查到的解决办法： 

Certain step counts for DDIM cause out of bounds error by MartinCairnsSQL · Pull Request #3917 · AUTOMATIC1111/stable-diffusion-webui · GitHub

When DDIM is set to any of the following step counts 3, 9, 27, 36, 37, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110,111 the last value in the ddim_timesteps array is 1000 which is out of bounds for the 1000 entries in the alphacums array.

The formula 999 / (1000 // num_steps) + 1 will give the next valid DDIM step if it's an integer value.

自行问题分析：

检查源码发现这里的计算逻辑是这样的：

首先计算ddim timesteps的间隔：(1000 // num_ddim_timesteps)

然后在（0，1000， (1000 // num_ddim_timesteps)）之间取序列，最后把序列中的每个值+1得到ddim timesteps

def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True):
    if ddim_discr_method == 'uniform':
        c = num_ddpm_timesteps // num_ddim_timesteps
        ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))
    elif ddim_discr_method == 'quad':
        ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int)
    else:
        raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"')
 
    # assert ddim_timesteps.shape[0] == num_ddim_timesteps
    # add one to get the final alpha values right (the ones from first scale to data during sampling)
    steps_out = ddim_timesteps + 1
    if verbose:
        print(f'Selected timesteps for ddim sampler: {steps_out}')
    return steps_out

根据刚刚计算出的ddim timesteps得到alphas

alphacums是一个（1000，）的向量，然后刚刚计算得到的ddim_timesteps索引得到alphas

def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True):
    # select alphas for computing the variance schedule
    alphas = alphacums[ddim_timesteps]
    alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
 
    # according the the formula provided in https://arxiv.org/abs/2010.02502
    sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))
    if verbose:
        print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}')
        print(f'For the chosen value of eta, which is {eta}, '
              f'this results in the following sigma_t schedule for ddim sampler {sigmas}')
    return sigmas, alphas, alphas_prev

所以如果num_ddim_timesteps设置为3，最后ddim timesteps[-1]是1000，所以取alphas会超出维度

解决方案

PoliticalCostlyAdministrators - Replit

根据给出的num_ddim_timesteps计算最接近的可行num_ddim_timesteps的方法：

def next_valid_ddim_step(ddim_num_steps):
    valid_step = 999 / (1000 // ddim_num_steps)
    if valid_step == floor(valid_step):
        return int(valid_step) + 1
    
    return ddim_num_steps

或者直接把 make_ddim_timesteps中的

c = num_ddpm_timesteps // num_ddim_timesteps

np.asarray(list(range(0, num_ddpm_timesteps, c)))

换成

c = num_ddpm_timesteps // num_ddim_timesteps

(np.arange(0, num_ddim_timesteps) * c).astype(int)

也就是：

def make_ddim_timesteps_hym(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True):
    if ddim_discr_method == 'uniform':
        c = num_ddpm_timesteps // num_ddim_timesteps
        # ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))
        ddim_timesteps = (np.arange(0, num_ddim_timesteps) * c).astype(int)
    elif ddim_discr_method == 'quad':
        ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int)
    else:
        raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"')
 
    # assert ddim_timesteps.shape[0] == num_ddim_timesteps
    # add one to get the final alpha values right (the ones from first scale to data during sampling)
    steps_out = ddim_timesteps + 1
    if verbose:
        print(f'Selected timesteps for ddim sampler: {steps_out}')
    return steps_out