CachedSearch: Training-Free Cached Exploration for Test-Time Search in Video Diffusion

Paper (arXiv 2607.23159) · Code · Project page · ComfyUI nodes · Diffusers community script

CachedSearch makes best-of-N seed search cheaper for video diffusion models. It generates every candidate with a training-free cache turned on, scores the cached drafts with a verifier, and regenerates only the winning seed at full compute.

The method needs no training, so this repository holds no weights. It holds the official implementation: the cachedsearch Python package and two example scripts, copied unchanged from the GitHub repository at commit cc2a637.

CachedSearch overview

How it works

  1. Generate every candidate with a training-free cache enabled.
  2. Score the cached drafts with the verifier the search uses.
  3. Select the seed with the highest draft score.
  4. Regenerate only that seed at full compute and return it.
drafts = [cached_generate(prompt, seed) for seed in seeds]
winner = seeds[argmax(verifier(video, prompt) for video in drafts)]
return full_generate(prompt, winner)

The sampler is deterministic in the seed, so the delivered video matches what full-compute search returns whenever both pick the same seed. Caching can change which candidate wins. It does not change the quality of the video you get back.

CachedSearch method

Results reported in the official README:

  • Ranking survives caching: median Spearman rho is 0.905, with 72% top-1 agreement on VBench.
  • Exploring under caching and committing at full compute retains 94.7% of the best-of-8 gain at 63% of the cost.
  • Ranking failures concentrate where candidates are nearly tied, so they cost little.
  • The approach works with any published cache engine and transfers across six models in four families.

Files

Path Contents
cachedsearch/api.py cached_search, cached_search_batch, calibrate_tau
cachedsearch/caching.py CacheConfig, CachedTransformer, wrap_pipeline (the EasyCache-style adaptive cache and a static baseline)
cachedsearch/verifiers.py the default ImageReward verifier
examples/run_wan.py best-of-8 search on Wan2.1
examples/calibrate_new_model.py threshold calibration for a new model family

Install

Install from GitHub:

pip install "git+https://github.com/shreshthsaini/CachedSearch.git@cc2a637c89485af723233aeb8ad76190b572a464"

Or fetch the package from this repository, with no install step:

import sys
from huggingface_hub import snapshot_download

repo_dir = snapshot_download("shreshthsaini/CachedSearch", allow_patterns=["cachedsearch/*"])
sys.path.insert(0, repo_dir)

Requirements: torch, diffusers, numpy, plus whatever your pipeline needs (transformers for the text encoder). The default verifier needs image-reward. Pass your own verifier= to skip that dependency.

Usage

cached_search takes any diffusers video pipeline that exposes .transformer. The verifier is any callable (frames, prompt) -> float where higher is better.

The example below runs on CPU in a few seconds. It uses hf-internal-testing/tiny-wan-pipe, a tiny Wan pipeline with random weights, and a toy verifier. It checks that the search loop and the cache run end to end. The frames are noise and the scores mean nothing.

import sys

import numpy as np
import torch
from diffusers import WanPipeline
from huggingface_hub import snapshot_download

# Skip these two lines if you installed the package with pip.
repo_dir = snapshot_download("shreshthsaini/CachedSearch", allow_patterns=["cachedsearch/*"])
sys.path.insert(0, repo_dir)

from cachedsearch import CacheConfig, cached_search
from cachedsearch.caching import wrap_pipeline

pipe = WanPipeline.from_pretrained("hf-internal-testing/tiny-wan-pipe", torch_dtype=torch.float32)
gen_kwargs = dict(height=32, width=32, num_frames=9, guidance_scale=5.0)


def toy_verifier(frames, prompt):
    # Stand-in scorer so the example runs without ImageReward. Not a quality metric.
    return float(np.asarray(frames).mean())


result = cached_search(
    pipe,
    "a red fox running through deep snow",
    verifier=toy_verifier,
    n=4,             # search width
    tau=0.10,        # caching threshold
    steps=20,
    mode="commit",   # regenerate the winner at full compute
    gen_kwargs=gen_kwargs,
)

print(result.seed)                      # the seed that won
print(result.draft_scores)              # verifier score of every cached draft
print(np.asarray(result.video).shape)   # (9, 32, 32, 3), the full-compute video

# Count how many transformer calls one cached rollout skips.
cache = wrap_pipeline(pipe, CacheConfig(mode="adaptive", tau=0.10, total_steps=20))
cache.reset()
pipe(prompt="a red fox running through deep snow", num_inference_steps=20,
     generator=torch.Generator().manual_seed(0), **gen_kwargs)
print(cache.stats())

Output on CPU with torch 2.14.1 and diffusers 0.41.0:

3
[0.4587..., 0.4588..., 0.4589..., 0.4590...]
(9, 32, 32, 3)
{'transformer0': {'branch0': {'skips': 8, 'computes': 12}, 'branch1': {'skips': 8, 'computes': 12}}}

With these settings the cache skips 8 of 20 transformer calls per classifier-free-guidance branch.

Real models

For a real model, load the checkpoint on a GPU and drop the verifier argument to use ImageReward averaged over 8 frames, the scorer behind the paper's numbers. examples/run_wan.py does this for Wan-AI/Wan2.1-T2V-1.3B-Diffusers:

pip install image-reward imageio imageio-ffmpeg
python examples/run_wan.py --prompt "a red fox running through deep snow"

This script and examples/calibrate_new_model.py are the official ones from the GitHub repository. They need a CUDA GPU and were not rerun when this card was written. Only the CPU example above was.

Thresholds for tested models

The caching threshold tau is the one model-specific number. The table below comes from the official README: the recommended tau is the most aggressive setting that still keeps at least 85% of the search gain in the calibration sweep.

Model Params Family Recommended tau Gain capture Exploration speedup
Wan2.1-T2V-1.3B 1.3B Wan 0.20 88.3% 2.41x
Wan2.1-T2V-14B 14B Wan 0.10 87.5% 2.05x
Wan2.2-TI2V-5B 5B Wan 0.10 86.0% 2.05x
CogVideoX-5B 5B CogVideoX 0.05 85.9% 1.78x
HunyuanVideo 13B HunyuanVideo 0.05 85.1% 1.77x
LTX-Video 2B LTX 0.02 79.6% 1.71x

No threshold measured for LTX-Video reached 85% capture. For a model that is not listed, run examples/calibrate_new_model.py once (about two GPU-hours) and use the threshold it recommends. Do not copy a threshold across model families.

Modes

  • mode="commit" (default): regenerate the winner with caching off. The delivered video is a full-compute sample.
  • mode="keep": return the cached draft of the winner. Cheaper, with mild motion dampening at aggressive thresholds.

Citation

@article{saini2026cachedsearch,
  title         = {CachedSearch: Training-Free Cached Exploration for Test-Time Search in Video Diffusion},
  author        = {Saini, Shreshth and Birkbeck, Neil and Wang, Yilin and Adsumilli, Balu and Bovik, Alan C.},
  year          = {2026},
  eprint        = {2607.23159},
  archivePrefix = {arXiv},
  primaryClass  = {cs.AI}
}

License

MIT. See LICENSE.

Downloads last month
-
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Collection including shreshthsaini/CachedSearch

Paper for shreshthsaini/CachedSearch