534 lines
21 KiB
Python
534 lines
21 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""End-to-end FastMetal-5B-QAD generation on Apple Silicon (MLX DiT + MLX TAEHV).
|
|
|
|
This is the Wan2.2 TI2V entrypoint. Use FastVideo/FastMetal-5B-QAD:
|
|
|
|
hf download FastVideo/FastMetal-5B-QAD --local-dir ./FastMetal-5B-QAD
|
|
python examples/inference/basic/mlx_wan22_generate.py \\
|
|
--mlx-checkpoint ./FastMetal-5B-QAD \\
|
|
--text-encoder-root ./FastMetal-5B-QAD \\
|
|
--vae-root ./FastMetal-5B-QAD/vae
|
|
|
|
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.
|
|
|
|
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.checkpoint_compat import (
|
|
UnsupportedMLXCheckpointError,
|
|
raise_if_unsupported_mlx_checkpoint,
|
|
resolve_mlx_checkpoint,
|
|
)
|
|
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="Packed FastMetal-5B-QAD MLX DiT directory (mlx_dit.json + mlx_dit.safetensors). "
|
|
"If omitted, a FastMetal directory passed as --text-encoder-root is used when it "
|
|
"already contains those files.",
|
|
)
|
|
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.mlx_checkpoint = resolve_mlx_checkpoint(args.mlx_checkpoint, args.text_encoder_root)
|
|
if args.mlx_checkpoint is not None:
|
|
if args.text_encoder_root is None and (args.mlx_checkpoint / "text_encoder").is_dir():
|
|
args.text_encoder_root = args.mlx_checkpoint
|
|
if args.vae_root is None and (args.mlx_checkpoint / "vae").is_dir():
|
|
args.vae_root = args.mlx_checkpoint / "vae"
|
|
try:
|
|
raise_if_unsupported_mlx_checkpoint(args.mlx_checkpoint, args.dit_checkpoint)
|
|
except UnsupportedMLXCheckpointError as exc:
|
|
raise SystemExit(str(exc)) from exc
|
|
|
|
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()
|