vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-08-04 05:10:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 53
unsloth/scripts/video_quality.py
Daniel Han 2b8998c8eb Studio: tighten comments across the image generation, video and training code 
Rewrite the comments added by this branch to be shorter and clearer:
collapse multi-line explanations into one or two lines, drop comments that
only restate the code, and keep the rationale that explains why a choice was
made. No code changes.
2026-07-31 17:47:47 +00:00

504 lines
19 KiB
Python
Raw Permalink Blame History

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Video quality-vs-cost harness for the Studio video backend.
The video analogue of scripts/diffusion_quality.py: hold the prompt + seed +
shape fixed, render one clip with a high-fidelity reference configuration
(default the family's BF16 artifact), then render the same clip with each
candidate configuration (a GGUF quant, a dense torchao quant, a speed profile,
a step cache) and measure how far the output drifts from the reference.
Per candidate it reports:
- mean PSNR / SSIM over evenly sampled frames (pixel + structural fidelity),
- a temporal-consistency deviation: the relative error between the reference's
and the candidate's frame-to-frame motion-energy series, which catches
flicker/juddering that per-frame SSIM alone can miss,
- black-frame and NaN collapse checks (the failure mode quant bugs actually
produce, per the image backend's qwen fp8 incident),
- an audio check for families that generate sound (LTX-2): RMS ratio vs the
reference and a silence trip-wire,
- wall time per generate and peak VRAM.
Verdict bands map the standing accuracy budget: a candidate that keeps mean
SSIM at or above 0.75 PASSes (a ~25 percent structural drift is acceptable for
a large speed/memory win), 0.50-0.75 WARNs, and below 0.50 or any black/NaN/
silence collapse FAILs regardless of how fast it is.
Runtime-budgeted: ONE short clip per candidate (default 33 frames at 480p-class
sizes) so a full family sweep stays in minutes, not hours. Metrics are pure
numpy; torch / diffusers / the backend load lazily so --help and --selftest run
on a host without them.
Examples:
# CPU metric sanity check (no GPU, no model):
python scripts/video_quality.py --selftest
# LTX-2.3 GGUF quants against the BF16 GGUF reference:
CUDA_VISIBLE_DEVICES=1 python scripts/video_quality.py \\
--model unsloth/LTX-2.3-GGUF --model-kind gguf \\
--reference "gguf_filename=distilled-1.1/ltx-2.3-22b-distilled-1.1-BF16.gguf" \\
--candidates "gguf_filename=distilled-1.1/ltx-2.3-22b-distilled-1.1-Q8_0.gguf" \\
"gguf_filename=distilled-1.1/ltx-2.3-22b-distilled-1.1-UD-Q4_K_M.gguf" \\
--steps 8 --guidance 1.0 --out-dir outputs/video_quality/ltx23
# Wan2.2-5B dense int8 + speed profiles against plain bf16:
CUDA_VISIBLE_DEVICES=1 python scripts/video_quality.py \\
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \\
--reference "" \\
--candidates "transformer_quant=int8" "speed_mode=max" \\
--steps 20 --out-dir outputs/video_quality/wan5b
"""
from __future__ import annotations
import argparse
import json
import math
import sys
import time
from pathlib import Path
from typing import Any, Optional
_BACKEND_ROOT = Path(__file__).resolve().parent.parent / "studio" / "backend"
if str(_BACKEND_ROOT) not in sys.path:
sys.path.insert(0, str(_BACKEND_ROOT))
DEFAULT_PROMPT = (
"a golden retriever puppy runs through shallow ocean waves at sunset, "
"splashing water, cinematic, camera tracking sideways"
)
# Finite PSNR (dB) cap for an identical clip, matching scripts/diffusion_quality.py.
_PERFECT_MATCH_PSNR = 100.0
# ── frame metrics (pure numpy; frames are uint8 HxWx3 arrays) ────────────────
def _gray(frame: Any) -> Any:
import numpy as np
f = np.asarray(frame, dtype = np.float64)
return f @ np.array([0.299, 0.587, 0.114])
def frame_psnr(a: Any, b: Any) -> float:
import numpy as np
a64 = np.asarray(a, dtype = np.float64)
b64 = np.asarray(b, dtype = np.float64)
if a64.shape != b64.shape:
return 0.0
mse = float(((a64 - b64) ** 2).mean())
if mse == 0.0:
return math.inf
return 20.0 * math.log10(255.0) - 10.0 * math.log10(mse)
def _box_mean(x: Any, w: int) -> Any:
import numpy as np
r = w // 2
xp = np.pad(x, r, mode = "edge")
ii = np.cumsum(np.cumsum(xp, axis = 0), axis = 1)
ii = np.pad(ii, ((1, 0), (1, 0)), mode = "constant")
h, wd = x.shape
total = ii[w : h + w, w : wd + w] - ii[0:h, w : wd + w] - ii[w : h + w, 0:wd] + ii[0:h, 0:wd]
return total / float(w * w)
def frame_ssim(
a: Any,
b: Any,
window: int = 7,
) -> float:
"""Pure numpy box-window SSIM on luminance (Wang et al. constants); identical
math to scripts/diffusion_quality.py so image and video budgets compare."""
ga, gb = _gray(a), _gray(b)
if ga.shape != gb.shape:
return 0.0
c1, c2 = (0.01 * 255) ** 2, (0.03 * 255) ** 2
mu_a, mu_b = _box_mean(ga, window), _box_mean(gb, window)
mu_a2, mu_b2, mu_ab = mu_a * mu_a, mu_b * mu_b, mu_a * mu_b
var_a = _box_mean(ga * ga, window) - mu_a2
var_b = _box_mean(gb * gb, window) - mu_b2
cov_ab = _box_mean(ga * gb, window) - mu_ab
ssim_map = ((2 * mu_ab + c1) * (2 * cov_ab + c2)) / (
(mu_a2 + mu_b2 + c1) * (var_a + var_b + c2)
)
return float(ssim_map.mean())
def motion_energy(frames: Any) -> list[float]:
"""Mean absolute frame-to-frame luminance difference, one value per frame
transition. The temporal signature of the clip: flicker inflates it, frozen
or smeared motion deflates it."""
import numpy as np
grays = [_gray(f) for f in frames]
return [float(np.abs(grays[i + 1] - grays[i]).mean()) for i in range(len(grays) - 1)]
def temporal_deviation(ref_frames: Any, cand_frames: Any) -> float:
"""Relative L1 error between the two motion-energy series (0 = identical
temporal behaviour). Series lengths must match (same frame count)."""
ref_series = motion_energy(ref_frames)
cand_series = motion_energy(cand_frames)
if len(ref_series) != len(cand_series) or not ref_series:
return math.inf
denom = sum(abs(v) for v in ref_series) + 1e-6
return sum(abs(r - c) for r, c in zip(ref_series, cand_series)) / denom
def clip_metrics(
ref_frames: Any,
cand_frames: Any,
sample_count: int = 5,
) -> dict[str, Any]:
"""All frame metrics for one candidate clip vs the reference clip."""
import numpy as np
# A truncated candidate is gated FAIL, not prefix-compared: good early frames would mask the missing tail.
ref_count, cand_count = len(ref_frames), len(cand_frames)
frame_count_mismatch = ref_count != cand_count
n = min(ref_count, cand_count)
if n == 0:
# An empty/corrupt decode must gate as FAIL, not crash the whole run.
return {
"frames_compared": 0,
"ref_frame_count": ref_count,
"cand_frame_count": cand_count,
"frame_count_mismatch": frame_count_mismatch,
"psnr_mean": 0.0,
"ssim_mean": 0.0,
"temporal_deviation": math.inf,
"min_luma": 0.0,
"has_nan": True,
}
idx = sorted({int(round(i * (n - 1) / max(1, sample_count - 1))) for i in range(sample_count)})
psnrs = [min(frame_psnr(ref_frames[i], cand_frames[i]), _PERFECT_MATCH_PSNR) for i in idx]
ssims = [frame_ssim(ref_frames[i], cand_frames[i]) for i in idx]
lumas = [float(_gray(cand_frames[i]).mean() / 255.0) for i in idx]
has_nan = any(
bool(np.isnan(np.asarray(f, dtype = np.float64)).any()) for f in (cand_frames[i] for i in idx)
)
return {
"frames_compared": len(idx),
"ref_frame_count": ref_count,
"cand_frame_count": cand_count,
"frame_count_mismatch": frame_count_mismatch,
"psnr_mean": sum(psnrs) / len(psnrs),
"ssim_mean": sum(ssims) / len(ssims),
"temporal_deviation": temporal_deviation(ref_frames[:n], cand_frames[:n]),
"min_luma": min(lumas),
"has_nan": has_nan,
}
def audio_metrics(ref_audio: Optional[Any], cand_audio: Optional[Any]) -> dict[str, Any]:
"""RMS comparison for families with sound. None audio on both sides is fine;
losing the track (or emitting silence) when the reference has one is not."""
import numpy as np
def _rms(a: Any) -> Optional[float]:
if a is None:
return None
arr = np.asarray(a, dtype = np.float64)
return float(np.sqrt((arr**2).mean())) if arr.size else 0.0
ref_rms, cand_rms = _rms(ref_audio), _rms(cand_audio)
# NaN compares False against any threshold, so call it out: a NaN track is a collapse.
silent_collapse = (
ref_rms is not None
and ref_rms >= 1e-3
and (cand_rms is None or math.isnan(cand_rms) or cand_rms < 1e-4)
)
return {"ref_rms": ref_rms, "cand_rms": cand_rms, "silent_collapse": silent_collapse}
def verdict(metrics: dict[str, Any], audio: dict[str, Any]) -> str:
"""PASS / WARN / FAIL per the standing accuracy budget (~25 percent structural
drift acceptable, 50 percent or a collapse never)."""
if (
metrics["has_nan"]
or metrics.get("frame_count_mismatch")
or metrics["min_luma"] < 0.02
or audio.get("silent_collapse")
):
return "FAIL"
if metrics["ssim_mean"] < 0.50 or metrics["temporal_deviation"] > 1.0:
return "FAIL"
if metrics["ssim_mean"] < 0.75 or metrics["temporal_deviation"] > 0.5:
return "WARN"
return "PASS"
# ── mp4 decode (PyAV, same dependency the backend encodes with) ─────────────
def decode_mp4(mp4_bytes: bytes, workdir: Path, name: str) -> tuple[list[Any], Optional[Any]]:
"""Frames (uint8 arrays) + mono audio samples (float array or None) from bytes."""
import av
import numpy as np
path = workdir / f"{name}.mp4"
path.write_bytes(mp4_bytes)
container = av.open(str(path))
frames = [f.to_ndarray(format = "rgb24") for f in container.decode(container.streams.video[0])]
audio = None
if container.streams.audio:
container.close()
container = av.open(str(path))
chunks = [c.to_ndarray() for c in container.decode(container.streams.audio[0])]
if chunks:
audio = np.concatenate([c.reshape(c.shape[0], -1).mean(axis = 0) for c in chunks])
container.close()
return frames, audio
# ── configuration plumbing ───────────────────────────────────────────────────
def parse_spec(spec: str) -> dict[str, str]:
"""'k=v;k=v' (or space-free 'k=v,k=v') -> dict; empty string -> {} (pure base)."""
out: dict[str, str] = {}
for part in spec.replace(",", ";").split(";"):
part = part.strip()
if not part:
continue
if "=" not in part:
raise ValueError(f"Bad candidate spec fragment '{part}' (expected key=value)")
key, value = part.split("=", 1)
out[key.strip()] = value.strip()
return out
def spec_label(spec: dict[str, str]) -> str:
if not spec:
return "base"
return ",".join(
f"{k}={Path(v).name if k == 'gguf_filename' else v}" for k, v in sorted(spec.items())
)
def run_config(
backend: Any, args: Any, spec: dict[str, str], workdir: Path, name: str
) -> dict[str, Any]:
"""Load per spec, generate the fixed clip, unload. Returns frames/audio/cost."""
import torch
load_kwargs: dict[str, Any] = {
"gguf_filename": spec.get("gguf_filename"),
"model_kind": spec.get("model_kind", args.model_kind),
"memory_mode": spec.get("memory_mode"),
"speed_mode": spec.get("speed_mode"),
"attention_backend": spec.get("attention_backend"),
"transformer_cache": spec.get("transformer_cache"),
"transformer_quant": spec.get("transformer_quant"),
}
t0 = time.monotonic()
status = backend.load_pipeline(args.model, **load_kwargs)
load_s = time.monotonic() - t0
if torch.cuda.is_available():
torch.cuda.reset_peak_memory_stats()
t0 = time.monotonic()
result = backend.generate(
prompt = args.prompt,
width = args.width,
height = args.height,
num_frames = args.frames,
fps = args.fps,
steps = args.steps,
guidance = args.guidance,
seed = args.seed,
)
generate_s = time.monotonic() - t0
peak_gib = torch.cuda.max_memory_allocated() / 2**30 if torch.cuda.is_available() else 0.0
backend.unload()
frames, audio = decode_mp4(result["mp4_bytes"], workdir, name)
return {
"frames": frames,
"audio": audio,
"load_s": round(load_s, 1),
"generate_s": round(generate_s, 1),
"peak_vram_gib": round(peak_gib, 2),
"resolved": {
k: v
for k, v in status.items()
if k
in (
"speed_mode",
"attention_backend",
"transformer_cache",
"transformer_quant",
"offload_policy",
"model_kind",
)
},
}
def run_gate(args: Any) -> int:
from core.inference.video import get_video_backend
out_dir = Path(args.out_dir)
out_dir.mkdir(parents = True, exist_ok = True)
backend = get_video_backend()
print(f"reference: {args.reference or 'base'}", flush = True)
ref = run_config(backend, args, parse_spec(args.reference), out_dir, "reference")
print(
f" load {ref['load_s']}s, generate {ref['generate_s']}s, "
f"peak {ref['peak_vram_gib']} GiB",
flush = True,
)
rows = []
for spec_str in args.candidates:
spec = parse_spec(spec_str)
label = spec_label(spec)
print(f"candidate: {label}", flush = True)
cand = run_config(backend, args, spec, out_dir, label.replace("/", "_").replace("=", "-"))
metrics = clip_metrics(ref["frames"], cand["frames"], sample_count = args.sample_frames)
audio = audio_metrics(ref["audio"], cand["audio"])
row = {
"candidate": label,
**{
k: (round(v, 4) if isinstance(v, float) and math.isfinite(v) else v)
for k, v in metrics.items()
},
**{f"audio_{k}": v for k, v in audio.items()},
"load_s": cand["load_s"],
"generate_s": cand["generate_s"],
"ref_generate_s": ref["generate_s"],
"peak_vram_gib": cand["peak_vram_gib"],
"resolved": cand["resolved"],
"verdict": verdict(metrics, audio),
}
rows.append(row)
print(
f" ssim {row['ssim_mean']:.3f} | psnr {row['psnr_mean']:.1f} dB | "
f"temporal {row['temporal_deviation']:.3f} | luma>={row['min_luma']:.3f} | "
f"gen {row['generate_s']}s (ref {ref['generate_s']}s) | "
f"vram {row['peak_vram_gib']} GiB | {row['verdict']}",
flush = True,
)
report = {
"model": args.model,
"reference": args.reference or "base",
"prompt": args.prompt,
"shape": [args.width, args.height, args.frames, args.fps],
"steps": args.steps,
"guidance": args.guidance,
"seed": args.seed,
"reference_cost": {k: ref[k] for k in ("load_s", "generate_s", "peak_vram_gib")},
"candidates": rows,
}
(out_dir / "report.json").write_text(json.dumps(report, indent = 1))
print(f"report: {out_dir / 'report.json'}", flush = True)
return 0 if all(r["verdict"] != "FAIL" for r in rows) else 1
# ── selftest (CPU-only, synthetic clips, no torch/model) ────────────────────
def selftest() -> int:
import numpy as np
rng = np.random.default_rng(0)
h, w, n = 64, 96, 12
def make_clip(
offset = 0.0,
noise = 0.0,
black = False,
):
frames = []
for t in range(n):
x = np.linspace(0, 1, w)[None, :] + t * 0.05 + offset
base = (np.sin(x * 6.283) * 0.5 + 0.5) * 255.0
frame = np.repeat(base[..., None], 3, axis = 2) * np.ones((h, 1, 1))
if noise:
frame = frame + rng.normal(0, noise, frame.shape)
if black:
frame = frame * 0.0
frames.append(np.clip(frame, 0, 255).astype(np.uint8))
return frames
ref = make_clip()
ok = True
def check(cond, msg):
nonlocal ok
print(("PASS: " if cond else "FAIL: ") + msg)
ok = ok and cond
same = clip_metrics(ref, make_clip())
check(
same["ssim_mean"] > 0.99 and same["temporal_deviation"] < 0.01,
f"identical clip scores ~1 (ssim {same['ssim_mean']:.3f})",
)
check(verdict(same, {"silent_collapse": False}) == "PASS", "identical clip verdict PASS")
noisy = clip_metrics(ref, make_clip(noise = 12.0))
check(0.3 < noisy["ssim_mean"] < 0.99, f"noisy clip degrades ssim ({noisy['ssim_mean']:.3f})")
black = clip_metrics(ref, make_clip(black = True))
check(
verdict(black, {"silent_collapse": False}) == "FAIL",
f"black clip verdict FAIL (min_luma {black['min_luma']:.3f})",
)
shifted = clip_metrics(ref, make_clip(offset = 0.5))
check(shifted["ssim_mean"] < same["ssim_mean"], "content shift lowers ssim")
# A truncated render with a pixel-identical prefix must still FAIL on the frame-count mismatch.
truncated = clip_metrics(ref, make_clip()[: n // 2])
check(
truncated["frame_count_mismatch"] is True
and verdict(truncated, {"silent_collapse": False}) == "FAIL",
f"truncated clip verdict FAIL ({truncated['cand_frame_count']}/{truncated['ref_frame_count']} frames)",
)
audio = audio_metrics(np.sin(np.linspace(0, 100, 16000)), np.zeros(16000))
check(audio["silent_collapse"] is True, "silent audio collapse detected")
audio_ok = audio_metrics(
np.sin(np.linspace(0, 100, 16000)), np.sin(np.linspace(0, 100, 16000)) * 0.8
)
check(audio_ok["silent_collapse"] is False, "attenuated audio is not a collapse")
print("VIDEO-QUALITY-SELFTEST", "PASS" if ok else "FAIL")
return 0 if ok else 1
def main() -> int:
parser = argparse.ArgumentParser(description = __doc__.split("\n")[0])
parser.add_argument("--selftest", action = "store_true", help = "CPU metric sanity check")
parser.add_argument("--model", help = "Repo id handed to the video backend")
parser.add_argument("--model-kind", default = None, help = "pipeline | gguf | single_file")
parser.add_argument(
"--reference", default = "", help = "Reference spec 'k=v;k=v' ('' = plain base load)"
)
parser.add_argument("--candidates", nargs = "+", default = [], help = "Candidate specs 'k=v;k=v'")
parser.add_argument("--prompt", default = DEFAULT_PROMPT)
parser.add_argument("--width", type = int, default = 768)
parser.add_argument("--height", type = int, default = 512)
parser.add_argument("--frames", type = int, default = 33)
parser.add_argument("--fps", type = int, default = 24)
parser.add_argument("--steps", type = int, default = None)
parser.add_argument("--guidance", type = float, default = None)
parser.add_argument("--seed", type = int, default = 7)
parser.add_argument("--sample-frames", type = int, default = 5)
parser.add_argument("--out-dir", default = "outputs/video_quality")
args = parser.parse_args()
if args.selftest:
return selftest()
if not args.model or not args.candidates:
parser.error("--model and --candidates are required (or use --selftest)")
return run_gate(args)
if __name__ == "__main__":
sys.exit(main())
Reference in a new issue View git blame Copy permalink