[feat]: Apple Silicon MLX runtime — INT8 Wan2.1 and Wan2.2 inference (#1638)
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>
This commit is contained in:
Aryan Kumar
co-authored by
Aryan Kumar
Aryan Kumar
parent
8208536cd1
commit
8537dcd6de
@@ -0,0 +1,46 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Tiny MLX RIFE frame-interpolation smoke test."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastvideo.mlx_runtime.rife_interp import interpolate, load_model
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="MLX RIFE 4.25 frame interpolation smoke test."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--self-test",
|
||||
action="store_true",
|
||||
help="Run a tiny two-frame interpolation test.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
if not args.self_test:
|
||||
raise SystemExit("Nothing to do; pass --self-test")
|
||||
|
||||
frame0 = np.zeros((64, 96, 3), dtype=np.uint8)
|
||||
frame1 = np.zeros((64, 96, 3), dtype=np.uint8)
|
||||
frame1[:, :, 0] = 255
|
||||
start = time.perf_counter()
|
||||
model = load_model()
|
||||
load_s = time.perf_counter() - start
|
||||
start = time.perf_counter()
|
||||
frames = interpolate([frame0, frame1], factor=2, model=model)
|
||||
interp_s = time.perf_counter() - start
|
||||
assert len(frames) == 3
|
||||
assert frames[1].shape == frame0.shape
|
||||
assert frames[1].dtype == np.uint8
|
||||
print(
|
||||
"MLX RIFE self-test passed: "
|
||||
f"load_s={load_s:.3f} interp_s={interp_s:.3f} shape={frames[1].shape}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,510 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""End-to-end Wan2.2-TI2V-5B generation on Apple Silicon (MLX DiT + MLX TAEHV).
|
||||
|
||||
Pipeline: torch/MPS UMT5 encode (shared with 1.3B) → MLXWan22DiT 3-step DMD
|
||||
(warped schedule, flow_shift=5) → MLX TAEHV decode (taew2_2.pth). Fully MLX
|
||||
on the heavy DiT + decode path.
|
||||
|
||||
PYTHONPATH=$PWD python examples/inference/basic/mlx_wan22_generate.py \
|
||||
--prompt "A red fox trotting through a snowy pine forest at golden hour" \
|
||||
--output-path video_samples/demo_5b/fox_5b_mlx.mp4
|
||||
|
||||
Decoder backends: ``taehv`` (default, MLX, ~seconds), ``taehv-torch`` (parity),
|
||||
``wan-vae`` (full AutoencoderKLWan on MPS, slow).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastvideo.mlx_runtime.fast_spatial import DEFAULT_FAST_SPATIAL_SHARPEN
|
||||
from fastvideo.mlx_runtime.frame_upsample import DEFAULT_PIXEL_UPSAMPLE_MODE, PIXEL_UPSAMPLE_MODES
|
||||
from fastvideo.mlx_runtime.memory import cleanup_mlx
|
||||
from fastvideo.mlx_runtime.prompt_cache import (
|
||||
fingerprint_digest,
|
||||
load_prompt_cache,
|
||||
save_prompt_cache,
|
||||
text_encoder_fingerprint,
|
||||
)
|
||||
from fastvideo.mlx_runtime.rife_interp import aligned_keyframe_count
|
||||
|
||||
FASTWAN21_MODEL_ID = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers"
|
||||
FASTWAN22_MODEL_ID = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
|
||||
DEFAULT_HEIGHT = 448
|
||||
DEFAULT_WIDTH = 832
|
||||
DEFAULT_NUM_FRAMES = 121
|
||||
|
||||
def _resolve_model_paths(
|
||||
*,
|
||||
text_encoder_root: Path | None,
|
||||
dit_checkpoint: Path | None,
|
||||
dit_config: Path | None,
|
||||
vae_root: Path | None,
|
||||
mlx_checkpoint: Path | None,
|
||||
decode_backend: str,
|
||||
) -> tuple[Path, Path | None, Path | None, Path | None]:
|
||||
"""Download only the missing assets required by the selected Wan2.2 path."""
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
if text_encoder_root is None:
|
||||
text_encoder_root = Path(snapshot_download(
|
||||
FASTWAN21_MODEL_ID,
|
||||
allow_patterns=["tokenizer/*", "text_encoder/*"],
|
||||
))
|
||||
if mlx_checkpoint is None and (dit_checkpoint is None or dit_config is None):
|
||||
patterns = []
|
||||
if dit_checkpoint is None:
|
||||
patterns.append("transformer/diffusion_pytorch_model.safetensors")
|
||||
if dit_config is None:
|
||||
patterns.append("transformer/config.json")
|
||||
model_root = Path(snapshot_download(FASTWAN22_MODEL_ID, allow_patterns=patterns))
|
||||
dit_checkpoint = dit_checkpoint or model_root / "transformer/diffusion_pytorch_model.safetensors"
|
||||
dit_config = dit_config or model_root / "transformer/config.json"
|
||||
if decode_backend == "wan-vae" and vae_root is None:
|
||||
model_root = Path(snapshot_download(FASTWAN22_MODEL_ID, allow_patterns=["vae/*"]))
|
||||
vae_root = model_root / "vae"
|
||||
return text_encoder_root, dit_checkpoint, dit_config, vae_root
|
||||
|
||||
|
||||
def _prompt_cache_fingerprint(
|
||||
*,
|
||||
prompt: str,
|
||||
prompt_used: str,
|
||||
enhance_prompt: bool,
|
||||
enhance_prompt_backend: str,
|
||||
text_encoder_root: Path,
|
||||
max_sequence_length: int,
|
||||
dtype: str,
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"prompt_used": prompt_used,
|
||||
"enhance_prompt": enhance_prompt,
|
||||
"enhance_prompt_backend": enhance_prompt_backend,
|
||||
"text_encoder": text_encoder_fingerprint(text_encoder_root),
|
||||
"max_sequence_length": max_sequence_length,
|
||||
"dtype": dtype,
|
||||
}
|
||||
|
||||
|
||||
def _default_prompt_cache_path(fingerprint: dict[str, object]) -> Path:
|
||||
"""Content-addressed default cache file for a prompt fingerprint.
|
||||
|
||||
The Wan2.1 entrypoint caches prompt embeddings by default; this one only
|
||||
did so when handed an explicit ``--prompt-embeds-cache`` path, so every 5B
|
||||
run paid a full UMT5 encode (~45s on an M4 Max) even for a repeat prompt.
|
||||
The fingerprint already covers everything that changes the embedding, so
|
||||
hash it for the filename.
|
||||
"""
|
||||
digest = fingerprint_digest(fingerprint)[:32]
|
||||
return Path.home() / ".cache" / "fastvideo" / "prompt_embeds" / f"wan22_{digest}.npy"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="MLX Wan2.2-5B T2V (encode → DiT DMD → TAEHV/VAE decode)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
default="A red fox trotting through a snowy pine forest at golden hour, cinematic",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-path",
|
||||
type=Path,
|
||||
default=Path("video_samples/demo_5b/fox_5b_mlx.mp4"),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text-encoder-root",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Root with text_encoder/ + tokenizer/",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-embeds-cache",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Explicit .npy UMT5 embedding cache file. Overrides the automatic "
|
||||
"content-addressed cache (--prompt-cache).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-cache",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Cache prompt embeddings under ~/.cache/fastvideo/prompt_embeds so "
|
||||
"repeat runs skip the text encoder entirely. Default: on.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text-encoder-device",
|
||||
choices=("auto", "cpu", "mps"),
|
||||
default="cpu",
|
||||
help="Device for UMT5 encoding. CPU is safest beside the 5B MLX DiT.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enhance-prompt",
|
||||
action="store_true",
|
||||
help="Apply deterministic local cinematic prompt enrichment before UMT5.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enhance-prompt-backend",
|
||||
choices=("template",),
|
||||
default="template",
|
||||
help="Prompt enrichment backend.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dit-checkpoint",
|
||||
type=Path,
|
||||
default=None,
|
||||
)
|
||||
parser.add_argument("--dit-config", type=Path, default=None)
|
||||
parser.add_argument(
|
||||
"--mlx-checkpoint",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Pre-quantized MLX DiT checkpoint directory. Rewrapped with Wan2.2 per-token conditioning.",
|
||||
)
|
||||
parser.add_argument("--vae-root", type=Path, default=None)
|
||||
parser.add_argument("--height", type=int, default=DEFAULT_HEIGHT)
|
||||
parser.add_argument("--width", type=int, default=DEFAULT_WIDTH)
|
||||
parser.add_argument(
|
||||
"--num-frames",
|
||||
type=int,
|
||||
default=DEFAULT_NUM_FRAMES,
|
||||
help="Pixel frames (121 at 24fps = 5.04 seconds)",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=1234)
|
||||
parser.add_argument("--renoise-seed", type=int, default=0)
|
||||
parser.add_argument("--fps", type=int, default=24)
|
||||
parser.add_argument("--flow-shift", type=float, default=5.0)
|
||||
parser.add_argument("--dmd-denoising-steps", default="1000,757,522")
|
||||
parser.add_argument(
|
||||
"--no-warp",
|
||||
action="store_true",
|
||||
help="Disable schedule warping (debug only).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--fast",
|
||||
action="store_true",
|
||||
help="Generate fewer frames then RIFE-interpolate to --num-frames.",
|
||||
)
|
||||
parser.add_argument("--fast-factor", type=int, default=2)
|
||||
parser.add_argument("--fast-sharpen", type=float, default=0.6)
|
||||
parser.add_argument(
|
||||
"--fast-spatial",
|
||||
action="store_true",
|
||||
help="Denoise and decode at reduced spatial resolution, then resample "
|
||||
"the decoded frames up to the target size.",
|
||||
)
|
||||
parser.add_argument("--fast-spatial-scale", type=int, default=2)
|
||||
parser.add_argument(
|
||||
"--fast-spatial-upsample-mode",
|
||||
choices=PIXEL_UPSAMPLE_MODES,
|
||||
default=DEFAULT_PIXEL_UPSAMPLE_MODE,
|
||||
)
|
||||
parser.add_argument("--fast-spatial-sharpen", type=float, default=DEFAULT_FAST_SPATIAL_SHARPEN)
|
||||
parser.add_argument(
|
||||
"--refine",
|
||||
action="store_true",
|
||||
help="Two-pass DMD: coarse denoise, upsample/re-noise, full-res denoise.",
|
||||
)
|
||||
parser.add_argument("--refine-scale", type=int, default=2)
|
||||
parser.add_argument(
|
||||
"--refine-upsample-mode",
|
||||
choices=("bilinear", "nearest"),
|
||||
default="bilinear",
|
||||
)
|
||||
parser.add_argument("--no-refine-add-noise", action="store_true")
|
||||
parser.add_argument(
|
||||
"--decode-backend",
|
||||
choices=("taehv", "taehv-torch", "wan-vae"),
|
||||
default="taehv",
|
||||
)
|
||||
parser.add_argument("--save-latents", type=Path, default=None)
|
||||
parser.add_argument("--metrics-json", type=Path, default=None,
|
||||
help="Write measured run metadata as JSON for reports or galleries.")
|
||||
parser.add_argument(
|
||||
"--compile",
|
||||
action="store_true",
|
||||
help="Compile the DiT forward with mx.compile; fallback to eager on failure.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.fast_factor < 2:
|
||||
parser.error("--fast-factor must be at least 2")
|
||||
# --fast-spatial used to be rejected here because it upsampled the completed
|
||||
# 48-channel latent, which is out of distribution for the decoder and gave
|
||||
# black or noisy video. The upsample now runs on decoded frames, so the
|
||||
# latent never leaves the grid it was denoised on and the mode is usable.
|
||||
if args.refine and args.fast_spatial:
|
||||
print("[wan22] --refine takes precedence over --fast-spatial")
|
||||
args.text_encoder_root, args.dit_checkpoint, args.dit_config, args.vae_root = _resolve_model_paths(
|
||||
text_encoder_root=args.text_encoder_root,
|
||||
dit_checkpoint=args.dit_checkpoint,
|
||||
dit_config=args.dit_config,
|
||||
vae_root=args.vae_root,
|
||||
mlx_checkpoint=args.mlx_checkpoint,
|
||||
decode_backend=args.decode_backend,
|
||||
)
|
||||
target_frames = args.num_frames
|
||||
if args.fast:
|
||||
args.num_frames = aligned_keyframe_count(target_frames, args.fast_factor)
|
||||
print(
|
||||
f"[wan22 fast] generating {args.num_frames} frames, "
|
||||
f"RIFE {args.fast_factor}x -> {target_frames}"
|
||||
)
|
||||
|
||||
import mlx.core as mx
|
||||
import torch
|
||||
|
||||
from examples.inference.basic.mlx_wan_prompt_to_video import (
|
||||
_postprocess_video,
|
||||
encode_prompt,
|
||||
make_rotary_embeddings,
|
||||
)
|
||||
from fastvideo.mlx_runtime.fast_spatial import plan_fast_spatial
|
||||
from fastvideo.mlx_runtime.refine import (
|
||||
default_refine_timesteps,
|
||||
plan_refine_resolutions,
|
||||
prepare_refine_latents,
|
||||
)
|
||||
from fastvideo.mlx_runtime.wan22 import (
|
||||
mlx_wan22_dit_from_diffusers_safetensors,
|
||||
mlx_wan22_dit_from_mlx_checkpoint,
|
||||
)
|
||||
from fastvideo.mlx_runtime.wan22_sample import build_wan22_dmd_schedule, sample_wan22_dmd
|
||||
from fastvideo.mlx_runtime.wan_vae import decode_latents_to_video
|
||||
|
||||
if args.mlx_checkpoint is not None:
|
||||
config = json.loads((args.mlx_checkpoint / "mlx_dit.json").read_text())["config"]
|
||||
else:
|
||||
config = json.loads(args.dit_config.read_text())
|
||||
patch_size = tuple(config.get("patch_size", (1, 2, 2)))
|
||||
if args.refine:
|
||||
active_plan = plan_refine_resolutions(
|
||||
height=args.height, width=args.width, num_frames=args.num_frames,
|
||||
spatial_scale=args.refine_scale, vae_spatial_compression=16,
|
||||
vae_temporal_compression=4, patch_size=patch_size, enabled=True,
|
||||
)
|
||||
spatial_mode = "refine"
|
||||
elif args.fast_spatial:
|
||||
fast_spatial_plan = plan_fast_spatial(
|
||||
height=args.height, width=args.width, num_frames=args.num_frames,
|
||||
spatial_scale=args.fast_spatial_scale, vae_spatial_compression=16,
|
||||
vae_temporal_compression=4, patch_size=patch_size,
|
||||
upsample_mode=args.fast_spatial_upsample_mode,
|
||||
sharpen=args.fast_spatial_sharpen, enabled=True,
|
||||
)
|
||||
active_plan = fast_spatial_plan.plan
|
||||
spatial_mode = "fast_spatial"
|
||||
else:
|
||||
active_plan = plan_refine_resolutions(
|
||||
height=args.height, width=args.width, num_frames=args.num_frames,
|
||||
spatial_scale=1, vae_spatial_compression=16, vae_temporal_compression=4,
|
||||
patch_size=patch_size, enabled=False,
|
||||
)
|
||||
spatial_mode = "off"
|
||||
lat_h, lat_w = active_plan.stage1_latent_height, active_plan.stage1_latent_width
|
||||
lat_t = active_plan.latent_frames
|
||||
in_ch = int(config["in_channels"])
|
||||
print(f"[5B] latent {in_ch}x{lat_t}x{lat_h}x{lat_w}", flush=True)
|
||||
|
||||
total_start = time.perf_counter()
|
||||
prompt_for_encode = args.prompt
|
||||
enhance_backend = None
|
||||
enhance_elapsed_s = 0.0
|
||||
if args.enhance_prompt:
|
||||
from fastvideo.mlx_runtime.prompt_enhance import enhance_prompt
|
||||
|
||||
enhancement = enhance_prompt(args.prompt, backend=args.enhance_prompt_backend)
|
||||
prompt_for_encode = enhancement.enhanced
|
||||
enhance_backend = enhancement.backend
|
||||
enhance_elapsed_s = enhancement.elapsed_s
|
||||
print(f"[enhance] backend={enhance_backend} in {enhance_elapsed_s:.2f}s", flush=True)
|
||||
print(f"[enhance] prompt: {prompt_for_encode}", flush=True)
|
||||
|
||||
t0 = time.perf_counter()
|
||||
prompt_cache_fingerprint = _prompt_cache_fingerprint(
|
||||
prompt=args.prompt,
|
||||
prompt_used=prompt_for_encode,
|
||||
enhance_prompt=args.enhance_prompt,
|
||||
enhance_prompt_backend=args.enhance_prompt_backend,
|
||||
text_encoder_root=args.text_encoder_root,
|
||||
max_sequence_length=512,
|
||||
dtype="fp16",
|
||||
)
|
||||
prompt_cache_path = args.prompt_embeds_cache
|
||||
if prompt_cache_path is None and args.prompt_cache:
|
||||
prompt_cache_path = _default_prompt_cache_path(prompt_cache_fingerprint)
|
||||
cached_embeds = load_prompt_cache(
|
||||
prompt_cache_path,
|
||||
prompt_cache_fingerprint,
|
||||
)
|
||||
if cached_embeds is not None:
|
||||
embeds = torch.from_numpy(cached_embeds).contiguous()
|
||||
else:
|
||||
embeds = encode_prompt(
|
||||
model_root=args.text_encoder_root,
|
||||
prompt=prompt_for_encode,
|
||||
max_sequence_length=512,
|
||||
device_arg=args.text_encoder_device,
|
||||
dtype_arg="fp16",
|
||||
)
|
||||
save_prompt_cache(
|
||||
prompt_cache_path,
|
||||
embeds.cpu().numpy(),
|
||||
prompt_cache_fingerprint,
|
||||
)
|
||||
ehs = mx.array(embeds.numpy()).astype(mx.float16)
|
||||
prompt_encode_s = time.perf_counter() - t0
|
||||
print(f"[5B] prompt encoded {tuple(ehs.shape)} in {prompt_encode_s:.1f}s", flush=True)
|
||||
|
||||
t1 = time.perf_counter()
|
||||
if args.mlx_checkpoint is not None:
|
||||
dit = mlx_wan22_dit_from_mlx_checkpoint(
|
||||
args.mlx_checkpoint,
|
||||
compile=args.compile,
|
||||
)
|
||||
else:
|
||||
dit = mlx_wan22_dit_from_diffusers_safetensors(
|
||||
args.dit_checkpoint,
|
||||
args.dit_config,
|
||||
dtype="fp16",
|
||||
compile=args.compile,
|
||||
)
|
||||
dit_load_s = time.perf_counter() - t1
|
||||
print(f"[5B] DiT loaded in {dit_load_s:.1f}s", flush=True)
|
||||
|
||||
freqs = make_rotary_embeddings(config, latent_frames=lat_t, latent_height=lat_h, latent_width=lat_w)
|
||||
gen = torch.Generator().manual_seed(args.seed)
|
||||
noise = mx.array(
|
||||
torch.randn(1, in_ch, lat_t, lat_h, lat_w, generator=gen, dtype=torch.float32).numpy()).astype(mx.float16)
|
||||
|
||||
steps = [int(s) for s in args.dmd_denoising_steps.split(",") if s.strip()]
|
||||
t2 = time.perf_counter()
|
||||
mx.reset_peak_memory()
|
||||
latents = sample_wan22_dmd(
|
||||
dit,
|
||||
ehs,
|
||||
noise,
|
||||
freqs,
|
||||
dmd_denoising_steps=steps,
|
||||
flow_shift=args.flow_shift,
|
||||
warp_denoising_step=not args.no_warp,
|
||||
seed=args.renoise_seed,
|
||||
)
|
||||
if spatial_mode == "refine":
|
||||
schedule, warped_steps = build_wan22_dmd_schedule(
|
||||
steps, flow_shift=args.flow_shift, warp_denoising_step=not args.no_warp,
|
||||
)
|
||||
# The grid opens at sigma == 1, where the hand-off
|
||||
# `(1 - sigma) * upsampled + sigma * noise` weights stage 1 at zero and
|
||||
# refine silently becomes a plain full-res run. Drop the leading
|
||||
# full-noise steps so stage 1 actually reaches stage 2.
|
||||
stage2_warped = default_refine_timesteps(schedule, warped_steps)
|
||||
stage2_steps = steps[len(warped_steps) - len(stage2_warped):]
|
||||
sigma = schedule.sigma_for(stage2_warped[0])
|
||||
print(f"[5B refine] stage-2 steps={stage2_steps} sigma={sigma:.4f} "
|
||||
f"(stage-1 weight {1.0 - sigma:.4f})", flush=True)
|
||||
latents = prepare_refine_latents(
|
||||
latents, scale=args.refine_scale, sigma=sigma,
|
||||
add_noise_flag=not args.no_refine_add_noise,
|
||||
upsample_mode=args.refine_upsample_mode, seed=args.renoise_seed + 1,
|
||||
)
|
||||
freqs_stage2 = make_rotary_embeddings(
|
||||
config, latent_frames=lat_t,
|
||||
latent_height=active_plan.stage2_latent_height,
|
||||
latent_width=active_plan.stage2_latent_width,
|
||||
)
|
||||
latents = sample_wan22_dmd(
|
||||
dit, ehs, latents, freqs_stage2, dmd_denoising_steps=stage2_steps,
|
||||
flow_shift=args.flow_shift, warp_denoising_step=not args.no_warp,
|
||||
seed=args.renoise_seed + 2,
|
||||
)
|
||||
# spatial_mode == "fast_spatial" leaves the latents on the stage-1 grid;
|
||||
# the resample happens after decode, in _postprocess_video.
|
||||
denoise_s = time.perf_counter() - t2
|
||||
peak = mx.get_peak_memory() / (1024**3)
|
||||
print(f"[5B] denoise {len(steps)} steps in {denoise_s:.1f}s, peak {peak:.2f} GiB", flush=True)
|
||||
|
||||
latents_np = np.array(latents.astype(mx.float32))
|
||||
if args.save_latents is not None:
|
||||
args.save_latents.parent.mkdir(parents=True, exist_ok=True)
|
||||
np.savez(args.save_latents, latents=latents_np, prompt=args.prompt, seed=args.seed)
|
||||
print(f"[5B] wrote latents {args.save_latents}", flush=True)
|
||||
|
||||
if spatial_mode == "refine":
|
||||
del freqs_stage2
|
||||
del dit, latents, ehs, noise, freqs
|
||||
cleanup_mlx()
|
||||
|
||||
metrics = decode_latents_to_video(
|
||||
latents_np,
|
||||
args.output_path,
|
||||
fps=args.fps,
|
||||
backend=args.decode_backend,
|
||||
vae_dir=args.vae_root if args.decode_backend == "wan-vae" else None,
|
||||
z_dim=in_ch,
|
||||
)
|
||||
# One h264 round-trip for both post-decode passes (see _postprocess_video).
|
||||
rife_s = 0.0
|
||||
rife_request = ({
|
||||
"factor": args.fast_factor,
|
||||
"target_frames": target_frames,
|
||||
"sharpen": args.fast_sharpen,
|
||||
} if args.fast else None)
|
||||
spatial_request = fast_spatial_plan if spatial_mode == "fast_spatial" else None
|
||||
if rife_request is not None or spatial_request is not None:
|
||||
rife_start = time.perf_counter()
|
||||
_postprocess_video(
|
||||
video_path=args.output_path, fps=args.fps,
|
||||
rife=rife_request, spatial=spatial_request,
|
||||
)
|
||||
rife_s = time.perf_counter() - rife_start
|
||||
print(f"[5B] decoded via {metrics['backend']} in {metrics['decode_s']:.1f}s → {args.output_path}", flush=True)
|
||||
summary = {
|
||||
"output_path": str(args.output_path.resolve()),
|
||||
"prompt": args.prompt,
|
||||
"prompt_used": prompt_for_encode,
|
||||
"enhance_prompt": args.enhance_prompt,
|
||||
"enhance_backend": enhance_backend,
|
||||
"enhance_elapsed_s": round(enhance_elapsed_s, 3),
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"fps": args.fps,
|
||||
"target_frames": target_frames,
|
||||
"generated_frames": args.num_frames,
|
||||
"seed": args.seed,
|
||||
"renoise_seed": args.renoise_seed,
|
||||
"dmd_denoising_steps": steps,
|
||||
"flow_shift": args.flow_shift,
|
||||
"warp": not args.no_warp,
|
||||
"spatial_mode": spatial_mode,
|
||||
"fast": args.fast,
|
||||
"fast_factor": args.fast_factor if args.fast else None,
|
||||
"fast_spatial_scale": args.fast_spatial_scale if args.fast_spatial else None,
|
||||
"refine_scale": args.refine_scale if args.refine else None,
|
||||
"decode_backend": args.decode_backend,
|
||||
"prompt_encode_s": round(prompt_encode_s, 3),
|
||||
"dit_load_s": round(dit_load_s, 3),
|
||||
"denoise_s": round(denoise_s, 3),
|
||||
"decode_s": round(metrics["decode_s"], 3),
|
||||
"rife_s": round(rife_s, 3),
|
||||
"wall_total_s": round(time.perf_counter() - total_start, 3),
|
||||
"peak_gib": round(peak, 3),
|
||||
"latent_shape": [in_ch, lat_t, lat_h, lat_w],
|
||||
"stage2_latent_shape": [in_ch, lat_t, active_plan.stage2_latent_height, active_plan.stage2_latent_width],
|
||||
"mlx_checkpoint": str(args.mlx_checkpoint.resolve()) if args.mlx_checkpoint else None,
|
||||
}
|
||||
if args.metrics_json is not None:
|
||||
args.metrics_json.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.metrics_json.write_text(json.dumps(summary, indent=2) + "\n")
|
||||
print(f"[5B] wrote metrics {args.metrics_json}", flush=True)
|
||||
print(json.dumps(summary, indent=2), flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,130 @@
|
||||
"""Compare Wan VAE and TAEHV decode on saved FastWan latents."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from examples.inference.basic.mlx_wan_prompt_to_video import DEFAULT_MODEL_ROOT, decode_latents_to_video
|
||||
|
||||
|
||||
def _torch_mps_memory() -> dict[str, int | None]:
|
||||
"""
|
||||
Report current and recommended memory usage for the MPS backend.
|
||||
|
||||
Returns:
|
||||
dict[str, int | None]: Memory metrics in bytes, or `None` values when
|
||||
PyTorch or MPS is unavailable.
|
||||
"""
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
return {
|
||||
"current_allocated_bytes": None,
|
||||
"driver_allocated_bytes": None,
|
||||
"recommended_max_bytes": None,
|
||||
}
|
||||
if not torch.backends.mps.is_available():
|
||||
return {
|
||||
"current_allocated_bytes": None,
|
||||
"driver_allocated_bytes": None,
|
||||
"recommended_max_bytes": None,
|
||||
}
|
||||
return {
|
||||
"current_allocated_bytes": int(torch.mps.current_allocated_memory()),
|
||||
"driver_allocated_bytes": int(torch.mps.driver_allocated_memory()),
|
||||
"recommended_max_bytes": int(torch.mps.recommended_max_memory()),
|
||||
}
|
||||
|
||||
|
||||
def _parse_backends(raw: str) -> list[str]:
|
||||
"""
|
||||
Parse and validate a comma-separated list of decoding backends.
|
||||
|
||||
Parameters:
|
||||
raw (str): Comma-separated backend names.
|
||||
|
||||
Returns:
|
||||
list[str]: Trimmed, supported backend names in input order.
|
||||
|
||||
Raises:
|
||||
ValueError: If the input contains an unsupported backend.
|
||||
"""
|
||||
backends = [backend.strip() for backend in raw.split(",") if backend.strip()]
|
||||
allowed = {"wan-vae", "taehv"}
|
||||
unknown = sorted(set(backends) - allowed)
|
||||
if unknown:
|
||||
raise ValueError(f"Unsupported decode backends: {unknown}")
|
||||
return backends
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""
|
||||
Benchmark selected Wan latent decoding backends and record their performance metrics.
|
||||
|
||||
Loads the specified latent array, decodes it with each selected backend, exports the
|
||||
results as MP4 files, and writes per-backend timing and Torch MPS memory metrics to
|
||||
`metrics.json`.
|
||||
"""
|
||||
parser = argparse.ArgumentParser(description="Benchmark decode backends on saved Wan/FastWan latents.")
|
||||
parser.add_argument("--model-root", type=Path, default=DEFAULT_MODEL_ROOT)
|
||||
parser.add_argument("--latents-path", type=Path, required=True)
|
||||
parser.add_argument("--output-dir", type=Path, default=Path("video_samples/mlx_decode_benchmark"))
|
||||
parser.add_argument("--backends", default="wan-vae,taehv")
|
||||
parser.add_argument("--fps", type=int, default=16)
|
||||
parser.add_argument("--torch-device", default="auto")
|
||||
parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16")
|
||||
parser.add_argument("--taehv-source-path", type=Path, default=None)
|
||||
parser.add_argument("--taehv-checkpoint-path", type=Path, default=None)
|
||||
parser.add_argument("--taehv-parallel", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
latents = np.load(args.latents_path)
|
||||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
rows = []
|
||||
|
||||
for backend in _parse_backends(args.backends):
|
||||
print(f"=== Decode backend: {backend} ===")
|
||||
before = _torch_mps_memory()
|
||||
start = time.perf_counter()
|
||||
output_path = args.output_dir / f"{args.latents_path.stem}_{backend}.mp4"
|
||||
decode_latents_to_video(
|
||||
model_root=args.model_root,
|
||||
latents_np=latents,
|
||||
output_path=output_path,
|
||||
fps=args.fps,
|
||||
device_arg=args.torch_device,
|
||||
dtype_arg=args.torch_dtype,
|
||||
backend=backend,
|
||||
taehv_source_path=args.taehv_source_path,
|
||||
taehv_checkpoint_path=args.taehv_checkpoint_path,
|
||||
taehv_parallel=args.taehv_parallel,
|
||||
)
|
||||
elapsed = time.perf_counter() - start
|
||||
after = _torch_mps_memory()
|
||||
metrics = {
|
||||
"backend": backend,
|
||||
"latents_path": str(args.latents_path),
|
||||
"latents_shape": list(latents.shape),
|
||||
"decode_export_s": elapsed,
|
||||
"torch_mps_current_before_bytes": before["current_allocated_bytes"],
|
||||
"torch_mps_current_after_bytes": after["current_allocated_bytes"],
|
||||
"torch_mps_driver_before_bytes": before["driver_allocated_bytes"],
|
||||
"torch_mps_driver_after_bytes": after["driver_allocated_bytes"],
|
||||
"torch_mps_recommended_max_bytes": after["recommended_max_bytes"],
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
rows.append(metrics)
|
||||
print(json.dumps(metrics, indent=2))
|
||||
|
||||
metrics_path = args.output_dir / "metrics.json"
|
||||
metrics_path.write_text(json.dumps(rows, indent=2))
|
||||
print(f"Wrote decode metrics to: {metrics_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,365 @@
|
||||
"""Benchmark MLX FastWan quantization modes with one shared prompt encode."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
import numpy as np
|
||||
|
||||
from examples.inference.basic.mlx_wan_prompt_to_video import (
|
||||
DEFAULT_MODEL_ROOT,
|
||||
decode_latents_to_video,
|
||||
encode_prompt,
|
||||
make_rotary_embeddings,
|
||||
)
|
||||
from fastvideo.mlx_runtime.memory import cleanup_mlx
|
||||
|
||||
|
||||
def _parse_modes(raw: str) -> list[str]:
|
||||
"""
|
||||
Parse and validate a comma-separated list of quantization modes.
|
||||
|
||||
Parameters:
|
||||
raw (str): Comma-separated mode names.
|
||||
|
||||
Returns:
|
||||
list[str]: Normalized, whitespace-trimmed mode names.
|
||||
|
||||
Raises:
|
||||
ValueError: If any mode is unsupported.
|
||||
"""
|
||||
modes = [mode.strip() for mode in raw.split(",") if mode.strip()]
|
||||
allowed = {"none", "int8", "int4", "mxfp8", "mxfp4", "nvfp4"}
|
||||
unknown = sorted(set(modes) - allowed)
|
||||
if unknown:
|
||||
raise ValueError(f"Unsupported modes: {unknown}")
|
||||
return modes
|
||||
|
||||
|
||||
def _latent_delta_metrics(candidate: np.ndarray, baseline: np.ndarray) -> dict[str, float]:
|
||||
"""
|
||||
Compare candidate and baseline latent arrays using error and signal-quality metrics.
|
||||
|
||||
Parameters:
|
||||
candidate (np.ndarray): Latent array to evaluate.
|
||||
baseline (np.ndarray): Reference latent array for comparison.
|
||||
|
||||
Returns:
|
||||
dict[str, float]: Mean squared error, mean absolute error, maximum absolute
|
||||
error, and signal-to-noise ratio in decibels between the arrays.
|
||||
"""
|
||||
diff = candidate.astype(np.float32) - baseline.astype(np.float32)
|
||||
mse = float(np.mean(np.square(diff)))
|
||||
mae = float(np.mean(np.abs(diff)))
|
||||
max_abs = float(np.max(np.abs(diff)))
|
||||
signal = float(np.mean(np.square(baseline.astype(np.float32))))
|
||||
return {
|
||||
"latent_mse_vs_fp16": mse,
|
||||
"latent_mae_vs_fp16": mae,
|
||||
"latent_max_abs_vs_fp16": max_abs,
|
||||
"latent_snr_db_vs_fp16": float(10.0 * np.log10(signal / mse)) if mse > 0 else float("inf"),
|
||||
}
|
||||
|
||||
|
||||
def _torch_mps_memory() -> dict[str, int | None]:
|
||||
"""
|
||||
Report PyTorch MPS memory statistics when PyTorch MPS is available.
|
||||
|
||||
Returns:
|
||||
dict[str, int | None]: A mapping of MPS memory metric names to byte counts, or `None` values when PyTorch or MPS is unavailable.
|
||||
"""
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
return {
|
||||
"torch_mps_current_allocated_bytes": None,
|
||||
"torch_mps_driver_allocated_bytes": None,
|
||||
"torch_mps_recommended_max_bytes": None,
|
||||
}
|
||||
if not torch.backends.mps.is_available():
|
||||
return {
|
||||
"torch_mps_current_allocated_bytes": None,
|
||||
"torch_mps_driver_allocated_bytes": None,
|
||||
"torch_mps_recommended_max_bytes": None,
|
||||
}
|
||||
return {
|
||||
"torch_mps_current_allocated_bytes": int(torch.mps.current_allocated_memory()),
|
||||
"torch_mps_driver_allocated_bytes": int(torch.mps.driver_allocated_memory()),
|
||||
"torch_mps_recommended_max_bytes": int(torch.mps.recommended_max_memory()),
|
||||
}
|
||||
|
||||
|
||||
def _decode_with_metrics(*, args, latents: np.ndarray, output_path: Path) -> dict[str, float | int | None | str]:
|
||||
"""
|
||||
Decode latents to a video and collect export timing and PyTorch MPS memory metrics.
|
||||
|
||||
Parameters:
|
||||
args: Configuration values for decoding and video export.
|
||||
latents (np.ndarray): Latent representation to decode.
|
||||
output_path (Path): Destination path for the exported video.
|
||||
|
||||
Returns:
|
||||
dict[str, float | int | None | str]: Video export duration and PyTorch MPS memory measurements.
|
||||
"""
|
||||
before = _torch_mps_memory()
|
||||
decode_start = time.perf_counter()
|
||||
decode_latents_to_video(
|
||||
model_root=args.model_root,
|
||||
latents_np=latents,
|
||||
output_path=output_path,
|
||||
fps=args.fps,
|
||||
device_arg=args.torch_device,
|
||||
dtype_arg=args.torch_dtype,
|
||||
backend=args.decode_backend,
|
||||
taehv_source_path=args.taehv_source_path,
|
||||
taehv_checkpoint_path=args.taehv_checkpoint_path,
|
||||
taehv_parallel=args.taehv_parallel,
|
||||
)
|
||||
decode_time = time.perf_counter() - decode_start
|
||||
after = _torch_mps_memory()
|
||||
return {
|
||||
"decode_export_s": decode_time,
|
||||
"decode_torch_mps_current_before_bytes": before["torch_mps_current_allocated_bytes"],
|
||||
"decode_torch_mps_current_after_bytes": after["torch_mps_current_allocated_bytes"],
|
||||
"decode_torch_mps_driver_before_bytes": before["torch_mps_driver_allocated_bytes"],
|
||||
"decode_torch_mps_driver_after_bytes": after["torch_mps_driver_allocated_bytes"],
|
||||
"decode_torch_mps_recommended_max_bytes": after["torch_mps_recommended_max_bytes"],
|
||||
}
|
||||
|
||||
|
||||
def _run_one_mode(
|
||||
*,
|
||||
mode: str,
|
||||
args,
|
||||
config: dict,
|
||||
checkpoint_path: Path,
|
||||
config_path: Path,
|
||||
prompt_embeds,
|
||||
freqs_cis,
|
||||
):
|
||||
"""
|
||||
Run denoising for one quantization mode and collect performance and memory metrics.
|
||||
|
||||
Parameters:
|
||||
mode (str): Quantization mode to benchmark.
|
||||
args: Benchmark configuration, including dtype, dimensions, seed, scheduler, and denoising settings.
|
||||
config (dict): Model configuration containing the input channel count.
|
||||
checkpoint_path (Path): Path to the transformer checkpoint.
|
||||
config_path (Path): Path to the transformer configuration.
|
||||
prompt_embeds: Encoded prompt embeddings shared across benchmark modes.
|
||||
freqs_cis: Rotary positional embeddings used during denoising.
|
||||
|
||||
Returns:
|
||||
dict: The mode name, generated latent array, and metrics for model loading,
|
||||
denoising, step timing, and MLX memory usage.
|
||||
"""
|
||||
import mlx.core as mx
|
||||
import torch
|
||||
|
||||
from fastvideo.benchmarks.mlx_fastwan_bench import denoise_dmd_on_device
|
||||
from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
|
||||
from fastvideo.mlx_runtime.fastwan import mlx_dit_from_diffusers_safetensors
|
||||
from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step
|
||||
|
||||
mx_dtype = mx.float16 if args.mlx_dtype == "fp16" else mx.float32
|
||||
quantization = None if mode == "none" else mode
|
||||
latent_frames = (args.num_frames - 1) // 4 + 1
|
||||
latent_height = args.height // 8
|
||||
latent_width = args.width // 8
|
||||
|
||||
load_start = time.perf_counter()
|
||||
mx.clear_cache()
|
||||
mx.reset_peak_memory()
|
||||
dit = mlx_dit_from_diffusers_safetensors(
|
||||
checkpoint_path,
|
||||
config_path,
|
||||
dtype=args.mlx_dtype,
|
||||
quantization=quantization,
|
||||
)
|
||||
load_time = time.perf_counter() - load_start
|
||||
load_peak_memory = mx.get_peak_memory()
|
||||
|
||||
scheduler = FlowMatchEulerDiscreteScheduler(shift=args.flow_shift)
|
||||
schedule = MLXDMDSchedule.from_torch_scheduler(scheduler)
|
||||
timesteps = [int(step.strip()) for step in args.dmd_denoising_steps.split(",") if step.strip()]
|
||||
# Same torch generator sequence as the original host-round-trip loop
|
||||
# (initial latents first, then one re-noise draw per intermediate step),
|
||||
# so every mode still shares identical stochasticity.
|
||||
generator = torch.Generator(device="cpu").manual_seed(args.seed)
|
||||
latents_seed = torch.randn(
|
||||
(1, int(config["in_channels"]), latent_frames, latent_height, latent_width),
|
||||
generator=generator,
|
||||
dtype=torch.float32,
|
||||
).numpy()
|
||||
renoise_by_step = [
|
||||
torch.randn(latents_seed.shape, generator=generator, dtype=torch.float32).numpy()
|
||||
for _ in range(max(0, len(timesteps) - 1))
|
||||
]
|
||||
latents = mx.array(latents_seed).astype(mx_dtype)
|
||||
encoder_hidden_states = mx.array(prompt_embeds.numpy()).astype(mx_dtype)
|
||||
|
||||
denoise_start = time.perf_counter()
|
||||
mx.reset_peak_memory()
|
||||
latents_np, step_times = denoise_dmd_on_device(
|
||||
mx=mx,
|
||||
dit=dit,
|
||||
latents=latents,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
freqs_cis=freqs_cis,
|
||||
timesteps=timesteps,
|
||||
renoise_by_step=renoise_by_step,
|
||||
schedule=schedule,
|
||||
dmd_step=dmd_step,
|
||||
mx_dtype=mx_dtype,
|
||||
)
|
||||
denoise_time = time.perf_counter() - denoise_start
|
||||
denoise_peak_memory = mx.get_peak_memory()
|
||||
active_memory = mx.get_active_memory()
|
||||
return {
|
||||
"mode": mode,
|
||||
"latents": latents_np,
|
||||
"metrics": {
|
||||
"mlx_dit_load_s": load_time,
|
||||
"mlx_denoise_s": denoise_time,
|
||||
"mlx_denoise_first_step_s": step_times[0] if step_times else None,
|
||||
"mlx_load_peak_bytes": int(load_peak_memory),
|
||||
"mlx_denoise_peak_bytes": int(denoise_peak_memory),
|
||||
"mlx_active_after_denoise_bytes": int(active_memory),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""
|
||||
Run the MLX FastWan quantization benchmark for the selected modes and write latency, memory, output, and latent-difference metrics to the output directory.
|
||||
"""
|
||||
parser = argparse.ArgumentParser(description="Benchmark MLX FastWan quantization modes.")
|
||||
parser.add_argument("--model-root", type=Path, default=DEFAULT_MODEL_ROOT)
|
||||
parser.add_argument("--prompt", default="A snow leopard walks across a windy mountain ridge.")
|
||||
parser.add_argument("--height", type=int, default=192)
|
||||
parser.add_argument("--width", type=int, default=320)
|
||||
parser.add_argument("--num-frames", type=int, default=17)
|
||||
parser.add_argument("--dmd-denoising-steps", default="1000,757,522")
|
||||
parser.add_argument("--flow-shift", type=float, default=8.0)
|
||||
parser.add_argument("--max-sequence-length", type=int, default=256)
|
||||
parser.add_argument("--seed", type=int, default=1024)
|
||||
parser.add_argument("--fps", type=int, default=16)
|
||||
parser.add_argument("--torch-device", default="auto")
|
||||
parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16")
|
||||
parser.add_argument("--mlx-dtype", choices=("fp16", "fp32"), default="fp16")
|
||||
parser.add_argument("--modes", default="none,int8,int4,mxfp8,mxfp4,nvfp4")
|
||||
parser.add_argument("--output-dir", type=Path, default=Path("video_samples/mlx_quant_benchmark"))
|
||||
parser.add_argument("--decode-backend", choices=("none", "wan-vae", "taehv"), default="taehv")
|
||||
parser.add_argument("--taehv-source-path", type=Path, default=None)
|
||||
parser.add_argument("--taehv-checkpoint-path", type=Path, default=None)
|
||||
parser.add_argument("--taehv-parallel", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
import mlx.core as mx
|
||||
import torch
|
||||
|
||||
mx.random.seed(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
config_path = args.model_root / "transformer/config.json"
|
||||
checkpoint_path = args.model_root / "transformer/diffusion_pytorch_model.safetensors"
|
||||
config = json.loads(config_path.read_text())
|
||||
latent_frames = (args.num_frames - 1) // 4 + 1
|
||||
latent_height = args.height // 8
|
||||
latent_width = args.width // 8
|
||||
|
||||
prompt_start = time.perf_counter()
|
||||
prompt_embeds = encode_prompt(
|
||||
model_root=args.model_root,
|
||||
prompt=args.prompt,
|
||||
max_sequence_length=args.max_sequence_length,
|
||||
device_arg=args.torch_device,
|
||||
dtype_arg=args.torch_dtype,
|
||||
)
|
||||
prompt_time = time.perf_counter() - prompt_start
|
||||
freqs_cis = make_rotary_embeddings(
|
||||
config,
|
||||
latent_frames=latent_frames,
|
||||
latent_height=latent_height,
|
||||
latent_width=latent_width,
|
||||
)
|
||||
|
||||
from fastvideo.mlx_runtime.fastwan import UnsupportedMLXQuantizationError
|
||||
|
||||
baseline_latents = None
|
||||
rows = []
|
||||
for mode in _parse_modes(args.modes):
|
||||
print(f"=== MLX quant mode: {mode} ===")
|
||||
mode_start = time.perf_counter()
|
||||
try:
|
||||
result = _run_one_mode(
|
||||
mode=mode,
|
||||
args=args,
|
||||
config=config,
|
||||
checkpoint_path=checkpoint_path,
|
||||
config_path=config_path,
|
||||
prompt_embeds=prompt_embeds,
|
||||
freqs_cis=freqs_cis,
|
||||
)
|
||||
except UnsupportedMLXQuantizationError as exc:
|
||||
print(f"skipping mode (unsupported by this MLX build): {exc}")
|
||||
rows.append({"mode": mode, "status": "unsupported_by_mlx", "error": str(exc)})
|
||||
continue
|
||||
cleanup_mlx(mx)
|
||||
latents = result["latents"]
|
||||
if baseline_latents is None:
|
||||
baseline_latents = latents
|
||||
latent_path = args.output_dir / f"latents_{mode}.npy"
|
||||
np.save(latent_path, latents)
|
||||
|
||||
decode_time = 0.0
|
||||
decode_metrics = {}
|
||||
output_path = None
|
||||
if args.decode_backend != "none":
|
||||
output_path = args.output_dir / f"video_{mode}_{args.decode_backend}_{args.height}x{args.width}x{args.num_frames}.mp4"
|
||||
decode_metrics = _decode_with_metrics(args=args, latents=latents, output_path=output_path)
|
||||
decode_time = cast(float, decode_metrics["decode_export_s"])
|
||||
|
||||
mode_total = time.perf_counter() - mode_start
|
||||
mlx_denoise_peak_bytes = int(result["metrics"]["mlx_denoise_peak_bytes"])
|
||||
mlx_active_bytes = int(result["metrics"]["mlx_active_after_denoise_bytes"])
|
||||
metrics = {
|
||||
"mode": mode,
|
||||
"status": "ok",
|
||||
"prompt_encode_shared_s": prompt_time,
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_frames": args.num_frames,
|
||||
"decode_backend": args.decode_backend,
|
||||
"decode_export_s": decode_time,
|
||||
"mode_total_excluding_shared_prompt_s": mode_total,
|
||||
"mode_total_including_shared_prompt_s": mode_total + prompt_time,
|
||||
"latents_path": str(latent_path),
|
||||
"output_path": str(output_path) if output_path else None,
|
||||
"mlx_denoise_peak_gib": mlx_denoise_peak_bytes / (1024**3),
|
||||
"mlx_active_after_denoise_gib": mlx_active_bytes / (1024**3),
|
||||
"mlx_dit_peak_under_16gb": mlx_denoise_peak_bytes < 16 * 1024**3,
|
||||
"mlx_dit_active_under_16gb": mlx_active_bytes < 16 * 1024**3,
|
||||
"mac_16gb_status": (
|
||||
"dit_memory_fits_16gb_measured_decode_separately"
|
||||
if mlx_denoise_peak_bytes < 16 * 1024**3 else "dit_memory_exceeds_16gb"
|
||||
),
|
||||
**result["metrics"],
|
||||
**decode_metrics,
|
||||
**_latent_delta_metrics(latents, baseline_latents),
|
||||
}
|
||||
rows.append(metrics)
|
||||
print(json.dumps(metrics, indent=2))
|
||||
|
||||
metrics_path = args.output_dir / "metrics.json"
|
||||
metrics_path.write_text(json.dumps(rows, indent=2))
|
||||
print(f"Wrote benchmark metrics to: {metrics_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,104 @@
|
||||
"""Compare generated MP4s against a reference MP4 with simple pixel metrics."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def _read_video(path: Path) -> np.ndarray:
|
||||
"""
|
||||
Read all frames from a video file as an RGB NumPy array.
|
||||
|
||||
Parameters:
|
||||
path (Path): Path to the video file.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Video frames stacked along the first axis.
|
||||
|
||||
Raises:
|
||||
ValueError: If the video contains no readable frames.
|
||||
"""
|
||||
import cv2
|
||||
|
||||
cap = cv2.VideoCapture(str(path))
|
||||
frames = []
|
||||
try:
|
||||
while True:
|
||||
ok, frame_bgr = cap.read()
|
||||
if not ok:
|
||||
break
|
||||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||||
frames.append(frame_rgb)
|
||||
finally:
|
||||
cap.release()
|
||||
if not frames:
|
||||
raise ValueError(f"No frames read from {path}")
|
||||
return np.stack(frames, axis=0)
|
||||
|
||||
|
||||
def _metrics(candidate: np.ndarray, reference: np.ndarray) -> dict[str, float | int | list[int]]:
|
||||
"""
|
||||
Compute pixel-level comparison metrics between candidate and reference video frames.
|
||||
|
||||
Parameters:
|
||||
candidate (np.ndarray): Candidate video frames in frame, height, width, and channel order.
|
||||
reference (np.ndarray): Reference video frames with the same shape as the candidate.
|
||||
|
||||
Returns:
|
||||
dict[str, float | int | list[int]]: Frame dimensions and pixel comparison metrics, including MSE, MAE, maximum absolute difference, and PSNR in decibels.
|
||||
|
||||
Raises:
|
||||
ValueError: If the candidate and reference arrays have different shapes.
|
||||
"""
|
||||
if candidate.shape != reference.shape:
|
||||
raise ValueError(f"Shape mismatch: candidate={candidate.shape}, reference={reference.shape}")
|
||||
candidate_f = candidate.astype(np.float32)
|
||||
reference_f = reference.astype(np.float32)
|
||||
diff = candidate_f - reference_f
|
||||
mse = float(np.mean(np.square(diff)))
|
||||
mae = float(np.mean(np.abs(diff)))
|
||||
max_abs = float(np.max(np.abs(diff)))
|
||||
psnr = float(20.0 * np.log10(255.0 / np.sqrt(mse))) if mse > 0 else float("inf")
|
||||
return {
|
||||
"frames": int(candidate.shape[0]),
|
||||
"height": int(candidate.shape[1]),
|
||||
"width": int(candidate.shape[2]),
|
||||
"channels": int(candidate.shape[3]),
|
||||
"mse_vs_reference": mse,
|
||||
"mae_vs_reference": mae,
|
||||
"max_abs_vs_reference": max_abs,
|
||||
"psnr_db_vs_reference": psnr,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Compare candidate MP4 videos with a reference and write pixel-level metrics to a JSON file."""
|
||||
parser = argparse.ArgumentParser(description="Compare MP4s against a reference MP4.")
|
||||
parser.add_argument("--reference", type=Path, required=True)
|
||||
parser.add_argument("--candidates", type=Path, nargs="+", required=True)
|
||||
parser.add_argument("--metrics-json", type=Path, required=True)
|
||||
args = parser.parse_args()
|
||||
|
||||
reference = _read_video(args.reference)
|
||||
rows = []
|
||||
for candidate_path in args.candidates:
|
||||
candidate = _read_video(candidate_path)
|
||||
row = {
|
||||
"reference_path": str(args.reference),
|
||||
"candidate_path": str(candidate_path),
|
||||
**_metrics(candidate, reference),
|
||||
}
|
||||
rows.append(row)
|
||||
print(json.dumps(row, indent=2))
|
||||
|
||||
args.metrics_json.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.metrics_json.write_text(json.dumps(rows, indent=2))
|
||||
print(f"Wrote video quality metrics to: {args.metrics_json}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user