Add direct H3 audio VAE decode

src/h3_blackwell_runtime/audio_vae_decoder.py

@@ -0,0 +1,191 @@
+"""Standalone decoder-only MiniMax H3 audio VAE."""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from torch import nn
+from torch.nn import functional as F
+
+
+def snake(x: torch.Tensor, alpha: torch.Tensor, beta: torch.Tensor) -> torch.Tensor:
+    t = torch.sin(alpha * x)
+    return t.mul_(t).mul_((beta + 1e-9).reciprocal()).add_(x)
+
+
+class SnakeBeta(nn.Module):
+    def __init__(self, channels: int, *, device=None):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.empty(channels, device=device))
+        self.beta = nn.Parameter(torch.empty(channels, device=device))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        alpha = torch.exp(self.alpha.to(device=x.device, dtype=x.dtype)).view(1, -1, 1)
+        beta = torch.exp(self.beta.to(device=x.device, dtype=x.dtype)).view(1, -1, 1)
+        return snake(x, alpha, beta)
+
+
+def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
+    even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+    delta_f = 4 * half_width
+    a = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+    if a > 50.0:
+        beta = 0.1102 * (a - 8.7)
+    elif a >= 21.0:
+        beta = 0.5842 * (a - 21) ** 0.4 + 0.07886 * (a - 21.0)
+    else:
+        beta = 0.0
+    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+    time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size
+    filt = 2 * cutoff * window * torch.sinc(2 * cutoff * time)
+    filt /= filt.sum()
+    return filt.view(1, 1, kernel_size)
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, ratio: int = 2, kernel_size: int = 12, *, device=None):
+        super().__init__()
+        self.ratio = ratio
+        self.stride = ratio
+        self.pad = kernel_size // ratio - 1
+        self.pad_left = self.pad * ratio + (kernel_size - ratio) // 2
+        self.pad_right = self.pad * ratio + (kernel_size - ratio + 1) // 2
+        self.register_buffer("filter", kaiser_sinc_filter1d(0.5 / ratio, 0.6 / ratio, kernel_size).to(device=device))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        _, channels, _ = x.shape
+        x = F.pad(x, (self.pad, self.pad), mode="replicate")
+        filt = self.filter.to(device=x.device, dtype=x.dtype).expand(channels, -1, -1)
+        x = F.conv_transpose1d(x, filt, stride=self.stride, groups=channels).mul_(self.ratio)
+        return x[..., self.pad_left:-self.pad_right]
+
+
+class LowPassFilter1d(nn.Module):
+    def __init__(self, cutoff: float = 0.5, half_width: float = 0.6, stride: int = 1, kernel_size: int = 12, *, device=None):
+        super().__init__()
+        self.pad_left = kernel_size // 2 - int(kernel_size % 2 == 0)
+        self.pad_right = kernel_size // 2
+        self.stride = stride
+        self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size).to(device=device))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        _, channels, _ = x.shape
+        x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate")
+        filt = self.filter.to(device=x.device, dtype=x.dtype).expand(channels, -1, -1)
+        return F.conv1d(x, filt, stride=self.stride, groups=channels)
+
+
+class DownSample1d(nn.Module):
+    def __init__(self, ratio: int = 2, kernel_size: int = 12, *, device=None):
+        super().__init__()
+        self.lowpass = LowPassFilter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, stride=ratio, kernel_size=kernel_size, device=device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.lowpass(x)
+
+
+class Activation1d(nn.Module):
+    def __init__(self, activation: nn.Module, *, device=None):
+        super().__init__()
+        self.act = activation
+        self.upsample = UpSample1d(device=device)
+        self.downsample = DownSample1d(device=device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.downsample(self.act(self.upsample(x)))
+
+
+def get_padding(kernel_size: int, dilation: int = 1) -> int:
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+class AMPBlock1(nn.Module):
+    def __init__(self, channels: int, kernel_size: int = 3, dilation: tuple[int, int, int] = (1, 3, 5), *, device=None):
+        super().__init__()
+        self.convs1 = nn.ModuleList([nn.Conv1d(channels, channels, kernel_size, dilation=d, padding=get_padding(kernel_size, d), device=device) for d in dilation])
+        self.convs2 = nn.ModuleList([nn.Conv1d(channels, channels, kernel_size, dilation=1, padding=get_padding(kernel_size, 1), device=device) for _ in dilation])
+        self.activations = nn.ModuleList([Activation1d(SnakeBeta(channels, device=device), device=device) for _ in range(len(dilation) * 2)])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        acts1, acts2 = self.activations[::2], self.activations[1::2]
+        for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2):
+            residual = conv2(act2(conv1(act1(x))))
+            x = residual.add_(x)
+        return x
+
+
+class BigVGAN(nn.Module):
+    def __init__(self, num_mels: int = 2048, upsample_initial_channel: int = 1024, *, device=None):
+        super().__init__()
+        upsample_rates = (5, 5, 2, 2, 2, 2, 2)
+        upsample_kernel_sizes = (9, 9, 4, 4, 4, 4, 4)
+        resblock_kernel_sizes = (3, 7, 11)
+        resblock_dilation_sizes = ((1, 3, 5), (1, 3, 5), (1, 3, 5))
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.conv_pre = nn.Conv1d(num_mels, upsample_initial_channel, 7, 1, padding=3, device=device)
+        self.ups = nn.ModuleList()
+        for index, (rate, kernel) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(nn.ModuleList([nn.ConvTranspose1d(upsample_initial_channel // (2 ** index), upsample_initial_channel // (2 ** (index + 1)), kernel, rate, padding=(kernel - rate) // 2, device=device)]))
+        self.resblocks = nn.ModuleList()
+        for index in range(len(self.ups)):
+            channels = upsample_initial_channel // (2 ** (index + 1))
+            for kernel, dilation in zip(resblock_kernel_sizes, resblock_dilation_sizes):
+                self.resblocks.append(AMPBlock1(channels, kernel, dilation, device=device))
+        self.activation_post = Activation1d(SnakeBeta(channels, device=device), device=device)
+        self.conv_post = nn.Conv1d(channels, 1, 7, 1, padding=3, bias=False, device=device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv_pre(x)
+        for index in range(self.num_upsamples):
+            for upsample in self.ups[index]:
+                x = upsample(x)
+            combined = None
+            for kernel_index in range(self.num_kernels):
+                value = self.resblocks[index * self.num_kernels + kernel_index](x)
+                combined = value if combined is None else combined + value
+            x = combined.div_(self.num_kernels)
+        return self.conv_post(self.activation_post(x)).clamp_(-1.0, 1.0)
+
+
+class MiniMaxH3AudioVAE(nn.Module):
+    """Decoder-only MiniMax H3 stereo audio VAE at 32 kHz."""
+
+    def __init__(self, latent_channels: int = 32, latent_dim: int = 2048, decoder_dim: int = 1024, *, device=None):
+        super().__init__()
+        self.sample_rate = 32000
+        self.output_sample_rate = self.sample_rate
+        self.samples_per_latent = 800
+        self.latents_per_second = 40
+        self.dec_in_proj = nn.Conv1d(latent_channels, latent_dim, 1, device=device)
+        self.decoder = BigVGAN(num_mels=latent_dim, upsample_initial_channel=decoder_dim, device=device)
+        self.register_buffer("latents_mean", torch.empty(latent_channels, device=device))
+        self.register_buffer("latents_std", torch.empty(latent_channels, device=device))
+
+    @classmethod
+    def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.float32) -> "MiniMaxH3AudioVAE":
+        model = cls(device="meta")
+        expected = model.state_dict()
+        with safe_open(str(path), framework="pt", device=str(device)) as checkpoint:
+            missing = sorted(set(expected) - set(checkpoint.keys()))
+            shape_errors = [(name, tuple(expected[name].shape), tuple(checkpoint.get_slice(name).get_shape())) for name in expected if name in checkpoint.keys() and tuple(expected[name].shape) != tuple(checkpoint.get_slice(name).get_shape())]
+            if missing or shape_errors:
+                details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else [])
+                raise ValueError("incompatible H3 audio VAE checkpoint; " + "; ".join(details))
+            weights = {name: checkpoint.get_tensor(name).to(dtype=dtype) for name in expected}
+        model.load_state_dict(weights, strict=True, assign=True)
+        return model
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        """Decode normalized latents `[B,32,2,T]` to waveform `[B,2,L]`."""
+        batch, channels, stereo, steps = z.shape
+        z = z.permute(0, 2, 1, 3).reshape(batch * stereo, channels, steps)
+        mean = self.latents_mean.view(1, -1, 1).to(device=z.device, dtype=z.dtype)
+        std = self.latents_std.view(1, -1, 1).to(device=z.device, dtype=z.dtype)
+        z = z * std + mean
+        waveform = self.decoder(self.dec_in_proj(z))
+        return waveform.reshape(batch, stereo, -1)

src/h3_blackwell_runtime/sampler.py

@@ -62,6 +62,7 @@ def sample_video_res_multistep(
     *,
     steps: int = 12,
     model_timesteps: list[torch.Tensor] | tuple[torch.Tensor, ...] | None = None,
+    return_audio: bool = False,
 ) -> torch.Tensor:
     """Direct H3 beta/RES sampling with Comfy-equivalent joint AV carry semantics."""
     sigmas = beta_sigmas(steps, device=video.device)
@@ -101,7 +102,7 @@ def sample_video_res_multistep(
             f"{time.perf_counter() - step_started:.1f}s, elapsed {elapsed:.1f}s, eta {eta:.1f}s",
             flush=True,
         )
-    return video
+    return (video, audio_carried) if return_audio else video
 
 
 @torch.inference_mode()

tools/decode_audio_latent.py

@@ -0,0 +1,38 @@
+"""Decode a saved H3 audio latent to WAV."""
+
+import argparse
+import subprocess
+from pathlib import Path
+
+import torch
+
+from h3_blackwell_runtime.audio_vae_decoder import MiniMaxH3AudioVAE
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--latent", type=Path, required=True)
+parser.add_argument("--output", type=Path, required=True)
+args = parser.parse_args()
+
+state = torch.load(args.latent, map_location="cuda", weights_only=False)
+if isinstance(state, dict):
+    latent = state.get("audio_latent", state.get("latent"))
+    if latent is None:
+        raise ValueError("saved state does not contain 'audio_latent' or 'latent'")
+else:
+    latent = state
+latent = latent.to("cuda")
+
+vae = MiniMaxH3AudioVAE.from_safetensors("/vae/minimax_h3_audio_vae_fp32.safetensors", device="cuda").eval()
+with torch.inference_mode():
+    waveform = vae.decode(latent.to(next(vae.parameters()).dtype)).clamp(-1, 1).cpu()[0]
+
+args.output.parent.mkdir(parents=True, exist_ok=True)
+raw = args.output.with_suffix(".f32le")
+waveform.transpose(0, 1).contiguous().numpy().tofile(raw)
+subprocess.run([
+    "ffmpeg", "-y", "-f", "f32le", "-ar", "32000", "-ac", "2",
+    "-i", str(raw), str(args.output),
+], check=True)
+raw.unlink()
+print({"output": str(args.output), "sample_rate": 32000, "shape": tuple(waveform.shape)})

tools/direct_t2v_preview.py

@@ -13,6 +13,7 @@ warnings.filterwarnings("ignore", message="Found GPU0 NVIDIA GB10 which is of cu
 import torch
 
 from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.audio_vae_decoder import MiniMaxH3AudioVAE
 from h3_blackwell_runtime.denoiser import H3PackedDenoiser
 from h3_blackwell_runtime.packing import H3PromptPacker
 from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner
@@ -34,6 +35,9 @@ parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default="
 parser.add_argument("--model-timesteps-capture", type=Path, help="Directory containing captured input_XX.pt H3 timesteps for strict parity checks.")
 parser.add_argument("--profile-memory", action="store_true")
 parser.add_argument("--save-latent", type=Path)
+parser.add_argument("--save-audio-latent", type=Path)
+parser.add_argument("--audio-output", type=Path)
+parser.add_argument("--mux-audio", action="store_true")
 parser.add_argument("--skip-decode", action="store_true")
 args = parser.parse_args()
 started = time.perf_counter()
@@ -77,12 +81,26 @@ if args.model_timesteps_capture is not None:
         torch.load(args.model_timesteps_capture / f"input_{index:02d}.pt", map_location="cuda", weights_only=False)["timesteps"]
         for index in range(args.steps)
     ]
-latent = sample_video_res_multistep(model, H3PromptPacker(checkpoint), text, video, audio, steps=args.steps, model_timesteps=model_timesteps)
+want_audio = args.save_audio_latent is not None or args.audio_output is not None or args.mux_audio
+sampled = sample_video_res_multistep(model, H3PromptPacker(checkpoint), text, video, audio, steps=args.steps, model_timesteps=model_timesteps, return_audio=want_audio)
+if want_audio:
+    latent, audio_latent = sampled
+else:
+    latent, audio_latent = sampled, None
 report_memory("sampled")
 if args.save_latent is not None:
     args.save_latent.parent.mkdir(parents=True, exist_ok=True)
-    torch.save({"latent": latent.detach().cpu(), "frames": frames, "width": args.width, "height": args.height, "prompt": args.prompt, "seed": args.seed}, args.save_latent)
+    state = {"latent": latent.detach().cpu(), "frames": frames, "width": args.width, "height": args.height, "prompt": args.prompt, "seed": args.seed}
+    if audio_latent is not None:
+        state["audio_latent"] = audio_latent.detach().cpu()
+    torch.save(state, args.save_latent)
     print({"latent": str(args.save_latent)}, flush=True)
+if args.save_audio_latent is not None:
+    if audio_latent is None:
+        raise RuntimeError("audio latent was not sampled")
+    args.save_audio_latent.parent.mkdir(parents=True, exist_ok=True)
+    torch.save({"audio_latent": audio_latent.detach().cpu(), "frames": frames, "prompt": args.prompt, "seed": args.seed}, args.save_audio_latent)
+    print({"audio_latent": str(args.save_audio_latent)}, flush=True)
 if args.skip_decode:
     raise SystemExit(0)
 vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda").eval()
@@ -90,11 +108,34 @@ report_memory("vae_loaded")
 with torch.inference_mode():
     pixels = vae.decode(latent.to(next(vae.parameters()).dtype))[:, :, :frames]
 report_memory("vae_decoded")
+del vae
+torch.cuda.empty_cache()
 pixels = ((pixels[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu()
 report_memory("pixels_cpu")
 args.output.parent.mkdir(parents=True, exist_ok=True)
 raw = args.output.with_suffix(".rgb")
+video_output = args.output.with_name(args.output.stem + ".video.mp4") if args.mux_audio else args.output
 pixels.numpy().tofile(raw)
-subprocess.run(["ffmpeg", "-y", "-f", "rawvideo", "-pixel_format", "rgb24", "-video_size", f"{pixels.shape[2]}x{pixels.shape[1]}", "-framerate", "24", "-i", str(raw), "-an", "-c:v", "libx264", "-pix_fmt", "yuv420p", str(args.output)], check=True)
+subprocess.run(["ffmpeg", "-y", "-f", "rawvideo", "-pixel_format", "rgb24", "-video_size", f"{pixels.shape[2]}x{pixels.shape[1]}", "-framerate", "24", "-i", str(raw), "-an", "-c:v", "libx264", "-pix_fmt", "yuv420p", str(video_output)], check=True)
 raw.unlink()
+audio_path = args.audio_output
+if args.mux_audio and audio_path is None:
+    audio_path = args.output.with_suffix(".wav")
+if audio_path is not None:
+    if audio_latent is None:
+        raise RuntimeError("audio latent was not sampled")
+    audio_vae = MiniMaxH3AudioVAE.from_safetensors("/vae/minimax_h3_audio_vae_fp32.safetensors", device="cuda").eval()
+    report_memory("audio_vae_loaded")
+    with torch.inference_mode():
+        waveform = audio_vae.decode(audio_latent.to(next(audio_vae.parameters()).dtype)).clamp(-1, 1).cpu()[0]
+    report_memory("audio_decoded")
+    audio_path.parent.mkdir(parents=True, exist_ok=True)
+    audio_raw = audio_path.with_suffix(".f32le")
+    waveform.transpose(0, 1).contiguous().numpy().tofile(audio_raw)
+    subprocess.run(["ffmpeg", "-y", "-f", "f32le", "-ar", "32000", "-ac", "2", "-i", str(audio_raw), str(audio_path)], check=True)
+    audio_raw.unlink()
+    print({"audio_output": str(audio_path), "sample_rate": 32000, "audio_shape": tuple(waveform.shape)}, flush=True)
+if args.mux_audio:
+    subprocess.run(["ffmpeg", "-y", "-i", str(video_output), "-i", str(audio_path), "-c:v", "copy", "-c:a", "aac", "-shortest", str(args.output)], check=True)
+    video_output.unlink()
 print({"output": str(args.output), "frames": frames, "shape": tuple(pixels.shape)})