Calibration Cache with Auto Caching (Pro)

This tutorial shows how to use calibration cache mode with auto caching on FLUX.1-dev. Unlike Taylor or reuse modes, calibration fits linear coefficients from real denoiser outputs on a small prompt set.

Attach a caption datamodule to SmashConfig before smashing a fresh model. Coefficients are stored in memory on pipe.cache_helper.calibrated_coeffs and persisted as caching_coeffs.pkl when you save. On load, inference reuses those coefficients only — recalibration is not supported (prompt data is not saved with the model).

1. Loading the Flux Model

First, load FLUX.1-dev.

import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.to("cuda")
# pipe.enable_model_cpu_offload()  # save VRAM; remove if you have enough GPU memory

2. Calibration Prompts and Smash Config

DrawBench is test-only in Pruna. We shuffle the test split with a fixed seed and keep a subset of prompts for coefficient fitting.

from pruna.data.pruna_datamodule import PrunaDataModule

from pruna_pro import SmashConfig

CALIBRATION_SEED = 42
NUM_CALIBRATION_PROMPTS = 50

drawbench = PrunaDataModule.from_string("DrawBench", seed=CALIBRATION_SEED)
calibration_prompts = drawbench.test_dataset.shuffle(seed=CALIBRATION_SEED).select(
    range(NUM_CALIBRATION_PROMPTS)
)
_ = calibration_prompts.select([0])

smash_config = SmashConfig({
    "auto": {
        "cache_mode": "calibration",
        "speed_factor": 0.4,
    },
    "torch_compile": {},
})
smash_config.add_data((calibration_prompts, _, _), "prompt_collate")

3. Smashing the Model

Replace the token with the one provided by PrunaAI.

from pruna_pro import smash

pipe = smash(
    model=pipe,
    token="<your_pruna_token>",
    smash_config=smash_config,
)

4. Running Inference

The first call with a given num_inference_steps calibrates coefficients (full generations over the prompt set, then a least-squares fit), then runs your request. Later calls with the same step count and schedule will reuse the in-memory coefficients.

prompt = (
    "An anime illustration of Sydney Opera House sitting next to Eiffel tower, under a blue night sky of "
    "roiling energy, exploding yellow stars, and radiating swirls of blue."
)
num_inference_steps = 50

# First call: calibrates coefficients for this step count (slow).
pipe(prompt, num_inference_steps=num_inference_steps).images[0]

# Second call: reuses calibrated coefficients (fast).
pipe(prompt, num_inference_steps=num_inference_steps).images[0]

5. Inspecting Calibrated Coefficients

Fitted coefficients live on the cache helper. Keys are (schedule_id, calibration_order) where schedule_id encodes the caching schedule (which steps compute vs reuse). Values are CalibratedCoeffs objects with the per-step linear weights used at inference.

After saving the smashed model, the same data is written to caching_coeffs.pkl and restored on load.

pipe.cache_helper.calibrated_coeffs

The speed factor can be adjusted on the fly. A new schedule triggers recalibration (the prompt dataset is reused; only the fit is rerun).

pipe.cache_helper.set_params(speed_factor=0.33)
pipe(prompt, num_inference_steps=num_inference_steps).images[0]

pipe.cache_helper.calibrated_coeffs

6. Calibration Order

calibration_order controls how many prior compute steps are used to approximate each skipped step (default 2). Changing it adds a new entry under calibrated_coeffs. Use a fixed generator seed when comparing orders on the same prompt.

pipe.cache_helper.set_params(calibration_order=3, speed_factor=0.25)
pipe(
    prompt,
    num_inference_steps=num_inference_steps,
    generator=torch.Generator(device="cuda").manual_seed(42),
).images[0]

pipe.cache_helper.calibrated_coeffs

calibration_order=None uses all prior compute steps in the fit. This can help on some schedules but is more expensive and may overfit small calibration sets.

pipe.cache_helper.set_params(calibration_order=None)
pipe(
    prompt,
    num_inference_steps=num_inference_steps,
    generator=torch.Generator(device="cuda").manual_seed(42),
).images[0]

pipe.cache_helper.calibrated_coeffs

Changing num_inference_steps builds or reuses a calibration dataset for that length, then fits coefficients for the current schedule and order.

pipe.cache_helper.set_params(calibration_order=15)
pipe(
    prompt,
    num_inference_steps=28,
    generator=torch.Generator(device="cuda").manual_seed(42),
).images[0]

pipe.cache_helper.calibrated_coeffs

7. Saving the Smashed Model

Saving writes caching_coeffs.pkl and registers caching_coeffs in smash_config.load_artifacts_fns so coefficients are restored automatically on load.

save_dir = "flux-calibration-cache"
pipe.save_pretrained(save_dir)

8. Loading and Inference-Only Use

Reload with from_pretrained. No prompt datamodule is required: load_caching_coeffs restores pipe.cache_helper.calibrated_coeffs from disk.

Before loading, free the in-memory model (del pipe, then gc.collect() and torch.cuda.empty_cache()) so reload does not run out of memory.

Loaded models are inference-only for calibration mode. Changing schedule/order without matching persisted coefficients raises an error.

import gc

from pruna_pro import PrunaProModel

del pipe
gc.collect()
torch.cuda.empty_cache()

loaded = PrunaProModel.from_pretrained(save_dir)
loaded.cache_helper.calibrated_coeffs

# Inference reuses persisted coefficients (no calibration dataset needed).
loaded(
    prompt,
    num_inference_steps=num_inference_steps,
    generator=torch.Generator(device="cuda").manual_seed(42),
).images[0]

Wrap Up

You enabled auto caching with calibration mode, calibrated on DrawBench prompts during a fresh smash, saved caching_coeffs.pkl, and loaded the model for inference-only use without re-attaching data.