File size: 1,474 Bytes
fa2d87b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | #!/usr/bin/env python3
from __future__ import annotations
from pathlib import Path
from typing import Any
import torch
FORMAT_VERSION = 1
def prepare_inference_model(model: torch.nn.Module) -> torch.nn.Module:
return model.eval().requires_grad_(False)
def save_latent_bundle(
path: Path,
*,
video_latents: torch.Tensor,
audio_latents: torch.Tensor,
metadata: dict[str, Any],
) -> None:
if video_latents.ndim != 5 or video_latents.shape[0] != 1:
raise ValueError(
"video latents must have shape (1, channels, frames, height, width)"
)
if audio_latents.ndim != 3 or audio_latents.shape[0] != 2:
raise ValueError("audio latents must have shape (2, channels, samples)")
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix(path.suffix + ".tmp")
torch.save(
{
"format_version": FORMAT_VERSION,
"metadata": dict(metadata),
"video_latents": video_latents.detach().to("cpu").contiguous(),
"audio_latents": audio_latents.detach().to("cpu").contiguous(),
},
temporary,
)
temporary.replace(path)
def load_latent_bundle(path: Path) -> dict[str, Any]:
payload = torch.load(path, map_location="cpu", weights_only=False)
if payload.get("format_version") != FORMAT_VERSION:
raise ValueError(f"unsupported latent bundle format: {payload.get('format_version')}")
return payload
|