Self-Forcing / scripts /generate_vbench8_extended_naive_baselines.py
Cccccz's picture
Upload Python scripts
bc29ee3 verified
Raw History Blame Contribute Delete
22.7 kB
#!/usr/bin/env python3
"""Generate Extended-251 videos for timestep-skipping and velocity-cache baselines.
For every prompt, an original four-step FFFF rollout is generated first with
the same seed and used as the PSNR/SSIM/LPIPS reference. Timestep-skipping
policies shorten the denoising schedule itself. Velocity-cache policies retain
the original four timesteps and, at each R, reuse the velocity from the most
recent F while converting it to x0 using the current noisy latent and timestep.
The cache pattern is applied independently to every autoregressive chunk.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import os
import sys
import tempfile
import time
from pathlib import Path
from typing import Any
def preparse_gpu() -> str:
parser = argparse.ArgumentParser(add_help=False)
parser.add_argument("--gpu", default="0")
args, _ = parser.parse_known_args()
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
os.environ.setdefault("MPLCONFIGDIR", tempfile.mkdtemp(prefix="self_forcing_mpl_"))
return str(args.gpu)
PHYSICAL_GPU = preparse_gpu()
import lpips
import torch
from omegaconf import OmegaConf
REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from scripts import evaluate_single_block_fppf as base
from scripts.naive_vbench_policies import (
EVALUATION_STRATEGY_NAMES,
POLICIES_BY_NAME,
REFERENCE_POLICY,
NaivePolicy,
)
from scripts.vbench8_protocol import SUITE_COUNTS
from utils.misc import set_seed
from utils.wan_wrapper import WanTextEncoder, WanVAEWrapper
NUM_CHUNKS = base.NUM_CHUNKS
FRAMES_PER_CHUNK = base.FRAMES_PER_CHUNK
MAPPING_DEFAULT = REPO_ROOT / "assets/vbench8_extended_subset_mapping.json"
OUTPUT_DEFAULT = REPO_ROOT / "evaluation_runs/vbench8_extended_naive_baselines"
def atomic_json(path: Path, value: Any) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix(path.suffix + ".tmp")
temporary.write_text(
json.dumps(value, ensure_ascii=False, indent=2, allow_nan=True) + "\n",
encoding="utf-8",
)
os.replace(temporary, path)
def atomic_video(frames: torch.Tensor, path: Path) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_name(path.name + ".partial.mp4")
if temporary.exists():
temporary.unlink()
base.save_mp4(frames, temporary)
os.replace(temporary, path)
def read_mapping(path: Path) -> list[dict[str, Any]]:
value = json.loads(path.read_text(encoding="utf-8"))
if not isinstance(value, list) or len(value) != 251:
raise ValueError(f"Expected a 251-row mapping list: {path}")
counts = {suite: 0 for suite in SUITE_COUNTS}
globals_seen: set[int] = set()
suite_seen: dict[str, set[int]] = {suite: set() for suite in SUITE_COUNTS}
for row in value:
suite = str(row["prompt_suite"])
global_index = int(row["global_index"])
suite_index = int(row["suite_index"])
if suite not in counts:
raise ValueError(f"Unknown prompt suite: {suite}")
if global_index in globals_seen or suite_index in suite_seen[suite]:
raise ValueError(
f"Duplicate mapping index: {global_index} / {suite}:{suite_index}"
)
if not str(row["extended_prompt"]).strip():
raise ValueError(f"Empty prompt at global index {global_index}")
counts[suite] += 1
globals_seen.add(global_index)
suite_seen[suite].add(suite_index)
if counts != SUITE_COUNTS:
raise ValueError(f"Unexpected prompt-suite counts: {counts}")
return sorted(value, key=lambda row: int(row["global_index"]))
def reset_runtime_caches(pipeline: Any, device: torch.device) -> None:
if pipeline.kv_cache1 is None:
pipeline._initialize_kv_cache(1, torch.bfloat16, device)
pipeline._initialize_crossattn_cache(1, torch.bfloat16, device)
for cache in pipeline.kv_cache1:
cache["global_end_index"].zero_()
cache["local_end_index"].zero_()
for cache in pipeline.crossattn_cache:
cache["is_init"] = False
def warped_timesteps(pipeline: Any, raw_timesteps: tuple[int, ...]) -> torch.Tensor:
raw = torch.tensor(raw_timesteps, dtype=torch.long)
if not bool(pipeline.args.warp_denoising_step):
return raw
scheduler_timesteps = torch.cat(
(pipeline.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))
)
return scheduler_timesteps[1000 - raw]
def policy_timesteps(pipeline: Any, policy: NaivePolicy) -> torch.Tensor:
if policy.family in {"full", "velocity_reuse"}:
actual = pipeline.denoising_step_list.detach().cpu()
else:
actual = warped_timesteps(pipeline, policy.raw_timesteps)
if len(actual) != policy.denoise_steps:
raise ValueError(
f"{policy.name}: expected {policy.denoise_steps} timesteps, got {len(actual)}"
)
return actual
def start_timing() -> tuple[torch.cuda.Event, torch.cuda.Event]:
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
return start, end
def finish_timing(
events_by_category: dict[str, list[tuple[torch.cuda.Event, torch.cuda.Event]]],
category: str,
events: tuple[torch.cuda.Event, torch.cuda.Event],
) -> None:
events[1].record()
events_by_category[category].append(events)
@torch.inference_mode()
def generate_rollout(
*,
pipeline: Any,
conditional_dict: dict[str, torch.Tensor],
seed: int,
device: torch.device,
policy: NaivePolicy,
) -> tuple[torch.Tensor, dict[str, Any]]:
reset_runtime_caches(pipeline, device)
set_seed(seed)
noise = torch.randn(
1,
NUM_CHUNKS * FRAMES_PER_CHUNK,
base.LATENT_CHANNELS,
base.LATENT_HEIGHT,
base.LATENT_WIDTH,
dtype=torch.bfloat16,
device=device,
)
timesteps_cpu = policy_timesteps(pipeline, policy)
actual_timestep_values = [float(value) for value in timesteps_cpu]
timesteps = timesteps_cpu.to(device=device)
output_chunks: list[torch.Tensor] = []
decisions: list[dict[str, Any]] = []
full_calls = 0
reuse_calls = 0
timing_events: dict[str, list[tuple[torch.cuda.Event, torch.cuda.Event]]] = {
"full_dit": [],
"context_dit": [],
}
started = time.perf_counter()
for chunk in range(NUM_CHUNKS):
noisy_input = noise[
:, chunk * FRAMES_PER_CHUNK : (chunk + 1) * FRAMES_PER_CHUNK
]
cached_flow: torch.Tensor | None = None
denoised_pred: torch.Tensor | None = None
timestep: torch.Tensor | None = None
for step, current_timestep in enumerate(timesteps):
timestep = torch.ones(
[1, FRAMES_PER_CHUNK], dtype=torch.int64, device=device
) * current_timestep
action = (
policy.cache_pattern[step]
if policy.family == "velocity_reuse" and policy.cache_pattern is not None
else "F"
)
if action == "R":
if cached_flow is None:
raise RuntimeError(
f"{policy.name}: reuse requested before a full step in chunk {chunk}"
)
denoised_pred = pipeline.generator._convert_flow_pred_to_x0(
flow_pred=cached_flow.flatten(0, 1),
xt=noisy_input.flatten(0, 1),
timestep=timestep.flatten(0, 1),
).unflatten(0, cached_flow.shape[:2])
reuse_calls += 1
else:
full_events = start_timing()
flow, denoised_pred = pipeline.generator(
noisy_image_or_video=noisy_input,
conditional_dict=conditional_dict,
timestep=timestep,
kv_cache=pipeline.kv_cache1,
crossattn_cache=pipeline.crossattn_cache,
current_start=chunk * base.TOKENS_PER_CHUNK,
)
finish_timing(timing_events, "full_dit", full_events)
cached_flow = flow.detach()
full_calls += 1
decisions.append(
{
"chunk": chunk,
"step": step,
"action": action,
"timestep": actual_timestep_values[step],
}
)
if step < len(timesteps) - 1:
if denoised_pred is None:
raise RuntimeError("Denoising step did not produce x0")
next_timestep = timesteps[step + 1]
flat = denoised_pred.flatten(0, 1)
noisy_input = pipeline.scheduler.add_noise(
flat,
torch.randn_like(flat),
next_timestep
* torch.ones(
[FRAMES_PER_CHUNK], dtype=torch.long, device=device
),
).unflatten(0, denoised_pred.shape[:2])
if denoised_pred is None or timestep is None:
raise RuntimeError(f"{policy.name}: chunk {chunk} produced no clean latent")
output_chunks.append(denoised_pred)
context_events = start_timing()
pipeline.generator(
noisy_image_or_video=denoised_pred,
conditional_dict=conditional_dict,
timestep=torch.ones_like(timestep) * pipeline.args.context_noise,
kv_cache=pipeline.kv_cache1,
crossattn_cache=pipeline.crossattn_cache,
current_start=chunk * base.TOKENS_PER_CHUNK,
)
finish_timing(timing_events, "context_dit", context_events)
torch.cuda.synchronize()
elapsed = {
category: sum(start.elapsed_time(end) for start, end in events)
for category, events in timing_events.items()
}
expected_full = NUM_CHUNKS * policy.full_calls_per_chunk
expected_reuse = NUM_CHUNKS * policy.reuse_calls_per_chunk
if (full_calls, reuse_calls) != (expected_full, expected_reuse):
raise RuntimeError(
f"{policy.name}: got calls F={full_calls}, R={reuse_calls}; "
f"expected F={expected_full}, R={expected_reuse}"
)
return torch.cat(output_chunks, dim=1), {
"generation_time_s": time.perf_counter() - started,
"full_calls": full_calls,
"reuse_calls": reuse_calls,
"predictor_calls": 0,
"accepted_predictor_calls": 0,
"rejected_predictor_calls": 0,
"full_dit_time_ms": elapsed["full_dit"],
"predictor_time_ms": 0.0,
"confidence_head_time_ms": 0.0,
"context_dit_time_ms": elapsed["context_dit"],
"actual_dit_time_ms": elapsed["full_dit"] + elapsed["context_dit"],
"model_path_time_ms": elapsed["full_dit"],
"raw_timesteps": list(policy.raw_timesteps),
"actual_timesteps": actual_timestep_values,
"decisions": decisions,
}
def compact_metrics(value: dict[str, Any]) -> dict[str, Any]:
keys = (
"pixel_mse",
"psnr",
"psnr_frame_mean",
"ssim",
"lpips",
"rollout_start_frame",
"rollout_pixel_mse",
"rollout_psnr",
"rollout_ssim",
"rollout_lpips",
"num_frames",
)
return {key: value[key] for key in keys}
def identity_metrics(num_frames: int) -> dict[str, Any]:
return {
"pixel_mse": 0.0,
"psnr": 120.0,
"psnr_frame_mean": 120.0,
"ssim": 1.0,
"lpips": 0.0,
"rollout_start_frame": base.PIXEL_FRAMES_FIRST_CHUNK,
"rollout_pixel_mse": 0.0,
"rollout_psnr": 120.0,
"rollout_ssim": 1.0,
"rollout_lpips": 0.0,
"num_frames": num_frames,
}
def load_models(device: torch.device) -> tuple[Any, Any, Any, Any]:
print("[setup] loading VAE", flush=True)
vae = WanVAEWrapper().to(device=device, dtype=torch.bfloat16).eval()
config = OmegaConf.merge(
OmegaConf.load(REPO_ROOT / "configs/default_config.yaml"),
OmegaConf.load(REPO_ROOT / "configs/self_forcing_sid.yaml"),
)
print("[setup] loading frozen Self-Forcing Teacher", flush=True)
pipeline = base.build_pipeline(
config, REPO_ROOT / "checkpoints/self_forcing_dmd.pt", vae, device
)
text_encoder = WanTextEncoder().to(device=device, dtype=torch.bfloat16).eval()
text_encoder.requires_grad_(False)
print("[setup] loading LPIPS", flush=True)
lpips_model = lpips.LPIPS(net="alex", verbose=False).to(device).eval()
lpips_model.requires_grad_(False)
return vae, pipeline, text_encoder, lpips_model
def artifact_paths(
output_root: Path, mapping_row: dict[str, Any], strategy: str
) -> tuple[Path, Path]:
suite = str(mapping_row["prompt_suite"])
suite_index = int(mapping_row["suite_index"])
global_index = int(mapping_row["global_index"])
return (
output_root / "generated_videos" / strategy / suite / f"{suite_index:03d}.mp4",
output_root
/ "generation_metrics/per_prompt"
/ strategy
/ f"global_{global_index:04d}.json",
)
def is_prompt_complete(
output_root: Path, mapping_row: dict[str, Any], policies: tuple[NaivePolicy, ...]
) -> bool:
for policy in (REFERENCE_POLICY, *policies):
video, record = artifact_paths(output_root, mapping_row, policy.name)
if not video.is_file() or not record.is_file():
return False
return True
def write_result(
*,
output_root: Path,
mapping_row: dict[str, Any],
policy: NaivePolicy,
seed: int,
diagnostic: dict[str, Any],
pixel_metrics: dict[str, Any],
frames: torch.Tensor,
) -> None:
video_path, record_path = artifact_paths(output_root, mapping_row, policy.name)
atomic_video(frames, video_path)
diagnostic_without_decisions = {
key: value for key, value in diagnostic.items() if key != "decisions"
}
record = {
"status": "complete",
"strategy": policy.name,
"label": policy.label,
"policy": policy.to_dict(),
"global_index": int(mapping_row["global_index"]),
"prompt_suite": str(mapping_row["prompt_suite"]),
"suite_index": int(mapping_row["suite_index"]),
"prompt": str(mapping_row["extended_prompt"]),
"seed": seed,
"generation": diagnostic_without_decisions,
"decisions": diagnostic["decisions"],
"pixel_metrics_vs_ffff": pixel_metrics,
"video": str(video_path.relative_to(output_root)),
}
atomic_json(record_path, record)
def run_prompt(
*,
mapping_row: dict[str, Any],
policies: tuple[NaivePolicy, ...],
output_root: Path,
seed: int,
vae: Any,
pipeline: Any,
text_encoder: Any,
lpips_model: Any,
device: torch.device,
) -> None:
global_index = int(mapping_row["global_index"])
suite = str(mapping_row["prompt_suite"])
suite_index = int(mapping_row["suite_index"])
prompt = str(mapping_row["extended_prompt"])
print(f"[prompt] global={global_index} suite={suite}/{suite_index}", flush=True)
conditional = text_encoder(text_prompts=[prompt])
reference_latent, reference_diagnostic = generate_rollout(
pipeline=pipeline,
conditional_dict=conditional,
seed=seed,
device=device,
policy=REFERENCE_POLICY,
)
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
reference_pixels = vae.decode_to_pixel(reference_latent, use_cache=False)
reference_u8 = base.pixels_to_u8(reference_pixels)
if reference_u8.ndim != 4 or reference_u8.shape[0] != 81:
raise RuntimeError(f"Unexpected reference shape: {tuple(reference_u8.shape)}")
write_result(
output_root=output_root,
mapping_row=mapping_row,
policy=REFERENCE_POLICY,
seed=seed,
diagnostic=reference_diagnostic,
pixel_metrics=identity_metrics(int(reference_u8.shape[0])),
frames=reference_u8,
)
for policy in policies:
latent, diagnostic = generate_rollout(
pipeline=pipeline,
conditional_dict=conditional,
seed=seed,
device=device,
policy=policy,
)
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
pixels = vae.decode_to_pixel(latent, use_cache=False)
prediction_u8 = base.pixels_to_u8(pixels)
metrics = base.frame_metrics(
reference_u8=reference_u8,
prediction_u8=prediction_u8,
lpips_model=lpips_model,
batch_size=4,
device=device,
)
write_result(
output_root=output_root,
mapping_row=mapping_row,
policy=policy,
seed=seed,
diagnostic=diagnostic,
pixel_metrics=compact_metrics(metrics),
frames=prediction_u8,
)
print(
f"[result] global={global_index} strategy={policy.name} "
f"F={diagnostic['full_calls']} R={diagnostic['reuse_calls']} "
f"psnr={metrics['psnr']:.4f} ssim={metrics['ssim']:.6f} "
f"lpips={metrics['lpips']:.6f}",
flush=True,
)
if hasattr(vae.model, "clear_cache"):
vae.model.clear_cache()
del latent, pixels, prediction_u8, metrics
torch.cuda.empty_cache()
if hasattr(vae.model, "clear_cache"):
vae.model.clear_cache()
del conditional, reference_latent, reference_pixels, reference_u8
torch.cuda.empty_cache()
def write_manifest(
*,
output_root: Path,
gpu: str,
shard_index: int,
num_shards: int,
rows: list[dict[str, Any]],
policies: tuple[NaivePolicy, ...],
seed: int,
status: str,
) -> None:
selected = "_".join(policy.name for policy in policies)
selection_id = hashlib.sha256(selected.encode("utf-8")).hexdigest()[:8]
manifest_name = f"naive_gpu{gpu}_shard{shard_index:02d}_{selection_id}.json"
atomic_json(
output_root / "generation_metrics/manifests" / manifest_name,
{
"status": status,
"physical_gpu": gpu,
"shard_index": shard_index,
"num_shards": num_shards,
"global_indices": [int(row["global_index"]) for row in rows],
"num_prompts": len(rows),
"strategies": [REFERENCE_POLICY.name, *(policy.name for policy in policies)],
"seed": seed,
},
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--gpu", default=PHYSICAL_GPU)
parser.add_argument("--mapping", type=Path, default=MAPPING_DEFAULT)
parser.add_argument("--output-root", type=Path, default=OUTPUT_DEFAULT)
parser.add_argument("--shard-index", type=int, default=0)
parser.add_argument("--num-shards", type=int, default=1)
parser.add_argument(
"--strategy",
action="append",
choices=EVALUATION_STRATEGY_NAMES,
default=None,
help="Baseline to generate; repeat to select several. Default: all six.",
)
parser.add_argument(
"--global-index",
dest="global_indices",
action="append",
type=int,
default=None,
help="Process only this global mapping index; may be repeated.",
)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--overwrite", action="store_true")
return parser.parse_args()
def main() -> None:
args = parse_args()
if args.num_shards < 1 or not 0 <= args.shard_index < args.num_shards:
raise ValueError("Invalid shard index/number of shards")
strategy_names = args.strategy or list(EVALUATION_STRATEGY_NAMES)
if len(set(strategy_names)) != len(strategy_names):
raise ValueError("--strategy values must be unique")
policies = tuple(POLICIES_BY_NAME[name] for name in strategy_names)
mapping = read_mapping(args.mapping.resolve())
if args.global_indices is None:
rows = [
row
for position, row in enumerate(mapping)
if position % args.num_shards == args.shard_index
]
else:
requested = set(args.global_indices)
mapping_by_global = {int(row["global_index"]): row for row in mapping}
unknown = sorted(requested - set(mapping_by_global))
if unknown:
raise ValueError(f"Unknown global indices: {unknown}")
rows = [row for row in mapping if int(row["global_index"]) in requested]
output_root = args.output_root.resolve()
output_root.mkdir(parents=True, exist_ok=True)
write_manifest(
output_root=output_root,
gpu=str(args.gpu),
shard_index=args.shard_index,
num_shards=args.num_shards,
rows=rows,
policies=policies,
seed=args.seed,
status="running",
)
print(
f"[setup] physical_gpu={args.gpu} shard={args.shard_index}/{args.num_shards} "
f"prompts={len(rows)} strategies={','.join(strategy_names)}",
flush=True,
)
device = torch.device("cuda")
torch.set_grad_enabled(False)
vae, pipeline, text_encoder, lpips_model = load_models(device)
completed = 0
try:
for row in rows:
if not args.overwrite and is_prompt_complete(output_root, row, policies):
completed += 1
print(
f"[cached] {completed}/{len(rows)} global={row['global_index']}",
flush=True,
)
continue
run_prompt(
mapping_row=row,
policies=policies,
output_root=output_root,
seed=args.seed,
vae=vae,
pipeline=pipeline,
text_encoder=text_encoder,
lpips_model=lpips_model,
device=device,
)
completed += 1
print(f"[progress] {completed}/{len(rows)} prompts", flush=True)
finally:
write_manifest(
output_root=output_root,
gpu=str(args.gpu),
shard_index=args.shard_index,
num_shards=args.num_shards,
rows=rows,
policies=policies,
seed=args.seed,
status="complete" if completed == len(rows) else "failed",
)
print(f"[complete] gpu={args.gpu} prompts={completed}/{len(rows)}", flush=True)
if __name__ == "__main__":
main()