Add CachedSearch implementation and model card

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/fig1_overview.png filter=lfs diff=lfs merge=lfs -text
+assets/fig2_method.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Shreshth Saini
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,192 @@
+---
+license: mit
+library_name: diffusers
+pipeline_tag: text-to-video
+tags:
+  - text-to-video
+  - video-generation
+  - test-time-search
+  - best-of-n
+  - caching
+  - training-free
+  - wan
+  - cogvideox
+  - hunyuanvideo
+  - ltx-video
+---
+
+# CachedSearch: Training-Free Cached Exploration for Test-Time Search in Video Diffusion
+
+[Paper (arXiv 2607.23159)](https://arxiv.org/abs/2607.23159) · [Code](https://github.com/shreshthsaini/CachedSearch) · [Project page](https://shreshthsaini.github.io/CachedSearch/) · [ComfyUI nodes](https://github.com/shreshthsaini/ComfyUI-CachedSearch) · [Diffusers community script](https://github.com/huggingface/diffusers/blob/main/examples/community/README_community_scripts.md#cachedsearch-cheaper-best-of-n-for-video-pipelines)
+
+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`](https://github.com/shreshthsaini/CachedSearch/tree/cc2a637c89485af723233aeb8ad76190b572a464).
+
+![CachedSearch overview](assets/fig1_overview.png)
+
+## 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.
+
+```text
+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](assets/fig2_method.png)
+
+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:
+
+```bash
+pip install "git+https://github.com/shreshthsaini/CachedSearch.git@cc2a637c89485af723233aeb8ad76190b572a464"
+```
+
+Or fetch the package from this repository, with no install step:
+
+```python
+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.
+
+```python
+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:
+
+```text
+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`:
+
+```bash
+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
+
+```bibtex
+@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](LICENSE).

cachedsearch/__init__.py

@@ -0,0 +1,23 @@
+"""CachedSearch: explore cheap, commit full.
+
+Drop-in test-time search for any diffusers video pipeline:
+
+    from cachedsearch import cached_search
+
+    result = cached_search(pipe, "a red fox in snow", n=8)
+
+The verifier defaults to ImageReward averaged over 8 frames, the scorer behind
+the paper's numbers. Pass ``verifier=`` to use your own.
+
+The exploration rollouts run under training-free caching; only the winning
+seed is regenerated at full compute, so the returned video is a genuine
+full-compute sample.
+"""
+
+from .api import CacheConfig, cached_search, cached_search_batch, calibrate_tau
+from .verifiers import default_verifier, imagereward_verifier
+
+__all__ = [
+    "cached_search", "cached_search_batch", "calibrate_tau", "CacheConfig",
+    "default_verifier", "imagereward_verifier",
+]

cachedsearch/api.py

@@ -0,0 +1,181 @@
+"""The user-facing CachedSearch API.
+
+Three functions:
+
+    cached_search(pipe, prompt, verifier, n=8)   run search on one prompt
+    cached_search_batch(pipe, prompts, verifier) run search on many prompts
+    calibrate_tau(pipe, prompts, verifier)       find the threshold for a new model
+
+Everything is training-free. The only model-specific quantity is the caching
+threshold tau, which `calibrate_tau` measures in about two GPU-hours.
+"""
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Callable, Iterable, Sequence
+
+import numpy as np
+import torch
+
+from .caching import CacheConfig, wrap_pipeline
+
+from .verifiers import Verifier, default_verifier, imagereward_verifier
+
+
+@dataclass
+class SearchResult:
+    video: object          # the delivered video (full-compute sample in commit mode)
+    seed: int              # the seed that produced it
+    draft_scores: list     # verifier score of every cached draft
+    explore_seconds: float
+    commit_seconds: float
+
+    @property
+    def total_seconds(self) -> float:
+        return self.explore_seconds + self.commit_seconds
+
+
+def _generate(pipe, prompt: str, seed: int, steps: int, gen_kwargs: dict):
+    generator = torch.Generator(device=pipe.device).manual_seed(seed)
+    out = pipe(prompt=prompt, num_inference_steps=steps, generator=generator, **gen_kwargs)
+    return out.frames[0]
+
+
+def cached_search(
+    pipe,
+    prompt: str,
+    verifier: Verifier | None = None,
+    n: int = 8,
+    tau: float = 0.10,
+    steps: int = 50,
+    mode: str = "commit",
+    seeds: Sequence[int] | None = None,
+    gen_kwargs: dict | None = None,
+) -> SearchResult:
+    """Best-of-n search where exploration runs under caching.
+
+    Args:
+        pipe: any diffusers video pipeline exposing `.transformer` (Wan, CogVideoX,
+            HunyuanVideo, LTX-Video all work).
+        prompt: the text prompt.
+        verifier: callable (frames, prompt) -> float, higher is better. Any
+            scorer works. Defaults to ImageReward averaged over 8 uniformly
+            spaced frames, the verifier behind the paper's headline numbers.
+        n: search width (number of candidates). Recommit pays off for n >= 4.
+        tau: caching threshold. 0.10 is calibrated for Wan2.1; run `calibrate_tau`
+            for a new architecture (do not copy this constant across families).
+        steps: denoising steps.
+        mode: "commit" regenerates the winner at full compute (recommended; the
+            delivered video is a genuine full-compute sample). "keep" returns the
+            cached draft directly: cheaper, but caching artifacts reach the user
+            and motion is mildly dampened.
+        seeds: explicit seeds; defaults to range(n).
+        gen_kwargs: extra pipeline kwargs (height, width, num_frames, guidance...).
+
+    Returns:
+        SearchResult with the delivered video, its seed, draft scores, and timings.
+    """
+    seeds = list(seeds if seeds is not None else range(n))
+    gen_kwargs = dict(gen_kwargs or {})
+    if verifier is None:
+        verifier = default_verifier()
+
+    # ---- explore: every candidate under caching -------------------------------
+    cache = wrap_pipeline(pipe, CacheConfig(mode="adaptive", tau=tau, total_steps=steps))
+    drafts, scores = [], []
+    t0 = time.time()
+    for seed in seeds:
+        cache.reset()                      # cache state must not leak across candidates
+        frames = _generate(pipe, prompt, seed, steps, gen_kwargs)
+        drafts.append(frames)
+        scores.append(float(verifier(frames, prompt)))
+    explore_s = time.time() - t0
+
+    best = int(np.argmax(scores))
+    winner_seed = seeds[best]
+
+    # ---- commit: regenerate only the winner at full compute -------------------
+    if mode == "keep":
+        return SearchResult(drafts[best], winner_seed, scores, explore_s, 0.0)
+
+    wrap_pipeline(pipe, CacheConfig(mode="off", total_steps=steps))
+    t0 = time.time()
+    video = _generate(pipe, prompt, winner_seed, steps, gen_kwargs)
+    commit_s = time.time() - t0
+    return SearchResult(video, winner_seed, scores, explore_s, commit_s)
+
+
+def cached_search_batch(
+    pipe, prompts: Iterable[str], verifier: Verifier | None = None, **kwargs
+) -> list[SearchResult]:
+    """cached_search over many prompts; returns one SearchResult per prompt."""
+    return [cached_search(pipe, p, verifier, **kwargs) for p in prompts]
+
+
+def calibrate_tau(
+    pipe,
+    prompts: Sequence[str],
+    verifier: Verifier | None = None,
+    taus: Sequence[float] = (0.05, 0.10, 0.20),
+    n: int = 8,
+    steps: int = 50,
+    target_capture: float = 0.90,
+    gen_kwargs: dict | None = None,
+) -> dict:
+    """Measure the speed-fidelity frontier of a NEW model and pick tau.
+
+    Protocol (this is the paper's protocol in miniature): generate every
+    (prompt, seed) pair twice, once at full compute and once under caching,
+    score both, and report how much of full-compute search's gain the cached
+    ranking retains. Roughly two GPU-hours with 25 prompts on a 1.3B model.
+
+    Returns {tau: {"speedup", "capture", "spearman"}} plus "recommended": the
+    largest tau whose capture is at or above target_capture.
+    """
+    prompts = list(prompts)
+    gen_kwargs = dict(gen_kwargs or {})
+    seeds = list(range(n))
+    if verifier is None:
+        verifier = default_verifier()
+
+    # full-compute references, shared across all tau values
+    wrap_pipeline(pipe, CacheConfig(mode="off", total_steps=steps))
+    full_scores, full_latency = {}, []
+    for p in prompts:
+        row = []
+        for s in seeds:
+            t0 = time.time()
+            frames = _generate(pipe, p, s, steps, gen_kwargs)
+            full_latency.append(time.time() - t0)
+            row.append(float(verifier(frames, p)))
+        full_scores[p] = np.asarray(row)
+    cf = float(np.mean(full_latency))
+
+    report = {}
+    for tau in taus:
+        cache = wrap_pipeline(pipe, CacheConfig(mode="adaptive", tau=tau, total_steps=steps))
+        rhos, captures, latency = [], [], []
+        for p in prompts:
+            row = []
+            for s in seeds:
+                cache.reset()
+                t0 = time.time()
+                frames = _generate(pipe, p, s, steps, gen_kwargs)
+                latency.append(time.time() - t0)
+                row.append(float(verifier(frames, p)))
+            cached = np.asarray(row)
+            full = full_scores[p]
+            ra, rb = np.argsort(np.argsort(full)), np.argsort(np.argsort(cached))
+            rhos.append(float(np.corrcoef(ra, rb)[0, 1]))
+            pick, spread = int(np.argmax(cached)), full.max() - full.mean()
+            captures.append((full[pick] - full.mean()) / spread if spread > 1e-9 else 1.0)
+        report[tau] = {
+            "speedup": cf / float(np.mean(latency)),
+            "capture": float(np.mean(captures)),
+            "spearman": float(np.median(rhos)),
+        }
+
+    ok = [t for t in taus if report[t]["capture"] >= target_capture]
+    report["recommended"] = max(ok) if ok else min(taus)
+    return report

cachedsearch/caching.py

@@ -0,0 +1,144 @@
+"""Training-free caching wrappers for video DiTs (Wan, CogVideoX, HunyuanVideo, LTX via diffusers).
+
+Two modes:
+  - "adaptive": EasyCache-style online rule. Cache the transformation vector
+    delta = v(x_t) - x_t; while an accumulated change indicator stays below
+    tau, skip the transformer forward and return x_t + delta.
+  - "static": fixed skip pattern over a step range (PAB/uniform-style baseline).
+
+CFG handling: diffusers WanPipeline calls the transformer once per branch per
+step. We key cache state on call parity (branch index = call_count % branches)
+so cond/uncond each get independent state.
+"""
+from __future__ import annotations
+import torch
+from dataclasses import dataclass, field
+
+
+@dataclass
+class _BranchState:
+    prev_input: torch.Tensor | None = None
+    prev_output: torch.Tensor | None = None
+    delta: torch.Tensor | None = None
+    accum: float = 0.0
+    skips: int = 0
+    computes: int = 0
+    step: int = 0
+
+
+@dataclass
+class CacheConfig:
+    mode: str = "adaptive"          # "adaptive" | "static" | "off"
+    tau: float = 0.10               # adaptive: accumulated-change tolerance
+    warmup_steps: int = 5           # always compute the first K steps
+    cooldown_steps: int = 5         # always compute the last K steps (needs total_steps)
+    total_steps: int = 50
+    static_keep_every: int = 2      # static: compute every k-th step in the middle
+    num_branches: int = 2           # CFG cond+uncond
+
+
+class CachedTransformer(torch.nn.Module):
+    """Wraps a diffusers transformer; drop-in via `pipe.transformer = CachedTransformer(pipe.transformer, cfg)`."""
+
+    def __init__(self, transformer: torch.nn.Module, cfg: CacheConfig):
+        super().__init__()
+        self.t = transformer
+        self.cfg = cfg
+        self.dtype = transformer.dtype
+        self.config = transformer.config  # pipeline reads .config
+        self._call = 0
+        self._states = [_BranchState() for _ in range(cfg.num_branches)]
+
+    # pipeline may call other attrs; delegate
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.t, name)
+
+    def reset(self):
+        self._call = 0
+        self._states = [_BranchState() for _ in range(self.cfg.num_branches)]
+
+    def stats(self):
+        return {
+            f"branch{i}": {"skips": s.skips, "computes": s.computes}
+            for i, s in enumerate(self._states)
+        }
+
+    def _should_compute(self, st: _BranchState, x: torch.Tensor) -> bool:
+        c = self.cfg
+        if c.mode == "off" or st.delta is None:
+            return True
+        if st.step < c.warmup_steps or st.step >= c.total_steps - c.cooldown_steps:
+            return True
+        if c.mode == "static":
+            return (st.step % c.static_keep_every) == 0
+        # adaptive: estimate output change rate from input change (EasyCache-style)
+        if st.prev_input is None:
+            return True
+        dx = (x - st.prev_input).float().norm() / max(st.prev_input.float().norm().item(), 1e-8)
+        st.accum += dx.item()
+        return st.accum > c.tau
+
+    def forward(self, hidden_states, *args, **kwargs):
+        st = self._states[self._call % self.cfg.num_branches]
+        self._call += 1
+        x = hidden_states
+
+        if self._should_compute(st, x):
+            out = self.t(hidden_states, *args, **kwargs)
+            sample = out[0] if isinstance(out, tuple) else out.sample
+            st.delta = (sample - x).detach()
+            st.prev_input = x.detach()
+            st.prev_output = sample.detach()
+            st.accum = 0.0
+            st.computes += 1
+            st.step += 1
+            return out
+        # skip: approximate v(x_t) ~= x_t + cached delta
+        st.skips += 1
+        st.step += 1
+        approx = x + st.delta
+        return (approx,)
+
+
+class CacheHandle:
+    """Handle over every wrapped transformer of a pipeline (Wan 2.2 has two)."""
+
+    def __init__(self, wrapped: list[CachedTransformer]):
+        self.wrapped = wrapped
+
+    def reset(self):
+        for w in self.wrapped:
+            w.reset()
+
+    def stats(self):
+        return {f"transformer{i}": w.stats() for i, w in enumerate(self.wrapped)}
+
+    @property
+    def cfg(self) -> CacheConfig:
+        return self.wrapped[0].cfg
+
+
+def wrap_pipeline(pipe, cfg: CacheConfig) -> CacheHandle:
+    """Attach caching to a diffusers video pipeline. Call .reset() between generations.
+
+    Wraps ``pipe.transformer`` and, when present, ``pipe.transformer_2`` (Wan 2.2).
+    Calling it again with a new config reconfigures the existing wrappers in place.
+    """
+    wrapped = []
+    for name in ("transformer", "transformer_2"):
+        model = getattr(pipe, name, None)
+        if model is None:
+            continue
+        if isinstance(model, CachedTransformer):
+            model.cfg = cfg
+            model.reset()
+        else:
+            model = CachedTransformer(model, cfg)
+            setattr(pipe, name, model)
+        wrapped.append(model)
+    if not wrapped:
+        raise AttributeError("The pipeline exposes no `transformer` to wrap.")
+    return CacheHandle(wrapped)

cachedsearch/verifiers.py

@@ -0,0 +1,52 @@
+"""Verifiers: the scorer that ranks candidates.
+
+CachedSearch is verifier-agnostic. Any callable ``(frames, prompt) -> float``
+works, higher meaning better. If you do not pass one, we use ImageReward
+averaged over uniformly spaced frames, which is the verifier all of the
+paper's headline numbers use.
+"""
+from __future__ import annotations
+
+from typing import Callable
+
+import numpy as np
+
+Verifier = Callable[[object, str], float]
+
+_DEFAULT = None
+
+
+def imagereward_verifier(device: str = "cuda", num_frames: int = 8) -> Verifier:
+    """ImageReward averaged over `num_frames` uniformly spaced frames.
+
+    Requires the `image-reward` package; the model downloads on first use.
+    """
+    import ImageReward as RM
+    from PIL import Image
+
+    model = RM.load("ImageReward-v1.0", device=device)
+
+    def score(frames, prompt: str) -> float:
+        idx = np.linspace(0, len(frames) - 1, min(num_frames, len(frames))).astype(int)
+        vals = [
+            model.score(prompt, Image.fromarray(np.asarray(frames[i]).astype("uint8")))
+            for i in idx
+        ]
+        return float(np.mean(vals))
+
+    return score
+
+
+def default_verifier(device: str = "cuda") -> Verifier:
+    """The default scorer, loaded once and reused across calls."""
+    global _DEFAULT
+    if _DEFAULT is None:
+        try:
+            _DEFAULT = imagereward_verifier(device=device)
+        except ImportError as exc:  # keep the failure actionable
+            raise ImportError(
+                "No verifier was passed and the default (ImageReward) is not "
+                "installed. Either `pip install image-reward`, or pass your own "
+                "scorer: cached_search(pipe, prompt, verifier=my_scorer)."
+            ) from exc
+    return _DEFAULT

examples/calibrate_new_model.py

@@ -0,0 +1,46 @@
+"""Calibrate the caching threshold for a model we did not test.
+
+    python examples/calibrate_new_model.py --model <hf-id> --prompts data/calib25.txt
+
+Costs about two GPU-hours for a small model (25 prompts x 8 seeds, twice).
+Fidelity tracks architecture family, not parameter size, so calibrate once per
+family and reuse the threshold across sizes within it.
+"""
+import argparse
+
+import torch
+from diffusers import DiffusionPipeline
+
+from cachedsearch import calibrate_tau
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--prompts", required=True, help="text file, one prompt per line")
+    ap.add_argument("--num-prompts", type=int, default=25)
+    ap.add_argument("--taus", default="0.02,0.05,0.10,0.20")
+    ap.add_argument("--target-capture", type=float, default=0.90)
+    args = ap.parse_args()
+
+    prompts = [l.strip() for l in open(args.prompts) if l.strip()][: args.num_prompts]
+    pipe = DiffusionPipeline.from_pretrained(args.model, torch_dtype=torch.bfloat16).to("cuda")
+
+    report = calibrate_tau(
+        pipe,
+        prompts,
+        taus=[float(t) for t in args.taus.split(",")],
+        target_capture=args.target_capture,
+    )
+
+    print(f"{'tau':>6} {'speedup':>9} {'capture':>9} {'median rho':>11}")
+    for tau, row in report.items():
+        if tau == "recommended":
+            continue
+        print(f"{tau:>6} {row['speedup']:>8.2f}x {100*row['capture']:>8.1f}% {row['spearman']:>11.3f}")
+    print(f"\nrecommended tau for this model: {report['recommended']}")
+    print("Use it as: cached_search(pipe, prompt, verifier, tau=<value>)")
+
+
+if __name__ == "__main__":
+    main()

examples/run_wan.py

@@ -0,0 +1,52 @@
+"""End-to-end example: best-of-8 search on Wan2.1 for the price of about 5 rollouts.
+
+    python examples/run_wan.py --prompt "a red fox running through deep snow"
+
+Swap the pipeline for CogVideoX, HunyuanVideo, or LTX-Video and the only thing
+that changes is tau (run examples/calibrate_new_model.py once for a new family).
+"""
+import argparse
+
+import torch
+from diffusers import AutoencoderKLWan, WanPipeline
+
+from cachedsearch import cached_search, imagereward_verifier
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--prompt", required=True)
+    ap.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    ap.add_argument("--n", type=int, default=8, help="search width")
+    ap.add_argument("--tau", type=float, default=0.10, help="caching threshold")
+    ap.add_argument("--mode", default="commit", choices=["commit", "keep"])
+    ap.add_argument("--out", default="delivered.mp4")
+    # Pass --verifier to swap in your own scorer; the default is ImageReward.
+    args = ap.parse_args()
+
+    vae = AutoencoderKLWan.from_pretrained(args.model, subfolder="vae", torch_dtype=torch.float32)
+    pipe = WanPipeline.from_pretrained(args.model, vae=vae, torch_dtype=torch.bfloat16).to("cuda")
+
+    result = cached_search(
+        pipe,
+        args.prompt,
+        n=args.n,
+        tau=args.tau,
+        mode=args.mode,
+        gen_kwargs=dict(height=480, width=832, num_frames=81, guidance_scale=5.0),
+    )
+
+    print(f"winner seed      : {result.seed}")
+    print(f"draft scores     : {[round(s, 3) for s in result.draft_scores]}")
+    print(f"explore / commit : {result.explore_seconds:.0f}s / {result.commit_seconds:.0f}s")
+    print(f"total            : {result.total_seconds:.0f}s "
+          f"(full best-of-{args.n} would cost about "
+          f"{args.n * result.commit_seconds:.0f}s)")
+
+    from diffusers.utils import export_to_video
+    export_to_video(result.video, args.out, fps=16)
+    print(f"wrote {args.out}")
+
+
+if __name__ == "__main__":
+    main()