Repository: huggingface/diffusers

Issue:
num_images_per_prompt=num_samples FAILS after latest commit in SD pipelines
### Describe the bug

This used to wwork:

images = pipe(
    prompt=prompt,
    image=image,
    mask_image=mask_image,
    guidance_scale=guidance_scale,
    generator=generator,
    num_images_per_prompt=num_samples,
    num_inference_steps =50,
    height=height,
    width=width,
    
).images

Now it doesn't.



### Reproduction

Ue the RUNWAYML inpainting notebook. IT fails in that!


https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/in_painting_with_stable_diffusion_using_diffusers.ipynb


<img width="1339" alt="image" src="https://user-images.githubusercontent.com/26677859/201589394-ffadb77b-28c4-4667-8e53-7d381f482261.png">
 

### Logs

```shell
RuntimeError                              Traceback (most recent call last)
<ipython-input-8-7352d4d77608> in <module>
     11     guidance_scale=guidance_scale,
     12     generator=generator,
---> 13     num_images_per_prompt=num_samples,
     14 ).images

1 frames
/usr/local/lib/python3.7/dist-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py in __call__(self, prompt, image, mask_image, height, width, num_inference_steps, guidance_scale, negative_prompt, num_images_per_prompt, eta, generator, latents, output_type, return_dict, callback, callback_steps, **kwargs)
    567 
    568             # concat latents, mask, masked_image_latents in the channel dimension
--> 569             latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)
    570 
    571             latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 6 but got size 2 for tensor number 1 in the list.
```


### System Info

The colab notebook in the runwayml inpainting page fails.


Discussion:
Marking @patrickvonplaten as the latest commits were made by you. 
It has to do something with num_samples because the error keeps changing with num_samples. 
Marking @patil-suraj @pcuenca Incase you guys can resolve this quickly
Thanks for reporting, looking into it.

Resolving patch:
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
@@ -536,7 +536,7 @@ def __call__(
         mask, masked_image_latents = self.prepare_mask_latents(
             mask,
             masked_image,
-            batch_size,
+            batch_size * num_images_per_prompt,
             height,
             width,
             text_embeddings.dtype,


Test patch:
diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py
--- a/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py
+++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py
@@ -215,6 +215,47 @@ def test_stable_diffusion_inpaint(self):
         assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
         assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2
 
+    def test_stable_diffusion_inpaint_with_num_images_per_prompt(self):
+        device = "cpu"
+        unet = self.dummy_cond_unet_inpaint
+        scheduler = PNDMScheduler(skip_prk_steps=True)
+        vae = self.dummy_vae
+        bert = self.dummy_text_encoder
+        tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
+
+        image = self.dummy_image.cpu().permute(0, 2, 3, 1)[0]
+        init_image = Image.fromarray(np.uint8(image)).convert("RGB").resize((128, 128))
+        mask_image = Image.fromarray(np.uint8(image + 4)).convert("RGB").resize((128, 128))
+
+        # make sure here that pndm scheduler skips prk
+        sd_pipe = StableDiffusionInpaintPipeline(
+            unet=unet,
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=bert,
+            tokenizer=tokenizer,
+            safety_checker=None,
+            feature_extractor=None,
+        )
+        sd_pipe = sd_pipe.to(device)
+        sd_pipe.set_progress_bar_config(disable=None)
+
+        prompt = "A painting of a squirrel eating a burger"
+        generator = torch.Generator(device=device).manual_seed(0)
+        images = sd_pipe(
+            [prompt],
+            generator=generator,
+            guidance_scale=6.0,
+            num_inference_steps=2,
+            output_type="np",
+            image=init_image,
+            mask_image=mask_image,
+            num_images_per_prompt=2,
+        ).images
+
+        # check if the output is a list of 2 images
+        assert len(images) == 2
+
     @unittest.skipIf(torch_device != "cuda", "This test requires a GPU")
     def test_stable_diffusion_inpaint_fp16(self):
         """Test that stable diffusion inpaint_legacy works with fp16"""