from transformers import PretrainedConfig from .speaker_embedding_config import SpeakerEmbeddingConfig class HiFiGANConfig(PretrainedConfig): model_type = "hifi_gan" def __init__( self, # Audio settings raw_sample_rate=22050, target_sample_rate=24000, n_fft=1024, hop_length=256, win_length=1024, mel_channels=96, # Speaker embedding speaker_embed_config: SpeakerEmbeddingConfig = None, speaker_embed_checkpoint: str = None, update_speaker_embedding: bool = False, # Input feature type: "mel" (default) or "wav2vec_cnn" input_type: str = "mel", wav2vec_cnn_model_name: str = "facebook/wav2vec2-base", wav2vec_cnn_frozen: bool = True, wav2vec_cnn_sample_rate: int = 16000, # HiFi-GAN Generator config hifigan_in_channels=96, hifigan_out_channels=1, hifigan_channels=512, hifigan_global_channels=256, # Speaker conditioning hifigan_kernel_size=7, #hifigan_upsample_scales=[8, 8, 2, 2], hifigan_upsample_kernel_sizes=[8,8,4,4,4,4], #hifigan_upsample_kernel_sizes=[16, 16, 4, 4], hifigan_upsample_scales=[4,4,2,2,2,2], hifigan_resblock_kernel_sizes=[3, 7, 11], hifigan_resblock_dilations=[[1, 3, 5], [1, 3, 5], [1, 3, 5]], hifigan_use_additional_convs=True, hifigan_bias=False, hifigan_nonlinear_activation="snakebeta", snake_logscale=True, hifigan_nonlinear_activation_params={"negative_slope": 0.1}, hifigan_use_weight_norm=True, mpd_reshapes=[2, 3, 5, 7, 11], mrd_resolutions=[ (1024, 120, 600), (2048, 240, 1200), (512, 50, 240), ], # VAE config block_out_channel_mults=[1, 2, 2, 3], # Loss weights lambda_adv=1.0, # Adversarial loss weight lambda_feat_match=10.0, # Feature matching loss weight add_mel_loss=True, lambda_mel=15.0, # Mel spectrogram loss weight lambda_wavlm=15.0, # WavLM perceptual loss weight add_wavlm_loss=False, add_wav2vec_loss=False, lambda_wav2vec=15.0, # Wav2Vec perceptual loss weight add_whisper_encoder_loss=False, add_whisper_kl_loss=False, lambda_whisper_encoder=15.0, # Whisper encoder hidden states L1 weight lambda_whisper_kl=1.0, # Whisper decoder KL divergence weight add_whisper_feature_loss=False, lambda_whisper_features=15.0, # Whisper conv feature extractor L1 weight use_wavlm_features: bool = False, # Use WavLM instead of MEL wavlm_model_name: str = "microsoft/wavlm-base-plus", wavlm_layer_idx: int = 6, # Which layer to extract features from (0-11 for base) train_discriminator: bool = True, # Watermark (FiLM) add_film_watermark: bool = False, watermark_bits: int = 32, lambda_watermark: float = 1.0, watermark_film_hidden: int = 128, # VocBulwark watermarking add_watermark_vocbulwark: bool = False, vocbulwark_bits: int = 100, vocbulwark_ta_hidden: int = 128, lambda_vocbulwark_ext: float = 1.0, lambda_vocbulwark_mel: float = 0.1, lambda_vocbulwark_mstft: float = 0.1, lambda_vocbulwark_wavlm: float = 0.0, vocbulwark_acc_tau1: float = 0.9, vocbulwark_acc_tau2: float = 0.95, vocbulwark_bootstrap_threshold: float = 0.0, vocbulwark_zero_conv_std: float = 0.02, vocbulwark_fidelity_vs_clean: bool = False, **kwargs ): super().__init__(**kwargs) # Audio self.raw_sample_rate = raw_sample_rate self.target_sample_rate = target_sample_rate self.n_fft = n_fft self.hop_length = hop_length self.win_length = win_length # Mel self.mel_channels = mel_channels # Speaker self.speaker_embed_config = speaker_embed_config self.speaker_embed_checkpoint = speaker_embed_checkpoint self.update_speaker_embedding = update_speaker_embedding # Input type self.input_type = input_type self.wav2vec_cnn_model_name = wav2vec_cnn_model_name self.wav2vec_cnn_frozen = wav2vec_cnn_frozen self.wav2vec_cnn_sample_rate = wav2vec_cnn_sample_rate # VAE parameters self.block_out_channel_mults = block_out_channel_mults # HiFi-GAN parameters self.hifigan_in_channels = hifigan_in_channels self.hifigan_out_channels = hifigan_out_channels self.hifigan_channels = hifigan_channels self.hifigan_global_channels = hifigan_global_channels self.hifigan_kernel_size = hifigan_kernel_size self.hifigan_upsample_scales = hifigan_upsample_scales self.hifigan_upsample_kernel_sizes = hifigan_upsample_kernel_sizes self.hifigan_resblock_kernel_sizes = hifigan_resblock_kernel_sizes self.hifigan_resblock_dilations = hifigan_resblock_dilations self.hifigan_use_additional_convs = hifigan_use_additional_convs self.hifigan_bias = hifigan_bias self.hifigan_nonlinear_activation = hifigan_nonlinear_activation self.hifigan_nonlinear_activation_params = hifigan_nonlinear_activation_params self.hifigan_use_weight_norm = hifigan_use_weight_norm self.snake_logscale = snake_logscale self.mpd_reshapes = mpd_reshapes self.mrd_resolutions = mrd_resolutions # Loss weights self.lambda_adv = lambda_adv self.lambda_feat_match = lambda_feat_match self.add_mel_loss = add_mel_loss self.lambda_mel = lambda_mel self.lambda_wavlm = lambda_wavlm self.add_wavlm_loss = add_wavlm_loss self.add_wav2vec_loss = add_wav2vec_loss self.lambda_wav2vec = lambda_wav2vec self.add_whisper_encoder_loss = add_whisper_encoder_loss self.add_whisper_kl_loss = add_whisper_kl_loss self.lambda_whisper_encoder = lambda_whisper_encoder self.lambda_whisper_kl = lambda_whisper_kl self.add_whisper_feature_loss = add_whisper_feature_loss self.lambda_whisper_features = lambda_whisper_features self.use_wavlm_features = use_wavlm_features self.wavlm_model_name = wavlm_model_name self.wavlm_layer_idx = wavlm_layer_idx # train discrimnator or not self.train_discriminator = train_discriminator # Watermark (FiLM) self.add_film_watermark = add_film_watermark self.watermark_bits = watermark_bits self.lambda_watermark = lambda_watermark self.watermark_film_hidden = watermark_film_hidden # VocBulwark self.add_watermark_vocbulwark = add_watermark_vocbulwark self.vocbulwark_bits = vocbulwark_bits self.vocbulwark_ta_hidden = vocbulwark_ta_hidden self.lambda_vocbulwark_ext = lambda_vocbulwark_ext self.lambda_vocbulwark_mel = lambda_vocbulwark_mel self.lambda_vocbulwark_mstft = lambda_vocbulwark_mstft self.lambda_vocbulwark_wavlm = lambda_vocbulwark_wavlm self.vocbulwark_acc_tau1 = vocbulwark_acc_tau1 self.vocbulwark_acc_tau2 = vocbulwark_acc_tau2 self.vocbulwark_bootstrap_threshold = vocbulwark_bootstrap_threshold self.vocbulwark_zero_conv_std = vocbulwark_zero_conv_std self.vocbulwark_fidelity_vs_clean = vocbulwark_fidelity_vs_clean