Ref2VA/audio_vae/minimax_h3_audio_vae.py
| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # Remote entry: self-contained MiniMax H3 audio VAE (DAC-lineage encoder + BigVGAN decoder). |
| 3 | # Loaded via config.json:auto_map with trust_remote_code; weights are safetensors-only. |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | import json |
| 7 | from pathlib import Path |
| 8 | |
| 9 | import torch.nn as nn |
| 10 | |
| 11 | # --- dependency manifest --- |
| 12 | # diffusers' dynamic-module loader only copies ONE level of relative |
| 13 | # imports into its cache; list every bundle module here so all files |
| 14 | # are copied, letting their own second-level imports resolve. |
| 15 | from .dac_activations import SnakeBeta as _dep_dac_activations # noqa: F401 |
| 16 | from .dac_alias_free_act import Activation1d as _dep_dac_alias_free_act # noqa: F401 |
| 17 | from .dac_alias_free_filter import kaiser_sinc_filter1d as _dep_dac_alias_free_filter # noqa: F401 |
| 18 | from .dac_alias_free_resample import UpSample1d as _dep_dac_alias_free_resample # noqa: F401 |
| 19 | from .dac_attn_proj import GeGluMlp as _dep_dac_attn_proj # noqa: F401 |
| 20 | from .dac_bigvgan import AttrDict as _dep_dac_bigvgan # noqa: F401 |
| 21 | from .dac_audio_vae import AttrDict as _dep_dac_audio_vae # noqa: F401 |
| 22 | from .dac_utils import init_weights as _dep_dac_utils # noqa: F401 |
| 23 | # --- end dependency manifest --- |
| 24 | from safetensors.torch import load_file |
| 25 | |
| 26 | from .dac_audio_vae import DacAudioVAE |
| 27 | |
| 28 | |
| 29 | def _load_yaml(path: Path) -> dict: |
| 30 | try: |
| 31 | import yaml |
| 32 | except ImportError as exc: |
| 33 | raise ImportError("MiniMax H3 audio VAE requires PyYAML.") from exc |
| 34 | with path.open("r", encoding="utf-8") as f: |
| 35 | return yaml.safe_load(f) |
| 36 | |
| 37 | |
| 38 | class MiniMaxH3AudioVAE(nn.Module): |
| 39 | def __init__(self, model: nn.Module) -> None: |
| 40 | super().__init__() |
| 41 | self.model = model |
| 42 | |
| 43 | @classmethod |
| 44 | def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs): |
| 45 | component_dir = Path(pretrained_model_name_or_path) |
| 46 | with (component_dir / "config.json").open("r", encoding="utf-8") as f: |
| 47 | config = json.load(f) |
| 48 | |
| 49 | audio_config = _load_yaml(component_dir / config["source_config_path"]) |
| 50 | if "source_safetensors_path" not in config: |
| 51 | raise KeyError( |
| 52 | "source_safetensors_path is required; pickle checkpoints are not supported" |
| 53 | ) |
| 54 | if "source_metadata_path" not in config: |
| 55 | raise KeyError( |
| 56 | "source_metadata_path is required when source_safetensors_path is set" |
| 57 | ) |
| 58 | state_dict = load_file( |
| 59 | component_dir / config["source_safetensors_path"], device="cpu" |
| 60 | ) |
| 61 | with (component_dir / config["source_metadata_path"]).open( |
| 62 | "r", encoding="utf-8" |
| 63 | ) as f: |
| 64 | metadata_doc = json.load(f) |
| 65 | metadata = metadata_doc["metadata"]["kwargs"] |
| 66 | |
| 67 | model = DacAudioVAE( |
| 68 | encoder_rates=metadata["encoder_rates"], |
| 69 | decoder_rates=metadata["decoder_rates"], |
| 70 | attn_proj=metadata["attn_proj"], |
| 71 | decoder_type=metadata["decoder_type"], |
| 72 | decoder_dim=audio_config["model_config"]["decoder_dim"], |
| 73 | vae_latent_channels=audio_config["model_config"]["vae_latent_channels"], |
| 74 | sample_rate=metadata["sample_rate"], |
| 75 | ) |
| 76 | model.load_state_dict(state_dict, strict=True) |
| 77 | return cls(model.eval()) |
| 78 | |
| 79 | def decode(self, *args, **kwargs): |
| 80 | return self.model.decode(*args, **kwargs) |
| 81 | |
| 82 | def __getattr__(self, name: str): |
| 83 | try: |
| 84 | return super().__getattr__(name) |
| 85 | except AttributeError: |
| 86 | return getattr(self.model, name) |
| 87 | |