How can I do image-to-image batched inference with FLUX.2 Klein, with one-to-one input image-prompt pairing?

[joshchristo]

I implemented something similar to the code below, but it appears that the model treats all provided images as shared context. As a result, each prompt is paired with all passed images (e.g., prompt_a is processed with image_a, image_b, and image_c).

What I am trying to achieve instead is batched inference with one-to-one pairing, so each prompt is processed only with its corresponding image (e.g., prompt_a with image_a, prompt_b with image_b, etc.).

What is the correct way to structure the inputs or batching logic to ensure this one-to-one mapping?

I would greatly appreciate any help with this. Thank you for your time.

import torch
from diffusers import Flux2KleinPipeline
from diffusers.utils import load_image

model_id = "black-forest-labs/FLUX.2-klein-9B"
pipe = Flux2KleinPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)

prompts = [
    "prompt_a",
    "prompt_b",
    "prompt_c"
]

images = [
    load_image("path/to/image_a.jpg"),
    load_image("path/to/image_b.jpg"),
    load_image("path/to/image_c.jpg")
]

outputs = pipe(
    prompt=prompts,
    image=images,
    num_inference_steps=4,  
    guidance_scale=1.0, 
).images

[missionfloyd]

You can iterate through the prompt/image pairs like so:

outputs = []

for image, prompt in zip(images, prompts):
    output = pipe(
        prompt=prompt,
        image=image,
        num_inference_steps=4,
        guidance_scale=1.0
    ).images[0]

    outputs.append(output)

[polinlin418-del]

I think, this is no batch processing.

[x-tahosin]

The issue you're hitting is fundamental to how Flux2KleinPipeline (and most diffusion pipelines) handle the image argument when it's a list: the pipeline broadcasts every image against every prompt, producing a Cartesian product, rather than zipping them one-to-one.

Why this happens

Internally, the pipeline stacks prompts and images into batch dimensions. When both are lists, the standard pipe(...) call doesn't have an explicit image_prompt_pairs concept — it treats prompt as batch dimension N and image as batch dimension M, giving you N×M outputs. That's why prompt_a gets paired with image_a, image_b, and image_c.

Solution 1: Manual one-to-one loop (simplest, slightly slower)

The most robust approach today is to call the pipeline individually for each pair and collect results:

outputs = []
for prompt, img in zip(prompts, images):
    out = pipe(
        prompt=prompt,
        image=img,
        num_inference_steps=4,
        guidance_scale=1.0,
    ).images[0]
    outputs.append(out)

This guarantees 1:1 pairing. The overhead is usually negligible for num_inference_steps=4 and a 9B model on a decent GPU, because the majority of cost is model forward passes, not Python loop overhead.

Solution 2: Collapse into a single batched call with paired indices

If you absolutely need one batched pipeline call (e.g. for VAE encode efficiency), you can duplicate prompts and images so each appears exactly once in the same batch order:

# Flatten: [prompt_a, prompt_b, prompt_c] and [img_a, img_b, img_c]
# become batch items 0,1,2 with direct correspondence
outputs = pipe(
    prompt=prompts,        # length 3
    image=images,          # length 3
    num_inference_steps=4,
    guidance_scale=1.0,
).images  # length 3

Wait — this is exactly what you tried. The key detail is whether Flux2KleinPipeline specifically respects list-length alignment or does implicit broadcasting. From the Diffusers source for most pipelines:

• prompt is processed through the text encoder → shape (batch_size, seq_len, dim)
• image is processed through the VAE encoder → shape (batch_size, channels, h, w)
• If both lists have the same length, they SHOULD zip naturally... unless the pipeline does an explicit itertools.product-style expansion internally.

For FluxPipeline and variants, the image-to-image path uses prepare_latents which typically expects image to be a single PIL image or a tensor. When a list is passed, it may be converting to a tensor stack via torch.stack([...]) — which DOES preserve order if the list lengths match.

The real fix

Check if you're using a custom wrapper or a community pipeline that adds broadcasting. The base DiffusionPipeline generally zips same-length lists. If you see cross-pairing, it's likely:

1. A custom pipeline implementation
2. You're passing a nested list structure that gets unpacked unexpectedly
3. There's a batch_size or num_images_per_prompt override being applied

Diagnostic step

Add a debug print right before the pipe() call:

print(f"Prompts: {len(prompts)}, Images: {len(images)}")
print(f"Prompt[0]: {prompts[0]}")
print(f"Image[0] size: {images[0].size}")

Then inspect the output length. If you get 9 images from 3 prompts + 3 images, you confirmed Cartesian-product behavior.

Recommended path forward

For production use with FLUX.2 Klein, I'd go with Solution 1 (explicit loop) until you profile it. With 4 steps and a 9B model, the loop overhead is likely <2% of total latency. If you need true batching for throughput, use the torch.compile() path on the pipeline and batch at the framework level rather than relying on list semantics.

If you can share:

• Exact diffusers version (diffusers.__version__)
• Whether you installed from main or a specific release
• The full traceback (if any) when running your original code

…I can verify whether this is a known issue in that specific version or a custom pipeline behavior.

[rootkiller6788]

This isn't a batching mistake on your side — it's just how this pipeline reads image. For Flux2KleinPipeline a list of
images isn't "one image per prompt": each image gets VAE-encoded and appended to a shared visual context that every
prompt in the call attends to. Internally prepare_image_latents concatenates all the reference images and repeats them
across the prompt batch, so 3 prompts + 3 images produces 3 outputs, each generated with all three images as context.
There's no input shape that tells it to pair prompt_a with only image_a — the batch axis is the prompts, and the
images are common context.

So for genuinely one-to-one you call it once per pair, like the loop in the earlier reply. It's not batched in the
throughput sense, but that isn't something to engineer around: one call can't carry a different image for each row, so
the loop is the only faithful way to give each prompt its own image.

One thing that might actually help you: the single batched call runs all N prompts in one pass, each against all M
images — handy when you want a group of prompts to share the same reference images (same style or subject), and it
saves VAE encoding of those images. It just can't give you per-prompt pairing.

[ZitouniNidhal]

To perform 1-to-1 paired batched inference in Diffusers pipelines, ensure that:

1. prompts and images are both standard Python lists of identical length: len(prompts) == len(images).
2. Each image in the list is pre-loaded as a PIL.Image (or a batched 4D PyTorch tensor [B, C, H, W]).

Here is the clean implementation for batching FLUX / Diffusers image-to-image:

import torch
from diffusers import Flux2KleinPipeline
from diffusers.utils import load_image

model_id = "black-forest-labs/FLUX.2-klein-9B"
pipe = Flux2KleinPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16).to("cuda")

prompts = [
    "A futuristic cyberpunk city",
    "A cozy wooden cabin in a snowy forest",
    "A watercolor painting of a sunflower"
]

images = [
    load_image("path/to/image_a.jpg").resize((1024, 1024)),
    load_image("path/to/image_b.jpg").resize((1024, 1024)),
    load_image("path/to/image_c.jpg").resize((1024, 1024))
]

# Pass lists directly; diffusers automatically batches item-by-item (batch_size = 3)
output_images = pipe(
    prompt=prompts,
    image=images,
    num_inference_steps=4,
    guidance_scale=1.0,
).images

# output_images[0] corresponds to prompts[0] + images[0], etc.

[tk1475]

Short answer: call it once per pair. A single call can't give each prompt its own image, so the zip loop from earlier in the thread is the correct way to get true one-to-one:

outputs = [
    pipe(prompt=p, image=img, num_inference_steps=4, guidance_scale=1.0).images[0]
    for p, img in zip(prompts, images)
]

Worth flagging that the answers saying equal-length lists pair item by item aren't right for this pipeline, which is exactly the cross-pairing you already saw. For FLUX.2 Klein a list of images is treated as shared reference context, so every prompt in the batch attends to all the images. That's useful when you want several prompts to share the same references, but it can't do per-prompt pairing.

Quick way to confirm on your setup: run your original batched call and check how many images come back and whether each one reflects only its matching input. If you need throughput later, batch groups that share the same reference and loop across the groups.