#!/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())