Runtime Error in Image Upscaling Code - Batch Size Mismatch

🧨 Diffusers
1

I encountered a runtime error while running the code in the Image-to-upscaler-to-super-resolution documentation.

# Start with an image-to-image pipeline:
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import make_image_grid, load_image

pipeline = AutoPipelineForImage2Image.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")
pipeline.enable_model_cpu_offload()

# prepare image
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-init.png"
init_image = load_image(url)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"

# pass prompt and image to pipeline
image_1 = pipeline(prompt, image=init_image, output_type="latent").images[0]

# Chain it to an upscaler pipeline to increase the image resolution:
from diffusers import StableDiffusionLatentUpscalePipeline

upscaler = StableDiffusionLatentUpscalePipeline.from_pretrained(
    "stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")
upscaler.enable_model_cpu_offload()

image_2 = upscaler(prompt, image=image_1, output_type="latent").images[0]

# Finally, chain it to a super-resolution pipeline to further enhance the resolution:
from diffusers import StableDiffusionUpscalePipeline

super_res = StableDiffusionUpscalePipeline.from_pretrained(
    "stabilityai/stable-diffusion-x4-upscaler", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")
super_res.enable_model_cpu_offload()

image_3 = super_res(prompt, image=image_2).images[0]
make_image_grid([init_image, image_3.resize((512, 512))], rows=1, cols=2)
image_3.show()

In particular, this line

image_3 = super_res(prompt, image=image_2).images[0]

throws the following error

Exception has occurred: ValueError
`prompt` has batch size 1 and `image` has batch size 4. Please make sure that passed `prompt` matches the batch size of `image`.
  File "/home/user/repos/diffusers/img2img.py", line 59, in <module>
    image_3 = super_res(prompt, image=image_2).images[0]
ValueError: `prompt` has batch size 1 and `image` has batch size 4. Please make sure that passed `prompt` matches the batch size of `image`.

I am wondering if there might be an issue with the code provided in the documentation.