chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,346 @@
|
||||
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License").
|
||||
# You may not use this file except in compliance with the License.
|
||||
# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""
|
||||
Decode saved latent files (.pt) to pixel video using a local LightVAE loader.
|
||||
|
||||
This mirrors `scripts/decode_vae_latents.py`, but only depends on the current
|
||||
repo's 5B VAE code plus a LightVAE checkpoint such as `MG-LightVAE_v2.pth`.
|
||||
|
||||
Examples:
|
||||
python scripts/decode_lightvae_latents.py \
|
||||
--input_dir /path/to/latents \
|
||||
--vae_path /path/to/MG-LightVAE_v2.pth
|
||||
|
||||
torchrun --nproc_per_node=8 scripts/decode_lightvae_latents.py \
|
||||
--input_dir /path/to/latents \
|
||||
--ckpt_dir /path/to/lightvae_ckpts \
|
||||
--vae_type mg_lightvae
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torchvision.io import write_video
|
||||
from tqdm import tqdm
|
||||
|
||||
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
if REPO_ROOT not in sys.path:
|
||||
sys.path.insert(0, REPO_ROOT)
|
||||
|
||||
|
||||
from utils.lightvae_5b_wrapper import LightVAE5BWrapper
|
||||
|
||||
|
||||
def init_distributed():
|
||||
"""Initialize distributed process group if launched via torchrun."""
|
||||
if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
|
||||
rank = int(os.environ["RANK"])
|
||||
world_size = int(os.environ["WORLD_SIZE"])
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", 0))
|
||||
torch.cuda.set_device(local_rank)
|
||||
dist.init_process_group(backend="nccl")
|
||||
return rank, world_size, local_rank
|
||||
return 0, 1, 0
|
||||
|
||||
|
||||
def decode_latent_to_video(
|
||||
vae, latent: torch.Tensor, device: torch.device, dtype: torch.dtype
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Decode latent to pixel video.
|
||||
|
||||
Args:
|
||||
vae: VAE wrapper exposing `decode_to_pixel(latent)`
|
||||
latent: shape (batch, T, C, H, W) or (T, C, H, W)
|
||||
device, dtype: device and dtype for computation
|
||||
|
||||
Returns:
|
||||
video: shape (batch, T, C, H, W), range [0, 1]
|
||||
"""
|
||||
if latent.dim() == 4:
|
||||
latent = latent.unsqueeze(0)
|
||||
latent = latent.to(device=device, dtype=dtype)
|
||||
video = vae.decode_to_pixel(latent)
|
||||
video = (video * 0.5 + 0.5).clamp(0, 1)
|
||||
return video
|
||||
|
||||
|
||||
def _normalize_requested_vae_type(value: str) -> str:
|
||||
normalized = value.strip().lower()
|
||||
if normalized == "wan":
|
||||
return "wan2.2"
|
||||
if normalized in {"wan2.2", "mg_lightvae", "mg_lightvae_v2"}:
|
||||
return normalized
|
||||
raise ValueError(
|
||||
f"Unsupported --vae_type '{value}'. "
|
||||
"Expected one of: wan2.2, wan, mg_lightvae, mg_lightvae_v2."
|
||||
)
|
||||
|
||||
|
||||
def _parse_lightvae_pruning_rate(value: Optional[str]) -> Optional[float]:
|
||||
if value is None:
|
||||
return None
|
||||
lowered = str(value).strip().lower()
|
||||
if lowered in {"", "auto", "none"}:
|
||||
return None
|
||||
return float(lowered)
|
||||
|
||||
|
||||
def _resolve_vae_paths(
|
||||
*,
|
||||
ckpt_dir: Optional[str],
|
||||
vae_path: Optional[str],
|
||||
requested_vae_type: str,
|
||||
lightvae_pruning_rate: Optional[str],
|
||||
):
|
||||
pruning_rate = _parse_lightvae_pruning_rate(lightvae_pruning_rate)
|
||||
ckpt_dir = os.path.abspath(ckpt_dir) if ckpt_dir else None
|
||||
|
||||
if vae_path is not None:
|
||||
resolved_vae_path = os.path.abspath(vae_path)
|
||||
if requested_vae_type == "wan2.2" and pruning_rate is None:
|
||||
pruning_rate = 0.0
|
||||
else:
|
||||
if ckpt_dir is None:
|
||||
raise ValueError("Either --ckpt_dir or --vae_path must be provided.")
|
||||
if requested_vae_type == "mg_lightvae":
|
||||
resolved_vae_path = os.path.join(ckpt_dir, "MG-LightVAE.pth")
|
||||
if pruning_rate is None:
|
||||
pruning_rate = 0.5
|
||||
elif requested_vae_type == "mg_lightvae_v2":
|
||||
resolved_vae_path = os.path.join(ckpt_dir, "MG-LightVAE_v2.pth")
|
||||
if pruning_rate is None:
|
||||
pruning_rate = 0.75
|
||||
else:
|
||||
resolved_vae_path = os.path.join(ckpt_dir, "Wan2.2_VAE.pth")
|
||||
if pruning_rate is None:
|
||||
pruning_rate = 0.0
|
||||
|
||||
return resolved_vae_path, pruning_rate
|
||||
|
||||
|
||||
def _load_latent_tensor(pt_path: str):
|
||||
try:
|
||||
data = torch.load(pt_path, map_location="cpu", weights_only=True)
|
||||
except TypeError:
|
||||
data = torch.load(pt_path, map_location="cpu")
|
||||
|
||||
if isinstance(data, torch.Tensor):
|
||||
return data
|
||||
if isinstance(data, dict):
|
||||
latent = next(iter(data.values()))
|
||||
if isinstance(latent, torch.Tensor):
|
||||
return latent
|
||||
raise TypeError(f"First dict value is not a tensor: {type(latent)}")
|
||||
raise TypeError(f"Unsupported latent file payload type: {type(data)}")
|
||||
|
||||
|
||||
def _parse_dtype(dtype_name: str) -> torch.dtype:
|
||||
mapping = {
|
||||
"float32": torch.float32,
|
||||
"float16": torch.float16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
}
|
||||
try:
|
||||
return mapping[dtype_name]
|
||||
except KeyError as exc:
|
||||
raise ValueError(
|
||||
f"Unsupported --dtype '{dtype_name}'. "
|
||||
"Expected one of: float32, float16, bfloat16."
|
||||
) from exc
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Decode saved latents to video with a local LightVAE loader."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_dir",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Directory containing .pt latent files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory to save decoded .mp4 videos. Default: input_dir/decoded_lightvae_videos",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--fps",
|
||||
type=int,
|
||||
default=24,
|
||||
help="FPS for output video (default: 24).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pattern",
|
||||
type=str,
|
||||
default="latents_*.pt",
|
||||
help="Glob pattern for latent files (default: latents_*.pt).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ckpt_dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory containing `MG-LightVAE*.pth` or `Wan2.2_VAE.pth`.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vae_path",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Explicit LightVAE/Wan2.2 checkpoint path. "
|
||||
"Use this when you only have a single VAE checkpoint file."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--matrix_game_root",
|
||||
type=str,
|
||||
default=None,
|
||||
help=argparse.SUPPRESS,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--teacher_vae_path",
|
||||
"--lightvae_encoder_path",
|
||||
dest="teacher_vae_path",
|
||||
type=str,
|
||||
default=None,
|
||||
help=argparse.SUPPRESS,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vae_type",
|
||||
type=str,
|
||||
default="mg_lightvae_v2",
|
||||
choices=["wan2.2", "wan", "mg_lightvae", "mg_lightvae_v2"],
|
||||
help=(
|
||||
"VAE variant to use. "
|
||||
"`mg_lightvae` maps to MG-LightVAE.pth, "
|
||||
"`mg_lightvae_v2` maps to MG-LightVAE_v2.pth."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lightvae_pruning_rate",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Override pruning rate. Use `auto`/`none` to let Wan2_2_VAE infer it "
|
||||
"when an explicit LightVAE checkpoint is provided."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
type=str,
|
||||
default="bfloat16",
|
||||
choices=["float32", "float16", "bfloat16"],
|
||||
help="Decode dtype (default: bfloat16).",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
rank, world_size, local_rank = init_distributed()
|
||||
is_main = rank == 0
|
||||
|
||||
output_dir = args.output_dir or os.path.join(
|
||||
args.input_dir, "decoded_lightvae_videos"
|
||||
)
|
||||
if is_main:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
if world_size > 1:
|
||||
dist.barrier()
|
||||
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
device = torch.device(
|
||||
f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu"
|
||||
)
|
||||
dtype = _parse_dtype(args.dtype)
|
||||
|
||||
requested_vae_type = _normalize_requested_vae_type(args.vae_type)
|
||||
resolved_vae_path, resolved_pruning_rate = _resolve_vae_paths(
|
||||
ckpt_dir=args.ckpt_dir,
|
||||
vae_path=args.vae_path,
|
||||
requested_vae_type=requested_vae_type,
|
||||
lightvae_pruning_rate=args.lightvae_pruning_rate,
|
||||
)
|
||||
|
||||
vae = LightVAE5BWrapper(
|
||||
vae_path=resolved_vae_path,
|
||||
pruning_rate=resolved_pruning_rate,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
).eval()
|
||||
|
||||
if is_main:
|
||||
print(
|
||||
"Using local VAE-only loader: "
|
||||
f"requested={requested_vae_type}, "
|
||||
f"pruning={vae.pruning_rate}, "
|
||||
f"vae_path={vae.vae_path}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
search_path = os.path.join(args.input_dir, args.pattern)
|
||||
pt_files = sorted(glob.glob(search_path))
|
||||
if not pt_files:
|
||||
if is_main:
|
||||
print(f"No files matching '{search_path}' found. Exiting.")
|
||||
return
|
||||
|
||||
pt_files_local = pt_files[rank::world_size]
|
||||
if is_main:
|
||||
print(
|
||||
f"Found {len(pt_files)} latent file(s), {world_size} GPU(s), "
|
||||
f"~{len(pt_files_local)} per GPU.",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
pbar = tqdm(pt_files_local, desc=f"[Rank {rank}] Decoding", disable=(not is_main))
|
||||
for pt_path in pbar:
|
||||
basename = os.path.splitext(os.path.basename(pt_path))[0]
|
||||
if basename.startswith("latents_"):
|
||||
video_name = basename.replace("latents_", "video_", 1)
|
||||
else:
|
||||
video_name = f"video_{basename}"
|
||||
out_path = os.path.join(output_dir, f"{video_name}.mp4")
|
||||
|
||||
try:
|
||||
latent = _load_latent_tensor(pt_path)
|
||||
except Exception as exc:
|
||||
print(f"[Rank {rank}] Failed to load {pt_path}: {exc}", flush=True)
|
||||
continue
|
||||
|
||||
try:
|
||||
video = decode_latent_to_video(vae, latent, device, dtype)
|
||||
video_frames = video[0].cpu()
|
||||
video_uint8 = (video_frames * 255.0).clamp(0, 255).to(torch.uint8)
|
||||
video_uint8 = video_uint8.permute(0, 2, 3, 1)
|
||||
write_video(out_path, video_uint8, fps=args.fps)
|
||||
except Exception as exc:
|
||||
print(f"[Rank {rank}] Failed to decode {pt_path}: {exc}", flush=True)
|
||||
|
||||
if world_size > 1:
|
||||
dist.barrier()
|
||||
if is_main:
|
||||
print(f"Done. Decoded {len(pt_files)} latent(s) to {output_dir}", flush=True)
|
||||
if world_size > 1:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user