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