🐛[`core`/`GLIGEN`]: TypeError when iterating over 0-d tensor with In-painting mode when EulerAncestralDiscreteScheduler is used

[chuzhdontcode]

Describe the bug

TypeError raised when trying to directly iterate over a 0-dimension tensor in the denoising stage of GLIGEN In-painting operation.

The error occurs when using diffusion noise schedulers that iterate over timesteps
(e.g. EulerAnchestralDiscreteScheduler, KDPM2AncestralDiscreteScheduler), during in-painting operation with the #4441 and #4777.

For further clarification, this operation of the add_noise function 🔽

diffusers/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py

Lines 387 to 388 in ae2fc01

step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]

in the affected noise schedulers expects the timesteps to be a non-0 dim torch Tensor. However, in the affected pipelines, timesteps is 0-dimension.

Reproduction

import torch
from diffusers import StableDiffusionGLIGENPipeline
from diffusers.utils import load_image
import matplotlib.pyplot as plt

# Generate an image described by the prompt and
# insert objects described by text at the region defined by bounding boxes
pipe = StableDiffusionGLIGENPipeline.from_pretrained(
    "masterful/gligen-1-4-generation-text-box", variant="fp16", torch_dtype=torch.float16
)
pipe = pipe.to("cuda")

prompt = "a waterfall and a modern high speed train running through the tunnel in a beautiful forest with fall foliage"
boxes = [[0.1387, 0.2051, 0.4277, 0.7090], [0.4980, 0.4355, 0.8516, 0.7266]]
phrases = ["a waterfall", "a modern high speed train running through the tunnel"]

images = pipe(
    prompt=prompt,
    gligen_phrases=phrases,
    gligen_boxes=boxes,
    gligen_scheduled_sampling_beta=1,
    output_type="pil",
    num_inference_steps=50,
).images

plt.imshow(images[0])

Logs

{
	"name": "TypeError",
	"message": "iteration over a 0-d tensor",
	"stack": "---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<path to python file> Cell 9 line 3
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=24'>25</a> pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=26'>27</a> # images = pipe(
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=27'>28</a> #     prompt=prompt,
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=28'>29</a> #     image=input_image,
   (...)
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=31'>32</a> #     num_images_per_prompt=2,
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=32'>33</a> # ).images
---> <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=33'>34</a> images = pipe(
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=34'>35</a>     prompt=prompt,
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=35'>36</a>     gligen_phrases=phrases,
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=36'>37</a>     gligen_boxes=boxes,
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=37'>38</a>     gligen_scheduled_sampling_beta=1.0,
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=38'>39</a>     gligen_inpaint_image = gligen_inpaint_image,
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=39'>40</a>     output_type=\"np\",
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=40'>41</a>     num_inference_steps=50,
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=41'>42</a>     num_images_per_prompt=1,
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=42'>43</a> ).images
     <a href='vscode-notebook-cell:<path to python file>#X11sdnNjb2RlLXJlbW90ZQ%3D%3D?line=44'>45</a> plt.imshow(images[0])

File ~/.conda/envs/conda_env/lib/python3.10/site-packages/torch/utils/_contextlib.py:115, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    112 @functools.wraps(func)
    113 def decorate_context(*args, **kwargs):
    114     with ctx_factory():
--> 115         return func(*args, **kwargs)

File ~/defect_generation/diffusers_chu/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_gligen.py:803, in StableDiffusionGLIGENPipeline.__call__(self, prompt, height, width, num_inference_steps, guidance_scale, gligen_scheduled_sampling_beta, gligen_phrases, gligen_boxes, gligen_inpaint_image, negative_prompt, num_images_per_prompt, eta, generator, latents, prompt_embeds, negative_prompt_embeds, output_type, return_dict, callback, callback_steps, cross_attention_kwargs, clip_skip)
    799     latents = torch.randn_like(latents[:, :4])
    801 if gligen_inpaint_image is not None:
    802     gligen_inpaint_latent_with_noise = (
--> 803         self.scheduler.add_noise(gligen_inpaint_latent, torch.randn_like(gligen_inpaint_latent), t)
    804         .expand(latents.shape[0], -1, -1, -1)
    805         .clone()
    806     )
    807     latents = gligen_inpaint_latent_with_noise * gligen_inpaint_mask + latents * (
    808         1 - gligen_inpaint_mask
    809     )
    811 # expand the latents if we are doing classifier free guidance

File ~/defect_generation/diffusers_chu/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py:387, in EulerAncestralDiscreteScheduler.add_noise(self, original_samples, noise, timesteps)
    384     schedule_timesteps = self.timesteps.to(original_samples.device)
    385     timesteps = timesteps.to(original_samples.device)
--> 387 step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
    389 sigma = sigmas[step_indices].flatten()
    390 while len(sigma.shape) < len(original_samples.shape):

File ~/.conda/envs/conda_env/lib/python3.10/site-packages/torch/_tensor.py:930, in Tensor.__iter__(self)
    920 def __iter__(self):
    921     # NB: we use 'imap' and not 'map' here, so that in Python 2 we get a
    922     # generator and don't eagerly perform all the indexes.  This could
   (...)
    927     # NB: We have intentionally skipped __torch_function__ dispatch here.
    928     # See gh-54457
    929     if self.dim() == 0:
--> 930         raise TypeError(\"iteration over a 0-d tensor\")
    931     if torch._C._get_tracing_state():
    932         warnings.warn(
    933             \"Iterating over a tensor might cause the trace to be incorrect. \"
    934             \"Passing a tensor of different shape won't change the number of \"
   (...)
    938             stacklevel=2,
    939         )

TypeError: iteration over a 0-d tensor" <=================ERROR!!!!!
}

System Info

• diffusers version: 0.22.0.dev0
• Platform: Linux-5.15.0-60-generic-x86_64-with-glibc2.31
• Python version: 3.10.13
• PyTorch version (GPU?): 2.0.1 (True)
• Huggingface_hub version: 0.16.4
• Transformers version: 4.33.1
• Accelerate version: not installed
• xFormers version: not installed
• Using GPU in script?: Yes
• Using distributed or parallel set-up in script?: No

Who can help?

cc: @sayakpaul , @nikhil-masterful, @tuanh123789