File size: 4,228 Bytes
38a51ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | #!/usr/bin/env python
"""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 # noqa: E402
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()
|