Add CachedSearch implementation and model card

914261a verified 14 days ago

7.04 kB

"""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

	"""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