File size: 5,912 Bytes
1fdc661 | 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """Self-contained inference model for the VocBulwark HiFi-GAN / BigVGAN vocoder.
Loadable with `AutoModel.from_pretrained(path, trust_remote_code=True)` without
the training repository. This is the lean vocoder: it takes a mel-spectrogram and
a precomputed speaker embedding, and does NOT include the speaker encoder (that
is a separate model), the training losses, discriminators, or watermark heads.
"""
from dataclasses import dataclass
import torch
import torch.nn.utils.parametrize as parametrize
from torch.nn.utils import remove_weight_norm
from transformers import PreTrainedModel
from transformers.modeling_outputs import ModelOutput
def _strip_weight_norm(root):
"""Remove every weight-norm reparametrization (both conventions), leaving
plain ``.weight`` parameters that match the folded export weights."""
for m in root.modules():
if parametrize.is_parametrized(m, "weight"):
parametrize.remove_parametrizations(m, "weight", leave_parametrized=True)
elif hasattr(m, "weight_g"):
try:
remove_weight_norm(m)
except (ValueError, RuntimeError):
pass
from .configuration_hifigan import HiFiGANConfig
from .bigvgan_model import BigVGAN
# transformers' trust_remote_code loader copies only modules *directly* imported
# by this entry file, and recognizes the `from .X import Y` form. Import a real
# symbol from every bundled leaf module so all files get copied into the
# dynamic-module cache (they are used transitively by the imports above).
from .bigvgan_activations import Snake as _Snake # noqa: F401
from .alias_free_act import Activation1d as _Activation1d # noqa: F401
from .alias_free_resample import UpSample1d as _UpSample1d # noqa: F401
from .alias_free_filter import LowPassFilter1d as _LowPassFilter1d # noqa: F401
from .temporal_adapter import TemporalAdapterBlock as _TemporalAdapterBlock # noqa: F401
@dataclass
class HiFiGANOutput(ModelOutput):
audio: torch.Tensor = None
class HiFiGANArchitecture(PreTrainedModel):
config_class = HiFiGANConfig
def __init__(self, config):
super().__init__(config)
# Lean vocoder: the speaker embedding is supplied at inference time, so the
# generator only needs its dimension to build the conditioning conv.
global_channels = int(getattr(config, "speaker_embedding_size", 0)) or -1
self.hifi_gan = BigVGAN(
num_mels=config.hifigan_in_channels,
upsample_initial_channel=config.hifigan_channels,
resblock_kernel_sizes=config.hifigan_resblock_kernel_sizes,
resblock_dilation_sizes=config.hifigan_resblock_dilations,
upsample_kernel_sizes=config.hifigan_upsample_kernel_sizes,
upsample_rates=config.hifigan_upsample_scales,
snake_logscale=config.snake_logscale,
activation=config.hifigan_nonlinear_activation,
use_bias_at_final=config.hifigan_bias,
global_channels=global_channels,
watermark_bits=config.watermark_bits if config.add_film_watermark else 0,
watermark_film_hidden=config.watermark_film_hidden,
vocbulwark_bits=config.vocbulwark_bits if config.add_watermark_vocbulwark else 0,
vocbulwark_ta_hidden=config.vocbulwark_ta_hidden,
vocbulwark_zero_conv_std=getattr(config, "vocbulwark_zero_conv_std", 0.02),
)
# Export folds weight-norm into plain .weight; strip the reparametrization
# here so module keys match (and inference is independent of torch's
# weight_norm convention).
_strip_weight_norm(self)
@torch.no_grad()
def forward(
self,
mel_spectrogram=None,
speaker_embedding=None,
input_features=None,
return_loss=False,
**kwargs,
):
"""Vocode a (log-)mel spectrogram into a waveform.
Args:
mel_spectrogram: [B, mel_channels, T] input features.
speaker_embedding: [B, speaker_embedding_size] speaker conditioning,
produced by the companion speaker-encoder model.
Returns:
HiFiGANOutput with `.audio` of shape [B, 1, samples] @ target_sample_rate.
The model's fixed provenance watermark is embedded in every generated clip
and cannot be changed or disabled through this interface.
"""
hifi_gan_device = self.hifi_gan.conv_pre.bias.device
if torch.is_autocast_enabled():
target_dtype = torch.get_autocast_gpu_dtype()
else:
target_dtype = self.hifi_gan.conv_pre.bias.dtype
if input_features is not None:
features = input_features.to(device=hifi_gan_device, dtype=target_dtype)
else:
features = mel_spectrogram.to(device=hifi_gan_device, dtype=target_dtype)
if speaker_embedding is None:
raise ValueError(
"speaker_embedding is required: pass a "
"[B, config.speaker_embedding_size] tensor from the companion "
"speaker-encoder model.")
g = speaker_embedding.to(device=hifi_gan_device, dtype=target_dtype)
if g.dim() == 2: # [B, E] -> [B, E, 1]
g = g.unsqueeze(2)
# Provenance watermark: the model's fixed signature is embedded in every
# generated clip. It is intentionally not overridable through this
# interface — there is no way to disable or change it here.
watermark = None
fw = getattr(self.config, "fixed_watermark", None)
if fw is not None:
watermark = torch.tensor(fw, dtype=target_dtype, device=hifi_gan_device)
watermark = watermark.unsqueeze(0).expand(features.shape[0], -1)
audio = self.hifi_gan(features, g=g, watermark=watermark)
return HiFiGANOutput(audio=audio)
|