scripts/decode_h3_latents.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import gc
import json
import resource
import time
from pathlib import Path

import av
import torch
from accelerate import cpu_offload
from diffusers import AutoencoderKLMiniMaxH3, AutoencoderKLMiniMaxH3Audio
from diffusers.modular_pipelines.minimax_h3.packing import (
    MINIMAX_H3_PIXEL_MEAN,
    MINIMAX_H3_PIXEL_STD,
)
from diffusers.utils.export_utils import _prepare_audio_stream, _write_audio, encode_video
from diffusers.video_processor import VideoProcessor

from checkpoint_io import atomic_json_write
from latent_io import load_latent_bundle, prepare_inference_model
from media_packaging import atomic_media_output


def encode_master_crf1(
    frames: list,
    *,
    fps: int,
    output_path: Path,
    audio: torch.Tensor,
    audio_sample_rate: int,
) -> None:
    """Write the pre-delivery master before any HEVC recompression."""
    container = av.open(str(output_path), mode="w", format="mp4")
    stream = container.add_stream(
        "libx264",
        rate=fps,
        options={"crf": "1", "preset": "slow"},
    )
    stream.width = frames[0].width
    stream.height = frames[0].height
    stream.pix_fmt = "yuv444p"
    audio_stream = _prepare_audio_stream(container, audio_sample_rate)
    try:
        for image in frames:
            frame = av.VideoFrame.from_image(image)
            for packet in stream.encode(frame):
                container.mux(packet)
        for packet in stream.encode():
            container.mux(packet)
        _write_audio(container, audio_stream, audio, audio_sample_rate, av)
    finally:
        container.close()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--latents", type=Path, required=True)
    parser.add_argument("--vae", type=Path, required=True)
    parser.add_argument("--audio-vae", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--preview-output", type=Path, required=True)
    parser.add_argument("--disable-spatial-tiling", action="store_true")
    args = parser.parse_args()

    report_path = args.output.with_suffix(".decode.metrics.json")
    report: dict[str, object] = {
        "status": "running",
        "latent_bundle": str(args.latents),
        "vae": str(args.vae),
        "vae_dtype": "fp32_source",
        "audio_vae": str(args.audio_vae),
        "audio_vae_dtype": "fp32_source",
        "offload": {
            "video_vae": "sequential_cpu_offload",
            "audio_vae": "resident_cuda_during_audio_stage",
        },
        "video_vae_tiling": not args.disable_spatial_tiling,
        "master_video": {
            "codec": "libx264",
            "crf": 1,
            "preset": "slow",
            "pixel_format": "yuv444p",
        },
        "gpu": torch.cuda.get_device_name(),
    }
    report_path.parent.mkdir(parents=True, exist_ok=True)
    atomic_json_write(report, report_path)
    try:
        bundle = load_latent_bundle(args.latents)
        metadata = dict(bundle["metadata"])
        report["metadata"] = metadata

        video_load_started = time.perf_counter()
        vae = prepare_inference_model(
            AutoencoderKLMiniMaxH3.from_pretrained(
                args.vae,
                low_cpu_mem_usage=True,
                dtype=torch.float32,
            )
        )
        if args.disable_spatial_tiling:
            vae.disable_tiling()
        else:
            vae.enable_tiling()
        cpu_offload(
            vae,
            execution_device=torch.device("cuda"),
            offload_buffers=True,
        )
        report["video_vae_load_seconds"] = time.perf_counter() - video_load_started

        torch.cuda.reset_peak_memory_stats()
        video_decode_started = time.perf_counter()
        video_latents = bundle["video_latents"].to("cuda")
        with torch.inference_mode(), torch.autocast(device_type="cuda", dtype=torch.float16):
            video = vae.decode(video_latents, return_dict=False)[0]
        pixel_mean = torch.tensor(MINIMAX_H3_PIXEL_MEAN, device="cuda").view(
            1, -1, 1, 1, 1
        )
        pixel_std = torch.tensor(MINIMAX_H3_PIXEL_STD, device="cuda").view(
            1, -1, 1, 1, 1
        )
        video = (video.float() * pixel_std + pixel_mean).clamp(0, 1)
        videos = VideoProcessor(vae_scale_factor=16, do_normalize=False).postprocess_video(
            video,
            output_type="pil",
        )
        torch.cuda.synchronize()
        report["video_decode_seconds"] = time.perf_counter() - video_decode_started
        report["video_decode_torch_peak_bytes"] = torch.cuda.max_memory_allocated()
        report["video_frames"] = len(videos[0])

        del vae, video_latents, video, pixel_mean, pixel_std
        gc.collect()
        torch.cuda.empty_cache()

        audio_load_started = time.perf_counter()
        audio_vae = prepare_inference_model(
            AutoencoderKLMiniMaxH3Audio.from_pretrained(
                args.audio_vae,
                low_cpu_mem_usage=True,
                dtype=torch.float32,
            )
        ).to("cuda")
        report["audio_vae_offload"] = "resident_cuda_during_audio_stage"
        report["audio_vae_load_seconds"] = time.perf_counter() - audio_load_started

        torch.cuda.reset_peak_memory_stats()
        audio_decode_started = time.perf_counter()
        audio_latents = bundle["audio_latents"].to("cuda")
        with torch.inference_mode():
            audio = audio_vae.decode(audio_latents, return_dict=False)[0]
        audio = audio.float().permute(1, 0, 2)
        if not torch.isfinite(audio).all() or torch.count_nonzero(audio) == 0:
            raise RuntimeError("source audio VAE produced invalid or silent audio")
        torch.cuda.synchronize()
        report["audio_decode_seconds"] = time.perf_counter() - audio_decode_started
        report["audio_decode_torch_peak_bytes"] = torch.cuda.max_memory_allocated()
        report["audio_min"] = float(audio.min())
        report["audio_max"] = float(audio.max())
        report["audio_std"] = float(audio.std())

        preview_started = time.perf_counter()
        with atomic_media_output(args.preview_output) as partial:
            encode_video(
                videos[0],
                fps=int(metadata["fps"]),
                output_path=str(partial),
                audio=audio[0],
                audio_sample_rate=int(metadata["sampling_rate"]),
            )
        report["preview_encode_seconds"] = time.perf_counter() - preview_started
        report["preview_output"] = str(args.preview_output)
        report["preview_output_bytes"] = args.preview_output.stat().st_size
        atomic_json_write(report, report_path)

        encode_started = time.perf_counter()
        with atomic_media_output(args.output) as partial:
            encode_master_crf1(
                videos[0],
                fps=int(metadata["fps"]),
                output_path=partial,
                audio=audio[0],
                audio_sample_rate=int(metadata["sampling_rate"]),
            )
        report["encode_seconds"] = time.perf_counter() - encode_started
        report["output_bytes"] = args.output.stat().st_size
        report["rss_peak_bytes"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024
        report["status"] = "pass"
    except Exception as error:
        report["status"] = "fail"
        report["error"] = f"{type(error).__name__}: {error}"
        raise
    finally:
        atomic_json_write(report, report_path)
        print(json.dumps(report))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())