RuntimeError with LyCORIS, Batch Inference and skip_guidance_layers

[japppie]

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:

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

[sayakpaul]

This is not a fully reproducible snippet. Please provide one.

Cc: @asomoza and @yiyixuxu for skip layer guidance.

[hlky]

Reproducible with:

pip install lycoris_lora

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]

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

[sayakpaul]

So, SLG seems to be the culprit here?

[japppie]

So, SLG seems to be the culprit here?

Indeed, problem seems to appear since I started using SLG

[hlky]

After some testing lycoris is unrelated, minimally reproducible with

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

diffusers/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py

Line 912 in 827b6c2

if self.do_classifier_free_guidance and skip_guidance_layers is None

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.

diffusers/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py

Lines 938 to 946 in 827b6c2

	noise_pred_skip_layers = self.transformer(
	hidden_states=latent_model_input,
	timestep=timestep,
	encoder_hidden_states=original_prompt_embeds,
	pooled_projections=original_pooled_prompt_embeds,
	joint_attention_kwargs=self.joint_attention_kwargs,
	return_dict=False,
	skip_layers=skip_guidance_layers,
	)[0]