File size: 3,046 Bytes
bc29ee3 | 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 | #!/usr/bin/env python3
"""Generate original-concat FPPP videos against the existing FFFF reference."""
from __future__ import annotations
import sys
from pathlib import Path
from typing import Any
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 import generate_vbench8_extended_disca as runner
import lpips
import torch
from omegaconf import OmegaConf
from utils.wan_wrapper import WanTextEncoder, WanVAEWrapper
WEIGHTS_DEFAULT = (
REPO_ROOT
/ "training_runs/layer17_stage1_1000p_4gpu_b16_2000steps"
/ "checkpoint_step_2000/predictor.safetensors"
)
OUTPUT_DEFAULT = (
REPO_ROOT
/ "evaluation_runs/vbench8_extended_concat_fppp_stage1_step2000_20260902"
)
def strategy_configs(patterns: list[str] | tuple[str, ...]) -> list[dict[str, Any]]:
if patterns != ["FPPP"] and tuple(patterns) != ("FPPP",):
raise ValueError("Original-concat evaluation supports exactly --pattern FPPP")
return [
{
"name": "concat_fppp",
"policy": "fppf",
"pattern": "FPPP",
"candidate_steps": [1, 2, 3],
"allow_chunk0_predictor": False,
"beta": None,
"threshold": None,
"target_accepts": 18,
"head": None,
"input_variant": "self_forcing",
}
]
def load_models(
weights: Path, device: torch.device
) -> tuple[Any, Any, Any, Any, None, None, 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 Teacher and original concat Predictor", flush=True)
pipeline = base.build_pipeline(
config, REPO_ROOT / "checkpoints/self_forcing_dmd.pt", vae, device
)
predictor = base.load_predictor(
pipeline.generator.model,
{
"source_layer": 17,
"weights": weights,
"input_variant": "self_forcing",
"gate_mode": "baseline",
},
device,
)
if predictor.input_variant != "self_forcing":
raise RuntimeError("Loaded Predictor is not the original concat variant")
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, predictor, text_encoder, None, None, lpips_model
def main() -> None:
runner.WEIGHTS_DEFAULT = WEIGHTS_DEFAULT
runner.OUTPUT_DEFAULT = OUTPUT_DEFAULT
runner.strategy_configs = strategy_configs
runner.load_models = load_models
runner.main()
if __name__ == "__main__":
main()
|