Repository: huggingface/diffusers
Issue category: Bug Report

Issue:
RuntimeError with LyCORIS, Batch Inference and skip_guidance_layers
### Describe the bug

A `RuntimeError` occurs when using the following combination:

* SD3 
* Batch inference (`num_images_per_prompt > 1`)
* LyCORIS 
* `skip_guidance_layers` is set

The error message is: `"RuntimeError: The size of tensor a (2) must match the size of tensor b (4) at non-singleton dimension 0"`

It seems that batch inference (`num_images_per_prompt > 1`) does not work in conjunction with `skip_guidance_layers`.


### Reproduction

This code snippet produces the error:

```python
self.pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers", 
    torch_dtype=torch.bfloat16
    )

self.pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
    self.pipe.scheduler.config,
    timestep_spacing="trailing",
    shift=3.0
    )
self.pipe.to("cuda")

lora_scale = 1.0
wrapper, _ = create_lycoris_from_weights(lora_scale, my_lora, self.pipe.transformer)
wrapper.merge_to()

image = self.pipe(
    prompt=request.prompt,
    num_inference_steps=request.num_inference_steps,
    num_images_per_prompt=2,  # Batch inference
    output_type="pil",
    generator=torch.Generator(device="cuda").manual_seed(42),
    guidance_scale=request.guidance_scale,
    width=request.width,
    height=request.height,
    skip_guidance_layers=[7, 8, 9],  # Doesn't seem to work with batching
).images[0]
```

Commenting out `skip_guidance_layers` resolves the error.

**Expected behavior**

Batch inference should work correctly even when `skip_guidance_layers` is used with LyCORIS.



### Logs

_No response_

### System Info

**Environment**

*  CUDA Version: 12.4
*  Python version: 3.12.1 (main, Jan 11 2024, 10:22:40) [GCC 10.2.1 20210110]
*  Diffusers version: https://github.com/huggingface/diffusers.git@99c0483b67427de467f11aa35d54678fd36a7ea2
*  The specific LyCORIS model and inference method used from Bghira: https://huggingface.co/bghira/sd35m-photo-mixedres-cL-sS3-noOverride?not-for-all-audiences=true

### Who can help?

@sayakpaul


Hints:
This is not a fully reproducible snippet. Please provide one. 

Cc: @asomoza and @yiyixuxu for skip layer guidance. 
Reproducible with:

```bash
pip install lycoris_lora
```

```python
from diffusers import StableDiffusion3Pipeline, FlowMatchEulerDiscreteScheduler
import torch
from huggingface_hub import hf_hub_download

from lycoris import create_lycoris_from_weights

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-medium", torch_dtype=torch.bfloat16
)
pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
    pipe.scheduler.config, timestep_spacing="trailing", shift=3.0
)
pipe.to("cuda")

weights_path = hf_hub_download(
    "bghira/sd35m-photo-mixedres-cL-sS3-noOverride",
    filename="pytorch_lora_weights.safetensors",
)

lora_scale = 1.0
wrapper, _ = create_lycoris_from_weights(lora_scale, weights_path, pipe.transformer)
wrapper.merge_to()

prompt = "A photo-realistic image of a cat"
num_inference_steps = 28
guidance_scale = 3.5
width = 1024
height = 1024

image = pipe(
    prompt=prompt,
    num_inference_steps=num_inference_steps,
    num_images_per_prompt=2,  # Batch inference
    output_type="pil",
    generator=torch.Generator(device="cuda").manual_seed(42),
    guidance_scale=guidance_scale,
    width=width,
    height=height,
    skip_guidance_layers=[7, 8, 9],  # Doesn't seem to work with batching
).images[0]
```

```python
Traceback (most recent call last):
  File "/workspace/test.py", line 30, in <module>
    image = pipe(
            ^^^^^
  File "/usr/local/lib/python3.11/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/diffusers/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py", line 918, in __call__
    noise_pred = self.transformer(
                 ^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/diffusers/src/diffusers/models/transformers/transformer_sd3.py", line 393, in forward
    temb = self.time_text_embed(timestep, pooled_projections)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/diffusers/src/diffusers/models/embeddings.py", line 1208, in forward
    conditioning = timesteps_emb + pooled_projections
                   ~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~
RuntimeError: The size of tensor a (2) must match the size of tensor b (4) at non-singleton dimension 0
```
So, SLG seems to be the culprit here? 
> So, SLG seems to be the culprit here?

Indeed, problem seems to appear since I started using SLG
After some testing lycoris is unrelated, minimally reproducible with

```python
from diffusers import StableDiffusion3Pipeline, FlowMatchEulerDiscreteScheduler
import torch

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-medium", torch_dtype=torch.bfloat16
)
pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
    pipe.scheduler.config, timestep_spacing="trailing", shift=3.0
)
pipe.to("cuda")

prompt = "A photo-realistic image of a cat"
num_inference_steps = 28
guidance_scale = 3.5
width = 1024
height = 1024

image = pipe(
    prompt=prompt,
    num_inference_steps=num_inference_steps,
    num_images_per_prompt=2,
    output_type="pil",
    generator=torch.Generator("cuda").manual_seed(42),
    guidance_scale=guidance_scale,
    width=width,
    height=height,
    skip_guidance_layers=[7, 8, 9],
).images[0]
image.save("minimal.png")
```

```
latent_model_input: torch.Size([2, 16, 128, 128])
timestep: torch.Size([2])
prompt_embeds: torch.Size([4, 333, 4096])
pooled_prompt_embeds: torch.Size([4, 2048])
```

Batch discrepancy, from 
https://github.com/huggingface/diffusers/blob/827b6c25f9b78a297345f356a7d152fd6faf27d8/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py#L912

Same discrepancy with num_images_per_prompt=1, I think in this case PyTorch broadcasting takes care of it
```
latent_model_input: torch.Size([1, 16, 128, 128])
timestep: torch.Size([1])
prompt_embeds: torch.Size([2, 333, 4096])
pooled_prompt_embeds: torch.Size([2, 2048])
```

Ends up with similar issue in the next section. I'll do more testing and put together a PR.
https://github.com/huggingface/diffusers/blob/827b6c25f9b78a297345f356a7d152fd6faf27d8/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py#L938-L946


Patch:
diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py
index a77231cdc02d..aee1ad8c75f5 100644
--- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py
+++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py
@@ -907,11 +907,7 @@ def __call__(
                     continue
 
                 # expand the latents if we are doing classifier free guidance
-                latent_model_input = (
-                    torch.cat([latents] * 2)
-                    if self.do_classifier_free_guidance and skip_guidance_layers is None
-                    else latents
-                )
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
                 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
                 timestep = t.expand(latent_model_input.shape[0])
 
@@ -935,6 +931,8 @@ def __call__(
                         else False
                     )
                     if skip_guidance_layers is not None and should_skip_layers:
+                        timestep = t.expand(latents.shape[0])
+                        latent_model_input = latents
                         noise_pred_skip_layers = self.transformer(
                             hidden_states=latent_model_input,
                             timestep=timestep,


Test patch: