FLUX.1-Kontext-dev batch inference throughput issue

[Django-Jiang]

Describe the bug

I am trying to use FLUX.1-Kontext-dev in batch inference, but the inference time for batch is linear to batch size, e.g., time of batch size 2 ~= 2* time of batch size 1. This makes me confused, did I get something wrong or this is the issue with this pipeline?

Reproduction

import torch
from diffusers import FluxKontextPipeline
from diffusers.utils import load_image


# Load the pipeline
pipe = FluxKontextPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Kontext-dev", 
    torch_dtype=torch.bfloat16
)
pipe.to("cuda")

input_images = [
    load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"),
    load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png")
]

prompts = [
    "Make it look like a watercolor painting",
    "Add dramatic lighting and shadows",
]
with torch.inference_mode():
    images = pipe(
        image=input_images,
        prompt=prompts,
        guidance_scale=2.5,
        num_inference_steps=28,
        generator=[torch.Generator("cuda").manual_seed(42),torch.Generator("cuda").manual_seed(123)]
    ).images

for i, img in enumerate(images):
    img.save(f"output_i2i_batch_{i}.png")

Logs

System Info

• 🤗 Diffusers version: 0.36.0.dev0
• Platform: Linux-6.8.0-63-generic-x86_64-with-glibc2.39
• Running on Google Colab?: No
• Python version: 3.10.18
• PyTorch version (GPU?): 2.8.0+cu128 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Huggingface_hub version: 0.35.3
• Transformers version: 4.57.0
• Accelerate version: 1.10.1
• PEFT version: not installed
• Bitsandbytes version: not installed
• Safetensors version: 0.6.2
• xFormers version: not installed
• Accelerator: NVIDIA H100 NVL, 95830 MiB
  NVIDIA H100 NVL, 95830 MiB
  NVIDIA H100 NVL, 95830 MiB
  NVIDIA H100 NVL, 95830 MiB
  NVIDIA H100 NVL, 95830 MiB
  NVIDIA H100 NVL, 95830 MiB
  NVIDIA H100 NVL, 95830 MiB
  NVIDIA H100 NVL, 95830 MiB

Who can help?

@yiyixuxu @DN6

[SahilCarterr]

Hey! @Django-Jiang

diffusers/src/diffusers/pipelines/flux/pipeline_flux_kontext.py

Lines 1087 to 1099 in 693d8a3

	noise_pred = self.transformer(
	hidden_states=latent_model_input,
	timestep=timestep / 1000,
	guidance=guidance,
	pooled_projections=pooled_prompt_embeds,
	encoder_hidden_states=prompt_embeds,
	txt_ids=text_ids,
	img_ids=latent_ids,
	joint_attention_kwargs=self.joint_attention_kwargs,
	return_dict=False,
	)[0]
	noise_pred = noise_pred[:, : latents.size(1)]

self.transformer() in FluxKontext (and similar Flux models) do not use true batch-parallelism for independent samples.It performs cross-attention and image conditioning that depends on each image’s token sequence.Still latents has shape [B, ...], the loops over B internally.
that is why

time(batch=2) ≈ 2 × time(batch=1)

[xidiancpy]

Hey! @Django-Jiang

diffusers/src/diffusers/pipelines/flux/pipeline_flux_kontext.py

Lines 1087 to 1099 in 693d8a3

noise_pred = self.transformer(
hidden_states=latent_model_input,
timestep=timestep / 1000,
guidance=guidance,
pooled_projections=pooled_prompt_embeds,
encoder_hidden_states=prompt_embeds,
txt_ids=text_ids,
img_ids=latent_ids,
joint_attention_kwargs=self.joint_attention_kwargs,
return_dict=False,
)[0]
noise_pred = noise_pred[:, : latents.size(1)]

self.transformer() in FluxKontext (and similar Flux models) do not use true batch-parallelism for independent samples.It performs cross-attention and image conditioning that depends on each image’s token sequence.Still latents has shape [B, ...], the loops over B internally.
that is why

time(batch=2) ≈ 2 × time(batch=1)

@SahilCarterr Your response was very helpful to me because I also encountered the same problem. I would like to ask where the code implementation for "loops over B" is located? Because I have never seen code like this loop.

[xidiancpy]

@Django-Jiang Is there any way you can share your thoughts on this issue? Thank you.

[Django-Jiang]

@Django-Jiang Is there any way you can share your thoughts on this issue? Thank you.

I am not an expert in this domain, but I think it is an implementation design issue, and it would be better to support this feature