Outrider brief — Add ChebBooster (Chebyshev feature cache) to 🤗 diffusers

Author: remyx-ai[bot]

benchmarks/benchmarking_chebyshev_cache.py

@@ -0,0 +1,183 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Control harness for the ChebBooster (Chebyshev cache) PR.
+
+Three-arm comparison on FLUX.1-dev with identical seeds/prompts and matched
+`cache_interval` + warmup:
+
+  (a) no-cache reference
+  (b) in-repo TaylorSeer cache (the control arm)
+  (c) Chebyshev cache (ChebBooster port)
+
+Reports wall-clock latency per arm and quality of (b) and (c) against the
+no-cache reference (PSNR always; LPIPS and CLIP score when `lpips` /
+`open_clip_torch` are installed). The claim this harness is meant to check:
+at matched latency, Chebyshev cache >= TaylorSeer cache on quality. If the
+delta is within noise, the cache variant is redundant and the PR should not
+be opened upstream.
+
+Usage:
+    python benchmarking_chebyshev_cache.py --num-prompts 8 --num-inference-steps 28
+"""
+
+import argparse
+import csv
+import time
+
+import torch
+
+from diffusers import ChebyshevCacheConfig, FluxPipeline, TaylorSeerCacheConfig
+
+
+CKPT_ID = "black-forest-labs/FLUX.1-dev"
+RESULT_FILENAME = "chebyshev_cache_control.csv"
+
+# Placeholder prompt set — replace with the DrawBench prompts shipped in the
+# reference repo (ChebBooster-FLUX/DrawBench.jsonl) for the full control run.
+PROMPTS = [
+    "A photograph of an astronaut riding a horse on Mars.",
+    "A bowl of fruit on a wooden table, studio lighting.",
+    "A cyberpunk city street at night in the rain.",
+    "An oil painting of a lighthouse during a storm.",
+]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num-prompts", type=int, default=len(PROMPTS))
+    parser.add_argument("--num-inference-steps", type=int, default=28)
+    parser.add_argument("--cache-interval", type=int, default=5)
+    parser.add_argument("--disable-cache-before-step", type=int, default=3)
+    parser.add_argument("--cheb-order", type=int, default=6)
+    parser.add_argument("--seed", type=int, default=0)
+    return parser.parse_args()
+
+
+def load_pipeline():
+    pipe = FluxPipeline.from_pretrained(CKPT_ID, torch_dtype=torch.bfloat16)
+    pipe.to("cuda")
+    pipe.set_progress_bar_config(disable=True)
+    return pipe
+
+
+def run_arm(pipe, prompts, num_inference_steps, seed, cache_config=None):
+    if cache_config is not None:
+        pipe.transformer.enable_cache(cache_config)
+
+    images = []
+    start = time.perf_counter()
+    for i, prompt in enumerate(prompts):
+        generator = torch.Generator(device="cuda").manual_seed(seed + i)
+        image = pipe(prompt, num_inference_steps=num_inference_steps, generator=generator).images[0]
+        images.append(image)
+    torch.cuda.synchronize()
+    latency = (time.perf_counter() - start) / len(prompts)
+
+    if cache_config is not None:
+        pipe.transformer.disable_cache()
+    return images, latency
+
+
+def psnr_vs_reference(images, reference_images):
+    values = []
+    for image, reference in zip(images, reference_images):
+        x = torch.tensor(list(image.getdata()), dtype=torch.float32) / 255.0
+        y = torch.tensor(list(reference.getdata()), dtype=torch.float32) / 255.0
+        mse = torch.mean((x - y) ** 2).item()
+        values.append(float("inf") if mse == 0 else 10.0 * torch.log10(torch.tensor(1.0 / mse)).item())
+    return sum(values) / len(values)
+
+
+def lpips_vs_reference(images, reference_images):
+    try:
+        import lpips
+        import numpy as np
+    except ImportError:
+        return None
+
+    loss_fn = lpips.LPIPS(net="vgg").to("cuda")
+    values = []
+    for image, reference in zip(images, reference_images):
+        x = torch.from_numpy(np.array(image)).permute(2, 0, 1)[None].float().to("cuda") / 127.5 - 1.0
+        y = torch.from_numpy(np.array(reference)).permute(2, 0, 1)[None].float().to("cuda") / 127.5 - 1.0
+        with torch.no_grad():
+            values.append(loss_fn(x, y).item())
+    return sum(values) / len(values)
+
+
+def main():
+    args = parse_args()
+    prompts = PROMPTS[: args.num_prompts]
+    shared_kwargs = {
+        "cache_interval": args.cache_interval,
+        "disable_cache_before_step": args.disable_cache_before_step,
+    }
+
+    pipe = load_pipeline()
+
+    reference_images, reference_latency = run_arm(pipe, prompts, args.num_inference_steps, args.seed)
+
+    taylorseer_images, taylorseer_latency = run_arm(
+        pipe,
+        prompts,
+        args.num_inference_steps,
+        args.seed,
+        cache_config=TaylorSeerCacheConfig(max_order=1, **shared_kwargs),
+    )
+
+    chebyshev_images, chebyshev_latency = run_arm(
+        pipe,
+        prompts,
+        args.num_inference_steps,
+        args.seed,
+        cache_config=ChebyshevCacheConfig(cheb_order=args.cheb_order, **shared_kwargs),
+    )
+
+    rows = [
+        {
+            "arm": "no-cache",
+            "latency_s_per_image": reference_latency,
+            "speedup": 1.0,
+            "psnr_vs_reference": float("inf"),
+            "lpips_vs_reference": 0.0,
+        },
+        {
+            "arm": "taylorseer",
+            "latency_s_per_image": taylorseer_latency,
+            "speedup": reference_latency / taylorseer_latency,
+            "psnr_vs_reference": psnr_vs_reference(taylorseer_images, reference_images),
+            "lpips_vs_reference": lpips_vs_reference(taylorseer_images, reference_images),
+        },
+        {
+            "arm": "chebyshev",
+            "latency_s_per_image": chebyshev_latency,
+            "speedup": reference_latency / chebyshev_latency,
+            "psnr_vs_reference": psnr_vs_reference(chebyshev_images, reference_images),
+            "lpips_vs_reference": lpips_vs_reference(chebyshev_images, reference_images),
+        },
+    ]
+
+    with open(RESULT_FILENAME, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+        writer.writeheader()
+        writer.writerows(rows)
+
+    for row in rows:
+        print(row)
+    print(f"\nResults written to {RESULT_FILENAME}")
+
+
+if __name__ == "__main__":
+    main()

docs/source/en/api/cache.md

@@ -41,6 +41,12 @@ Cache methods speedup diffusion transformers by storing and reusing intermediate
 
 [[autodoc]] apply_taylorseer_cache
 
+## ChebyshevCacheConfig
+
+[[autodoc]] ChebyshevCacheConfig
+
+[[autodoc]] apply_chebyshev_cache
+
 ## MagCacheConfig
 
 [[autodoc]] MagCacheConfig

docs/source/en/optimization/cache.md

@@ -112,6 +112,34 @@ config = TaylorSeerCacheConfig(
 pipe.transformer.enable_cache(config)
 ```
 
+## ChebBooster (Chebyshev) Cache
+
+[ChebBooster](https://arxiv.org/abs/2608.23429) uses the same cache machinery as TaylorSeer Cache but replaces the Taylor-series extrapolation with a numerically-stable barycentric Chebyshev interpolation over a rolling window of full-compute activations. Taylor extrapolation degrades as the order and prediction horizon grow; the Chebyshev interpolant stays stable at higher orders, which the authors report translates into better quality at the same speedup (up to 3.68x latency and 5.12x FLOPs reduction on DiT-XL/2, PixArt-Sigma, and FLUX.1-dev, per the paper).
+
+To enable Chebyshev Cache, create a [`ChebyshevCacheConfig`] and pass it to your pipeline's transformer:
+
+- `cache_interval`: Number of steps to reuse cached outputs before performing a full forward pass
+- `disable_cache_before_step`: Initial steps that use full computations to gather the interpolation history window
+- `cheb_order`: Maximum number of full-compute activations kept in the rolling history window (the paper's `n`, defaults to `6`)
+
+```python
+import torch
+from diffusers import FluxPipeline, ChebyshevCacheConfig
+
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+
+config = ChebyshevCacheConfig(
+    cache_interval=5,
+    cheb_order=6,
+    disable_cache_before_step=3,
+    cheb_factors_dtype=torch.float32,
+)
+pipe.transformer.enable_cache(config)
+```
+
 ## MagCache
 
 [MagCache](https://github.com/Zehong-Ma/MagCache) accelerates inference by skipping transformer blocks based on the magnitude of the residual update. It observes that the magnitude of updates (Output - Input) decays predictably over the diffusion process. By accumulating an "error budget" based on pre-computed magnitude ratios, it dynamically decides when to skip computation and reuse the previous residual.

src/diffusers/__init__.py

@@ -175,6 +175,7 @@
     )
     _import_structure["hooks"].extend(
         [
+            "ChebyshevCacheConfig",
             "FasterCacheConfig",
             "FirstBlockCacheConfig",
             "HookRegistry",
@@ -184,6 +185,7 @@
             "SmoothedEnergyGuidanceConfig",
             "TaylorSeerCacheConfig",
             "TextKVCacheConfig",
+            "apply_chebyshev_cache",
             "apply_faster_cache",
             "apply_first_block_cache",
             "apply_layer_skip",
@@ -1036,6 +1038,7 @@
             TangentialClassifierFreeGuidance,
         )
         from .hooks import (
+            ChebyshevCacheConfig,
             FasterCacheConfig,
             FirstBlockCacheConfig,
             HookRegistry,
@@ -1045,6 +1048,7 @@
             SmoothedEnergyGuidanceConfig,
             TaylorSeerCacheConfig,
             TextKVCacheConfig,
+            apply_chebyshev_cache,
             apply_faster_cache,
             apply_first_block_cache,
             apply_layer_skip,

src/diffusers/hooks/__init__.py

@@ -16,6 +16,7 @@
 
 
 if is_torch_available():
+    from .chebyshev_cache import ChebyshevCacheConfig, apply_chebyshev_cache
     from .context_parallel import apply_context_parallel
     from .faster_cache import FasterCacheConfig, apply_faster_cache
     from .first_block_cache import FirstBlockCacheConfig, apply_first_block_cache