| |
| """Fit the TeaCache rescale polynomial for a base model. |
| |
| TeaCache ships per-model 4th-degree coefficients that map the cheap input-side |
| relative L1 distance onto the output-side distance. Those were fitted for |
| bidirectional Wan on a 50-step schedule with an ``e0``-only indicator, so they do |
| not transfer to a 4-step block-causal model with the modulated-input indicator. |
| This script reruns TeaCache's own fitting procedure on the target model. |
| |
| python calibrate.py --base self_forcing --num-prompts 8 \ |
| --out results/coeff_self_forcing.json |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import sys |
|
|
| ROOT = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, ROOT) |
|
|
| from harness import BASES, enter_base, load_pipeline |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--base", choices=sorted(BASES), required=True) |
| p.add_argument("--indicator", default="modulated_input", |
| choices=["modulated_input", "e0"]) |
| p.add_argument("--num-prompts", type=int, default=8) |
| p.add_argument("--num-output-frames", type=int, default=21) |
| p.add_argument("--degree", type=int, default=1, |
| help="Degree 4 (TeaCache default) is ill-conditioned on the narrow " |
| "distance range a 4-step schedule produces; degree 1 stays " |
| "monotone and is safe to extrapolate per token.") |
| p.add_argument("--seed", type=int, default=0) |
| p.add_argument("--out", required=True) |
| args = p.parse_args() |
|
|
| out_path = args.out if os.path.isabs(args.out) else os.path.join(ROOT, args.out) |
| base = enter_base(args.base) |
|
|
| import numpy as np |
| import torch |
|
|
| from utils.dataset import TextDataset |
| from utils.misc import set_seed |
|
|
| from cachelib import CacheController, build_method, cached_inference, install |
|
|
| set_seed(args.seed) |
| torch.set_grad_enabled(False) |
| device = torch.device("cuda") |
| pipeline = load_pipeline(base) |
|
|
| method = build_method("calibrate", indicator=args.indicator) |
| num_steps = len(pipeline.denoising_step_list) |
| controller = CacheController(method, num_steps=num_steps, |
| forced_steps=(0, num_steps - 1)) |
| install(pipeline.generator.model, controller) |
|
|
| prompt_file = ("prompts/MovieGenVideoBench.txt" |
| if os.path.exists("prompts/MovieGenVideoBench.txt") |
| else "prompts/demos.txt") |
| dataset = TextDataset(prompt_path=prompt_file) |
|
|
| for i in range(min(args.num_prompts, len(dataset))): |
| set_seed(args.seed) |
| noise = torch.randn([1, args.num_output_frames, 16, 60, 104], |
| device=device, dtype=torch.bfloat16) |
| cached_inference(pipeline, controller, noise, [dataset[i]["prompts"]], |
| decode=False) |
| print(f"[{i}] samples={len(method.samples)}", flush=True) |
|
|
| xs = np.array([s["x"] for s in method.samples], dtype=np.float64) |
| ys = np.array([s["y"] for s in method.samples], dtype=np.float64) |
| coeffs = np.polyfit(xs, ys, args.degree) |
| pred = np.poly1d(coeffs)(xs) |
|
|
| result = { |
| "base": args.base, |
| "indicator": args.indicator, |
| "degree": args.degree, |
| "num_samples": int(xs.size), |
| "coefficients": [float(c) for c in coeffs], |
| "x_range": [float(xs.min()), float(xs.max())], |
| "y_range": [float(ys.min()), float(ys.max())], |
| "fit_r2": float(1 - ((ys - pred) ** 2).sum() / ((ys - ys.mean()) ** 2).sum()), |
| "per_step_mean_x": { |
| str(k): float(xs[[s["step"] == k for s in method.samples]].mean()) |
| for k in sorted({s["step"] for s in method.samples}) |
| }, |
| "per_step_mean_y": { |
| str(k): float(ys[[s["step"] == k for s in method.samples]].mean()) |
| for k in sorted({s["step"] for s in method.samples}) |
| }, |
| "samples": method.samples, |
| } |
| print(json.dumps({k: v for k, v in result.items() if k != "samples"}, indent=2)) |
|
|
| os.makedirs(os.path.dirname(out_path), exist_ok=True) |
| with open(out_path, "w") as f: |
| json.dump(result, f, indent=2) |
| print(f"wrote {out_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
