Skip to content

[Bug] FLUX.1-Kontext-dev multiple-batch inference #864

Description

@dangph-alala

Checklist

Describe the Bug

Multi-batch inference (batch_size=2) is slower than running 2 separate single-batch inferences
Although I saw: #289 (comment)
Is this caused by the gpu or by my code?
omg: #398
@lmxyy

Environment

nunchaku 1.0.0+torch2.6
torch 2.6.0+cu124
python 3.11
vram rtx 4070ti super 16gb
ram 32gb

Reproduction Steps

transformer = NunchakuFluxTransformer2dModel.from_pretrained(
    f"mit-han-lab/nunchaku-flux.1-kontext-dev/svdq-{get_precision()}_r32-flux.1-kontext-dev.safetensors"
)
transformer.set_attention_impl("nunchaku-fp16")  # set attention implementation to fp16

text_encoder = NunchakuT5EncoderModel.from_pretrained(
    "mit-han-lab/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors"
)

pipeline = FluxKontextPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Kontext-dev", 
    transformer=transformer, text_encoder_2=text_encoder,
    torch_dtype=torch.bfloat16
).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",
]

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions