diff --git a/src/h3_blackwell_runtime/audio_vae_decoder.py b/src/h3_blackwell_runtime/audio_vae_decoder.py new file mode 100644 index 0000000..756b638 --- /dev/null +++ b/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) diff --git a/src/h3_blackwell_runtime/sampler.py b/src/h3_blackwell_runtime/sampler.py index 3962745..a959d5f 100644 --- a/src/h3_blackwell_runtime/sampler.py +++ b/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() diff --git a/tools/decode_audio_latent.py b/tools/decode_audio_latent.py new file mode 100644 index 0000000..74d4415 --- /dev/null +++ b/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)}) diff --git a/tools/direct_t2v_preview.py b/tools/direct_t2v_preview.py index 1a20928..78be53d 100644 --- a/tools/direct_t2v_preview.py +++ b/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)})