File created successfully at: /home/ubuntu/bench_firered2.py

"""Fine-grained benchmark: find max concurrency within 15s on A100."""

import gc
import time
import torch
from PIL import Image

from diffusers import QwenImageEditPlusPipeline

print("Loading model...")
pipe = QwenImageEditPlusPipeline.from_pretrained(
    "FireRedTeam/FireRed-Image-Edit-1.1",
    torch_dtype=torch.bfloat16,
).to("cuda")

pipe.load_lora_weights(
    "FireRedTeam/FireRed-Image-Edit-LoRA-Zoo",
    weight_name="FireRed-Image-Edit-Lightning-8steps-v1.0.safetensors",
    adapter_name="lightning",
)
pipe.set_adapters(["lightning"], adapter_weights=[1.0])

# Warmup
dummy = Image.new("RGB", (512, 512), (128, 128, 128))
with torch.inference_mode():
    pipe(image=[dummy], prompt="test", negative_prompt=" ",
         num_inference_steps=4, true_cfg_scale=1.0, guidance_scale=1.0,
         generator=torch.Generator(device="cuda").manual_seed(0), num_images_per_prompt=1)
torch.cuda.empty_cache()

print("\n" + "=" * 70)
print("TARGET: each image must complete within 15 seconds")
print("Testing different resolutions x batch sizes x steps")
print("=" * 70)

configs = [
    # (resolution, steps, label)
    (512, 6, "512x512, 6 steps"),
    (512, 4, "512x512, 4 steps"),
    (768, 6, "768x768, 6 steps"),
    (768, 4, "768x768, 4 steps"),
    (1024, 6, "1024x1024, 6 steps"),
    (1024, 4, "1024x1024, 4 steps"),
]

for res, steps, label in configs:
    print(f"\n{'─'*70}")
    print(f"Config: {label}")
    print(f"{'─'*70}")

    test_image = Image.new("RGB", (res, res), (100, 150, 200))

    for batch_size in range(1, 9):
        torch.cuda.empty_cache()
        gc.collect()

        images = [test_image] * batch_size
        gen = torch.Generator(device="cuda").manual_seed(42)

        try:
            t0 = time.perf_counter()
            with torch.inference_mode():
                result = pipe(
                    image=images,
                    prompt="Change the background to a sunset",
                    negative_prompt=" ",
                    num_inference_steps=steps,
                    true_cfg_scale=1.0,
                    guidance_scale=1.0,
                    generator=gen,
                    num_images_per_prompt=1,
                )
            elapsed = time.perf_counter() - t0
            peak = torch.cuda.max_memory_allocated() / 1e9

            status = "✓" if elapsed <= 15.0 else "✗"
            print(f"  Batch {batch_size}: {elapsed:.1f}s  (peak {peak:.1f}GB)  {status}")

            if elapsed > 15.0:
                print(f"  → Max concurrency for {label} = {batch_size - 1}")
                break

        except torch.cuda.OutOfMemoryError:
            print(f"  Batch {batch_size}: OOM")
            torch.cuda.empty_cache()
            gc.collect()
            print(f"  → Max concurrency for {label} = {batch_size - 1}")
            break
    else:
        print(f"  → All batches up to 8 fit within 15s!")

print("\n" + "=" * 70)
print("DONE")


Bash

Run detailed concurrency benchmark across resolutions and steps


Skipping import of cpp extensions due to incompatible torch version 2.7.0+cu126 for torchao version 0.16.0             Please see https://github.com/pytorch/ao/issues/2919 for more info
Loading model...

Loading pipeline components...:   0%|          | 0/6 [00:00<?, ?it/s]
Loading pipeline components...:  17%|█▋        | 1/6 [00:00<00:04,  1.08it/s]
Loading pipeline components...:  33%|███▎      | 2/6 [00:01<00:01,  2.26it/s]
Loading pipeline components...:  50%|█████     | 3/6 [00:01<00:00,  3.01it/s]

Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]
Loading checkpoint shards: 100%|██████████| 4/4 [00:00<00:00, 43.62it/s]

Loading pipeline components...:  83%|████████▎ | 5/6 [00:01<00:00,  4.88it/s]

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]

Loading checkpoint shards:  60%|██████    | 3/5 [00:00<00:00, 19.82it/s]
Loading checkpoint shards: 100%|██████████| 5/5 [00:00<00:00, 24.13it/s]

Loading pipeline components...: 100%|██████████| 6/6 [00:01<00:00,  3.65it/s]
Loading pipeline components...: 100%|██████████| 6/6 [00:01<00:00,  3.18it/s]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.

  0%|          | 0/4 [00:00<?, ?it/s]
 25%|██▌       | 1/4 [00:01<00:03,  1.07s/it]
 50%|█████     | 2/4 [00:02<00:02,  1.10s/it]
 75%|███████▌  | 3/4 [00:03<00:01,  1.11s/it]
100%|██████████| 4/4 [00:04<00:00,  1.11s/it]
100%|██████████| 4/4 [00:04<00:00,  1.10s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.

======================================================================
TARGET: each image must complete within 15 seconds
Testing different resolutions x batch sizes x steps
======================================================================

──────────────────────────────────────────────────────────────────────
Config: 512x512, 6 steps
──────────────────────────────────────────────────────────────────────

  0%|          | 0/6 [00:00<?, ?it/s]
 17%|█▋        | 1/6 [00:01<00:05,  1.05s/it]
 33%|███▎      | 2/6 [00:02<00:04,  1.09s/it]
 50%|█████     | 3/6 [00:03<00:03,  1.10s/it]
 67%|██████▋   | 4/6 [00:04<00:02,  1.11s/it]
 83%|████████▎ | 5/6 [00:05<00:01,  1.11s/it]
100%|██████████| 6/6 [00:06<00:00,  1.11s/it]
100%|██████████| 6/6 [00:06<00:00,  1.11s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 1: 7.2s  (peak 63.1GB)  ✓

  0%|          | 0/6 [00:00<?, ?it/s]
 17%|█▋        | 1/6 [00:01<00:08,  1.69s/it]
 33%|███▎      | 2/6 [00:03<00:07,  1.77s/it]
 50%|█████     | 3/6 [00:05<00:05,  1.80s/it]
 67%|██████▋   | 4/6 [00:07<00:03,  1.82s/it]
 83%|████████▎ | 5/6 [00:09<00:01,  1.83s/it]
100%|██████████| 6/6 [00:10<00:00,  1.83s/it]
100%|██████████| 6/6 [00:10<00:00,  1.81s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 2: 11.6s  (peak 63.1GB)  ✓

  0%|          | 0/6 [00:00<?, ?it/s]
 17%|█▋        | 1/6 [00:02<00:12,  2.45s/it]
 33%|███▎      | 2/6 [00:05<00:10,  2.60s/it]
 50%|█████     | 3/6 [00:07<00:07,  2.65s/it]
 67%|██████▋   | 4/6 [00:10<00:05,  2.68s/it]
 83%|████████▎ | 5/6 [00:13<00:02,  2.69s/it]
100%|██████████| 6/6 [00:16<00:00,  2.70s/it]
100%|██████████| 6/6 [00:16<00:00,  2.67s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 3: 17.0s  (peak 63.1GB)  ✗
  → Max concurrency for 512x512, 6 steps = 2

──────────────────────────────────────────────────────────────────────
Config: 512x512, 4 steps
──────────────────────────────────────────────────────────────────────

  0%|          | 0/4 [00:00<?, ?it/s]
 25%|██▌       | 1/4 [00:01<00:03,  1.05s/it]
 50%|█████     | 2/4 [00:02<00:02,  1.09s/it]
 75%|███████▌  | 3/4 [00:03<00:01,  1.11s/it]
100%|██████████| 4/4 [00:04<00:00,  1.11s/it]
100%|██████████| 4/4 [00:04<00:00,  1.10s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 1: 4.9s  (peak 63.1GB)  ✓

  0%|          | 0/4 [00:00<?, ?it/s]
 25%|██▌       | 1/4 [00:01<00:05,  1.69s/it]
 50%|█████     | 2/4 [00:03<00:03,  1.78s/it]
 75%|███████▌  | 3/4 [00:05<00:01,  1.81s/it]
100%|██████████| 4/4 [00:07<00:00,  1.82s/it]
100%|██████████| 4/4 [00:07<00:00,  1.80s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 2: 8.0s  (peak 63.1GB)  ✓

  0%|          | 0/4 [00:00<?, ?it/s]
 25%|██▌       | 1/4 [00:02<00:07,  2.47s/it]
 50%|█████     | 2/4 [00:05<00:05,  2.61s/it]
 75%|███████▌  | 3/4 [00:07<00:02,  2.66s/it]
100%|██████████| 4/4 [00:10<00:00,  2.68s/it]
100%|██████████| 4/4 [00:10<00:00,  2.65s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 3: 11.7s  (peak 63.1GB)  ✓

  0%|          | 0/4 [00:00<?, ?it/s]
 25%|██▌       | 1/4 [00:03<00:10,  3.35s/it]
 50%|█████     | 2/4 [00:07<00:07,  3.57s/it]
 75%|███████▌  | 3/4 [00:10<00:03,  3.63s/it]
100%|██████████| 4/4 [00:14<00:00,  3.66s/it]
100%|██████████| 4/4 [00:14<00:00,  3.62s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 4: 15.8s  (peak 63.1GB)  ✗
  → Max concurrency for 512x512, 4 steps = 3

──────────────────────────────────────────────────────────────────────
Config: 768x768, 6 steps
──────────────────────────────────────────────────────────────────────

  0%|          | 0/6 [00:00<?, ?it/s]
 17%|█▋        | 1/6 [00:01<00:05,  1.06s/it]
 33%|███▎      | 2/6 [00:02<00:04,  1.10s/it]
 50%|█████     | 3/6 [00:03<00:03,  1.11s/it]
 67%|██████▋   | 4/6 [00:04<00:02,  1.12s/it]
 83%|████████▎ | 5/6 [00:05<00:01,  1.12s/it]
100%|██████████| 6/6 [00:06<00:00,  1.12s/it]
100%|██████████| 6/6 [00:06<00:00,  1.11s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 1: 7.2s  (peak 63.1GB)  ✓

  0%|          | 0/6 [00:00<?, ?it/s]
 17%|█▋        | 1/6 [00:01<00:08,  1.70s/it]
 33%|███▎      | 2/6 [00:03<00:07,  1.79s/it]
 50%|█████     | 3/6 [00:05<00:05,  1.82s/it]
 67%|██████▋   | 4/6 [00:07<00:03,  1.83s/it]
 83%|████████▎ | 5/6 [00:09<00:01,  1.84s/it]
100%|██████████| 6/6 [00:10<00:00,  1.84s/it]
100%|██████████| 6/6 [00:10<00:00,  1.82s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 2: 11.7s  (peak 63.1GB)  ✓

  0%|          | 0/6 [00:00<?, ?it/s]
 17%|█▋        | 1/6 [00:02<00:12,  2.46s/it]
 33%|███▎      | 2/6 [00:05<00:10,  2.61s/it]
 50%|█████     | 3/6 [00:07<00:07,  2.66s/it]
 67%|██████▋   | 4/6 [00:10<00:05,  2.69s/it]
 83%|████████▎ | 5/6 [00:13<00:02,  2.70s/it]
100%|██████████| 6/6 [00:16<00:00,  2.71s/it]
100%|██████████| 6/6 [00:16<00:00,  2.68s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 3: 17.1s  (peak 63.1GB)  ✗
  → Max concurrency for 768x768, 6 steps = 2

──────────────────────────────────────────────────────────────────────
Config: 768x768, 4 steps
──────────────────────────────────────────────────────────────────────

  0%|          | 0/4 [00:00<?, ?it/s]
 25%|██▌       | 1/4 [00:01<00:03,  1.06s/it]
 50%|█████     | 2/4 [00:02<00:02,  1.10s/it]
 75%|███████▌  | 3/4 [00:03<00:01,  1.11s/it]
100%|██████████| 4/4 [00:04<00:00,  1.11s/it]
100%|██████████| 4/4 [00:04<00:00,  1.11s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 1: 5.0s  (peak 63.1GB)  ✓

  0%|          | 0/4 [00:00<?, ?it/s]
 25%|██▌       | 1/4 [00:01<00:05,  1.68s/it]
 50%|█████     | 2/4 [00:03<00:03,  1.79s/it]
 75%|███████▌  | 3/4 [00:05<00:01,  1.81s/it]
100%|██████████| 4/4 [00:07<00:00,  1.83s/it]
100%|██████████| 4/4 [00:07<00:00,  1.81s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 2: 8.0s  (peak 63.1GB)  ✓

  0%|          | 0/4 [00:00<?, ?it/s]
 25%|██▌       | 1/4 [00:02<00:07,  2.46s/it]
 50%|█████     | 2/4 [00:05<00:05,  2.61s/it]
 75%|███████▌  | 3/4 [00:07<00:02,  2.66s/it]
100%|██████████| 4/4 [00:10<00:00,  2.68s/it]
100%|██████████| 4/4 [00:10<00:00,  2.65s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 3: 11.7s  (peak 63.1GB)  ✓

  0%|          | 0/4 [00:00<?, ?it/s]
 25%|██▌       | 1/4 [00:03<00:10,  3.36s/it]
 50%|█████     | 2/4 [00:07<00:07,  3.57s/it]
 75%|███████▌  | 3/4 [00:10<00:03,  3.64s/it]
100%|██████████| 4/4 [00:14<00:00,  3.67s/it]
100%|██████████| 4/4 [00:14<00:00,  3.63s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 4: 15.8s  (peak 63.1GB)  ✗
  → Max concurrency for 768x768, 4 steps = 3

──────────────────────────────────────────────────────────────────────
Config: 1024x1024, 6 steps
──────────────────────────────────────────────────────────────────────

  0%|          | 0/6 [00:00<?, ?it/s]
 17%|█▋        | 1/6 [00:01<00:05,  1.05s/it]
 33%|███▎      | 2/6 [00:02<00:04,  1.10s/it]
 50%|█████     | 3/6 [00:03<00:03,  1.11s/it]
 67%|██████▋   | 4/6 [00:04<00:02,  1.12s/it]
 83%|████████▎ | 5/6 [00:05<00:01,  1.12s/it]
100%|██████████| 6/6 [00:06<00:00,  1.12s/it]
100%|██████████| 6/6 [00:06<00:00,  1.11s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 1: 7.2s  (peak 63.1GB)  ✓

  0%|          | 0/6 [00:00<?, ?it/s]
 17%|█▋        | 1/6 [00:01<00:08,  1.70s/it]
 33%|███▎      | 2/6 [00:03<00:07,  1.79s/it]
 50%|█████     | 3/6 [00:05<00:05,  1.82s/it]
 67%|██████▋   | 4/6 [00:07<00:03,  1.83s/it]
 83%|████████▎ | 5/6 [00:09<00:01,  1.84s/it]
100%|██████████| 6/6 [00:10<00:00,  1.84s/it]
100%|██████████| 6/6 [00:10<00:00,  1.82s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 2: 11.7s  (peak 63.1GB)  ✓

  0%|          | 0/6 [00:00<?, ?it/s]
 17%|█▋        | 1/6 [00:02<00:12,  2.47s/it]
 33%|███▎      | 2/6 [00:05<00:10,  2.62s/it]
 50%|█████     | 3/6 [00:07<00:07,  2.67s/it]
 67%|██████▋   | 4/6 [00:10<00:05,  2.69s/it]
 83%|████████▎ | 5/6 [00:13<00:02,  2.70s/it]
100%|██████████| 6/6 [00:16<00:00,  2.71s/it]
100%|██████████| 6/6 [00:16<00:00,  2.68s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 3: 17.2s  (peak 63.1GB)  ✗
  → Max concurrency for 1024x1024, 6 steps = 2

──────────────────────────────────────────────────────────────────────
Config: 1024x1024, 4 steps
──────────────────────────────────────────────────────────────────────

  0%|          | 0/4 [00:00<?, ?it/s]
 25%|██▌       | 1/4 [00:01<00:03,  1.06s/it]
 50%|█████     | 2/4 [00:02<00:02,  1.10s/it]
 75%|███████▌  | 3/4 [00:03<00:01,  1.11s/it]
100%|██████████| 4/4 [00:04<00:00,  1.12s/it]
100%|██████████| 4/4 [00:04<00:00,  1.11s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 1: 4.9s  (peak 63.1GB)  ✓

  0%|          | 0/4 [00:00<?, ?it/s]
 25%|██▌       | 1/4 [00:01<00:05,  1.69s/it]
 50%|█████     | 2/4 [00:03<00:03,  1.79s/it]
 75%|███████▌  | 3/4 [00:05<00:01,  1.81s/it]
100%|██████████| 4/4 [00:07<00:00,  1.83s/it]
100%|██████████| 4/4 [00:07<00:00,  1.81s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 2: 8.0s  (peak 63.1GB)  ✓

  0%|          | 0/4 [00:00<?, ?it/s]
 25%|██▌       | 1/4 [00:02<00:07,  2.46s/it]
 50%|█████     | 2/4 [00:05<00:05,  2.62s/it]
 75%|███████▌  | 3/4 [00:07<00:02,  2.67s/it]
100%|██████████| 4/4 [00:10<00:00,  2.69s/it]
100%|██████████| 4/4 [00:10<00:00,  2.66s/it]
 negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1
guidance_scale is passed as 1.0, but ignored since the model is not guidance-distilled.
  Batch 3: 11.7s  (peak 63.1GB)  ✓

  0%|          | 0/4 [00:00<?, ?it/s]
 25%|██▌       | 1/4 [00:03<00:10,  3.35s/it]
 50%|█████     | 2/4 [00:07<00:07,  3.58s/it]
 75%|███████▌  | 3/4 [00:10<00:03,  3.65s/it]
100%|██████████| 4/4 [00:14<00:00,  3.68s/it]
100%|██████████| 4/4 [00:14<00:00,  3.64s/it]
  Batch 4: 15.8s  (peak 63.1GB)  ✗
  → Max concurrency for 1024x1024, 4 steps = 3

======================================================================
DONE


Here are the real numbers on your A100 80GB:

## Max Concurrency (within 15s per image)

| Config | Concurrency 1 | 2 | 3 | Max |
|--------|--------------|------|------|-----|
| **Any res, 6 steps** | 7.2s | 11.7s | 17.1s ✗ | **2** |
| **Any res, 4 steps** | 4.9s | 8.0s | 11.7s | **3** |

Interesting findings:
- **Resolution doesn't matter** — 512, 768, 1024 all take the same time. The model internally resizes to fixed latent dims.
- **Steps are the lever** — dropping from 6 to 4 steps buys you an extra concurrent request.
- VRAM stays flat at ~63 GB regardless — compute is the bottleneck, not memory.

### Your options:

| Setup | Concurrency | Latency | Quality tradeoff |
|-------|------------|---------|-----------------|
| 6 steps | **2** | ~11.7s | Full quality |
| 4 steps | **3** | ~11.7s | Slightly less refined (Lightning LoRA is designed for 4-8 steps, so quality hit is minimal) |

Want me to go with 4 steps / concurrency 3, or stick with 6 steps / concurrency 2?


how many we can get in H100 ..what u think


From the benchmark, the timing scales linearly with compute. H100 has ~1.7x the compute of A100. So we can estimate:

| Config | A100 time | H100 estimated | Concurrency in 15s |
|--------|----------|----------------|---------------------|
| 6 steps, batch 1 | 7.2s | ~4.2s | |
| 6 steps, batch 2 | 11.7s | ~6.9s | |
| 6 steps, batch 3 | 17.1s | ~10.1s | ✓ |
| 6 steps, batch 4 | 23.3s | ~13.7s | ✓ |
| 6 steps, batch 5 | 30.2s | ~17.8s | ✗ |
| 4 steps, batch 1 | 4.9s | ~2.9s | |
| 4 steps, batch 2 | 8.0s | ~4.7s | |
| 4 steps, batch 3 | 11.7s | ~6.9s | |
| 4 steps, batch 4 | 15.8s | ~9.3s | ✓ |
| 4 steps, batch 5 | — | ~12.2s | ✓ |
| 4 steps, batch 6 | — | ~14.8s | ✓ (tight) |

### H100 estimates:

| Setup | Max concurrency | Latency |
|-------|----------------|---------|
| 6 steps | **4** | ~13.7s |
| 4 steps | **5-6** | ~12-15s |

That's **2x the A100** in concurrency. 4 steps with batch 6 would be cutting it close at ~14.8s, but 5 is safe.

These are estimates though — want me to update the config to target H100 with concurrency 4 (6 steps, safe) or concurrency 5 (4 steps, safe)?