[Feature Request] Add DiCache for sample-adaptive caching in FLUX transformers #15134

Diffusers already supports caching techniques such as FirstBlockCache, but additional speed–quality trade-offs may be possible through sample-adaptive cache decisions and trajectory-aligned residual reuse.

I would like to propose integrating [DiCache (ICLR 2026)](https://proceedings.iclr.cc/paper_files/paper/2026/hash/78288ef33b18a351c3cd679dc9a15c8d-Abstract-Conference.html), which uses shallow transformer probes to determine whether a full forward pass is necessary and applies trajectory-aligned reuse when skipping computation.

Describe the solution you'd like.

Add a DiCacheConfig compatible with Diffusers' existing caching API:

from diffusers import DiCacheConfig

pipe.transformer.enable_cache(
    DiCacheConfig(threshold=0.2, probe_depth=1)
)

# Generate images

pipe.transformer.disable_cache()

This is a proposed API, not an existing implementation. I propose an initial FLUX integration using the existing cache-hook and state-management architecture, without pipeline-specific monkey-patching.

Describe alternatives you've considered.

I benchmarked DiCache against Diffusers' FirstBlockCache and two DiCache ablations on pretrained FLUX.1-dev.

These results suggest that DiCache provides an additional operating region in the speed–fidelity trade-off. They do not establish superiority over a fully threshold-tuned FirstBlockCache Pareto frontier.

Additional context.

Experimental setup:

• Model: aniketppanchal/flux.1-dev-nf4-pkg, revision dde39613d8f3d7abe59d7c59ff16838b14bcebf8
• Hardware: Tesla T4
• Resolution: 512×512, 30 denoising steps
• Evaluation: 10 prompts × 3 seeds × 7 variants, with 2 timing repetitions each
• Versions: Diffusers 0.35.1, Transformers 4.56.1
• Measurements: denoising-only runtime, SSIM, PSNR, LPIPS, CLIP alignment, latent error, and GPU memory

Validation results:

All 210 method–prompt–seed configurations completed with zero recorded failures. Cache-disabled inference achieved exact tensor parity with native FLUX inference. Deterministic cache-reset and interleaved-generation tests also passed.

Per-timestep instrumentation demonstrated sample-dependent scheduling: DiCache produced 6 distinct cache schedules at threshold 0.2 and 3 at threshold 0.3 across 30 trajectories, compared with 1 schedule for the fixed-schedule control.

Ablations indicated improvements from both adaptive decisions and residual-history reuse.

Limitations: The experiments use an NF4-quantized checkpoint, not the paper's original full-precision setup. Fidelity can vary significantly between samples, with minimum SSIM values of 0.8120 and 0.7362 for DiCache thresholds 0.2 and 0.3, respectively. Additional matched-fidelity comparisons and broader model coverage would be valuable.

References:

• Paper: https://proceedings.iclr.cc/paper_files/paper/2026/hash/78288ef33b18a351c3cd679dc9a15c8d-Abstract-Conference.html
• Official implementation: https://github.com/Bujiazi/DiCache

I have reproducible benchmark scripts, paired images, cache trajectories, and validation records available.

I can prepare a minimal FLUX-first integration with the existing Diffusers caching infrastructure, together with regression tests and documentation. I would be happy to align the implementation with the maintainers' preferred design.

Would the maintainers be interested in this contribution? Any guidance on the preferred integration approach would be greatly appreciated.

If approved, I am willing to implement this feature and submit a pull request. I would also request this issue be assigned to me if possible.