Bug in DDIMInverseScheduler function step

[Rashfu]

Describe the bug

The step function call in DDIMInverseScheduler is causing an array out-of-bounds bug. The error is located here.

diffusers/src/diffusers/schedulers/scheduling_ddim_inverse.py

Line 343 in 872ae1d

alpha_prod_t_prev = self.alphas_cumprod[prev_timestep]

The next time step(prev_timestep) for 980 is 1000, which is is out of bounds for self.alphas_cumprod. I believe that adding a conditional statement afterwards to check if it exceeds the length of the array and retrieve the last value of the array would fix the issue.

alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep < self.alphas_cumprod.shape[0] else self.alphas_cumprod[-1]

Reproduction

from diffusers import LDMPipeline, StableDiffusionPipeline, DDIMInverseScheduler, DDIMScheduler
from torch import autocast, inference_mode
import torch

ddim = LDMPipeline.from_pretrained("CompVis/ldm-celebahq-256")
device = "cuda" if torch.cuda.is_available() else "cpu"
ddim.to(device)

from tqdm.auto import tqdm
from PIL import Image

def show_lat(latents):  
  with torch.inference_mode():
    image = ddim.vqvae.decode(latents).sample  #№ latents
    image = (image / 2 + 0.5).clamp(0, 1)
    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
    images = (image * 255).round().astype("uint8")
    pil_images = [Image.fromarray(image) for image in images]
    return pil_images[0]

batch_size = 1
orig_latents = torch.randn((batch_size, ddim.unet.config.in_channels, ddim.unet.config.sample_size, ddim.unet.config.sample_size)).to(device)

latents = orig_latents.clone()
with autocast("cuda"), inference_mode():
    for i, e in enumerate(tqdm(ddim.scheduler.timesteps)):
        latents = ddim.scheduler.step(ddim.unet(latents, e).sample, e, latents, eta=0).prev_sample
show_lat(latents)

decoded_latents = latents.clone()
inverse_scheduler = DDIMInverseScheduler.from_config(ddim.scheduler.config)
inverse_scheduler.set_timesteps(ddim.scheduler.num_inference_steps)

with autocast("cuda"), inference_mode():
    for i, e in enumerate(tqdm(ddim.scheduler.timesteps.flip(0))):
        decoded_latents = inverse_scheduler.step(ddim.unet(decoded_latents, e).sample, e, decoded_latents).prev_sample
show_lat(decoded_latents)

Logs

IndexError                                Traceback (most recent call last)
Cell In[7], line 8
      6 with autocast("cuda"), inference_mode():
      7     for i, e in enumerate(tqdm(ddim.scheduler.timesteps.flip(0))):
----> 8         decoded_latents = inverse_scheduler.step(ddim.unet(decoded_latents, e).sample, e, decoded_latents).prev_sample
      9 show_lat(decoded_latents)

File ~/anaconda3/envs/zjy/lib/python3.8/site-packages/diffusers/schedulers/scheduling_ddim_inverse.py:343, in DDIMInverseScheduler.step(self, model_output, timestep, sample, eta, use_clipped_model_output, variance_noise, return_dict)
    340 # 2. compute alphas, betas
    341 # change original implementation to exactly match noise levels for analogous forward process
    342 alpha_prod_t = self.alphas_cumprod[timestep] if timestep >= 0 else self.initial_alpha_cumprod
--> 343 alpha_prod_t_prev = self.alphas_cumprod[prev_timestep]
    345 beta_prod_t = 1 - alpha_prod_t
    347 # 3. compute predicted original sample from predicted noise also called
    348 # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf

IndexError: index 1000 is out of bounds for dimension 0 with size 1000

System Info

• diffusers version: 0.21.4
• Platform: Linux-4.15.0-213-generic-x86_64-with-glibc2.17
• Python version: 3.8.18
• PyTorch version (GPU?): 1.12.1+cu113 (True)
• Huggingface_hub version: 0.17.3
• Transformers version: not installed
• Accelerate version: 0.23.0

Who can help?

@patrickvonplaten I guess this issue is linked to the previous one #702.

[patrickvonplaten]

Is this related to #5145 ?
cc @richardSHkim any thoughts on this one?

[richardSHkim]

Hi, @Rashfu. Could you try this error with the latest version of diffusers?
It seems like v0.21.4 uses old version of ddim inverse scheduler.
I found no error with the latest version which uses fixed ddim inverse scheduler (#5145)

[Rashfu]

Thanks for your advice! It works for me by installing diffusers from source.

pip install git+https://github.com/huggingface/diffusers