chore: import upstream snapshot with attribution

scripts/compute_sp_vae_chunk_halo.py

@@ -0,0 +1,163 @@
+# 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
+
+"""Compute SP chunk-halo VAE frame windows for a training config."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+from omegaconf import OmegaConf
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from utils.config import normalize_config, wan_default_config
+
+DEFAULT_SP_VAE_HALO_LATENTS = 28
+
+
+def latent_range_to_raw_window(
+    latent_start: int,
+    latent_end: int,
+    *,
+    temporal_compression_ratio: int,
+) -> tuple[int, int, int]:
+    """Map a latent range to the raw-frame window needed by Wan VAE encode."""
+    if latent_end <= latent_start:
+        raise ValueError(f"latent_end must be > latent_start, got {latent_start}, {latent_end}")
+    ratio = int(temporal_compression_ratio)
+    if latent_start == 0:
+        return 0, 1 + ratio * (latent_end - 1), 0
+    return ratio * (latent_start - 1), 1 + ratio * (latent_end - 1), 1
+
+
+def compute_chunk_halo_metas(
+    *,
+    total_latent_frames: int,
+    total_raw_frames: int,
+    sp_size: int,
+    halo_latents: int,
+    temporal_compression_ratio: int,
+) -> list[dict[str, int]]:
+    if sp_size <= 0:
+        raise ValueError("sp_size must be positive")
+    if halo_latents < 0:
+        raise ValueError("halo_latents must be non-negative")
+    if total_latent_frames % sp_size != 0:
+        raise ValueError(
+            f"total_latent_frames={total_latent_frames} must be divisible by sp_size={sp_size}"
+        )
+
+    local_latent_frames = total_latent_frames // sp_size
+    metas = []
+    for sp_rank in range(sp_size):
+        keep_start = sp_rank * local_latent_frames
+        keep_end = keep_start + local_latent_frames
+        halo_start = max(0, keep_start - halo_latents)
+        raw_start, raw_end, pseudo_prefix_latents = latent_range_to_raw_window(
+            halo_start,
+            keep_end,
+            temporal_compression_ratio=temporal_compression_ratio,
+        )
+        raw_start = max(0, raw_start)
+        raw_end = min(total_raw_frames, raw_end)
+        drop_latents = pseudo_prefix_latents + (keep_start - halo_start)
+        metas.append(
+            {
+                "sp_rank": sp_rank,
+                "keep_start": keep_start,
+                "keep_end": keep_end,
+                "halo_start": halo_start,
+                "raw_start": raw_start,
+                "raw_end": raw_end,
+                "raw_frames": raw_end - raw_start,
+                "drop_latents": drop_latents,
+                "local_latent_frames": local_latent_frames,
+            }
+        )
+    return metas
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default="configs/train_ar.yaml", help="Training config path")
+    parser.add_argument("--sp-size", type=int, default=None, help="Override sequence_parallel_size")
+    parser.add_argument("--halo-latents", type=int, default=None, help="Override vae_halo_latents")
+    parser.add_argument("--latent-frames", type=int, default=None, help="Override image_or_video_shape[1]")
+    parser.add_argument("--raw-frames", type=int, default=None, help="Override total raw frame count")
+    parser.add_argument("--json", action="store_true", help="Print JSON instead of a table")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    cfg = normalize_config(OmegaConf.load(args.config))
+    model_name = cfg.model_kwargs.model_name
+    temporal_ratio = int(wan_default_config[model_name]["temporal_compression_ratio"])
+    total_latent_frames = int(
+        args.latent_frames if args.latent_frames is not None else cfg.image_or_video_shape[1]
+    )
+    total_raw_frames = int(
+        args.raw_frames
+        if args.raw_frames is not None
+        else ((total_latent_frames - 1) * temporal_ratio + 1)
+    )
+    sp_size = int(
+        args.sp_size if args.sp_size is not None else getattr(cfg, "sequence_parallel_size", 1)
+    )
+    halo_latents = int(
+        args.halo_latents
+        if args.halo_latents is not None
+        else getattr(cfg, "vae_halo_latents", DEFAULT_SP_VAE_HALO_LATENTS)
+    )
+
+    metas = compute_chunk_halo_metas(
+        total_latent_frames=total_latent_frames,
+        total_raw_frames=total_raw_frames,
+        sp_size=sp_size,
+        halo_latents=halo_latents,
+        temporal_compression_ratio=temporal_ratio,
+    )
+    payload = {
+        "config": str(Path(args.config)),
+        "model_name": model_name,
+        "temporal_compression_ratio": temporal_ratio,
+        "total_latent_frames": total_latent_frames,
+        "total_raw_frames": total_raw_frames,
+        "sp_size": sp_size,
+        "halo_latents": halo_latents,
+        "rank_metas": metas,
+    }
+    if args.json:
+        print(json.dumps(payload, indent=2))
+        return
+
+    print(
+        f"config={payload['config']} model={model_name} "
+        f"latent_frames={total_latent_frames} raw_frames={total_raw_frames} "
+        f"sp_size={sp_size} halo_latents={halo_latents}"
+    )
+    print("rank  keep_latents  halo_start  raw_window  raw_frames  drop_latents")
+    for meta in metas:
+        keep_latents = f"[{meta['keep_start']},{meta['keep_end']})".ljust(12)
+        raw_window = f"[{meta['raw_start']},{meta['raw_end']})".ljust(11)
+        print(
+            f"{meta['sp_rank']:>4}  {keep_latents}  {meta['halo_start']:>10}  "
+            f"{raw_window}  {meta['raw_frames']:>10}  {meta['drop_latents']:>12}"
+        )
+
+
+if __name__ == "__main__":
+    main()

scripts/decode_lightvae_latents.py

@@ -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()

scripts/decode_vae_latents.py

@@ -0,0 +1,167 @@
+# 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 and write .mp4 using the specified VAE.
+
+Supports multi-GPU via torchrun:
+    torchrun --nproc_per_node=8 scripts/decode_vae_latents.py --input_dir /path/to/latents
+Single-GPU also works:
+    python scripts/decode_vae_latents.py --input_dir /path/to/latents
+
+Uses the Wan2.2-TI2V-5B VAE.
+"""
+import sys
+import os
+import argparse
+import glob
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import torch
+import torch.distributed as dist
+from torchvision.io import write_video
+from tqdm import tqdm
+
+
+def init_distributed():
+    """Initialize distributed process group if launched via torchrun. Returns (rank, world_size)."""
+    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 get_vae():
+    """Return the Wan2.2-TI2V-5B VAE wrapper."""
+    from utils.wan_5b_wrapper import WanVAEWrapper
+    return WanVAEWrapper()
+
+
+def decode_latent_to_video(vae, latent: torch.Tensor, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    """
+    Decode latent to pixel video; same logic as pipeline/causal_diffusion_inference.py.
+
+    Args:
+        vae: Wan2.2-TI2V-5B VAE wrapper
+        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 main():
+    parser = argparse.ArgumentParser(description="Decode saved latents to video.")
+    parser.add_argument(
+        "--input_dir",
+        type=str,
+        required=True,
+        help="Directory containing .pt latent files (e.g. latents_rank00_idx000000.pt).",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="Directory to save decoded .mp4 videos. Default: input_dir/decoded_videos",
+    )
+    parser.add_argument(
+        "--fps",
+        type=int,
+        default=24,
+        help="FPS for output video (default: 16).",
+    )
+    parser.add_argument(
+        "--pattern",
+        type=str,
+        default="latents_*.pt",
+        help="Glob pattern for latent files (default: latents_*.pt).",
+    )
+    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_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 = torch.bfloat16
+
+    vae = get_vae()
+    vae = vae.to(device=device, dtype=dtype).eval()
+
+    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
+
+    # Shard files across ranks
+    pt_files_local = pt_files[rank::world_size]
+    if is_main:
+        print(f"Found {len(pt_files)} latent file(s), {world_size} GPU(s), ~{len(pt_files_local)} per GPU.")
+
+    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]
+        video_name = basename.replace("latents_", "video_", 1) if basename.startswith("latents_") else f"video_{basename}"
+        out_path = os.path.join(output_dir, f"{video_name}.mp4")
+
+        try:
+            data = torch.load(pt_path, map_location="cpu", weights_only=True)
+            if isinstance(data, torch.Tensor):
+                latent = data
+            elif isinstance(data, dict):
+                latent = next(iter(data.values()))
+                if not isinstance(latent, torch.Tensor):
+                    print(f"[Rank {rank}] Skip {pt_path}: dict value is not a tensor.")
+                    continue
+            else:
+                print(f"[Rank {rank}] Skip {pt_path}: unsupported type {type(data)}.")
+                continue
+        except Exception as e:
+            print(f"[Rank {rank}] Failed to load {pt_path}: {e}")
+            continue
+
+        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)  # (T, C, H, W) -> (T, H, W, C)
+        write_video(out_path, video_uint8, fps=args.fps)
+
+    if world_size > 1:
+        dist.barrier()
+    if is_main:
+        print(f"Done. Decoded {len(pt_files)} latent(s) to {output_dir}")
+    if world_size > 1:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

scripts/merge_lora_generator.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+# 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
+"""Merge a LongLive generator checkpoint with LoRA weights for simple inference."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+
+def _torch_load(path: str):
+    import torch
+
+    try:
+        return torch.load(path, map_location="cpu", weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location="cpu")
+
+
+def _load_lora_state(path: str):
+    checkpoint = _torch_load(path)
+    if isinstance(checkpoint, dict) and "generator_lora" in checkpoint:
+        return checkpoint["generator_lora"]
+    return checkpoint
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config_path", required=True, help="Inference yaml containing model, checkpoint, and adapter settings.")
+    parser.add_argument("--output_path", required=True, help="Path to save the merged generator checkpoint.")
+    parser.add_argument("--generator_ckpt", default=None, help="Override checkpoints.generator_ckpt from the yaml.")
+    parser.add_argument("--lora_ckpt", default=None, help="Override checkpoints.lora_ckpt from the yaml.")
+    parser.add_argument("--device", default="cuda:0", help="Device used for merging, e.g. cuda:0 or cpu.")
+    parser.add_argument("--dtype", choices=("bf16", "fp32"), default="bf16", help="Save merged weights in this dtype.")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    import torch
+    from omegaconf import OmegaConf
+
+    from utils.config import normalize_config
+    from utils.inference_utils import load_generator_checkpoint
+    from utils.lora_utils import configure_lora_for_model
+    from utils.nvfp4_checkpoint import cpu_state_dict
+    from utils.wan_5b_wrapper import WanDiffusionWrapper
+
+    config = normalize_config(OmegaConf.load(args.config_path))
+    generator_ckpt = args.generator_ckpt or getattr(config, "generator_ckpt", None)
+    lora_ckpt = args.lora_ckpt or getattr(config, "lora_ckpt", None)
+    if not generator_ckpt:
+        raise ValueError("Missing generator checkpoint. Set checkpoints.generator_ckpt or pass --generator_ckpt.")
+    if not lora_ckpt:
+        raise ValueError("Missing LoRA checkpoint. Set checkpoints.lora_ckpt or pass --lora_ckpt.")
+    if not getattr(config, "adapter", None):
+        raise ValueError("Missing adapter config. The merge script needs the LoRA rank/alpha/dropout settings.")
+
+    device = torch.device(args.device)
+    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32
+
+    print(f"Building generator: {config.model_kwargs}")
+    generator = WanDiffusionWrapper(**getattr(config, "model_kwargs", {}), is_causal=True)
+    generator.eval().requires_grad_(False)
+
+    print(f"Loading generator checkpoint: {generator_ckpt}")
+    incompatible = load_generator_checkpoint(
+        generator,
+        generator_ckpt,
+        use_ema=bool(getattr(config, "use_ema", False)),
+    )
+    missing = getattr(incompatible, "missing_keys", [])
+    unexpected = getattr(incompatible, "unexpected_keys", [])
+    if missing:
+        print(f"[Warning] Missing generator keys: {missing[:8]} ...")
+    if unexpected:
+        print(f"[Warning] Unexpected generator keys: {unexpected[:8]} ...")
+
+    print(f"Applying LoRA config: {config.adapter}")
+    generator.model = configure_lora_for_model(
+        generator.model,
+        model_name="generator",
+        lora_config=config.adapter,
+        is_main_process=True,
+    )
+
+    import peft
+
+    print(f"Loading LoRA checkpoint: {lora_ckpt}")
+    peft.set_peft_model_state_dict(generator.model, _load_lora_state(lora_ckpt))  # type: ignore[arg-type]
+
+    print(f"Merging LoRA on {device} in {dtype}...")
+    generator.to(device=device, dtype=dtype)
+    generator.model = generator.model.merge_and_unload(safe_merge=True)
+    generator.eval().requires_grad_(False)
+
+    output_path = Path(args.output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    checkpoint = {
+        "generator": cpu_state_dict(generator),
+        "checkpoint_format": "longlive_generator_merged_lora",
+        "source_generator_ckpt": str(generator_ckpt),
+        "source_lora_ckpt": str(lora_ckpt),
+        "model_name": getattr(config.model_kwargs, "model_name", None),
+        "dtype": str(dtype).replace("torch.", ""),
+        "merged_lora": True,
+    }
+    torch.save(checkpoint, output_path)
+    size_gib = os.path.getsize(output_path) / (1024 ** 3)
+    print(f"Saved merged generator to {output_path} ({size_gib:.2f} GiB).")
+    print("Use this file as checkpoints.generator_ckpt for inference and remove adapter/lora_ckpt from the inference config.")
+
+
+if __name__ == "__main__":
+    main()

scripts/save_merged_nvfp4_generator.py

@@ -0,0 +1,246 @@
+#!/usr/bin/env python3
+# 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
+"""Save a merged LoRA generator as a reusable checkpoint."""
+from __future__ import annotations
+
+import argparse
+import gc
+import os
+import sys
+from pathlib import Path
+
+import torch
+from omegaconf import OmegaConf
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from utils.config import normalize_config
+from utils.lora_utils import configure_lora_for_model
+from utils.nvfp4_checkpoint import (
+    NVFP4_CHECKPOINT_FORMAT,
+    NVFP4_CHECKPOINT_VERSION,
+    clean_fsdp_state_dict_keys,
+    cpu_state_dict,
+    is_nvfp4_state_dict,
+    quantize_model_for_fouroversix_nvfp4,
+    unwrap_generator_state_dict,
+)
+from utils.quant import _materialize_quantized_weights_for_inference
+from utils.wan_5b_wrapper import WanDiffusionWrapper
+
+
+def _torch_load(path: str):
+    try:
+        return torch.load(path, map_location="cpu", weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location="cpu")
+
+
+def _load_generator_checkpoint(generator: WanDiffusionWrapper, checkpoint_path: str, use_ema: bool) -> None:
+    checkpoint = _torch_load(checkpoint_path)
+    state_dict = unwrap_generator_state_dict(checkpoint, use_ema=use_ema)
+    if is_nvfp4_state_dict(state_dict):
+        raise ValueError(
+            f"{checkpoint_path} is already a materialized NVFP4 checkpoint; "
+            "use it directly as checkpoints.generator_ckpt."
+        )
+    if use_ema:
+        state_dict = clean_fsdp_state_dict_keys(state_dict)
+    incompatible = generator.load_state_dict(state_dict, strict=not use_ema)
+    missing = getattr(incompatible, "missing_keys", [])
+    unexpected = getattr(incompatible, "unexpected_keys", [])
+    if missing:
+        print(f"[Warning] Missing generator keys while loading base checkpoint: {missing[:8]} ...")
+    if unexpected:
+        print(f"[Warning] Unexpected generator keys while loading base checkpoint: {unexpected[:8]} ...")
+
+
+def _load_lora_state(lora_ckpt_path: str):
+    checkpoint = _torch_load(lora_ckpt_path)
+    if isinstance(checkpoint, dict) and "generator_lora" in checkpoint:
+        return checkpoint["generator_lora"]
+    return checkpoint
+
+
+def _merge_lora(generator: WanDiffusionWrapper, config, lora_ckpt_path: str) -> WanDiffusionWrapper:
+    if not getattr(config, "adapter", None):
+        raise ValueError("LoRA merge was requested, but config.adapter is missing.")
+    if not lora_ckpt_path:
+        raise ValueError("LoRA merge was requested, but no lora_ckpt was provided.")
+
+    print(f"Applying LoRA config: {config.adapter}")
+    generator.model = configure_lora_for_model(
+        generator.model,
+        model_name="generator",
+        lora_config=config.adapter,
+        is_main_process=True,
+    )
+
+    import peft
+
+    print(f"Loading LoRA weights: {lora_ckpt_path}")
+    peft.set_peft_model_state_dict(generator.model, _load_lora_state(lora_ckpt_path))  # type: ignore[arg-type]
+    print("Merging LoRA into generator...")
+    generator.model = generator.model.merge_and_unload(safe_merge=True)
+    return generator
+
+
+def _metadata(
+    config,
+    args: argparse.Namespace,
+    *,
+    backend: str,
+    matched_modules: list[str],
+    materialized_modules: list[str],
+):
+    checkpoint_format = (
+        "longlive_generator_merged_bf16"
+        if backend == "transformer_engine"
+        else NVFP4_CHECKPOINT_FORMAT
+    )
+    quant_format = "bf16" if backend == "transformer_engine" else "nvfp4"
+    quant_backend = "transformer_engine_runtime" if backend == "transformer_engine" else "fouroversix"
+    return {
+        "checkpoint_format": checkpoint_format,
+        "checkpoint_version": NVFP4_CHECKPOINT_VERSION,
+        "source_generator_ckpt": args.generator_ckpt,
+        "source_lora_ckpt": args.lora_ckpt,
+        "merged_lora": bool(args.lora_ckpt and not args.no_merge_lora),
+        "model_name": getattr(config.model_kwargs, "model_name", None),
+        "quantization": {
+            "format": quant_format,
+            "backend": quant_backend,
+            "materialized": backend == "fouroversix",
+            "dtype": "bfloat16",
+            "scale_rule": getattr(config, "model_quant_scale_rule", "static_6"),
+            "activation_scale_rule": getattr(config, "model_quant_activation_scale_rule", None),
+            "weight_scale_rule": getattr(config, "model_quant_weight_scale_rule", None),
+            "gradient_scale_rule": getattr(config, "model_quant_gradient_scale_rule", None),
+            "te_inference_only": bool(
+                getattr(config, "model_quant_te_inference_only", backend == "transformer_engine")
+            ),
+            "te_low_precision_weights": bool(
+                getattr(config, "model_quant_te_low_precision_weights", backend == "transformer_engine")
+            ),
+            "te_fallback_to_fouroversix": bool(
+                getattr(config, "model_quant_te_fallback_to_fouroversix", False)
+            ),
+            "matched_filtered_modules": matched_modules,
+            "materialized_modules": materialized_modules,
+        },
+    }
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Merge generator LoRA weights and save either packed FourOverSix NVFP4 or TE-ready bf16."
+    )
+    parser.add_argument("--config_path", required=True, help="Inference yaml that contains model/adapter/quant settings.")
+    parser.add_argument("--output_path", required=True, help="Path to save the generator .pt file.")
+    parser.add_argument("--generator_ckpt", default=None, help="Override checkpoints.generator_ckpt from the yaml.")
+    parser.add_argument("--lora_ckpt", default=None, help="Override checkpoints.lora_ckpt from the yaml.")
+    parser.add_argument("--device", default="cuda:0", help="Device used for quantization, e.g. cuda:0 or cpu.")
+    parser.add_argument(
+        "--backend",
+        choices=("fouroversix", "transformer_engine"),
+        default="fouroversix",
+        help=(
+            "fouroversix saves packed/materialized NVFP4. "
+            "transformer_engine saves merged bf16 for TE runtime quantization."
+        ),
+    )
+    parser.add_argument("--no_merge_lora", action="store_true", help="Quantize the base generator without merging LoRA.")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    config = normalize_config(OmegaConf.load(args.config_path))
+    args.generator_ckpt = args.generator_ckpt or getattr(config, "generator_ckpt", None)
+    args.lora_ckpt = args.lora_ckpt or getattr(config, "lora_ckpt", None)
+
+    if not args.generator_ckpt:
+        raise ValueError("Missing generator checkpoint. Set checkpoints.generator_ckpt or pass --generator_ckpt.")
+    config.model_quant_use_transformer_engine = args.backend == "transformer_engine"
+
+    device = torch.device(args.device)
+    print(f"Building generator on CPU: {config.model_kwargs}")
+    generator = WanDiffusionWrapper(**getattr(config, "model_kwargs", {}), is_causal=True)
+    generator.eval().requires_grad_(False)
+
+    print(f"Loading base generator checkpoint: {args.generator_ckpt}")
+    _load_generator_checkpoint(generator, args.generator_ckpt, use_ema=bool(getattr(config, "use_ema", False)))
+
+    should_merge_lora = bool(getattr(config, "merge_lora", False)) and not args.no_merge_lora
+    if should_merge_lora:
+        generator = _merge_lora(generator, config, args.lora_ckpt)
+    else:
+        print("Skipping LoRA merge; quantizing the loaded generator as-is.")
+
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    print(f"Moving generator to {device} and casting to bfloat16...")
+    generator.to(device=device, dtype=torch.bfloat16)
+    materialized_modules = []
+    if args.backend == "transformer_engine":
+        matched_modules = []
+        print(
+            "Saving merged bf16 weights for TransformerEngine runtime quantization. "
+            "TransformerEngine state_dict is not a packed NVFP4 storage format."
+        )
+    else:
+        generator.model, matched_modules = quantize_model_for_fouroversix_nvfp4(
+            generator.model,
+            config=config,
+            keep_master_weights=False,
+            verbose=True,
+        )
+
+        print("Materializing NVFP4 weights and dropping bf16 master weights...")
+        materialized_modules, master_bytes, quantized_bytes = _materialize_quantized_weights_for_inference(
+            generator.model,
+            target_device=device,
+        )
+        print(
+            "[NVFP4] "
+            f"materialized_modules={len(materialized_modules)}, "
+            f"master_weight={master_bytes / (1024 ** 3):.3f} GiB, "
+            f"quantized_weight={quantized_bytes / (1024 ** 3):.3f} GiB"
+        )
+
+    print("Copying checkpoint tensors to CPU for saving...")
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    output_path = Path(args.output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    checkpoint = {
+        "generator": cpu_state_dict(generator),
+        **_metadata(
+            config,
+            args,
+            backend=args.backend,
+            matched_modules=matched_modules,
+            materialized_modules=materialized_modules,
+        ),
+    }
+    torch.save(checkpoint, output_path)
+    size_gib = os.path.getsize(output_path) / (1024 ** 3)
+    print(f"Saved {args.backend} generator checkpoint to {output_path} ({size_gib:.3f} GiB)")
+
+
+if __name__ == "__main__":
+    with torch.no_grad():
+        main()