How do I set up multiple samples to display in a colab notebook with K-Diffusion?

Hi! I’m working on a notebook right now using the new K-Diffusion pipeline support, does anyone have a code example of how to call for multiple samples and display them all successively in a colab notebook?

Passing num_images_per_prompt above 1 in a pipe gives me this error:

RuntimeError: The size of tensor a (4) must match the size of tensor b (2) at non-singleton dimension 0

Thank you in advance! <3

Interesting! Can you please post a code snippet to show how you’re using it? num_images_per_prompt should work, so there might be a bug there.

Happy to! Thanks! <3

The below example works as expected.

from diffusers import StableDiffusionKDiffusionPipeline
import torch
from PIL import Image
import os

model_id = "stabilityai/stable-diffusion-2-1-base" 

pipe = StableDiffusionKDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to(


prompt="mountain landscape in the style of thomas kinkade" 

negative_prompt="3d, monochromatic" 

img = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=steps, guidance_scale=cfg, width=width, height=height).images[0]

If I add num_images_per_prompt=2 into the pipe at the bottom, it errors out for me like mentioned above. It could be that I’m just missing something silly, but appreciate you looking into it for me!

Yes, this looks like a bug to me!

Currently being discussed in this PR.

Thank you!

Thank you so much for looking into this!

1 Like

It should be fixed now if you install from github:

pip uninstall diffusers -y
pip install git+