diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 0000000000000000000000000000000000000000..9e86607446c706ccfb8b630440d785592ed79a92
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1,27 @@
+campplus_cn_common.bin filter=lfs diff=lfs merge=lfs -text
+examples/reference/azuma_0.wav filter=lfs diff=lfs merge=lfs -text
+examples/reference/dingzhen_0.wav filter=lfs diff=lfs merge=lfs -text
+examples/reference/kobe_0.wav filter=lfs diff=lfs merge=lfs -text
+examples/reference/s1p1.wav filter=lfs diff=lfs merge=lfs -text
+examples/reference/s1p2.wav filter=lfs diff=lfs merge=lfs -text
+examples/reference/s2p1.wav filter=lfs diff=lfs merge=lfs -text
+examples/reference/s2p2.wav filter=lfs diff=lfs merge=lfs -text
+examples/reference/s3p1.wav filter=lfs diff=lfs merge=lfs -text
+examples/reference/s3p2.wav filter=lfs diff=lfs merge=lfs -text
+examples/reference/s4p1.wav filter=lfs diff=lfs merge=lfs -text
+examples/reference/s4p2.wav filter=lfs diff=lfs merge=lfs -text
+examples/reference/teio_0.wav filter=lfs diff=lfs merge=lfs -text
+examples/reference/trump_0.wav filter=lfs diff=lfs merge=lfs -text
+examples/source/glados_0.wav filter=lfs diff=lfs merge=lfs -text
+examples/source/jay_0.wav filter=lfs diff=lfs merge=lfs -text
+examples/source/source_s1.wav filter=lfs diff=lfs merge=lfs -text
+examples/source/source_s2.wav filter=lfs diff=lfs merge=lfs -text
+examples/source/source_s3.wav filter=lfs diff=lfs merge=lfs -text
+examples/source/source_s4.wav filter=lfs diff=lfs merge=lfs -text
+examples/source/TECHNOPOLIS[[:space:]]-[[:space:]]2085[[:space:]]\[vocals\]_\[cut_14sec\].wav filter=lfs diff=lfs merge=lfs -text
+examples/source/Wiz[[:space:]]Khalifa,Charlie[[:space:]]Puth[[:space:]]-[[:space:]]See[[:space:]]You[[:space:]]Again[[:space:]]\[vocals\]_\[cut_28sec\].wav filter=lfs diff=lfs merge=lfs -text
+examples/source/yae_0.wav filter=lfs diff=lfs merge=lfs -text
+modules/bigvgan/alias_free_activation/cuda/build/.ninja_deps filter=lfs diff=lfs merge=lfs -text
+modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation_cuda.cuda.o filter=lfs diff=lfs merge=lfs -text
+modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation_cuda.pyd filter=lfs diff=lfs merge=lfs -text
+modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation.o filter=lfs diff=lfs merge=lfs -text
diff --git a/README.md b/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..2180fcafd567aefcd41e271e150a52a743affb88
--- /dev/null
+++ b/README.md
@@ -0,0 +1,13 @@
+---
+title: Quantum Voice Clone
+emoji: 🎤🔄
+colorFrom: green
+colorTo: green
+sdk: gradio
+sdk_version: 5.24.0
+app_file: app_v1v2.py
+pinned: false
+license: apache-2.0
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
\ No newline at end of file
diff --git a/app.py b/app.py
new file mode 100644
index 0000000000000000000000000000000000000000..0d4bd9ab3941730b125bb5901d3fd77117798e00
--- /dev/null
+++ b/app.py
@@ -0,0 +1,372 @@
+import spaces
+import gradio as gr
+import torch
+import torchaudio
+import librosa
+from modules.commons import build_model, load_checkpoint, recursive_munch
+import yaml
+from hf_utils import load_custom_model_from_hf
+import numpy as np
+from pydub import AudioSegment
+
+# Load model and configuration
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+dit_checkpoint_path, dit_config_path = load_custom_model_from_hf("Plachta/Seed-VC",
+ "DiT_seed_v2_uvit_whisper_small_wavenet_bigvgan_pruned.pth",
+ "config_dit_mel_seed_uvit_whisper_small_wavenet.yml")
+
+config = yaml.safe_load(open(dit_config_path, 'r'))
+model_params = recursive_munch(config['model_params'])
+model = build_model(model_params, stage='DiT')
+hop_length = config['preprocess_params']['spect_params']['hop_length']
+sr = config['preprocess_params']['sr']
+
+# Load checkpoints
+model, _, _, _ = load_checkpoint(model, None, dit_checkpoint_path,
+ load_only_params=True, ignore_modules=[], is_distributed=False)
+for key in model:
+ model[key].eval()
+ model[key].to(device)
+model.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192)
+
+# Load additional modules
+from modules.campplus.DTDNN import CAMPPlus
+
+campplus_ckpt_path = load_custom_model_from_hf("funasr/campplus", "campplus_cn_common.bin", config_filename=None)
+campplus_model = CAMPPlus(feat_dim=80, embedding_size=192)
+campplus_model.load_state_dict(torch.load(campplus_ckpt_path, map_location="cpu"))
+campplus_model.eval()
+campplus_model.to(device)
+
+from modules.bigvgan import bigvgan
+
+bigvgan_model = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_22khz_80band_256x', use_cuda_kernel=False)
+
+# remove weight norm in the model and set to eval mode
+bigvgan_model.remove_weight_norm()
+bigvgan_model = bigvgan_model.eval().to(device)
+
+ckpt_path, config_path = load_custom_model_from_hf("Plachta/FAcodec", 'pytorch_model.bin', 'config.yml')
+
+codec_config = yaml.safe_load(open(config_path))
+codec_model_params = recursive_munch(codec_config['model_params'])
+codec_encoder = build_model(codec_model_params, stage="codec")
+
+ckpt_params = torch.load(ckpt_path, map_location="cpu")
+
+for key in codec_encoder:
+ codec_encoder[key].load_state_dict(ckpt_params[key], strict=False)
+_ = [codec_encoder[key].eval() for key in codec_encoder]
+_ = [codec_encoder[key].to(device) for key in codec_encoder]
+
+# whisper
+from transformers import AutoFeatureExtractor, WhisperModel
+
+whisper_name = model_params.speech_tokenizer.whisper_name if hasattr(model_params.speech_tokenizer,
+ 'whisper_name') else "openai/whisper-small"
+whisper_model = WhisperModel.from_pretrained(whisper_name, torch_dtype=torch.float16).to(device)
+del whisper_model.decoder
+whisper_feature_extractor = AutoFeatureExtractor.from_pretrained(whisper_name)
+
+# Generate mel spectrograms
+mel_fn_args = {
+ "n_fft": config['preprocess_params']['spect_params']['n_fft'],
+ "win_size": config['preprocess_params']['spect_params']['win_length'],
+ "hop_size": config['preprocess_params']['spect_params']['hop_length'],
+ "num_mels": config['preprocess_params']['spect_params']['n_mels'],
+ "sampling_rate": sr,
+ "fmin": 0,
+ "fmax": None,
+ "center": False
+}
+from modules.audio import mel_spectrogram
+
+to_mel = lambda x: mel_spectrogram(x, **mel_fn_args)
+
+# f0 conditioned model
+dit_checkpoint_path, dit_config_path = load_custom_model_from_hf("Plachta/Seed-VC",
+ "DiT_seed_v2_uvit_whisper_base_f0_44k_bigvgan_pruned_ft_ema.pth",
+ "config_dit_mel_seed_uvit_whisper_base_f0_44k.yml")
+
+config = yaml.safe_load(open(dit_config_path, 'r'))
+model_params = recursive_munch(config['model_params'])
+model_f0 = build_model(model_params, stage='DiT')
+hop_length = config['preprocess_params']['spect_params']['hop_length']
+sr = config['preprocess_params']['sr']
+
+# Load checkpoints
+model_f0, _, _, _ = load_checkpoint(model_f0, None, dit_checkpoint_path,
+ load_only_params=True, ignore_modules=[], is_distributed=False)
+for key in model_f0:
+ model_f0[key].eval()
+ model_f0[key].to(device)
+model_f0.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192)
+
+# f0 extractor
+from modules.rmvpe import RMVPE
+
+model_path = load_custom_model_from_hf("lj1995/VoiceConversionWebUI", "rmvpe.pt", None)
+rmvpe = RMVPE(model_path, is_half=False, device=device)
+
+mel_fn_args_f0 = {
+ "n_fft": config['preprocess_params']['spect_params']['n_fft'],
+ "win_size": config['preprocess_params']['spect_params']['win_length'],
+ "hop_size": config['preprocess_params']['spect_params']['hop_length'],
+ "num_mels": config['preprocess_params']['spect_params']['n_mels'],
+ "sampling_rate": sr,
+ "fmin": 0,
+ "fmax": None,
+ "center": False
+}
+to_mel_f0 = lambda x: mel_spectrogram(x, **mel_fn_args_f0)
+bigvgan_44k_model = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_44khz_128band_512x', use_cuda_kernel=False)
+
+# remove weight norm in the model and set to eval mode
+bigvgan_44k_model.remove_weight_norm()
+bigvgan_44k_model = bigvgan_44k_model.eval().to(device)
+
+def adjust_f0_semitones(f0_sequence, n_semitones):
+ factor = 2 ** (n_semitones / 12)
+ return f0_sequence * factor
+
+def crossfade(chunk1, chunk2, overlap):
+ fade_out = np.cos(np.linspace(0, np.pi / 2, overlap)) ** 2
+ fade_in = np.cos(np.linspace(np.pi / 2, 0, overlap)) ** 2
+ chunk2[:overlap] = chunk2[:overlap] * fade_in + chunk1[-overlap:] * fade_out
+ return chunk2
+
+# streaming and chunk processing related params
+bitrate = "320k"
+overlap_frame_len = 16
+@spaces.GPU
+@torch.no_grad()
+@torch.inference_mode()
+def voice_conversion(source, target, diffusion_steps, length_adjust, inference_cfg_rate, f0_condition, auto_f0_adjust, pitch_shift):
+ inference_module = model if not f0_condition else model_f0
+ mel_fn = to_mel if not f0_condition else to_mel_f0
+ bigvgan_fn = bigvgan_model if not f0_condition else bigvgan_44k_model
+ sr = 22050 if not f0_condition else 44100
+ hop_length = 256 if not f0_condition else 512
+ max_context_window = sr // hop_length * 30
+ overlap_wave_len = overlap_frame_len * hop_length
+ # Load audio
+ source_audio = librosa.load(source, sr=sr)[0]
+ ref_audio = librosa.load(target, sr=sr)[0]
+
+ # Process audio
+ source_audio = torch.tensor(source_audio).unsqueeze(0).float().to(device)
+ ref_audio = torch.tensor(ref_audio[:sr * 25]).unsqueeze(0).float().to(device)
+
+ # Resample
+ ref_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000)
+ converted_waves_16k = torchaudio.functional.resample(source_audio, sr, 16000)
+ # if source audio less than 30 seconds, whisper can handle in one forward
+ if converted_waves_16k.size(-1) <= 16000 * 30:
+ alt_inputs = whisper_feature_extractor([converted_waves_16k.squeeze(0).cpu().numpy()],
+ return_tensors="pt",
+ return_attention_mask=True,
+ sampling_rate=16000)
+ alt_input_features = whisper_model._mask_input_features(
+ alt_inputs.input_features, attention_mask=alt_inputs.attention_mask).to(device)
+ alt_outputs = whisper_model.encoder(
+ alt_input_features.to(whisper_model.encoder.dtype),
+ head_mask=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ )
+ S_alt = alt_outputs.last_hidden_state.to(torch.float32)
+ S_alt = S_alt[:, :converted_waves_16k.size(-1) // 320 + 1]
+ else:
+ overlapping_time = 5 # 5 seconds
+ S_alt_list = []
+ buffer = None
+ traversed_time = 0
+ while traversed_time < converted_waves_16k.size(-1):
+ if buffer is None: # first chunk
+ chunk = converted_waves_16k[:, traversed_time:traversed_time + 16000 * 30]
+ else:
+ chunk = torch.cat([buffer, converted_waves_16k[:, traversed_time:traversed_time + 16000 * (30 - overlapping_time)]], dim=-1)
+ alt_inputs = whisper_feature_extractor([chunk.squeeze(0).cpu().numpy()],
+ return_tensors="pt",
+ return_attention_mask=True,
+ sampling_rate=16000)
+ alt_input_features = whisper_model._mask_input_features(
+ alt_inputs.input_features, attention_mask=alt_inputs.attention_mask).to(device)
+ alt_outputs = whisper_model.encoder(
+ alt_input_features.to(whisper_model.encoder.dtype),
+ head_mask=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ )
+ S_alt = alt_outputs.last_hidden_state.to(torch.float32)
+ S_alt = S_alt[:, :chunk.size(-1) // 320 + 1]
+ if traversed_time == 0:
+ S_alt_list.append(S_alt)
+ else:
+ S_alt_list.append(S_alt[:, 50 * overlapping_time:])
+ buffer = chunk[:, -16000 * overlapping_time:]
+ traversed_time += 30 * 16000 if traversed_time == 0 else chunk.size(-1) - 16000 * overlapping_time
+ S_alt = torch.cat(S_alt_list, dim=1)
+
+ ori_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000)
+ ori_inputs = whisper_feature_extractor([ori_waves_16k.squeeze(0).cpu().numpy()],
+ return_tensors="pt",
+ return_attention_mask=True)
+ ori_input_features = whisper_model._mask_input_features(
+ ori_inputs.input_features, attention_mask=ori_inputs.attention_mask).to(device)
+ with torch.no_grad():
+ ori_outputs = whisper_model.encoder(
+ ori_input_features.to(whisper_model.encoder.dtype),
+ head_mask=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ )
+ S_ori = ori_outputs.last_hidden_state.to(torch.float32)
+ S_ori = S_ori[:, :ori_waves_16k.size(-1) // 320 + 1]
+
+ mel = mel_fn(source_audio.to(device).float())
+ mel2 = mel_fn(ref_audio.to(device).float())
+
+ target_lengths = torch.LongTensor([int(mel.size(2) * length_adjust)]).to(mel.device)
+ target2_lengths = torch.LongTensor([mel2.size(2)]).to(mel2.device)
+
+ feat2 = torchaudio.compliance.kaldi.fbank(ref_waves_16k,
+ num_mel_bins=80,
+ dither=0,
+ sample_frequency=16000)
+ feat2 = feat2 - feat2.mean(dim=0, keepdim=True)
+ style2 = campplus_model(feat2.unsqueeze(0))
+
+ if f0_condition:
+ F0_ori = rmvpe.infer_from_audio(ref_waves_16k[0], thred=0.5)
+ F0_alt = rmvpe.infer_from_audio(converted_waves_16k[0], thred=0.5)
+
+ F0_ori = torch.from_numpy(F0_ori).to(device)[None]
+ F0_alt = torch.from_numpy(F0_alt).to(device)[None]
+
+ voiced_F0_ori = F0_ori[F0_ori > 1]
+ voiced_F0_alt = F0_alt[F0_alt > 1]
+
+ log_f0_alt = torch.log(F0_alt + 1e-5)
+ voiced_log_f0_ori = torch.log(voiced_F0_ori + 1e-5)
+ voiced_log_f0_alt = torch.log(voiced_F0_alt + 1e-5)
+ median_log_f0_ori = torch.median(voiced_log_f0_ori)
+ median_log_f0_alt = torch.median(voiced_log_f0_alt)
+
+ # shift alt log f0 level to ori log f0 level
+ shifted_log_f0_alt = log_f0_alt.clone()
+ if auto_f0_adjust:
+ shifted_log_f0_alt[F0_alt > 1] = log_f0_alt[F0_alt > 1] - median_log_f0_alt + median_log_f0_ori
+ shifted_f0_alt = torch.exp(shifted_log_f0_alt)
+ if pitch_shift != 0:
+ shifted_f0_alt[F0_alt > 1] = adjust_f0_semitones(shifted_f0_alt[F0_alt > 1], pitch_shift)
+ else:
+ F0_ori = None
+ F0_alt = None
+ shifted_f0_alt = None
+
+ # Length regulation
+ cond, _, codes, commitment_loss, codebook_loss = inference_module.length_regulator(S_alt, ylens=target_lengths, n_quantizers=3, f0=shifted_f0_alt)
+ prompt_condition, _, codes, commitment_loss, codebook_loss = inference_module.length_regulator(S_ori, ylens=target2_lengths, n_quantizers=3, f0=F0_ori)
+
+ max_source_window = max_context_window - mel2.size(2)
+ # split source condition (cond) into chunks
+ processed_frames = 0
+ generated_wave_chunks = []
+ # generate chunk by chunk and stream the output
+ while processed_frames < cond.size(1):
+ chunk_cond = cond[:, processed_frames:processed_frames + max_source_window]
+ is_last_chunk = processed_frames + max_source_window >= cond.size(1)
+ cat_condition = torch.cat([prompt_condition, chunk_cond], dim=1)
+ with torch.autocast(device_type='cuda', dtype=torch.float16):
+ # Voice Conversion
+ vc_target = inference_module.cfm.inference(cat_condition,
+ torch.LongTensor([cat_condition.size(1)]).to(mel2.device),
+ mel2, style2, None, diffusion_steps,
+ inference_cfg_rate=inference_cfg_rate)
+ vc_target = vc_target[:, :, mel2.size(-1):]
+ vc_wave = bigvgan_fn(vc_target.float())[0]
+ if processed_frames == 0:
+ if is_last_chunk:
+ output_wave = vc_wave[0].cpu().numpy()
+ generated_wave_chunks.append(output_wave)
+ output_wave = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave.tobytes(), frame_rate=sr,
+ sample_width=output_wave.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=bitrate).read()
+ yield mp3_bytes, (sr, np.concatenate(generated_wave_chunks))
+ break
+ output_wave = vc_wave[0, :-overlap_wave_len].cpu().numpy()
+ generated_wave_chunks.append(output_wave)
+ previous_chunk = vc_wave[0, -overlap_wave_len:]
+ processed_frames += vc_target.size(2) - overlap_frame_len
+ output_wave = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave.tobytes(), frame_rate=sr,
+ sample_width=output_wave.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=bitrate).read()
+ yield mp3_bytes, None
+ elif is_last_chunk:
+ output_wave = crossfade(previous_chunk.cpu().numpy(), vc_wave[0].cpu().numpy(), overlap_wave_len)
+ generated_wave_chunks.append(output_wave)
+ processed_frames += vc_target.size(2) - overlap_frame_len
+ output_wave = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave.tobytes(), frame_rate=sr,
+ sample_width=output_wave.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=bitrate).read()
+ yield mp3_bytes, (sr, np.concatenate(generated_wave_chunks))
+ break
+ else:
+ output_wave = crossfade(previous_chunk.cpu().numpy(), vc_wave[0, :-overlap_wave_len].cpu().numpy(), overlap_wave_len)
+ generated_wave_chunks.append(output_wave)
+ previous_chunk = vc_wave[0, -overlap_wave_len:]
+ processed_frames += vc_target.size(2) - overlap_frame_len
+ output_wave = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave.tobytes(), frame_rate=sr,
+ sample_width=output_wave.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=bitrate).read()
+ yield mp3_bytes, None
+
+
+if __name__ == "__main__":
+ description = ("State-of-the-Art zero-shot voice conversion/singing voice conversion."
+ "for details and updates.
Note that any reference audio will be forcefully clipped to 25s if beyond this length.
"
+ "If total duration of source and reference audio exceeds 30s, source audio will be processed in chunks.
")
+ inputs = [
+ gr.Audio(type="filepath", label="Source Audio"),
+ gr.Audio(type="filepath", label="Reference Audio"),
+ gr.Slider(minimum=1, maximum=200, value=25, step=1, label="Diffusion Steps ", info="25 by default, 50~100 for best quality"),
+ gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.0, label="Length Adjust", info="<1.0 for speed-up speech, >1.0 for slow-down speech"),
+ gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.7, label="Inference CFG Rate", info="has subtle influence"),
+ gr.Checkbox(label="Use F0 conditioned model", value=False, info="Must set to true for singing voice conversion"),
+ gr.Checkbox(label="Auto F0 adjust", value=True,
+ info="Roughly adjust F0 to match target voice. Only works when F0 conditioned model is used."),
+ gr.Slider(label='Pitch shift', minimum=-24, maximum=24, step=1, value=0, info="Pitch shift in semitones, only works when F0 conditioned model is used"),
+ ]
+
+ examples = [["examples/source/yae_0.wav", "examples/reference/dingzhen_0.wav", 25, 1.0, 0.7, False, True, 0],
+ ["examples/source/jay_0.wav", "examples/reference/azuma_0.wav", 25, 1.0, 0.7, False, True, 0],
+ ["examples/source/Wiz Khalifa,Charlie Puth - See You Again [vocals]_[cut_28sec].wav",
+ "examples/reference/kobe_0.wav", 50, 1.0, 0.7, True, False, -6],
+ ["examples/source/TECHNOPOLIS - 2085 [vocals]_[cut_14sec].wav",
+ "examples/reference/trump_0.wav", 50, 1.0, 0.7, True, False, -12],
+ ]
+
+ outputs = [gr.Audio(label="Stream Output Audio", streaming=True, format='mp3'),
+ gr.Audio(label="Full Output Audio", streaming=False, format='wav')]
+
+ gr.Interface(fn=voice_conversion,
+ description=description,
+ inputs=inputs,
+ outputs=outputs,
+ title="Seed Voice Conversion",
+ examples=examples,
+ cache_examples=False,
+ ).launch()
\ No newline at end of file
diff --git a/app_v1v2.py b/app_v1v2.py
new file mode 100644
index 0000000000000000000000000000000000000000..262152c52a99bd36b1c0a60f26ffeedb22dc8545
--- /dev/null
+++ b/app_v1v2.py
@@ -0,0 +1,217 @@
+import spaces
+import gradio as gr
+import torch
+import yaml
+import argparse
+from seed_vc_wrapper import SeedVCWrapper
+from modules.v2.vc_wrapper import VoiceConversionWrapper
+
+# Set up device and torch configurations
+if torch.cuda.is_available():
+ device = torch.device("cuda")
+elif torch.backends.mps.is_available():
+ device = torch.device("mps")
+else:
+ device = torch.device("cpu")
+
+torch._inductor.config.coordinate_descent_tuning = True
+torch._inductor.config.triton.unique_kernel_names = True
+
+if hasattr(torch._inductor.config, "fx_graph_cache"):
+ # Experimental feature to reduce compilation times, will be on by default in future
+ torch._inductor.config.fx_graph_cache = True
+
+dtype = torch.float16
+
+
+def load_v2_models():
+ from hydra.utils import instantiate
+ from omegaconf import DictConfig
+ cfg = DictConfig(yaml.safe_load(open("configs/v2/vc_wrapper.yaml", "r")))
+ vc_wrapper = instantiate(cfg)
+ vc_wrapper.load_checkpoints()
+ vc_wrapper.to(device)
+ vc_wrapper.eval()
+
+ vc_wrapper.setup_ar_caches(max_batch_size=1, max_seq_len=4096, dtype=dtype, device=device)
+
+ return vc_wrapper
+
+# Global variables to store model instances
+vc_wrapper_v1 = SeedVCWrapper()
+vc_wrapper_v2 = load_v2_models()
+
+@spaces.GPU
+def convert_voice_v1_wrapper(source_audio_path, target_audio_path, diffusion_steps=10,
+ length_adjust=1.0, inference_cfg_rate=0.7, f0_condition=False,
+ auto_f0_adjust=True, pitch_shift=0, stream_output=True):
+ """
+ Wrapper function for vc_wrapper.convert_voice that can be decorated with @spaces.GPU
+ """
+
+ # Use yield from to properly handle the generator
+ yield from vc_wrapper_v1.convert_voice(
+ source=source_audio_path,
+ target=target_audio_path,
+ diffusion_steps=diffusion_steps,
+ length_adjust=length_adjust,
+ inference_cfg_rate=inference_cfg_rate,
+ f0_condition=f0_condition,
+ auto_f0_adjust=auto_f0_adjust,
+ pitch_shift=pitch_shift,
+ stream_output=stream_output
+ )
+
+@spaces.GPU
+def convert_voice_v2_wrapper(source_audio_path, target_audio_path, diffusion_steps=30,
+ length_adjust=1.0, intelligebility_cfg_rate=0.7, similarity_cfg_rate=0.7,
+ top_p=0.7, temperature=0.7, repetition_penalty=1.5,
+ convert_style=False, anonymization_only=False, stream_output=True):
+ """
+ Wrapper function for vc_wrapper.convert_voice_with_streaming that can be decorated with @spaces.GPU
+ """
+
+ # Use yield from to properly handle the generator
+ yield from vc_wrapper_v2.convert_voice_with_streaming(
+ source_audio_path=source_audio_path,
+ target_audio_path=target_audio_path,
+ diffusion_steps=diffusion_steps,
+ length_adjust=length_adjust,
+ intelligebility_cfg_rate=intelligebility_cfg_rate,
+ similarity_cfg_rate=similarity_cfg_rate,
+ top_p=top_p,
+ temperature=temperature,
+ repetition_penalty=repetition_penalty,
+ convert_style=convert_style,
+ anonymization_only=anonymization_only,
+ device=device,
+ dtype=dtype,
+ stream_output=stream_output
+ )
+
+
+def create_v1_interface():
+ # Set up Gradio interface
+ description = (
+ "Zero-shot voice conversion with in-context learning. "
+ "for details and updates.
Note that any reference audio will be forcefully clipped to 25s if beyond this length.
"
+ "If total duration of source and reference audio exceeds 30s, source audio will be processed in chunks.
")
+
+ inputs = [
+ gr.Audio(type="filepath", label="Source Audio"),
+ gr.Audio(type="filepath", label="Reference Audio"),
+ gr.Slider(minimum=1, maximum=200, value=10, step=1, label="Diffusion Steps",
+ info="10 by default, 50~100 for best quality"),
+ gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.0, label="Length Adjust",
+ info="<1.0 for speed-up speech, >1.0 for slow-down speech"),
+ gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.7, label="Inference CFG Rate",
+ info="has subtle influence"),
+ gr.Checkbox(label="Use F0 conditioned model", value=False,
+ info="Must set to true for singing voice conversion"),
+ gr.Checkbox(label="Auto F0 adjust", value=True,
+ info="Roughly adjust F0 to match target voice. Only works when F0 conditioned model is used."),
+ gr.Slider(label='Pitch shift', minimum=-24, maximum=24, step=1, value=0,
+ info="Pitch shift in semitones, only works when F0 conditioned model is used"),
+ ]
+
+ examples = [
+ ["examples/source/yae_0.wav", "examples/reference/dingzhen_0.wav", 25, 1.0, 0.7, False, True, 0],
+ ["examples/source/jay_0.wav", "examples/reference/azuma_0.wav", 25, 1.0, 0.7, True, True, 0],
+ ["examples/source/Wiz Khalifa,Charlie Puth - See You Again [vocals]_[cut_28sec].wav",
+ "examples/reference/teio_0.wav", 100, 1.0, 0.7, True, False, 0],
+ ["examples/source/TECHNOPOLIS - 2085 [vocals]_[cut_14sec].wav",
+ "examples/reference/trump_0.wav", 50, 1.0, 0.7, True, False, -12],
+ ]
+
+ outputs = [
+ gr.Audio(label="Stream Output Audio", streaming=True, format='mp3'),
+ gr.Audio(label="Full Output Audio", streaming=False, format='wav')
+ ]
+
+ return gr.Interface(
+ fn=convert_voice_v1_wrapper,
+ description=description,
+ inputs=inputs,
+ outputs=outputs,
+ title="Seed Voice Conversion V1 (Voice & Singing Voice Conversion)",
+ examples=examples,
+ cache_examples=False,
+ )
+
+
+def create_v2_interface():
+ # Set up Gradio interface
+ description = (
+ "Zero-shot voice/style conversion with in-context learning."
+ "for details and updates.
Note that any reference audio will be forcefully clipped to 25s if beyond this length.
"
+ "If total duration of source and reference audio exceeds 30s, source audio will be processed in chunks.
"
+ "Please click the 'convert style/emotion/accent' checkbox to convert the style, emotion, or accent of the source audio, or else only timbre conversion will be performed.
"
+ "Click the 'anonymization only' checkbox will ignore reference audio but convert source to an 'average voice' determined by model itself.
")
+ inputs = [
+ gr.Audio(type="filepath", label="Source Audio"),
+ gr.Audio(type="filepath", label="Reference Audio"),
+ gr.Slider(minimum=1, maximum=200, value=30, step=1, label="Diffusion Steps",
+ info="30 by default, 50~100 for best quality"),
+ gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.0, label="Length Adjust",
+ info="<1.0 for speed-up speech, >1.0 for slow-down speech"),
+ gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.0, label="Intelligibility CFG Rate",
+ info="controls pronunciation intelligibility"),
+ gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.7, label="Similarity CFG Rate",
+ info="controls similarity to reference audio"),
+ gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.9, label="Top-p",
+ info="AR model sampling top P"),
+ gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Temperature",
+ info="AR model sampling temperature"),
+ gr.Slider(minimum=1.0, maximum=3.0, step=0.1, value=1.0, label="Repetition Penalty",
+ info="AR model sampling repetition penalty"),
+ gr.Checkbox(label="convert style/emotion/accent", value=False),
+ gr.Checkbox(label="anonymization only", value=False),
+ ]
+
+ examples = [
+ ["examples/source/yae_0.wav", "examples/reference/dingzhen_0.wav", 25, 1.0, 0.7, 0.7, 0.9, 1.0, 1.0, True,
+ False],
+ ["examples/source/jay_0.wav", "examples/reference/azuma_0.wav", 25, 1.0, 0.7, 0.7, 0.9, 1.0, 1.0, True, False],
+ ]
+
+ outputs = [
+ gr.Audio(label="Stream Output Audio", streaming=True, format='mp3'),
+ gr.Audio(label="Full Output Audio", streaming=False, format='wav')
+ ]
+
+ return gr.Interface(
+ fn=convert_voice_v2_wrapper,
+ description=description,
+ inputs=inputs,
+ outputs=outputs,
+ title="Seed Voice Conversion V2 (Voice & Style Conversion)",
+ examples=examples,
+ cache_examples=False,
+ )
+
+
+def main(args):
+ # Create interfaces
+ v1_interface = create_v1_interface()
+ v2_interface = create_v2_interface()
+
+ # Create tabs
+ with gr.Blocks(title="Seed Voice Conversion") as demo:
+ gr.Markdown("# Seed Voice Conversion")
+ gr.Markdown("Choose between V1 (Voice & Singing Voice Conversion) or V2 (Voice & Style Conversion)")
+
+ with gr.Tabs():
+ with gr.TabItem("V2 - Voice & Style Conversion"):
+ v2_interface.render()
+ with gr.TabItem("V1 - Voice & Singing Voice Conversion"):
+ v1_interface.render()
+
+ # Launch the combined interface
+ demo.launch()
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--compile", type=bool, default=True)
+ args = parser.parse_args()
+ main(args)
\ No newline at end of file
diff --git a/campplus_cn_common.bin b/campplus_cn_common.bin
new file mode 100644
index 0000000000000000000000000000000000000000..ccc729ec7f373f8d2351be91c9efdc7b21c00384
--- /dev/null
+++ b/campplus_cn_common.bin
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3388cf5fd3493c9ac9c69851d8e7a8badcfb4f3dc631020c4961371646d5ada8
+size 28036335
diff --git a/configs/astral_quantization/default_2048.yml b/configs/astral_quantization/default_2048.yml
new file mode 100644
index 0000000000000000000000000000000000000000..54f91e7cea7722a8cdb85c9855bffeeedb84e5db
--- /dev/null
+++ b/configs/astral_quantization/default_2048.yml
@@ -0,0 +1,40 @@
+_target_: modules.astral_quantization.default_model.AstralQuantizer
+tokenizer_name: "openai/whisper-small"
+ssl_model_name: "facebook/hubert-large-ll60k"
+ssl_output_layer: 18
+encoder:
+ _target_: modules.astral_quantization.convnext.ConvNeXtV2Stage
+ dim: 512
+ num_blocks: 12
+ intermediate_dim: 1536
+ dilation: 1
+ input_dim: 1024
+quantizer:
+ _target_: modules.astral_quantization.bsq.BinarySphericalQuantize
+ codebook_size: 2048 # codebook size, must be a power of 2
+ dim: 512
+ entropy_loss_weight: 0.1
+ diversity_gamma: 1.0
+ spherical: True
+ enable_entropy_loss: True
+ soft_entropy_loss: True
+decoder:
+ _target_: modules.astral_quantization.convnext.ConvNeXtV2Stage
+ dim: 512
+ num_blocks: 12
+ intermediate_dim: 1536
+ dilation: 1
+ output_dim: 1024
+ gin_channels: 192
+asr_decoder:
+ _target_: modules.astral_quantization.asr_decoder.ASRDecoder
+ hidden_dim: 768
+ num_heads: 12
+ depth: 12
+ block_size: 4096
+ in_channels: 512
+ n_vocab: 51866
+ bos_id: 50528
+ eos_id: 50527
+ dropout_rate: 0.0
+ attn_dropout_rate: 0.0
\ No newline at end of file
diff --git a/configs/astral_quantization/default_32.yml b/configs/astral_quantization/default_32.yml
new file mode 100644
index 0000000000000000000000000000000000000000..bf160129893fb893eed26eabcb8a2da9c42a7159
--- /dev/null
+++ b/configs/astral_quantization/default_32.yml
@@ -0,0 +1,40 @@
+_target_: default_model.AstralQuantizer
+tokenizer_name: "openai/whisper-small"
+ssl_model_name: "facebook/hubert-large-ll60k"
+ssl_output_layer: 18
+encoder:
+ _target_: modules.convnext.ConvNeXtV2Stage
+ dim: 512
+ num_blocks: 12
+ intermediate_dim: 1536
+ dilation: 1
+ input_dim: 1024
+quantizer:
+ _target_: modules.bsq.BinarySphericalQuantize
+ codebook_size: 32 # codebook size, must be a power of 2
+ dim: 512
+ entropy_loss_weight: 0.1
+ diversity_gamma: 1.0
+ spherical: True
+ enable_entropy_loss: True
+ soft_entropy_loss: True
+decoder:
+ _target_: modules.convnext.ConvNeXtV2Stage
+ dim: 512
+ num_blocks: 12
+ intermediate_dim: 1536
+ dilation: 1
+ output_dim: 1024
+ gin_channels: 192
+asr_decoder:
+ _target_: modules.asr_decoder.ASRDecoder
+ hidden_dim: 768
+ num_heads: 12
+ depth: 12
+ block_size: 4096
+ in_channels: 512
+ n_vocab: 51866
+ bos_id: 50528
+ eos_id: 50527
+ dropout_rate: 0.0
+ attn_dropout_rate: 0.0
\ No newline at end of file
diff --git a/configs/config.json b/configs/config.json
new file mode 100644
index 0000000000000000000000000000000000000000..e74f0b4898f6e47e1d198b62cdba989784ce2bb0
--- /dev/null
+++ b/configs/config.json
@@ -0,0 +1 @@
+{"reference_audio_path": "D:/FAcodec/test_waves/kobe_0.wav", "sg_hostapi": "MME", "sg_wasapi_exclusive": false, "sg_input_device": "\u9ea6\u514b\u98ce (Razer BlackShark V2 HS 2.4", "sg_output_device": "\u626c\u58f0\u5668 (Razer BlackShark V2 HS 2.4", "sr_type": "sr_model", "diffusion_steps": 10.0, "inference_cfg_rate": 0.0, "max_prompt_length": 3.0, "block_time": 0.7, "crossfade_length": 0.04, "extra_time": 0.5, "extra_time_right": 0.02}
\ No newline at end of file
diff --git a/configs/config_dit_mel_seed.yml b/configs/config_dit_mel_seed.yml
new file mode 100644
index 0000000000000000000000000000000000000000..84a72160afdf25c0d62319314f275edde4fd8ef9
--- /dev/null
+++ b/configs/config_dit_mel_seed.yml
@@ -0,0 +1,79 @@
+log_dir: "./runs/run_dit_mel_seed"
+save_freq: 1
+log_interval: 10
+save_interval: 1000
+device: "cuda"
+epochs: 1000 # number of epochs for first stage training (pre-training)
+batch_size: 4
+batch_length: 100 # maximum duration of audio in a batch (in seconds)
+max_len: 80 # maximum number of frames
+pretrained_model: ""
+pretrained_encoder: ""
+load_only_params: False # set to true if do not want to load epoch numbers and optimizer parameters
+
+F0_path: "modules/JDC/bst.t7"
+
+preprocess_params:
+ sr: 22050
+ spect_params:
+ n_fft: 1024
+ win_length: 1024
+ hop_length: 256
+ n_mels: 80
+
+model_params:
+ dit_type: "DiT" # uDiT or DiT
+ reg_loss_type: "l2" # l1 or l2
+
+ speech_tokenizer:
+ path: "speech_tokenizer_v1.onnx"
+
+ style_encoder:
+ dim: 192
+ campplus_path: "campplus_cn_common.bin"
+
+ DAC:
+ encoder_dim: 64
+ encoder_rates: [2, 5, 5, 6]
+ decoder_dim: 1536
+ decoder_rates: [ 6, 5, 5, 2 ]
+ sr: 24000
+
+ length_regulator:
+ channels: 768
+ is_discrete: true
+ content_codebook_size: 4096
+ in_frame_rate: 50
+ out_frame_rate: 80
+ sampling_ratios: [1, 1, 1, 1]
+
+ DiT:
+ hidden_dim: 768
+ num_heads: 12
+ depth: 12
+ class_dropout_prob: 0.1
+ block_size: 4096
+ in_channels: 80
+ style_condition: true
+ final_layer_type: 'wavenet'
+ target: 'mel' # mel or codec
+ content_dim: 768
+ content_codebook_size: 1024
+ content_type: 'discrete'
+ f0_condition: false
+ n_f0_bins: 512
+ content_codebooks: 1
+ is_causal: false
+ long_skip_connection: true
+ zero_prompt_speech_token: false # for prompt component, do not input corresponding speech token
+
+ wavenet:
+ hidden_dim: 768
+ num_layers: 8
+ kernel_size: 5
+ dilation_rate: 1
+ p_dropout: 0.2
+ style_condition: true
+
+loss_params:
+ base_lr: 0.0001
\ No newline at end of file
diff --git a/configs/config_dit_mel_seed_facodec_small.yml b/configs/config_dit_mel_seed_facodec_small.yml
new file mode 100644
index 0000000000000000000000000000000000000000..eeb0690e56d6c7bf9b46bc9394a883825931ae06
--- /dev/null
+++ b/configs/config_dit_mel_seed_facodec_small.yml
@@ -0,0 +1,97 @@
+log_dir: "./runs/run_dit_mel_seed_facodec_small"
+save_freq: 1
+log_interval: 10
+save_interval: 1000
+device: "cuda"
+epochs: 1000 # number of epochs for first stage training (pre-training)
+batch_size: 2
+batch_length: 100 # maximum duration of audio in a batch (in seconds)
+max_len: 80 # maximum number of frames
+pretrained_model: ""
+pretrained_encoder: ""
+load_only_params: False # set to true if do not want to load epoch numbers and optimizer parameters
+
+F0_path: "modules/JDC/bst.t7"
+
+data_params:
+ train_data: "./data/train.txt"
+ val_data: "./data/val.txt"
+ root_path: "./data/"
+
+preprocess_params:
+ sr: 22050
+ spect_params:
+ n_fft: 1024
+ win_length: 1024
+ hop_length: 256
+ n_mels: 80
+
+model_params:
+ dit_type: "DiT" # uDiT or DiT
+ reg_loss_type: "l1" # l1 or l2
+
+ speech_tokenizer:
+ type: 'facodec' # facodec or cosyvoice
+ path: "checkpoints/speech_tokenizer_v1.onnx"
+
+ style_encoder:
+ dim: 192
+ campplus_path: "checkpoints/campplus_cn_common.bin"
+
+ DAC:
+ encoder_dim: 64
+ encoder_rates: [2, 5, 5, 6]
+ decoder_dim: 1536
+ decoder_rates: [ 6, 5, 5, 2 ]
+ sr: 24000
+
+ length_regulator:
+ channels: 512
+ is_discrete: true
+ content_codebook_size: 1024
+ in_frame_rate: 80
+ out_frame_rate: 80
+ sampling_ratios: [1, 1, 1, 1]
+ token_dropout_prob: 0.3 # probability of performing token dropout
+ token_dropout_range: 1.0 # maximum percentage of tokens to drop out
+ n_codebooks: 3
+ quantizer_dropout: 0.5
+ f0_condition: false
+ n_f0_bins: 512
+
+ DiT:
+ hidden_dim: 512
+ num_heads: 8
+ depth: 13
+ class_dropout_prob: 0.1
+ block_size: 8192
+ in_channels: 80
+ style_condition: true
+ final_layer_type: 'wavenet'
+ target: 'mel' # mel or codec
+ content_dim: 512
+ content_codebook_size: 1024
+ content_type: 'discrete'
+ f0_condition: true
+ n_f0_bins: 512
+ content_codebooks: 1
+ is_causal: false
+ long_skip_connection: true
+ zero_prompt_speech_token: false # for prompt component, do not input corresponding speech token
+ time_as_token: false
+ style_as_token: false
+ uvit_skip_connection: true
+ add_resblock_in_transformer: false
+
+ wavenet:
+ hidden_dim: 512
+ num_layers: 8
+ kernel_size: 5
+ dilation_rate: 1
+ p_dropout: 0.2
+ style_condition: true
+
+loss_params:
+ base_lr: 0.0001
+ lambda_mel: 45
+ lambda_kl: 1.0
\ No newline at end of file
diff --git a/configs/config_dit_mel_seed_wavenet.yml b/configs/config_dit_mel_seed_wavenet.yml
new file mode 100644
index 0000000000000000000000000000000000000000..6d8b8aab93162d30b23f8d0dc5634d29523dee43
--- /dev/null
+++ b/configs/config_dit_mel_seed_wavenet.yml
@@ -0,0 +1,79 @@
+log_dir: "./runs/run_dit_mel_seed"
+save_freq: 1
+log_interval: 10
+save_interval: 1000
+device: "cuda"
+epochs: 1000 # number of epochs for first stage training (pre-training)
+batch_size: 4
+batch_length: 100 # maximum duration of audio in a batch (in seconds)
+max_len: 80 # maximum number of frames
+pretrained_model: ""
+pretrained_encoder: ""
+load_only_params: False # set to true if do not want to load epoch numbers and optimizer parameters
+
+F0_path: "modules/JDC/bst.t7"
+
+preprocess_params:
+ sr: 22050
+ spect_params:
+ n_fft: 1024
+ win_length: 1024
+ hop_length: 256
+ n_mels: 80
+
+model_params:
+ dit_type: "DiT" # uDiT or DiT
+ reg_loss_type: "l2" # l1 or l2
+
+ speech_tokenizer:
+ path: "checkpoints/speech_tokenizer_v1.onnx"
+
+ style_encoder:
+ dim: 192
+ campplus_path: "campplus_cn_common.bin"
+
+ DAC:
+ encoder_dim: 64
+ encoder_rates: [2, 5, 5, 6]
+ decoder_dim: 1536
+ decoder_rates: [ 6, 5, 5, 2 ]
+ sr: 24000
+
+ length_regulator:
+ channels: 768
+ is_discrete: true
+ content_codebook_size: 4096
+ in_frame_rate: 50
+ out_frame_rate: 80
+ sampling_ratios: [1, 1, 1, 1]
+
+ DiT:
+ hidden_dim: 768
+ num_heads: 12
+ depth: 12
+ class_dropout_prob: 0.1
+ block_size: 8192
+ in_channels: 80
+ style_condition: true
+ final_layer_type: 'wavenet'
+ target: 'mel' # mel or codec
+ content_dim: 768
+ content_codebook_size: 1024
+ content_type: 'discrete'
+ f0_condition: false
+ n_f0_bins: 512
+ content_codebooks: 1
+ is_causal: false
+ long_skip_connection: true
+ zero_prompt_speech_token: false # for prompt component, do not input corresponding speech token
+
+ wavenet:
+ hidden_dim: 768
+ num_layers: 8
+ kernel_size: 5
+ dilation_rate: 1
+ p_dropout: 0.2
+ style_condition: true
+
+loss_params:
+ base_lr: 0.0001
diff --git a/configs/hifigan.yml b/configs/hifigan.yml
new file mode 100644
index 0000000000000000000000000000000000000000..eaddafa313af0adf24ddd5cdec3bd457fcbbc274
--- /dev/null
+++ b/configs/hifigan.yml
@@ -0,0 +1,25 @@
+hift:
+ in_channels: 80
+ base_channels: 512
+ nb_harmonics: 8
+ sampling_rate: 22050
+ nsf_alpha: 0.1
+ nsf_sigma: 0.003
+ nsf_voiced_threshold: 10
+ upsample_rates: [8, 8]
+ upsample_kernel_sizes: [16, 16]
+ istft_params:
+ n_fft: 16
+ hop_len: 4
+ resblock_kernel_sizes: [3, 7, 11]
+ resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
+ source_resblock_kernel_sizes: [7, 11]
+ source_resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5]]
+ lrelu_slope: 0.1
+ audio_limit: 0.99
+f0_predictor:
+ num_class: 1
+ in_channels: 80
+ cond_channels: 512
+
+pretrained_model_path: "checkpoints/hift.pt"
diff --git a/configs/inuse/.gitignore b/configs/inuse/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/configs/inuse/config.json b/configs/inuse/config.json
new file mode 100644
index 0000000000000000000000000000000000000000..17d2df311cd2b19c1888c8fbb82effbfbc6edef3
--- /dev/null
+++ b/configs/inuse/config.json
@@ -0,0 +1 @@
+{"reference_audio_path": "D:/seed-vc/examples/reference/trump_0.wav", "sg_hostapi": "MME", "sg_wasapi_exclusive": false, "sg_input_device": "\u9ea6\u514b\u98ce (Razer BlackShark V2 HS USB", "sg_output_device": "\u626c\u58f0\u5668 (Razer BlackShark V2 HS USB", "sr_type": "sr_model", "diffusion_steps": 8.0, "inference_cfg_rate": 0.7, "max_prompt_length": 3.0, "block_time": 0.58, "crossfade_length": 0.04, "extra_time_ce": 2.5, "extra_time": 0.5, "extra_time_right": 0.02}
\ No newline at end of file
diff --git a/configs/presets/config_dit_mel_seed_uvit_whisper_base_f0_44k.yml b/configs/presets/config_dit_mel_seed_uvit_whisper_base_f0_44k.yml
new file mode 100644
index 0000000000000000000000000000000000000000..0ec7ef4f6d8c7cf160687747ae01e0d39f6c128d
--- /dev/null
+++ b/configs/presets/config_dit_mel_seed_uvit_whisper_base_f0_44k.yml
@@ -0,0 +1,98 @@
+log_dir: "./runs"
+save_freq: 1
+log_interval: 10
+save_interval: 1000
+device: "cuda"
+epochs: 1000 # number of epochs for first stage training (pre-training)
+batch_size: 1
+batch_length: 100 # maximum duration of audio in a batch (in seconds)
+max_len: 80 # maximum number of frames
+pretrained_model: "DiT_seed_v2_uvit_whisper_base_f0_44k_bigvgan_pruned_ft_ema.pth"
+pretrained_encoder: ""
+load_only_params: False # set to true if do not want to load epoch numbers and optimizer parameters
+
+preprocess_params:
+ sr: 44100
+ spect_params:
+ n_fft: 2048
+ win_length: 2048
+ hop_length: 512
+ n_mels: 128
+ fmin: 0
+ fmax: "None"
+
+model_params:
+ dit_type: "DiT" # uDiT or DiT
+ reg_loss_type: "l1" # l1 or l2
+
+ timbre_shifter:
+ se_db_path: "./modules/openvoice/checkpoints_v2/converter/se_db.pt"
+ ckpt_path: './modules/openvoice/checkpoints_v2/converter'
+
+ vocoder:
+ type: "bigvgan"
+ name: "nvidia/bigvgan_v2_44khz_128band_512x"
+
+ speech_tokenizer:
+ type: 'whisper'
+ name: "openai/whisper-small"
+
+ style_encoder:
+ dim: 192
+ campplus_path: "campplus_cn_common.bin"
+
+ DAC:
+ encoder_dim: 64
+ encoder_rates: [2, 5, 5, 6]
+ decoder_dim: 1536
+ decoder_rates: [ 6, 5, 5, 2 ]
+ sr: 24000
+
+ length_regulator:
+ channels: 768
+ is_discrete: false
+ in_channels: 768
+ content_codebook_size: 2048
+ sampling_ratios: [1, 1, 1, 1]
+ vector_quantize: false
+ n_codebooks: 1
+ quantizer_dropout: 0.0
+ f0_condition: true
+ n_f0_bins: 256
+
+ DiT:
+ hidden_dim: 768
+ num_heads: 12
+ depth: 17
+ class_dropout_prob: 0.1
+ block_size: 8192
+ in_channels: 128
+ style_condition: true
+ final_layer_type: 'mlp'
+ target: 'mel' # mel or codec
+ content_dim: 768
+ content_codebook_size: 1024
+ content_type: 'discrete'
+ f0_condition: true
+ n_f0_bins: 256
+ content_codebooks: 1
+ is_causal: false
+ long_skip_connection: false
+ zero_prompt_speech_token: false # for prompt component, do not input corresponding speech token
+ time_as_token: false
+ style_as_token: false
+ uvit_skip_connection: true
+ add_resblock_in_transformer: false
+
+ wavenet:
+ hidden_dim: 768
+ num_layers: 8
+ kernel_size: 5
+ dilation_rate: 1
+ p_dropout: 0.2
+ style_condition: true
+
+loss_params:
+ base_lr: 0.0001
+ lambda_mel: 45
+ lambda_kl: 1.0
\ No newline at end of file
diff --git a/configs/presets/config_dit_mel_seed_uvit_whisper_small_wavenet.yml b/configs/presets/config_dit_mel_seed_uvit_whisper_small_wavenet.yml
new file mode 100644
index 0000000000000000000000000000000000000000..492910d163c7773c64f846ee55384e3e8b81ac00
--- /dev/null
+++ b/configs/presets/config_dit_mel_seed_uvit_whisper_small_wavenet.yml
@@ -0,0 +1,91 @@
+log_dir: "./runs"
+save_freq: 1
+log_interval: 10
+save_interval: 1000
+device: "cuda"
+epochs: 1000 # number of epochs for first stage training (pre-training)
+batch_size: 2
+batch_length: 100 # maximum duration of audio in a batch (in seconds)
+max_len: 80 # maximum number of frames
+pretrained_model: "DiT_seed_v2_uvit_whisper_small_wavenet_bigvgan_pruned.pth"
+pretrained_encoder: ""
+load_only_params: False # set to true if do not want to load epoch numbers and optimizer parameters
+
+preprocess_params:
+ sr: 22050
+ spect_params:
+ n_fft: 1024
+ win_length: 1024
+ hop_length: 256
+ n_mels: 80
+ fmin: 0
+ fmax: "None"
+
+model_params:
+ dit_type: "DiT" # uDiT or DiT
+ reg_loss_type: "l1" # l1 or l2
+
+ timbre_shifter:
+ se_db_path: "./modules/openvoice/checkpoints_v2/converter/se_db.pt"
+ ckpt_path: './modules/openvoice/checkpoints_v2/converter'
+
+ speech_tokenizer:
+ type: 'whisper'
+ name: "openai/whisper-small"
+
+ style_encoder:
+ dim: 192
+ campplus_path: "campplus_cn_common.bin"
+
+ vocoder:
+ type: "bigvgan"
+ name: "nvidia/bigvgan_v2_22khz_80band_256x"
+
+ length_regulator:
+ channels: 512
+ is_discrete: false
+ in_channels: 768
+ content_codebook_size: 2048
+ sampling_ratios: [1, 1, 1, 1]
+ vector_quantize: false
+ n_codebooks: 1
+ quantizer_dropout: 0.0
+ f0_condition: false
+ n_f0_bins: 512
+
+ DiT:
+ hidden_dim: 512
+ num_heads: 8
+ depth: 13
+ class_dropout_prob: 0.1
+ block_size: 8192
+ in_channels: 80
+ style_condition: true
+ final_layer_type: 'wavenet'
+ target: 'mel' # mel or codec
+ content_dim: 512
+ content_codebook_size: 1024
+ content_type: 'discrete'
+ f0_condition: false
+ n_f0_bins: 512
+ content_codebooks: 1
+ is_causal: false
+ long_skip_connection: true
+ zero_prompt_speech_token: false # for prompt component, do not input corresponding speech token
+ time_as_token: false
+ style_as_token: false
+ uvit_skip_connection: true
+ add_resblock_in_transformer: false
+
+ wavenet:
+ hidden_dim: 512
+ num_layers: 8
+ kernel_size: 5
+ dilation_rate: 1
+ p_dropout: 0.2
+ style_condition: true
+
+loss_params:
+ base_lr: 0.0001
+ lambda_mel: 45
+ lambda_kl: 1.0
\ No newline at end of file
diff --git a/configs/presets/config_dit_mel_seed_uvit_xlsr_tiny.yml b/configs/presets/config_dit_mel_seed_uvit_xlsr_tiny.yml
new file mode 100644
index 0000000000000000000000000000000000000000..e0677397377158dd30ffdf905946fbd297b36bd5
--- /dev/null
+++ b/configs/presets/config_dit_mel_seed_uvit_xlsr_tiny.yml
@@ -0,0 +1,82 @@
+log_dir: "./runs/"
+save_freq: 1
+log_interval: 10
+save_interval: 500
+device: "cuda"
+epochs: 1000 # number of epochs for first stage training (pre-training)
+batch_size: 2
+batch_length: 100 # maximum duration of audio in a batch (in seconds)
+max_len: 80 # maximum number of frames
+pretrained_model: "DiT_uvit_tat_xlsr_ema.pth"
+pretrained_encoder: ""
+load_only_params: False # set to true if do not want to load epoch numbers and optimizer parameters
+
+preprocess_params:
+ sr: 22050
+ spect_params:
+ n_fft: 1024
+ win_length: 1024
+ hop_length: 256
+ n_mels: 80
+ fmin: 0
+ fmax: 8000
+
+model_params:
+ dit_type: "DiT" # uDiT or DiT
+ reg_loss_type: "l1" # l1 or l2
+ diffusion_type: "flow"
+
+ timbre_shifter:
+ se_db_path: "./modules/openvoice/checkpoints_v2/converter/se_db.pt"
+ ckpt_path: './modules/openvoice/checkpoints_v2/converter'
+
+ vocoder:
+ type: "hifigan"
+
+ speech_tokenizer:
+ type: 'xlsr'
+ output_layer: 12
+ name: 'facebook/wav2vec2-xls-r-300m'
+
+ style_encoder:
+ dim: 192
+ campplus_path: "campplus_cn_common.bin"
+
+ length_regulator:
+ channels: 384
+ is_discrete: false
+ in_channels: 1024
+ content_codebook_size: 1024
+ sampling_ratios: [1, 1, 1, 1]
+ vector_quantize: false
+ n_codebooks: 2
+ quantizer_dropout: 0.0
+ f0_condition: false
+ n_f0_bins: 512
+
+ DiT:
+ hidden_dim: 384
+ num_heads: 6
+ depth: 9
+ class_dropout_prob: 0.1
+ block_size: 8192
+ in_channels: 80
+ style_condition: true
+ final_layer_type: 'mlp'
+ target: 'mel' # mel or betavae
+ content_dim: 384
+ content_codebook_size: 1024
+ content_type: 'discrete'
+ f0_condition: false
+ n_f0_bins: 512
+ content_codebooks: 1
+ is_causal: false
+ long_skip_connection: false
+ zero_prompt_speech_token: false # for prompt component, do not input corresponding speech token
+ time_as_token: true
+ style_as_token: true
+ uvit_skip_connection: true
+ add_resblock_in_transformer: false
+
+loss_params:
+ base_lr: 0.0001
\ No newline at end of file
diff --git a/configs/v2/ar_base.yaml b/configs/v2/ar_base.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/configs/v2/dit_small.yaml b/configs/v2/dit_small.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..60b93e953d7ec19039af147b7a52002b757a822a
--- /dev/null
+++ b/configs/v2/dit_small.yaml
@@ -0,0 +1,17 @@
+_target_: modules.v2.cfm.CFM
+estimator:
+ _target_: modules.v2.dit_wrapper.DiT
+ time_as_token: true
+ style_as_token: true
+ uvit_skip_connection: false
+ block_size: 8192
+ depth: 13
+ num_heads: 8
+ hidden_dim: 512
+ in_channels: 80
+ content_dim: 512
+ style_encoder_dim: 192
+ class_dropout_prob: 0.1
+ dropout_rate: 0.0
+ attn_dropout_rate: 0.0
+
diff --git a/configs/v2/vc_wrapper.yaml b/configs/v2/vc_wrapper.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..c3fe5b84431f53ebd12ec60663f40f61f4a8c231
--- /dev/null
+++ b/configs/v2/vc_wrapper.yaml
@@ -0,0 +1,105 @@
+_target_: modules.v2.vc_wrapper.VoiceConversionWrapper
+sr: 22050
+hop_size: 256
+mel_fn:
+ _target_: modules.audio.mel_spectrogram
+ _partial_: true
+ n_fft: 1024
+ win_size: 1024
+ hop_size: 256
+ num_mels: 80
+ sampling_rate: 22050
+ fmin: 0
+ fmax: null
+ center: False
+cfm:
+ _target_: modules.v2.cfm.CFM
+ estimator:
+ _target_: modules.v2.dit_wrapper.DiT
+ time_as_token: true
+ style_as_token: true
+ uvit_skip_connection: false
+ block_size: 8192
+ depth: 13
+ num_heads: 8
+ hidden_dim: 512
+ in_channels: 80
+ content_dim: 512
+ style_encoder_dim: 192
+ class_dropout_prob: 0.1
+ dropout_rate: 0.0
+ attn_dropout_rate: 0.0
+cfm_length_regulator:
+ _target_: modules.v2.length_regulator.InterpolateRegulator
+ channels: 512
+ is_discrete: true
+ codebook_size: 2048
+ sampling_ratios: [ 1, 1, 1, 1 ]
+ f0_condition: false
+ar:
+ _target_: modules.v2.ar.NaiveWrapper
+ model:
+ _target_: modules.v2.ar.NaiveTransformer
+ config:
+ _target_: modules.v2.ar.NaiveModelArgs
+ dropout: 0.0
+ rope_base: 10000.0
+ dim: 768
+ head_dim: 64
+ n_local_heads: 2
+ intermediate_size: 2304
+ n_head: 12
+ n_layer: 12
+ vocab_size: 2049 # 1 + 1 for eos
+ar_length_regulator:
+ _target_: modules.v2.length_regulator.InterpolateRegulator
+ channels: 768
+ is_discrete: true
+ codebook_size: 32
+ sampling_ratios: [ ]
+ f0_condition: false
+style_encoder:
+ _target_: modules.campplus.DTDNN.CAMPPlus
+ feat_dim: 80
+ embedding_size: 192
+content_extractor_narrow:
+ _target_: modules.astral_quantization.default_model.AstralQuantizer
+ tokenizer_name: "openai/whisper-small"
+ ssl_model_name: "facebook/hubert-large-ll60k"
+ ssl_output_layer: 18
+ skip_ssl: true
+ encoder: &bottleneck_encoder
+ _target_: modules.astral_quantization.convnext.ConvNeXtV2Stage
+ dim: 512
+ num_blocks: 12
+ intermediate_dim: 1536
+ dilation: 1
+ input_dim: 1024
+ quantizer:
+ _target_: modules.astral_quantization.bsq.BinarySphericalQuantize
+ codebook_size: 32 # codebook size, must be a power of 2
+ dim: 512
+ entropy_loss_weight: 0.1
+ diversity_gamma: 1.0
+ spherical: True
+ enable_entropy_loss: True
+ soft_entropy_loss: True
+content_extractor_wide:
+ _target_: modules.astral_quantization.default_model.AstralQuantizer
+ tokenizer_name: "openai/whisper-small"
+ ssl_model_name: "facebook/hubert-large-ll60k"
+ ssl_output_layer: 18
+ encoder: *bottleneck_encoder
+ quantizer:
+ _target_: modules.astral_quantization.bsq.BinarySphericalQuantize
+ codebook_size: 2048 # codebook size, must be a power of 2
+ dim: 512
+ entropy_loss_weight: 0.1
+ diversity_gamma: 1.0
+ spherical: True
+ enable_entropy_loss: True
+ soft_entropy_loss: True
+vocoder:
+ _target_: modules.bigvgan.bigvgan.BigVGAN.from_pretrained
+ pretrained_model_name_or_path: "nvidia/bigvgan_v2_22khz_80band_256x"
+ use_cuda_kernel: false
diff --git a/dac/__init__.py b/dac/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a5d03388fa63486960c783ebe7f1bd411b95b1d
--- /dev/null
+++ b/dac/__init__.py
@@ -0,0 +1,16 @@
+__version__ = "1.0.0"
+
+# preserved here for legacy reasons
+__model_version__ = "latest"
+
+import audiotools
+
+audiotools.ml.BaseModel.INTERN += ["dac.**"]
+audiotools.ml.BaseModel.EXTERN += ["einops"]
+
+
+from . import nn
+from . import model
+from . import utils
+from .model import DAC
+from .model import DACFile
diff --git a/dac/__main__.py b/dac/__main__.py
new file mode 100644
index 0000000000000000000000000000000000000000..393698e7da671bce1478b4d78b2f685d2640636b
--- /dev/null
+++ b/dac/__main__.py
@@ -0,0 +1,36 @@
+import sys
+
+import argbind
+
+from dac.utils import download
+from dac.utils.decode import decode
+from dac.utils.encode import encode
+
+STAGES = ["encode", "decode", "download"]
+
+
+def run(stage: str):
+ """Run stages.
+
+ Parameters
+ ----------
+ stage : str
+ Stage to run
+ """
+ if stage not in STAGES:
+ raise ValueError(f"Unknown command: {stage}. Allowed commands are {STAGES}")
+ stage_fn = globals()[stage]
+
+ if stage == "download":
+ stage_fn()
+ return
+
+ stage_fn()
+
+
+if __name__ == "__main__":
+ group = sys.argv.pop(1)
+ args = argbind.parse_args(group=group)
+
+ with argbind.scope(args):
+ run(group)
diff --git a/dac/__pycache__/__init__.cpython-310.pyc b/dac/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f2fd16c8e7911de0771a9fb02dfea1deb03a7a6d
Binary files /dev/null and b/dac/__pycache__/__init__.cpython-310.pyc differ
diff --git a/dac/model/__init__.py b/dac/model/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..58b47475d39e4249d2cd45577503bb68ffdacd00
--- /dev/null
+++ b/dac/model/__init__.py
@@ -0,0 +1,4 @@
+from .base import CodecMixin
+from .base import DACFile
+from .dac import DAC
+from .discriminator import Discriminator
diff --git a/dac/model/__pycache__/__init__.cpython-310.pyc b/dac/model/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f14e5f53c061ea4f23f4d54f78b95c89d53ffc75
Binary files /dev/null and b/dac/model/__pycache__/__init__.cpython-310.pyc differ
diff --git a/dac/model/__pycache__/base.cpython-310.pyc b/dac/model/__pycache__/base.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b7674799655ce4c06ea3dc15789b9bb9269f9247
Binary files /dev/null and b/dac/model/__pycache__/base.cpython-310.pyc differ
diff --git a/dac/model/__pycache__/dac.cpython-310.pyc b/dac/model/__pycache__/dac.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..df37f2e2bf777a88ec15b95aedcb7ab27553de98
Binary files /dev/null and b/dac/model/__pycache__/dac.cpython-310.pyc differ
diff --git a/dac/model/__pycache__/discriminator.cpython-310.pyc b/dac/model/__pycache__/discriminator.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d8f1a372d7c48a502389a8148adc06db1b8d3f5f
Binary files /dev/null and b/dac/model/__pycache__/discriminator.cpython-310.pyc differ
diff --git a/dac/model/__pycache__/encodec.cpython-310.pyc b/dac/model/__pycache__/encodec.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9b89f08e09105cac6ffd2ab25f62f27b29641ac3
Binary files /dev/null and b/dac/model/__pycache__/encodec.cpython-310.pyc differ
diff --git a/dac/model/base.py b/dac/model/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..2ef5a44074ca2bd5d726fe90fa7c8c87da1b3b7a
--- /dev/null
+++ b/dac/model/base.py
@@ -0,0 +1,294 @@
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+import torch
+import tqdm
+from audiotools import AudioSignal
+from torch import nn
+
+SUPPORTED_VERSIONS = ["1.0.0"]
+
+
+@dataclass
+class DACFile:
+ codes: torch.Tensor
+
+ # Metadata
+ chunk_length: int
+ original_length: int
+ input_db: float
+ channels: int
+ sample_rate: int
+ padding: bool
+ dac_version: str
+
+ def save(self, path):
+ artifacts = {
+ "codes": self.codes.numpy().astype(np.uint16),
+ "metadata": {
+ "input_db": self.input_db.numpy().astype(np.float32),
+ "original_length": self.original_length,
+ "sample_rate": self.sample_rate,
+ "chunk_length": self.chunk_length,
+ "channels": self.channels,
+ "padding": self.padding,
+ "dac_version": SUPPORTED_VERSIONS[-1],
+ },
+ }
+ path = Path(path).with_suffix(".dac")
+ with open(path, "wb") as f:
+ np.save(f, artifacts)
+ return path
+
+ @classmethod
+ def load(cls, path):
+ artifacts = np.load(path, allow_pickle=True)[()]
+ codes = torch.from_numpy(artifacts["codes"].astype(int))
+ if artifacts["metadata"].get("dac_version", None) not in SUPPORTED_VERSIONS:
+ raise RuntimeError(
+ f"Given file {path} can't be loaded with this version of descript-audio-codec."
+ )
+ return cls(codes=codes, **artifacts["metadata"])
+
+
+class CodecMixin:
+ @property
+ def padding(self):
+ if not hasattr(self, "_padding"):
+ self._padding = True
+ return self._padding
+
+ @padding.setter
+ def padding(self, value):
+ assert isinstance(value, bool)
+
+ layers = [
+ l for l in self.modules() if isinstance(l, (nn.Conv1d, nn.ConvTranspose1d))
+ ]
+
+ for layer in layers:
+ if value:
+ if hasattr(layer, "original_padding"):
+ layer.padding = layer.original_padding
+ else:
+ layer.original_padding = layer.padding
+ layer.padding = tuple(0 for _ in range(len(layer.padding)))
+
+ self._padding = value
+
+ def get_delay(self):
+ # Any number works here, delay is invariant to input length
+ l_out = self.get_output_length(0)
+ L = l_out
+
+ layers = []
+ for layer in self.modules():
+ if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)):
+ layers.append(layer)
+
+ for layer in reversed(layers):
+ d = layer.dilation[0]
+ k = layer.kernel_size[0]
+ s = layer.stride[0]
+
+ if isinstance(layer, nn.ConvTranspose1d):
+ L = ((L - d * (k - 1) - 1) / s) + 1
+ elif isinstance(layer, nn.Conv1d):
+ L = (L - 1) * s + d * (k - 1) + 1
+
+ L = math.ceil(L)
+
+ l_in = L
+
+ return (l_in - l_out) // 2
+
+ def get_output_length(self, input_length):
+ L = input_length
+ # Calculate output length
+ for layer in self.modules():
+ if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)):
+ d = layer.dilation[0]
+ k = layer.kernel_size[0]
+ s = layer.stride[0]
+
+ if isinstance(layer, nn.Conv1d):
+ L = ((L - d * (k - 1) - 1) / s) + 1
+ elif isinstance(layer, nn.ConvTranspose1d):
+ L = (L - 1) * s + d * (k - 1) + 1
+
+ L = math.floor(L)
+ return L
+
+ @torch.no_grad()
+ def compress(
+ self,
+ audio_path_or_signal: Union[str, Path, AudioSignal],
+ win_duration: float = 1.0,
+ verbose: bool = False,
+ normalize_db: float = -16,
+ n_quantizers: int = None,
+ ) -> DACFile:
+ """Processes an audio signal from a file or AudioSignal object into
+ discrete codes. This function processes the signal in short windows,
+ using constant GPU memory.
+
+ Parameters
+ ----------
+ audio_path_or_signal : Union[str, Path, AudioSignal]
+ audio signal to reconstruct
+ win_duration : float, optional
+ window duration in seconds, by default 5.0
+ verbose : bool, optional
+ by default False
+ normalize_db : float, optional
+ normalize db, by default -16
+
+ Returns
+ -------
+ DACFile
+ Object containing compressed codes and metadata
+ required for decompression
+ """
+ audio_signal = audio_path_or_signal
+ if isinstance(audio_signal, (str, Path)):
+ audio_signal = AudioSignal.load_from_file_with_ffmpeg(str(audio_signal))
+
+ self.eval()
+ original_padding = self.padding
+ original_device = audio_signal.device
+
+ audio_signal = audio_signal.clone()
+ original_sr = audio_signal.sample_rate
+
+ resample_fn = audio_signal.resample
+ loudness_fn = audio_signal.loudness
+
+ # If audio is > 10 minutes long, use the ffmpeg versions
+ if audio_signal.signal_duration >= 10 * 60 * 60:
+ resample_fn = audio_signal.ffmpeg_resample
+ loudness_fn = audio_signal.ffmpeg_loudness
+
+ original_length = audio_signal.signal_length
+ resample_fn(self.sample_rate)
+ input_db = loudness_fn()
+
+ if normalize_db is not None:
+ audio_signal.normalize(normalize_db)
+ audio_signal.ensure_max_of_audio()
+
+ nb, nac, nt = audio_signal.audio_data.shape
+ audio_signal.audio_data = audio_signal.audio_data.reshape(nb * nac, 1, nt)
+ win_duration = (
+ audio_signal.signal_duration if win_duration is None else win_duration
+ )
+
+ if audio_signal.signal_duration <= win_duration:
+ # Unchunked compression (used if signal length < win duration)
+ self.padding = True
+ n_samples = nt
+ hop = nt
+ else:
+ # Chunked inference
+ self.padding = False
+ # Zero-pad signal on either side by the delay
+ audio_signal.zero_pad(self.delay, self.delay)
+ n_samples = int(win_duration * self.sample_rate)
+ # Round n_samples to nearest hop length multiple
+ n_samples = int(math.ceil(n_samples / self.hop_length) * self.hop_length)
+ hop = self.get_output_length(n_samples)
+
+ codes = []
+ range_fn = range if not verbose else tqdm.trange
+
+ for i in range_fn(0, nt, hop):
+ x = audio_signal[..., i : i + n_samples]
+ x = x.zero_pad(0, max(0, n_samples - x.shape[-1]))
+
+ audio_data = x.audio_data.to(self.device)
+ audio_data = self.preprocess(audio_data, self.sample_rate)
+ _, c, _, _, _ = self.encode(audio_data, n_quantizers)
+ codes.append(c.to(original_device))
+ chunk_length = c.shape[-1]
+
+ codes = torch.cat(codes, dim=-1)
+
+ dac_file = DACFile(
+ codes=codes,
+ chunk_length=chunk_length,
+ original_length=original_length,
+ input_db=input_db,
+ channels=nac,
+ sample_rate=original_sr,
+ padding=self.padding,
+ dac_version=SUPPORTED_VERSIONS[-1],
+ )
+
+ if n_quantizers is not None:
+ codes = codes[:, :n_quantizers, :]
+
+ self.padding = original_padding
+ return dac_file
+
+ @torch.no_grad()
+ def decompress(
+ self,
+ obj: Union[str, Path, DACFile],
+ verbose: bool = False,
+ ) -> AudioSignal:
+ """Reconstruct audio from a given .dac file
+
+ Parameters
+ ----------
+ obj : Union[str, Path, DACFile]
+ .dac file location or corresponding DACFile object.
+ verbose : bool, optional
+ Prints progress if True, by default False
+
+ Returns
+ -------
+ AudioSignal
+ Object with the reconstructed audio
+ """
+ self.eval()
+ if isinstance(obj, (str, Path)):
+ obj = DACFile.load(obj)
+
+ original_padding = self.padding
+ self.padding = obj.padding
+
+ range_fn = range if not verbose else tqdm.trange
+ codes = obj.codes
+ original_device = codes.device
+ chunk_length = obj.chunk_length
+ recons = []
+
+ for i in range_fn(0, codes.shape[-1], chunk_length):
+ c = codes[..., i : i + chunk_length].to(self.device)
+ z = self.quantizer.from_codes(c)[0]
+ r = self.decode(z)
+ recons.append(r.to(original_device))
+
+ recons = torch.cat(recons, dim=-1)
+ recons = AudioSignal(recons, self.sample_rate)
+
+ resample_fn = recons.resample
+ loudness_fn = recons.loudness
+
+ # If audio is > 10 minutes long, use the ffmpeg versions
+ if recons.signal_duration >= 10 * 60 * 60:
+ resample_fn = recons.ffmpeg_resample
+ loudness_fn = recons.ffmpeg_loudness
+
+ recons.normalize(obj.input_db)
+ resample_fn(obj.sample_rate)
+ recons = recons[..., : obj.original_length]
+ loudness_fn()
+ recons.audio_data = recons.audio_data.reshape(
+ -1, obj.channels, obj.original_length
+ )
+
+ self.padding = original_padding
+ return recons
diff --git a/dac/model/dac.py b/dac/model/dac.py
new file mode 100644
index 0000000000000000000000000000000000000000..272229486aaf6d915ae8d1d0e81b11d685354a2b
--- /dev/null
+++ b/dac/model/dac.py
@@ -0,0 +1,400 @@
+import math
+from typing import List
+from typing import Union
+
+import numpy as np
+import torch
+from audiotools import AudioSignal
+from audiotools.ml import BaseModel
+from torch import nn
+
+from .base import CodecMixin
+from dac.nn.layers import Snake1d
+from dac.nn.layers import WNConv1d
+from dac.nn.layers import WNConvTranspose1d
+from dac.nn.quantize import ResidualVectorQuantize
+from .encodec import SConv1d, SConvTranspose1d, SLSTM
+
+
+def init_weights(m):
+ if isinstance(m, nn.Conv1d):
+ nn.init.trunc_normal_(m.weight, std=0.02)
+ nn.init.constant_(m.bias, 0)
+
+
+class ResidualUnit(nn.Module):
+ def __init__(self, dim: int = 16, dilation: int = 1, causal: bool = False):
+ super().__init__()
+ conv1d_type = SConv1d# if causal else WNConv1d
+ pad = ((7 - 1) * dilation) // 2
+ self.block = nn.Sequential(
+ Snake1d(dim),
+ conv1d_type(dim, dim, kernel_size=7, dilation=dilation, padding=pad, causal=causal, norm='weight_norm'),
+ Snake1d(dim),
+ conv1d_type(dim, dim, kernel_size=1, causal=causal, norm='weight_norm'),
+ )
+
+ def forward(self, x):
+ y = self.block(x)
+ pad = (x.shape[-1] - y.shape[-1]) // 2
+ if pad > 0:
+ x = x[..., pad:-pad]
+ return x + y
+
+
+class EncoderBlock(nn.Module):
+ def __init__(self, dim: int = 16, stride: int = 1, causal: bool = False):
+ super().__init__()
+ conv1d_type = SConv1d# if causal else WNConv1d
+ self.block = nn.Sequential(
+ ResidualUnit(dim // 2, dilation=1, causal=causal),
+ ResidualUnit(dim // 2, dilation=3, causal=causal),
+ ResidualUnit(dim // 2, dilation=9, causal=causal),
+ Snake1d(dim // 2),
+ conv1d_type(
+ dim // 2,
+ dim,
+ kernel_size=2 * stride,
+ stride=stride,
+ padding=math.ceil(stride / 2),
+ causal=causal,
+ norm='weight_norm',
+ ),
+ )
+
+ def forward(self, x):
+ return self.block(x)
+
+
+class Encoder(nn.Module):
+ def __init__(
+ self,
+ d_model: int = 64,
+ strides: list = [2, 4, 8, 8],
+ d_latent: int = 64,
+ causal: bool = False,
+ lstm: int = 2,
+ ):
+ super().__init__()
+ conv1d_type = SConv1d# if causal else WNConv1d
+ # Create first convolution
+ self.block = [conv1d_type(1, d_model, kernel_size=7, padding=3, causal=causal, norm='weight_norm')]
+
+ # Create EncoderBlocks that double channels as they downsample by `stride`
+ for stride in strides:
+ d_model *= 2
+ self.block += [EncoderBlock(d_model, stride=stride, causal=causal)]
+
+ # Add LSTM if needed
+ self.use_lstm = lstm
+ if lstm:
+ self.block += [SLSTM(d_model, lstm)]
+
+ # Create last convolution
+ self.block += [
+ Snake1d(d_model),
+ conv1d_type(d_model, d_latent, kernel_size=3, padding=1, causal=causal, norm='weight_norm'),
+ ]
+
+ # Wrap black into nn.Sequential
+ self.block = nn.Sequential(*self.block)
+ self.enc_dim = d_model
+
+ def forward(self, x):
+ return self.block(x)
+
+ def reset_cache(self):
+ # recursively find all submodules named SConv1d in self.block and use their reset_cache method
+ def reset_cache(m):
+ if isinstance(m, SConv1d) or isinstance(m, SLSTM):
+ m.reset_cache()
+ return
+ for child in m.children():
+ reset_cache(child)
+
+ reset_cache(self.block)
+
+
+class DecoderBlock(nn.Module):
+ def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1, causal: bool = False):
+ super().__init__()
+ conv1d_type = SConvTranspose1d #if causal else WNConvTranspose1d
+ self.block = nn.Sequential(
+ Snake1d(input_dim),
+ conv1d_type(
+ input_dim,
+ output_dim,
+ kernel_size=2 * stride,
+ stride=stride,
+ padding=math.ceil(stride / 2),
+ causal=causal,
+ norm='weight_norm'
+ ),
+ ResidualUnit(output_dim, dilation=1, causal=causal),
+ ResidualUnit(output_dim, dilation=3, causal=causal),
+ ResidualUnit(output_dim, dilation=9, causal=causal),
+ )
+
+ def forward(self, x):
+ return self.block(x)
+
+
+class Decoder(nn.Module):
+ def __init__(
+ self,
+ input_channel,
+ channels,
+ rates,
+ d_out: int = 1,
+ causal: bool = False,
+ lstm: int = 2,
+ ):
+ super().__init__()
+ conv1d_type = SConv1d# if causal else WNConv1d
+ # Add first conv layer
+ layers = [conv1d_type(input_channel, channels, kernel_size=7, padding=3, causal=causal, norm='weight_norm')]
+
+ if lstm:
+ layers += [SLSTM(channels, num_layers=lstm)]
+
+ # Add upsampling + MRF blocks
+ for i, stride in enumerate(rates):
+ input_dim = channels // 2**i
+ output_dim = channels // 2 ** (i + 1)
+ layers += [DecoderBlock(input_dim, output_dim, stride, causal=causal)]
+
+ # Add final conv layer
+ layers += [
+ Snake1d(output_dim),
+ conv1d_type(output_dim, d_out, kernel_size=7, padding=3, causal=causal, norm='weight_norm'),
+ nn.Tanh(),
+ ]
+
+ self.model = nn.Sequential(*layers)
+
+ def forward(self, x):
+ return self.model(x)
+
+
+class DAC(BaseModel, CodecMixin):
+ def __init__(
+ self,
+ encoder_dim: int = 64,
+ encoder_rates: List[int] = [2, 4, 8, 8],
+ latent_dim: int = None,
+ decoder_dim: int = 1536,
+ decoder_rates: List[int] = [8, 8, 4, 2],
+ n_codebooks: int = 9,
+ codebook_size: int = 1024,
+ codebook_dim: Union[int, list] = 8,
+ quantizer_dropout: bool = False,
+ sample_rate: int = 44100,
+ lstm: int = 2,
+ causal: bool = False,
+ ):
+ super().__init__()
+
+ self.encoder_dim = encoder_dim
+ self.encoder_rates = encoder_rates
+ self.decoder_dim = decoder_dim
+ self.decoder_rates = decoder_rates
+ self.sample_rate = sample_rate
+
+ if latent_dim is None:
+ latent_dim = encoder_dim * (2 ** len(encoder_rates))
+
+ self.latent_dim = latent_dim
+
+ self.hop_length = np.prod(encoder_rates)
+ self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim, causal=causal, lstm=lstm)
+
+ self.n_codebooks = n_codebooks
+ self.codebook_size = codebook_size
+ self.codebook_dim = codebook_dim
+ self.quantizer = ResidualVectorQuantize(
+ input_dim=latent_dim,
+ n_codebooks=n_codebooks,
+ codebook_size=codebook_size,
+ codebook_dim=codebook_dim,
+ quantizer_dropout=quantizer_dropout,
+ )
+
+ self.decoder = Decoder(
+ latent_dim,
+ decoder_dim,
+ decoder_rates,
+ lstm=lstm,
+ causal=causal,
+ )
+ self.sample_rate = sample_rate
+ self.apply(init_weights)
+
+ self.delay = self.get_delay()
+
+ def preprocess(self, audio_data, sample_rate):
+ if sample_rate is None:
+ sample_rate = self.sample_rate
+ assert sample_rate == self.sample_rate
+
+ length = audio_data.shape[-1]
+ right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
+ audio_data = nn.functional.pad(audio_data, (0, right_pad))
+
+ return audio_data
+
+ def encode(
+ self,
+ audio_data: torch.Tensor,
+ n_quantizers: int = None,
+ ):
+ """Encode given audio data and return quantized latent codes
+
+ Parameters
+ ----------
+ audio_data : Tensor[B x 1 x T]
+ Audio data to encode
+ n_quantizers : int, optional
+ Number of quantizers to use, by default None
+ If None, all quantizers are used.
+
+ Returns
+ -------
+ dict
+ A dictionary with the following keys:
+ "z" : Tensor[B x D x T]
+ Quantized continuous representation of input
+ "codes" : Tensor[B x N x T]
+ Codebook indices for each codebook
+ (quantized discrete representation of input)
+ "latents" : Tensor[B x N*D x T]
+ Projected latents (continuous representation of input before quantization)
+ "vq/commitment_loss" : Tensor[1]
+ Commitment loss to train encoder to predict vectors closer to codebook
+ entries
+ "vq/codebook_loss" : Tensor[1]
+ Codebook loss to update the codebook
+ "length" : int
+ Number of samples in input audio
+ """
+ z = self.encoder(audio_data)
+ z, codes, latents, commitment_loss, codebook_loss = self.quantizer(
+ z, n_quantizers
+ )
+ return z, codes, latents, commitment_loss, codebook_loss
+
+ def decode(self, z: torch.Tensor):
+ """Decode given latent codes and return audio data
+
+ Parameters
+ ----------
+ z : Tensor[B x D x T]
+ Quantized continuous representation of input
+ length : int, optional
+ Number of samples in output audio, by default None
+
+ Returns
+ -------
+ dict
+ A dictionary with the following keys:
+ "audio" : Tensor[B x 1 x length]
+ Decoded audio data.
+ """
+ return self.decoder(z)
+
+ def forward(
+ self,
+ audio_data: torch.Tensor,
+ sample_rate: int = None,
+ n_quantizers: int = None,
+ ):
+ """Model forward pass
+
+ Parameters
+ ----------
+ audio_data : Tensor[B x 1 x T]
+ Audio data to encode
+ sample_rate : int, optional
+ Sample rate of audio data in Hz, by default None
+ If None, defaults to `self.sample_rate`
+ n_quantizers : int, optional
+ Number of quantizers to use, by default None.
+ If None, all quantizers are used.
+
+ Returns
+ -------
+ dict
+ A dictionary with the following keys:
+ "z" : Tensor[B x D x T]
+ Quantized continuous representation of input
+ "codes" : Tensor[B x N x T]
+ Codebook indices for each codebook
+ (quantized discrete representation of input)
+ "latents" : Tensor[B x N*D x T]
+ Projected latents (continuous representation of input before quantization)
+ "vq/commitment_loss" : Tensor[1]
+ Commitment loss to train encoder to predict vectors closer to codebook
+ entries
+ "vq/codebook_loss" : Tensor[1]
+ Codebook loss to update the codebook
+ "length" : int
+ Number of samples in input audio
+ "audio" : Tensor[B x 1 x length]
+ Decoded audio data.
+ """
+ length = audio_data.shape[-1]
+ audio_data = self.preprocess(audio_data, sample_rate)
+ z, codes, latents, commitment_loss, codebook_loss = self.encode(
+ audio_data, n_quantizers
+ )
+
+ x = self.decode(z)
+ return {
+ "audio": x[..., :length],
+ "z": z,
+ "codes": codes,
+ "latents": latents,
+ "vq/commitment_loss": commitment_loss,
+ "vq/codebook_loss": codebook_loss,
+ }
+
+
+if __name__ == "__main__":
+ import numpy as np
+ from functools import partial
+
+ model = DAC().to("cpu")
+
+ for n, m in model.named_modules():
+ o = m.extra_repr()
+ p = sum([np.prod(p.size()) for p in m.parameters()])
+ fn = lambda o, p: o + f" {p/1e6:<.3f}M params."
+ setattr(m, "extra_repr", partial(fn, o=o, p=p))
+ print(model)
+ print("Total # of params: ", sum([np.prod(p.size()) for p in model.parameters()]))
+
+ length = 88200 * 2
+ x = torch.randn(1, 1, length).to(model.device)
+ x.requires_grad_(True)
+ x.retain_grad()
+
+ # Make a forward pass
+ out = model(x)["audio"]
+ print("Input shape:", x.shape)
+ print("Output shape:", out.shape)
+
+ # Create gradient variable
+ grad = torch.zeros_like(out)
+ grad[:, :, grad.shape[-1] // 2] = 1
+
+ # Make a backward pass
+ out.backward(grad)
+
+ # Check non-zero values
+ gradmap = x.grad.squeeze(0)
+ gradmap = (gradmap != 0).sum(0) # sum across features
+ rf = (gradmap != 0).sum()
+
+ print(f"Receptive field: {rf.item()}")
+
+ x = AudioSignal(torch.randn(1, 1, 44100 * 60), 44100)
+ model.decompress(model.compress(x, verbose=True), verbose=True)
diff --git a/dac/model/discriminator.py b/dac/model/discriminator.py
new file mode 100644
index 0000000000000000000000000000000000000000..5aa1a307e1dcff19ed36c681eac3dba39ddabb59
--- /dev/null
+++ b/dac/model/discriminator.py
@@ -0,0 +1,228 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from audiotools import AudioSignal
+from audiotools import ml
+from audiotools import STFTParams
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+
+def WNConv1d(*args, **kwargs):
+ act = kwargs.pop("act", True)
+ conv = weight_norm(nn.Conv1d(*args, **kwargs))
+ if not act:
+ return conv
+ return nn.Sequential(conv, nn.LeakyReLU(0.1))
+
+
+def WNConv2d(*args, **kwargs):
+ act = kwargs.pop("act", True)
+ conv = weight_norm(nn.Conv2d(*args, **kwargs))
+ if not act:
+ return conv
+ return nn.Sequential(conv, nn.LeakyReLU(0.1))
+
+
+class MPD(nn.Module):
+ def __init__(self, period):
+ super().__init__()
+ self.period = period
+ self.convs = nn.ModuleList(
+ [
+ WNConv2d(1, 32, (5, 1), (3, 1), padding=(2, 0)),
+ WNConv2d(32, 128, (5, 1), (3, 1), padding=(2, 0)),
+ WNConv2d(128, 512, (5, 1), (3, 1), padding=(2, 0)),
+ WNConv2d(512, 1024, (5, 1), (3, 1), padding=(2, 0)),
+ WNConv2d(1024, 1024, (5, 1), 1, padding=(2, 0)),
+ ]
+ )
+ self.conv_post = WNConv2d(
+ 1024, 1, kernel_size=(3, 1), padding=(1, 0), act=False
+ )
+
+ def pad_to_period(self, x):
+ t = x.shape[-1]
+ x = F.pad(x, (0, self.period - t % self.period), mode="reflect")
+ return x
+
+ def forward(self, x):
+ fmap = []
+
+ x = self.pad_to_period(x)
+ x = rearrange(x, "b c (l p) -> b c l p", p=self.period)
+
+ for layer in self.convs:
+ x = layer(x)
+ fmap.append(x)
+
+ x = self.conv_post(x)
+ fmap.append(x)
+
+ return fmap
+
+
+class MSD(nn.Module):
+ def __init__(self, rate: int = 1, sample_rate: int = 44100):
+ super().__init__()
+ self.convs = nn.ModuleList(
+ [
+ WNConv1d(1, 16, 15, 1, padding=7),
+ WNConv1d(16, 64, 41, 4, groups=4, padding=20),
+ WNConv1d(64, 256, 41, 4, groups=16, padding=20),
+ WNConv1d(256, 1024, 41, 4, groups=64, padding=20),
+ WNConv1d(1024, 1024, 41, 4, groups=256, padding=20),
+ WNConv1d(1024, 1024, 5, 1, padding=2),
+ ]
+ )
+ self.conv_post = WNConv1d(1024, 1, 3, 1, padding=1, act=False)
+ self.sample_rate = sample_rate
+ self.rate = rate
+
+ def forward(self, x):
+ x = AudioSignal(x, self.sample_rate)
+ x.resample(self.sample_rate // self.rate)
+ x = x.audio_data
+
+ fmap = []
+
+ for l in self.convs:
+ x = l(x)
+ fmap.append(x)
+ x = self.conv_post(x)
+ fmap.append(x)
+
+ return fmap
+
+
+BANDS = [(0.0, 0.1), (0.1, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.0)]
+
+
+class MRD(nn.Module):
+ def __init__(
+ self,
+ window_length: int,
+ hop_factor: float = 0.25,
+ sample_rate: int = 44100,
+ bands: list = BANDS,
+ ):
+ """Complex multi-band spectrogram discriminator.
+ Parameters
+ ----------
+ window_length : int
+ Window length of STFT.
+ hop_factor : float, optional
+ Hop factor of the STFT, defaults to ``0.25 * window_length``.
+ sample_rate : int, optional
+ Sampling rate of audio in Hz, by default 44100
+ bands : list, optional
+ Bands to run discriminator over.
+ """
+ super().__init__()
+
+ self.window_length = window_length
+ self.hop_factor = hop_factor
+ self.sample_rate = sample_rate
+ self.stft_params = STFTParams(
+ window_length=window_length,
+ hop_length=int(window_length * hop_factor),
+ match_stride=True,
+ )
+
+ n_fft = window_length // 2 + 1
+ bands = [(int(b[0] * n_fft), int(b[1] * n_fft)) for b in bands]
+ self.bands = bands
+
+ ch = 32
+ convs = lambda: nn.ModuleList(
+ [
+ WNConv2d(2, ch, (3, 9), (1, 1), padding=(1, 4)),
+ WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
+ WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
+ WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
+ WNConv2d(ch, ch, (3, 3), (1, 1), padding=(1, 1)),
+ ]
+ )
+ self.band_convs = nn.ModuleList([convs() for _ in range(len(self.bands))])
+ self.conv_post = WNConv2d(ch, 1, (3, 3), (1, 1), padding=(1, 1), act=False)
+
+ def spectrogram(self, x):
+ x = AudioSignal(x, self.sample_rate, stft_params=self.stft_params)
+ x = torch.view_as_real(x.stft())
+ x = rearrange(x, "b 1 f t c -> (b 1) c t f")
+ # Split into bands
+ x_bands = [x[..., b[0] : b[1]] for b in self.bands]
+ return x_bands
+
+ def forward(self, x):
+ x_bands = self.spectrogram(x)
+ fmap = []
+
+ x = []
+ for band, stack in zip(x_bands, self.band_convs):
+ for layer in stack:
+ band = layer(band)
+ fmap.append(band)
+ x.append(band)
+
+ x = torch.cat(x, dim=-1)
+ x = self.conv_post(x)
+ fmap.append(x)
+
+ return fmap
+
+
+class Discriminator(nn.Module):
+ def __init__(
+ self,
+ rates: list = [],
+ periods: list = [2, 3, 5, 7, 11],
+ fft_sizes: list = [2048, 1024, 512],
+ sample_rate: int = 44100,
+ bands: list = BANDS,
+ ):
+ """Discriminator that combines multiple discriminators.
+
+ Parameters
+ ----------
+ rates : list, optional
+ sampling rates (in Hz) to run MSD at, by default []
+ If empty, MSD is not used.
+ periods : list, optional
+ periods (of samples) to run MPD at, by default [2, 3, 5, 7, 11]
+ fft_sizes : list, optional
+ Window sizes of the FFT to run MRD at, by default [2048, 1024, 512]
+ sample_rate : int, optional
+ Sampling rate of audio in Hz, by default 44100
+ bands : list, optional
+ Bands to run MRD at, by default `BANDS`
+ """
+ super().__init__()
+ discs = []
+ discs += [MPD(p) for p in periods]
+ discs += [MSD(r, sample_rate=sample_rate) for r in rates]
+ discs += [MRD(f, sample_rate=sample_rate, bands=bands) for f in fft_sizes]
+ self.discriminators = nn.ModuleList(discs)
+
+ def preprocess(self, y):
+ # Remove DC offset
+ y = y - y.mean(dim=-1, keepdims=True)
+ # Peak normalize the volume of input audio
+ y = 0.8 * y / (y.abs().max(dim=-1, keepdim=True)[0] + 1e-9)
+ return y
+
+ def forward(self, x):
+ x = self.preprocess(x)
+ fmaps = [d(x) for d in self.discriminators]
+ return fmaps
+
+
+if __name__ == "__main__":
+ disc = Discriminator()
+ x = torch.zeros(1, 1, 44100)
+ results = disc(x)
+ for i, result in enumerate(results):
+ print(f"disc{i}")
+ for i, r in enumerate(result):
+ print(r.shape, r.mean(), r.min(), r.max())
+ print()
diff --git a/dac/model/encodec.py b/dac/model/encodec.py
new file mode 100644
index 0000000000000000000000000000000000000000..38c9cb3609322cde1997f456e2bd0a256dcd0fc7
--- /dev/null
+++ b/dac/model/encodec.py
@@ -0,0 +1,320 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Convolutional layers wrappers and utilities."""
+
+import math
+import typing as tp
+import warnings
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.utils import spectral_norm, weight_norm
+
+import typing as tp
+
+import einops
+
+
+class ConvLayerNorm(nn.LayerNorm):
+ """
+ Convolution-friendly LayerNorm that moves channels to last dimensions
+ before running the normalization and moves them back to original position right after.
+ """
+ def __init__(self, normalized_shape: tp.Union[int, tp.List[int], torch.Size], **kwargs):
+ super().__init__(normalized_shape, **kwargs)
+
+ def forward(self, x):
+ x = einops.rearrange(x, 'b ... t -> b t ...')
+ x = super().forward(x)
+ x = einops.rearrange(x, 'b t ... -> b ... t')
+ return
+
+
+CONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm',
+ 'time_layer_norm', 'layer_norm', 'time_group_norm'])
+
+
+def apply_parametrization_norm(module: nn.Module, norm: str = 'none') -> nn.Module:
+ assert norm in CONV_NORMALIZATIONS
+ if norm == 'weight_norm':
+ return weight_norm(module)
+ elif norm == 'spectral_norm':
+ return spectral_norm(module)
+ else:
+ # We already check was in CONV_NORMALIZATION, so any other choice
+ # doesn't need reparametrization.
+ return module
+
+
+def get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs) -> nn.Module:
+ """Return the proper normalization module. If causal is True, this will ensure the returned
+ module is causal, or return an error if the normalization doesn't support causal evaluation.
+ """
+ assert norm in CONV_NORMALIZATIONS
+ if norm == 'layer_norm':
+ assert isinstance(module, nn.modules.conv._ConvNd)
+ return ConvLayerNorm(module.out_channels, **norm_kwargs)
+ elif norm == 'time_group_norm':
+ if causal:
+ raise ValueError("GroupNorm doesn't support causal evaluation.")
+ assert isinstance(module, nn.modules.conv._ConvNd)
+ return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
+ else:
+ return nn.Identity()
+
+
+def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int,
+ padding_total: int = 0) -> int:
+ """See `pad_for_conv1d`.
+ """
+ length = x.shape[-1]
+ n_frames = (length - kernel_size + padding_total) / stride + 1
+ ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
+ return ideal_length - length
+
+
+def pad_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0):
+ """Pad for a convolution to make sure that the last window is full.
+ Extra padding is added at the end. This is required to ensure that we can rebuild
+ an output of the same length, as otherwise, even with padding, some time steps
+ might get removed.
+ For instance, with total padding = 4, kernel size = 4, stride = 2:
+ 0 0 1 2 3 4 5 0 0 # (0s are padding)
+ 1 2 3 # (output frames of a convolution, last 0 is never used)
+ 0 0 1 2 3 4 5 0 # (output of tr. conv., but pos. 5 is going to get removed as padding)
+ 1 2 3 4 # once you removed padding, we are missing one time step !
+ """
+ extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+ return F.pad(x, (0, extra_padding))
+
+
+def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'zero', value: float = 0.):
+ """Tiny wrapper around F.pad, just to allow for reflect padding on small input.
+ If this is the case, we insert extra 0 padding to the right before the reflection happen.
+ """
+ length = x.shape[-1]
+ padding_left, padding_right = paddings
+ assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+ if mode == 'reflect':
+ max_pad = max(padding_left, padding_right)
+ extra_pad = 0
+ if length <= max_pad:
+ extra_pad = max_pad - length + 1
+ x = F.pad(x, (0, extra_pad))
+ padded = F.pad(x, paddings, mode, value)
+ end = padded.shape[-1] - extra_pad
+ return padded[..., :end]
+ else:
+ return F.pad(x, paddings, mode, value)
+
+
+def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
+ """Remove padding from x, handling properly zero padding. Only for 1d!"""
+ padding_left, padding_right = paddings
+ assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+ assert (padding_left + padding_right) <= x.shape[-1]
+ end = x.shape[-1] - padding_right
+ return x[..., padding_left: end]
+
+
+class NormConv1d(nn.Module):
+ """Wrapper around Conv1d and normalization applied to this conv
+ to provide a uniform interface across normalization approaches.
+ """
+ def __init__(self, *args, causal: bool = False, norm: str = 'none',
+ norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+ super().__init__()
+ self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
+ self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
+ self.norm_type = norm
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.norm(x)
+ return x
+
+
+class NormConv2d(nn.Module):
+ """Wrapper around Conv2d and normalization applied to this conv
+ to provide a uniform interface across normalization approaches.
+ """
+ def __init__(self, *args, norm: str = 'none',
+ norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+ super().__init__()
+ self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)
+ self.norm = get_norm_module(self.conv, causal=False, norm=norm, **norm_kwargs)
+ self.norm_type = norm
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.norm(x)
+ return x
+
+
+class NormConvTranspose1d(nn.Module):
+ """Wrapper around ConvTranspose1d and normalization applied to this conv
+ to provide a uniform interface across normalization approaches.
+ """
+ def __init__(self, *args, causal: bool = False, norm: str = 'none',
+ norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+ super().__init__()
+ self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm)
+ self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)
+ self.norm_type = norm
+
+ def forward(self, x):
+ x = self.convtr(x)
+ x = self.norm(x)
+ return x
+
+
+class NormConvTranspose2d(nn.Module):
+ """Wrapper around ConvTranspose2d and normalization applied to this conv
+ to provide a uniform interface across normalization approaches.
+ """
+ def __init__(self, *args, norm: str = 'none',
+ norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+ super().__init__()
+ self.convtr = apply_parametrization_norm(nn.ConvTranspose2d(*args, **kwargs), norm)
+ self.norm = get_norm_module(self.convtr, causal=False, norm=norm, **norm_kwargs)
+
+ def forward(self, x):
+ x = self.convtr(x)
+ x = self.norm(x)
+ return x
+
+
+class SConv1d(nn.Module):
+ """Conv1d with some builtin handling of asymmetric or causal padding
+ and normalization.
+ """
+ def __init__(self, in_channels: int, out_channels: int,
+ kernel_size: int, stride: int = 1, dilation: int = 1,
+ groups: int = 1, bias: bool = True, causal: bool = False,
+ norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {},
+ pad_mode: str = 'reflect', **kwargs):
+ super().__init__()
+ # warn user on unusual setup between dilation and stride
+ if stride > 1 and dilation > 1:
+ warnings.warn('SConv1d has been initialized with stride > 1 and dilation > 1'
+ f' (kernel_size={kernel_size} stride={stride}, dilation={dilation}).')
+ self.conv = NormConv1d(in_channels, out_channels, kernel_size, stride,
+ dilation=dilation, groups=groups, bias=bias, causal=causal,
+ norm=norm, norm_kwargs=norm_kwargs)
+ self.causal = causal
+ self.pad_mode = pad_mode
+
+ self.cache_enabled = False
+
+ def reset_cache(self):
+ """Reset the cache when starting a new stream."""
+ self.cache = None
+ self.cache_enabled = True
+
+ def forward(self, x):
+ B, C, T = x.shape
+ kernel_size = self.conv.conv.kernel_size[0]
+ stride = self.conv.conv.stride[0]
+ dilation = self.conv.conv.dilation[0]
+ kernel_size = (kernel_size - 1) * dilation + 1 # effective kernel size with dilations
+ padding_total = kernel_size - stride
+ extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+
+ if self.causal:
+ # Left padding for causal
+ if self.cache_enabled and self.cache is not None:
+ # Concatenate the cache (previous inputs) with the new input for streaming
+ x = torch.cat([self.cache, x], dim=2)
+ else:
+ x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
+ else:
+ # Asymmetric padding required for odd strides
+ padding_right = padding_total // 2
+ padding_left = padding_total - padding_right
+ x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode)
+
+ # Store the most recent input frames for future cache use
+ if self.cache_enabled:
+ if self.cache is None:
+ # Initialize cache with zeros (at the start of streaming)
+ self.cache = torch.zeros(B, C, kernel_size - 1, device=x.device)
+ # Update the cache by storing the latest input frames
+ if kernel_size > 1:
+ self.cache = x[:, :, -kernel_size + 1:].detach() # Only store the necessary frames
+
+ return self.conv(x)
+
+
+
+class SConvTranspose1d(nn.Module):
+ """ConvTranspose1d with some builtin handling of asymmetric or causal padding
+ and normalization.
+ """
+ def __init__(self, in_channels: int, out_channels: int,
+ kernel_size: int, stride: int = 1, causal: bool = False,
+ norm: str = 'none', trim_right_ratio: float = 1.,
+ norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+ super().__init__()
+ self.convtr = NormConvTranspose1d(in_channels, out_channels, kernel_size, stride,
+ causal=causal, norm=norm, norm_kwargs=norm_kwargs)
+ self.causal = causal
+ self.trim_right_ratio = trim_right_ratio
+ assert self.causal or self.trim_right_ratio == 1., \
+ "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
+ assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1.
+
+ def forward(self, x):
+ kernel_size = self.convtr.convtr.kernel_size[0]
+ stride = self.convtr.convtr.stride[0]
+ padding_total = kernel_size - stride
+
+ y = self.convtr(x)
+
+ # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be
+ # removed at the very end, when keeping only the right length for the output,
+ # as removing it here would require also passing the length at the matching layer
+ # in the encoder.
+ if self.causal:
+ # Trim the padding on the right according to the specified ratio
+ # if trim_right_ratio = 1.0, trim everything from right
+ padding_right = math.ceil(padding_total * self.trim_right_ratio)
+ padding_left = padding_total - padding_right
+ y = unpad1d(y, (padding_left, padding_right))
+ else:
+ # Asymmetric padding required for odd strides
+ padding_right = padding_total // 2
+ padding_left = padding_total - padding_right
+ y = unpad1d(y, (padding_left, padding_right))
+ return y
+
+class SLSTM(nn.Module):
+ """
+ LSTM without worrying about the hidden state, nor the layout of the data.
+ Expects input as convolutional layout.
+ """
+ def __init__(self, dimension: int, num_layers: int = 2, skip: bool = True):
+ super().__init__()
+ self.skip = skip
+ self.lstm = nn.LSTM(dimension, dimension, num_layers)
+ self.hidden = None
+ self.cache_enabled = False
+
+ def forward(self, x):
+ x = x.permute(2, 0, 1)
+ if self.training or not self.cache_enabled:
+ y, _ = self.lstm(x)
+ else:
+ y, self.hidden = self.lstm(x, self.hidden)
+ if self.skip:
+ y = y + x
+ y = y.permute(1, 2, 0)
+ return y
+
+ def reset_cache(self):
+ self.hidden = None
+ self.cache_enabled = True
\ No newline at end of file
diff --git a/dac/nn/__init__.py b/dac/nn/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ec2f6b972aeb8a21764ff73dae2095eb94bb8ba4
--- /dev/null
+++ b/dac/nn/__init__.py
@@ -0,0 +1,3 @@
+from . import layers
+from . import loss
+from . import quantize
diff --git a/dac/nn/__pycache__/__init__.cpython-310.pyc b/dac/nn/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ea5881530369892e17f5064112b1fd62a41cbc09
Binary files /dev/null and b/dac/nn/__pycache__/__init__.cpython-310.pyc differ
diff --git a/dac/nn/__pycache__/layers.cpython-310.pyc b/dac/nn/__pycache__/layers.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..937cedd877b9c037bdb45436351ed24c39766a4f
Binary files /dev/null and b/dac/nn/__pycache__/layers.cpython-310.pyc differ
diff --git a/dac/nn/__pycache__/loss.cpython-310.pyc b/dac/nn/__pycache__/loss.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..01753711afd094b933fc8e56aa23c0810d673e60
Binary files /dev/null and b/dac/nn/__pycache__/loss.cpython-310.pyc differ
diff --git a/dac/nn/__pycache__/quantize.cpython-310.pyc b/dac/nn/__pycache__/quantize.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2abd5a78cb6883450392148012504034757e90c6
Binary files /dev/null and b/dac/nn/__pycache__/quantize.cpython-310.pyc differ
diff --git a/dac/nn/layers.py b/dac/nn/layers.py
new file mode 100644
index 0000000000000000000000000000000000000000..63c53a88b881deab94fb06b6a395951cf9cc995a
--- /dev/null
+++ b/dac/nn/layers.py
@@ -0,0 +1,33 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+
+def WNConv1d(*args, **kwargs):
+ return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+ return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+# Scripting this brings model speed up 1.4x
+@torch.jit.script
+def snake(x, alpha):
+ shape = x.shape
+ x = x.reshape(shape[0], shape[1], -1)
+ x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+ x = x.reshape(shape)
+ return x
+
+
+class Snake1d(nn.Module):
+ def __init__(self, channels):
+ super().__init__()
+ self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+ def forward(self, x):
+ return snake(x, self.alpha)
diff --git a/dac/nn/loss.py b/dac/nn/loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..7a26e62f885f0bb7c6774e95660a7848e10a8f87
--- /dev/null
+++ b/dac/nn/loss.py
@@ -0,0 +1,368 @@
+import typing
+from typing import List
+
+import torch
+import torch.nn.functional as F
+from audiotools import AudioSignal
+from audiotools import STFTParams
+from torch import nn
+
+
+class L1Loss(nn.L1Loss):
+ """L1 Loss between AudioSignals. Defaults
+ to comparing ``audio_data``, but any
+ attribute of an AudioSignal can be used.
+
+ Parameters
+ ----------
+ attribute : str, optional
+ Attribute of signal to compare, defaults to ``audio_data``.
+ weight : float, optional
+ Weight of this loss, defaults to 1.0.
+
+ Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/distance.py
+ """
+
+ def __init__(self, attribute: str = "audio_data", weight: float = 1.0, **kwargs):
+ self.attribute = attribute
+ self.weight = weight
+ super().__init__(**kwargs)
+
+ def forward(self, x: AudioSignal, y: AudioSignal):
+ """
+ Parameters
+ ----------
+ x : AudioSignal
+ Estimate AudioSignal
+ y : AudioSignal
+ Reference AudioSignal
+
+ Returns
+ -------
+ torch.Tensor
+ L1 loss between AudioSignal attributes.
+ """
+ if isinstance(x, AudioSignal):
+ x = getattr(x, self.attribute)
+ y = getattr(y, self.attribute)
+ return super().forward(x, y)
+
+
+class SISDRLoss(nn.Module):
+ """
+ Computes the Scale-Invariant Source-to-Distortion Ratio between a batch
+ of estimated and reference audio signals or aligned features.
+
+ Parameters
+ ----------
+ scaling : int, optional
+ Whether to use scale-invariant (True) or
+ signal-to-noise ratio (False), by default True
+ reduction : str, optional
+ How to reduce across the batch (either 'mean',
+ 'sum', or none).], by default ' mean'
+ zero_mean : int, optional
+ Zero mean the references and estimates before
+ computing the loss, by default True
+ clip_min : int, optional
+ The minimum possible loss value. Helps network
+ to not focus on making already good examples better, by default None
+ weight : float, optional
+ Weight of this loss, defaults to 1.0.
+
+ Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/distance.py
+ """
+
+ def __init__(
+ self,
+ scaling: int = True,
+ reduction: str = "mean",
+ zero_mean: int = True,
+ clip_min: int = None,
+ weight: float = 1.0,
+ ):
+ self.scaling = scaling
+ self.reduction = reduction
+ self.zero_mean = zero_mean
+ self.clip_min = clip_min
+ self.weight = weight
+ super().__init__()
+
+ def forward(self, x: AudioSignal, y: AudioSignal):
+ eps = 1e-8
+ # nb, nc, nt
+ if isinstance(x, AudioSignal):
+ references = x.audio_data
+ estimates = y.audio_data
+ else:
+ references = x
+ estimates = y
+
+ nb = references.shape[0]
+ references = references.reshape(nb, 1, -1).permute(0, 2, 1)
+ estimates = estimates.reshape(nb, 1, -1).permute(0, 2, 1)
+
+ # samples now on axis 1
+ if self.zero_mean:
+ mean_reference = references.mean(dim=1, keepdim=True)
+ mean_estimate = estimates.mean(dim=1, keepdim=True)
+ else:
+ mean_reference = 0
+ mean_estimate = 0
+
+ _references = references - mean_reference
+ _estimates = estimates - mean_estimate
+
+ references_projection = (_references**2).sum(dim=-2) + eps
+ references_on_estimates = (_estimates * _references).sum(dim=-2) + eps
+
+ scale = (
+ (references_on_estimates / references_projection).unsqueeze(1)
+ if self.scaling
+ else 1
+ )
+
+ e_true = scale * _references
+ e_res = _estimates - e_true
+
+ signal = (e_true**2).sum(dim=1)
+ noise = (e_res**2).sum(dim=1)
+ sdr = -10 * torch.log10(signal / noise + eps)
+
+ if self.clip_min is not None:
+ sdr = torch.clamp(sdr, min=self.clip_min)
+
+ if self.reduction == "mean":
+ sdr = sdr.mean()
+ elif self.reduction == "sum":
+ sdr = sdr.sum()
+ return sdr
+
+
+class MultiScaleSTFTLoss(nn.Module):
+ """Computes the multi-scale STFT loss from [1].
+
+ Parameters
+ ----------
+ window_lengths : List[int], optional
+ Length of each window of each STFT, by default [2048, 512]
+ loss_fn : typing.Callable, optional
+ How to compare each loss, by default nn.L1Loss()
+ clamp_eps : float, optional
+ Clamp on the log magnitude, below, by default 1e-5
+ mag_weight : float, optional
+ Weight of raw magnitude portion of loss, by default 1.0
+ log_weight : float, optional
+ Weight of log magnitude portion of loss, by default 1.0
+ pow : float, optional
+ Power to raise magnitude to before taking log, by default 2.0
+ weight : float, optional
+ Weight of this loss, by default 1.0
+ match_stride : bool, optional
+ Whether to match the stride of convolutional layers, by default False
+
+ References
+ ----------
+
+ 1. Engel, Jesse, Chenjie Gu, and Adam Roberts.
+ "DDSP: Differentiable Digital Signal Processing."
+ International Conference on Learning Representations. 2019.
+
+ Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py
+ """
+
+ def __init__(
+ self,
+ window_lengths: List[int] = [2048, 512],
+ loss_fn: typing.Callable = nn.L1Loss(),
+ clamp_eps: float = 1e-5,
+ mag_weight: float = 1.0,
+ log_weight: float = 1.0,
+ pow: float = 2.0,
+ weight: float = 1.0,
+ match_stride: bool = False,
+ window_type: str = None,
+ ):
+ super().__init__()
+ self.stft_params = [
+ STFTParams(
+ window_length=w,
+ hop_length=w // 4,
+ match_stride=match_stride,
+ window_type=window_type,
+ )
+ for w in window_lengths
+ ]
+ self.loss_fn = loss_fn
+ self.log_weight = log_weight
+ self.mag_weight = mag_weight
+ self.clamp_eps = clamp_eps
+ self.weight = weight
+ self.pow = pow
+
+ def forward(self, x: AudioSignal, y: AudioSignal):
+ """Computes multi-scale STFT between an estimate and a reference
+ signal.
+
+ Parameters
+ ----------
+ x : AudioSignal
+ Estimate signal
+ y : AudioSignal
+ Reference signal
+
+ Returns
+ -------
+ torch.Tensor
+ Multi-scale STFT loss.
+ """
+ loss = 0.0
+ for s in self.stft_params:
+ x.stft(s.window_length, s.hop_length, s.window_type)
+ y.stft(s.window_length, s.hop_length, s.window_type)
+ loss += self.log_weight * self.loss_fn(
+ x.magnitude.clamp(self.clamp_eps).pow(self.pow).log10(),
+ y.magnitude.clamp(self.clamp_eps).pow(self.pow).log10(),
+ )
+ loss += self.mag_weight * self.loss_fn(x.magnitude, y.magnitude)
+ return loss
+
+
+class MelSpectrogramLoss(nn.Module):
+ """Compute distance between mel spectrograms. Can be used
+ in a multi-scale way.
+
+ Parameters
+ ----------
+ n_mels : List[int]
+ Number of mels per STFT, by default [150, 80],
+ window_lengths : List[int], optional
+ Length of each window of each STFT, by default [2048, 512]
+ loss_fn : typing.Callable, optional
+ How to compare each loss, by default nn.L1Loss()
+ clamp_eps : float, optional
+ Clamp on the log magnitude, below, by default 1e-5
+ mag_weight : float, optional
+ Weight of raw magnitude portion of loss, by default 1.0
+ log_weight : float, optional
+ Weight of log magnitude portion of loss, by default 1.0
+ pow : float, optional
+ Power to raise magnitude to before taking log, by default 2.0
+ weight : float, optional
+ Weight of this loss, by default 1.0
+ match_stride : bool, optional
+ Whether to match the stride of convolutional layers, by default False
+
+ Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py
+ """
+
+ def __init__(
+ self,
+ n_mels: List[int] = [150, 80],
+ window_lengths: List[int] = [2048, 512],
+ loss_fn: typing.Callable = nn.L1Loss(),
+ clamp_eps: float = 1e-5,
+ mag_weight: float = 1.0,
+ log_weight: float = 1.0,
+ pow: float = 2.0,
+ weight: float = 1.0,
+ match_stride: bool = False,
+ mel_fmin: List[float] = [0.0, 0.0],
+ mel_fmax: List[float] = [None, None],
+ window_type: str = None,
+ ):
+ super().__init__()
+ self.stft_params = [
+ STFTParams(
+ window_length=w,
+ hop_length=w // 4,
+ match_stride=match_stride,
+ window_type=window_type,
+ )
+ for w in window_lengths
+ ]
+ self.n_mels = n_mels
+ self.loss_fn = loss_fn
+ self.clamp_eps = clamp_eps
+ self.log_weight = log_weight
+ self.mag_weight = mag_weight
+ self.weight = weight
+ self.mel_fmin = mel_fmin
+ self.mel_fmax = mel_fmax
+ self.pow = pow
+
+ def forward(self, x: AudioSignal, y: AudioSignal):
+ """Computes mel loss between an estimate and a reference
+ signal.
+
+ Parameters
+ ----------
+ x : AudioSignal
+ Estimate signal
+ y : AudioSignal
+ Reference signal
+
+ Returns
+ -------
+ torch.Tensor
+ Mel loss.
+ """
+ loss = 0.0
+ for n_mels, fmin, fmax, s in zip(
+ self.n_mels, self.mel_fmin, self.mel_fmax, self.stft_params
+ ):
+ kwargs = {
+ "window_length": s.window_length,
+ "hop_length": s.hop_length,
+ "window_type": s.window_type,
+ }
+ x_mels = x.mel_spectrogram(n_mels, mel_fmin=fmin, mel_fmax=fmax, **kwargs)
+ y_mels = y.mel_spectrogram(n_mels, mel_fmin=fmin, mel_fmax=fmax, **kwargs)
+
+ loss += self.log_weight * self.loss_fn(
+ x_mels.clamp(self.clamp_eps).pow(self.pow).log10(),
+ y_mels.clamp(self.clamp_eps).pow(self.pow).log10(),
+ )
+ loss += self.mag_weight * self.loss_fn(x_mels, y_mels)
+ return loss
+
+
+class GANLoss(nn.Module):
+ """
+ Computes a discriminator loss, given a discriminator on
+ generated waveforms/spectrograms compared to ground truth
+ waveforms/spectrograms. Computes the loss for both the
+ discriminator and the generator in separate functions.
+ """
+
+ def __init__(self, discriminator):
+ super().__init__()
+ self.discriminator = discriminator
+
+ def forward(self, fake, real):
+ d_fake = self.discriminator(fake.audio_data)
+ d_real = self.discriminator(real.audio_data)
+ return d_fake, d_real
+
+ def discriminator_loss(self, fake, real):
+ d_fake, d_real = self.forward(fake.clone().detach(), real)
+
+ loss_d = 0
+ for x_fake, x_real in zip(d_fake, d_real):
+ loss_d += torch.mean(x_fake[-1] ** 2)
+ loss_d += torch.mean((1 - x_real[-1]) ** 2)
+ return loss_d
+
+ def generator_loss(self, fake, real):
+ d_fake, d_real = self.forward(fake, real)
+
+ loss_g = 0
+ for x_fake in d_fake:
+ loss_g += torch.mean((1 - x_fake[-1]) ** 2)
+
+ loss_feature = 0
+
+ for i in range(len(d_fake)):
+ for j in range(len(d_fake[i]) - 1):
+ loss_feature += F.l1_loss(d_fake[i][j], d_real[i][j].detach())
+ return loss_g, loss_feature
diff --git a/dac/nn/quantize.py b/dac/nn/quantize.py
new file mode 100644
index 0000000000000000000000000000000000000000..962029dc668fc98cbbea5bd19bf5ea6f41467875
--- /dev/null
+++ b/dac/nn/quantize.py
@@ -0,0 +1,339 @@
+from typing import Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+from dac.nn.layers import WNConv1d
+
+class VectorQuantizeLegacy(nn.Module):
+ """
+ Implementation of VQ similar to Karpathy's repo:
+ https://github.com/karpathy/deep-vector-quantization
+ removed in-out projection
+ """
+
+ def __init__(self, input_dim: int, codebook_size: int):
+ super().__init__()
+ self.codebook_size = codebook_size
+ self.codebook = nn.Embedding(codebook_size, input_dim)
+
+ def forward(self, z, z_mask=None):
+ """Quantized the input tensor using a fixed codebook and returns
+ the corresponding codebook vectors
+
+ Parameters
+ ----------
+ z : Tensor[B x D x T]
+
+ Returns
+ -------
+ Tensor[B x D x T]
+ Quantized continuous representation of input
+ Tensor[1]
+ Commitment loss to train encoder to predict vectors closer to codebook
+ entries
+ Tensor[1]
+ Codebook loss to update the codebook
+ Tensor[B x T]
+ Codebook indices (quantized discrete representation of input)
+ Tensor[B x D x T]
+ Projected latents (continuous representation of input before quantization)
+ """
+
+ z_e = z
+ z_q, indices = self.decode_latents(z)
+
+ if z_mask is not None:
+ commitment_loss = (F.mse_loss(z_e, z_q.detach(), reduction="none").mean(1) * z_mask).sum() / z_mask.sum()
+ codebook_loss = (F.mse_loss(z_q, z_e.detach(), reduction="none").mean(1) * z_mask).sum() / z_mask.sum()
+ else:
+ commitment_loss = F.mse_loss(z_e, z_q.detach())
+ codebook_loss = F.mse_loss(z_q, z_e.detach())
+ z_q = (
+ z_e + (z_q - z_e).detach()
+ ) # noop in forward pass, straight-through gradient estimator in backward pass
+
+ return z_q, indices, z_e, commitment_loss, codebook_loss
+
+ def embed_code(self, embed_id):
+ return F.embedding(embed_id, self.codebook.weight)
+
+ def decode_code(self, embed_id):
+ return self.embed_code(embed_id).transpose(1, 2)
+
+ def decode_latents(self, latents):
+ encodings = rearrange(latents, "b d t -> (b t) d")
+ codebook = self.codebook.weight # codebook: (N x D)
+
+ # L2 normalize encodings and codebook (ViT-VQGAN)
+ encodings = F.normalize(encodings)
+ codebook = F.normalize(codebook)
+
+ # Compute euclidean distance with codebook
+ dist = (
+ encodings.pow(2).sum(1, keepdim=True)
+ - 2 * encodings @ codebook.t()
+ + codebook.pow(2).sum(1, keepdim=True).t()
+ )
+ indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+ z_q = self.decode_code(indices)
+ return z_q, indices
+
+class VectorQuantize(nn.Module):
+ """
+ Implementation of VQ similar to Karpathy's repo:
+ https://github.com/karpathy/deep-vector-quantization
+ Additionally uses following tricks from Improved VQGAN
+ (https://arxiv.org/pdf/2110.04627.pdf):
+ 1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space
+ for improved codebook usage
+ 2. l2-normalized codes: Converts euclidean distance to cosine similarity which
+ improves training stability
+ """
+
+ def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int):
+ super().__init__()
+ self.codebook_size = codebook_size
+ self.codebook_dim = codebook_dim
+
+ self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1)
+ self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1)
+ self.codebook = nn.Embedding(codebook_size, codebook_dim)
+
+ def forward(self, z, z_mask=None):
+ """Quantized the input tensor using a fixed codebook and returns
+ the corresponding codebook vectors
+
+ Parameters
+ ----------
+ z : Tensor[B x D x T]
+
+ Returns
+ -------
+ Tensor[B x D x T]
+ Quantized continuous representation of input
+ Tensor[1]
+ Commitment loss to train encoder to predict vectors closer to codebook
+ entries
+ Tensor[1]
+ Codebook loss to update the codebook
+ Tensor[B x T]
+ Codebook indices (quantized discrete representation of input)
+ Tensor[B x D x T]
+ Projected latents (continuous representation of input before quantization)
+ """
+
+ # Factorized codes (ViT-VQGAN) Project input into low-dimensional space
+ z_e = self.in_proj(z) # z_e : (B x D x T)
+ z_q, indices = self.decode_latents(z_e)
+
+ if z_mask is not None:
+ commitment_loss = (F.mse_loss(z_e, z_q.detach(), reduction="none").mean(1) * z_mask).sum() / z_mask.sum()
+ codebook_loss = (F.mse_loss(z_q, z_e.detach(), reduction="none").mean(1) * z_mask).sum() / z_mask.sum()
+ else:
+ commitment_loss = F.mse_loss(z_e, z_q.detach())
+ codebook_loss = F.mse_loss(z_q, z_e.detach())
+
+ z_q = (
+ z_e + (z_q - z_e).detach()
+ ) # noop in forward pass, straight-through gradient estimator in backward pass
+
+ z_q = self.out_proj(z_q)
+
+ return z_q, commitment_loss, codebook_loss, indices, z_e
+
+ def embed_code(self, embed_id):
+ return F.embedding(embed_id, self.codebook.weight)
+
+ def decode_code(self, embed_id):
+ return self.embed_code(embed_id).transpose(1, 2)
+
+ def decode_latents(self, latents):
+ encodings = rearrange(latents, "b d t -> (b t) d")
+ codebook = self.codebook.weight # codebook: (N x D)
+
+ # L2 normalize encodings and codebook (ViT-VQGAN)
+ encodings = F.normalize(encodings)
+ codebook = F.normalize(codebook)
+
+ # Compute euclidean distance with codebook
+ dist = (
+ encodings.pow(2).sum(1, keepdim=True)
+ - 2 * encodings @ codebook.t()
+ + codebook.pow(2).sum(1, keepdim=True).t()
+ )
+ indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+ z_q = self.decode_code(indices)
+ return z_q, indices
+
+
+class ResidualVectorQuantize(nn.Module):
+ """
+ Introduced in SoundStream: An end2end neural audio codec
+ https://arxiv.org/abs/2107.03312
+ """
+
+ def __init__(
+ self,
+ input_dim: int = 512,
+ n_codebooks: int = 9,
+ codebook_size: int = 1024,
+ codebook_dim: Union[int, list] = 8,
+ quantizer_dropout: float = 0.0,
+ ):
+ super().__init__()
+ if isinstance(codebook_dim, int):
+ codebook_dim = [codebook_dim for _ in range(n_codebooks)]
+
+ self.n_codebooks = n_codebooks
+ self.codebook_dim = codebook_dim
+ self.codebook_size = codebook_size
+
+ self.quantizers = nn.ModuleList(
+ [
+ VectorQuantize(input_dim, codebook_size, codebook_dim[i])
+ for i in range(n_codebooks)
+ ]
+ )
+ self.quantizer_dropout = quantizer_dropout
+
+ def forward(self, z, n_quantizers: int = None):
+ """Quantized the input tensor using a fixed set of `n` codebooks and returns
+ the corresponding codebook vectors
+ Parameters
+ ----------
+ z : Tensor[B x D x T]
+ n_quantizers : int, optional
+ No. of quantizers to use
+ (n_quantizers < self.n_codebooks ex: for quantizer dropout)
+ Note: if `self.quantizer_dropout` is True, this argument is ignored
+ when in training mode, and a random number of quantizers is used.
+ Returns
+ -------
+ dict
+ A dictionary with the following keys:
+
+ "z" : Tensor[B x D x T]
+ Quantized continuous representation of input
+ "codes" : Tensor[B x N x T]
+ Codebook indices for each codebook
+ (quantized discrete representation of input)
+ "latents" : Tensor[B x N*D x T]
+ Projected latents (continuous representation of input before quantization)
+ "vq/commitment_loss" : Tensor[1]
+ Commitment loss to train encoder to predict vectors closer to codebook
+ entries
+ "vq/codebook_loss" : Tensor[1]
+ Codebook loss to update the codebook
+ """
+ z_q = 0
+ residual = z
+ commitment_loss = 0
+ codebook_loss = 0
+
+ codebook_indices = []
+ latents = []
+
+ if n_quantizers is None:
+ n_quantizers = self.n_codebooks
+ if self.training:
+ n_quantizers = torch.ones((z.shape[0],)) * self.n_codebooks + 1
+ dropout = torch.randint(1, self.n_codebooks + 1, (z.shape[0],))
+ n_dropout = int(z.shape[0] * self.quantizer_dropout)
+ n_quantizers[:n_dropout] = dropout[:n_dropout]
+ n_quantizers = n_quantizers.to(z.device)
+
+ for i, quantizer in enumerate(self.quantizers):
+ if self.training is False and i >= n_quantizers:
+ break
+
+ z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(
+ residual
+ )
+
+ # Create mask to apply quantizer dropout
+ mask = (
+ torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers
+ )
+ z_q = z_q + z_q_i * mask[:, None, None]
+ residual = residual - z_q_i
+
+ # Sum losses
+ commitment_loss += (commitment_loss_i * mask).mean()
+ codebook_loss += (codebook_loss_i * mask).mean()
+
+ codebook_indices.append(indices_i)
+ latents.append(z_e_i)
+
+ codes = torch.stack(codebook_indices, dim=1)
+ latents = torch.cat(latents, dim=1)
+
+ return z_q, codes, latents, commitment_loss, codebook_loss
+
+ def from_codes(self, codes: torch.Tensor):
+ """Given the quantized codes, reconstruct the continuous representation
+ Parameters
+ ----------
+ codes : Tensor[B x N x T]
+ Quantized discrete representation of input
+ Returns
+ -------
+ Tensor[B x D x T]
+ Quantized continuous representation of input
+ """
+ z_q = 0.0
+ z_p = []
+ n_codebooks = codes.shape[1]
+ for i in range(n_codebooks):
+ z_p_i = self.quantizers[i].decode_code(codes[:, i, :])
+ z_p.append(z_p_i)
+
+ z_q_i = self.quantizers[i].out_proj(z_p_i)
+ z_q = z_q + z_q_i
+ return z_q, torch.cat(z_p, dim=1), codes
+
+ def from_latents(self, latents: torch.Tensor):
+ """Given the unquantized latents, reconstruct the
+ continuous representation after quantization.
+
+ Parameters
+ ----------
+ latents : Tensor[B x N x T]
+ Continuous representation of input after projection
+
+ Returns
+ -------
+ Tensor[B x D x T]
+ Quantized representation of full-projected space
+ Tensor[B x D x T]
+ Quantized representation of latent space
+ """
+ z_q = 0
+ z_p = []
+ codes = []
+ dims = np.cumsum([0] + [q.codebook_dim for q in self.quantizers])
+
+ n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[
+ 0
+ ]
+ for i in range(n_codebooks):
+ j, k = dims[i], dims[i + 1]
+ z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :])
+ z_p.append(z_p_i)
+ codes.append(codes_i)
+
+ z_q_i = self.quantizers[i].out_proj(z_p_i)
+ z_q = z_q + z_q_i
+
+ return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1)
+
+
+if __name__ == "__main__":
+ rvq = ResidualVectorQuantize(quantizer_dropout=True)
+ x = torch.randn(16, 512, 80)
+ y = rvq(x)
+ print(y["latents"].shape)
diff --git a/dac/utils/__init__.py b/dac/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f6dbede7b3e57655b4a732c92c84d7c72fbc9059
--- /dev/null
+++ b/dac/utils/__init__.py
@@ -0,0 +1,123 @@
+from pathlib import Path
+
+import argbind
+from audiotools import ml
+
+import dac
+
+DAC = dac.model.DAC
+Accelerator = ml.Accelerator
+
+__MODEL_LATEST_TAGS__ = {
+ ("44khz", "8kbps"): "0.0.1",
+ ("24khz", "8kbps"): "0.0.4",
+ ("16khz", "8kbps"): "0.0.5",
+ ("44khz", "16kbps"): "1.0.0",
+}
+
+__MODEL_URLS__ = {
+ (
+ "44khz",
+ "0.0.1",
+ "8kbps",
+ ): "https://github.com/descriptinc/descript-audio-codec/releases/download/0.0.1/weights.pth",
+ (
+ "24khz",
+ "0.0.4",
+ "8kbps",
+ ): "https://github.com/descriptinc/descript-audio-codec/releases/download/0.0.4/weights_24khz.pth",
+ (
+ "16khz",
+ "0.0.5",
+ "8kbps",
+ ): "https://github.com/descriptinc/descript-audio-codec/releases/download/0.0.5/weights_16khz.pth",
+ (
+ "44khz",
+ "1.0.0",
+ "16kbps",
+ ): "https://github.com/descriptinc/descript-audio-codec/releases/download/1.0.0/weights_44khz_16kbps.pth",
+}
+
+
+@argbind.bind(group="download", positional=True, without_prefix=True)
+def download(
+ model_type: str = "44khz", model_bitrate: str = "8kbps", tag: str = "latest"
+):
+ """
+ Function that downloads the weights file from URL if a local cache is not found.
+
+ Parameters
+ ----------
+ model_type : str
+ The type of model to download. Must be one of "44khz", "24khz", or "16khz". Defaults to "44khz".
+ model_bitrate: str
+ Bitrate of the model. Must be one of "8kbps", or "16kbps". Defaults to "8kbps".
+ Only 44khz model supports 16kbps.
+ tag : str
+ The tag of the model to download. Defaults to "latest".
+
+ Returns
+ -------
+ Path
+ Directory path required to load model via audiotools.
+ """
+ model_type = model_type.lower()
+ tag = tag.lower()
+
+ assert model_type in [
+ "44khz",
+ "24khz",
+ "16khz",
+ ], "model_type must be one of '44khz', '24khz', or '16khz'"
+
+ assert model_bitrate in [
+ "8kbps",
+ "16kbps",
+ ], "model_bitrate must be one of '8kbps', or '16kbps'"
+
+ if tag == "latest":
+ tag = __MODEL_LATEST_TAGS__[(model_type, model_bitrate)]
+
+ download_link = __MODEL_URLS__.get((model_type, tag, model_bitrate), None)
+
+ if download_link is None:
+ raise ValueError(
+ f"Could not find model with tag {tag} and model type {model_type}"
+ )
+
+ local_path = (
+ Path.home()
+ / ".cache"
+ / "descript"
+ / "dac"
+ / f"weights_{model_type}_{model_bitrate}_{tag}.pth"
+ )
+ if not local_path.exists():
+ local_path.parent.mkdir(parents=True, exist_ok=True)
+
+ # Download the model
+ import requests
+
+ response = requests.get(download_link)
+
+ if response.status_code != 200:
+ raise ValueError(
+ f"Could not download model. Received response code {response.status_code}"
+ )
+ local_path.write_bytes(response.content)
+
+ return local_path
+
+
+def load_model(
+ model_type: str = "44khz",
+ model_bitrate: str = "8kbps",
+ tag: str = "latest",
+ load_path: str = None,
+):
+ if not load_path:
+ load_path = download(
+ model_type=model_type, model_bitrate=model_bitrate, tag=tag
+ )
+ generator = DAC.load(load_path)
+ return generator
diff --git a/dac/utils/__pycache__/__init__.cpython-310.pyc b/dac/utils/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3f72b4aa117efc7129a22dad1faca05a33489cec
Binary files /dev/null and b/dac/utils/__pycache__/__init__.cpython-310.pyc differ
diff --git a/dac/utils/decode.py b/dac/utils/decode.py
new file mode 100644
index 0000000000000000000000000000000000000000..48c25298fd0f9a53f19c1d6d260f1f0563c7b7e4
--- /dev/null
+++ b/dac/utils/decode.py
@@ -0,0 +1,95 @@
+import warnings
+from pathlib import Path
+
+import argbind
+import numpy as np
+import torch
+from audiotools import AudioSignal
+from tqdm import tqdm
+
+from dac import DACFile
+from dac.utils import load_model
+
+warnings.filterwarnings("ignore", category=UserWarning)
+
+
+@argbind.bind(group="decode", positional=True, without_prefix=True)
+@torch.inference_mode()
+@torch.no_grad()
+def decode(
+ input: str,
+ output: str = "",
+ weights_path: str = "",
+ model_tag: str = "latest",
+ model_bitrate: str = "8kbps",
+ device: str = "cuda",
+ model_type: str = "44khz",
+ verbose: bool = False,
+):
+ """Decode audio from codes.
+
+ Parameters
+ ----------
+ input : str
+ Path to input directory or file
+ output : str, optional
+ Path to output directory, by default "".
+ If `input` is a directory, the directory sub-tree relative to `input` is re-created in `output`.
+ weights_path : str, optional
+ Path to weights file, by default "". If not specified, the weights file will be downloaded from the internet using the
+ model_tag and model_type.
+ model_tag : str, optional
+ Tag of the model to use, by default "latest". Ignored if `weights_path` is specified.
+ model_bitrate: str
+ Bitrate of the model. Must be one of "8kbps", or "16kbps". Defaults to "8kbps".
+ device : str, optional
+ Device to use, by default "cuda". If "cpu", the model will be loaded on the CPU.
+ model_type : str, optional
+ The type of model to use. Must be one of "44khz", "24khz", or "16khz". Defaults to "44khz". Ignored if `weights_path` is specified.
+ """
+ generator = load_model(
+ model_type=model_type,
+ model_bitrate=model_bitrate,
+ tag=model_tag,
+ load_path=weights_path,
+ )
+ generator.to(device)
+ generator.eval()
+
+ # Find all .dac files in input directory
+ _input = Path(input)
+ input_files = list(_input.glob("**/*.dac"))
+
+ # If input is a .dac file, add it to the list
+ if _input.suffix == ".dac":
+ input_files.append(_input)
+
+ # Create output directory
+ output = Path(output)
+ output.mkdir(parents=True, exist_ok=True)
+
+ for i in tqdm(range(len(input_files)), desc=f"Decoding files"):
+ # Load file
+ artifact = DACFile.load(input_files[i])
+
+ # Reconstruct audio from codes
+ recons = generator.decompress(artifact, verbose=verbose)
+
+ # Compute output path
+ relative_path = input_files[i].relative_to(input)
+ output_dir = output / relative_path.parent
+ if not relative_path.name:
+ output_dir = output
+ relative_path = input_files[i]
+ output_name = relative_path.with_suffix(".wav").name
+ output_path = output_dir / output_name
+ output_path.parent.mkdir(parents=True, exist_ok=True)
+
+ # Write to file
+ recons.write(output_path)
+
+
+if __name__ == "__main__":
+ args = argbind.parse_args()
+ with argbind.scope(args):
+ decode()
diff --git a/dac/utils/encode.py b/dac/utils/encode.py
new file mode 100644
index 0000000000000000000000000000000000000000..7c9a8b582b24c194ba82c2d280c4df4fa2cfff87
--- /dev/null
+++ b/dac/utils/encode.py
@@ -0,0 +1,94 @@
+import math
+import warnings
+from pathlib import Path
+
+import argbind
+import numpy as np
+import torch
+from audiotools import AudioSignal
+from audiotools.core import util
+from tqdm import tqdm
+
+from dac.utils import load_model
+
+warnings.filterwarnings("ignore", category=UserWarning)
+
+
+@argbind.bind(group="encode", positional=True, without_prefix=True)
+@torch.inference_mode()
+@torch.no_grad()
+def encode(
+ input: str,
+ output: str = "",
+ weights_path: str = "",
+ model_tag: str = "latest",
+ model_bitrate: str = "8kbps",
+ n_quantizers: int = None,
+ device: str = "cuda",
+ model_type: str = "44khz",
+ win_duration: float = 5.0,
+ verbose: bool = False,
+):
+ """Encode audio files in input path to .dac format.
+
+ Parameters
+ ----------
+ input : str
+ Path to input audio file or directory
+ output : str, optional
+ Path to output directory, by default "". If `input` is a directory, the directory sub-tree relative to `input` is re-created in `output`.
+ weights_path : str, optional
+ Path to weights file, by default "". If not specified, the weights file will be downloaded from the internet using the
+ model_tag and model_type.
+ model_tag : str, optional
+ Tag of the model to use, by default "latest". Ignored if `weights_path` is specified.
+ model_bitrate: str
+ Bitrate of the model. Must be one of "8kbps", or "16kbps". Defaults to "8kbps".
+ n_quantizers : int, optional
+ Number of quantizers to use, by default None. If not specified, all the quantizers will be used and the model will compress at maximum bitrate.
+ device : str, optional
+ Device to use, by default "cuda"
+ model_type : str, optional
+ The type of model to use. Must be one of "44khz", "24khz", or "16khz". Defaults to "44khz". Ignored if `weights_path` is specified.
+ """
+ generator = load_model(
+ model_type=model_type,
+ model_bitrate=model_bitrate,
+ tag=model_tag,
+ load_path=weights_path,
+ )
+ generator.to(device)
+ generator.eval()
+ kwargs = {"n_quantizers": n_quantizers}
+
+ # Find all audio files in input path
+ input = Path(input)
+ audio_files = util.find_audio(input)
+
+ output = Path(output)
+ output.mkdir(parents=True, exist_ok=True)
+
+ for i in tqdm(range(len(audio_files)), desc="Encoding files"):
+ # Load file
+ signal = AudioSignal(audio_files[i])
+
+ # Encode audio to .dac format
+ artifact = generator.compress(signal, win_duration, verbose=verbose, **kwargs)
+
+ # Compute output path
+ relative_path = audio_files[i].relative_to(input)
+ output_dir = output / relative_path.parent
+ if not relative_path.name:
+ output_dir = output
+ relative_path = audio_files[i]
+ output_name = relative_path.with_suffix(".dac").name
+ output_path = output_dir / output_name
+ output_path.parent.mkdir(parents=True, exist_ok=True)
+
+ artifact.save(output_path)
+
+
+if __name__ == "__main__":
+ args = argbind.parse_args()
+ with argbind.scope(args):
+ encode()
diff --git a/examples/reference/azuma_0.wav b/examples/reference/azuma_0.wav
new file mode 100644
index 0000000000000000000000000000000000000000..c2c64e1584de88bc53ac0ac85f96e0c8005d3437
--- /dev/null
+++ b/examples/reference/azuma_0.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3930141e927be50e7f3d666db5890ef9a4bda0623645483a6afaad241c82fb70
+size 628910
diff --git a/examples/reference/dingzhen_0.wav b/examples/reference/dingzhen_0.wav
new file mode 100644
index 0000000000000000000000000000000000000000..b260e1b4ed2d9ea746c711dd0f4983407dd4ee28
--- /dev/null
+++ b/examples/reference/dingzhen_0.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3db260824d11f56cdf2fccf2b84ad83c95a732ddfa2f8cb8a20b68ca06ea9ff8
+size 1088420
diff --git a/examples/reference/kobe_0.wav b/examples/reference/kobe_0.wav
new file mode 100644
index 0000000000000000000000000000000000000000..0e8d469a26c617d8552595682c15cf926d66f1a5
--- /dev/null
+++ b/examples/reference/kobe_0.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:075b46b422effe7908a20571821a9a769a434a4bd02d7c3a4b4a5aeccd668506
+size 643022
diff --git a/examples/reference/s1p1.wav b/examples/reference/s1p1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..0f024507d01b537c4f8a609641399649eeef71c9
--- /dev/null
+++ b/examples/reference/s1p1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be291f04e9c239218082552e9e7c0dba9bad5ce6306d2c9d2104195840214b5a
+size 700714
diff --git a/examples/reference/s1p2.wav b/examples/reference/s1p2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..eb88cecdea0d8493bb7fa9fc30d0f7e765f02d51
--- /dev/null
+++ b/examples/reference/s1p2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:07e937b52271380eeb52415428dc3fab6d530f767bf34b5e5cb52337ad294b17
+size 526116
diff --git a/examples/reference/s2p1.wav b/examples/reference/s2p1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..8e2bf143e8fc7887178cc564ed3dd35e2b50d20c
--- /dev/null
+++ b/examples/reference/s2p1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:308a38a5ac0a133a7d4742a95494d9670042fd9a233e60448f8370a9eba98a98
+size 664838
diff --git a/examples/reference/s2p2.wav b/examples/reference/s2p2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..04f6cbaa53901d2d784be3bc7703b1ec76890d19
--- /dev/null
+++ b/examples/reference/s2p2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e6f73146fcfd25044317b377b0f2e8ae03ab43bd43ae8713f37889357babbd5e
+size 564012
diff --git a/examples/reference/s3p1.wav b/examples/reference/s3p1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..49b4ab9fa82848be1586100fdefb37bd684f62fd
--- /dev/null
+++ b/examples/reference/s3p1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0cc05439406e44eeaa7be15742aa71e2ca8b9c43adbcd496ea1672d9283aa9b0
+size 557386
diff --git a/examples/reference/s3p2.wav b/examples/reference/s3p2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..0c5c30b861bdb58ece69e6bc17e9e713ddaa9e59
--- /dev/null
+++ b/examples/reference/s3p2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d28df338203ad8b3c7485474fac41d9ee2891cf27bc0c0239e3249e6c0efadb
+size 1140390
diff --git a/examples/reference/s4p1.wav b/examples/reference/s4p1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..67c3fe1cd5d24002567f3e6bacd8d7862f629249
--- /dev/null
+++ b/examples/reference/s4p1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:53892ece071342958403bc5643f84169a30b89cc0fc79eb69508bfa11dd85e68
+size 618528
diff --git a/examples/reference/s4p2.wav b/examples/reference/s4p2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..ecb0a82f481e0c456775aacbf549e5e83ff06356
--- /dev/null
+++ b/examples/reference/s4p2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a2c71e9e60b1eb772c9bc50ceea66566c9797a58a265289486715afe37bea813
+size 651098
diff --git a/examples/reference/teio_0.wav b/examples/reference/teio_0.wav
new file mode 100644
index 0000000000000000000000000000000000000000..99bc0ee47dc741e017304daad727284387b004a6
--- /dev/null
+++ b/examples/reference/teio_0.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:97af2b2fa7a950243f800bd5920d58841a1f602f7ccd23fc6479369956168842
+size 366402
diff --git a/examples/reference/trump_0.wav b/examples/reference/trump_0.wav
new file mode 100644
index 0000000000000000000000000000000000000000..c24f3495db37ef1884d4441723f3506355ac5ff6
--- /dev/null
+++ b/examples/reference/trump_0.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:716becc9daf00351dfe324398edea9e8378f9453408b27612d92b6721f80ddbc
+size 1379484
diff --git a/examples/source/TECHNOPOLIS - 2085 [vocals]_[cut_14sec].wav b/examples/source/TECHNOPOLIS - 2085 [vocals]_[cut_14sec].wav
new file mode 100644
index 0000000000000000000000000000000000000000..4dcdd668892a77d1ed8a392671c6b9d541b8c83f
--- /dev/null
+++ b/examples/source/TECHNOPOLIS - 2085 [vocals]_[cut_14sec].wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:87087ca5260ce96659b01a647edb30bb08527ed7d0c074fb5ae1e8338cc733e5
+size 2796016
diff --git a/examples/source/Wiz Khalifa,Charlie Puth - See You Again [vocals]_[cut_28sec].wav b/examples/source/Wiz Khalifa,Charlie Puth - See You Again [vocals]_[cut_28sec].wav
new file mode 100644
index 0000000000000000000000000000000000000000..38780402102a7a97d9e2d5ff19d445d8dfdff6a5
--- /dev/null
+++ b/examples/source/Wiz Khalifa,Charlie Puth - See You Again [vocals]_[cut_28sec].wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c7ff68178df50fc5f4f497099af8efe0b4508e7ff13665ee72d780159e6d1875
+size 5411452
diff --git a/examples/source/glados_0.wav b/examples/source/glados_0.wav
new file mode 100644
index 0000000000000000000000000000000000000000..58c9013feef828973f4eba0b06f852eca05e5ad8
--- /dev/null
+++ b/examples/source/glados_0.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb84353440fe0bcd44968fce96916074e9d573bfcf0da34a6d6c9a24fd3075e0
+size 640058
diff --git a/examples/source/jay_0.wav b/examples/source/jay_0.wav
new file mode 100644
index 0000000000000000000000000000000000000000..0c029b001844933049dbffc6c937b51720aae449
--- /dev/null
+++ b/examples/source/jay_0.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d30f1500acacb597c3b27d7a5937dd088b8029b27e9db8bf5982085f26f4457
+size 1270124
diff --git a/examples/source/source_s1.wav b/examples/source/source_s1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..a95e181b99e65464ab6e00d43bed2e5ad0efb9cc
--- /dev/null
+++ b/examples/source/source_s1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b91167bc3b5b310da093f8dfd228c209398b180930fb889ab9388180518a3475
+size 599038
diff --git a/examples/source/source_s2.wav b/examples/source/source_s2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..052862f5d00c2ef045208a80e89c62bb8176495c
--- /dev/null
+++ b/examples/source/source_s2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32bf74dfb8342ccc8284f47a9e5fe1f41b458a983a00938ed170f76951ce638f
+size 675496
diff --git a/examples/source/source_s3.wav b/examples/source/source_s3.wav
new file mode 100644
index 0000000000000000000000000000000000000000..e60db9c66f7fd43c471708754e5413241e4ca4df
--- /dev/null
+++ b/examples/source/source_s3.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e784fc21f6eade4633db4050528110a9fc307686487fcec257ad3db35a65c8d
+size 1168780
diff --git a/examples/source/source_s4.wav b/examples/source/source_s4.wav
new file mode 100644
index 0000000000000000000000000000000000000000..1ae47aa2d5f8271cbdf60220f06f87a21a31544b
--- /dev/null
+++ b/examples/source/source_s4.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee7ccf45748bf8918c17e9bf80f79af423cefd78d9c8eebb717a395305272b9f
+size 1068360
diff --git a/examples/source/yae_0.wav b/examples/source/yae_0.wav
new file mode 100644
index 0000000000000000000000000000000000000000..ab97d8d456e03946945705bebecdab7a59792892
--- /dev/null
+++ b/examples/source/yae_0.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:543afcfdbf20c6fc9fdf04d97faa6d545f4b99a92e3fc90655f1cf75172b3abd
+size 528044
diff --git a/hf_utils.py b/hf_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..4ae986c1d88d5b1214c9d3101249e90c5c370ca9
--- /dev/null
+++ b/hf_utils.py
@@ -0,0 +1,12 @@
+import os
+from huggingface_hub import hf_hub_download
+
+
+def load_custom_model_from_hf(repo_id, model_filename="pytorch_model.bin", config_filename=None):
+ os.makedirs("./checkpoints", exist_ok=True)
+ model_path = hf_hub_download(repo_id=repo_id, filename=model_filename, cache_dir="./checkpoints")
+ if config_filename is None:
+ return model_path
+ config_path = hf_hub_download(repo_id=repo_id, filename=config_filename, cache_dir="./checkpoints")
+
+ return model_path, config_path
\ No newline at end of file
diff --git a/modules/__pycache__/audio.cpython-310.pyc b/modules/__pycache__/audio.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9a74df30a404c4b9428c6a258fc08acbcd0c01c9
Binary files /dev/null and b/modules/__pycache__/audio.cpython-310.pyc differ
diff --git a/modules/__pycache__/commons.cpython-310.pyc b/modules/__pycache__/commons.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..410fdb13fad631d72c8c08f3faee7c88c4420a80
Binary files /dev/null and b/modules/__pycache__/commons.cpython-310.pyc differ
diff --git a/modules/__pycache__/commons.cpython-38.pyc b/modules/__pycache__/commons.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cfe6a34f14ee453f4a79de9b39d80fc62b4d7fda
Binary files /dev/null and b/modules/__pycache__/commons.cpython-38.pyc differ
diff --git a/modules/__pycache__/diffusion_transformer.cpython-310.pyc b/modules/__pycache__/diffusion_transformer.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2510e66289c2ab6bdc195bb268add8df05bccb5d
Binary files /dev/null and b/modules/__pycache__/diffusion_transformer.cpython-310.pyc differ
diff --git a/modules/__pycache__/encodec.cpython-310.pyc b/modules/__pycache__/encodec.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..49602d36e9c75bd422dd03297058c91d424c9a74
Binary files /dev/null and b/modules/__pycache__/encodec.cpython-310.pyc differ
diff --git a/modules/__pycache__/flow_matching.cpython-310.pyc b/modules/__pycache__/flow_matching.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ba5ba00a574e2bef2b3f9d7bb6436c679510ba97
Binary files /dev/null and b/modules/__pycache__/flow_matching.cpython-310.pyc differ
diff --git a/modules/__pycache__/length_regulator.cpython-310.pyc b/modules/__pycache__/length_regulator.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9aa75158b5960c81116d5e799344a670a19c7567
Binary files /dev/null and b/modules/__pycache__/length_regulator.cpython-310.pyc differ
diff --git a/modules/__pycache__/rmvpe.cpython-310.pyc b/modules/__pycache__/rmvpe.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6141174d53932fb187bdd318a1d62b7381fc4631
Binary files /dev/null and b/modules/__pycache__/rmvpe.cpython-310.pyc differ
diff --git a/modules/__pycache__/wavenet.cpython-310.pyc b/modules/__pycache__/wavenet.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7a435b3b1bb4a5037e856fa210329721c0a17c17
Binary files /dev/null and b/modules/__pycache__/wavenet.cpython-310.pyc differ
diff --git a/modules/alias_free_torch/__init__.py b/modules/alias_free_torch/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..080beecd700d74c9c915b7c41c5ab13c79103902
--- /dev/null
+++ b/modules/alias_free_torch/__init__.py
@@ -0,0 +1,5 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+
+from .filter import *
+from .resample import *
+from .act import *
diff --git a/modules/alias_free_torch/__pycache__/__init__.cpython-310.pyc b/modules/alias_free_torch/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..483eb08a0aef45cd7a84bc7a59461dff4d325b20
Binary files /dev/null and b/modules/alias_free_torch/__pycache__/__init__.cpython-310.pyc differ
diff --git a/modules/alias_free_torch/__pycache__/act.cpython-310.pyc b/modules/alias_free_torch/__pycache__/act.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b87f1678a289c9847ba3bc500711deecf51eb942
Binary files /dev/null and b/modules/alias_free_torch/__pycache__/act.cpython-310.pyc differ
diff --git a/modules/alias_free_torch/__pycache__/filter.cpython-310.pyc b/modules/alias_free_torch/__pycache__/filter.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c7b6208efe145606a7cf5bc2dac54a8e1123cecb
Binary files /dev/null and b/modules/alias_free_torch/__pycache__/filter.cpython-310.pyc differ
diff --git a/modules/alias_free_torch/__pycache__/resample.cpython-310.pyc b/modules/alias_free_torch/__pycache__/resample.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..17baf4e7acc2793941af10879a2c81cb58dccdd4
Binary files /dev/null and b/modules/alias_free_torch/__pycache__/resample.cpython-310.pyc differ
diff --git a/modules/alias_free_torch/act.py b/modules/alias_free_torch/act.py
new file mode 100644
index 0000000000000000000000000000000000000000..adc7eb7469c6fcb25e200abb78425f977d655cd8
--- /dev/null
+++ b/modules/alias_free_torch/act.py
@@ -0,0 +1,29 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+
+import torch.nn as nn
+from .resample import UpSample1d, DownSample1d
+
+
+class Activation1d(nn.Module):
+ def __init__(
+ self,
+ activation,
+ up_ratio: int = 2,
+ down_ratio: int = 2,
+ up_kernel_size: int = 12,
+ down_kernel_size: int = 12,
+ ):
+ super().__init__()
+ self.up_ratio = up_ratio
+ self.down_ratio = down_ratio
+ self.act = activation
+ self.upsample = UpSample1d(up_ratio, up_kernel_size)
+ self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+ # x: [B,C,T]
+ def forward(self, x):
+ x = self.upsample(x)
+ x = self.act(x)
+ x = self.downsample(x)
+
+ return x
diff --git a/modules/alias_free_torch/filter.py b/modules/alias_free_torch/filter.py
new file mode 100644
index 0000000000000000000000000000000000000000..8e18bda585d0006c4f34f5ae04bdea854c4e37b7
--- /dev/null
+++ b/modules/alias_free_torch/filter.py
@@ -0,0 +1,96 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+if "sinc" in dir(torch):
+ sinc = torch.sinc
+else:
+ # This code is adopted from adefossez's julius.core.sinc under the MIT License
+ # https://adefossez.github.io/julius/julius/core.html
+ def sinc(x: torch.Tensor):
+ """
+ Implementation of sinc, i.e. sin(pi * x) / (pi * x)
+ __Warning__: Different to julius.sinc, the input is multiplied by `pi`!
+ """
+ return torch.where(
+ x == 0,
+ torch.tensor(1.0, device=x.device, dtype=x.dtype),
+ torch.sin(math.pi * x) / math.pi / x,
+ )
+
+
+# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License
+# https://adefossez.github.io/julius/julius/lowpass.html
+def kaiser_sinc_filter1d(
+ cutoff, half_width, kernel_size
+): # return filter [1,1,kernel_size]
+ even = kernel_size % 2 == 0
+ half_size = kernel_size // 2
+
+ # For kaiser window
+ 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)
+
+ # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio
+ if even:
+ time = torch.arange(-half_size, half_size) + 0.5
+ else:
+ time = torch.arange(kernel_size) - half_size
+ if cutoff == 0:
+ filter_ = torch.zeros_like(time)
+ else:
+ filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
+ # Normalize filter to have sum = 1, otherwise we will have a small leakage
+ # of the constant component in the input signal.
+ filter_ /= filter_.sum()
+ filter = filter_.view(1, 1, kernel_size)
+
+ return filter
+
+
+class LowPassFilter1d(nn.Module):
+ def __init__(
+ self,
+ cutoff=0.5,
+ half_width=0.6,
+ stride: int = 1,
+ padding: bool = True,
+ padding_mode: str = "replicate",
+ kernel_size: int = 12,
+ ):
+ # kernel_size should be even number for stylegan3 setup,
+ # in this implementation, odd number is also possible.
+ super().__init__()
+ if cutoff < -0.0:
+ raise ValueError("Minimum cutoff must be larger than zero.")
+ if cutoff > 0.5:
+ raise ValueError("A cutoff above 0.5 does not make sense.")
+ self.kernel_size = kernel_size
+ self.even = kernel_size % 2 == 0
+ self.pad_left = kernel_size // 2 - int(self.even)
+ self.pad_right = kernel_size // 2
+ self.stride = stride
+ self.padding = padding
+ self.padding_mode = padding_mode
+ filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+ self.register_buffer("filter", filter)
+
+ # input [B, C, T]
+ def forward(self, x):
+ _, C, _ = x.shape
+
+ if self.padding:
+ x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+ out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
+
+ return out
diff --git a/modules/alias_free_torch/resample.py b/modules/alias_free_torch/resample.py
new file mode 100644
index 0000000000000000000000000000000000000000..dd35bc32b4b30a293ded20f4a66c1578ffbfc54e
--- /dev/null
+++ b/modules/alias_free_torch/resample.py
@@ -0,0 +1,57 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+
+import torch.nn as nn
+from torch.nn import functional as F
+from .filter import LowPassFilter1d
+from .filter import kaiser_sinc_filter1d
+
+
+class UpSample1d(nn.Module):
+ def __init__(self, ratio=2, kernel_size=None):
+ super().__init__()
+ self.ratio = ratio
+ self.kernel_size = (
+ int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+ )
+ self.stride = ratio
+ self.pad = self.kernel_size // ratio - 1
+ self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+ self.pad_right = (
+ self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+ )
+ filter = kaiser_sinc_filter1d(
+ cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size
+ )
+ self.register_buffer("filter", filter)
+
+ # x: [B, C, T]
+ def forward(self, x):
+ _, C, _ = x.shape
+
+ x = F.pad(x, (self.pad, self.pad), mode="replicate")
+ x = self.ratio * F.conv_transpose1d(
+ x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C
+ )
+ x = x[..., self.pad_left : -self.pad_right]
+
+ return x
+
+
+class DownSample1d(nn.Module):
+ def __init__(self, ratio=2, kernel_size=None):
+ super().__init__()
+ self.ratio = ratio
+ self.kernel_size = (
+ int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+ )
+ self.lowpass = LowPassFilter1d(
+ cutoff=0.5 / ratio,
+ half_width=0.6 / ratio,
+ stride=ratio,
+ kernel_size=self.kernel_size,
+ )
+
+ def forward(self, x):
+ xx = self.lowpass(x)
+
+ return xx
diff --git a/modules/astral_quantization/__pycache__/bsq.cpython-310.pyc b/modules/astral_quantization/__pycache__/bsq.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1b130ed261bf95a37afe14a33f4c91c645a70c7a
Binary files /dev/null and b/modules/astral_quantization/__pycache__/bsq.cpython-310.pyc differ
diff --git a/modules/astral_quantization/__pycache__/convnext.cpython-310.pyc b/modules/astral_quantization/__pycache__/convnext.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e49d9274013448f41582eeff7f8ed84ffa40c196
Binary files /dev/null and b/modules/astral_quantization/__pycache__/convnext.cpython-310.pyc differ
diff --git a/modules/astral_quantization/__pycache__/default_model.cpython-310.pyc b/modules/astral_quantization/__pycache__/default_model.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1a4da7177dacdfc5546f7aa70cf3bc0370475205
Binary files /dev/null and b/modules/astral_quantization/__pycache__/default_model.cpython-310.pyc differ
diff --git a/modules/astral_quantization/bsq.py b/modules/astral_quantization/bsq.py
new file mode 100644
index 0000000000000000000000000000000000000000..1b70f3401dafa187fc19666d31a0877f533abe33
--- /dev/null
+++ b/modules/astral_quantization/bsq.py
@@ -0,0 +1,569 @@
+"""
+Lookup Free Quantization
+Proposed in https://arxiv.org/abs/2310.05737
+
+In the simplest setup, each dimension is quantized into {-1, 1}.
+An entropy penalty is used to encourage utilization.
+"""
+
+from math import log2, ceil
+from functools import partial, cache
+from collections import namedtuple
+from contextlib import nullcontext
+
+import torch.distributed as dist
+from torch.distributed import nn as dist_nn
+
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+from torch.nn import Module
+from torch.amp import autocast
+
+from einops import rearrange, reduce, pack, unpack
+
+# constants
+
+Return = namedtuple('Return', ['quantized', 'indices', 'entropy_aux_loss'])
+
+LossBreakdown = namedtuple('LossBreakdown', ['per_sample_entropy', 'batch_entropy', 'commitment'])
+
+# distributed helpers
+
+@cache
+def is_distributed():
+ return dist.is_initialized() and dist.get_world_size() > 1
+
+def maybe_distributed_mean(t):
+ if not is_distributed():
+ return t
+
+ dist_nn.all_reduce(t)
+ t = t / dist.get_world_size()
+ return t
+
+# helper functions
+
+def exists(v):
+ return v is not None
+
+def identity(t):
+ return t
+
+def default(*args):
+ for arg in args:
+ if exists(arg):
+ return arg() if callable(arg) else arg
+ return None
+
+def pack_one(t, pattern):
+ return pack([t], pattern)
+
+def unpack_one(t, ps, pattern):
+ return unpack(t, ps, pattern)[0]
+
+def l2norm(t):
+ return F.normalize(t, dim = -1)
+
+# entropy
+
+def log(t, eps = 1e-5):
+ return t.clamp(min = eps).log()
+
+def entropy(prob):
+ return (-prob * log(prob)).sum(dim=-1)
+
+# cosine sim linear
+
+class CosineSimLinear(Module):
+ def __init__(
+ self,
+ dim_in,
+ dim_out,
+ scale = 1.
+ ):
+ super().__init__()
+ self.scale = scale
+ self.weight = nn.Parameter(torch.randn(dim_in, dim_out))
+
+ def forward(self, x):
+ x = F.normalize(x, dim = -1)
+ w = F.normalize(self.weight, dim = 0)
+ return (x @ w) * self.scale
+
+def soft_entropy_loss(u, tau=1.0, gamma=1.0):
+ """
+ Compute the soft entropy loss for Binary Spherical Quantization (BSQ).
+
+ Args:
+ u (torch.Tensor): Input latent embeddings of shape (batch_size, L).
+ tau (float): Temperature scaling factor.
+ gamma (float): Weight for the second entropy term.
+
+ Returns:
+ torch.Tensor: Soft entropy loss.
+ """
+ # Binary quantization: Generate implicit codebook corners
+ L = u.size(1) # Dimensionality of codebook
+ corners = torch.tensor([-1.0, 1.0], device=u.device) / (L**0.5)
+
+ # Compute soft quantization probabilities for all dimensions
+ # q_hat(c|u) for each dimension
+ prob_matrix = torch.sigmoid(2 * tau * corners.unsqueeze(1) * u.unsqueeze(2)) # Shape: (batch_size, L, 2)
+
+ # Entropy of q_hat(c|u) (independent along each dimension)
+ entropy_per_dim = -torch.sum(prob_matrix * prob_matrix.log(), dim=-1) # Shape: (batch_size, L)
+ entropy_term1 = entropy_per_dim.mean()
+
+ # Expected probabilities for dataset entropy (approximation)
+ expected_probs = prob_matrix.mean(dim=0) # Mean across batch, shape: (L, 2)
+ entropy_term2 = -torch.sum(expected_probs * expected_probs.log(), dim=-1).mean()
+
+ # Final entropy loss
+ loss = entropy_term1 - gamma * entropy_term2
+ return loss
+
+# class
+
+class BinarySphericalQuantize(Module):
+ def __init__(
+ self,
+ *,
+ dim = None,
+ codebook_size = None,
+ entropy_loss_weight = 0.1,
+ commitment_loss_weight = 0.,
+ diversity_gamma = 1.,
+ straight_through_activation = nn.Identity(),
+ num_codebooks = 1,
+ keep_num_codebooks_dim = None,
+ codebook_scale = 1., # for residual LFQ, codebook scaled down by 2x at each layer
+ frac_per_sample_entropy = 0.25, # make less than 1. to only use a random fraction of the probs for per sample entropy
+ has_projections = None,
+ projection_has_bias = True,
+ soft_clamp_input_value = None,
+ cosine_sim_project_in = False,
+ cosine_sim_project_in_scale = None,
+ channel_first = None,
+ experimental_softplus_entropy_loss = False,
+ entropy_loss_offset = 5., # how much to shift the loss before softplus
+ spherical = True, # from https://arxiv.org/abs/2406.07548
+ force_quantization_f32 = True, # will force the quantization step to be full precision
+ enable_entropy_loss = True,
+ soft_entropy_loss = True,
+ ):
+ super().__init__()
+
+ # some assert validations
+
+ assert exists(dim) or exists(codebook_size), 'either dim or codebook_size must be specified for LFQ'
+ assert not exists(codebook_size) or log2(codebook_size).is_integer(), f'your codebook size must be a power of 2 for lookup free quantization (suggested {2 ** ceil(log2(codebook_size))})'
+
+ codebook_size = default(codebook_size, lambda: 2 ** dim)
+ self.codebook_size = codebook_size
+
+ codebook_dim = int(log2(codebook_size))
+ codebook_dims = codebook_dim * num_codebooks
+ dim = default(dim, codebook_dims)
+
+ has_projections = default(has_projections, dim != codebook_dims)
+
+ if cosine_sim_project_in:
+ cosine_sim_project_in = default(cosine_sim_project_in_scale, codebook_scale)
+ project_in_klass = partial(CosineSimLinear, scale = cosine_sim_project_in)
+ else:
+ project_in_klass = partial(nn.Linear, bias = projection_has_bias)
+
+ self.project_in = project_in_klass(dim, codebook_dims) if has_projections else nn.Identity()
+ self.project_out = nn.Linear(codebook_dims, dim, bias = projection_has_bias) if has_projections else nn.Identity()
+ self.has_projections = has_projections
+
+ self.dim = dim
+ self.codebook_dim = codebook_dim
+ self.num_codebooks = num_codebooks
+
+ keep_num_codebooks_dim = default(keep_num_codebooks_dim, num_codebooks > 1)
+ assert not (num_codebooks > 1 and not keep_num_codebooks_dim)
+ self.keep_num_codebooks_dim = keep_num_codebooks_dim
+
+ # channel first
+
+ self.channel_first = channel_first
+
+ # straight through activation
+
+ self.activation = straight_through_activation
+
+ # whether to use BSQ (binary spherical quantization)
+
+ self.spherical = spherical
+ self.maybe_l2norm = (lambda t: l2norm(t) * self.codebook_scale) if spherical else identity
+
+ # entropy aux loss related weights
+
+ assert 0 < frac_per_sample_entropy <= 1.
+ self.frac_per_sample_entropy = frac_per_sample_entropy
+
+ self.diversity_gamma = diversity_gamma
+ self.entropy_loss_weight = entropy_loss_weight
+
+ # codebook scale
+
+ self.codebook_scale = codebook_scale
+
+ # commitment loss
+
+ self.commitment_loss_weight = commitment_loss_weight
+
+ # whether to soft clamp the input value from -value to value
+
+ self.soft_clamp_input_value = soft_clamp_input_value
+ assert not exists(soft_clamp_input_value) or soft_clamp_input_value >= codebook_scale
+
+ # whether to make the entropy loss positive through a softplus (experimental, please report if this worked or not in discussions)
+
+ self.entropy_loss_offset = entropy_loss_offset
+ self.experimental_softplus_entropy_loss = experimental_softplus_entropy_loss
+
+ # for no auxiliary loss, during inference
+
+ self.register_buffer('mask', 2 ** torch.arange(codebook_dim - 1, -1, -1))
+ self.register_buffer('zero', torch.tensor(0.), persistent = False)
+
+ # whether to force quantization step to be f32
+
+ self.force_quantization_f32 = force_quantization_f32
+
+ # codes
+ self.enable_entropy_loss = enable_entropy_loss
+ self.soft_entropy_loss = soft_entropy_loss
+ if codebook_size <= 100000:
+ all_codes = torch.arange(codebook_size)
+ bits = ((all_codes[..., None].int() & self.mask) != 0).float()
+ codebook = self.bits_to_codes(bits)
+
+ self.register_buffer('codebook', codebook.float(), persistent = False)
+ else:
+ all_codes = torch.arange(pow(2, 16))
+ mask = 2 ** torch.arange(16 - 1, -1, -1)
+ bits = ((all_codes[..., None].int() & mask) != 0).float()
+ codebook = self.bits_to_codes(bits)
+
+ self.register_buffer('codebook', codebook.float(), persistent = False)
+
+ def bits_to_codes(self, bits):
+ return bits * self.codebook_scale * 2 - self.codebook_scale
+
+ @property
+ def dtype(self):
+ return self.codebook.dtype
+
+ def indices_to_codes(
+ self,
+ indices,
+ project_out = True
+ ):
+ is_img_or_video = indices.ndim >= (3 + int(self.keep_num_codebooks_dim))
+ should_transpose = default(self.channel_first, is_img_or_video)
+
+ if not self.keep_num_codebooks_dim:
+ indices = rearrange(indices, '... -> ... 1')
+
+ # indices to codes, which are bits of either -1 or 1
+
+ bits = ((indices[..., None].int() & self.mask) != 0).to(self.dtype)
+
+ codes = self.bits_to_codes(bits)
+
+ codes = self.maybe_l2norm(codes)
+
+ codes = rearrange(codes, '... c d -> ... (c d)')
+
+ # whether to project codes out to original dimensions
+ # if the input feature dimensions were not log2(codebook size)
+
+ if project_out:
+ codes = self.project_out(codes)
+
+ # rearrange codes back to original shape
+
+ if should_transpose:
+ codes = rearrange(codes, 'b ... d -> b d ...')
+
+ return codes
+
+ def bits_to_z(self, bits):
+ # assert bits must contain only -1 and 1
+ assert torch.all(bits.abs() == 1)
+ quantized = bits.float()
+ quantized = self.maybe_l2norm(quantized)
+ z = self.project_out(quantized)
+ return z
+
+ def forward(
+ self,
+ x,
+ inv_temperature = 100.,
+ return_loss_breakdown = False,
+ mask = None,
+ return_bits = False
+ ):
+ """
+ einstein notation
+ b - batch
+ n - sequence (or flattened spatial dimensions)
+ d - feature dimension, which is also log2(codebook size)
+ c - number of codebook dim
+ """
+
+ is_img_or_video = x.ndim >= 4
+ should_transpose = default(self.channel_first, is_img_or_video)
+
+ # standardize image or video into (batch, seq, dimension)
+
+ if should_transpose:
+ x = rearrange(x, 'b d ... -> b ... d')
+ x, ps = pack_one(x, 'b * d')
+
+ assert x.shape[-1] == self.dim, f'expected dimension of {self.dim} but received {x.shape[-1]}'
+
+ x = self.project_in(x)
+
+ # maybe soft clamp
+
+ if exists(self.soft_clamp_input_value):
+ clamp_value = self.soft_clamp_input_value
+ x = (x / clamp_value).tanh() * clamp_value
+
+ # split out number of codebooks
+
+ x = rearrange(x, 'b n (c d) -> b n c d', c = self.num_codebooks)
+
+ # maybe l2norm
+
+ x = self.maybe_l2norm(x)
+
+ # whether to force quantization step to be full precision or not
+
+ force_f32 = self.force_quantization_f32
+
+ quantization_context = partial(autocast, 'cuda', enabled = False) if force_f32 else nullcontext
+
+ with quantization_context():
+
+ if force_f32:
+ orig_dtype = x.dtype
+ x = x.float()
+
+ # quantize by eq 3.
+
+ original_input = x
+
+ codebook_value = torch.ones_like(x) * self.codebook_scale
+ quantized = torch.where(x > 0, codebook_value, -codebook_value)
+ if return_bits:
+ return quantized
+
+ # calculate indices
+
+ indices = reduce((quantized > 0).int() * self.mask.int(), 'b n c d -> b n c', 'sum')
+
+ # maybe l2norm
+
+ quantized = self.maybe_l2norm(quantized)
+
+ # use straight-through gradients (optionally with custom activation fn) if training
+
+ if self.training:
+ x = self.activation(x)
+ x = x + (quantized - x).detach()
+ else:
+ x = quantized
+
+ # entropy aux loss
+ if self.soft_entropy_loss:
+ entropy_aux_loss = soft_entropy_loss(x, tau=1.0, gamma=1.0)
+ elif self.training and self.enable_entropy_loss:
+
+ if force_f32:
+ codebook = self.codebook.float()
+
+ codebook = self.maybe_l2norm(codebook)
+
+ # whether to only use a fraction of probs, for reducing memory
+
+ if self.frac_per_sample_entropy < 1.:
+ # account for mask
+ if exists(mask):
+ original_input = original_input[mask]
+ original_input = rearrange(original_input, 'b n ... -> (b n) ...')
+
+ rand_mask = torch.randn(self.codebook_dim).argsort(dim = -1) < 16
+
+ sampled_input = original_input[..., rand_mask]
+
+ sampled_distance = -2 * einsum('... i d, j d -> ... i j', sampled_input, codebook)
+
+ sampled_prob = (-sampled_distance * inv_temperature).softmax(dim = -1)
+
+ per_sample_probs = sampled_prob
+ else:
+ if exists(mask):
+ original_input = original_input[mask]
+ original_input = rearrange(original_input, 'b n ... -> (b n) ...')
+ # the same as euclidean distance up to a constant
+ distance = -2 * einsum('... i d, j d -> ... i j', original_input, codebook)
+
+ prob = (-distance * inv_temperature).softmax(dim = -1)
+
+ per_sample_probs = prob
+
+ # calculate per sample entropy
+
+ per_sample_entropy = entropy(per_sample_probs).mean()
+
+ # distribution over all available tokens in the batch
+
+ avg_prob = reduce(per_sample_probs, '... c d -> c d', 'mean')
+
+ avg_prob = maybe_distributed_mean(avg_prob)
+
+ codebook_entropy = entropy(avg_prob).mean()
+
+ # 1. entropy will be nudged to be low for each code, to encourage the network to output confident predictions
+ # 2. codebook entropy will be nudged to be high, to encourage all codes to be uniformly used within the batch
+
+ entropy_aux_loss = per_sample_entropy - self.diversity_gamma * codebook_entropy
+ else:
+ # if not training, just return dummy 0
+ entropy_aux_loss = per_sample_entropy = codebook_entropy = self.zero
+
+ # whether to make the entropy loss positive or not through a (shifted) softplus
+
+ if self.training and self.experimental_softplus_entropy_loss:
+ entropy_aux_loss = F.softplus(entropy_aux_loss + self.entropy_loss_offset)
+
+ # commit loss
+
+ if self.training and self.commitment_loss_weight > 0.:
+
+ commit_loss = F.mse_loss(original_input, quantized.detach(), reduction = 'none')
+
+ if exists(mask):
+ commit_loss = commit_loss[mask]
+
+ commit_loss = commit_loss.mean()
+ else:
+ commit_loss = self.zero
+
+ # input back to original dtype if needed
+
+ if force_f32:
+ x = x.type(orig_dtype)
+
+ # merge back codebook dim
+
+ x = rearrange(x, 'b n c d -> b n (c d)')
+
+ # project out to feature dimension if needed
+
+ x = self.project_out(x)
+
+ # reconstitute image or video dimensions
+
+ if should_transpose:
+ x = unpack_one(x, ps, 'b * d')
+ x = rearrange(x, 'b ... d -> b d ...')
+
+ indices = unpack_one(indices, ps, 'b * c')
+
+ # whether to remove single codebook dim
+
+ if not self.keep_num_codebooks_dim:
+ indices = rearrange(indices, '... 1 -> ...')
+
+ # complete aux loss
+
+ aux_loss = entropy_aux_loss * self.entropy_loss_weight + commit_loss * self.commitment_loss_weight
+
+ # returns
+
+ ret = Return(x, indices, aux_loss)
+
+ if not return_loss_breakdown:
+ return ret
+
+ return ret, LossBreakdown(per_sample_entropy, codebook_entropy, commit_loss)
+
+class GroupedResidualBSQ(Module):
+ def __init__(
+ self,
+ *,
+ dim,
+ groups = 1,
+ accept_image_fmap = False,
+ **kwargs
+ ):
+ super().__init__()
+ self.dim = dim
+ self.groups = groups
+ assert (dim % groups) == 0
+ dim_per_group = dim // groups
+
+ self.accept_image_fmap = accept_image_fmap
+
+ self.rvqs = nn.ModuleList([])
+
+ for _ in range(groups):
+ self.rvqs.append(LFQ(
+ dim = dim_per_group,
+ **kwargs
+ ))
+
+ self.codebook_size = self.rvqs[0].codebook_size
+
+ @property
+ def codebooks(self):
+ return torch.stack(tuple(rvq.codebooks for rvq in self.rvqs))
+
+ @property
+ def split_dim(self):
+ return 1 if self.accept_image_fmap else -1
+
+ def get_codes_from_indices(self, indices):
+ codes = tuple(rvq.get_codes_from_indices(chunk_indices) for rvq, chunk_indices in zip(self.rvqs, indices))
+ return torch.stack(codes)
+
+ def get_output_from_indices(self, indices):
+ outputs = tuple(rvq.get_output_from_indices(chunk_indices) for rvq, chunk_indices in zip(self.rvqs, indices))
+ return torch.cat(outputs, dim = self.split_dim)
+
+ def forward(
+ self,
+ x,
+ return_all_codes = False
+ ):
+ shape, split_dim = x.shape, self.split_dim
+ assert shape[split_dim] == self.dim
+
+ # split the feature dimension into groups
+
+ x = x.chunk(self.groups, dim = split_dim)
+
+ forward_kwargs = dict(
+ )
+
+ # invoke residual vq on each group
+
+ out = tuple(rvq(chunk, **forward_kwargs) for rvq, chunk in zip(self.rvqs, x))
+ out = tuple(zip(*out))
+
+ # otherwise, get all the zipped outputs and combine them
+
+ quantized, all_indices, *maybe_aux_loss = out
+
+ quantized = torch.cat(quantized, dim = split_dim)
+ all_indices = torch.stack(all_indices)
+
+ ret = (quantized, all_indices, *maybe_aux_loss)
+ return ret
diff --git a/modules/astral_quantization/convnext.py b/modules/astral_quantization/convnext.py
new file mode 100644
index 0000000000000000000000000000000000000000..7bef9e282ac332b7169339eeda082d0581d739d1
--- /dev/null
+++ b/modules/astral_quantization/convnext.py
@@ -0,0 +1,209 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import List
+
+
+class ConvNextV2LayerNorm(nn.Module):
+ r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
+ The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height,
+ width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width).
+ """
+
+ def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(normalized_shape))
+ self.bias = nn.Parameter(torch.zeros(normalized_shape))
+ self.eps = eps
+ self.data_format = data_format
+ if self.data_format not in ["channels_last", "channels_first"]:
+ raise NotImplementedError(f"Unsupported data format: {self.data_format}")
+ self.normalized_shape = (normalized_shape,)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if self.data_format == "channels_last":
+ x = torch.nn.functional.layer_norm(
+ x, self.normalized_shape, self.weight, self.bias, self.eps
+ )
+ elif self.data_format == "channels_first":
+ input_dtype = x.dtype
+ x = x.float()
+ u = x.mean(1, keepdim=True)
+ s = (x - u).pow(2).mean(1, keepdim=True)
+ x = (x - u) / torch.sqrt(s + self.eps)
+ x = x.to(dtype=input_dtype)
+ x = self.weight[None, :, None] * x + self.bias[None, :, None]
+ return x
+
+
+class GRN(nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+ self.gamma = nn.Parameter(torch.zeros(1, 1, dim))
+ self.beta = nn.Parameter(torch.zeros(1, 1, dim))
+
+ def forward(self, x):
+ Gx = torch.norm(x, p=2, dim=1, keepdim=True)
+ Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+ return self.gamma * (x * Nx) + self.beta + x
+
+class InterpolationLayer(nn.Module):
+ def __init__(self, ): # this is a default of 1 / 50 * (44100 / 512) / 4
+ super().__init__()
+ pass
+
+ def forward(self, x: torch.Tensor, target_len: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+ x = F.interpolate(x, size=target_len, mode='linear')
+ return x
+
+class ConvNeXtV2Stage(nn.Module):
+ def __init__(
+ self,
+ dim: int = 512,
+ intermediate_dim: int = 2048,
+ num_blocks: int = 1,
+ dilation: int = 1,
+ downsample_layer_indices: List[int] = None,
+ downsample_factors: List[int] = None,
+ upsample_layer_indices: List[int] = None,
+ upsample_factors: List[int] = None,
+ interpolation_layer_indices: List[int] = None,
+ input_dim: int = None,
+ output_dim: int = None,
+ gin_channels: int = 0,
+ ):
+ super().__init__()
+ # maybe downsample layers
+ if downsample_layer_indices is not None:
+ assert downsample_factors is not None
+ self.downsample_blocks = nn.ModuleList(
+ [
+ nn.Sequential(
+ ConvNextV2LayerNorm(dim, data_format="channels_first"),
+ nn.Conv1d(
+ dim, dim, kernel_size=downsample_factor, stride=downsample_factor
+ ),
+ ) for _, downsample_factor in zip(downsample_layer_indices, downsample_factors)
+ ]
+ )
+ self.downsample_layer_indices = downsample_layer_indices
+ else:
+ self.downsample_blocks = nn.ModuleList()
+ self.downsample_layer_indices = []
+
+ # maybe upsample layers
+ if upsample_layer_indices is not None:
+ assert upsample_factors is not None
+ self.upsample_blocks = nn.ModuleList(
+ [
+ nn.Sequential(
+ ConvNextV2LayerNorm(dim, data_format="channels_first"),
+ nn.ConvTranspose1d(
+ dim, dim, kernel_size=upsample_factor, stride=upsample_factor
+ ),
+ ) for _, upsample_factor in zip(upsample_layer_indices, upsample_factors)
+ ]
+ )
+ self.upsample_layer_indices = upsample_layer_indices
+ else:
+ self.upsample_blocks = nn.ModuleList()
+ self.upsample_layer_indices = []
+
+ # maybe interpolation layers
+ if interpolation_layer_indices is not None:
+ self.interpolation_blocks = nn.ModuleList(
+ [
+ InterpolationLayer()
+ for _ in interpolation_layer_indices
+ ]
+ )
+ self.interpolation_layer_indices = interpolation_layer_indices
+ else:
+ self.interpolation_blocks = nn.ModuleList()
+ self.interpolation_layer_indices = []
+
+ # main blocks
+ self.blocks = nn.ModuleList(
+ [
+ ConvNeXtV2Block(
+ dim=dim,
+ intermediate_dim=intermediate_dim,
+ dilation=dilation,
+ )
+ for _ in range(num_blocks)
+ ]
+ )
+ # maybe input and output projections
+ if input_dim is not None and input_dim != dim:
+ self.input_projection = nn.Conv1d(input_dim, dim, kernel_size=1)
+ else:
+ self.input_projection = nn.Identity()
+ if output_dim is not None and output_dim != dim:
+ self.output_projection = nn.Conv1d(dim, output_dim, kernel_size=1)
+ else:
+ self.output_projection = nn.Identity()
+
+ if gin_channels > 0:
+ self.gin = nn.Conv1d(gin_channels, dim, kernel_size=1)
+
+ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+ x = self.input_projection(x) # B, D, T
+ if hasattr(self, 'gin'):
+ g = kwargs['g']
+ x = x + self.gin(g)
+ # pad to a multiple of cumprod(downsample_factors)
+ if len(self.downsample_blocks) > 0:
+ downsample_factor = 1
+ for factor in self.downsample_blocks:
+ downsample_factor *= factor[1].stride[0]
+ pad_len = downsample_factor - x.size(-1) % downsample_factor
+ if pad_len > 0:
+ x = torch.cat([x, torch.zeros_like(x[:, :, :pad_len])], dim=-1)
+
+ # main blocks
+ for layer_idx, block in enumerate(self.blocks):
+ if layer_idx in self.downsample_layer_indices:
+ x = self.downsample_blocks[self.downsample_layer_indices.index(layer_idx)](x)
+ if layer_idx in self.upsample_layer_indices:
+ x = self.upsample_blocks[self.upsample_layer_indices.index(layer_idx)](x)
+ if layer_idx in self.interpolation_layer_indices:
+ x = self.interpolation_blocks[self.interpolation_layer_indices.index(layer_idx)](x, target_len=kwargs['target_len'])
+ x = block(x)
+ x = self.output_projection(x)
+ return x
+
+ def setup_caches(self, *args, **kwargs):
+ pass
+
+
+class ConvNeXtV2Block(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ intermediate_dim: int,
+ dilation: int = 1,
+ ):
+ super().__init__()
+ padding = (dilation * (7 - 1)) // 2
+ self.dwconv = nn.Conv1d(
+ dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation
+ ) # depthwise conv
+ self.norm = ConvNextV2LayerNorm(dim, data_format="channels_first")
+ self.pwconv1 = nn.Linear(
+ dim, intermediate_dim
+ ) # pointwise/1x1 convs, implemented with linear layers
+ self.act = nn.GELU()
+ self.grn = GRN(intermediate_dim)
+ self.pwconv2 = nn.Linear(intermediate_dim, dim)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ residual = x
+ x = self.dwconv(x)
+ x = self.norm(x)
+ x = x.transpose(1, 2) # b d n -> b n d
+ x = self.pwconv1(x)
+ x = self.act(x)
+ x = self.grn(x)
+ x = self.pwconv2(x)
+ x = x.transpose(1, 2) # b n d -> b d n
+ return residual + x
\ No newline at end of file
diff --git a/modules/astral_quantization/default_model.py b/modules/astral_quantization/default_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..ca50d1cb1b177a89955c9f13451e5ed159b753b5
--- /dev/null
+++ b/modules/astral_quantization/default_model.py
@@ -0,0 +1,73 @@
+import torch
+from transformers import AutoTokenizer, AutoModel, Wav2Vec2FeatureExtractor
+
+class AstralQuantizer(torch.nn.Module):
+ def __init__(
+ self,
+ tokenizer_name: str,
+ ssl_model_name: str,
+ ssl_output_layer: int,
+ encoder: torch.nn.Module,
+ quantizer: torch.nn.Module,
+ skip_ssl: bool = False,
+ ):
+ super().__init__()
+ self.encoder = encoder
+ self.quantizer = quantizer
+ self.tokenizer_name = tokenizer_name
+ self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+
+ # Load SSL model from Huggingface
+ self.ssl_model_name = ssl_model_name
+ self.ssl_output_layer = ssl_output_layer
+ self.ssl_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(ssl_model_name)
+
+ if skip_ssl: # in case the same SSL model has been loaded somewhere else
+ self.ssl_model = None
+ else:
+ self.ssl_model = AutoModel.from_pretrained(ssl_model_name).eval()
+ self.ssl_model.encoder.layers = self.ssl_model.encoder.layers[:ssl_output_layer]
+ self.ssl_model.encoder.layer_norm = torch.nn.Identity()
+
+ def load_separate_checkpoint(self, checkpoint_path):
+ params = torch.load(checkpoint_path, map_location='cpu')['net']
+ for key in params.keys():
+ for k in list(params[key].keys()):
+ if k.startswith("module."):
+ params[key][k[len("module."):]] = params[key][k]
+ del params[key][k]
+ self.encoder.load_state_dict(params['encoder'])
+ self.quantizer.load_state_dict(params['vq'])
+ if self.decoder is not None:
+ self.decoder.load_state_dict(params['decoder'])
+ if self.asr_decoder is not None:
+ self.asr_decoder.load_state_dict(params['predictor'], strict=False)
+
+ def forward(self, waves_16k, wave_16k_lens, ssl_model=None):
+ ssl_fn = self.ssl_model if self.ssl_model else ssl_model
+ assert ssl_fn is not None, "In case in-class SSL model loading is skipped, external ssl_model must be provided"
+ waves_16k_input_list = [
+ waves_16k[bib, :wave_16k_lens[bib]].cpu().numpy()
+ for bib in range(len(waves_16k))
+ ]
+ alt_inputs = self.ssl_feature_extractor(
+ waves_16k_input_list,
+ return_tensors='pt',
+ return_attention_mask=True,
+ padding=True,
+ sampling_rate=16000
+ ).to(waves_16k.device)
+ feature_lens = alt_inputs.data['attention_mask'].sum(-1) // 320 # frame rate of hubert is 50 Hz
+
+ outputs = ssl_fn(
+ alt_inputs.input_values,
+ attention_mask=alt_inputs.attention_mask,
+ )
+ last_hidden_states = outputs.last_hidden_state
+ last_hidden_states = last_hidden_states[:, :feature_lens.max(), :]
+ feature_lens = feature_lens.clamp(max=last_hidden_states.size(1))
+ last_hidden_states = last_hidden_states.transpose(1, 2)
+ x_hidden = self.encoder(last_hidden_states, feature_lens)
+ x_hidden = x_hidden.transpose(1, 2)
+ x_quantized, indices = self.quantizer(x_hidden)[:2]
+ return x_quantized, indices, feature_lens
\ No newline at end of file
diff --git a/modules/astral_quantization/transformer.py b/modules/astral_quantization/transformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..015ec417d02e3442d8ea120c3802807c73e89bb4
--- /dev/null
+++ b/modules/astral_quantization/transformer.py
@@ -0,0 +1,254 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import functional as F
+import time
+
+def find_multiple(n: int, k: int) -> int:
+ if n % k == 0:
+ return n
+ return n + k - (n % k)
+
+class AdaptiveLayerNorm(nn.Module):
+ r"""Adaptive Layer Normalization"""
+
+ def __init__(self, d_model, norm) -> None:
+ super(AdaptiveLayerNorm, self).__init__()
+ self.project_layer = nn.Linear(d_model, 2 * d_model)
+ self.norm = norm
+ self.d_model = d_model
+ self.eps = self.norm.eps
+
+ def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor:
+ if embedding is None:
+ return self.norm(input)
+ weight, bias = torch.split(
+ self.project_layer(embedding),
+ split_size_or_sections=self.d_model,
+ dim=-1,
+ )
+ return weight * self.norm(input) + bias
+
+
+@dataclass
+class ModelArgs:
+ block_size: int = 2048
+ vocab_size: int = 32000
+ n_layer: int = 32
+ n_head: int = 32
+ dim: int = 4096
+ intermediate_size: int = None
+ n_local_heads: int = -1
+ head_dim: int = 64
+ rope_base: float = 10000
+ norm_eps: float = 1e-5
+ has_cross_attention: bool = False
+ context_dim: int = 0
+ is_causal: bool = False
+ dropout_rate: float = 0.1
+ attn_dropout_rate: float = 0.1
+
+ def __post_init__(self):
+ if self.n_local_heads == -1:
+ self.n_local_heads = self.n_head
+ if self.intermediate_size is None:
+ hidden_dim = 4 * self.dim
+ n_hidden = int(2 * hidden_dim / 3)
+ self.intermediate_size = find_multiple(n_hidden, 256)
+ # self.head_dim = self.dim // self.n_head
+
+class Transformer(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.config = config
+
+ self.layers = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer))
+ self.norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+
+ self.max_batch_size = -1
+ self.max_seq_length = config.block_size
+ freqs_cis = precompute_freqs_cis(self.config.block_size, self.config.head_dim,
+ self.config.rope_base)
+ self.register_buffer("freqs_cis", freqs_cis)
+
+ causal_mask = torch.tril(
+ torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool)
+ )
+ self.register_buffer("causal_mask", causal_mask)
+
+ def forward(self,
+ x: Tensor,
+ c: Tensor,
+ input_pos: Optional[Tensor] = None,
+ mask: Optional[Tensor] = None,
+ context: Optional[Tensor] = None,
+ context_input_pos: Optional[Tensor] = None,
+ cross_attention_mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ if mask is None:
+ mask = self.causal_mask[:x.size(1), :x.size(1)]
+ else:
+ mask = mask[..., input_pos]
+ freqs_cis = self.freqs_cis[input_pos]
+ if context is not None:
+ context_freqs_cis = self.freqs_cis[context_input_pos]
+ else:
+ context_freqs_cis = None
+ skip_in_x_list = []
+ for i, layer in enumerate(self.layers):
+ x = layer(x, c, freqs_cis, mask, context, context_freqs_cis, cross_attention_mask)
+ x = self.norm(x, c)
+ return x
+
+
+class TransformerBlock(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.attention = Attention(config)
+ self.feed_forward = FeedForward(config)
+ self.ffn_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+ self.attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+
+ if config.has_cross_attention:
+ self.has_cross_attention = True
+ self.cross_attention = Attention(config, is_cross_attention=True)
+ self.cross_attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+ else:
+ self.has_cross_attention = False
+
+ def forward(self,
+ x: Tensor,
+ c: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ context: Optional[Tensor] = None,
+ context_freqs_cis: Optional[Tensor] = None,
+ cross_attention_mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ #time_attn_start = time.time()
+ h = x + self.attention(self.attention_norm(x, c), freqs_cis, mask)
+ #print(f"time take for attention of sequence length {x.shape[1]} is {time.time() - time_attn_start}")
+ if self.has_cross_attention:
+ h = h + self.cross_attention(self.cross_attention_norm(h, c), freqs_cis, cross_attention_mask, context, context_freqs_cis)
+ out = h + self.feed_forward(self.ffn_norm(h, c))
+ return out
+
+
+class Attention(nn.Module):
+ def __init__(self, config: ModelArgs, is_cross_attention: bool = False):
+ super().__init__()
+ assert config.dim % config.n_head == 0
+
+ total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim
+ # key, query, value projections for all heads, but in a batch
+ if is_cross_attention:
+ self.wq = nn.Linear(config.dim, config.n_head * config.head_dim, bias=False)
+ self.wkv = nn.Linear(config.context_dim, 2 * config.n_local_heads * config.head_dim, bias=False)
+ else:
+ self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False)
+ self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False)
+ self.kv_cache = None
+
+ self.n_head = config.n_head
+ self.head_dim = config.head_dim
+ self.n_local_heads = config.n_local_heads
+ self.dim = config.dim
+ self.attn_dropout_rate = config.attn_dropout_rate
+
+ def forward(self,
+ x: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ context: Optional[Tensor] = None,
+ context_freqs_cis: Optional[Tensor] = None,
+ ) -> Tensor:
+ bsz, seqlen, _ = x.shape
+
+ kv_size = self.n_local_heads * self.head_dim
+ if context is None:
+ q, k, v = self.wqkv(x).split([kv_size, kv_size, kv_size], dim=-1)
+ context_seqlen = seqlen
+ else:
+ q = self.wq(x)
+ k, v = self.wkv(context).split([kv_size, kv_size], dim=-1)
+ context_seqlen = context.shape[1]
+
+ q = q.view(bsz, seqlen, self.n_head, self.head_dim)
+ k = k.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+ v = v.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+
+ q = apply_rotary_emb(q, freqs_cis)
+ k = apply_rotary_emb(k, context_freqs_cis if context_freqs_cis is not None else freqs_cis)
+
+ q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
+
+ k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=self.attn_dropout_rate if self.training else 0.0)
+
+ y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.head_dim * self.n_head)
+
+ y = self.wo(y)
+ return y
+
+
+class FeedForward(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
+ self.dropout = nn.Dropout(config.dropout_rate)
+
+ def forward(self, x: Tensor) -> Tensor:
+ return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x)))
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-5):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def _norm(self, x):
+ return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+ def forward(self, x: Tensor) -> Tensor:
+ output = self._norm(x.float()).type_as(x)
+ return output * self.weight
+
+
+def precompute_freqs_cis(
+ seq_len: int, n_elem: int, base: int = 10000,
+ dtype: torch.dtype = torch.bfloat16
+) -> Tensor:
+ freqs = 1.0 / (base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem))
+ t = torch.arange(seq_len, device=freqs.device)
+ freqs = torch.outer(t, freqs)
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+ cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
+ return cache.to(dtype=dtype)
+
+
+def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
+ xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+ freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2)
+ x_out2 = torch.stack(
+ [
+ xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],
+ xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],
+ ],
+ -1,
+ )
+
+ x_out2 = x_out2.flatten(3)
+ return x_out2.type_as(x)
+
diff --git a/modules/audio.py b/modules/audio.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae677ffb1c124b557b3dbe0343ae415f5281cddb
--- /dev/null
+++ b/modules/audio.py
@@ -0,0 +1,82 @@
+import numpy as np
+import torch
+import torch.utils.data
+from librosa.filters import mel as librosa_mel_fn
+from scipy.io.wavfile import read
+
+MAX_WAV_VALUE = 32768.0
+
+
+def load_wav(full_path):
+ sampling_rate, data = read(full_path)
+ return data, sampling_rate
+
+
+def dynamic_range_compression(x, C=1, clip_val=1e-5):
+ return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
+
+
+def dynamic_range_decompression(x, C=1):
+ return np.exp(x) / C
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+ return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression_torch(x, C=1):
+ return torch.exp(x) / C
+
+
+def spectral_normalize_torch(magnitudes):
+ output = dynamic_range_compression_torch(magnitudes)
+ return output
+
+
+def spectral_de_normalize_torch(magnitudes):
+ output = dynamic_range_decompression_torch(magnitudes)
+ return output
+
+
+mel_basis = {}
+hann_window = {}
+
+
+def mel_spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False):
+ if torch.min(y) < -1.0:
+ print("min value is ", torch.min(y))
+ if torch.max(y) > 1.0:
+ print("max value is ", torch.max(y))
+
+ global mel_basis, hann_window # pylint: disable=global-statement
+ if f"{str(sampling_rate)}_{str(fmax)}_{str(y.device)}" not in mel_basis:
+ mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
+ mel_basis[str(sampling_rate) + "_" + str(fmax) + "_" + str(y.device)] = torch.from_numpy(mel).float().to(y.device)
+ hann_window[str(sampling_rate) + "_" + str(y.device)] = torch.hann_window(win_size).to(y.device)
+
+ y = torch.nn.functional.pad(
+ y.unsqueeze(1), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), mode="reflect"
+ )
+ y = y.squeeze(1)
+
+ spec = torch.view_as_real(
+ torch.stft(
+ y,
+ n_fft,
+ hop_length=hop_size,
+ win_length=win_size,
+ window=hann_window[str(sampling_rate) + "_" + str(y.device)],
+ center=center,
+ pad_mode="reflect",
+ normalized=False,
+ onesided=True,
+ return_complex=True,
+ )
+ )
+
+ spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))
+
+ spec = torch.matmul(mel_basis[str(sampling_rate) + "_" + str(fmax) + "_" + str(y.device)], spec)
+ spec = spectral_normalize_torch(spec)
+
+ return spec
diff --git a/modules/bigvgan/__pycache__/activations.cpython-310.pyc b/modules/bigvgan/__pycache__/activations.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..eee8dcac54f37639e7d8449cfe49e9a76d24aa98
Binary files /dev/null and b/modules/bigvgan/__pycache__/activations.cpython-310.pyc differ
diff --git a/modules/bigvgan/__pycache__/bigvgan.cpython-310.pyc b/modules/bigvgan/__pycache__/bigvgan.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6662b4a784f2db88191eec18ce00078d5b6c96eb
Binary files /dev/null and b/modules/bigvgan/__pycache__/bigvgan.cpython-310.pyc differ
diff --git a/modules/bigvgan/__pycache__/env.cpython-310.pyc b/modules/bigvgan/__pycache__/env.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ec5a2ad3c6a1f3d4d9d60db17ec78b5d5526f4bc
Binary files /dev/null and b/modules/bigvgan/__pycache__/env.cpython-310.pyc differ
diff --git a/modules/bigvgan/__pycache__/meldataset.cpython-310.pyc b/modules/bigvgan/__pycache__/meldataset.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0523c2c3702e699c82a522c1b33a6e5d67bee591
Binary files /dev/null and b/modules/bigvgan/__pycache__/meldataset.cpython-310.pyc differ
diff --git a/modules/bigvgan/__pycache__/utils.cpython-310.pyc b/modules/bigvgan/__pycache__/utils.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0a62420da71e1718f2e4605c8727f1484eb33ab8
Binary files /dev/null and b/modules/bigvgan/__pycache__/utils.cpython-310.pyc differ
diff --git a/modules/bigvgan/activations.py b/modules/bigvgan/activations.py
new file mode 100644
index 0000000000000000000000000000000000000000..9d5420c1994f064635a3bb283d83f6cec762d542
--- /dev/null
+++ b/modules/bigvgan/activations.py
@@ -0,0 +1,120 @@
+# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license.
+# LICENSE is in incl_licenses directory.
+
+import torch
+from torch import nn, sin, pow
+from torch.nn import Parameter
+
+
+class Snake(nn.Module):
+ '''
+ Implementation of a sine-based periodic activation function
+ Shape:
+ - Input: (B, C, T)
+ - Output: (B, C, T), same shape as the input
+ Parameters:
+ - alpha - trainable parameter
+ References:
+ - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+ https://arxiv.org/abs/2006.08195
+ Examples:
+ >>> a1 = snake(256)
+ >>> x = torch.randn(256)
+ >>> x = a1(x)
+ '''
+ def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):
+ '''
+ Initialization.
+ INPUT:
+ - in_features: shape of the input
+ - alpha: trainable parameter
+ alpha is initialized to 1 by default, higher values = higher-frequency.
+ alpha will be trained along with the rest of your model.
+ '''
+ super(Snake, self).__init__()
+ self.in_features = in_features
+
+ # initialize alpha
+ self.alpha_logscale = alpha_logscale
+ if self.alpha_logscale: # log scale alphas initialized to zeros
+ self.alpha = Parameter(torch.zeros(in_features) * alpha)
+ else: # linear scale alphas initialized to ones
+ self.alpha = Parameter(torch.ones(in_features) * alpha)
+
+ self.alpha.requires_grad = alpha_trainable
+
+ self.no_div_by_zero = 0.000000001
+
+ def forward(self, x):
+ '''
+ Forward pass of the function.
+ Applies the function to the input elementwise.
+ Snake ∶= x + 1/a * sin^2 (xa)
+ '''
+ alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T]
+ if self.alpha_logscale:
+ alpha = torch.exp(alpha)
+ x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+ return x
+
+
+class SnakeBeta(nn.Module):
+ '''
+ A modified Snake function which uses separate parameters for the magnitude of the periodic components
+ Shape:
+ - Input: (B, C, T)
+ - Output: (B, C, T), same shape as the input
+ Parameters:
+ - alpha - trainable parameter that controls frequency
+ - beta - trainable parameter that controls magnitude
+ References:
+ - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+ https://arxiv.org/abs/2006.08195
+ Examples:
+ >>> a1 = snakebeta(256)
+ >>> x = torch.randn(256)
+ >>> x = a1(x)
+ '''
+ def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):
+ '''
+ Initialization.
+ INPUT:
+ - in_features: shape of the input
+ - alpha - trainable parameter that controls frequency
+ - beta - trainable parameter that controls magnitude
+ alpha is initialized to 1 by default, higher values = higher-frequency.
+ beta is initialized to 1 by default, higher values = higher-magnitude.
+ alpha will be trained along with the rest of your model.
+ '''
+ super(SnakeBeta, self).__init__()
+ self.in_features = in_features
+
+ # initialize alpha
+ self.alpha_logscale = alpha_logscale
+ if self.alpha_logscale: # log scale alphas initialized to zeros
+ self.alpha = Parameter(torch.zeros(in_features) * alpha)
+ self.beta = Parameter(torch.zeros(in_features) * alpha)
+ else: # linear scale alphas initialized to ones
+ self.alpha = Parameter(torch.ones(in_features) * alpha)
+ self.beta = Parameter(torch.ones(in_features) * alpha)
+
+ self.alpha.requires_grad = alpha_trainable
+ self.beta.requires_grad = alpha_trainable
+
+ self.no_div_by_zero = 0.000000001
+
+ def forward(self, x):
+ '''
+ Forward pass of the function.
+ Applies the function to the input elementwise.
+ SnakeBeta ∶= x + 1/b * sin^2 (xa)
+ '''
+ alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T]
+ beta = self.beta.unsqueeze(0).unsqueeze(-1)
+ if self.alpha_logscale:
+ alpha = torch.exp(alpha)
+ beta = torch.exp(beta)
+ x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+ return x
\ No newline at end of file
diff --git a/modules/bigvgan/alias_free_activation/cuda/__init__.py b/modules/bigvgan/alias_free_activation/cuda/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/modules/bigvgan/alias_free_activation/cuda/__pycache__/__init__.cpython-310.pyc b/modules/bigvgan/alias_free_activation/cuda/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7692a9a3df20d2fe2b7923398589e05741607fd2
Binary files /dev/null and b/modules/bigvgan/alias_free_activation/cuda/__pycache__/__init__.cpython-310.pyc differ
diff --git a/modules/bigvgan/alias_free_activation/cuda/__pycache__/activation1d.cpython-310.pyc b/modules/bigvgan/alias_free_activation/cuda/__pycache__/activation1d.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1e13bedd060a666d59013e80b43573fad05e3516
Binary files /dev/null and b/modules/bigvgan/alias_free_activation/cuda/__pycache__/activation1d.cpython-310.pyc differ
diff --git a/modules/bigvgan/alias_free_activation/cuda/__pycache__/load.cpython-310.pyc b/modules/bigvgan/alias_free_activation/cuda/__pycache__/load.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..089a26db09f702d6a6984743f4befb52b1c80fbf
Binary files /dev/null and b/modules/bigvgan/alias_free_activation/cuda/__pycache__/load.cpython-310.pyc differ
diff --git a/modules/bigvgan/alias_free_activation/cuda/activation1d.py b/modules/bigvgan/alias_free_activation/cuda/activation1d.py
new file mode 100644
index 0000000000000000000000000000000000000000..76797ef1424b99a160803d53cb6c24fe20599bd2
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/cuda/activation1d.py
@@ -0,0 +1,77 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+import torch
+import torch.nn as nn
+from ..torch.resample import UpSample1d, DownSample1d
+
+# load fused CUDA kernel: this enables importing anti_alias_activation_cuda
+from ..cuda import load
+
+anti_alias_activation_cuda = load.load()
+
+
+class FusedAntiAliasActivation(torch.autograd.Function):
+ """
+ Assumes filter size 12, replication padding on upsampling/downsampling, and logscale alpha/beta parameters as inputs.
+ The hyperparameters are hard-coded in the kernel to maximize speed.
+ NOTE: The fused kenrel is incorrect for Activation1d with different hyperparameters.
+ """
+
+ @staticmethod
+ def forward(ctx, inputs, up_ftr, down_ftr, alpha, beta):
+ activation_results = anti_alias_activation_cuda.forward(
+ inputs, up_ftr, down_ftr, alpha, beta
+ )
+
+ return activation_results
+
+ @staticmethod
+ def backward(ctx, output_grads):
+ raise NotImplementedError
+ return output_grads, None, None
+
+
+class Activation1d(nn.Module):
+ def __init__(
+ self,
+ activation,
+ up_ratio: int = 2,
+ down_ratio: int = 2,
+ up_kernel_size: int = 12,
+ down_kernel_size: int = 12,
+ fused: bool = True,
+ ):
+ super().__init__()
+ self.up_ratio = up_ratio
+ self.down_ratio = down_ratio
+ self.act = activation
+ self.upsample = UpSample1d(up_ratio, up_kernel_size)
+ self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+ self.fused = fused # Whether to use fused CUDA kernel or not
+
+ def forward(self, x):
+ if not self.fused:
+ x = self.upsample(x)
+ x = self.act(x)
+ x = self.downsample(x)
+ return x
+ else:
+ if self.act.__class__.__name__ == "Snake":
+ beta = self.act.alpha.data # Snake uses same params for alpha and beta
+ else:
+ beta = (
+ self.act.beta.data
+ ) # Snakebeta uses different params for alpha and beta
+ alpha = self.act.alpha.data
+ if (
+ not self.act.alpha_logscale
+ ): # Exp baked into cuda kernel, cancel it out with a log
+ alpha = torch.log(alpha)
+ beta = torch.log(beta)
+
+ x = FusedAntiAliasActivation.apply(
+ x, self.upsample.filter, self.downsample.lowpass.filter, alpha, beta
+ )
+ return x
diff --git a/modules/bigvgan/alias_free_activation/cuda/anti_alias_activation.cpp b/modules/bigvgan/alias_free_activation/cuda/anti_alias_activation.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..94fd90da386e66ce12a64ef243e4d125926dfd2a
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/cuda/anti_alias_activation.cpp
@@ -0,0 +1,23 @@
+/* coding=utf-8
+ * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+ #include
+
+extern "C" torch::Tensor fwd_cuda(torch::Tensor const &input, torch::Tensor const &up_filter, torch::Tensor const &down_filter, torch::Tensor const &alpha, torch::Tensor const &beta);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+ m.def("forward", &fwd_cuda, "Anti-Alias Activation forward (CUDA)");
+}
\ No newline at end of file
diff --git a/modules/bigvgan/alias_free_activation/cuda/anti_alias_activation_cuda.cu b/modules/bigvgan/alias_free_activation/cuda/anti_alias_activation_cuda.cu
new file mode 100644
index 0000000000000000000000000000000000000000..7ee97492984a92c753f2357f03e7c04252060824
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/cuda/anti_alias_activation_cuda.cu
@@ -0,0 +1,246 @@
+/* coding=utf-8
+ * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include "type_shim.h"
+#include
+#include
+#include
+#include
+#include
+
+namespace
+{
+ // Hard-coded hyperparameters
+ // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and
+ constexpr int ELEMENTS_PER_LDG_STG = 1; //(WARP_ITERATIONS < 4) ? 1 : 4;
+ constexpr int BUFFER_SIZE = 32;
+ constexpr int FILTER_SIZE = 12;
+ constexpr int HALF_FILTER_SIZE = 6;
+ constexpr int UPSAMPLE_REPLICATION_PAD = 5; // 5 on each side, matching torch impl
+ constexpr int DOWNSAMPLE_REPLICATION_PAD_LEFT = 5; // matching torch impl
+ constexpr int DOWNSAMPLE_REPLICATION_PAD_RIGHT = 6; // matching torch impl
+
+ template
+ __global__ void anti_alias_activation_forward(
+ output_t *dst,
+ const input_t *src,
+ const input_t *up_ftr,
+ const input_t *down_ftr,
+ const input_t *alpha,
+ const input_t *beta,
+ int batch_size,
+ int channels,
+ int seq_len)
+ {
+ // Up and downsample filters
+ input_t up_filter[FILTER_SIZE];
+ input_t down_filter[FILTER_SIZE];
+
+ // Load data from global memory including extra indices reserved for replication paddings
+ input_t elements[2 * FILTER_SIZE + 2 * BUFFER_SIZE + 2 * UPSAMPLE_REPLICATION_PAD] = {0};
+ input_t intermediates[2 * FILTER_SIZE + 2 * BUFFER_SIZE + DOWNSAMPLE_REPLICATION_PAD_LEFT + DOWNSAMPLE_REPLICATION_PAD_RIGHT] = {0};
+
+ // Output stores downsampled output before writing to dst
+ output_t output[BUFFER_SIZE];
+
+ // blockDim/threadIdx = (128, 1, 1)
+ // gridDim/blockIdx = (seq_blocks, channels, batches)
+ int block_offset = (blockIdx.x * 128 * BUFFER_SIZE + seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
+ int local_offset = threadIdx.x * BUFFER_SIZE;
+ int seq_offset = blockIdx.x * 128 * BUFFER_SIZE + local_offset;
+
+ // intermediate have double the seq_len
+ int intermediate_local_offset = threadIdx.x * BUFFER_SIZE * 2;
+ int intermediate_seq_offset = blockIdx.x * 128 * BUFFER_SIZE * 2 + intermediate_local_offset;
+
+ // Get values needed for replication padding before moving pointer
+ const input_t *right_most_pntr = src + (seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
+ input_t seq_left_most_value = right_most_pntr[0];
+ input_t seq_right_most_value = right_most_pntr[seq_len - 1];
+
+ // Move src and dst pointers
+ src += block_offset + local_offset;
+ dst += block_offset + local_offset;
+
+ // Alpha and beta values for snake activatons. Applies exp by default
+ alpha = alpha + blockIdx.y;
+ input_t alpha_val = expf(alpha[0]);
+ beta = beta + blockIdx.y;
+ input_t beta_val = expf(beta[0]);
+
+ #pragma unroll
+ for (int it = 0; it < FILTER_SIZE; it += 1)
+ {
+ up_filter[it] = up_ftr[it];
+ down_filter[it] = down_ftr[it];
+ }
+
+ // Apply replication padding for upsampling, matching torch impl
+ #pragma unroll
+ for (int it = -HALF_FILTER_SIZE; it < BUFFER_SIZE + HALF_FILTER_SIZE; it += 1)
+ {
+ int element_index = seq_offset + it; // index for element
+ if ((element_index < 0) && (element_index >= -UPSAMPLE_REPLICATION_PAD))
+ {
+ elements[2 * (HALF_FILTER_SIZE + it)] = 2 * seq_left_most_value;
+ }
+ if ((element_index >= seq_len) && (element_index < seq_len + UPSAMPLE_REPLICATION_PAD))
+ {
+ elements[2 * (HALF_FILTER_SIZE + it)] = 2 * seq_right_most_value;
+ }
+ if ((element_index >= 0) && (element_index < seq_len))
+ {
+ elements[2 * (HALF_FILTER_SIZE + it)] = 2 * src[it];
+ }
+ }
+
+ // Apply upsampling strided convolution and write to intermediates. It reserves DOWNSAMPLE_REPLICATION_PAD_LEFT for replication padding of the downsampilng conv later
+ #pragma unroll
+ for (int it = 0; it < (2 * BUFFER_SIZE + 2 * FILTER_SIZE); it += 1)
+ {
+ input_t acc = 0.0;
+ int element_index = intermediate_seq_offset + it; // index for intermediate
+ #pragma unroll
+ for (int f_idx = 0; f_idx < FILTER_SIZE; f_idx += 1)
+ {
+ if ((element_index + f_idx) >= 0)
+ {
+ acc += up_filter[f_idx] * elements[it + f_idx];
+ }
+ }
+ intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] = acc;
+ }
+
+ // Apply activation function. It reserves DOWNSAMPLE_REPLICATION_PAD_LEFT and DOWNSAMPLE_REPLICATION_PAD_RIGHT for replication padding of the downsampilng conv later
+ double no_div_by_zero = 0.000000001;
+ #pragma unroll
+ for (int it = 0; it < 2 * BUFFER_SIZE + 2 * FILTER_SIZE; it += 1)
+ {
+ intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] += (1.0 / (beta_val + no_div_by_zero)) * sinf(intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] * alpha_val) * sinf(intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] * alpha_val);
+ }
+
+ // Apply replication padding before downsampling conv from intermediates
+ #pragma unroll
+ for (int it = 0; it < DOWNSAMPLE_REPLICATION_PAD_LEFT; it += 1)
+ {
+ intermediates[it] = intermediates[DOWNSAMPLE_REPLICATION_PAD_LEFT];
+ }
+ #pragma unroll
+ for (int it = DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE; it < DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE + DOWNSAMPLE_REPLICATION_PAD_RIGHT; it += 1)
+ {
+ intermediates[it] = intermediates[DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE - 1];
+ }
+
+ // Apply downsample strided convolution (assuming stride=2) from intermediates
+ #pragma unroll
+ for (int it = 0; it < BUFFER_SIZE; it += 1)
+ {
+ input_t acc = 0.0;
+ #pragma unroll
+ for (int f_idx = 0; f_idx < FILTER_SIZE; f_idx += 1)
+ {
+ // Add constant DOWNSAMPLE_REPLICATION_PAD_RIGHT to match torch implementation
+ acc += down_filter[f_idx] * intermediates[it * 2 + f_idx + DOWNSAMPLE_REPLICATION_PAD_RIGHT];
+ }
+ output[it] = acc;
+ }
+
+ // Write output to dst
+ #pragma unroll
+ for (int it = 0; it < BUFFER_SIZE; it += ELEMENTS_PER_LDG_STG)
+ {
+ int element_index = seq_offset + it;
+ if (element_index < seq_len)
+ {
+ dst[it] = output[it];
+ }
+ }
+
+ }
+
+ template
+ void dispatch_anti_alias_activation_forward(
+ output_t *dst,
+ const input_t *src,
+ const input_t *up_ftr,
+ const input_t *down_ftr,
+ const input_t *alpha,
+ const input_t *beta,
+ int batch_size,
+ int channels,
+ int seq_len)
+ {
+ if (seq_len == 0)
+ {
+ return;
+ }
+ else
+ {
+ // Use 128 threads per block to maximimize gpu utilization
+ constexpr int threads_per_block = 128;
+ constexpr int seq_len_per_block = 4096;
+ int blocks_per_seq_len = (seq_len + seq_len_per_block - 1) / seq_len_per_block;
+ dim3 blocks(blocks_per_seq_len, channels, batch_size);
+ dim3 threads(threads_per_block, 1, 1);
+
+ anti_alias_activation_forward
+ <<>>(dst, src, up_ftr, down_ftr, alpha, beta, batch_size, channels, seq_len);
+ }
+ }
+}
+
+extern "C" torch::Tensor fwd_cuda(torch::Tensor const &input, torch::Tensor const &up_filter, torch::Tensor const &down_filter, torch::Tensor const &alpha, torch::Tensor const &beta)
+{
+ // Input is a 3d tensor with dimensions [batches, channels, seq_len]
+ const int batches = input.size(0);
+ const int channels = input.size(1);
+ const int seq_len = input.size(2);
+
+ // Output
+ auto act_options = input.options().requires_grad(false);
+
+ torch::Tensor anti_alias_activation_results =
+ torch::empty({batches, channels, seq_len}, act_options);
+
+ void *input_ptr = static_cast(input.data_ptr());
+ void *up_filter_ptr = static_cast(up_filter.data_ptr());
+ void *down_filter_ptr = static_cast(down_filter.data_ptr());
+ void *alpha_ptr = static_cast(alpha.data_ptr());
+ void *beta_ptr = static_cast(beta.data_ptr());
+ void *anti_alias_activation_results_ptr = static_cast(anti_alias_activation_results.data_ptr());
+
+ DISPATCH_FLOAT_HALF_AND_BFLOAT(
+ input.scalar_type(),
+ "dispatch anti alias activation_forward",
+ dispatch_anti_alias_activation_forward(
+ reinterpret_cast(anti_alias_activation_results_ptr),
+ reinterpret_cast(input_ptr),
+ reinterpret_cast(up_filter_ptr),
+ reinterpret_cast(down_filter_ptr),
+ reinterpret_cast(alpha_ptr),
+ reinterpret_cast(beta_ptr),
+ batches,
+ channels,
+ seq_len););
+ return anti_alias_activation_results;
+}
\ No newline at end of file
diff --git a/modules/bigvgan/alias_free_activation/cuda/build/.ninja_deps b/modules/bigvgan/alias_free_activation/cuda/build/.ninja_deps
new file mode 100644
index 0000000000000000000000000000000000000000..ce9bc6ec886c4bda48c73677998abe0b2b73bfcc
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/cuda/build/.ninja_deps
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e233713716a5778577f244b0f310944ff26d3079ce0e42491791da7d42e363c1
+size 522068
diff --git a/modules/bigvgan/alias_free_activation/cuda/build/.ninja_log b/modules/bigvgan/alias_free_activation/cuda/build/.ninja_log
new file mode 100644
index 0000000000000000000000000000000000000000..bd3c097a5622dde1a5c17fd152e04750a1dedded
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/cuda/build/.ninja_log
@@ -0,0 +1,7 @@
+# ninja log v5
+9 39554 7516864785377831 anti_alias_activation.o 3a177f31dd72c43c
+13 152601 7516865914203767 anti_alias_activation_cuda.cuda.o 2d613e7382d803fd
+152628 153062 7516865920541751 anti_alias_activation_cuda.pyd f6366e9bdfb27f7
+128 50503 7654004565901584 anti_alias_activation.o 9ed3213f2e0d0858
+133 176837 7654005827401976 anti_alias_activation_cuda.cuda.o a679b6661c609136
+176839 177401 7654005835005523 anti_alias_activation_cuda.pyd f6366e9bdfb27f7
diff --git a/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation.o b/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation.o
new file mode 100644
index 0000000000000000000000000000000000000000..812f06975323c9e1937fa01c943e31ae02322145
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation.o
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:74c2824b05582070b69f51ec588aadb268c4fddf18fbb4590f901d1cdf32185c
+size 3246655
diff --git a/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation_cuda.cuda.o b/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation_cuda.cuda.o
new file mode 100644
index 0000000000000000000000000000000000000000..329fb7a9b147a0af665ff7f7686bd0cc915ecc84
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation_cuda.cuda.o
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86c48de557041de7ebaff7926b5f346cc5e4e2dddc6cf5b88409f6cb161db0f4
+size 4724513
diff --git a/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation_cuda.exp b/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation_cuda.exp
new file mode 100644
index 0000000000000000000000000000000000000000..3093a741ef126748042cafaef5c368f3ec5e2d3f
Binary files /dev/null and b/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation_cuda.exp differ
diff --git a/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation_cuda.lib b/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation_cuda.lib
new file mode 100644
index 0000000000000000000000000000000000000000..1be22a5e2a68606c56c961333b67a251bf40d8ea
Binary files /dev/null and b/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation_cuda.lib differ
diff --git a/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation_cuda.pyd b/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation_cuda.pyd
new file mode 100644
index 0000000000000000000000000000000000000000..dc51b91fc3d147e24ad0155ee31809556bb7208a
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/cuda/build/anti_alias_activation_cuda.pyd
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db37ea2dd31dfe67e68ee6019877d14638c41724ff9342c55f638f4d2cda3d03
+size 2454528
diff --git a/modules/bigvgan/alias_free_activation/cuda/build/build.ninja b/modules/bigvgan/alias_free_activation/cuda/build/build.ninja
new file mode 100644
index 0000000000000000000000000000000000000000..8c41cf88948be2657b26c226365a86b99278764a
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/cuda/build/build.ninja
@@ -0,0 +1,38 @@
+ninja_required_version = 1.3
+cxx = cl
+nvcc = C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8\bin\nvcc
+
+cflags = -DTORCH_EXTENSION_NAME=anti_alias_activation_cuda -DTORCH_API_INCLUDE_EXTENSION_H -ID:\Anaconda\envs\vocos\lib\site-packages\torch\include -ID:\Anaconda\envs\vocos\lib\site-packages\torch\include\torch\csrc\api\include "-IC:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8\include" -ID:\Anaconda\envs\vocos\Include /std:c++17 -O3 /MD /wd4819 /wd4251 /wd4244 /wd4267 /wd4275 /wd4018 /wd4190 /wd4624 /wd4067 /wd4068 /EHsc
+post_cflags =
+cuda_cflags = -Xcudafe --diag_suppress=dll_interface_conflict_dllexport_assumed -Xcudafe --diag_suppress=dll_interface_conflict_none_assumed -Xcudafe --diag_suppress=field_without_dll_interface -Xcudafe --diag_suppress=base_class_has_different_dll_interface -Xcompiler /EHsc -Xcompiler /wd4068 -Xcompiler /wd4067 -Xcompiler /wd4624 -Xcompiler /wd4190 -Xcompiler /wd4018 -Xcompiler /wd4275 -Xcompiler /wd4267 -Xcompiler /wd4244 -Xcompiler /wd4251 -Xcompiler /wd4819 -Xcompiler /MD -DTORCH_EXTENSION_NAME=anti_alias_activation_cuda -DTORCH_API_INCLUDE_EXTENSION_H -ID:\Anaconda\envs\vocos\lib\site-packages\torch\include -ID:\Anaconda\envs\vocos\lib\site-packages\torch\include\torch\csrc\api\include "-IC:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8\include" -ID:\Anaconda\envs\vocos\Include -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_86,code=sm_86 -std=c++17 -O3 -gencode arch=compute_70,code=sm_70 --use_fast_math -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ --expt-relaxed-constexpr --expt-extended-lambda -gencode arch=compute_80,code=sm_80
+cuda_post_cflags =
+cuda_dlink_post_cflags =
+sycl_dlink_post_cflags =
+ldflags = /DLL c10.lib c10_cuda.lib torch_cpu.lib torch_cuda.lib -INCLUDE:?warp_size@cuda@at@@YAHXZ torch.lib /LIBPATH:D:\Anaconda\envs\vocos\lib\site-packages\torch\lib torch_python.lib /LIBPATH:D:\Anaconda\envs\vocos\libs "/LIBPATH:C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8\lib\x64" cudart.lib
+
+rule compile
+ command = cl /showIncludes $cflags -c $in /Fo$out $post_cflags
+ deps = msvc
+
+rule cuda_compile
+ depfile = $out.d
+ deps = gcc
+ command = $nvcc --generate-dependencies-with-compile --dependency-output $out.d $cuda_cflags -c $in -o $out $cuda_post_cflags
+
+
+
+
+
+rule link
+ command = "D$:\Visual Studio\VC\Tools\MSVC\14.29.30133\bin\Hostx86\x64/link.exe" $in /nologo $ldflags /out:$out
+
+build anti_alias_activation.o: compile D$:\seed-vc\modules\bigvgan\alias_free_activation\cuda\anti_alias_activation.cpp
+build anti_alias_activation_cuda.cuda.o: cuda_compile D$:\seed-vc\modules\bigvgan\alias_free_activation\cuda\anti_alias_activation_cuda.cu
+
+
+
+
+
+build anti_alias_activation_cuda.pyd: link anti_alias_activation.o anti_alias_activation_cuda.cuda.o
+
+default anti_alias_activation_cuda.pyd
diff --git a/modules/bigvgan/alias_free_activation/cuda/compat.h b/modules/bigvgan/alias_free_activation/cuda/compat.h
new file mode 100644
index 0000000000000000000000000000000000000000..0f93af5700470e7f6066af7dbe56aced98ea32d9
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/cuda/compat.h
@@ -0,0 +1,29 @@
+/* coding=utf-8
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*This code is copied fron NVIDIA apex:
+ * https://github.com/NVIDIA/apex
+ * with minor changes. */
+
+#ifndef TORCH_CHECK
+#define TORCH_CHECK AT_CHECK
+#endif
+
+#ifdef VERSION_GE_1_3
+#define DATA_PTR data_ptr
+#else
+#define DATA_PTR data
+#endif
diff --git a/modules/bigvgan/alias_free_activation/cuda/load.py b/modules/bigvgan/alias_free_activation/cuda/load.py
new file mode 100644
index 0000000000000000000000000000000000000000..82afde3d73dda72b06af28a622fdab1954825a28
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/cuda/load.py
@@ -0,0 +1,86 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+import os
+import pathlib
+import subprocess
+
+from torch.utils import cpp_extension
+
+"""
+Setting this param to a list has a problem of generating different compilation commands (with diferent order of architectures) and leading to recompilation of fused kernels.
+Set it to empty stringo avoid recompilation and assign arch flags explicity in extra_cuda_cflags below
+"""
+os.environ["TORCH_CUDA_ARCH_LIST"] = ""
+
+
+def load():
+ # Check if cuda 11 is installed for compute capability 8.0
+ cc_flag = []
+ _, bare_metal_major, _ = _get_cuda_bare_metal_version(cpp_extension.CUDA_HOME)
+ if int(bare_metal_major) >= 11:
+ cc_flag.append("-gencode")
+ cc_flag.append("arch=compute_80,code=sm_80")
+
+ # Build path
+ srcpath = pathlib.Path(__file__).parent.absolute()
+ buildpath = srcpath / "build"
+ _create_build_dir(buildpath)
+
+ # Helper function to build the kernels.
+ def _cpp_extention_load_helper(name, sources, extra_cuda_flags):
+ return cpp_extension.load(
+ name=name,
+ sources=sources,
+ build_directory=buildpath,
+ extra_cflags=[
+ "-O3",
+ ],
+ extra_cuda_cflags=[
+ "-O3",
+ "-gencode",
+ "arch=compute_70,code=sm_70",
+ "--use_fast_math",
+ ]
+ + extra_cuda_flags
+ + cc_flag,
+ verbose=True,
+ )
+
+ extra_cuda_flags = [
+ "-U__CUDA_NO_HALF_OPERATORS__",
+ "-U__CUDA_NO_HALF_CONVERSIONS__",
+ "--expt-relaxed-constexpr",
+ "--expt-extended-lambda",
+ ]
+
+ sources = [
+ srcpath / "anti_alias_activation.cpp",
+ srcpath / "anti_alias_activation_cuda.cu",
+ ]
+ anti_alias_activation_cuda = _cpp_extention_load_helper(
+ "anti_alias_activation_cuda", sources, extra_cuda_flags
+ )
+
+ return anti_alias_activation_cuda
+
+
+def _get_cuda_bare_metal_version(cuda_dir):
+ raw_output = subprocess.check_output(
+ [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
+ )
+ output = raw_output.split()
+ release_idx = output.index("release") + 1
+ release = output[release_idx].split(".")
+ bare_metal_major = release[0]
+ bare_metal_minor = release[1][0]
+
+ return raw_output, bare_metal_major, bare_metal_minor
+
+
+def _create_build_dir(buildpath):
+ try:
+ os.mkdir(buildpath)
+ except OSError:
+ if not os.path.isdir(buildpath):
+ print(f"Creation of the build directory {buildpath} failed")
diff --git a/modules/bigvgan/alias_free_activation/cuda/type_shim.h b/modules/bigvgan/alias_free_activation/cuda/type_shim.h
new file mode 100644
index 0000000000000000000000000000000000000000..4328d0369a5fb8730cdf236d9f267453f4201d84
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/cuda/type_shim.h
@@ -0,0 +1,92 @@
+/* coding=utf-8
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include
+#include "compat.h"
+
+#define DISPATCH_FLOAT_HALF_AND_BFLOAT(TYPE, NAME, ...) \
+ switch (TYPE) \
+ { \
+ case at::ScalarType::Float: \
+ { \
+ using scalar_t = float; \
+ __VA_ARGS__; \
+ break; \
+ } \
+ case at::ScalarType::Half: \
+ { \
+ using scalar_t = at::Half; \
+ __VA_ARGS__; \
+ break; \
+ } \
+ case at::ScalarType::BFloat16: \
+ { \
+ using scalar_t = at::BFloat16; \
+ __VA_ARGS__; \
+ break; \
+ } \
+ default: \
+ AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \
+ }
+
+#define DISPATCH_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(TYPEIN, TYPEOUT, NAME, ...) \
+ switch (TYPEIN) \
+ { \
+ case at::ScalarType::Float: \
+ { \
+ using scalar_t_in = float; \
+ switch (TYPEOUT) \
+ { \
+ case at::ScalarType::Float: \
+ { \
+ using scalar_t_out = float; \
+ __VA_ARGS__; \
+ break; \
+ } \
+ case at::ScalarType::Half: \
+ { \
+ using scalar_t_out = at::Half; \
+ __VA_ARGS__; \
+ break; \
+ } \
+ case at::ScalarType::BFloat16: \
+ { \
+ using scalar_t_out = at::BFloat16; \
+ __VA_ARGS__; \
+ break; \
+ } \
+ default: \
+ AT_ERROR(#NAME, " not implemented for '", toString(TYPEOUT), "'"); \
+ } \
+ break; \
+ } \
+ case at::ScalarType::Half: \
+ { \
+ using scalar_t_in = at::Half; \
+ using scalar_t_out = at::Half; \
+ __VA_ARGS__; \
+ break; \
+ } \
+ case at::ScalarType::BFloat16: \
+ { \
+ using scalar_t_in = at::BFloat16; \
+ using scalar_t_out = at::BFloat16; \
+ __VA_ARGS__; \
+ break; \
+ } \
+ default: \
+ AT_ERROR(#NAME, " not implemented for '", toString(TYPEIN), "'"); \
+ }
diff --git a/modules/bigvgan/alias_free_activation/torch/__init__.py b/modules/bigvgan/alias_free_activation/torch/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..7bb0ad84ef184dcb15464c8ca827ae1c284f8990
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/torch/__init__.py
@@ -0,0 +1,6 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+# LICENSE is in incl_licenses directory.
+
+from .filter import *
+from .resample import *
+from .act import *
diff --git a/modules/bigvgan/alias_free_activation/torch/__pycache__/__init__.cpython-310.pyc b/modules/bigvgan/alias_free_activation/torch/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cb684d13c78024166a004ef5072e674a5164233f
Binary files /dev/null and b/modules/bigvgan/alias_free_activation/torch/__pycache__/__init__.cpython-310.pyc differ
diff --git a/modules/bigvgan/alias_free_activation/torch/__pycache__/act.cpython-310.pyc b/modules/bigvgan/alias_free_activation/torch/__pycache__/act.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4d909a23a036ef4db92f0523ca51047f492a89e9
Binary files /dev/null and b/modules/bigvgan/alias_free_activation/torch/__pycache__/act.cpython-310.pyc differ
diff --git a/modules/bigvgan/alias_free_activation/torch/__pycache__/filter.cpython-310.pyc b/modules/bigvgan/alias_free_activation/torch/__pycache__/filter.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e8ec3a951f28adb444e01fcf8532d9a2218bf7f7
Binary files /dev/null and b/modules/bigvgan/alias_free_activation/torch/__pycache__/filter.cpython-310.pyc differ
diff --git a/modules/bigvgan/alias_free_activation/torch/__pycache__/resample.cpython-310.pyc b/modules/bigvgan/alias_free_activation/torch/__pycache__/resample.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..63a1166f483a4b45fce21ad41893444f5d71dfd5
Binary files /dev/null and b/modules/bigvgan/alias_free_activation/torch/__pycache__/resample.cpython-310.pyc differ
diff --git a/modules/bigvgan/alias_free_activation/torch/act.py b/modules/bigvgan/alias_free_activation/torch/act.py
new file mode 100644
index 0000000000000000000000000000000000000000..c118fa7194f9a37d8c4fae2535700d2a96af4f9c
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/torch/act.py
@@ -0,0 +1,30 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+# LICENSE is in incl_licenses directory.
+
+import torch.nn as nn
+from .resample import UpSample1d, DownSample1d
+
+
+class Activation1d(nn.Module):
+ def __init__(
+ self,
+ activation,
+ up_ratio: int = 2,
+ down_ratio: int = 2,
+ up_kernel_size: int = 12,
+ down_kernel_size: int = 12,
+ ):
+ super().__init__()
+ self.up_ratio = up_ratio
+ self.down_ratio = down_ratio
+ self.act = activation
+ self.upsample = UpSample1d(up_ratio, up_kernel_size)
+ self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+ # x: [B,C,T]
+ def forward(self, x):
+ x = self.upsample(x)
+ x = self.act(x)
+ x = self.downsample(x)
+
+ return x
diff --git a/modules/bigvgan/alias_free_activation/torch/filter.py b/modules/bigvgan/alias_free_activation/torch/filter.py
new file mode 100644
index 0000000000000000000000000000000000000000..81a4a9a7cefb457f8880f54385335180fbd43f1b
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/torch/filter.py
@@ -0,0 +1,101 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+# LICENSE is in incl_licenses directory.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+if "sinc" in dir(torch):
+ sinc = torch.sinc
+else:
+ # This code is adopted from adefossez's julius.core.sinc under the MIT License
+ # https://adefossez.github.io/julius/julius/core.html
+ # LICENSE is in incl_licenses directory.
+ def sinc(x: torch.Tensor):
+ """
+ Implementation of sinc, i.e. sin(pi * x) / (pi * x)
+ __Warning__: Different to julius.sinc, the input is multiplied by `pi`!
+ """
+ return torch.where(
+ x == 0,
+ torch.tensor(1.0, device=x.device, dtype=x.dtype),
+ torch.sin(math.pi * x) / math.pi / x,
+ )
+
+
+# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License
+# https://adefossez.github.io/julius/julius/lowpass.html
+# LICENSE is in incl_licenses directory.
+def kaiser_sinc_filter1d(
+ cutoff, half_width, kernel_size
+): # return filter [1,1,kernel_size]
+ even = kernel_size % 2 == 0
+ half_size = kernel_size // 2
+
+ # For kaiser window
+ 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)
+
+ # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio
+ if even:
+ time = torch.arange(-half_size, half_size) + 0.5
+ else:
+ time = torch.arange(kernel_size) - half_size
+ if cutoff == 0:
+ filter_ = torch.zeros_like(time)
+ else:
+ filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
+ """
+ Normalize filter to have sum = 1, otherwise we will have a small leakage of the constant component in the input signal.
+ """
+ filter_ /= filter_.sum()
+ filter = filter_.view(1, 1, kernel_size)
+
+ return filter
+
+
+class LowPassFilter1d(nn.Module):
+ def __init__(
+ self,
+ cutoff=0.5,
+ half_width=0.6,
+ stride: int = 1,
+ padding: bool = True,
+ padding_mode: str = "replicate",
+ kernel_size: int = 12,
+ ):
+ """
+ kernel_size should be even number for stylegan3 setup, in this implementation, odd number is also possible.
+ """
+ super().__init__()
+ if cutoff < -0.0:
+ raise ValueError("Minimum cutoff must be larger than zero.")
+ if cutoff > 0.5:
+ raise ValueError("A cutoff above 0.5 does not make sense.")
+ self.kernel_size = kernel_size
+ self.even = kernel_size % 2 == 0
+ self.pad_left = kernel_size // 2 - int(self.even)
+ self.pad_right = kernel_size // 2
+ self.stride = stride
+ self.padding = padding
+ self.padding_mode = padding_mode
+ filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+ self.register_buffer("filter", filter)
+
+ # Input [B, C, T]
+ def forward(self, x):
+ _, C, _ = x.shape
+
+ if self.padding:
+ x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+ out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
+
+ return out
diff --git a/modules/bigvgan/alias_free_activation/torch/resample.py b/modules/bigvgan/alias_free_activation/torch/resample.py
new file mode 100644
index 0000000000000000000000000000000000000000..8e17883696490e5f3cdd8dcfce347a3318061eb8
--- /dev/null
+++ b/modules/bigvgan/alias_free_activation/torch/resample.py
@@ -0,0 +1,58 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+# LICENSE is in incl_licenses directory.
+
+import torch.nn as nn
+from torch.nn import functional as F
+from .filter import LowPassFilter1d
+from .filter import kaiser_sinc_filter1d
+
+
+class UpSample1d(nn.Module):
+ def __init__(self, ratio=2, kernel_size=None):
+ super().__init__()
+ self.ratio = ratio
+ self.kernel_size = (
+ int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+ )
+ self.stride = ratio
+ self.pad = self.kernel_size // ratio - 1
+ self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+ self.pad_right = (
+ self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+ )
+ filter = kaiser_sinc_filter1d(
+ cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size
+ )
+ self.register_buffer("filter", filter)
+
+ # x: [B, C, T]
+ def forward(self, x):
+ _, C, _ = x.shape
+
+ x = F.pad(x, (self.pad, self.pad), mode="replicate")
+ x = self.ratio * F.conv_transpose1d(
+ x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C
+ )
+ x = x[..., self.pad_left : -self.pad_right]
+
+ return x
+
+
+class DownSample1d(nn.Module):
+ def __init__(self, ratio=2, kernel_size=None):
+ super().__init__()
+ self.ratio = ratio
+ self.kernel_size = (
+ int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+ )
+ self.lowpass = LowPassFilter1d(
+ cutoff=0.5 / ratio,
+ half_width=0.6 / ratio,
+ stride=ratio,
+ kernel_size=self.kernel_size,
+ )
+
+ def forward(self, x):
+ xx = self.lowpass(x)
+
+ return xx
diff --git a/modules/bigvgan/bigvgan.py b/modules/bigvgan/bigvgan.py
new file mode 100644
index 0000000000000000000000000000000000000000..41d6e44a2cb59a39d51cb4994d05804dc497dda7
--- /dev/null
+++ b/modules/bigvgan/bigvgan.py
@@ -0,0 +1,492 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
+# LICENSE is in incl_licenses directory.
+
+import os
+import json
+from pathlib import Path
+from typing import Optional, Union, Dict
+
+import torch
+import torch.nn as nn
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+from . import activations
+from .utils import init_weights, get_padding
+from .alias_free_activation.torch.act import Activation1d as TorchActivation1d
+from .env import AttrDict
+
+from huggingface_hub import PyTorchModelHubMixin, hf_hub_download
+
+
+def load_hparams_from_json(path) -> AttrDict:
+ with open(path) as f:
+ data = f.read()
+ return AttrDict(json.loads(data))
+
+
+class AMPBlock1(torch.nn.Module):
+ """
+ AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer.
+ AMPBlock1 has additional self.convs2 that contains additional Conv1d layers with a fixed dilation=1 followed by each layer in self.convs1
+
+ Args:
+ h (AttrDict): Hyperparameters.
+ channels (int): Number of convolution channels.
+ kernel_size (int): Size of the convolution kernel. Default is 3.
+ dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5).
+ activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Default is None.
+ """
+
+ def __init__(
+ self,
+ h: AttrDict,
+ channels: int,
+ kernel_size: int = 3,
+ dilation: tuple = (1, 3, 5),
+ activation: str = None,
+ ):
+ super().__init__()
+
+ self.h = h
+
+ self.convs1 = nn.ModuleList(
+ [
+ weight_norm(
+ Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ stride=1,
+ dilation=d,
+ padding=get_padding(kernel_size, d),
+ )
+ )
+ for d in dilation
+ ]
+ )
+ self.convs1.apply(init_weights)
+
+ self.convs2 = nn.ModuleList(
+ [
+ weight_norm(
+ Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ stride=1,
+ dilation=1,
+ padding=get_padding(kernel_size, 1),
+ )
+ )
+ for _ in range(len(dilation))
+ ]
+ )
+ self.convs2.apply(init_weights)
+
+ self.num_layers = len(self.convs1) + len(
+ self.convs2
+ ) # Total number of conv layers
+
+ # Select which Activation1d, lazy-load cuda version to ensure backward compatibility
+ if self.h.get("use_cuda_kernel", False):
+ from .alias_free_activation.cuda.activation1d import (
+ Activation1d as CudaActivation1d,
+ )
+
+ Activation1d = CudaActivation1d
+ else:
+ Activation1d = TorchActivation1d
+
+ # Activation functions
+ if activation == "snake":
+ self.activations = nn.ModuleList(
+ [
+ Activation1d(
+ activation=activations.Snake(
+ channels, alpha_logscale=h.snake_logscale
+ )
+ )
+ for _ in range(self.num_layers)
+ ]
+ )
+ elif activation == "snakebeta":
+ self.activations = nn.ModuleList(
+ [
+ Activation1d(
+ activation=activations.SnakeBeta(
+ channels, alpha_logscale=h.snake_logscale
+ )
+ )
+ for _ in range(self.num_layers)
+ ]
+ )
+ else:
+ raise NotImplementedError(
+ "activation incorrectly specified. check the config file and look for 'activation'."
+ )
+
+ def forward(self, x):
+ acts1, acts2 = self.activations[::2], self.activations[1::2]
+ for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2):
+ xt = a1(x)
+ xt = c1(xt)
+ xt = a2(xt)
+ xt = c2(xt)
+ x = xt + x
+
+ return x
+
+ def remove_weight_norm(self):
+ for l in self.convs1:
+ remove_weight_norm(l)
+ for l in self.convs2:
+ remove_weight_norm(l)
+
+
+class AMPBlock2(torch.nn.Module):
+ """
+ AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer.
+ Unlike AMPBlock1, AMPBlock2 does not contain extra Conv1d layers with fixed dilation=1
+
+ Args:
+ h (AttrDict): Hyperparameters.
+ channels (int): Number of convolution channels.
+ kernel_size (int): Size of the convolution kernel. Default is 3.
+ dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5).
+ activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Default is None.
+ """
+
+ def __init__(
+ self,
+ h: AttrDict,
+ channels: int,
+ kernel_size: int = 3,
+ dilation: tuple = (1, 3, 5),
+ activation: str = None,
+ ):
+ super().__init__()
+
+ self.h = h
+
+ self.convs = nn.ModuleList(
+ [
+ weight_norm(
+ Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ stride=1,
+ dilation=d,
+ padding=get_padding(kernel_size, d),
+ )
+ )
+ for d in dilation
+ ]
+ )
+ self.convs.apply(init_weights)
+
+ self.num_layers = len(self.convs) # Total number of conv layers
+
+ # Select which Activation1d, lazy-load cuda version to ensure backward compatibility
+ if self.h.get("use_cuda_kernel", False):
+ from .alias_free_activation.cuda.activation1d import (
+ Activation1d as CudaActivation1d,
+ )
+
+ Activation1d = CudaActivation1d
+ else:
+ Activation1d = TorchActivation1d
+
+ # Activation functions
+ if activation == "snake":
+ self.activations = nn.ModuleList(
+ [
+ Activation1d(
+ activation=activations.Snake(
+ channels, alpha_logscale=h.snake_logscale
+ )
+ )
+ for _ in range(self.num_layers)
+ ]
+ )
+ elif activation == "snakebeta":
+ self.activations = nn.ModuleList(
+ [
+ Activation1d(
+ activation=activations.SnakeBeta(
+ channels, alpha_logscale=h.snake_logscale
+ )
+ )
+ for _ in range(self.num_layers)
+ ]
+ )
+ else:
+ raise NotImplementedError(
+ "activation incorrectly specified. check the config file and look for 'activation'."
+ )
+
+ def forward(self, x):
+ for c, a in zip(self.convs, self.activations):
+ xt = a(x)
+ xt = c(xt)
+ x = xt + x
+
+ def remove_weight_norm(self):
+ for l in self.convs:
+ remove_weight_norm(l)
+
+
+class BigVGAN(
+ torch.nn.Module,
+ PyTorchModelHubMixin,
+ library_name="bigvgan",
+ repo_url="https://github.com/NVIDIA/BigVGAN",
+ docs_url="https://github.com/NVIDIA/BigVGAN/blob/main/README.md",
+ pipeline_tag="audio-to-audio",
+ license="mit",
+ tags=["neural-vocoder", "audio-generation", "arxiv:2206.04658"],
+):
+ """
+ BigVGAN is a neural vocoder model that applies anti-aliased periodic activation for residual blocks (resblocks).
+ New in BigVGAN-v2: it can optionally use optimized CUDA kernels for AMP (anti-aliased multi-periodicity) blocks.
+
+ Args:
+ h (AttrDict): Hyperparameters.
+ use_cuda_kernel (bool): If set to True, loads optimized CUDA kernels for AMP. This should be used for inference only, as training is not supported with CUDA kernels.
+
+ Note:
+ - The `use_cuda_kernel` parameter should be used for inference only, as training with CUDA kernels is not supported.
+ - Ensure that the activation function is correctly specified in the hyperparameters (h.activation).
+ """
+
+ def __init__(self, h: AttrDict, use_cuda_kernel: bool = False):
+ super().__init__()
+ self.h = h
+ self.h["use_cuda_kernel"] = use_cuda_kernel
+
+ # Select which Activation1d, lazy-load cuda version to ensure backward compatibility
+ if self.h.get("use_cuda_kernel", False):
+ from .alias_free_activation.cuda.activation1d import (
+ Activation1d as CudaActivation1d,
+ )
+
+ Activation1d = CudaActivation1d
+ else:
+ Activation1d = TorchActivation1d
+
+ self.num_kernels = len(h.resblock_kernel_sizes)
+ self.num_upsamples = len(h.upsample_rates)
+
+ # Pre-conv
+ self.conv_pre = weight_norm(
+ Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)
+ )
+
+ # Define which AMPBlock to use. BigVGAN uses AMPBlock1 as default
+ if h.resblock == "1":
+ resblock_class = AMPBlock1
+ elif h.resblock == "2":
+ resblock_class = AMPBlock2
+ else:
+ raise ValueError(
+ f"Incorrect resblock class specified in hyperparameters. Got {h.resblock}"
+ )
+
+ # Transposed conv-based upsamplers. does not apply anti-aliasing
+ self.ups = nn.ModuleList()
+ for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
+ self.ups.append(
+ nn.ModuleList(
+ [
+ weight_norm(
+ ConvTranspose1d(
+ h.upsample_initial_channel // (2 ** i),
+ h.upsample_initial_channel // (2 ** (i + 1)),
+ k,
+ u,
+ padding=(k - u) // 2,
+ )
+ )
+ ]
+ )
+ )
+
+ # Residual blocks using anti-aliased multi-periodicity composition modules (AMP)
+ self.resblocks = nn.ModuleList()
+ for i in range(len(self.ups)):
+ ch = h.upsample_initial_channel // (2 ** (i + 1))
+ for j, (k, d) in enumerate(
+ zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)
+ ):
+ self.resblocks.append(
+ resblock_class(h, ch, k, d, activation=h.activation)
+ )
+
+ # Post-conv
+ activation_post = (
+ activations.Snake(ch, alpha_logscale=h.snake_logscale)
+ if h.activation == "snake"
+ else (
+ activations.SnakeBeta(ch, alpha_logscale=h.snake_logscale)
+ if h.activation == "snakebeta"
+ else None
+ )
+ )
+ if activation_post is None:
+ raise NotImplementedError(
+ "activation incorrectly specified. check the config file and look for 'activation'."
+ )
+
+ self.activation_post = Activation1d(activation=activation_post)
+
+ # Whether to use bias for the final conv_post. Default to True for backward compatibility
+ self.use_bias_at_final = h.get("use_bias_at_final", True)
+ self.conv_post = weight_norm(
+ Conv1d(ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final)
+ )
+
+ # Weight initialization
+ for i in range(len(self.ups)):
+ self.ups[i].apply(init_weights)
+ self.conv_post.apply(init_weights)
+
+ # Final tanh activation. Defaults to True for backward compatibility
+ self.use_tanh_at_final = h.get("use_tanh_at_final", True)
+
+ def forward(self, x):
+ # Pre-conv
+ x = self.conv_pre(x)
+
+ for i in range(self.num_upsamples):
+ # Upsampling
+ for i_up in range(len(self.ups[i])):
+ x = self.ups[i][i_up](x)
+ # AMP blocks
+ xs = None
+ for j in range(self.num_kernels):
+ if xs is None:
+ xs = self.resblocks[i * self.num_kernels + j](x)
+ else:
+ xs += self.resblocks[i * self.num_kernels + j](x)
+ x = xs / self.num_kernels
+
+ # Post-conv
+ x = self.activation_post(x)
+ x = self.conv_post(x)
+ # Final tanh activation
+ if self.use_tanh_at_final:
+ x = torch.tanh(x)
+ else:
+ x = torch.clamp(x, min=-1.0, max=1.0) # Bound the output to [-1, 1]
+
+ return x
+
+ def remove_weight_norm(self):
+ try:
+ print("Removing weight norm...")
+ for l in self.ups:
+ for l_i in l:
+ remove_weight_norm(l_i)
+ for l in self.resblocks:
+ l.remove_weight_norm()
+ remove_weight_norm(self.conv_pre)
+ remove_weight_norm(self.conv_post)
+ except ValueError:
+ print("[INFO] Model already removed weight norm. Skipping!")
+ pass
+
+ # Additional methods for huggingface_hub support
+ def _save_pretrained(self, save_directory: Path) -> None:
+ """Save weights and config.json from a Pytorch model to a local directory."""
+
+ model_path = save_directory / "bigvgan_generator.pt"
+ torch.save({"generator": self.state_dict()}, model_path)
+
+ config_path = save_directory / "config.json"
+ with open(config_path, "w") as config_file:
+ json.dump(self.h, config_file, indent=4)
+
+ @classmethod
+ def _from_pretrained(
+ cls,
+ *,
+ model_id: str,
+ revision: str,
+ cache_dir: str,
+ force_download: bool,
+ proxies: Optional[Dict],
+ resume_download: bool,
+ local_files_only: bool,
+ token: Union[str, bool, None],
+ map_location: str = "cpu", # Additional argument
+ strict: bool = False, # Additional argument
+ use_cuda_kernel: bool = False,
+ **model_kwargs,
+ ):
+ """Load Pytorch pretrained weights and return the loaded model."""
+
+ # Download and load hyperparameters (h) used by BigVGAN
+ if os.path.isdir(model_id):
+ print("Loading config.json from local directory")
+ config_file = os.path.join(model_id, "config.json")
+ else:
+ config_file = hf_hub_download(
+ repo_id=model_id,
+ filename="config.json",
+ revision=revision,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ resume_download=resume_download,
+ token=token,
+ local_files_only=local_files_only,
+ )
+ h = load_hparams_from_json(config_file)
+
+ # instantiate BigVGAN using h
+ if use_cuda_kernel:
+ print(
+ f"[WARNING] You have specified use_cuda_kernel=True during BigVGAN.from_pretrained(). Only inference is supported (training is not implemented)!"
+ )
+ print(
+ f"[WARNING] You need nvcc and ninja installed in your system that matches your PyTorch build is using to build the kernel. If not, the model will fail to initialize or generate incorrect waveform!"
+ )
+ print(
+ f"[WARNING] For detail, see the official GitHub repository: https://github.com/NVIDIA/BigVGAN?tab=readme-ov-file#using-custom-cuda-kernel-for-synthesis"
+ )
+ model = cls(h, use_cuda_kernel=use_cuda_kernel)
+
+ # Download and load pretrained generator weight
+ if os.path.isdir(model_id):
+ print("Loading weights from local directory")
+ model_file = os.path.join(model_id, "bigvgan_generator.pt")
+ else:
+ print(f"Loading weights from {model_id}")
+ model_file = hf_hub_download(
+ repo_id=model_id,
+ filename="bigvgan_generator.pt",
+ revision=revision,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ resume_download=resume_download,
+ token=token,
+ local_files_only=local_files_only,
+ )
+
+ checkpoint_dict = torch.load(model_file, map_location=map_location)
+
+ try:
+ model.load_state_dict(checkpoint_dict["generator"])
+ except RuntimeError:
+ print(
+ f"[INFO] the pretrained checkpoint does not contain weight norm. Loading the checkpoint after removing weight norm!"
+ )
+ model.remove_weight_norm()
+ model.load_state_dict(checkpoint_dict["generator"])
+
+ return model
\ No newline at end of file
diff --git a/modules/bigvgan/config.json b/modules/bigvgan/config.json
new file mode 100644
index 0000000000000000000000000000000000000000..f3c53e8e4b2eb735082e9577f6b26a16ccd67a3b
--- /dev/null
+++ b/modules/bigvgan/config.json
@@ -0,0 +1,63 @@
+{
+ "resblock": "1",
+ "num_gpus": 0,
+ "batch_size": 32,
+ "learning_rate": 0.0001,
+ "adam_b1": 0.8,
+ "adam_b2": 0.99,
+ "lr_decay": 0.9999996,
+ "seed": 1234,
+
+ "upsample_rates": [4,4,2,2,2,2],
+ "upsample_kernel_sizes": [8,8,4,4,4,4],
+ "upsample_initial_channel": 1536,
+ "resblock_kernel_sizes": [3,7,11],
+ "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+
+ "use_tanh_at_final": false,
+ "use_bias_at_final": false,
+
+ "activation": "snakebeta",
+ "snake_logscale": true,
+
+ "use_cqtd_instead_of_mrd": true,
+ "cqtd_filters": 128,
+ "cqtd_max_filters": 1024,
+ "cqtd_filters_scale": 1,
+ "cqtd_dilations": [1, 2, 4],
+ "cqtd_hop_lengths": [512, 256, 256],
+ "cqtd_n_octaves": [9, 9, 9],
+ "cqtd_bins_per_octaves": [24, 36, 48],
+
+ "mpd_reshapes": [2, 3, 5, 7, 11],
+ "use_spectral_norm": false,
+ "discriminator_channel_mult": 1,
+
+ "use_multiscale_melloss": true,
+ "lambda_melloss": 15,
+
+ "clip_grad_norm": 500,
+
+ "segment_size": 65536,
+ "num_mels": 80,
+ "num_freq": 1025,
+ "n_fft": 1024,
+ "hop_size": 256,
+ "win_size": 1024,
+
+ "sampling_rate": 22050,
+
+ "fmin": 0,
+ "fmax": null,
+ "fmax_for_loss": null,
+
+ "normalize_volume": true,
+
+ "num_workers": 4,
+
+ "dist_config": {
+ "dist_backend": "nccl",
+ "dist_url": "tcp://localhost:54321",
+ "world_size": 1
+ }
+}
diff --git a/modules/bigvgan/env.py b/modules/bigvgan/env.py
new file mode 100644
index 0000000000000000000000000000000000000000..6f22811bf32ed67ddb58a44641ff7ff122fab107
--- /dev/null
+++ b/modules/bigvgan/env.py
@@ -0,0 +1,18 @@
+# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
+# LICENSE is in incl_licenses directory.
+
+import os
+import shutil
+
+
+class AttrDict(dict):
+ def __init__(self, *args, **kwargs):
+ super(AttrDict, self).__init__(*args, **kwargs)
+ self.__dict__ = self
+
+
+def build_env(config, config_name, path):
+ t_path = os.path.join(path, config_name)
+ if config != t_path:
+ os.makedirs(path, exist_ok=True)
+ shutil.copyfile(config, os.path.join(path, config_name))
\ No newline at end of file
diff --git a/modules/bigvgan/meldataset.py b/modules/bigvgan/meldataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..154811c6f8c668e7ffed3318bbaf7d25ffdca0a3
--- /dev/null
+++ b/modules/bigvgan/meldataset.py
@@ -0,0 +1,354 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
+# LICENSE is in incl_licenses directory.
+
+import math
+import os
+import random
+import torch
+import torch.utils.data
+import numpy as np
+from librosa.util import normalize
+from scipy.io.wavfile import read
+from librosa.filters import mel as librosa_mel_fn
+import pathlib
+from tqdm import tqdm
+
+MAX_WAV_VALUE = 32767.0 # NOTE: 32768.0 -1 to prevent int16 overflow (results in popping sound in corner cases)
+
+
+def load_wav(full_path, sr_target):
+ sampling_rate, data = read(full_path)
+ if sampling_rate != sr_target:
+ raise RuntimeError(
+ f"Sampling rate of the file {full_path} is {sampling_rate} Hz, but the model requires {sr_target} Hz"
+ )
+ return data, sampling_rate
+
+
+def dynamic_range_compression(x, C=1, clip_val=1e-5):
+ return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
+
+
+def dynamic_range_decompression(x, C=1):
+ return np.exp(x) / C
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+ return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression_torch(x, C=1):
+ return torch.exp(x) / C
+
+
+def spectral_normalize_torch(magnitudes):
+ return dynamic_range_compression_torch(magnitudes)
+
+
+def spectral_de_normalize_torch(magnitudes):
+ return dynamic_range_decompression_torch(magnitudes)
+
+
+mel_basis_cache = {}
+hann_window_cache = {}
+
+
+def mel_spectrogram(
+ y: torch.Tensor,
+ n_fft: int,
+ num_mels: int,
+ sampling_rate: int,
+ hop_size: int,
+ win_size: int,
+ fmin: int,
+ fmax: int = None,
+ center: bool = False,
+) -> torch.Tensor:
+ """
+ Calculate the mel spectrogram of an input signal.
+ This function uses slaney norm for the librosa mel filterbank (using librosa.filters.mel) and uses Hann window for STFT (using torch.stft).
+
+ Args:
+ y (torch.Tensor): Input signal.
+ n_fft (int): FFT size.
+ num_mels (int): Number of mel bins.
+ sampling_rate (int): Sampling rate of the input signal.
+ hop_size (int): Hop size for STFT.
+ win_size (int): Window size for STFT.
+ fmin (int): Minimum frequency for mel filterbank.
+ fmax (int): Maximum frequency for mel filterbank. If None, defaults to half the sampling rate (fmax = sr / 2.0) inside librosa_mel_fn
+ center (bool): Whether to pad the input to center the frames. Default is False.
+
+ Returns:
+ torch.Tensor: Mel spectrogram.
+ """
+ if torch.min(y) < -1.0:
+ print(f"[WARNING] Min value of input waveform signal is {torch.min(y)}")
+ if torch.max(y) > 1.0:
+ print(f"[WARNING] Max value of input waveform signal is {torch.max(y)}")
+
+ device = y.device
+ key = f"{n_fft}_{num_mels}_{sampling_rate}_{hop_size}_{win_size}_{fmin}_{fmax}_{device}"
+
+ if key not in mel_basis_cache:
+ mel = librosa_mel_fn(
+ sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax
+ )
+ mel_basis_cache[key] = torch.from_numpy(mel).float().to(device)
+ hann_window_cache[key] = torch.hann_window(win_size).to(device)
+
+ mel_basis = mel_basis_cache[key]
+ hann_window = hann_window_cache[key]
+
+ padding = (n_fft - hop_size) // 2
+ y = torch.nn.functional.pad(
+ y.unsqueeze(1), (padding, padding), mode="reflect"
+ ).squeeze(1)
+
+ spec = torch.stft(
+ y,
+ n_fft,
+ hop_length=hop_size,
+ win_length=win_size,
+ window=hann_window,
+ center=center,
+ pad_mode="reflect",
+ normalized=False,
+ onesided=True,
+ return_complex=True,
+ )
+ spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
+
+ mel_spec = torch.matmul(mel_basis, spec)
+ mel_spec = spectral_normalize_torch(mel_spec)
+
+ return mel_spec
+
+
+def get_mel_spectrogram(wav, h):
+ """
+ Generate mel spectrogram from a waveform using given hyperparameters.
+
+ Args:
+ wav (torch.Tensor): Input waveform.
+ h: Hyperparameters object with attributes n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax.
+
+ Returns:
+ torch.Tensor: Mel spectrogram.
+ """
+ return mel_spectrogram(
+ wav,
+ h.n_fft,
+ h.num_mels,
+ h.sampling_rate,
+ h.hop_size,
+ h.win_size,
+ h.fmin,
+ h.fmax,
+ )
+
+
+def get_dataset_filelist(a):
+ training_files = []
+ validation_files = []
+ list_unseen_validation_files = []
+
+ with open(a.input_training_file, "r", encoding="utf-8") as fi:
+ training_files = [
+ os.path.join(a.input_wavs_dir, x.split("|")[0] + ".wav")
+ for x in fi.read().split("\n")
+ if len(x) > 0
+ ]
+ print(f"first training file: {training_files[0]}")
+
+ with open(a.input_validation_file, "r", encoding="utf-8") as fi:
+ validation_files = [
+ os.path.join(a.input_wavs_dir, x.split("|")[0] + ".wav")
+ for x in fi.read().split("\n")
+ if len(x) > 0
+ ]
+ print(f"first validation file: {validation_files[0]}")
+
+ for i in range(len(a.list_input_unseen_validation_file)):
+ with open(a.list_input_unseen_validation_file[i], "r", encoding="utf-8") as fi:
+ unseen_validation_files = [
+ os.path.join(a.list_input_unseen_wavs_dir[i], x.split("|")[0] + ".wav")
+ for x in fi.read().split("\n")
+ if len(x) > 0
+ ]
+ print(
+ f"first unseen {i}th validation fileset: {unseen_validation_files[0]}"
+ )
+ list_unseen_validation_files.append(unseen_validation_files)
+
+ return training_files, validation_files, list_unseen_validation_files
+
+
+class MelDataset(torch.utils.data.Dataset):
+ def __init__(
+ self,
+ training_files,
+ hparams,
+ segment_size,
+ n_fft,
+ num_mels,
+ hop_size,
+ win_size,
+ sampling_rate,
+ fmin,
+ fmax,
+ split=True,
+ shuffle=True,
+ n_cache_reuse=1,
+ device=None,
+ fmax_loss=None,
+ fine_tuning=False,
+ base_mels_path=None,
+ is_seen=True,
+ ):
+ self.audio_files = training_files
+ random.seed(1234)
+ if shuffle:
+ random.shuffle(self.audio_files)
+ self.hparams = hparams
+ self.is_seen = is_seen
+ if self.is_seen:
+ self.name = pathlib.Path(self.audio_files[0]).parts[0]
+ else:
+ self.name = "-".join(pathlib.Path(self.audio_files[0]).parts[:2]).strip("/")
+
+ self.segment_size = segment_size
+ self.sampling_rate = sampling_rate
+ self.split = split
+ self.n_fft = n_fft
+ self.num_mels = num_mels
+ self.hop_size = hop_size
+ self.win_size = win_size
+ self.fmin = fmin
+ self.fmax = fmax
+ self.fmax_loss = fmax_loss
+ self.cached_wav = None
+ self.n_cache_reuse = n_cache_reuse
+ self._cache_ref_count = 0
+ self.device = device
+ self.fine_tuning = fine_tuning
+ self.base_mels_path = base_mels_path
+
+ print("[INFO] checking dataset integrity...")
+ for i in tqdm(range(len(self.audio_files))):
+ assert os.path.exists(
+ self.audio_files[i]
+ ), f"{self.audio_files[i]} not found"
+
+ def __getitem__(self, index):
+ filename = self.audio_files[index]
+ if self._cache_ref_count == 0:
+ audio, sampling_rate = load_wav(filename, self.sampling_rate)
+ audio = audio / MAX_WAV_VALUE
+ if not self.fine_tuning:
+ audio = normalize(audio) * 0.95
+ self.cached_wav = audio
+ if sampling_rate != self.sampling_rate:
+ raise ValueError(
+ f"{sampling_rate} SR doesn't match target {self.sampling_rate} SR"
+ )
+ self._cache_ref_count = self.n_cache_reuse
+ else:
+ audio = self.cached_wav
+ self._cache_ref_count -= 1
+
+ audio = torch.FloatTensor(audio)
+ audio = audio.unsqueeze(0)
+
+ if not self.fine_tuning:
+ if self.split:
+ if audio.size(1) >= self.segment_size:
+ max_audio_start = audio.size(1) - self.segment_size
+ audio_start = random.randint(0, max_audio_start)
+ audio = audio[:, audio_start : audio_start + self.segment_size]
+ else:
+ audio = torch.nn.functional.pad(
+ audio, (0, self.segment_size - audio.size(1)), "constant"
+ )
+
+ mel = mel_spectrogram(
+ audio,
+ self.n_fft,
+ self.num_mels,
+ self.sampling_rate,
+ self.hop_size,
+ self.win_size,
+ self.fmin,
+ self.fmax,
+ center=False,
+ )
+ else: # Validation step
+ # Match audio length to self.hop_size * n for evaluation
+ if (audio.size(1) % self.hop_size) != 0:
+ audio = audio[:, : -(audio.size(1) % self.hop_size)]
+ mel = mel_spectrogram(
+ audio,
+ self.n_fft,
+ self.num_mels,
+ self.sampling_rate,
+ self.hop_size,
+ self.win_size,
+ self.fmin,
+ self.fmax,
+ center=False,
+ )
+ assert (
+ audio.shape[1] == mel.shape[2] * self.hop_size
+ ), f"audio shape {audio.shape} mel shape {mel.shape}"
+
+ else:
+ mel = np.load(
+ os.path.join(
+ self.base_mels_path,
+ os.path.splitext(os.path.split(filename)[-1])[0] + ".npy",
+ )
+ )
+ mel = torch.from_numpy(mel)
+
+ if len(mel.shape) < 3:
+ mel = mel.unsqueeze(0)
+
+ if self.split:
+ frames_per_seg = math.ceil(self.segment_size / self.hop_size)
+
+ if audio.size(1) >= self.segment_size:
+ mel_start = random.randint(0, mel.size(2) - frames_per_seg - 1)
+ mel = mel[:, :, mel_start : mel_start + frames_per_seg]
+ audio = audio[
+ :,
+ mel_start
+ * self.hop_size : (mel_start + frames_per_seg)
+ * self.hop_size,
+ ]
+ else:
+ mel = torch.nn.functional.pad(
+ mel, (0, frames_per_seg - mel.size(2)), "constant"
+ )
+ audio = torch.nn.functional.pad(
+ audio, (0, self.segment_size - audio.size(1)), "constant"
+ )
+
+ mel_loss = mel_spectrogram(
+ audio,
+ self.n_fft,
+ self.num_mels,
+ self.sampling_rate,
+ self.hop_size,
+ self.win_size,
+ self.fmin,
+ self.fmax_loss,
+ center=False,
+ )
+
+ return (mel.squeeze(), audio.squeeze(0), filename, mel_loss.squeeze())
+
+ def __len__(self):
+ return len(self.audio_files)
diff --git a/modules/bigvgan/utils.py b/modules/bigvgan/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..538591e1cabb7d72caf35820125e08d7789a1461
--- /dev/null
+++ b/modules/bigvgan/utils.py
@@ -0,0 +1,99 @@
+# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
+# LICENSE is in incl_licenses directory.
+
+import glob
+import os
+import matplotlib
+import torch
+from torch.nn.utils import weight_norm
+
+matplotlib.use("Agg")
+import matplotlib.pylab as plt
+from .meldataset import MAX_WAV_VALUE
+from scipy.io.wavfile import write
+
+
+def plot_spectrogram(spectrogram):
+ fig, ax = plt.subplots(figsize=(10, 2))
+ im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
+ plt.colorbar(im, ax=ax)
+
+ fig.canvas.draw()
+ plt.close()
+
+ return fig
+
+
+def plot_spectrogram_clipped(spectrogram, clip_max=2.0):
+ fig, ax = plt.subplots(figsize=(10, 2))
+ im = ax.imshow(
+ spectrogram,
+ aspect="auto",
+ origin="lower",
+ interpolation="none",
+ vmin=1e-6,
+ vmax=clip_max,
+ )
+ plt.colorbar(im, ax=ax)
+
+ fig.canvas.draw()
+ plt.close()
+
+ return fig
+
+
+def init_weights(m, mean=0.0, std=0.01):
+ classname = m.__class__.__name__
+ if classname.find("Conv") != -1:
+ m.weight.data.normal_(mean, std)
+
+
+def apply_weight_norm(m):
+ classname = m.__class__.__name__
+ if classname.find("Conv") != -1:
+ weight_norm(m)
+
+
+def get_padding(kernel_size, dilation=1):
+ return int((kernel_size * dilation - dilation) / 2)
+
+
+def load_checkpoint(filepath, device):
+ assert os.path.isfile(filepath)
+ print(f"Loading '{filepath}'")
+ checkpoint_dict = torch.load(filepath, map_location=device)
+ print("Complete.")
+ return checkpoint_dict
+
+
+def save_checkpoint(filepath, obj):
+ print(f"Saving checkpoint to {filepath}")
+ torch.save(obj, filepath)
+ print("Complete.")
+
+
+def scan_checkpoint(cp_dir, prefix, renamed_file=None):
+ # Fallback to original scanning logic first
+ pattern = os.path.join(cp_dir, prefix + "????????")
+ cp_list = glob.glob(pattern)
+
+ if len(cp_list) > 0:
+ last_checkpoint_path = sorted(cp_list)[-1]
+ print(f"[INFO] Resuming from checkpoint: '{last_checkpoint_path}'")
+ return last_checkpoint_path
+
+ # If no pattern-based checkpoints are found, check for renamed file
+ if renamed_file:
+ renamed_path = os.path.join(cp_dir, renamed_file)
+ if os.path.isfile(renamed_path):
+ print(f"[INFO] Resuming from renamed checkpoint: '{renamed_file}'")
+ return renamed_path
+
+ return None
+
+
+def save_audio(audio, path, sr):
+ # wav: torch with 1d shape
+ audio = audio * MAX_WAV_VALUE
+ audio = audio.cpu().numpy().astype("int16")
+ write(path, sr, audio)
diff --git a/modules/campplus/DTDNN.py b/modules/campplus/DTDNN.py
new file mode 100644
index 0000000000000000000000000000000000000000..0ba03bd3033da3eb0a72e6acb569b15a3188dae5
--- /dev/null
+++ b/modules/campplus/DTDNN.py
@@ -0,0 +1,115 @@
+# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+
+from collections import OrderedDict
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from modules.campplus.layers import DenseLayer, StatsPool, TDNNLayer, CAMDenseTDNNBlock, TransitLayer, BasicResBlock, get_nonlinear
+
+
+class FCM(nn.Module):
+ def __init__(self,
+ block=BasicResBlock,
+ num_blocks=[2, 2],
+ m_channels=32,
+ feat_dim=80):
+ super(FCM, self).__init__()
+ self.in_planes = m_channels
+ self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(m_channels)
+
+ self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2)
+ self.layer2 = self._make_layer(block, m_channels, num_blocks[1], stride=2)
+
+ self.conv2 = nn.Conv2d(m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False)
+ self.bn2 = nn.BatchNorm2d(m_channels)
+ self.out_channels = m_channels * (feat_dim // 8)
+
+ def _make_layer(self, block, planes, num_blocks, stride):
+ strides = [stride] + [1] * (num_blocks - 1)
+ layers = []
+ for stride in strides:
+ layers.append(block(self.in_planes, planes, stride))
+ self.in_planes = planes * block.expansion
+ return nn.Sequential(*layers)
+
+ def forward(self, x):
+ x = x.unsqueeze(1)
+ out = F.relu(self.bn1(self.conv1(x)))
+ out = self.layer1(out)
+ out = self.layer2(out)
+ out = F.relu(self.bn2(self.conv2(out)))
+
+ shape = out.shape
+ out = out.reshape(shape[0], shape[1]*shape[2], shape[3])
+ return out
+
+class CAMPPlus(nn.Module):
+ def __init__(self,
+ feat_dim=80,
+ embedding_size=512,
+ growth_rate=32,
+ bn_size=4,
+ init_channels=128,
+ config_str='batchnorm-relu',
+ memory_efficient=True):
+ super(CAMPPlus, self).__init__()
+
+ self.head = FCM(feat_dim=feat_dim)
+ channels = self.head.out_channels
+
+ self.xvector = nn.Sequential(
+ OrderedDict([
+
+ ('tdnn',
+ TDNNLayer(channels,
+ init_channels,
+ 5,
+ stride=2,
+ dilation=1,
+ padding=-1,
+ config_str=config_str)),
+ ]))
+ channels = init_channels
+ for i, (num_layers, kernel_size,
+ dilation) in enumerate(zip((12, 24, 16), (3, 3, 3), (1, 2, 2))):
+ block = CAMDenseTDNNBlock(num_layers=num_layers,
+ in_channels=channels,
+ out_channels=growth_rate,
+ bn_channels=bn_size * growth_rate,
+ kernel_size=kernel_size,
+ dilation=dilation,
+ config_str=config_str,
+ memory_efficient=memory_efficient)
+ self.xvector.add_module('block%d' % (i + 1), block)
+ channels = channels + num_layers * growth_rate
+ self.xvector.add_module(
+ 'transit%d' % (i + 1),
+ TransitLayer(channels,
+ channels // 2,
+ bias=False,
+ config_str=config_str))
+ channels //= 2
+
+ self.xvector.add_module(
+ 'out_nonlinear', get_nonlinear(config_str, channels))
+
+ self.xvector.add_module('stats', StatsPool())
+ self.xvector.add_module(
+ 'dense',
+ DenseLayer(channels * 2, embedding_size, config_str='batchnorm_'))
+
+ for m in self.modules():
+ if isinstance(m, (nn.Conv1d, nn.Linear)):
+ nn.init.kaiming_normal_(m.weight.data)
+ if m.bias is not None:
+ nn.init.zeros_(m.bias)
+
+ def forward(self, x):
+ x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
+ x = self.head(x)
+ x = self.xvector(x)
+ return x
\ No newline at end of file
diff --git a/modules/campplus/__pycache__/DTDNN.cpython-310.pyc b/modules/campplus/__pycache__/DTDNN.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..af5a2747ecabbbc92afe18b634b7bb419a860c05
Binary files /dev/null and b/modules/campplus/__pycache__/DTDNN.cpython-310.pyc differ
diff --git a/modules/campplus/__pycache__/layers.cpython-310.pyc b/modules/campplus/__pycache__/layers.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c219af4ba58ae41cdb69bf30ea3c7d0864b85177
Binary files /dev/null and b/modules/campplus/__pycache__/layers.cpython-310.pyc differ
diff --git a/modules/campplus/classifier.py b/modules/campplus/classifier.py
new file mode 100644
index 0000000000000000000000000000000000000000..0f1c03e94737ddde0bd72ca850f43f7d760b4e73
--- /dev/null
+++ b/modules/campplus/classifier.py
@@ -0,0 +1,70 @@
+# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from modules.campplus.layers import DenseLayer
+
+
+class CosineClassifier(nn.Module):
+ def __init__(
+ self,
+ input_dim,
+ num_blocks=0,
+ inter_dim=512,
+ out_neurons=1000,
+ ):
+
+ super().__init__()
+ self.blocks = nn.ModuleList()
+
+ for index in range(num_blocks):
+ self.blocks.append(
+ DenseLayer(input_dim, inter_dim, config_str='batchnorm')
+ )
+ input_dim = inter_dim
+
+ self.weight = nn.Parameter(
+ torch.FloatTensor(out_neurons, input_dim)
+ )
+ nn.init.xavier_uniform_(self.weight)
+
+ def forward(self, x):
+ # x: [B, dim]
+ for layer in self.blocks:
+ x = layer(x)
+
+ # normalized
+ x = F.linear(F.normalize(x), F.normalize(self.weight))
+ return x
+
+class LinearClassifier(nn.Module):
+ def __init__(
+ self,
+ input_dim,
+ num_blocks=0,
+ inter_dim=512,
+ out_neurons=1000,
+ ):
+
+ super().__init__()
+ self.blocks = nn.ModuleList()
+
+ self.nonlinear = nn.ReLU(inplace=True)
+ for index in range(num_blocks):
+ self.blocks.append(
+ DenseLayer(input_dim, inter_dim, bias=True)
+ )
+ input_dim = inter_dim
+
+ self.linear = nn.Linear(input_dim, out_neurons, bias=True)
+
+ def forward(self, x):
+ # x: [B, dim]
+ x = self.nonlinear(x)
+ for layer in self.blocks:
+ x = layer(x)
+ x = self.linear(x)
+ return x
\ No newline at end of file
diff --git a/modules/campplus/layers.py b/modules/campplus/layers.py
new file mode 100644
index 0000000000000000000000000000000000000000..5403a698f0afd143c1c80043238a5d90a166a555
--- /dev/null
+++ b/modules/campplus/layers.py
@@ -0,0 +1,253 @@
+# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint as cp
+from torch import nn
+
+
+def get_nonlinear(config_str, channels):
+ nonlinear = nn.Sequential()
+ for name in config_str.split('-'):
+ if name == 'relu':
+ nonlinear.add_module('relu', nn.ReLU(inplace=True))
+ elif name == 'prelu':
+ nonlinear.add_module('prelu', nn.PReLU(channels))
+ elif name == 'batchnorm':
+ nonlinear.add_module('batchnorm', nn.BatchNorm1d(channels))
+ elif name == 'batchnorm_':
+ nonlinear.add_module('batchnorm',
+ nn.BatchNorm1d(channels, affine=False))
+ else:
+ raise ValueError('Unexpected module ({}).'.format(name))
+ return nonlinear
+
+def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, eps=1e-2):
+ mean = x.mean(dim=dim)
+ std = x.std(dim=dim, unbiased=unbiased)
+ stats = torch.cat([mean, std], dim=-1)
+ if keepdim:
+ stats = stats.unsqueeze(dim=dim)
+ return stats
+
+
+class StatsPool(nn.Module):
+ def forward(self, x):
+ return statistics_pooling(x)
+
+
+class TDNNLayer(nn.Module):
+ def __init__(self,
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride=1,
+ padding=0,
+ dilation=1,
+ bias=False,
+ config_str='batchnorm-relu'):
+ super(TDNNLayer, self).__init__()
+ if padding < 0:
+ assert kernel_size % 2 == 1, 'Expect equal paddings, but got even kernel size ({})'.format(
+ kernel_size)
+ padding = (kernel_size - 1) // 2 * dilation
+ self.linear = nn.Conv1d(in_channels,
+ out_channels,
+ kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ bias=bias)
+ self.nonlinear = get_nonlinear(config_str, out_channels)
+
+ def forward(self, x):
+ x = self.linear(x)
+ x = self.nonlinear(x)
+ return x
+
+
+class CAMLayer(nn.Module):
+ def __init__(self,
+ bn_channels,
+ out_channels,
+ kernel_size,
+ stride,
+ padding,
+ dilation,
+ bias,
+ reduction=2):
+ super(CAMLayer, self).__init__()
+ self.linear_local = nn.Conv1d(bn_channels,
+ out_channels,
+ kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ bias=bias)
+ self.linear1 = nn.Conv1d(bn_channels, bn_channels // reduction, 1)
+ self.relu = nn.ReLU(inplace=True)
+ self.linear2 = nn.Conv1d(bn_channels // reduction, out_channels, 1)
+ self.sigmoid = nn.Sigmoid()
+
+ def forward(self, x):
+ y = self.linear_local(x)
+ context = x.mean(-1, keepdim=True)+self.seg_pooling(x)
+ context = self.relu(self.linear1(context))
+ m = self.sigmoid(self.linear2(context))
+ return y*m
+
+ def seg_pooling(self, x, seg_len=100, stype='avg'):
+ if stype == 'avg':
+ seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
+ elif stype == 'max':
+ seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
+ else:
+ raise ValueError('Wrong segment pooling type.')
+ shape = seg.shape
+ seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1)
+ seg = seg[..., :x.shape[-1]]
+ return seg
+
+
+class CAMDenseTDNNLayer(nn.Module):
+ def __init__(self,
+ in_channels,
+ out_channels,
+ bn_channels,
+ kernel_size,
+ stride=1,
+ dilation=1,
+ bias=False,
+ config_str='batchnorm-relu',
+ memory_efficient=False):
+ super(CAMDenseTDNNLayer, self).__init__()
+ assert kernel_size % 2 == 1, 'Expect equal paddings, but got even kernel size ({})'.format(
+ kernel_size)
+ padding = (kernel_size - 1) // 2 * dilation
+ self.memory_efficient = memory_efficient
+ self.nonlinear1 = get_nonlinear(config_str, in_channels)
+ self.linear1 = nn.Conv1d(in_channels, bn_channels, 1, bias=False)
+ self.nonlinear2 = get_nonlinear(config_str, bn_channels)
+ self.cam_layer = CAMLayer(bn_channels,
+ out_channels,
+ kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ bias=bias)
+
+ def bn_function(self, x):
+ return self.linear1(self.nonlinear1(x))
+
+ def forward(self, x):
+ if self.training and self.memory_efficient:
+ x = cp.checkpoint(self.bn_function, x)
+ else:
+ x = self.bn_function(x)
+ x = self.cam_layer(self.nonlinear2(x))
+ return x
+
+
+class CAMDenseTDNNBlock(nn.ModuleList):
+ def __init__(self,
+ num_layers,
+ in_channels,
+ out_channels,
+ bn_channels,
+ kernel_size,
+ stride=1,
+ dilation=1,
+ bias=False,
+ config_str='batchnorm-relu',
+ memory_efficient=False):
+ super(CAMDenseTDNNBlock, self).__init__()
+ for i in range(num_layers):
+ layer = CAMDenseTDNNLayer(in_channels=in_channels + i * out_channels,
+ out_channels=out_channels,
+ bn_channels=bn_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ dilation=dilation,
+ bias=bias,
+ config_str=config_str,
+ memory_efficient=memory_efficient)
+ self.add_module('tdnnd%d' % (i + 1), layer)
+
+ def forward(self, x):
+ for layer in self:
+ x = torch.cat([x, layer(x)], dim=1)
+ return x
+
+
+class TransitLayer(nn.Module):
+ def __init__(self,
+ in_channels,
+ out_channels,
+ bias=True,
+ config_str='batchnorm-relu'):
+ super(TransitLayer, self).__init__()
+ self.nonlinear = get_nonlinear(config_str, in_channels)
+ self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
+
+ def forward(self, x):
+ x = self.nonlinear(x)
+ x = self.linear(x)
+ return x
+
+
+class DenseLayer(nn.Module):
+ def __init__(self,
+ in_channels,
+ out_channels,
+ bias=False,
+ config_str='batchnorm-relu'):
+ super(DenseLayer, self).__init__()
+ self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
+ self.nonlinear = get_nonlinear(config_str, out_channels)
+
+ def forward(self, x):
+ if len(x.shape) == 2:
+ x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1)
+ else:
+ x = self.linear(x)
+ x = self.nonlinear(x)
+ return x
+
+
+class BasicResBlock(nn.Module):
+ expansion = 1
+
+ def __init__(self, in_planes, planes, stride=1):
+ super(BasicResBlock, self).__init__()
+ self.conv1 = nn.Conv2d(in_planes,
+ planes,
+ kernel_size=3,
+ stride=(stride, 1),
+ padding=1,
+ bias=False)
+ self.bn1 = nn.BatchNorm2d(planes)
+ self.conv2 = nn.Conv2d(planes,
+ planes,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ bias=False)
+ self.bn2 = nn.BatchNorm2d(planes)
+
+ self.shortcut = nn.Sequential()
+ if stride != 1 or in_planes != self.expansion * planes:
+ self.shortcut = nn.Sequential(
+ nn.Conv2d(in_planes,
+ self.expansion * planes,
+ kernel_size=1,
+ stride=(stride, 1),
+ bias=False),
+ nn.BatchNorm2d(self.expansion * planes))
+
+ def forward(self, x):
+ out = F.relu(self.bn1(self.conv1(x)))
+ out = self.bn2(self.conv2(out))
+ out += self.shortcut(x)
+ out = F.relu(out)
+ return out
\ No newline at end of file
diff --git a/modules/commons.py b/modules/commons.py
new file mode 100644
index 0000000000000000000000000000000000000000..34e8b5ea2c08d96f08e17c300199f71daac2cdf0
--- /dev/null
+++ b/modules/commons.py
@@ -0,0 +1,476 @@
+import math
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+from munch import Munch
+import json
+import argparse
+
+def str2bool(v):
+ if isinstance(v, bool):
+ return v
+ if v.lower() in ("yes", "true", "t", "y", "1"):
+ return True
+ elif v.lower() in ("no", "false", "f", "n", "0"):
+ return False
+ else:
+ raise argparse.ArgumentTypeError("Boolean value expected.")
+
+class AttrDict(dict):
+ def __init__(self, *args, **kwargs):
+ super(AttrDict, self).__init__(*args, **kwargs)
+ self.__dict__ = self
+
+
+def init_weights(m, mean=0.0, std=0.01):
+ classname = m.__class__.__name__
+ if classname.find("Conv") != -1:
+ m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+ return int((kernel_size * dilation - dilation) / 2)
+
+
+def convert_pad_shape(pad_shape):
+ l = pad_shape[::-1]
+ pad_shape = [item for sublist in l for item in sublist]
+ return pad_shape
+
+
+def intersperse(lst, item):
+ result = [item] * (len(lst) * 2 + 1)
+ result[1::2] = lst
+ return result
+
+
+def kl_divergence(m_p, logs_p, m_q, logs_q):
+ """KL(P||Q)"""
+ kl = (logs_q - logs_p) - 0.5
+ kl += (
+ 0.5 * (torch.exp(2.0 * logs_p) + ((m_p - m_q) ** 2)) * torch.exp(-2.0 * logs_q)
+ )
+ return kl
+
+
+def rand_gumbel(shape):
+ """Sample from the Gumbel distribution, protect from overflows."""
+ uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
+ return -torch.log(-torch.log(uniform_samples))
+
+
+def rand_gumbel_like(x):
+ g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
+ return g
+
+
+def slice_segments(x, ids_str, segment_size=4):
+ ret = torch.zeros_like(x[:, :, :segment_size])
+ for i in range(x.size(0)):
+ idx_str = ids_str[i]
+ idx_end = idx_str + segment_size
+ ret[i] = x[i, :, idx_str:idx_end]
+ return ret
+
+
+def slice_segments_audio(x, ids_str, segment_size=4):
+ ret = torch.zeros_like(x[:, :segment_size])
+ for i in range(x.size(0)):
+ idx_str = ids_str[i]
+ idx_end = idx_str + segment_size
+ ret[i] = x[i, idx_str:idx_end]
+ return ret
+
+
+def rand_slice_segments(x, x_lengths=None, segment_size=4):
+ b, d, t = x.size()
+ if x_lengths is None:
+ x_lengths = t
+ ids_str_max = x_lengths - segment_size + 1
+ ids_str = ((torch.rand([b]).to(device=x.device) * ids_str_max).clip(0)).to(
+ dtype=torch.long
+ )
+ ret = slice_segments(x, ids_str, segment_size)
+ return ret, ids_str
+
+
+def get_timing_signal_1d(length, channels, min_timescale=1.0, max_timescale=1.0e4):
+ position = torch.arange(length, dtype=torch.float)
+ num_timescales = channels // 2
+ log_timescale_increment = math.log(float(max_timescale) / float(min_timescale)) / (
+ num_timescales - 1
+ )
+ inv_timescales = min_timescale * torch.exp(
+ torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment
+ )
+ scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
+ signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
+ signal = F.pad(signal, [0, 0, 0, channels % 2])
+ signal = signal.view(1, channels, length)
+ return signal
+
+
+def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4):
+ b, channels, length = x.size()
+ signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+ return x + signal.to(dtype=x.dtype, device=x.device)
+
+
+def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1):
+ b, channels, length = x.size()
+ signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+ return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)
+
+
+def subsequent_mask(length):
+ mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
+ return mask
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+ n_channels_int = n_channels[0]
+ in_act = input_a + input_b
+ t_act = torch.tanh(in_act[:, :n_channels_int, :])
+ s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+ acts = t_act * s_act
+ return acts
+
+
+def convert_pad_shape(pad_shape):
+ l = pad_shape[::-1]
+ pad_shape = [item for sublist in l for item in sublist]
+ return pad_shape
+
+
+def shift_1d(x):
+ x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
+ return x
+
+
+def sequence_mask(length, max_length=None):
+ if max_length is None:
+ max_length = length.max()
+ x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+ return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+def avg_with_mask(x, mask):
+ assert mask.dtype == torch.float, "Mask should be float"
+
+ if mask.ndim == 2:
+ mask = mask.unsqueeze(1)
+
+ if mask.shape[1] == 1:
+ mask = mask.expand_as(x)
+
+ return (x * mask).sum() / mask.sum()
+
+
+def generate_path(duration, mask):
+ """
+ duration: [b, 1, t_x]
+ mask: [b, 1, t_y, t_x]
+ """
+ device = duration.device
+
+ b, _, t_y, t_x = mask.shape
+ cum_duration = torch.cumsum(duration, -1)
+
+ cum_duration_flat = cum_duration.view(b * t_x)
+ path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
+ path = path.view(b, t_x, t_y)
+ path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+ path = path.unsqueeze(1).transpose(2, 3) * mask
+ return path
+
+
+def clip_grad_value_(parameters, clip_value, norm_type=2):
+ if isinstance(parameters, torch.Tensor):
+ parameters = [parameters]
+ parameters = list(filter(lambda p: p.grad is not None, parameters))
+ norm_type = float(norm_type)
+ if clip_value is not None:
+ clip_value = float(clip_value)
+
+ total_norm = 0
+ for p in parameters:
+ param_norm = p.grad.data.norm(norm_type)
+ total_norm += param_norm.item() ** norm_type
+ if clip_value is not None:
+ p.grad.data.clamp_(min=-clip_value, max=clip_value)
+ total_norm = total_norm ** (1.0 / norm_type)
+ return total_norm
+
+
+def log_norm(x, mean=-4, std=4, dim=2):
+ """
+ normalized log mel -> mel -> norm -> log(norm)
+ """
+ x = torch.log(torch.exp(x * std + mean).norm(dim=dim))
+ return x
+
+
+def load_F0_models(path):
+ # load F0 model
+ from .JDC.model import JDCNet
+
+ F0_model = JDCNet(num_class=1, seq_len=192)
+ params = torch.load(path, map_location="cpu")["net"]
+ F0_model.load_state_dict(params)
+ _ = F0_model.train()
+
+ return F0_model
+
+
+def modify_w2v_forward(self, output_layer=15):
+ """
+ change forward method of w2v encoder to get its intermediate layer output
+ :param self:
+ :param layer:
+ :return:
+ """
+ from transformers.modeling_outputs import BaseModelOutput
+
+ def forward(
+ hidden_states,
+ attention_mask=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ ):
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attentions = () if output_attentions else None
+
+ conv_attention_mask = attention_mask
+ if attention_mask is not None:
+ # make sure padded tokens output 0
+ hidden_states = hidden_states.masked_fill(
+ ~attention_mask.bool().unsqueeze(-1), 0.0
+ )
+
+ # extend attention_mask
+ attention_mask = 1.0 - attention_mask[:, None, None, :].to(
+ dtype=hidden_states.dtype
+ )
+ attention_mask = attention_mask * torch.finfo(hidden_states.dtype).min
+ attention_mask = attention_mask.expand(
+ attention_mask.shape[0],
+ 1,
+ attention_mask.shape[-1],
+ attention_mask.shape[-1],
+ )
+
+ hidden_states = self.dropout(hidden_states)
+
+ if self.embed_positions is not None:
+ relative_position_embeddings = self.embed_positions(hidden_states)
+ else:
+ relative_position_embeddings = None
+
+ deepspeed_zero3_is_enabled = False
+
+ for i, layer in enumerate(self.layers):
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
+ dropout_probability = torch.rand([])
+
+ skip_the_layer = (
+ True
+ if self.training and (dropout_probability < self.config.layerdrop)
+ else False
+ )
+ if not skip_the_layer or deepspeed_zero3_is_enabled:
+ # under deepspeed zero3 all gpus must run in sync
+ if self.gradient_checkpointing and self.training:
+ layer_outputs = self._gradient_checkpointing_func(
+ layer.__call__,
+ hidden_states,
+ attention_mask,
+ relative_position_embeddings,
+ output_attentions,
+ conv_attention_mask,
+ )
+ else:
+ layer_outputs = layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ relative_position_embeddings=relative_position_embeddings,
+ output_attentions=output_attentions,
+ conv_attention_mask=conv_attention_mask,
+ )
+ hidden_states = layer_outputs[0]
+
+ if skip_the_layer:
+ layer_outputs = (None, None)
+
+ if output_attentions:
+ all_self_attentions = all_self_attentions + (layer_outputs[1],)
+
+ if i == output_layer - 1:
+ break
+
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ if not return_dict:
+ return tuple(
+ v
+ for v in [hidden_states, all_hidden_states, all_self_attentions]
+ if v is not None
+ )
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ )
+
+ return forward
+
+
+MATPLOTLIB_FLAG = False
+
+
+def plot_spectrogram_to_numpy(spectrogram):
+ global MATPLOTLIB_FLAG
+ if not MATPLOTLIB_FLAG:
+ import matplotlib
+ import logging
+
+ matplotlib.use("Agg")
+ MATPLOTLIB_FLAG = True
+ mpl_logger = logging.getLogger("matplotlib")
+ mpl_logger.setLevel(logging.WARNING)
+ import matplotlib.pylab as plt
+ import numpy as np
+
+ fig, ax = plt.subplots(figsize=(10, 2))
+ im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
+ plt.colorbar(im, ax=ax)
+ plt.xlabel("Frames")
+ plt.ylabel("Channels")
+ plt.tight_layout()
+
+ fig.canvas.draw()
+ data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
+ data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+ plt.close()
+ return data
+
+
+def normalize_f0(f0_sequence):
+ # Remove unvoiced frames (replace with -1)
+ voiced_indices = np.where(f0_sequence > 0)[0]
+ f0_voiced = f0_sequence[voiced_indices]
+
+ # Convert to log scale
+ log_f0 = np.log2(f0_voiced)
+
+ # Calculate mean and standard deviation
+ mean_f0 = np.mean(log_f0)
+ std_f0 = np.std(log_f0)
+
+ # Normalize the F0 sequence
+ normalized_f0 = (log_f0 - mean_f0) / std_f0
+
+ # Create the normalized F0 sequence with unvoiced frames
+ normalized_sequence = np.zeros_like(f0_sequence)
+ normalized_sequence[voiced_indices] = normalized_f0
+ normalized_sequence[f0_sequence <= 0] = -1 # Assign -1 to unvoiced frames
+
+ return normalized_sequence
+
+
+def build_model(args, stage="DiT"):
+ if stage == "DiT":
+ from modules.flow_matching import CFM
+ from modules.length_regulator import InterpolateRegulator
+
+ length_regulator = InterpolateRegulator(
+ channels=args.length_regulator.channels,
+ sampling_ratios=args.length_regulator.sampling_ratios,
+ is_discrete=args.length_regulator.is_discrete,
+ in_channels=args.length_regulator.in_channels if hasattr(args.length_regulator, "in_channels") else None,
+ codebook_size=args.length_regulator.content_codebook_size,
+ f0_condition=args.length_regulator.f0_condition if hasattr(args.length_regulator, "f0_condition") else False,
+ n_f0_bins=args.length_regulator.n_f0_bins if hasattr(args.length_regulator, "n_f0_bins") else 512,
+ )
+ cfm = CFM(args)
+ nets = Munch(
+ cfm=cfm,
+ length_regulator=length_regulator,
+ )
+ else:
+ raise ValueError(f"Unknown stage: {stage}")
+
+ return nets
+
+
+def load_checkpoint(
+ model,
+ optimizer,
+ path,
+ load_only_params=True,
+ ignore_modules=[],
+ is_distributed=False,
+ load_ema=False,
+):
+ state = torch.load(path, map_location="cpu")
+ params = state["net"]
+ if load_ema and "ema" in state:
+ print("Loading EMA")
+ for key in model:
+ i = 0
+ for param_name in params[key]:
+ if "input_pos" in param_name:
+ continue
+ assert params[key][param_name].shape == state["ema"][key][0][i].shape
+ params[key][param_name] = state["ema"][key][0][i].clone()
+ i += 1
+ for key in model:
+ if key in params and key not in ignore_modules:
+ if not is_distributed:
+ # strip prefix of DDP (module.), create a new OrderedDict that does not contain the prefix
+ for k in list(params[key].keys()):
+ if k.startswith("module."):
+ params[key][k[len("module.") :]] = params[key][k]
+ del params[key][k]
+ model_state_dict = model[key].state_dict()
+
+ filtered_state_dict = {
+ k: v
+ for k, v in params[key].items()
+ if k in model_state_dict and v.shape == model_state_dict[k].shape
+ }
+ skipped_keys = set(params[key].keys()) - set(filtered_state_dict.keys())
+ if skipped_keys:
+ print(
+ f"Warning: Skipped loading some keys due to shape mismatch: {skipped_keys}"
+ )
+ print("%s loaded" % key)
+ model[key].load_state_dict(filtered_state_dict, strict=False)
+ _ = [model[key].eval() for key in model]
+
+ if not load_only_params:
+ epoch = state["epoch"] + 1
+ iters = state["iters"]
+ optimizer.load_state_dict(state["optimizer"])
+ optimizer.load_scheduler_state_dict(state["scheduler"])
+
+ else:
+ epoch = 0
+ iters = 0
+
+ return model, optimizer, epoch, iters
+
+
+def recursive_munch(d):
+ if isinstance(d, dict):
+ return Munch((k, recursive_munch(v)) for k, v in d.items())
+ elif isinstance(d, list):
+ return [recursive_munch(v) for v in d]
+ else:
+ return d
diff --git a/modules/cosyvoice_tokenizer/__pycache__/frontend.cpython-310.pyc b/modules/cosyvoice_tokenizer/__pycache__/frontend.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0a1490526d5d1dce3865f1b21f084c5a8d4a5528
Binary files /dev/null and b/modules/cosyvoice_tokenizer/__pycache__/frontend.cpython-310.pyc differ
diff --git a/modules/cosyvoice_tokenizer/frontend.py b/modules/cosyvoice_tokenizer/frontend.py
new file mode 100644
index 0000000000000000000000000000000000000000..a812eda9849b59ac394772c362c3104449898bd9
--- /dev/null
+++ b/modules/cosyvoice_tokenizer/frontend.py
@@ -0,0 +1,54 @@
+# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from functools import partial
+import onnxruntime
+import torch
+import numpy as np
+import whisper
+import torchaudio.compliance.kaldi as kaldi
+
+class CosyVoiceFrontEnd:
+
+ def __init__(self, speech_tokenizer_model: str, device: str = 'cuda', device_id: int = 0):
+ self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+ option = onnxruntime.SessionOptions()
+ option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+ option.intra_op_num_threads = 1
+ self.speech_tokenizer_session = onnxruntime.InferenceSession(speech_tokenizer_model, sess_options=option, providers=["CUDAExecutionProvider"if device == "cuda" else "CPUExecutionProvider"])
+ if device == 'cuda':
+ self.speech_tokenizer_session.set_providers(['CUDAExecutionProvider'], [ {'device_id': device_id}])
+
+ def extract_speech_token(self, speech):
+ feat = whisper.log_mel_spectrogram(speech, n_mels=128)
+ speech_token = self.speech_tokenizer_session.run(None, {self.speech_tokenizer_session.get_inputs()[0].name: feat.detach().cpu().numpy(),
+ self.speech_tokenizer_session.get_inputs()[1].name: np.array([feat.shape[2]], dtype=np.int32)})[0].flatten().tolist()
+ speech_token = torch.tensor([speech_token], dtype=torch.int32).to(self.device)
+ speech_token_len = torch.tensor([speech_token.shape[1]], dtype=torch.int32).to(self.device)
+ return speech_token, speech_token_len
+
+ def _extract_spk_embedding(self, speech):
+ feat = kaldi.fbank(speech,
+ num_mel_bins=80,
+ dither=0,
+ sample_frequency=16000)
+ feat = feat - feat.mean(dim=0, keepdim=True)
+ embedding = self.campplus_session.run(None, {self.campplus_session.get_inputs()[0].name: feat.unsqueeze(dim=0).cpu().numpy()})[0].flatten().tolist()
+ embedding = torch.tensor([embedding]).to(self.device)
+ return embedding
+
+ def _extract_speech_feat(self, speech):
+ speech_feat = self.feat_extractor(speech).squeeze(dim=0).transpose(0, 1).to(self.device)
+ speech_feat = speech_feat.unsqueeze(dim=0)
+ speech_feat_len = torch.tensor([speech_feat.shape[1]], dtype=torch.int32).to(self.device)
+ return speech_feat, speech_feat_len
\ No newline at end of file
diff --git a/modules/diffusion_transformer.py b/modules/diffusion_transformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..f9b468fa6701a72fab3e55e31dadc814e10c78f1
--- /dev/null
+++ b/modules/diffusion_transformer.py
@@ -0,0 +1,537 @@
+import torch
+from torch import nn
+import math
+
+# from modules.torchscript_modules.gpt_fast_model import ModelArgs, Transformer
+from modules.wavenet import WN
+from modules.commons import sequence_mask
+
+from torch.nn.utils import weight_norm
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import functional as F
+
+
+def find_multiple(n: int, k: int) -> int:
+ if n % k == 0:
+ return n
+ return n + k - (n % k)
+
+class AdaptiveLayerNorm(nn.Module):
+ r"""Adaptive Layer Normalization"""
+
+ def __init__(self, d_model, norm) -> None:
+ super(AdaptiveLayerNorm, self).__init__()
+ self.project_layer = nn.Linear(d_model, 2 * d_model)
+ self.norm = norm
+ self.d_model = d_model
+ self.eps = self.norm.eps
+
+ def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor:
+ if embedding is None:
+ return self.norm(input)
+ weight, bias = torch.split(
+ self.project_layer(embedding),
+ split_size_or_sections=self.d_model,
+ dim=-1,
+ )
+ return weight * self.norm(input) + bias
+
+
+@dataclass
+class ModelArgs:
+ block_size: int = 2048
+ vocab_size: int = 32000
+ n_layer: int = 32
+ n_head: int = 32
+ dim: int = 4096
+ intermediate_size: int = None
+ n_local_heads: int = -1
+ head_dim: int = 64
+ rope_base: float = 10000
+ norm_eps: float = 1e-5
+ has_cross_attention: bool = False
+ context_dim: int = 0
+ uvit_skip_connection: bool = False
+ time_as_token: bool = False
+
+ def __post_init__(self):
+ if self.n_local_heads == -1:
+ self.n_local_heads = self.n_head
+ if self.intermediate_size is None:
+ hidden_dim = 4 * self.dim
+ n_hidden = int(2 * hidden_dim / 3)
+ self.intermediate_size = find_multiple(n_hidden, 256)
+ # self.head_dim = self.dim // self.n_head
+
+class Transformer(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.config = config
+
+ self.layers = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer))
+ self.norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+
+ self.freqs_cis: Optional[Tensor] = None
+ self.mask_cache: Optional[Tensor] = None
+ self.max_batch_size = -1
+ self.max_seq_length = -1
+
+ def setup_caches(self, max_batch_size, max_seq_length, use_kv_cache=False):
+ if self.max_seq_length >= max_seq_length and self.max_batch_size >= max_batch_size:
+ return
+ head_dim = self.config.dim // self.config.n_head
+ max_seq_length = find_multiple(max_seq_length, 8)
+ self.max_seq_length = max_seq_length
+ self.max_batch_size = max_batch_size
+ dtype = self.norm.project_layer.weight.dtype
+ device = self.norm.project_layer.weight.device
+
+ self.freqs_cis = precompute_freqs_cis(self.config.block_size, self.config.head_dim,
+ self.config.rope_base, dtype).to(device)
+ self.causal_mask = torch.tril(torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool)).to(device)
+ self.use_kv_cache = use_kv_cache
+ self.uvit_skip_connection = self.config.uvit_skip_connection
+ if self.uvit_skip_connection:
+ self.layers_emit_skip = [i for i in range(self.config.n_layer) if i < self.config.n_layer // 2]
+ self.layers_receive_skip = [i for i in range(self.config.n_layer) if i > self.config.n_layer // 2]
+ else:
+ self.layers_emit_skip = []
+ self.layers_receive_skip = []
+
+ def forward(self,
+ x: Tensor,
+ c: Tensor,
+ input_pos: Optional[Tensor] = None,
+ mask: Optional[Tensor] = None,
+ context: Optional[Tensor] = None,
+ context_input_pos: Optional[Tensor] = None,
+ cross_attention_mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ assert self.freqs_cis is not None, "Caches must be initialized first"
+ if mask is None: # in case of non-causal model
+ if not self.training and self.use_kv_cache:
+ mask = self.causal_mask[None, None, input_pos]
+ else:
+ mask = self.causal_mask[None, None, input_pos]
+ mask = mask[..., input_pos]
+ freqs_cis = self.freqs_cis[input_pos]
+ if context is not None:
+ context_freqs_cis = self.freqs_cis[context_input_pos]
+ else:
+ context_freqs_cis = None
+ skip_in_x_list = []
+ for i, layer in enumerate(self.layers):
+ if self.uvit_skip_connection and i in self.layers_receive_skip:
+ skip_in_x = skip_in_x_list.pop(-1)
+ else:
+ skip_in_x = None
+ x = layer(x, c, input_pos, freqs_cis, mask, context, context_freqs_cis, cross_attention_mask, skip_in_x)
+ if self.uvit_skip_connection and i in self.layers_emit_skip:
+ skip_in_x_list.append(x)
+ x = self.norm(x, c)
+ return x
+
+ @classmethod
+ def from_name(cls, name: str):
+ return cls(ModelArgs.from_name(name))
+
+
+class TransformerBlock(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.attention = Attention(config)
+ self.feed_forward = FeedForward(config)
+ self.ffn_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+ self.attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+
+ if config.has_cross_attention:
+ self.has_cross_attention = True
+ self.cross_attention = Attention(config, is_cross_attention=True)
+ self.cross_attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+ else:
+ self.has_cross_attention = False
+
+ if config.uvit_skip_connection:
+ self.skip_in_linear = nn.Linear(config.dim * 2, config.dim)
+ self.uvit_skip_connection = True
+ else:
+ self.uvit_skip_connection = False
+
+ self.time_as_token = config.time_as_token
+
+ def forward(self,
+ x: Tensor,
+ c: Tensor,
+ input_pos: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ context: Optional[Tensor] = None,
+ context_freqs_cis: Optional[Tensor] = None,
+ cross_attention_mask: Optional[Tensor] = None,
+ skip_in_x: Optional[Tensor] = None,
+ ) -> Tensor:
+ c = None if self.time_as_token else c
+ if self.uvit_skip_connection and skip_in_x is not None:
+ x = self.skip_in_linear(torch.cat([x, skip_in_x], dim=-1))
+ h = x + self.attention(self.attention_norm(x, c), freqs_cis, mask, input_pos)
+ if self.has_cross_attention:
+ h = h + self.cross_attention(self.cross_attention_norm(h, c), freqs_cis, cross_attention_mask, input_pos, context, context_freqs_cis)
+ out = h + self.feed_forward(self.ffn_norm(h, c))
+ return out
+
+
+class Attention(nn.Module):
+ def __init__(self, config: ModelArgs, is_cross_attention: bool = False):
+ super().__init__()
+ assert config.dim % config.n_head == 0
+
+ total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim
+ # key, query, value projections for all heads, but in a batch
+ if is_cross_attention:
+ self.wq = nn.Linear(config.dim, config.n_head * config.head_dim, bias=False)
+ self.wkv = nn.Linear(config.context_dim, 2 * config.n_local_heads * config.head_dim, bias=False)
+ else:
+ self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False)
+ self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False)
+ self.kv_cache = None
+
+ self.n_head = config.n_head
+ self.head_dim = config.head_dim
+ self.n_local_heads = config.n_local_heads
+ self.dim = config.dim
+ # self._register_load_state_dict_pre_hook(self.load_hook)
+
+ # def load_hook(self, state_dict, prefix, *args):
+ # if prefix + "wq.weight" in state_dict:
+ # wq = state_dict.pop(prefix + "wq.weight")
+ # wk = state_dict.pop(prefix + "wk.weight")
+ # wv = state_dict.pop(prefix + "wv.weight")
+ # state_dict[prefix + "wqkv.weight"] = torch.cat([wq, wk, wv])
+
+ def forward(self,
+ x: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ input_pos: Optional[Tensor] = None,
+ context: Optional[Tensor] = None,
+ context_freqs_cis: Optional[Tensor] = None,
+ ) -> Tensor:
+ bsz, seqlen, _ = x.shape
+
+ kv_size = self.n_local_heads * self.head_dim
+ if context is None:
+ q, k, v = self.wqkv(x).split([kv_size, kv_size, kv_size], dim=-1)
+ context_seqlen = seqlen
+ else:
+ q = self.wq(x)
+ k, v = self.wkv(context).split([kv_size, kv_size], dim=-1)
+ context_seqlen = context.shape[1]
+
+ q = q.view(bsz, seqlen, self.n_head, self.head_dim)
+ k = k.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+ v = v.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+
+ q = apply_rotary_emb(q, freqs_cis)
+ k = apply_rotary_emb(k, context_freqs_cis if context_freqs_cis is not None else freqs_cis)
+
+ q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
+
+ if self.kv_cache is not None:
+ k, v = self.kv_cache.update(input_pos, k, v)
+
+ k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
+
+ y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.head_dim * self.n_head)
+
+ y = self.wo(y)
+ return y
+
+
+class FeedForward(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
+
+ def forward(self, x: Tensor) -> Tensor:
+ return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-5):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def _norm(self, x):
+ return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+ def forward(self, x: Tensor) -> Tensor:
+ output = self._norm(x.float()).type_as(x)
+ return output * self.weight
+
+
+def precompute_freqs_cis(
+ seq_len: int, n_elem: int, base: int = 10000,
+ dtype: torch.dtype = torch.bfloat16
+) -> Tensor:
+ freqs = 1.0 / (base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem))
+ t = torch.arange(seq_len, device=freqs.device)
+ freqs = torch.outer(t, freqs)
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+ cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
+ return cache.to(dtype=dtype)
+
+
+def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
+ xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+ freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2)
+ x_out2 = torch.stack(
+ [
+ xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],
+ xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],
+ ],
+ -1,
+ )
+
+ x_out2 = x_out2.flatten(3)
+ return x_out2.type_as(x)
+
+
+def modulate(x, shift, scale):
+ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+#################################################################################
+# Embedding Layers for Timesteps and Class Labels #
+#################################################################################
+
+class TimestepEmbedder(nn.Module):
+ """
+ Embeds scalar timesteps into vector representations.
+ """
+ def __init__(self, hidden_size, frequency_embedding_size=256):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+ nn.SiLU(),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+ self.max_period = 10000
+ self.scale = 1000
+
+ half = frequency_embedding_size // 2
+ freqs = torch.exp(
+ -math.log(self.max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+ )
+ self.register_buffer("freqs", freqs)
+
+ def timestep_embedding(self, t):
+ """
+ Create sinusoidal timestep embeddings.
+ :param t: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ :param dim: the dimension of the output.
+ :param max_period: controls the minimum frequency of the embeddings.
+ :return: an (N, D) Tensor of positional embeddings.
+ """
+ # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+
+ args = self.scale * t[:, None].float() * self.freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if self.frequency_embedding_size % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t):
+ t_freq = self.timestep_embedding(t)
+ t_emb = self.mlp(t_freq)
+ return t_emb
+
+
+class StyleEmbedder(nn.Module):
+ """
+ Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
+ """
+ def __init__(self, input_size, hidden_size, dropout_prob):
+ super().__init__()
+ use_cfg_embedding = dropout_prob > 0
+ self.embedding_table = nn.Embedding(int(use_cfg_embedding), hidden_size)
+ self.style_in = weight_norm(nn.Linear(input_size, hidden_size, bias=True))
+ self.input_size = input_size
+ self.dropout_prob = dropout_prob
+
+ def forward(self, labels, train, force_drop_ids=None):
+ use_dropout = self.dropout_prob > 0
+ if (train and use_dropout) or (force_drop_ids is not None):
+ labels = self.token_drop(labels, force_drop_ids)
+ else:
+ labels = self.style_in(labels)
+ embeddings = labels
+ return embeddings
+
+class FinalLayer(nn.Module):
+ """
+ The final layer of DiT.
+ """
+ def __init__(self, hidden_size, patch_size, out_channels):
+ super().__init__()
+ self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.linear = weight_norm(nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True))
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+ )
+
+ def forward(self, x, c):
+ shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+ x = modulate(self.norm_final(x), shift, scale)
+ x = self.linear(x)
+ return x
+
+class DiT(torch.nn.Module):
+ def __init__(
+ self,
+ args
+ ):
+ super(DiT, self).__init__()
+ self.time_as_token = args.DiT.time_as_token if hasattr(args.DiT, 'time_as_token') else False
+ self.style_as_token = args.DiT.style_as_token if hasattr(args.DiT, 'style_as_token') else False
+ self.uvit_skip_connection = args.DiT.uvit_skip_connection if hasattr(args.DiT, 'uvit_skip_connection') else False
+ model_args = ModelArgs(
+ block_size=16384,#args.DiT.block_size,
+ n_layer=args.DiT.depth,
+ n_head=args.DiT.num_heads,
+ dim=args.DiT.hidden_dim,
+ head_dim=args.DiT.hidden_dim // args.DiT.num_heads,
+ vocab_size=1024,
+ uvit_skip_connection=self.uvit_skip_connection,
+ time_as_token=self.time_as_token,
+ )
+ self.transformer = Transformer(model_args)
+ self.in_channels = args.DiT.in_channels
+ self.out_channels = args.DiT.in_channels
+ self.num_heads = args.DiT.num_heads
+
+ self.x_embedder = weight_norm(nn.Linear(args.DiT.in_channels, args.DiT.hidden_dim, bias=True))
+
+ self.content_type = args.DiT.content_type # 'discrete' or 'continuous'
+ self.content_codebook_size = args.DiT.content_codebook_size # for discrete content
+ self.content_dim = args.DiT.content_dim # for continuous content
+ self.cond_embedder = nn.Embedding(args.DiT.content_codebook_size, args.DiT.hidden_dim) # discrete content
+ self.cond_projection = nn.Linear(args.DiT.content_dim, args.DiT.hidden_dim, bias=True) # continuous content
+
+ self.is_causal = args.DiT.is_causal
+
+ self.t_embedder = TimestepEmbedder(args.DiT.hidden_dim)
+
+ input_pos = torch.arange(16384)
+ self.register_buffer("input_pos", input_pos)
+
+ self.final_layer_type = args.DiT.final_layer_type # mlp or wavenet
+ if self.final_layer_type == 'wavenet':
+ self.t_embedder2 = TimestepEmbedder(args.wavenet.hidden_dim)
+ self.conv1 = nn.Linear(args.DiT.hidden_dim, args.wavenet.hidden_dim)
+ self.conv2 = nn.Conv1d(args.wavenet.hidden_dim, args.DiT.in_channels, 1)
+ self.wavenet = WN(hidden_channels=args.wavenet.hidden_dim,
+ kernel_size=args.wavenet.kernel_size,
+ dilation_rate=args.wavenet.dilation_rate,
+ n_layers=args.wavenet.num_layers,
+ gin_channels=args.wavenet.hidden_dim,
+ p_dropout=args.wavenet.p_dropout,
+ causal=False)
+ self.final_layer = FinalLayer(args.wavenet.hidden_dim, 1, args.wavenet.hidden_dim)
+ self.res_projection = nn.Linear(args.DiT.hidden_dim,
+ args.wavenet.hidden_dim) # residual connection from tranformer output to final output
+ self.wavenet_style_condition = args.wavenet.style_condition
+ assert args.DiT.style_condition == args.wavenet.style_condition
+ else:
+ self.final_mlp = nn.Sequential(
+ nn.Linear(args.DiT.hidden_dim, args.DiT.hidden_dim),
+ nn.SiLU(),
+ nn.Linear(args.DiT.hidden_dim, args.DiT.in_channels),
+ )
+ self.transformer_style_condition = args.DiT.style_condition
+
+
+ self.class_dropout_prob = args.DiT.class_dropout_prob
+ self.content_mask_embedder = nn.Embedding(1, args.DiT.hidden_dim)
+
+ self.long_skip_connection = args.DiT.long_skip_connection
+ self.skip_linear = nn.Linear(args.DiT.hidden_dim + args.DiT.in_channels, args.DiT.hidden_dim)
+
+ self.cond_x_merge_linear = nn.Linear(args.DiT.hidden_dim + args.DiT.in_channels * 2 +
+ args.style_encoder.dim * self.transformer_style_condition * (not self.style_as_token),
+ args.DiT.hidden_dim)
+ if self.style_as_token:
+ self.style_in = nn.Linear(args.style_encoder.dim, args.DiT.hidden_dim)
+
+ def setup_caches(self, max_batch_size, max_seq_length):
+ self.transformer.setup_caches(max_batch_size, max_seq_length, use_kv_cache=False)
+ def forward(self, x, prompt_x, x_lens, t, style, cond, mask_content=False):
+ class_dropout = False
+ if self.training and torch.rand(1) < self.class_dropout_prob:
+ class_dropout = True
+ if not self.training and mask_content:
+ class_dropout = True
+ # cond_in_module = self.cond_embedder if self.content_type == 'discrete' else self.cond_projection
+ cond_in_module = self.cond_projection
+
+ B, _, T = x.size()
+
+
+ t1 = self.t_embedder(t) # (N, D)
+
+ cond = cond_in_module(cond)
+
+ x = x.transpose(1, 2)
+ prompt_x = prompt_x.transpose(1, 2)
+
+ x_in = torch.cat([x, prompt_x, cond], dim=-1)
+ if self.transformer_style_condition and not self.style_as_token:
+ x_in = torch.cat([x_in, style[:, None, :].repeat(1, T, 1)], dim=-1)
+ if class_dropout:
+ x_in[..., self.in_channels:] = x_in[..., self.in_channels:] * 0
+ x_in = self.cond_x_merge_linear(x_in) # (N, T, D)
+
+ if self.style_as_token:
+ style = self.style_in(style)
+ style = torch.zeros_like(style) if class_dropout else style
+ x_in = torch.cat([style.unsqueeze(1), x_in], dim=1)
+ if self.time_as_token:
+ x_in = torch.cat([t1.unsqueeze(1), x_in], dim=1)
+ x_mask = sequence_mask(x_lens + self.style_as_token + self.time_as_token).to(x.device).unsqueeze(1)
+ input_pos = self.input_pos[:x_in.size(1)] # (T,)
+ x_mask_expanded = x_mask[:, None, :].repeat(1, 1, x_in.size(1), 1) if not self.is_causal else None
+ x_res = self.transformer(x_in, t1.unsqueeze(1), input_pos, x_mask_expanded)
+ x_res = x_res[:, 1:] if self.time_as_token else x_res
+ x_res = x_res[:, 1:] if self.style_as_token else x_res
+ if self.long_skip_connection:
+ x_res = self.skip_linear(torch.cat([x_res, x], dim=-1))
+ if self.final_layer_type == 'wavenet':
+ x = self.conv1(x_res)
+ x = x.transpose(1, 2)
+ t2 = self.t_embedder2(t)
+ x = self.wavenet(x, x_mask, g=t2.unsqueeze(2)).transpose(1, 2) + self.res_projection(
+ x_res) # long residual connection
+ x = self.final_layer(x, t1).transpose(1, 2)
+ x = self.conv2(x)
+ else:
+ x = self.final_mlp(x_res)
+ x = x.transpose(1, 2)
+ return x
\ No newline at end of file
diff --git a/modules/encodec.py b/modules/encodec.py
new file mode 100644
index 0000000000000000000000000000000000000000..bf4ab1ab1629640cc5db4cfd9097d24c3a06f552
--- /dev/null
+++ b/modules/encodec.py
@@ -0,0 +1,292 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Convolutional layers wrappers and utilities."""
+
+import math
+import typing as tp
+import warnings
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.utils import spectral_norm, weight_norm
+
+import typing as tp
+
+import einops
+
+
+class ConvLayerNorm(nn.LayerNorm):
+ """
+ Convolution-friendly LayerNorm that moves channels to last dimensions
+ before running the normalization and moves them back to original position right after.
+ """
+ def __init__(self, normalized_shape: tp.Union[int, tp.List[int], torch.Size], **kwargs):
+ super().__init__(normalized_shape, **kwargs)
+
+ def forward(self, x):
+ x = einops.rearrange(x, 'b ... t -> b t ...')
+ x = super().forward(x)
+ x = einops.rearrange(x, 'b t ... -> b ... t')
+ return
+
+
+CONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm',
+ 'time_layer_norm', 'layer_norm', 'time_group_norm'])
+
+
+def apply_parametrization_norm(module: nn.Module, norm: str = 'none') -> nn.Module:
+ assert norm in CONV_NORMALIZATIONS
+ if norm == 'weight_norm':
+ return weight_norm(module)
+ elif norm == 'spectral_norm':
+ return spectral_norm(module)
+ else:
+ # We already check was in CONV_NORMALIZATION, so any other choice
+ # doesn't need reparametrization.
+ return module
+
+
+def get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs) -> nn.Module:
+ """Return the proper normalization module. If causal is True, this will ensure the returned
+ module is causal, or return an error if the normalization doesn't support causal evaluation.
+ """
+ assert norm in CONV_NORMALIZATIONS
+ if norm == 'layer_norm':
+ assert isinstance(module, nn.modules.conv._ConvNd)
+ return ConvLayerNorm(module.out_channels, **norm_kwargs)
+ elif norm == 'time_group_norm':
+ if causal:
+ raise ValueError("GroupNorm doesn't support causal evaluation.")
+ assert isinstance(module, nn.modules.conv._ConvNd)
+ return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
+ else:
+ return nn.Identity()
+
+
+def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int,
+ padding_total: int = 0) -> int:
+ """See `pad_for_conv1d`.
+ """
+ length = x.shape[-1]
+ n_frames = (length - kernel_size + padding_total) / stride + 1
+ ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
+ return ideal_length - length
+
+
+def pad_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0):
+ """Pad for a convolution to make sure that the last window is full.
+ Extra padding is added at the end. This is required to ensure that we can rebuild
+ an output of the same length, as otherwise, even with padding, some time steps
+ might get removed.
+ For instance, with total padding = 4, kernel size = 4, stride = 2:
+ 0 0 1 2 3 4 5 0 0 # (0s are padding)
+ 1 2 3 # (output frames of a convolution, last 0 is never used)
+ 0 0 1 2 3 4 5 0 # (output of tr. conv., but pos. 5 is going to get removed as padding)
+ 1 2 3 4 # once you removed padding, we are missing one time step !
+ """
+ extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+ return F.pad(x, (0, extra_padding))
+
+
+def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'zero', value: float = 0.):
+ """Tiny wrapper around F.pad, just to allow for reflect padding on small input.
+ If this is the case, we insert extra 0 padding to the right before the reflection happen.
+ """
+ length = x.shape[-1]
+ padding_left, padding_right = paddings
+ assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+ if mode == 'reflect':
+ max_pad = max(padding_left, padding_right)
+ extra_pad = 0
+ if length <= max_pad:
+ extra_pad = max_pad - length + 1
+ x = F.pad(x, (0, extra_pad))
+ padded = F.pad(x, paddings, mode, value)
+ end = padded.shape[-1] - extra_pad
+ return padded[..., :end]
+ else:
+ return F.pad(x, paddings, mode, value)
+
+
+def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
+ """Remove padding from x, handling properly zero padding. Only for 1d!"""
+ padding_left, padding_right = paddings
+ assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+ assert (padding_left + padding_right) <= x.shape[-1]
+ end = x.shape[-1] - padding_right
+ return x[..., padding_left: end]
+
+
+class NormConv1d(nn.Module):
+ """Wrapper around Conv1d and normalization applied to this conv
+ to provide a uniform interface across normalization approaches.
+ """
+ def __init__(self, *args, causal: bool = False, norm: str = 'none',
+ norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+ super().__init__()
+ self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
+ self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
+ self.norm_type = norm
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.norm(x)
+ return x
+
+
+class NormConv2d(nn.Module):
+ """Wrapper around Conv2d and normalization applied to this conv
+ to provide a uniform interface across normalization approaches.
+ """
+ def __init__(self, *args, norm: str = 'none',
+ norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+ super().__init__()
+ self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)
+ self.norm = get_norm_module(self.conv, causal=False, norm=norm, **norm_kwargs)
+ self.norm_type = norm
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.norm(x)
+ return x
+
+
+class NormConvTranspose1d(nn.Module):
+ """Wrapper around ConvTranspose1d and normalization applied to this conv
+ to provide a uniform interface across normalization approaches.
+ """
+ def __init__(self, *args, causal: bool = False, norm: str = 'none',
+ norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+ super().__init__()
+ self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm)
+ self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)
+ self.norm_type = norm
+
+ def forward(self, x):
+ x = self.convtr(x)
+ x = self.norm(x)
+ return x
+
+
+class NormConvTranspose2d(nn.Module):
+ """Wrapper around ConvTranspose2d and normalization applied to this conv
+ to provide a uniform interface across normalization approaches.
+ """
+ def __init__(self, *args, norm: str = 'none',
+ norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+ super().__init__()
+ self.convtr = apply_parametrization_norm(nn.ConvTranspose2d(*args, **kwargs), norm)
+ self.norm = get_norm_module(self.convtr, causal=False, norm=norm, **norm_kwargs)
+
+ def forward(self, x):
+ x = self.convtr(x)
+ x = self.norm(x)
+ return x
+
+
+class SConv1d(nn.Module):
+ """Conv1d with some builtin handling of asymmetric or causal padding
+ and normalization.
+ """
+ def __init__(self, in_channels: int, out_channels: int,
+ kernel_size: int, stride: int = 1, dilation: int = 1,
+ groups: int = 1, bias: bool = True, causal: bool = False,
+ norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {},
+ pad_mode: str = 'reflect', **kwargs):
+ super().__init__()
+ # warn user on unusual setup between dilation and stride
+ if stride > 1 and dilation > 1:
+ warnings.warn('SConv1d has been initialized with stride > 1 and dilation > 1'
+ f' (kernel_size={kernel_size} stride={stride}, dilation={dilation}).')
+ self.conv = NormConv1d(in_channels, out_channels, kernel_size, stride,
+ dilation=dilation, groups=groups, bias=bias, causal=causal,
+ norm=norm, norm_kwargs=norm_kwargs)
+ self.causal = causal
+ self.pad_mode = pad_mode
+
+ def forward(self, x):
+ B, C, T = x.shape
+ kernel_size = self.conv.conv.kernel_size[0]
+ stride = self.conv.conv.stride[0]
+ dilation = self.conv.conv.dilation[0]
+ kernel_size = (kernel_size - 1) * dilation + 1 # effective kernel size with dilations
+ padding_total = kernel_size - stride
+ extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+ if self.causal:
+ # Left padding for causal
+ x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
+ else:
+ # Asymmetric padding required for odd strides
+ padding_right = padding_total // 2
+ padding_left = padding_total - padding_right
+ x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode)
+ return self.conv(x)
+
+
+class SConvTranspose1d(nn.Module):
+ """ConvTranspose1d with some builtin handling of asymmetric or causal padding
+ and normalization.
+ """
+ def __init__(self, in_channels: int, out_channels: int,
+ kernel_size: int, stride: int = 1, causal: bool = False,
+ norm: str = 'none', trim_right_ratio: float = 1.,
+ norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+ super().__init__()
+ self.convtr = NormConvTranspose1d(in_channels, out_channels, kernel_size, stride,
+ causal=causal, norm=norm, norm_kwargs=norm_kwargs)
+ self.causal = causal
+ self.trim_right_ratio = trim_right_ratio
+ assert self.causal or self.trim_right_ratio == 1., \
+ "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
+ assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1.
+
+ def forward(self, x):
+ kernel_size = self.convtr.convtr.kernel_size[0]
+ stride = self.convtr.convtr.stride[0]
+ padding_total = kernel_size - stride
+
+ y = self.convtr(x)
+
+ # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be
+ # removed at the very end, when keeping only the right length for the output,
+ # as removing it here would require also passing the length at the matching layer
+ # in the encoder.
+ if self.causal:
+ # Trim the padding on the right according to the specified ratio
+ # if trim_right_ratio = 1.0, trim everything from right
+ padding_right = math.ceil(padding_total * self.trim_right_ratio)
+ padding_left = padding_total - padding_right
+ y = unpad1d(y, (padding_left, padding_right))
+ else:
+ # Asymmetric padding required for odd strides
+ padding_right = padding_total // 2
+ padding_left = padding_total - padding_right
+ y = unpad1d(y, (padding_left, padding_right))
+ return y
+
+class SLSTM(nn.Module):
+ """
+ LSTM without worrying about the hidden state, nor the layout of the data.
+ Expects input as convolutional layout.
+ """
+ def __init__(self, dimension: int, num_layers: int = 2, skip: bool = True):
+ super().__init__()
+ self.skip = skip
+ self.lstm = nn.LSTM(dimension, dimension, num_layers)
+ self.hidden = None
+
+ def forward(self, x):
+ x = x.permute(2, 0, 1)
+ if self.training:
+ y, _ = self.lstm(x)
+ else:
+ y, self.hidden = self.lstm(x, self.hidden)
+ if self.skip:
+ y = y + x
+ y = y.permute(1, 2, 0)
+ return y
\ No newline at end of file
diff --git a/modules/flow_matching.py b/modules/flow_matching.py
new file mode 100644
index 0000000000000000000000000000000000000000..61389183c6604edf80a9517ead197dd6aa097740
--- /dev/null
+++ b/modules/flow_matching.py
@@ -0,0 +1,167 @@
+from abc import ABC
+
+import torch
+import torch.nn.functional as F
+
+from modules.diffusion_transformer import DiT
+from modules.commons import sequence_mask
+
+from tqdm import tqdm
+
+class BASECFM(torch.nn.Module, ABC):
+ def __init__(
+ self,
+ args,
+ ):
+ super().__init__()
+ self.sigma_min = 1e-6
+
+ self.estimator = None
+
+ self.in_channels = args.DiT.in_channels
+
+ self.criterion = torch.nn.MSELoss() if args.reg_loss_type == "l2" else torch.nn.L1Loss()
+
+ if hasattr(args.DiT, 'zero_prompt_speech_token'):
+ self.zero_prompt_speech_token = args.DiT.zero_prompt_speech_token
+ else:
+ self.zero_prompt_speech_token = False
+
+ @torch.inference_mode()
+ def inference(self, mu, x_lens, prompt, style, f0, n_timesteps, temperature=1.0, inference_cfg_rate=0.5):
+ """Forward diffusion
+
+ Args:
+ mu (torch.Tensor): output of encoder
+ shape: (batch_size, n_feats, mel_timesteps)
+ mask (torch.Tensor): output_mask
+ shape: (batch_size, 1, mel_timesteps)
+ n_timesteps (int): number of diffusion steps
+ temperature (float, optional): temperature for scaling noise. Defaults to 1.0.
+ spks (torch.Tensor, optional): speaker ids. Defaults to None.
+ shape: (batch_size, spk_emb_dim)
+ cond: Not used but kept for future purposes
+
+ Returns:
+ sample: generated mel-spectrogram
+ shape: (batch_size, n_feats, mel_timesteps)
+ """
+ B, T = mu.size(0), mu.size(1)
+ z = torch.randn([B, self.in_channels, T], device=mu.device) * temperature
+ t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device)
+ # t_span = t_span + (-1) * (torch.cos(torch.pi / 2 * t_span) - 1 + t_span)
+ return self.solve_euler(z, x_lens, prompt, mu, style, f0, t_span, inference_cfg_rate)
+
+ def solve_euler(self, x, x_lens, prompt, mu, style, f0, t_span, inference_cfg_rate=0.5):
+ """
+ Fixed euler solver for ODEs.
+ Args:
+ x (torch.Tensor): random noise
+ t_span (torch.Tensor): n_timesteps interpolated
+ shape: (n_timesteps + 1,)
+ mu (torch.Tensor): output of encoder
+ shape: (batch_size, n_feats, mel_timesteps)
+ mask (torch.Tensor): output_mask
+ shape: (batch_size, 1, mel_timesteps)
+ spks (torch.Tensor, optional): speaker ids. Defaults to None.
+ shape: (batch_size, spk_emb_dim)
+ cond: Not used but kept for future purposes
+ """
+ t, _, _ = t_span[0], t_span[-1], t_span[1] - t_span[0]
+
+ # I am storing this because I can later plot it by putting a debugger here and saving it to a file
+ # Or in future might add like a return_all_steps flag
+ sol = []
+ # apply prompt
+ prompt_len = prompt.size(-1)
+ prompt_x = torch.zeros_like(x)
+ prompt_x[..., :prompt_len] = prompt[..., :prompt_len]
+ x[..., :prompt_len] = 0
+ if self.zero_prompt_speech_token:
+ mu[..., :prompt_len] = 0
+ for step in tqdm(range(1, len(t_span))):
+ dt = t_span[step] - t_span[step - 1]
+ if inference_cfg_rate > 0:
+ # Stack original and CFG (null) inputs for batched processing
+ stacked_prompt_x = torch.cat([prompt_x, torch.zeros_like(prompt_x)], dim=0)
+ stacked_style = torch.cat([style, torch.zeros_like(style)], dim=0)
+ stacked_mu = torch.cat([mu, torch.zeros_like(mu)], dim=0)
+ stacked_x = torch.cat([x, x], dim=0)
+ stacked_t = torch.cat([t.unsqueeze(0), t.unsqueeze(0)], dim=0)
+
+ # Perform a single forward pass for both original and CFG inputs
+ stacked_dphi_dt = self.estimator(
+ stacked_x, stacked_prompt_x, x_lens, stacked_t, stacked_style, stacked_mu,
+ )
+
+ # Split the output back into the original and CFG components
+ dphi_dt, cfg_dphi_dt = stacked_dphi_dt.chunk(2, dim=0)
+
+ # Apply CFG formula
+ dphi_dt = (1.0 + inference_cfg_rate) * dphi_dt - inference_cfg_rate * cfg_dphi_dt
+ else:
+ dphi_dt = self.estimator(x, prompt_x, x_lens, t.unsqueeze(0), style, mu)
+
+ x = x + dt * dphi_dt
+ t = t + dt
+ sol.append(x)
+ if step < len(t_span) - 1:
+ dt = t_span[step + 1] - t
+ x[:, :, :prompt_len] = 0
+
+ return sol[-1]
+ def forward(self, x1, x_lens, prompt_lens, mu, style):
+ """Computes diffusion loss
+
+ Args:
+ x1 (torch.Tensor): Target
+ shape: (batch_size, n_feats, mel_timesteps)
+ mask (torch.Tensor): target mask
+ shape: (batch_size, 1, mel_timesteps)
+ mu (torch.Tensor): output of encoder
+ shape: (batch_size, n_feats, mel_timesteps)
+ spks (torch.Tensor, optional): speaker embedding. Defaults to None.
+ shape: (batch_size, spk_emb_dim)
+
+ Returns:
+ loss: conditional flow matching loss
+ y: conditional flow
+ shape: (batch_size, n_feats, mel_timesteps)
+ """
+ b, _, t = x1.shape
+
+ # random timestep
+ t = torch.rand([b, 1, 1], device=mu.device, dtype=x1.dtype)
+ # sample noise p(x_0)
+ z = torch.randn_like(x1)
+
+ y = (1 - (1 - self.sigma_min) * t) * z + t * x1
+ u = x1 - (1 - self.sigma_min) * z
+
+ prompt = torch.zeros_like(x1)
+ for bib in range(b):
+ prompt[bib, :, :prompt_lens[bib]] = x1[bib, :, :prompt_lens[bib]]
+ # range covered by prompt are set to 0
+ y[bib, :, :prompt_lens[bib]] = 0
+ if self.zero_prompt_speech_token:
+ mu[bib, :, :prompt_lens[bib]] = 0
+
+ estimator_out = self.estimator(y, prompt, x_lens, t.squeeze(1).squeeze(1), style, mu, prompt_lens)
+ loss = 0
+ for bib in range(b):
+ loss += self.criterion(estimator_out[bib, :, prompt_lens[bib]:x_lens[bib]], u[bib, :, prompt_lens[bib]:x_lens[bib]])
+ loss /= b
+
+ return loss, estimator_out + (1 - self.sigma_min) * z
+
+
+
+class CFM(BASECFM):
+ def __init__(self, args):
+ super().__init__(
+ args
+ )
+ if args.dit_type == "DiT":
+ self.estimator = DiT(args)
+ else:
+ raise NotImplementedError(f"Unknown diffusion type {args.dit_type}")
diff --git a/modules/gpt_fast/__pycache__/model.cpython-310.pyc b/modules/gpt_fast/__pycache__/model.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..32c5caea63627dd61f2e21bcd088bf10621307ae
Binary files /dev/null and b/modules/gpt_fast/__pycache__/model.cpython-310.pyc differ
diff --git a/modules/gpt_fast/generate.py b/modules/gpt_fast/generate.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f09fb28540826ca1151e6bc0af3da304f52b858
--- /dev/null
+++ b/modules/gpt_fast/generate.py
@@ -0,0 +1,436 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+import itertools
+import sys
+import time
+from pathlib import Path
+from typing import Optional, Tuple
+
+import torch
+import torch._dynamo.config
+import torch._inductor.config
+
+def device_sync(device):
+ if "cuda" in device:
+ torch.cuda.synchronize(device)
+ elif ("cpu" in device) or ("mps" in device):
+ pass
+ else:
+ print(f"device={device} is not yet suppported")
+
+
+torch._inductor.config.coordinate_descent_tuning = True
+torch._inductor.config.triton.unique_kernel_names = True
+torch._inductor.config.fx_graph_cache = True # Experimental feature to reduce compilation times, will be on by default in future
+
+default_device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# support running without installing as a package
+wd = Path(__file__).parent.parent.resolve()
+sys.path.append(str(wd))
+
+from model import Transformer
+from tokenizer import get_tokenizer
+
+def multinomial_sample_one_no_sync(probs_sort): # Does multinomial sampling without a cuda synchronization
+ q = torch.empty_like(probs_sort).exponential_(1)
+ return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.int)
+
+def logits_to_probs(logits, temperature: float = 1.0, top_k: Optional[int] = None):
+ logits = logits / max(temperature, 1e-5)
+
+ if top_k is not None:
+ v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+ pivot = v.select(-1, -1).unsqueeze(-1)
+ logits = torch.where(logits < pivot, -float("Inf"), logits)
+ probs = torch.nn.functional.softmax(logits, dim=-1)
+ return probs
+
+def sample(logits, temperature: float = 1.0, top_k: Optional[int] = None):
+ probs = logits_to_probs(logits[0, -1], temperature, top_k)
+ idx_next = multinomial_sample_one_no_sync(probs)
+ return idx_next, probs
+
+def prefill(model: Transformer, x: torch.Tensor, input_pos: torch.Tensor, **sampling_kwargs) -> torch.Tensor:
+ # input_pos: [B, S]
+ logits = model(x, input_pos)
+ return sample(logits, **sampling_kwargs)[0]
+
+def decode_one_token(model: Transformer, x: torch.Tensor, input_pos: torch.Tensor, **sampling_kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
+ # input_pos: [B, 1]
+ assert input_pos.shape[-1] == 1
+ logits = model(x, input_pos)
+ return sample(logits, **sampling_kwargs)
+
+def decode_n_tokens(model: Transformer, cur_token: torch.Tensor, input_pos: torch.Tensor, num_new_tokens: int, callback=lambda _: _, **sampling_kwargs):
+ new_tokens, new_probs = [], []
+ for i in range(num_new_tokens):
+ with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True): # Actually better for Inductor to codegen attention here
+ next_token, next_prob = decode_one_token(
+ model, cur_token, input_pos, **sampling_kwargs
+ )
+ input_pos += 1
+ new_tokens.append(next_token.clone())
+ callback(new_tokens[-1])
+ new_probs.append(next_prob.clone())
+ cur_token = next_token.view(1, -1)
+
+ return new_tokens, new_probs
+
+
+def model_forward(model, x, input_pos):
+ return model(x, input_pos)
+
+def speculative_decode(
+ model: Transformer,
+ draft_model: Transformer,
+ cur_token: torch.Tensor,
+ input_pos: int,
+ speculate_k: int,
+ **sampling_kwargs
+) -> torch.Tensor:
+ # draft model inference sequentially
+ device = cur_token.device
+ orig_input_pos = torch.tensor([input_pos], dtype=torch.int64, device=cur_token.device)
+ draft_tokens, draft_probs = decode_n_tokens(draft_model, cur_token.view(1, -1), orig_input_pos.clone(), speculate_k, **sampling_kwargs)
+
+ draft_tokens = torch.cat(draft_tokens)
+ # parallel inference on target model using draft tokens
+ target_logits = model_forward(
+ model,
+ torch.cat([cur_token.view(1), draft_tokens]).view(1, -1),
+ torch.arange(input_pos, input_pos + speculate_k + 1, device=cur_token.device)
+ )
+ target_probs = logits_to_probs(target_logits[0], **sampling_kwargs)
+ draft_probs = torch.stack(draft_probs)
+ # q: target prob, p: draft prob
+ # q >= p: always accept draft token
+ # q < p: q/p prob to accept draft token
+ p = draft_probs[torch.arange(0, speculate_k, device=device), draft_tokens]
+ q = target_probs[torch.arange(0, speculate_k, device=device), draft_tokens]
+ accept_draft_prob = torch.minimum(torch.ones(()), q[:speculate_k]/ p)
+ rejected_locations = (torch.rand_like(accept_draft_prob) > accept_draft_prob).nonzero()
+
+ if rejected_locations.shape[0] == 0: # All draft tokens have been accepted
+ accept_length = speculate_k + 1
+ last_token = multinomial_sample_one_no_sync(target_probs[-1])
+ # fill last token into draft model
+ model_forward(
+ draft_model,
+ draft_tokens[-1].view(1, -1),
+ orig_input_pos + speculate_k,
+ )
+ return torch.cat([draft_tokens, last_token])
+ else:
+ accept_length = rejected_locations[0].item()
+ p = draft_probs[accept_length]
+ q = target_probs[accept_length]
+ new = q - p
+ new = torch.where(new > 0, new, 0.0)
+ new = new / new.sum()
+ next_token = multinomial_sample_one_no_sync(new)
+ return torch.cat([draft_tokens[:accept_length], next_token])
+
+@torch.no_grad()
+def generate(
+ model: Transformer,
+ prompt: torch.Tensor,
+ max_new_tokens: int,
+ *,
+ interactive: bool,
+ draft_model: Transformer,
+ speculate_k: Optional[int] = 8,
+ callback = lambda x: x,
+ **sampling_kwargs
+) -> torch.Tensor:
+ """
+ Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.
+ """
+
+ is_speculative = draft_model is not None
+ # create an empty tensor of the expected final shape and fill in the current tokens
+ T = prompt.size(0)
+ T_new = T + max_new_tokens
+ if interactive:
+ max_seq_length = 350
+ else:
+ max_seq_length = min(T_new, model.config.block_size)
+
+ device, dtype = prompt.device, prompt.dtype
+ max_seq_length = max_seq_length + speculate_k + 1 if is_speculative else max_seq_length
+ with torch.device(device):
+ model.setup_caches(max_batch_size=1, max_seq_length=max_seq_length)
+ if is_speculative and draft_model is not model:
+ draft_model.setup_caches(max_batch_size=1, max_seq_length=max_seq_length)
+
+ # create an empty tensor of the expected final shape and fill in the current tokens
+ empty = torch.empty(T_new, dtype=dtype, device=device)
+ empty[:T] = prompt
+ seq = empty
+ input_pos = torch.arange(0, T, device=device)
+
+ next_token = prefill(model, prompt.view(1, -1), input_pos, **sampling_kwargs).clone()
+ if is_speculative:
+ prefill(draft_model, prompt.view(1, -1), input_pos, **sampling_kwargs)
+ seq[T] = next_token
+
+ input_pos = torch.tensor([T], device=device, dtype=torch.int)
+ accept_counts = [0] * (speculate_k + 1)
+
+ if is_speculative:
+ input_pos = input_pos.item() # for speculative decoding easier to keep on host
+ while input_pos < T_new - 1:
+ cur_token = next_token.view(())
+
+ next_tokens = speculative_decode(
+ model, draft_model, cur_token, input_pos, speculate_k, **sampling_kwargs
+ )
+
+ accept_counts[len(next_tokens) - 1] += 1
+ num_added = min(T_new - input_pos - 1, len(next_tokens))
+ seq[input_pos + 1 : input_pos + num_added + 1] = next_tokens[: num_added]
+ for i in next_tokens[: num_added,]:
+ callback(i)
+ input_pos = input_pos + num_added
+ next_token = next_tokens[-1]
+ else:
+ generated_tokens, _ = decode_n_tokens(model, next_token.view(1, -1), input_pos, max_new_tokens - 1, callback=callback, **sampling_kwargs)
+ seq[T + 1:] = torch.cat(generated_tokens)
+
+ generate_stats = {
+ 'accept_counts': accept_counts
+ }
+ return seq, generate_stats
+
+def encode_tokens(tokenizer, string, bos=True, device=default_device):
+ tokens = tokenizer.encode(string)
+ if bos:
+ tokens = [tokenizer.bos_id()] + tokens
+ return torch.tensor(tokens, dtype=torch.int, device=device)
+
+def _load_model(checkpoint_path, device, precision, use_tp):
+ use_cuda = 'cuda' in device
+ with torch.device('meta'):
+ model = Transformer.from_name(checkpoint_path.parent.name)
+
+ if "int8" in str(checkpoint_path):
+ print("Using int8 weight-only quantization!")
+ from quantize import WeightOnlyInt8QuantHandler
+ simple_quantizer = WeightOnlyInt8QuantHandler(model)
+ model = simple_quantizer.convert_for_runtime()
+
+ if "int4" in str(checkpoint_path):
+ print("Using int4 weight-only quantization!")
+ path_comps = checkpoint_path.name.split(".")
+ groupsize = int(path_comps[-2][1:])
+ from quantize import WeightOnlyInt4QuantHandler
+ simple_quantizer = WeightOnlyInt4QuantHandler(model, groupsize)
+ model = simple_quantizer.convert_for_runtime()
+
+ checkpoint = torch.load(str(checkpoint_path), mmap=True, weights_only=True)
+ if "model" in checkpoint and "stories" in str(checkpoint_path):
+ checkpoint = checkpoint["model"]
+ model.load_state_dict(checkpoint, assign=True)
+
+ if use_tp:
+ from tp import apply_tp
+ print("Applying tensor parallel to model ...")
+ apply_tp(model)
+
+ model = model.to(device=device, dtype=precision)
+ return model.eval()
+
+def _get_model_size(model):
+ model_size = 0
+ for name, child in model.named_children():
+ if not isinstance(child, torch.nn.Embedding):
+ model_size += sum(
+ [
+ p.numel() * p.dtype.itemsize
+ for p in itertools.chain(child.parameters(), child.buffers())
+ ]
+ )
+ return model_size
+
+B_INST, E_INST = "[INST]", "[/INST]"
+
+def main(
+ prompt: str = "Hello, my name is",
+ interactive: bool = False,
+ num_samples: int = 5,
+ max_new_tokens: int = 100,
+ top_k: int = 200,
+ temperature: float = 0.8,
+ checkpoint_path: Path = Path("checkpoints/meta-Transformer/Transformer-2-7b-chat-hf/model.pth"),
+ compile: bool = True,
+ compile_prefill: bool = False,
+ profile: Optional[Path] = None,
+ draft_checkpoint_path: Optional[Path] = None,
+ speculate_k: int = 5,
+ device=default_device,
+) -> None:
+ """Generates text samples based on a pre-trained Transformer model and tokenizer.
+ """
+ assert checkpoint_path.is_file(), checkpoint_path
+
+ tokenizer_path = checkpoint_path.parent / "tokenizer.model"
+ assert tokenizer_path.is_file(), str(tokenizer_path)
+
+ global print
+ from tp import maybe_init_dist
+ rank = maybe_init_dist()
+ use_tp = rank is not None
+ if use_tp:
+ if rank != 0:
+ # only print on rank 0
+ print = lambda *args, **kwargs: None
+
+ print(f"Using device={device}")
+ precision = torch.bfloat16
+ is_speculative = draft_checkpoint_path is not None
+ is_chat = "chat" in str(checkpoint_path)
+
+ print("Loading model ...")
+ t0 = time.time()
+ model = _load_model(checkpoint_path, device, precision, use_tp)
+
+ if is_speculative:
+ draft_model = _load_model(draft_checkpoint_path, device, precision, use_tp)
+ else:
+ draft_model = None
+
+ device_sync(device=device) # MKG
+ print(f"Time to load model: {time.time() - t0:.02f} seconds")
+
+ tokenizer = get_tokenizer(tokenizer_path, checkpoint_path)
+
+ encoded = encode_tokens(tokenizer, prompt, bos=True, device=device)
+ prompt_length = encoded.size(0)
+
+ torch.manual_seed(1234)
+ model_size = _get_model_size(model)
+ if compile:
+ if is_speculative and use_tp: # and ("cuda" in device):
+ torch._inductor.config.triton.cudagraph_trees = False # Bug with cudagraph trees in this case
+
+ if is_speculative:
+ global model_forward, logits_to_prob
+ model_forward = torch.compile(model_forward, mode="reduce-overhead", fullgraph=True)
+
+ global decode_one_token, prefill
+ decode_one_token = torch.compile(decode_one_token, mode="reduce-overhead", fullgraph=True)
+
+ # Uncomment to squeeze more perf out of prefill
+ if compile_prefill:
+ prefill = torch.compile(prefill, fullgraph=True, dynamic=True)
+
+
+ aggregate_metrics = {
+ 'tokens_per_sec': [],
+ 'accept_counts': [],
+ }
+ start = -1 if compile else 0
+
+ for i in range(start, num_samples):
+ device_sync(device=device) # MKG
+ if i >= 0 and interactive:
+ prompt = input("What is your prompt? ")
+ if is_chat:
+ prompt = f"{B_INST} {prompt.strip()} {E_INST}"
+ encoded = encode_tokens(tokenizer, prompt, bos=True, device=device)
+
+ if interactive and i >= 0:
+ buffer = []
+ period_id = tokenizer.encode('.')[0]
+ done_generating = False
+ def callback(x):
+ nonlocal done_generating
+ if done_generating:
+ return
+ buffer.append(tokenizer.decode([period_id] + x.tolist())[1:])
+ if x.item() == tokenizer.eos_id():
+ done_generating = True
+ if len(buffer) == 4 or done_generating:
+ print(''.join(buffer), end='', flush=True)
+ buffer.clear()
+ # print(, end='', flush=True)
+ else:
+ callback = lambda x : x
+ t0 = time.perf_counter()
+ import contextlib
+ if (i != num_samples - 1 or not profile) or (use_tp and rank != 0):
+ prof = contextlib.nullcontext()
+ else:
+ torch.profiler._utils._init_for_cuda_graphs()
+ prof = torch.profiler.profile()
+ with prof:
+ y, metrics = generate(
+ model,
+ encoded,
+ max_new_tokens,
+ draft_model=draft_model,
+ speculate_k=speculate_k,
+ interactive=interactive,
+ callback=callback,
+ temperature=temperature,
+ top_k=top_k,
+ )
+ aggregate_metrics['accept_counts'].append(metrics['accept_counts'])
+ if i == -1:
+ print(f"Compilation time: {time.perf_counter() - t0:.2f} seconds")
+ continue
+ if hasattr(prof, "export_chrome_trace"):
+ if use_tp:
+ prof.export_chrome_trace(f"{profile}_rank_{rank}.json")
+ else:
+ prof.export_chrome_trace(f"{profile}.json")
+ device_sync(device=device) # MKG
+ t = time.perf_counter() - t0
+
+ if not interactive:
+ print(tokenizer.decode(y.tolist()))
+ else:
+ print()
+ tokens_generated = y.size(0) - prompt_length
+ tokens_sec = tokens_generated / t
+ aggregate_metrics['tokens_per_sec'].append(tokens_sec)
+ print(f"Time for inference {i + 1}: {t:.02f} sec total, {tokens_sec:.02f} tokens/sec")
+ print(f"Bandwidth achieved: {model_size * tokens_sec / 1e9:.02f} GB/s")
+ print("==========")
+ if is_speculative:
+ counts_aggregated = [sum(i) for i in zip(*aggregate_metrics['accept_counts'])]
+ acceptance_probs = [i/sum(counts_aggregated) for i in counts_aggregated]
+ print(f"Acceptance probs: {acceptance_probs}")
+ print(f"Mean Accepted: {sum([idx * i for idx, i in enumerate(counts_aggregated)])/sum(counts_aggregated)}")
+
+ print(f"Average tokens/sec: {torch.mean(torch.tensor(aggregate_metrics['tokens_per_sec'])).item():.2f}")
+ print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")
+
+
+if __name__ == '__main__':
+ import argparse
+ parser = argparse.ArgumentParser(description='Your CLI description.')
+
+ parser.add_argument('--prompt', type=str, default="Hello, my name is", help='Input prompt.')
+ parser.add_argument('--interactive', action='store_true', help='Whether to launch in interactive mode')
+ parser.add_argument('--num_samples', type=int, default=5, help='Number of samples.')
+ parser.add_argument('--max_new_tokens', type=int, default=200, help='Maximum number of new tokens.')
+ parser.add_argument('--top_k', type=int, default=200, help='Top-k for sampling.')
+ parser.add_argument('--temperature', type=float, default=0.8, help='Temperature for sampling.')
+ parser.add_argument('--checkpoint_path', type=Path, default=Path("checkpoints/meta-Transformer/Transformer-2-7b-chat-hf/model.pth"), help='Model checkpoint path.')
+ parser.add_argument('--compile', action='store_true', help='Whether to compile the model.')
+ parser.add_argument('--compile_prefill', action='store_true', help='Whether to compile the prefill (improves prefill perf, but higher compile times)')
+ parser.add_argument('--profile', type=Path, default=None, help='Profile path.')
+ parser.add_argument('--speculate_k', type=int, default=5, help='Speculative execution depth.')
+ parser.add_argument('--draft_checkpoint_path', type=Path, default=None, help='Draft checkpoint path.')
+ parser.add_argument('--device', type=str, default=default_device, help='Device to use')
+
+ args = parser.parse_args()
+ main(
+ args.prompt, args.interactive, args.num_samples, args.max_new_tokens, args.top_k,
+ args.temperature, args.checkpoint_path, args.compile, args.compile_prefill, args.profile, args.draft_checkpoint_path,
+ args.speculate_k, args.device
+ )
diff --git a/modules/gpt_fast/model.py b/modules/gpt_fast/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..c2a08d4dc48ab8ac9bc1790673da966abfc10b91
--- /dev/null
+++ b/modules/gpt_fast/model.py
@@ -0,0 +1,356 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import functional as F
+
+
+def find_multiple(n: int, k: int) -> int:
+ if n % k == 0:
+ return n
+ return n + k - (n % k)
+
+class AdaptiveLayerNorm(nn.Module):
+ r"""Adaptive Layer Normalization"""
+
+ def __init__(self, d_model, norm) -> None:
+ super(AdaptiveLayerNorm, self).__init__()
+ self.project_layer = nn.Linear(d_model, 2 * d_model)
+ self.norm = norm
+ self.d_model = d_model
+ self.eps = self.norm.eps
+
+ def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor:
+ if embedding is None:
+ return self.norm(input)
+ weight, bias = torch.split(
+ self.project_layer(embedding),
+ split_size_or_sections=self.d_model,
+ dim=-1,
+ )
+ return weight * self.norm(input) + bias
+
+
+@dataclass
+class ModelArgs:
+ block_size: int = 2048
+ vocab_size: int = 32000
+ n_layer: int = 32
+ n_head: int = 32
+ dim: int = 4096
+ intermediate_size: int = None
+ n_local_heads: int = -1
+ head_dim: int = 64
+ rope_base: float = 10000
+ norm_eps: float = 1e-5
+ has_cross_attention: bool = False
+ context_dim: int = 0
+ uvit_skip_connection: bool = False
+
+ def __post_init__(self):
+ if self.n_local_heads == -1:
+ self.n_local_heads = self.n_head
+ if self.intermediate_size is None:
+ hidden_dim = 4 * self.dim
+ n_hidden = int(2 * hidden_dim / 3)
+ self.intermediate_size = find_multiple(n_hidden, 256)
+ # self.head_dim = self.dim // self.n_head
+
+ @classmethod
+ def from_name(cls, name: str):
+ if name in transformer_configs:
+ return cls(**transformer_configs[name])
+ # fuzzy search
+ config = [config for config in transformer_configs if config.lower() in str(name).lower()]
+
+ # We may have two or more configs matched (e.g. "7B" and "Mistral-7B"). Find the best config match,
+ # take longer name (as it have more symbols matched)
+ if len(config) > 1:
+ config.sort(key=len, reverse=True)
+ assert len(config[0]) != len(config[1]), name # make sure only one 'best' match
+
+ return cls(**transformer_configs[config[0]])
+
+
+transformer_configs = {
+ "CodeLlama-7b-Python-hf": dict(block_size=16384, vocab_size=32000, n_layer=32, dim=4096, rope_base=1000000),
+ "7B": dict(n_layer=32, n_head=32, dim=4096),
+ "13B": dict(n_layer=40, n_head=40, dim=5120),
+ "30B": dict(n_layer=60, n_head=52, dim=6656),
+ "34B": dict(n_layer=48, n_head=64, dim=8192, vocab_size=32000, n_local_heads=8, intermediate_size=22016,
+ rope_base=1000000), # CodeLlama-34B-Python-hf
+ "70B": dict(n_layer=80, n_head=64, dim=8192, n_local_heads=8, intermediate_size=28672),
+ "Mistral-7B": dict(n_layer=32, n_head=32, n_local_heads=8, dim=4096, intermediate_size=14336, vocab_size=32000),
+ "stories15M": dict(n_layer=6, n_head=6, dim=288),
+ "stories110M": dict(n_layer=12, n_head=12, dim=768),
+
+ "llama-3-8b": dict(block_size=8192, n_layer=32, n_head=32, n_local_heads=8, dim=4096, intermediate_size=14336,
+ vocab_size=128256, rope_base=500000),
+ "llama-3-70b": dict(block_size=8192, n_layer=80, n_head=64, n_local_heads=8, dim=8192, intermediate_size=28672,
+ vocab_size=128256, rope_base=500000),
+}
+
+
+class KVCache(nn.Module):
+ def __init__(self, max_batch_size, max_seq_length, n_heads, head_dim, dtype=torch.bfloat16):
+ super().__init__()
+ cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim)
+ self.register_buffer('k_cache', torch.zeros(cache_shape, dtype=dtype))
+ self.register_buffer('v_cache', torch.zeros(cache_shape, dtype=dtype))
+
+ def update(self, input_pos, k_val, v_val):
+ # input_pos: [S], k_val: [B, H, S, D]
+ assert input_pos.shape[0] == k_val.shape[2]
+
+ k_out = self.k_cache
+ v_out = self.v_cache
+ k_out[:, :, input_pos] = k_val
+ v_out[:, :, input_pos] = v_val
+
+ return k_out, v_out
+
+
+class Transformer(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.config = config
+
+ self.layers = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer))
+ self.norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+
+ self.freqs_cis: Optional[Tensor] = None
+ self.mask_cache: Optional[Tensor] = None
+ self.max_batch_size = -1
+ self.max_seq_length = -1
+
+ def setup_caches(self, max_batch_size, max_seq_length, use_kv_cache=True):
+ if self.max_seq_length >= max_seq_length and self.max_batch_size >= max_batch_size:
+ return
+ head_dim = self.config.dim // self.config.n_head
+ max_seq_length = find_multiple(max_seq_length, 8)
+ self.max_seq_length = max_seq_length
+ self.max_batch_size = max_batch_size
+ dtype = self.norm.project_layer.weight.dtype
+ device = self.norm.project_layer.weight.device
+
+ if not self.training and use_kv_cache:
+ for b in self.layers:
+ b.attention.kv_cache = KVCache(max_batch_size, max_seq_length, self.config.n_local_heads, head_dim, dtype).to(device)
+
+ self.freqs_cis = precompute_freqs_cis(self.config.block_size, self.config.head_dim,
+ self.config.rope_base, dtype).to(device)
+ self.causal_mask = torch.tril(torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool)).to(device)
+ self.use_kv_cache = use_kv_cache
+ self.uvit_skip_connection = self.config.uvit_skip_connection
+ if self.uvit_skip_connection:
+ self.layers_emit_skip = [i for i in range(self.config.n_layer) if i < self.config.n_layer // 2]
+ self.layers_receive_skip = [i for i in range(self.config.n_layer) if i > self.config.n_layer // 2]
+ else:
+ self.layers_emit_skip = []
+ self.layers_receive_skip = []
+
+ def forward(self,
+ x: Tensor,
+ c: Tensor,
+ input_pos: Optional[Tensor] = None,
+ mask: Optional[Tensor] = None,
+ context: Optional[Tensor] = None,
+ context_input_pos: Optional[Tensor] = None,
+ cross_attention_mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ assert self.freqs_cis is not None, "Caches must be initialized first"
+ if mask is None: # in case of non-causal model
+ if not self.training and self.use_kv_cache:
+ mask = self.causal_mask[None, None, input_pos]
+ else:
+ mask = self.causal_mask[None, None, input_pos]
+ mask = mask[..., input_pos]
+ freqs_cis = self.freqs_cis[input_pos]
+ if context is not None:
+ context_freqs_cis = self.freqs_cis[context_input_pos]
+ else:
+ context_freqs_cis = None
+ skip_in_x_list = []
+ for i, layer in enumerate(self.layers):
+ if self.uvit_skip_connection and i in self.layers_receive_skip:
+ skip_in_x = skip_in_x_list.pop(-1)
+ else:
+ skip_in_x = None
+ x = layer(x, c, input_pos, freqs_cis, mask, context, context_freqs_cis, cross_attention_mask, skip_in_x)
+ if self.uvit_skip_connection and i in self.layers_emit_skip:
+ skip_in_x_list.append(x)
+ x = self.norm(x, c)
+ return x
+
+ @classmethod
+ def from_name(cls, name: str):
+ return cls(ModelArgs.from_name(name))
+
+
+class TransformerBlock(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.attention = Attention(config)
+ self.feed_forward = FeedForward(config)
+ self.ffn_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+ self.attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+
+ if config.has_cross_attention:
+ self.has_cross_attention = True
+ self.cross_attention = Attention(config, is_cross_attention=True)
+ self.cross_attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+ else:
+ self.has_cross_attention = False
+
+ if config.uvit_skip_connection:
+ self.skip_in_linear = nn.Linear(config.dim * 2, config.dim)
+ self.uvit_skip_connection = True
+ else:
+ self.uvit_skip_connection = False
+
+ def forward(self,
+ x: Tensor,
+ c: Tensor,
+ input_pos: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ context: Optional[Tensor] = None,
+ context_freqs_cis: Optional[Tensor] = None,
+ cross_attention_mask: Optional[Tensor] = None,
+ skip_in_x: Optional[Tensor] = None,
+ ) -> Tensor:
+ if self.uvit_skip_connection and skip_in_x is not None:
+ x = self.skip_in_linear(torch.cat([x, skip_in_x], dim=-1))
+ h = x + self.attention(self.attention_norm(x, c), freqs_cis, mask, input_pos)
+ if self.has_cross_attention:
+ h = h + self.cross_attention(self.cross_attention_norm(h, c), freqs_cis, cross_attention_mask, input_pos, context, context_freqs_cis)
+ out = h + self.feed_forward(self.ffn_norm(h, c))
+ return out
+
+
+class Attention(nn.Module):
+ def __init__(self, config: ModelArgs, is_cross_attention: bool = False):
+ super().__init__()
+ assert config.dim % config.n_head == 0
+
+ total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim
+ # key, query, value projections for all heads, but in a batch
+ if is_cross_attention:
+ self.wq = nn.Linear(config.dim, config.n_head * config.head_dim, bias=False)
+ self.wkv = nn.Linear(config.context_dim, 2 * config.n_local_heads * config.head_dim, bias=False)
+ else:
+ self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False)
+ self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False)
+ self.kv_cache = None
+
+ self.n_head = config.n_head
+ self.head_dim = config.head_dim
+ self.n_local_heads = config.n_local_heads
+ self.dim = config.dim
+ # self._register_load_state_dict_pre_hook(self.load_hook)
+
+ # def load_hook(self, state_dict, prefix, *args):
+ # if prefix + "wq.weight" in state_dict:
+ # wq = state_dict.pop(prefix + "wq.weight")
+ # wk = state_dict.pop(prefix + "wk.weight")
+ # wv = state_dict.pop(prefix + "wv.weight")
+ # state_dict[prefix + "wqkv.weight"] = torch.cat([wq, wk, wv])
+
+ def forward(self,
+ x: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ input_pos: Optional[Tensor] = None,
+ context: Optional[Tensor] = None,
+ context_freqs_cis: Optional[Tensor] = None,
+ ) -> Tensor:
+ bsz, seqlen, _ = x.shape
+
+ kv_size = self.n_local_heads * self.head_dim
+ if context is None:
+ q, k, v = self.wqkv(x).split([kv_size, kv_size, kv_size], dim=-1)
+ context_seqlen = seqlen
+ else:
+ q = self.wq(x)
+ k, v = self.wkv(context).split([kv_size, kv_size], dim=-1)
+ context_seqlen = context.shape[1]
+
+ q = q.view(bsz, seqlen, self.n_head, self.head_dim)
+ k = k.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+ v = v.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+
+ q = apply_rotary_emb(q, freqs_cis)
+ k = apply_rotary_emb(k, context_freqs_cis if context_freqs_cis is not None else freqs_cis)
+
+ q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
+
+ if self.kv_cache is not None:
+ k, v = self.kv_cache.update(input_pos, k, v)
+
+ k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
+
+ y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.head_dim * self.n_head)
+
+ y = self.wo(y)
+ return y
+
+
+class FeedForward(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
+
+ def forward(self, x: Tensor) -> Tensor:
+ return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-5):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def _norm(self, x):
+ return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+ def forward(self, x: Tensor) -> Tensor:
+ output = self._norm(x.float()).type_as(x)
+ return output * self.weight
+
+
+def precompute_freqs_cis(
+ seq_len: int, n_elem: int, base: int = 10000,
+ dtype: torch.dtype = torch.bfloat16
+) -> Tensor:
+ freqs = 1.0 / (base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem))
+ t = torch.arange(seq_len, device=freqs.device)
+ freqs = torch.outer(t, freqs)
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+ cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
+ return cache.to(dtype=dtype)
+
+
+def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
+ xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+ freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2)
+ x_out2 = torch.stack(
+ [
+ xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],
+ xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],
+ ],
+ -1,
+ )
+
+ x_out2 = x_out2.flatten(3)
+ return x_out2.type_as(x)
diff --git a/modules/gpt_fast/quantize.py b/modules/gpt_fast/quantize.py
new file mode 100644
index 0000000000000000000000000000000000000000..431710aa4384f395a52b9920b057c7d04f2147c5
--- /dev/null
+++ b/modules/gpt_fast/quantize.py
@@ -0,0 +1,622 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+import time
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tokenizer import get_tokenizer
+
+try:
+ from GPTQ import GenericGPTQRunner, InputRecorder
+ from eval import get_task_dict, evaluate, lm_eval
+except:
+ pass
+
+from model import Transformer
+
+##### Quantization Primitives ######
+
+def dynamically_quantize_per_channel(x, quant_min, quant_max, target_dtype):
+ # assumes symmetric quantization
+ # assumes axis == 0
+ # assumes dense memory format
+ # TODO(future): relax ^ as needed
+
+ # default setup for affine quantization of activations
+ eps = torch.finfo(torch.float32).eps
+
+ # get min and max
+ min_val, max_val = torch.aminmax(x, dim=1)
+
+ # calculate scales and zero_points based on min and max
+ # reference: https://fburl.com/code/srbiybme
+ min_val_neg = torch.min(min_val, torch.zeros_like(min_val))
+ max_val_pos = torch.max(max_val, torch.zeros_like(max_val))
+ device = min_val_neg.device
+
+ # reference: https://fburl.com/code/4wll53rk
+ max_val_pos = torch.max(-min_val_neg, max_val_pos)
+ scales = max_val_pos / (float(quant_max - quant_min) / 2)
+ # ensure scales is the same dtype as the original tensor
+ scales = torch.clamp(scales, min=eps).to(x.dtype)
+ zero_points = torch.zeros(min_val_neg.size(), dtype=torch.int64, device=device)
+
+ # quantize based on qmin/qmax/scales/zp
+ # reference: https://www.internalfb.com/code/fbsource/[8edc275012b1]/fbcode/caffe2/torch/ao/quantization/fx/_decomposed.py?lines=63
+ x_div = x / scales.unsqueeze(-1)
+ x_round = torch.round(x_div)
+ x_zp = x_round + zero_points.unsqueeze(-1)
+ quant = torch.clamp(x_zp, quant_min, quant_max).to(target_dtype)
+
+ return quant, scales, zero_points
+
+def get_group_qparams(w, n_bit=4, groupsize=128):
+ # needed for GPTQ with padding
+ if groupsize > w.shape[-1]:
+ groupsize = w.shape[-1]
+ assert groupsize > 1
+ assert w.shape[-1] % groupsize == 0
+ assert w.dim() == 2
+
+ to_quant = w.reshape(-1, groupsize)
+ assert torch.isnan(to_quant).sum() == 0
+
+ max_val = to_quant.amax(dim=1, keepdim=True)
+ min_val = to_quant.amin(dim=1, keepdim=True)
+ max_int = 2**n_bit - 1
+ scales = (max_val - min_val).clamp(min=1e-6) / max_int
+ zeros = min_val + scales * (2 ** (n_bit - 1))
+ return scales.to(torch.bfloat16).reshape(w.shape[0], -1), zeros.to(
+ torch.bfloat16
+ ).reshape(w.shape[0], -1)
+
+
+def pack_scales_and_zeros(scales, zeros):
+ assert scales.shape == zeros.shape
+ assert scales.dtype == torch.bfloat16
+ assert zeros.dtype == torch.bfloat16
+ return (
+ torch.cat(
+ [
+ scales.reshape(scales.size(0), scales.size(1), 1),
+ zeros.reshape(zeros.size(0), zeros.size(1), 1),
+ ],
+ 2,
+ )
+ .transpose(0, 1)
+ .contiguous()
+ )
+
+
+def unpack_scales_and_zeros(scales_and_zeros):
+ assert len(scales_and_zeros.shape) == 3 and scales_and_zeros.shape[2] == 2
+ assert scales_and_zeros.dtype == torch.float
+ return torch.split(scales_and_zeros.transpose(0, 1), 1, 2)
+
+
+def group_quantize_tensor_from_qparams(w, scales, zeros, n_bit=4, groupsize=128):
+ assert groupsize > 1
+ # needed for GPTQ single column quantize
+ if groupsize > w.shape[-1] and scales.shape[-1] == 1:
+ groupsize = w.shape[-1]
+
+ assert w.shape[-1] % groupsize == 0
+ assert w.dim() == 2
+
+ to_quant = w.reshape(-1, groupsize)
+ assert torch.isnan(to_quant).sum() == 0
+
+ scales = scales.reshape(-1, 1)
+ zeros = zeros.reshape(-1, 1)
+ min_val = zeros - scales * (2 ** (n_bit - 1))
+ max_int = 2**n_bit - 1
+ min_int = 0
+ w_int32 = (
+ to_quant.sub(min_val)
+ .div(scales)
+ .round()
+ .clamp_(min_int, max_int)
+ .to(torch.int32)
+ .reshape_as(w)
+ )
+
+ return w_int32
+
+
+def group_quantize_tensor(w, n_bit=4, groupsize=128):
+ scales, zeros = get_group_qparams(w, n_bit, groupsize)
+ w_int32 = group_quantize_tensor_from_qparams(w, scales, zeros, n_bit, groupsize)
+ scales_and_zeros = pack_scales_and_zeros(scales, zeros)
+ return w_int32, scales_and_zeros
+
+
+def group_dequantize_tensor_from_qparams(
+ w_int32, scales, zeros, n_bit=4, groupsize=128
+):
+ assert groupsize > 1
+ # needed for GPTQ single column dequantize
+ if groupsize > w_int32.shape[-1] and scales.shape[-1] == 1:
+ groupsize = w_int32.shape[-1]
+ assert w_int32.shape[-1] % groupsize == 0
+ assert w_int32.dim() == 2
+
+ w_int32_grouped = w_int32.reshape(-1, groupsize)
+ scales = scales.reshape(-1, 1)
+ zeros = zeros.reshape(-1, 1)
+
+ w_dq = (
+ w_int32_grouped.sub(2 ** (n_bit - 1)).mul(scales).add(zeros).reshape_as(w_int32)
+ )
+ return w_dq
+
+
+def group_dequantize_tensor(w_int32, scales_and_zeros, n_bit=4, groupsize=128):
+ scales, zeros = unpack_scales_and_zeros(scales_and_zeros)
+ return group_dequantize_tensor_from_qparams(
+ w_int32, scales, zeros, n_bit, groupsize
+ )
+
+class QuantHandler:
+ def __init__(self, mod):
+ self.mod = mod
+
+ def create_quantized_state_dict(self) -> "StateDict":
+ pass
+
+ def convert_for_runtime(self) -> "nn.Module":
+ pass
+
+class GPTQQuantHandler(QuantHandler):
+ """
+ This class implements a GPTQ QuantHandler that can be used to apply GPTQ to a model in concert with the GenericGPTQRunner class.
+ Unlike the base QuantHandler class, the user does not need to implement the create_quantized_state_dict, instead they have to reimplement
+ __init__ such that it defines the functions for the quantization mode. User is expected to reimplement convert_for_runtime.
+
+ The following functions (which must be defined in __init__) are used to define the quantization mode for both GPTQ and
+ create_quantized_state_dict. Here is a description of each function.
+
+ get_qparams_func:
+ A function that calculates the quantization qparams for an input tensor.
+ Args:
+ weight: A 2d weight tensor with non-integer dtype.
+ Returns:
+ qparams: it can have any format but will need to be handled by the other defined functions below.
+
+ quantize_func:
+ A function that applies quantization to an input tensor. It should be noted
+ that this function needs to be able to handle quantizing the entire weight tensor, a single group,
+ or a single column.
+ Args:
+ weight: A 2d weight tensor with non-integer dtype.
+ qparams: the output from get_qparams_func
+ Returns:
+ quantized_weight: A 2d quantized weight tensor (generally with an integer dtype)
+
+
+ dequantize_func:
+ A function that dequantizes an input quantized weight tensor. It should be noted
+ that this function needs to be able to handle dequantizing the entire weight tensor, a single group,
+ or a single column.
+ Args:
+ quantized_weight: A 2d quantized weight tensor (generally with an integer dtype)
+ qparams: the output from get_qparams_func
+ Returns:
+ weight: A 2d weight tensor with non-integer dtype.
+
+ combine_qparams_list_func:
+ A function that combines several qparams into one qparam.
+ Args:
+ qparams_list: a list of qparams objects, each obtained by calling get_qparams_func
+ on a single group from a weight tensor
+ Returns:
+ qparams: an object of the same format as the qparams above.
+
+ skip_layer_func:
+ A function that determines which linear layers should be skipped during GPTQ
+ Args:
+ weight: A 2d weight tensor with non-integer dtype.
+ Returns:
+ skip: boolean indicating whether layer should be skipped
+
+ make_names_and_values_dict_func:
+ A function that prepares the qparams and quantized_weight and creates a dictionary indicating how they
+ should be inserted into the state_dict. Generally any packing of the weight and qparams should be done here.
+ Args:
+ quantized_weight: A 2d quantized weight tensor (generally with an integer dtype)
+ qparams: the output from get_qparams_func
+ Returns:
+ names_and_values_dict: a dictionary mapping the name of the parameters of the quantized module to the
+ corresponding quantized weights and qparams.
+ """
+ def __init__(self):
+ assert self.mod is not None
+ assert self.get_qparams_func is not None
+ assert self.quantize_func is not None
+ assert self.dequantize_func is not None
+ assert self.combine_qparams_list_func is not None
+ assert self.make_names_and_values_dict_func is not None
+
+ @staticmethod
+ def get_inputs(model, tokenizer, calibration_tasks, calibration_limit, calibration_seq_length, pad_calibration_inputs) -> "MultiInput":
+ input_recorder = InputRecorder(
+ model,
+ tokenizer,
+ calibration_seq_length,
+ pad_calibration_inputs,
+ )
+
+ try:
+ lm_eval.tasks.initialize_tasks()
+ except:
+ pass
+ task_dict = get_task_dict(calibration_tasks)
+ print("Obtaining GPTQ calibration inputs on: ", calibration_tasks)
+
+ evaluate(
+ input_recorder,
+ task_dict,
+ limit=calibration_limit,
+ )
+ inputs = input_recorder.get_recorded_inputs()
+ assert inputs is not None, (
+ f"No inputs were collected, use a task other than {calibration_tasks}, "+
+ f"use option pad_calibration_inputs, or decrease calibration_sequence_length (currently "+
+ f"{calibration_seq_length})"
+ )
+ print(f"Obtained {len(inputs[0].values)} calibration samples")
+ return inputs
+
+ @torch.no_grad()
+ def create_quantized_state_dict(
+ self,
+ tokenizer,
+ blocksize,
+ percdamp,
+ groupsize,
+ calibration_tasks,
+ calibration_limit,
+ calibration_seq_length,
+ pad_calibration_inputs,
+ ) -> "StateDict":
+ inputs = GPTQQuantHandler.get_inputs(self.mod, tokenizer, calibration_tasks, calibration_limit, calibration_seq_length, pad_calibration_inputs)
+ print("Tracing model for GPTQ")
+ GPTQ_runner = GenericGPTQRunner(
+ self.mod,
+ inputs,
+ blocksize,
+ percdamp,
+ groupsize,
+ ).configure_quantization_mode(
+ self.get_qparams_func,
+ self.quantize_func,
+ self.dequantize_func,
+ self.combine_qparams_list_func,
+ self.make_names_and_values_dict_func,
+ self.skip_layer_func
+ )
+
+ print("Applying GPTQ to weights")
+ GPTQ_runner.run()
+ return GPTQ_runner.get_quantized_state_dict()
+
+ def convert_for_runtime(self) -> "nn.Module":
+ pass
+
+##### Weight-only int8 per-channel quantized code ######
+
+def replace_linear_weight_only_int8_per_channel(module):
+ for name, child in module.named_children():
+ if isinstance(child, nn.Linear):
+ setattr(module, name, WeightOnlyInt8Linear(child.in_features, child.out_features))
+ else:
+ replace_linear_weight_only_int8_per_channel(child)
+
+class WeightOnlyInt8QuantHandler:
+ def __init__(self, mod):
+ self.mod = mod
+
+ @torch.no_grad()
+ def create_quantized_state_dict(self):
+ cur_state_dict = self.mod.state_dict()
+ for fqn, mod in self.mod.named_modules():
+ if isinstance(mod, torch.nn.Linear):
+ int8_weight, scales, _ = dynamically_quantize_per_channel(mod.weight.float(), -128, 127, torch.int8)
+ cur_state_dict[f"{fqn}.weight"] = int8_weight
+ cur_state_dict[f"{fqn}.scales"] = scales.to(mod.weight.dtype)
+
+ return cur_state_dict
+
+ def convert_for_runtime(self):
+ replace_linear_weight_only_int8_per_channel(self.mod)
+ return self.mod
+
+
+class WeightOnlyInt8Linear(torch.nn.Module):
+ __constants__ = ['in_features', 'out_features']
+ in_features: int
+ out_features: int
+ weight: torch.Tensor
+
+ def __init__(self, in_features: int, out_features: int, bias: bool = True,
+ device=None, dtype=None) -> None:
+ factory_kwargs = {'device': device, 'dtype': dtype}
+ super().__init__()
+ self.in_features = in_features
+ self.out_features = out_features
+ self.register_buffer("weight", torch.empty((out_features, in_features), dtype=torch.int8))
+ self.register_buffer("scales", torch.ones(out_features, dtype=torch.bfloat16))
+
+ def forward(self, input: torch.Tensor) -> torch.Tensor:
+ return F.linear(input, self.weight.to(dtype=input.dtype)) * self.scales
+
+##### weight only int4 per channel groupwise quantized code ######
+
+def prepare_int4_weight_and_scales_and_zeros(weight_bf16, groupsize, inner_k_tiles):
+ weight_int32, scales_and_zeros = group_quantize_tensor(
+ weight_bf16, n_bit=4, groupsize=groupsize
+ )
+ weight_int4pack = torch.ops.aten._convert_weight_to_int4pack(weight_int32, inner_k_tiles)
+ return weight_int4pack, scales_and_zeros
+
+
+def linear_forward_int4(x, weight_int4pack, scales_and_zeros, out_features, groupsize):
+ origin_x_size = x.size()
+ x = x.reshape(-1, origin_x_size[-1])
+ c = torch.ops.aten._weight_int4pack_mm(x, weight_int4pack, groupsize, scales_and_zeros)
+ new_shape = origin_x_size[:-1] + (out_features,)
+ c = c.reshape(new_shape)
+ return c
+
+
+def _check_linear_int4_k(k, groupsize = 1, inner_k_tiles = 1):
+ return k % groupsize == 0 and k % (inner_k_tiles * 16) == 0
+
+def replace_linear_int4(module, groupsize, inner_k_tiles, padding):
+ for name, child in module.named_children():
+ if isinstance(child, nn.Linear):
+ if _check_linear_int4_k(child.in_features, groupsize, inner_k_tiles):
+ setattr(module, name, WeightOnlyInt4Linear(
+ child.in_features, child.out_features, bias=False,
+ groupsize=groupsize, inner_k_tiles=inner_k_tiles, padding=False,
+ ))
+ elif padding:
+ setattr(module, name, WeightOnlyInt4Linear(
+ child.in_features, child.out_features, bias=False,
+ groupsize=groupsize, inner_k_tiles=inner_k_tiles, padding=True,
+ ))
+ else:
+ replace_linear_int4(child, groupsize, inner_k_tiles, padding)
+
+
+class WeightOnlyInt4QuantHandler:
+ def __init__(self, mod, groupsize=128, inner_k_tiles=8, padding=True):
+ self.mod = mod
+ self.groupsize = groupsize
+ self.inner_k_tiles = inner_k_tiles
+ self.padding = padding
+ assert groupsize in [32, 64, 128, 256]
+ assert inner_k_tiles in [2, 4, 8]
+
+ @torch.no_grad()
+ def create_quantized_state_dict(self, use_cuda = True):
+ if use_cuda:
+ device="cuda"
+ else:
+ device="cpu"
+
+ cur_state_dict = self.mod.state_dict()
+ for fqn, mod in self.mod.named_modules():
+ if isinstance(mod, torch.nn.Linear):
+ assert not mod.bias
+ out_features = mod.out_features
+ in_features = mod.in_features
+ assert out_features % 8 == 0, "require out_features % 8 == 0"
+ print(f"linear: {fqn}, in={in_features}, out={out_features}")
+
+ weight = mod.weight.data
+ if not _check_linear_int4_k(in_features, self.groupsize, self.inner_k_tiles):
+ if self.padding:
+ from model import find_multiple
+ import torch.nn.functional as F
+ print(f"warning: {fqn} is padded to satisfy in_features % 1024 == 0")
+ padded_in_features = find_multiple(in_features, 1024)
+ weight = F.pad(weight, pad=(0, padded_in_features - in_features))
+ else:
+ print(f"warning: {fqn} is skipped, int4 requires that in_features is 32, 64, or is divisible by 1024, " +
+ "and that groupsize and inner_k_tiles*16 evenly divide into it")
+ continue
+ weight_int4pack, scales_and_zeros = prepare_int4_weight_and_scales_and_zeros(
+ weight.to(torch.bfloat16).to(device=device), self.groupsize, self.inner_k_tiles
+ )
+ cur_state_dict[f"{fqn}.weight"] = weight_int4pack.to('cpu')
+ cur_state_dict[f"{fqn}.scales_and_zeros"] = scales_and_zeros.to('cpu')
+
+ return cur_state_dict
+
+ def convert_for_runtime(self):
+ replace_linear_int4(self.mod, self.groupsize, self.inner_k_tiles, self.padding)
+ return self.mod
+
+class WeightOnlyInt4GPTQQuantHandler(GPTQQuantHandler):
+ def __init__(self, mod, groupsize=128, inner_k_tiles=8, padding=True):
+ from model import find_multiple
+ self.mod = mod
+ self.groupsize = groupsize
+ self.inner_k_tiles = inner_k_tiles
+ self.padding = padding
+ self.get_qparams_func = lambda w: get_group_qparams(w, 4, groupsize)
+ self.quantize_func = lambda w, qparams: \
+ group_quantize_tensor_from_qparams(w, qparams[0], qparams[1], 4, groupsize)
+ self.dequantize_func = lambda q, qparams: \
+ group_dequantize_tensor_from_qparams(q, qparams[0], qparams[1], 4, groupsize).float()
+ self.combine_qparams_list_func = lambda qparams_list: \
+ [torch.cat(x, dim=1) for x in zip(*qparams_list)]
+ # skip unless padding=True or its correctly sized
+ self.skip_layer_func = lambda linear_weight: not (
+ _check_linear_int4_k(linear_weight.shape[-1], groupsize, inner_k_tiles) or padding
+ )
+ # we need to do the padding here, both for q and the qparams if necessary
+ def make_names_and_values_dict_func(q, qparams):
+ k = q.shape[1]
+ new_k = find_multiple(k, 1024)
+ # how much we need to pad the weight
+ delta_k = new_k - q.shape[1]
+ final_q = torch.ops.aten._convert_weight_to_int4pack(F.pad(q, pad=(0, delta_k)), inner_k_tiles)
+ scales_and_zeros = pack_scales_and_zeros(*qparams)
+ # how many new groups we need for padded weight
+ delta_groups = new_k // groupsize - scales_and_zeros.shape[0]
+ final_s_and_z = F.pad(scales_and_zeros, pad=(0,0,0,0,0, delta_groups), value=1)
+ return {"weight": final_q, "scales_and_zeros": final_s_and_z}
+ self.make_names_and_values_dict_func = make_names_and_values_dict_func
+ super().__init__()
+
+
+ def convert_for_runtime(self):
+ replace_linear_int4(self.mod, self.groupsize, self.inner_k_tiles, self.padding)
+ return self.mod
+
+class WeightOnlyInt4Linear(torch.nn.Module):
+ __constants__ = ['in_features', 'out_features']
+ in_features: int
+ out_features: int
+ weight: torch.Tensor
+
+ def __init__(
+ self, in_features: int, out_features: int,
+ bias=True, device=None, dtype=None, groupsize: int = 128, inner_k_tiles: int = 8, padding: bool = True,
+ ) -> None:
+ super().__init__()
+ self.padding = padding
+ if padding:
+ from model import find_multiple
+ self.origin_in_features = in_features
+ in_features = find_multiple(in_features, 1024)
+
+ self.in_features = in_features
+ self.out_features = out_features
+ assert not bias, "require bias=False"
+ self.groupsize = groupsize
+ self.inner_k_tiles = inner_k_tiles
+
+ assert out_features % 8 == 0, "require out_features % 8 == 0"
+ assert in_features % (inner_k_tiles * 16) == 0, "require in_features % (innerKTiles * 16) == 0"
+ self.register_buffer(
+ "weight",
+ torch.empty((out_features // 8, in_features // (inner_k_tiles * 16), 32, inner_k_tiles // 2), dtype=torch.int32)
+ )
+ self.register_buffer(
+ "scales_and_zeros",
+ torch.empty((in_features // groupsize, out_features, 2), dtype=torch.bfloat16)
+ )
+
+ def forward(self, input: torch.Tensor) -> torch.Tensor:
+ input = input.to(torch.bfloat16)
+ if self.padding:
+ import torch.nn.functional as F
+ input = F.pad(input, pad=(0, self.in_features - self.origin_in_features))
+ return linear_forward_int4(
+ input,
+ self.weight, self.scales_and_zeros, self.out_features, self.groupsize
+ )
+
+
+def quantize(
+ checkpoint_path: Path = Path("checkpoints/meta-llama/Llama-2-7b-chat-hf/model.pth"),
+ mode: str = 'int8',
+ # following arguments only available when setting int4 quantization.
+ groupsize: int = 128,
+ # following arguments only used for GPTQ
+ calibration_tasks: list = ["hellaswag"],
+ calibration_limit: int = 1000,
+ calibration_seq_length: int = 100,
+ pad_calibration_inputs: bool = False,
+ percdamp: float = .01,
+ blocksize: int = 128,
+ label: str = '',
+) -> None:
+ assert checkpoint_path.is_file(), checkpoint_path
+
+ device = 'cpu'
+ precision = torch.bfloat16
+
+ print("Loading model ...")
+ t0 = time.time()
+
+ with torch.device('meta'):
+ model = Transformer.from_name(checkpoint_path.parent.name)
+
+ checkpoint = torch.load(str(checkpoint_path), mmap=True, weights_only=True)
+ model.load_state_dict(checkpoint, assign=True)
+ model = model.to(dtype=precision, device=device)
+
+ if mode == 'int8':
+ print("Quantizing model weights for int8 weight-only symmetric per-channel quantization")
+ quant_handler = WeightOnlyInt8QuantHandler(model)
+ quantized_state_dict = quant_handler.create_quantized_state_dict()
+
+ dir_name = checkpoint_path.parent
+ base_name = checkpoint_path.name
+ new_base_name = base_name.replace('.pth', f'{label}int8.pth')
+
+ elif mode == 'int4':
+ print("Quantizing model weights for int4 weight-only affine per-channel groupwise quantization")
+ quant_handler = WeightOnlyInt4QuantHandler(model, groupsize)
+ quantized_state_dict = quant_handler.create_quantized_state_dict()
+
+ dir_name = checkpoint_path.parent
+ base_name = checkpoint_path.name
+ new_base_name = base_name.replace('.pth', f"{label}int4.g{groupsize}.pth")
+
+ elif mode == 'int4-gptq':
+ print("Quantizing model weights for int4 weight-only affine per-channel groupwise quantization using GPTQ...")
+ quant_handler = WeightOnlyInt4GPTQQuantHandler(model, groupsize)
+
+ tokenizer_path = checkpoint_path.parent / "tokenizer.model"
+ assert tokenizer_path.is_file(), str(tokenizer_path)
+ tokenizer = get_tokenizer(tokenizer_path, checkpoint_path)
+
+ quantized_state_dict = quant_handler.create_quantized_state_dict(
+ tokenizer,
+ blocksize,
+ percdamp,
+ groupsize,
+ calibration_tasks,
+ calibration_limit,
+ calibration_seq_length,
+ pad_calibration_inputs
+ )
+
+ dir_name = checkpoint_path.parent
+ base_name = checkpoint_path.name
+ new_base_name = base_name.replace('.pth', f"{label}int4-gptq.g{groupsize}.pth")
+ else:
+ raise ValueError(f"Invalid quantization mode {mode} needs to be one of [int8, int4, int4-gpptq]")
+
+ quantize_path = dir_name / new_base_name
+ print(f"Writing quantized weights to {quantize_path}")
+ quantize_path.unlink(missing_ok=True) # remove existing file if one already there
+ torch.save(quantized_state_dict, quantize_path)
+ print(f"Quantization complete took {time.time() - t0:.02f} seconds")
+ return
+
+if __name__ == '__main__':
+ import argparse
+ parser = argparse.ArgumentParser(description='Quantize a model.')
+ parser.add_argument('--checkpoint_path', type=Path, default=Path("checkpoints/meta-llama/Llama-2-7b-chat-hf/model.pth"), help='Path to the model checkpoint to be quantized.')
+ parser.add_argument('--mode', '-q', type=str, default='int8', choices=['int8', 'int4', 'int4-gptq'], help='type of quantization to perform')
+ parser.add_argument('--groupsize', type=int, default=32, help='Group size for int4 quantization.')
+ parser.add_argument('--calibration_tasks', type=str, nargs='+', default=['wikitext'], help='tasks to do gptq calibration on, if doing gptq')
+ parser.add_argument('--calibration_limit', type=int, default=1000, help='number of samples to use for gptq calibration')
+ parser.add_argument('--calibration_seq_length', type=int, default=100, help='length of sequences to use for gptq calibration')
+ parser.add_argument('--pad_calibration_inputs', type=bool, default=False, help='pads sequences shorter than calibration_seq_length to that length, yielding more calibration inputs but running much slower')
+ parser.add_argument('--percdamp', type=float, default=.01, help='gptq percentage dampening')
+ parser.add_argument('--blocksize', type=int, default=128, help='blocksize for gptq')
+ parser.add_argument('--label', type=str, default='_', help='label to add to output filename')
+
+ args = parser.parse_args()
+ quantize(args.checkpoint_path, args.mode, args.groupsize, args.calibration_tasks, args.calibration_limit, args.calibration_seq_length, args.pad_calibration_inputs, args.percdamp, args.blocksize, args.label)
diff --git a/modules/hifigan/__pycache__/f0_predictor.cpython-310.pyc b/modules/hifigan/__pycache__/f0_predictor.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bf0eeed7feb3d694da20afa9c2f181021bf14877
Binary files /dev/null and b/modules/hifigan/__pycache__/f0_predictor.cpython-310.pyc differ
diff --git a/modules/hifigan/__pycache__/generator.cpython-310.pyc b/modules/hifigan/__pycache__/generator.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7bda6faedfb76ae79bcaf88ff8818df26eb048c8
Binary files /dev/null and b/modules/hifigan/__pycache__/generator.cpython-310.pyc differ
diff --git a/modules/hifigan/f0_predictor.py b/modules/hifigan/f0_predictor.py
new file mode 100644
index 0000000000000000000000000000000000000000..1daec26f3adbf6f7fc7b1c7088505ea763bdf167
--- /dev/null
+++ b/modules/hifigan/f0_predictor.py
@@ -0,0 +1,55 @@
+# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+import torch.nn as nn
+from torch.nn.utils import weight_norm
+
+
+class ConvRNNF0Predictor(nn.Module):
+ def __init__(self,
+ num_class: int = 1,
+ in_channels: int = 80,
+ cond_channels: int = 512
+ ):
+ super().__init__()
+
+ self.num_class = num_class
+ self.condnet = nn.Sequential(
+ weight_norm(
+ nn.Conv1d(in_channels, cond_channels, kernel_size=3, padding=1)
+ ),
+ nn.ELU(),
+ weight_norm(
+ nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)
+ ),
+ nn.ELU(),
+ weight_norm(
+ nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)
+ ),
+ nn.ELU(),
+ weight_norm(
+ nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)
+ ),
+ nn.ELU(),
+ weight_norm(
+ nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)
+ ),
+ nn.ELU(),
+ )
+ self.classifier = nn.Linear(in_features=cond_channels, out_features=self.num_class)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.condnet(x)
+ x = x.transpose(1, 2)
+ return torch.abs(self.classifier(x).squeeze(-1))
diff --git a/modules/hifigan/generator.py b/modules/hifigan/generator.py
new file mode 100644
index 0000000000000000000000000000000000000000..21d5bc7a9417cd840bedbdf77bc8fe41bc001b50
--- /dev/null
+++ b/modules/hifigan/generator.py
@@ -0,0 +1,454 @@
+# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""HIFI-GAN"""
+
+import typing as tp
+import numpy as np
+from scipy.signal import get_window
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import Conv1d
+from torch.nn import ConvTranspose1d
+from torch.nn.utils import remove_weight_norm
+from torch.nn.utils import weight_norm
+from torch.distributions.uniform import Uniform
+
+from torch import sin
+from torch.nn.parameter import Parameter
+
+
+"""hifigan based generator implementation.
+
+This code is modified from https://github.com/jik876/hifi-gan
+ ,https://github.com/kan-bayashi/ParallelWaveGAN and
+ https://github.com/NVIDIA/BigVGAN
+
+"""
+class Snake(nn.Module):
+ '''
+ Implementation of a sine-based periodic activation function
+ Shape:
+ - Input: (B, C, T)
+ - Output: (B, C, T), same shape as the input
+ Parameters:
+ - alpha - trainable parameter
+ References:
+ - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+ https://arxiv.org/abs/2006.08195
+ Examples:
+ >>> a1 = snake(256)
+ >>> x = torch.randn(256)
+ >>> x = a1(x)
+ '''
+ def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):
+ '''
+ Initialization.
+ INPUT:
+ - in_features: shape of the input
+ - alpha: trainable parameter
+ alpha is initialized to 1 by default, higher values = higher-frequency.
+ alpha will be trained along with the rest of your model.
+ '''
+ super(Snake, self).__init__()
+ self.in_features = in_features
+
+ # initialize alpha
+ self.alpha_logscale = alpha_logscale
+ if self.alpha_logscale: # log scale alphas initialized to zeros
+ self.alpha = Parameter(torch.zeros(in_features) * alpha)
+ else: # linear scale alphas initialized to ones
+ self.alpha = Parameter(torch.ones(in_features) * alpha)
+
+ self.alpha.requires_grad = alpha_trainable
+
+ self.no_div_by_zero = 0.000000001
+
+ def forward(self, x):
+ '''
+ Forward pass of the function.
+ Applies the function to the input elementwise.
+ Snake ∶= x + 1/a * sin^2 (xa)
+ '''
+ alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T]
+ if self.alpha_logscale:
+ alpha = torch.exp(alpha)
+ x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+ return x
+
+def get_padding(kernel_size, dilation=1):
+ return int((kernel_size * dilation - dilation) / 2)
+
+
+def init_weights(m, mean=0.0, std=0.01):
+ classname = m.__class__.__name__
+ if classname.find("Conv") != -1:
+ m.weight.data.normal_(mean, std)
+
+
+
+class ResBlock(torch.nn.Module):
+ """Residual block module in HiFiGAN/BigVGAN."""
+ def __init__(
+ self,
+ channels: int = 512,
+ kernel_size: int = 3,
+ dilations: tp.List[int] = [1, 3, 5],
+ ):
+ super(ResBlock, self).__init__()
+ self.convs1 = nn.ModuleList()
+ self.convs2 = nn.ModuleList()
+
+ for dilation in dilations:
+ self.convs1.append(
+ weight_norm(
+ Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation,
+ padding=get_padding(kernel_size, dilation)
+ )
+ )
+ )
+ self.convs2.append(
+ weight_norm(
+ Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=1,
+ padding=get_padding(kernel_size, 1)
+ )
+ )
+ )
+ self.convs1.apply(init_weights)
+ self.convs2.apply(init_weights)
+ self.activations1 = nn.ModuleList([
+ Snake(channels, alpha_logscale=False)
+ for _ in range(len(self.convs1))
+ ])
+ self.activations2 = nn.ModuleList([
+ Snake(channels, alpha_logscale=False)
+ for _ in range(len(self.convs2))
+ ])
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ for idx in range(len(self.convs1)):
+ xt = self.activations1[idx](x)
+ xt = self.convs1[idx](xt)
+ xt = self.activations2[idx](xt)
+ xt = self.convs2[idx](xt)
+ x = xt + x
+ return x
+
+ def remove_weight_norm(self):
+ for idx in range(len(self.convs1)):
+ remove_weight_norm(self.convs1[idx])
+ remove_weight_norm(self.convs2[idx])
+
+class SineGen(torch.nn.Module):
+ """ Definition of sine generator
+ SineGen(samp_rate, harmonic_num = 0,
+ sine_amp = 0.1, noise_std = 0.003,
+ voiced_threshold = 0,
+ flag_for_pulse=False)
+ samp_rate: sampling rate in Hz
+ harmonic_num: number of harmonic overtones (default 0)
+ sine_amp: amplitude of sine-wavefrom (default 0.1)
+ noise_std: std of Gaussian noise (default 0.003)
+ voiced_thoreshold: F0 threshold for U/V classification (default 0)
+ flag_for_pulse: this SinGen is used inside PulseGen (default False)
+ Note: when flag_for_pulse is True, the first time step of a voiced
+ segment is always sin(np.pi) or cos(0)
+ """
+
+ def __init__(self, samp_rate, harmonic_num=0,
+ sine_amp=0.1, noise_std=0.003,
+ voiced_threshold=0):
+ super(SineGen, self).__init__()
+ self.sine_amp = sine_amp
+ self.noise_std = noise_std
+ self.harmonic_num = harmonic_num
+ self.sampling_rate = samp_rate
+ self.voiced_threshold = voiced_threshold
+
+ def _f02uv(self, f0):
+ # generate uv signal
+ uv = (f0 > self.voiced_threshold).type(torch.float32)
+ return uv
+
+ @torch.no_grad()
+ def forward(self, f0):
+ """
+ :param f0: [B, 1, sample_len], Hz
+ :return: [B, 1, sample_len]
+ """
+
+ F_mat = torch.zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1))).to(f0.device)
+ for i in range(self.harmonic_num + 1):
+ F_mat[:, i: i + 1, :] = f0 * (i + 1) / self.sampling_rate
+
+ theta_mat = 2 * np.pi * (torch.cumsum(F_mat, dim=-1) % 1)
+ u_dist = Uniform(low=-np.pi, high=np.pi)
+ phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device)
+ phase_vec[:, 0, :] = 0
+
+ # generate sine waveforms
+ sine_waves = self.sine_amp * torch.sin(theta_mat + phase_vec)
+
+ # generate uv signal
+ uv = self._f02uv(f0)
+
+ # noise: for unvoiced should be similar to sine_amp
+ # std = self.sine_amp/3 -> max value ~ self.sine_amp
+ # . for voiced regions is self.noise_std
+ noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
+ noise = noise_amp * torch.randn_like(sine_waves)
+
+ # first: set the unvoiced part to 0 by uv
+ # then: additive noise
+ sine_waves = sine_waves * uv + noise
+ return sine_waves, uv, noise
+
+
+class SourceModuleHnNSF(torch.nn.Module):
+ """ SourceModule for hn-nsf
+ SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
+ add_noise_std=0.003, voiced_threshod=0)
+ sampling_rate: sampling_rate in Hz
+ harmonic_num: number of harmonic above F0 (default: 0)
+ sine_amp: amplitude of sine source signal (default: 0.1)
+ add_noise_std: std of additive Gaussian noise (default: 0.003)
+ note that amplitude of noise in unvoiced is decided
+ by sine_amp
+ voiced_threshold: threhold to set U/V given F0 (default: 0)
+ Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+ F0_sampled (batchsize, length, 1)
+ Sine_source (batchsize, length, 1)
+ noise_source (batchsize, length 1)
+ uv (batchsize, length, 1)
+ """
+
+ def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1,
+ add_noise_std=0.003, voiced_threshod=0):
+ super(SourceModuleHnNSF, self).__init__()
+
+ self.sine_amp = sine_amp
+ self.noise_std = add_noise_std
+
+ # to produce sine waveforms
+ self.l_sin_gen = SineGen(sampling_rate, harmonic_num,
+ sine_amp, add_noise_std, voiced_threshod)
+
+ # to merge source harmonics into a single excitation
+ self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
+ self.l_tanh = torch.nn.Tanh()
+
+ def forward(self, x):
+ """
+ Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+ F0_sampled (batchsize, length, 1)
+ Sine_source (batchsize, length, 1)
+ noise_source (batchsize, length 1)
+ """
+ # source for harmonic branch
+ with torch.no_grad():
+ sine_wavs, uv, _ = self.l_sin_gen(x.transpose(1, 2))
+ sine_wavs = sine_wavs.transpose(1, 2)
+ uv = uv.transpose(1, 2)
+ sine_merge = self.l_tanh(self.l_linear(sine_wavs))
+
+ # source for noise branch, in the same shape as uv
+ noise = torch.randn_like(uv) * self.sine_amp / 3
+ return sine_merge, noise, uv
+
+
+class HiFTGenerator(nn.Module):
+ """
+ HiFTNet Generator: Neural Source Filter + ISTFTNet
+ https://arxiv.org/abs/2309.09493
+ """
+ def __init__(
+ self,
+ in_channels: int = 80,
+ base_channels: int = 512,
+ nb_harmonics: int = 8,
+ sampling_rate: int = 22050,
+ nsf_alpha: float = 0.1,
+ nsf_sigma: float = 0.003,
+ nsf_voiced_threshold: float = 10,
+ upsample_rates: tp.List[int] = [8, 8],
+ upsample_kernel_sizes: tp.List[int] = [16, 16],
+ istft_params: tp.Dict[str, int] = {"n_fft": 16, "hop_len": 4},
+ resblock_kernel_sizes: tp.List[int] = [3, 7, 11],
+ resblock_dilation_sizes: tp.List[tp.List[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+ source_resblock_kernel_sizes: tp.List[int] = [7, 11],
+ source_resblock_dilation_sizes: tp.List[tp.List[int]] = [[1, 3, 5], [1, 3, 5]],
+ lrelu_slope: float = 0.1,
+ audio_limit: float = 0.99,
+ f0_predictor: torch.nn.Module = None,
+ ):
+ super(HiFTGenerator, self).__init__()
+
+ self.out_channels = 1
+ self.nb_harmonics = nb_harmonics
+ self.sampling_rate = sampling_rate
+ self.istft_params = istft_params
+ self.lrelu_slope = lrelu_slope
+ self.audio_limit = audio_limit
+
+ self.num_kernels = len(resblock_kernel_sizes)
+ self.num_upsamples = len(upsample_rates)
+ self.m_source = SourceModuleHnNSF(
+ sampling_rate=sampling_rate,
+ upsample_scale=np.prod(upsample_rates) * istft_params["hop_len"],
+ harmonic_num=nb_harmonics,
+ sine_amp=nsf_alpha,
+ add_noise_std=nsf_sigma,
+ voiced_threshod=nsf_voiced_threshold)
+ self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * istft_params["hop_len"])
+
+ self.conv_pre = weight_norm(
+ Conv1d(in_channels, base_channels, 7, 1, padding=3)
+ )
+
+ # Up
+ self.ups = nn.ModuleList()
+ for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+ self.ups.append(
+ weight_norm(
+ ConvTranspose1d(
+ base_channels // (2**i),
+ base_channels // (2**(i + 1)),
+ k,
+ u,
+ padding=(k - u) // 2,
+ )
+ )
+ )
+
+ # Down
+ self.source_downs = nn.ModuleList()
+ self.source_resblocks = nn.ModuleList()
+ downsample_rates = [1] + upsample_rates[::-1][:-1]
+ downsample_cum_rates = np.cumprod(downsample_rates)
+ for i, (u, k, d) in enumerate(zip(downsample_cum_rates[::-1], source_resblock_kernel_sizes,
+ source_resblock_dilation_sizes)):
+ if u == 1:
+ self.source_downs.append(
+ Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), 1, 1)
+ )
+ else:
+ self.source_downs.append(
+ Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), u * 2, u, padding=(u // 2))
+ )
+
+ self.source_resblocks.append(
+ ResBlock(base_channels // (2 ** (i + 1)), k, d)
+ )
+
+ self.resblocks = nn.ModuleList()
+ for i in range(len(self.ups)):
+ ch = base_channels // (2**(i + 1))
+ for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
+ self.resblocks.append(ResBlock(ch, k, d))
+
+ self.conv_post = weight_norm(Conv1d(ch, istft_params["n_fft"] + 2, 7, 1, padding=3))
+ self.ups.apply(init_weights)
+ self.conv_post.apply(init_weights)
+ self.reflection_pad = nn.ReflectionPad1d((1, 0))
+ self.stft_window = torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32))
+ self.f0_predictor = f0_predictor
+
+ def _f02source(self, f0: torch.Tensor) -> torch.Tensor:
+ f0 = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t
+
+ har_source, _, _ = self.m_source(f0)
+ return har_source.transpose(1, 2)
+
+ def _stft(self, x):
+ spec = torch.stft(
+ x,
+ self.istft_params["n_fft"], self.istft_params["hop_len"], self.istft_params["n_fft"], window=self.stft_window.to(x.device),
+ return_complex=True)
+ spec = torch.view_as_real(spec) # [B, F, TT, 2]
+ return spec[..., 0], spec[..., 1]
+
+ def _istft(self, magnitude, phase):
+ magnitude = torch.clip(magnitude, max=1e2)
+ real = magnitude * torch.cos(phase)
+ img = magnitude * torch.sin(phase)
+ inverse_transform = torch.istft(torch.complex(real, img), self.istft_params["n_fft"], self.istft_params["hop_len"], self.istft_params["n_fft"], window=self.stft_window.to(magnitude.device))
+ return inverse_transform
+
+ def forward(self, x: torch.Tensor, f0=None) -> torch.Tensor:
+ if f0 is None:
+ f0 = self.f0_predictor(x)
+ s = self._f02source(f0)
+
+ s_stft_real, s_stft_imag = self._stft(s.squeeze(1))
+ s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1)
+
+ x = self.conv_pre(x)
+ for i in range(self.num_upsamples):
+ x = F.leaky_relu(x, self.lrelu_slope)
+ x = self.ups[i](x)
+
+ if i == self.num_upsamples - 1:
+ x = self.reflection_pad(x)
+
+ # fusion
+ si = self.source_downs[i](s_stft)
+ si = self.source_resblocks[i](si)
+ x = x + si
+
+ xs = None
+ for j in range(self.num_kernels):
+ if xs is None:
+ xs = self.resblocks[i * self.num_kernels + j](x)
+ else:
+ xs += self.resblocks[i * self.num_kernels + j](x)
+ x = xs / self.num_kernels
+
+ x = F.leaky_relu(x)
+ x = self.conv_post(x)
+ magnitude = torch.exp(x[:, :self.istft_params["n_fft"] // 2 + 1, :])
+ phase = torch.sin(x[:, self.istft_params["n_fft"] // 2 + 1:, :]) # actually, sin is redundancy
+
+ x = self._istft(magnitude, phase)
+ x = torch.clamp(x, -self.audio_limit, self.audio_limit)
+ return x
+
+ def remove_weight_norm(self):
+ print('Removing weight norm...')
+ for l in self.ups:
+ remove_weight_norm(l)
+ for l in self.resblocks:
+ l.remove_weight_norm()
+ remove_weight_norm(self.conv_pre)
+ remove_weight_norm(self.conv_post)
+ self.source_module.remove_weight_norm()
+ for l in self.source_downs:
+ remove_weight_norm(l)
+ for l in self.source_resblocks:
+ l.remove_weight_norm()
+
+ @torch.inference_mode()
+ def inference(self, mel: torch.Tensor, f0=None) -> torch.Tensor:
+ return self.forward(x=mel, f0=f0)
diff --git a/modules/layers.py b/modules/layers.py
new file mode 100644
index 0000000000000000000000000000000000000000..dc218d864e1b3d605220d989cd2f8f1a317ae4c3
--- /dev/null
+++ b/modules/layers.py
@@ -0,0 +1,354 @@
+import math
+import torch
+from torch import nn
+from typing import Optional, Any
+from torch import Tensor
+import torch.nn.functional as F
+import torchaudio
+import torchaudio.functional as audio_F
+
+import random
+random.seed(0)
+
+
+def _get_activation_fn(activ):
+ if activ == 'relu':
+ return nn.ReLU()
+ elif activ == 'lrelu':
+ return nn.LeakyReLU(0.2)
+ elif activ == 'swish':
+ return lambda x: x*torch.sigmoid(x)
+ else:
+ raise RuntimeError('Unexpected activ type %s, expected [relu, lrelu, swish]' % activ)
+
+class LinearNorm(torch.nn.Module):
+ def __init__(self, in_dim, out_dim, bias=True, w_init_gain='linear'):
+ super(LinearNorm, self).__init__()
+ self.linear_layer = torch.nn.Linear(in_dim, out_dim, bias=bias)
+
+ torch.nn.init.xavier_uniform_(
+ self.linear_layer.weight,
+ gain=torch.nn.init.calculate_gain(w_init_gain))
+
+ def forward(self, x):
+ return self.linear_layer(x)
+
+
+class ConvNorm(torch.nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size=1, stride=1,
+ padding=None, dilation=1, bias=True, w_init_gain='linear', param=None):
+ super(ConvNorm, self).__init__()
+ if padding is None:
+ assert(kernel_size % 2 == 1)
+ padding = int(dilation * (kernel_size - 1) / 2)
+
+ self.conv = torch.nn.Conv1d(in_channels, out_channels,
+ kernel_size=kernel_size, stride=stride,
+ padding=padding, dilation=dilation,
+ bias=bias)
+
+ torch.nn.init.xavier_uniform_(
+ self.conv.weight, gain=torch.nn.init.calculate_gain(w_init_gain, param=param))
+
+ def forward(self, signal):
+ conv_signal = self.conv(signal)
+ return conv_signal
+
+class CausualConv(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=1, dilation=1, bias=True, w_init_gain='linear', param=None):
+ super(CausualConv, self).__init__()
+ if padding is None:
+ assert(kernel_size % 2 == 1)
+ padding = int(dilation * (kernel_size - 1) / 2) * 2
+ else:
+ self.padding = padding * 2
+ self.conv = nn.Conv1d(in_channels, out_channels,
+ kernel_size=kernel_size, stride=stride,
+ padding=self.padding,
+ dilation=dilation,
+ bias=bias)
+
+ torch.nn.init.xavier_uniform_(
+ self.conv.weight, gain=torch.nn.init.calculate_gain(w_init_gain, param=param))
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = x[:, :, :-self.padding]
+ return x
+
+class CausualBlock(nn.Module):
+ def __init__(self, hidden_dim, n_conv=3, dropout_p=0.2, activ='lrelu'):
+ super(CausualBlock, self).__init__()
+ self.blocks = nn.ModuleList([
+ self._get_conv(hidden_dim, dilation=3**i, activ=activ, dropout_p=dropout_p)
+ for i in range(n_conv)])
+
+ def forward(self, x):
+ for block in self.blocks:
+ res = x
+ x = block(x)
+ x += res
+ return x
+
+ def _get_conv(self, hidden_dim, dilation, activ='lrelu', dropout_p=0.2):
+ layers = [
+ CausualConv(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
+ _get_activation_fn(activ),
+ nn.BatchNorm1d(hidden_dim),
+ nn.Dropout(p=dropout_p),
+ CausualConv(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
+ _get_activation_fn(activ),
+ nn.Dropout(p=dropout_p)
+ ]
+ return nn.Sequential(*layers)
+
+class ConvBlock(nn.Module):
+ def __init__(self, hidden_dim, n_conv=3, dropout_p=0.2, activ='relu'):
+ super().__init__()
+ self._n_groups = 8
+ self.blocks = nn.ModuleList([
+ self._get_conv(hidden_dim, dilation=3**i, activ=activ, dropout_p=dropout_p)
+ for i in range(n_conv)])
+
+
+ def forward(self, x):
+ for block in self.blocks:
+ res = x
+ x = block(x)
+ x += res
+ return x
+
+ def _get_conv(self, hidden_dim, dilation, activ='relu', dropout_p=0.2):
+ layers = [
+ ConvNorm(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
+ _get_activation_fn(activ),
+ nn.GroupNorm(num_groups=self._n_groups, num_channels=hidden_dim),
+ nn.Dropout(p=dropout_p),
+ ConvNorm(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
+ _get_activation_fn(activ),
+ nn.Dropout(p=dropout_p)
+ ]
+ return nn.Sequential(*layers)
+
+class LocationLayer(nn.Module):
+ def __init__(self, attention_n_filters, attention_kernel_size,
+ attention_dim):
+ super(LocationLayer, self).__init__()
+ padding = int((attention_kernel_size - 1) / 2)
+ self.location_conv = ConvNorm(2, attention_n_filters,
+ kernel_size=attention_kernel_size,
+ padding=padding, bias=False, stride=1,
+ dilation=1)
+ self.location_dense = LinearNorm(attention_n_filters, attention_dim,
+ bias=False, w_init_gain='tanh')
+
+ def forward(self, attention_weights_cat):
+ processed_attention = self.location_conv(attention_weights_cat)
+ processed_attention = processed_attention.transpose(1, 2)
+ processed_attention = self.location_dense(processed_attention)
+ return processed_attention
+
+
+class Attention(nn.Module):
+ def __init__(self, attention_rnn_dim, embedding_dim, attention_dim,
+ attention_location_n_filters, attention_location_kernel_size):
+ super(Attention, self).__init__()
+ self.query_layer = LinearNorm(attention_rnn_dim, attention_dim,
+ bias=False, w_init_gain='tanh')
+ self.memory_layer = LinearNorm(embedding_dim, attention_dim, bias=False,
+ w_init_gain='tanh')
+ self.v = LinearNorm(attention_dim, 1, bias=False)
+ self.location_layer = LocationLayer(attention_location_n_filters,
+ attention_location_kernel_size,
+ attention_dim)
+ self.score_mask_value = -float("inf")
+
+ def get_alignment_energies(self, query, processed_memory,
+ attention_weights_cat):
+ """
+ PARAMS
+ ------
+ query: decoder output (batch, n_mel_channels * n_frames_per_step)
+ processed_memory: processed encoder outputs (B, T_in, attention_dim)
+ attention_weights_cat: cumulative and prev. att weights (B, 2, max_time)
+ RETURNS
+ -------
+ alignment (batch, max_time)
+ """
+
+ processed_query = self.query_layer(query.unsqueeze(1))
+ processed_attention_weights = self.location_layer(attention_weights_cat)
+ energies = self.v(torch.tanh(
+ processed_query + processed_attention_weights + processed_memory))
+
+ energies = energies.squeeze(-1)
+ return energies
+
+ def forward(self, attention_hidden_state, memory, processed_memory,
+ attention_weights_cat, mask):
+ """
+ PARAMS
+ ------
+ attention_hidden_state: attention rnn last output
+ memory: encoder outputs
+ processed_memory: processed encoder outputs
+ attention_weights_cat: previous and cummulative attention weights
+ mask: binary mask for padded data
+ """
+ alignment = self.get_alignment_energies(
+ attention_hidden_state, processed_memory, attention_weights_cat)
+
+ if mask is not None:
+ alignment.data.masked_fill_(mask, self.score_mask_value)
+
+ attention_weights = F.softmax(alignment, dim=1)
+ attention_context = torch.bmm(attention_weights.unsqueeze(1), memory)
+ attention_context = attention_context.squeeze(1)
+
+ return attention_context, attention_weights
+
+
+class ForwardAttentionV2(nn.Module):
+ def __init__(self, attention_rnn_dim, embedding_dim, attention_dim,
+ attention_location_n_filters, attention_location_kernel_size):
+ super(ForwardAttentionV2, self).__init__()
+ self.query_layer = LinearNorm(attention_rnn_dim, attention_dim,
+ bias=False, w_init_gain='tanh')
+ self.memory_layer = LinearNorm(embedding_dim, attention_dim, bias=False,
+ w_init_gain='tanh')
+ self.v = LinearNorm(attention_dim, 1, bias=False)
+ self.location_layer = LocationLayer(attention_location_n_filters,
+ attention_location_kernel_size,
+ attention_dim)
+ self.score_mask_value = -float(1e20)
+
+ def get_alignment_energies(self, query, processed_memory,
+ attention_weights_cat):
+ """
+ PARAMS
+ ------
+ query: decoder output (batch, n_mel_channels * n_frames_per_step)
+ processed_memory: processed encoder outputs (B, T_in, attention_dim)
+ attention_weights_cat: prev. and cumulative att weights (B, 2, max_time)
+ RETURNS
+ -------
+ alignment (batch, max_time)
+ """
+
+ processed_query = self.query_layer(query.unsqueeze(1))
+ processed_attention_weights = self.location_layer(attention_weights_cat)
+ energies = self.v(torch.tanh(
+ processed_query + processed_attention_weights + processed_memory))
+
+ energies = energies.squeeze(-1)
+ return energies
+
+ def forward(self, attention_hidden_state, memory, processed_memory,
+ attention_weights_cat, mask, log_alpha):
+ """
+ PARAMS
+ ------
+ attention_hidden_state: attention rnn last output
+ memory: encoder outputs
+ processed_memory: processed encoder outputs
+ attention_weights_cat: previous and cummulative attention weights
+ mask: binary mask for padded data
+ """
+ log_energy = self.get_alignment_energies(
+ attention_hidden_state, processed_memory, attention_weights_cat)
+
+ #log_energy =
+
+ if mask is not None:
+ log_energy.data.masked_fill_(mask, self.score_mask_value)
+
+ #attention_weights = F.softmax(alignment, dim=1)
+
+ #content_score = log_energy.unsqueeze(1) #[B, MAX_TIME] -> [B, 1, MAX_TIME]
+ #log_alpha = log_alpha.unsqueeze(2) #[B, MAX_TIME] -> [B, MAX_TIME, 1]
+
+ #log_total_score = log_alpha + content_score
+
+ #previous_attention_weights = attention_weights_cat[:,0,:]
+
+ log_alpha_shift_padded = []
+ max_time = log_energy.size(1)
+ for sft in range(2):
+ shifted = log_alpha[:,:max_time-sft]
+ shift_padded = F.pad(shifted, (sft,0), 'constant', self.score_mask_value)
+ log_alpha_shift_padded.append(shift_padded.unsqueeze(2))
+
+ biased = torch.logsumexp(torch.cat(log_alpha_shift_padded,2), 2)
+
+ log_alpha_new = biased + log_energy
+
+ attention_weights = F.softmax(log_alpha_new, dim=1)
+
+ attention_context = torch.bmm(attention_weights.unsqueeze(1), memory)
+ attention_context = attention_context.squeeze(1)
+
+ return attention_context, attention_weights, log_alpha_new
+
+
+class PhaseShuffle2d(nn.Module):
+ def __init__(self, n=2):
+ super(PhaseShuffle2d, self).__init__()
+ self.n = n
+ self.random = random.Random(1)
+
+ def forward(self, x, move=None):
+ # x.size = (B, C, M, L)
+ if move is None:
+ move = self.random.randint(-self.n, self.n)
+
+ if move == 0:
+ return x
+ else:
+ left = x[:, :, :, :move]
+ right = x[:, :, :, move:]
+ shuffled = torch.cat([right, left], dim=3)
+ return shuffled
+
+class PhaseShuffle1d(nn.Module):
+ def __init__(self, n=2):
+ super(PhaseShuffle1d, self).__init__()
+ self.n = n
+ self.random = random.Random(1)
+
+ def forward(self, x, move=None):
+ # x.size = (B, C, M, L)
+ if move is None:
+ move = self.random.randint(-self.n, self.n)
+
+ if move == 0:
+ return x
+ else:
+ left = x[:, :, :move]
+ right = x[:, :, move:]
+ shuffled = torch.cat([right, left], dim=2)
+
+ return shuffled
+
+class MFCC(nn.Module):
+ def __init__(self, n_mfcc=40, n_mels=80):
+ super(MFCC, self).__init__()
+ self.n_mfcc = n_mfcc
+ self.n_mels = n_mels
+ self.norm = 'ortho'
+ dct_mat = audio_F.create_dct(self.n_mfcc, self.n_mels, self.norm)
+ self.register_buffer('dct_mat', dct_mat)
+
+ def forward(self, mel_specgram):
+ if len(mel_specgram.shape) == 2:
+ mel_specgram = mel_specgram.unsqueeze(0)
+ unsqueezed = True
+ else:
+ unsqueezed = False
+ # (channel, n_mels, time).tranpose(...) dot (n_mels, n_mfcc)
+ # -> (channel, time, n_mfcc).tranpose(...)
+ mfcc = torch.matmul(mel_specgram.transpose(1, 2), self.dct_mat).transpose(1, 2)
+
+ # unpack batch
+ if unsqueezed:
+ mfcc = mfcc.squeeze(0)
+ return mfcc
diff --git a/modules/length_regulator.py b/modules/length_regulator.py
new file mode 100644
index 0000000000000000000000000000000000000000..8bc875326f8b846a09fbb9602d3ebf3ba6cc3b0f
--- /dev/null
+++ b/modules/length_regulator.py
@@ -0,0 +1,141 @@
+from typing import Tuple
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from modules.commons import sequence_mask
+import numpy as np
+from dac.nn.quantize import VectorQuantize
+
+# f0_bin = 256
+f0_max = 1100.0
+f0_min = 50.0
+f0_mel_min = 1127 * np.log(1 + f0_min / 700)
+f0_mel_max = 1127 * np.log(1 + f0_max / 700)
+
+def f0_to_coarse(f0, f0_bin):
+ f0_mel = 1127 * (1 + f0 / 700).log()
+ a = (f0_bin - 2) / (f0_mel_max - f0_mel_min)
+ b = f0_mel_min * a - 1.
+ f0_mel = torch.where(f0_mel > 0, f0_mel * a - b, f0_mel)
+ # torch.clip_(f0_mel, min=1., max=float(f0_bin - 1))
+ f0_coarse = torch.round(f0_mel).long()
+ f0_coarse = f0_coarse * (f0_coarse > 0)
+ f0_coarse = f0_coarse + ((f0_coarse < 1) * 1)
+ f0_coarse = f0_coarse * (f0_coarse < f0_bin)
+ f0_coarse = f0_coarse + ((f0_coarse >= f0_bin) * (f0_bin - 1))
+ return f0_coarse
+
+class InterpolateRegulator(nn.Module):
+ def __init__(
+ self,
+ channels: int,
+ sampling_ratios: Tuple,
+ is_discrete: bool = False,
+ in_channels: int = None, # only applies to continuous input
+ vector_quantize: bool = False, # whether to use vector quantization, only applies to continuous input
+ codebook_size: int = 1024, # for discrete only
+ out_channels: int = None,
+ groups: int = 1,
+ n_codebooks: int = 1, # number of codebooks
+ quantizer_dropout: float = 0.0, # dropout for quantizer
+ f0_condition: bool = False,
+ n_f0_bins: int = 512,
+ ):
+ super().__init__()
+ self.sampling_ratios = sampling_ratios
+ out_channels = out_channels or channels
+ model = nn.ModuleList([])
+ if len(sampling_ratios) > 0:
+ self.interpolate = True
+ for _ in sampling_ratios:
+ module = nn.Conv1d(channels, channels, 3, 1, 1)
+ norm = nn.GroupNorm(groups, channels)
+ act = nn.Mish()
+ model.extend([module, norm, act])
+ else:
+ self.interpolate = False
+ model.append(
+ nn.Conv1d(channels, out_channels, 1, 1)
+ )
+ self.model = nn.Sequential(*model)
+ self.embedding = nn.Embedding(codebook_size, channels)
+ self.is_discrete = is_discrete
+
+ self.mask_token = nn.Parameter(torch.zeros(1, channels))
+
+ self.n_codebooks = n_codebooks
+ if n_codebooks > 1:
+ self.extra_codebooks = nn.ModuleList([
+ nn.Embedding(codebook_size, channels) for _ in range(n_codebooks - 1)
+ ])
+ self.extra_codebook_mask_tokens = nn.ParameterList([
+ nn.Parameter(torch.zeros(1, channels)) for _ in range(n_codebooks - 1)
+ ])
+ self.quantizer_dropout = quantizer_dropout
+
+ if f0_condition:
+ self.f0_embedding = nn.Embedding(n_f0_bins, channels)
+ self.f0_condition = f0_condition
+ self.n_f0_bins = n_f0_bins
+ self.f0_bins = torch.arange(2, 1024, 1024 // n_f0_bins)
+ self.f0_mask = nn.Parameter(torch.zeros(1, channels))
+ else:
+ self.f0_condition = False
+
+ if not is_discrete:
+ self.content_in_proj = nn.Linear(in_channels, channels)
+ if vector_quantize:
+ self.vq = VectorQuantize(channels, codebook_size, 8)
+
+ def forward(self, x, ylens=None, n_quantizers=None, f0=None):
+ # apply token drop
+ if self.training:
+ n_quantizers = torch.ones((x.shape[0],)) * self.n_codebooks
+ dropout = torch.randint(1, self.n_codebooks + 1, (x.shape[0],))
+ n_dropout = int(x.shape[0] * self.quantizer_dropout)
+ n_quantizers[:n_dropout] = dropout[:n_dropout]
+ n_quantizers = n_quantizers.to(x.device)
+ # decide whether to drop for each sample in batch
+ else:
+ n_quantizers = torch.ones((x.shape[0],), device=x.device) * (self.n_codebooks if n_quantizers is None else n_quantizers)
+ if self.is_discrete:
+ if self.n_codebooks > 1:
+ assert len(x.size()) == 3
+ x_emb = self.embedding(x[:, 0])
+ for i, emb in enumerate(self.extra_codebooks):
+ x_emb = x_emb + (n_quantizers > i+1)[..., None, None] * emb(x[:, i+1])
+ # add mask token if not using this codebook
+ # x_emb = x_emb + (n_quantizers <= i+1)[..., None, None] * self.extra_codebook_mask_tokens[i]
+ x = x_emb
+ elif self.n_codebooks == 1:
+ if len(x.size()) == 2:
+ x = self.embedding(x)
+ else:
+ x = self.embedding(x[:, 0])
+ else:
+ x = self.content_in_proj(x)
+ # x in (B, T, D)
+ mask = sequence_mask(ylens).unsqueeze(-1)
+ if self.interpolate:
+ x = F.interpolate(x.transpose(1, 2).contiguous(), size=ylens.max(), mode='nearest')
+ else:
+ x = x.transpose(1, 2).contiguous()
+ mask = mask[:, :x.size(2), :]
+ ylens = ylens.clamp(max=x.size(2)).long()
+ if self.f0_condition:
+ if f0 is None:
+ x = x + self.f0_mask.unsqueeze(-1)
+ else:
+ #quantized_f0 = torch.bucketize(f0, self.f0_bins.to(f0.device)) # (N, T)
+ quantized_f0 = f0_to_coarse(f0, self.n_f0_bins)
+ quantized_f0 = quantized_f0.clamp(0, self.n_f0_bins - 1).long()
+ f0_emb = self.f0_embedding(quantized_f0)
+ f0_emb = F.interpolate(f0_emb.transpose(1, 2).contiguous(), size=ylens.max(), mode='nearest')
+ x = x + f0_emb
+ out = self.model(x).transpose(1, 2).contiguous()
+ if hasattr(self, 'vq'):
+ out_q, commitment_loss, codebook_loss, codes, out, = self.vq(out.transpose(1, 2))
+ out_q = out_q.transpose(1, 2)
+ return out_q * mask, ylens, codes, commitment_loss, codebook_loss
+ olens = ylens
+ return out * mask, olens, None, None, None
diff --git a/modules/quantize.py b/modules/quantize.py
new file mode 100644
index 0000000000000000000000000000000000000000..1736ea33b72ce7292e068ad418e411c1b76f4b73
--- /dev/null
+++ b/modules/quantize.py
@@ -0,0 +1,229 @@
+from dac.nn.quantize import ResidualVectorQuantize
+from torch import nn
+from modules.wavenet import WN
+import torch
+import torchaudio
+import torchaudio.functional as audio_F
+import numpy as np
+from .alias_free_torch import *
+from torch.nn.utils import weight_norm
+from torch import nn, sin, pow
+from einops.layers.torch import Rearrange
+from dac.model.encodec import SConv1d
+
+def init_weights(m):
+ if isinstance(m, nn.Conv1d):
+ nn.init.trunc_normal_(m.weight, std=0.02)
+ nn.init.constant_(m.bias, 0)
+
+
+def WNConv1d(*args, **kwargs):
+ return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+ return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+class SnakeBeta(nn.Module):
+ """
+ A modified Snake function which uses separate parameters for the magnitude of the periodic components
+ Shape:
+ - Input: (B, C, T)
+ - Output: (B, C, T), same shape as the input
+ Parameters:
+ - alpha - trainable parameter that controls frequency
+ - beta - trainable parameter that controls magnitude
+ References:
+ - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+ https://arxiv.org/abs/2006.08195
+ Examples:
+ >>> a1 = snakebeta(256)
+ >>> x = torch.randn(256)
+ >>> x = a1(x)
+ """
+
+ def __init__(
+ self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+ ):
+ """
+ Initialization.
+ INPUT:
+ - in_features: shape of the input
+ - alpha - trainable parameter that controls frequency
+ - beta - trainable parameter that controls magnitude
+ alpha is initialized to 1 by default, higher values = higher-frequency.
+ beta is initialized to 1 by default, higher values = higher-magnitude.
+ alpha will be trained along with the rest of your model.
+ """
+ super(SnakeBeta, self).__init__()
+ self.in_features = in_features
+
+ # initialize alpha
+ self.alpha_logscale = alpha_logscale
+ if self.alpha_logscale: # log scale alphas initialized to zeros
+ self.alpha = nn.Parameter(torch.zeros(in_features) * alpha)
+ self.beta = nn.Parameter(torch.zeros(in_features) * alpha)
+ else: # linear scale alphas initialized to ones
+ self.alpha = nn.Parameter(torch.ones(in_features) * alpha)
+ self.beta = nn.Parameter(torch.ones(in_features) * alpha)
+
+ self.alpha.requires_grad = alpha_trainable
+ self.beta.requires_grad = alpha_trainable
+
+ self.no_div_by_zero = 0.000000001
+
+ def forward(self, x):
+ """
+ Forward pass of the function.
+ Applies the function to the input elementwise.
+ SnakeBeta := x + 1/b * sin^2 (xa)
+ """
+ alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T]
+ beta = self.beta.unsqueeze(0).unsqueeze(-1)
+ if self.alpha_logscale:
+ alpha = torch.exp(alpha)
+ beta = torch.exp(beta)
+ x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+ return x
+
+class ResidualUnit(nn.Module):
+ def __init__(self, dim: int = 16, dilation: int = 1):
+ super().__init__()
+ pad = ((7 - 1) * dilation) // 2
+ self.block = nn.Sequential(
+ Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)),
+ WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+ Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)),
+ WNConv1d(dim, dim, kernel_size=1),
+ )
+
+ def forward(self, x):
+ return x + self.block(x)
+
+class CNNLSTM(nn.Module):
+ def __init__(self, indim, outdim, head, global_pred=False):
+ super().__init__()
+ self.global_pred = global_pred
+ self.model = nn.Sequential(
+ ResidualUnit(indim, dilation=1),
+ ResidualUnit(indim, dilation=2),
+ ResidualUnit(indim, dilation=3),
+ Activation1d(activation=SnakeBeta(indim, alpha_logscale=True)),
+ Rearrange("b c t -> b t c"),
+ )
+ self.heads = nn.ModuleList([nn.Linear(indim, outdim) for i in range(head)])
+
+ def forward(self, x):
+ # x: [B, C, T]
+ x = self.model(x)
+ if self.global_pred:
+ x = torch.mean(x, dim=1, keepdim=False)
+ outs = [head(x) for head in self.heads]
+ return outs
+
+def sequence_mask(length, max_length=None):
+ if max_length is None:
+ max_length = length.max()
+ x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+ return x.unsqueeze(0) < length.unsqueeze(1)
+class FAquantizer(nn.Module):
+ def __init__(self, in_dim=1024,
+ n_p_codebooks=1,
+ n_c_codebooks=2,
+ n_t_codebooks=2,
+ n_r_codebooks=3,
+ codebook_size=1024,
+ codebook_dim=8,
+ quantizer_dropout=0.5,
+ causal=False,
+ separate_prosody_encoder=False,
+ timbre_norm=False,):
+ super(FAquantizer, self).__init__()
+ conv1d_type = SConv1d# if causal else nn.Conv1d
+ self.prosody_quantizer = ResidualVectorQuantize(
+ input_dim=in_dim,
+ n_codebooks=n_p_codebooks,
+ codebook_size=codebook_size,
+ codebook_dim=codebook_dim,
+ quantizer_dropout=quantizer_dropout,
+ )
+
+ self.content_quantizer = ResidualVectorQuantize(
+ input_dim=in_dim,
+ n_codebooks=n_c_codebooks,
+ codebook_size=codebook_size,
+ codebook_dim=codebook_dim,
+ quantizer_dropout=quantizer_dropout,
+ )
+
+ self.residual_quantizer = ResidualVectorQuantize(
+ input_dim=in_dim,
+ n_codebooks=n_r_codebooks,
+ codebook_size=codebook_size,
+ codebook_dim=codebook_dim,
+ quantizer_dropout=quantizer_dropout,
+ )
+
+ self.melspec_linear = conv1d_type(in_channels=20, out_channels=256, kernel_size=1, causal=causal)
+ self.melspec_encoder = WN(hidden_channels=256, kernel_size=5, dilation_rate=1, n_layers=8, gin_channels=0, p_dropout=0.2, causal=causal)
+ self.melspec_linear2 = conv1d_type(in_channels=256, out_channels=1024, kernel_size=1, causal=causal)
+
+ self.prob_random_mask_residual = 0.75
+
+ SPECT_PARAMS = {
+ "n_fft": 2048,
+ "win_length": 1200,
+ "hop_length": 300,
+ }
+ MEL_PARAMS = {
+ "n_mels": 80,
+ }
+
+ self.to_mel = torchaudio.transforms.MelSpectrogram(
+ n_mels=MEL_PARAMS["n_mels"], sample_rate=24000, **SPECT_PARAMS
+ )
+ self.mel_mean, self.mel_std = -4, 4
+ self.frame_rate = 24000 / 300
+ self.hop_length = 300
+
+ def preprocess(self, wave_tensor, n_bins=20):
+ mel_tensor = self.to_mel(wave_tensor.squeeze(1))
+ mel_tensor = (torch.log(1e-5 + mel_tensor) - self.mel_mean) / self.mel_std
+ return mel_tensor[:, :n_bins, :int(wave_tensor.size(-1) / self.hop_length)]
+
+ def forward(self, x, wave_segments):
+ outs = 0
+ prosody_feature = self.preprocess(wave_segments)
+
+ f0_input = prosody_feature # (B, T, 20)
+ f0_input = self.melspec_linear(f0_input)
+ f0_input = self.melspec_encoder(f0_input, torch.ones(f0_input.shape[0], 1, f0_input.shape[2]).to(
+ f0_input.device).bool())
+ f0_input = self.melspec_linear2(f0_input)
+
+ common_min_size = min(f0_input.size(2), x.size(2))
+ f0_input = f0_input[:, :, :common_min_size]
+
+ x = x[:, :, :common_min_size]
+
+ z_p, codes_p, latents_p, commitment_loss_p, codebook_loss_p = self.prosody_quantizer(
+ f0_input, 1
+ )
+ outs += z_p.detach()
+
+ z_c, codes_c, latents_c, commitment_loss_c, codebook_loss_c = self.content_quantizer(
+ x, 2
+ )
+ outs += z_c.detach()
+
+ residual_feature = x - z_p.detach() - z_c.detach()
+
+ z_r, codes_r, latents_r, commitment_loss_r, codebook_loss_r = self.residual_quantizer(
+ residual_feature, 3
+ )
+
+ quantized = [z_p, z_c, z_r]
+ codes = [codes_p, codes_c, codes_r]
+
+ return quantized, codes
\ No newline at end of file
diff --git a/modules/rmvpe.py b/modules/rmvpe.py
new file mode 100644
index 0000000000000000000000000000000000000000..44ae2e0ec9fde661dd8360d3fd731fe66b5ab51c
--- /dev/null
+++ b/modules/rmvpe.py
@@ -0,0 +1,637 @@
+from io import BytesIO
+import os
+from typing import List, Optional, Tuple
+import numpy as np
+import torch
+
+import torch.nn as nn
+import torch.nn.functional as F
+from librosa.util import normalize, pad_center, tiny
+from scipy.signal import get_window
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class STFT(torch.nn.Module):
+ def __init__(
+ self, filter_length=1024, hop_length=512, win_length=None, window="hann"
+ ):
+ """
+ This module implements an STFT using 1D convolution and 1D transpose convolutions.
+ This is a bit tricky so there are some cases that probably won't work as working
+ out the same sizes before and after in all overlap add setups is tough. Right now,
+ this code should work with hop lengths that are half the filter length (50% overlap
+ between frames).
+
+ Keyword Arguments:
+ filter_length {int} -- Length of filters used (default: {1024})
+ hop_length {int} -- Hop length of STFT (restrict to 50% overlap between frames) (default: {512})
+ win_length {[type]} -- Length of the window function applied to each frame (if not specified, it
+ equals the filter length). (default: {None})
+ window {str} -- Type of window to use (options are bartlett, hann, hamming, blackman, blackmanharris)
+ (default: {'hann'})
+ """
+ super(STFT, self).__init__()
+ self.filter_length = filter_length
+ self.hop_length = hop_length
+ self.win_length = win_length if win_length else filter_length
+ self.window = window
+ self.forward_transform = None
+ self.pad_amount = int(self.filter_length / 2)
+ fourier_basis = np.fft.fft(np.eye(self.filter_length))
+
+ cutoff = int((self.filter_length / 2 + 1))
+ fourier_basis = np.vstack(
+ [np.real(fourier_basis[:cutoff, :]), np.imag(fourier_basis[:cutoff, :])]
+ )
+ forward_basis = torch.FloatTensor(fourier_basis)
+ inverse_basis = torch.FloatTensor(np.linalg.pinv(fourier_basis))
+
+ assert filter_length >= self.win_length
+ # get window and zero center pad it to filter_length
+ fft_window = get_window(window, self.win_length, fftbins=True)
+ fft_window = pad_center(fft_window, size=filter_length)
+ fft_window = torch.from_numpy(fft_window).float()
+
+ # window the bases
+ forward_basis *= fft_window
+ inverse_basis = (inverse_basis.T * fft_window).T
+
+ self.register_buffer("forward_basis", forward_basis.float())
+ self.register_buffer("inverse_basis", inverse_basis.float())
+ self.register_buffer("fft_window", fft_window.float())
+
+ def transform(self, input_data, return_phase=False):
+ """Take input data (audio) to STFT domain.
+
+ Arguments:
+ input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)
+
+ Returns:
+ magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
+ num_frequencies, num_frames)
+ phase {tensor} -- Phase of STFT with shape (num_batch,
+ num_frequencies, num_frames)
+ """
+ input_data = F.pad(
+ input_data,
+ (self.pad_amount, self.pad_amount),
+ mode="reflect",
+ )
+ forward_transform = input_data.unfold(
+ 1, self.filter_length, self.hop_length
+ ).permute(0, 2, 1)
+ forward_transform = torch.matmul(self.forward_basis, forward_transform)
+ cutoff = int((self.filter_length / 2) + 1)
+ real_part = forward_transform[:, :cutoff, :]
+ imag_part = forward_transform[:, cutoff:, :]
+ magnitude = torch.sqrt(real_part**2 + imag_part**2)
+ if return_phase:
+ phase = torch.atan2(imag_part.data, real_part.data)
+ return magnitude, phase
+ else:
+ return magnitude
+
+ def inverse(self, magnitude, phase):
+ """Call the inverse STFT (iSTFT), given magnitude and phase tensors produced
+ by the ```transform``` function.
+
+ Arguments:
+ magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
+ num_frequencies, num_frames)
+ phase {tensor} -- Phase of STFT with shape (num_batch,
+ num_frequencies, num_frames)
+
+ Returns:
+ inverse_transform {tensor} -- Reconstructed audio given magnitude and phase. Of
+ shape (num_batch, num_samples)
+ """
+ cat = torch.cat(
+ [magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1
+ )
+ fold = torch.nn.Fold(
+ output_size=(1, (cat.size(-1) - 1) * self.hop_length + self.filter_length),
+ kernel_size=(1, self.filter_length),
+ stride=(1, self.hop_length),
+ )
+ inverse_transform = torch.matmul(self.inverse_basis, cat)
+ inverse_transform = fold(inverse_transform)[
+ :, 0, 0, self.pad_amount : -self.pad_amount
+ ]
+ window_square_sum = (
+ self.fft_window.pow(2).repeat(cat.size(-1), 1).T.unsqueeze(0)
+ )
+ window_square_sum = fold(window_square_sum)[
+ :, 0, 0, self.pad_amount : -self.pad_amount
+ ]
+ inverse_transform /= window_square_sum
+ return inverse_transform
+
+ def forward(self, input_data):
+ """Take input data (audio) to STFT domain and then back to audio.
+
+ Arguments:
+ input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)
+
+ Returns:
+ reconstruction {tensor} -- Reconstructed audio given magnitude and phase. Of
+ shape (num_batch, num_samples)
+ """
+ self.magnitude, self.phase = self.transform(input_data, return_phase=True)
+ reconstruction = self.inverse(self.magnitude, self.phase)
+ return reconstruction
+
+
+from time import time as ttime
+
+
+class BiGRU(nn.Module):
+ def __init__(self, input_features, hidden_features, num_layers):
+ super(BiGRU, self).__init__()
+ self.gru = nn.GRU(
+ input_features,
+ hidden_features,
+ num_layers=num_layers,
+ batch_first=True,
+ bidirectional=True,
+ )
+
+ def forward(self, x):
+ return self.gru(x)[0]
+
+
+class ConvBlockRes(nn.Module):
+ def __init__(self, in_channels, out_channels, momentum=0.01):
+ super(ConvBlockRes, self).__init__()
+ self.conv = nn.Sequential(
+ nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=(3, 3),
+ stride=(1, 1),
+ padding=(1, 1),
+ bias=False,
+ ),
+ nn.BatchNorm2d(out_channels, momentum=momentum),
+ nn.ReLU(),
+ nn.Conv2d(
+ in_channels=out_channels,
+ out_channels=out_channels,
+ kernel_size=(3, 3),
+ stride=(1, 1),
+ padding=(1, 1),
+ bias=False,
+ ),
+ nn.BatchNorm2d(out_channels, momentum=momentum),
+ nn.ReLU(),
+ )
+ # self.shortcut:Optional[nn.Module] = None
+ if in_channels != out_channels:
+ self.shortcut = nn.Conv2d(in_channels, out_channels, (1, 1))
+
+ def forward(self, x: torch.Tensor):
+ if not hasattr(self, "shortcut"):
+ return self.conv(x) + x
+ else:
+ return self.conv(x) + self.shortcut(x)
+
+
+class Encoder(nn.Module):
+ def __init__(
+ self,
+ in_channels,
+ in_size,
+ n_encoders,
+ kernel_size,
+ n_blocks,
+ out_channels=16,
+ momentum=0.01,
+ ):
+ super(Encoder, self).__init__()
+ self.n_encoders = n_encoders
+ self.bn = nn.BatchNorm2d(in_channels, momentum=momentum)
+ self.layers = nn.ModuleList()
+ self.latent_channels = []
+ for i in range(self.n_encoders):
+ self.layers.append(
+ ResEncoderBlock(
+ in_channels, out_channels, kernel_size, n_blocks, momentum=momentum
+ )
+ )
+ self.latent_channels.append([out_channels, in_size])
+ in_channels = out_channels
+ out_channels *= 2
+ in_size //= 2
+ self.out_size = in_size
+ self.out_channel = out_channels
+
+ def forward(self, x: torch.Tensor):
+ concat_tensors: List[torch.Tensor] = []
+ x = self.bn(x)
+ for i, layer in enumerate(self.layers):
+ t, x = layer(x)
+ concat_tensors.append(t)
+ return x, concat_tensors
+
+
+class ResEncoderBlock(nn.Module):
+ def __init__(
+ self, in_channels, out_channels, kernel_size, n_blocks=1, momentum=0.01
+ ):
+ super(ResEncoderBlock, self).__init__()
+ self.n_blocks = n_blocks
+ self.conv = nn.ModuleList()
+ self.conv.append(ConvBlockRes(in_channels, out_channels, momentum))
+ for i in range(n_blocks - 1):
+ self.conv.append(ConvBlockRes(out_channels, out_channels, momentum))
+ self.kernel_size = kernel_size
+ if self.kernel_size is not None:
+ self.pool = nn.AvgPool2d(kernel_size=kernel_size)
+
+ def forward(self, x):
+ for i, conv in enumerate(self.conv):
+ x = conv(x)
+ if self.kernel_size is not None:
+ return x, self.pool(x)
+ else:
+ return x
+
+
+class Intermediate(nn.Module): #
+ def __init__(self, in_channels, out_channels, n_inters, n_blocks, momentum=0.01):
+ super(Intermediate, self).__init__()
+ self.n_inters = n_inters
+ self.layers = nn.ModuleList()
+ self.layers.append(
+ ResEncoderBlock(in_channels, out_channels, None, n_blocks, momentum)
+ )
+ for i in range(self.n_inters - 1):
+ self.layers.append(
+ ResEncoderBlock(out_channels, out_channels, None, n_blocks, momentum)
+ )
+
+ def forward(self, x):
+ for i, layer in enumerate(self.layers):
+ x = layer(x)
+ return x
+
+
+class ResDecoderBlock(nn.Module):
+ def __init__(self, in_channels, out_channels, stride, n_blocks=1, momentum=0.01):
+ super(ResDecoderBlock, self).__init__()
+ out_padding = (0, 1) if stride == (1, 2) else (1, 1)
+ self.n_blocks = n_blocks
+ self.conv1 = nn.Sequential(
+ nn.ConvTranspose2d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=(3, 3),
+ stride=stride,
+ padding=(1, 1),
+ output_padding=out_padding,
+ bias=False,
+ ),
+ nn.BatchNorm2d(out_channels, momentum=momentum),
+ nn.ReLU(),
+ )
+ self.conv2 = nn.ModuleList()
+ self.conv2.append(ConvBlockRes(out_channels * 2, out_channels, momentum))
+ for i in range(n_blocks - 1):
+ self.conv2.append(ConvBlockRes(out_channels, out_channels, momentum))
+
+ def forward(self, x, concat_tensor):
+ x = self.conv1(x)
+ x = torch.cat((x, concat_tensor), dim=1)
+ for i, conv2 in enumerate(self.conv2):
+ x = conv2(x)
+ return x
+
+
+class Decoder(nn.Module):
+ def __init__(self, in_channels, n_decoders, stride, n_blocks, momentum=0.01):
+ super(Decoder, self).__init__()
+ self.layers = nn.ModuleList()
+ self.n_decoders = n_decoders
+ for i in range(self.n_decoders):
+ out_channels = in_channels // 2
+ self.layers.append(
+ ResDecoderBlock(in_channels, out_channels, stride, n_blocks, momentum)
+ )
+ in_channels = out_channels
+
+ def forward(self, x: torch.Tensor, concat_tensors: List[torch.Tensor]):
+ for i, layer in enumerate(self.layers):
+ x = layer(x, concat_tensors[-1 - i])
+ return x
+
+
+class DeepUnet(nn.Module):
+ def __init__(
+ self,
+ kernel_size,
+ n_blocks,
+ en_de_layers=5,
+ inter_layers=4,
+ in_channels=1,
+ en_out_channels=16,
+ ):
+ super(DeepUnet, self).__init__()
+ self.encoder = Encoder(
+ in_channels, 128, en_de_layers, kernel_size, n_blocks, en_out_channels
+ )
+ self.intermediate = Intermediate(
+ self.encoder.out_channel // 2,
+ self.encoder.out_channel,
+ inter_layers,
+ n_blocks,
+ )
+ self.decoder = Decoder(
+ self.encoder.out_channel, en_de_layers, kernel_size, n_blocks
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x, concat_tensors = self.encoder(x)
+ x = self.intermediate(x)
+ x = self.decoder(x, concat_tensors)
+ return x
+
+
+class E2E(nn.Module):
+ def __init__(
+ self,
+ n_blocks,
+ n_gru,
+ kernel_size,
+ en_de_layers=5,
+ inter_layers=4,
+ in_channels=1,
+ en_out_channels=16,
+ ):
+ super(E2E, self).__init__()
+ self.unet = DeepUnet(
+ kernel_size,
+ n_blocks,
+ en_de_layers,
+ inter_layers,
+ in_channels,
+ en_out_channels,
+ )
+ self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1))
+ if n_gru:
+ self.fc = nn.Sequential(
+ BiGRU(3 * 128, 256, n_gru),
+ nn.Linear(512, 360),
+ nn.Dropout(0.25),
+ nn.Sigmoid(),
+ )
+ else:
+ self.fc = nn.Sequential(
+ nn.Linear(3 * nn.N_MELS, nn.N_CLASS), nn.Dropout(0.25), nn.Sigmoid()
+ )
+
+ def forward(self, mel):
+ # print(mel.shape)
+ mel = mel.transpose(-1, -2).unsqueeze(1)
+ x = self.cnn(self.unet(mel)).transpose(1, 2).flatten(-2)
+ x = self.fc(x)
+ # print(x.shape)
+ return x
+
+
+from librosa.filters import mel
+
+
+class MelSpectrogram(torch.nn.Module):
+ def __init__(
+ self,
+ is_half,
+ n_mel_channels,
+ sampling_rate,
+ win_length,
+ hop_length,
+ n_fft=None,
+ mel_fmin=0,
+ mel_fmax=None,
+ clamp=1e-5,
+ ):
+ super().__init__()
+ n_fft = win_length if n_fft is None else n_fft
+ self.hann_window = {}
+ mel_basis = mel(
+ sr=sampling_rate,
+ n_fft=n_fft,
+ n_mels=n_mel_channels,
+ fmin=mel_fmin,
+ fmax=mel_fmax,
+ htk=True,
+ )
+ mel_basis = torch.from_numpy(mel_basis).float()
+ self.register_buffer("mel_basis", mel_basis)
+ self.n_fft = win_length if n_fft is None else n_fft
+ self.hop_length = hop_length
+ self.win_length = win_length
+ self.sampling_rate = sampling_rate
+ self.n_mel_channels = n_mel_channels
+ self.clamp = clamp
+ self.is_half = is_half
+
+ def forward(self, audio, keyshift=0, speed=1, center=True):
+ factor = 2 ** (keyshift / 12)
+ n_fft_new = int(np.round(self.n_fft * factor))
+ win_length_new = int(np.round(self.win_length * factor))
+ hop_length_new = int(np.round(self.hop_length * speed))
+ keyshift_key = str(keyshift) + "_" + str(audio.device)
+ if keyshift_key not in self.hann_window:
+ self.hann_window[keyshift_key] = torch.hann_window(win_length_new).to(
+ audio.device
+ )
+ if "privateuseone" in str(audio.device):
+ if not hasattr(self, "stft"):
+ self.stft = STFT(
+ filter_length=n_fft_new,
+ hop_length=hop_length_new,
+ win_length=win_length_new,
+ window="hann",
+ ).to(audio.device)
+ magnitude = self.stft.transform(audio)
+ else:
+ fft = torch.stft(
+ audio,
+ n_fft=n_fft_new,
+ hop_length=hop_length_new,
+ win_length=win_length_new,
+ window=self.hann_window[keyshift_key],
+ center=center,
+ return_complex=True,
+ )
+ magnitude = torch.sqrt(fft.real.pow(2) + fft.imag.pow(2))
+ if keyshift != 0:
+ size = self.n_fft // 2 + 1
+ resize = magnitude.size(1)
+ if resize < size:
+ magnitude = F.pad(magnitude, (0, 0, 0, size - resize))
+ magnitude = magnitude[:, :size, :] * self.win_length / win_length_new
+ mel_output = torch.matmul(self.mel_basis, magnitude)
+ if self.is_half == True:
+ mel_output = mel_output.half()
+ log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
+ return log_mel_spec
+
+
+class RMVPE:
+ def __init__(self, model_path: str, is_half, device=None, use_jit=False):
+ self.resample_kernel = {}
+ self.resample_kernel = {}
+ self.is_half = is_half
+ if device is None:
+ #device = "cuda:0" if torch.cuda.is_available() else "cpu"
+ if torch.cuda.is_available():
+ device = "cuda:0"
+ elif torch.backends.mps.is_available():
+ device = "mps"
+ else:
+ device = "cpu"
+ self.device = device
+ self.mel_extractor = MelSpectrogram(
+ is_half, 128, 16000, 1024, 160, None, 30, 8000
+ ).to(device)
+ if "privateuseone" in str(device):
+ import onnxruntime as ort
+
+ ort_session = ort.InferenceSession(
+ "%s/rmvpe.onnx" % os.environ["rmvpe_root"],
+ providers=["DmlExecutionProvider"],
+ )
+ self.model = ort_session
+ else:
+ if str(self.device) == "cuda":
+ self.device = torch.device("cuda:0")
+
+ def get_default_model():
+ model = E2E(4, 1, (2, 2))
+ ckpt = torch.load(model_path, map_location="cpu")
+ model.load_state_dict(ckpt)
+ model.eval()
+ if is_half:
+ model = model.half()
+ else:
+ model = model.float()
+ return model
+
+ self.model = get_default_model()
+
+ self.model = self.model.to(device)
+ cents_mapping = 20 * np.arange(360) + 1997.3794084376191
+ self.cents_mapping = np.pad(cents_mapping, (4, 4)) # 368
+
+ def mel2hidden(self, mel):
+ with torch.no_grad():
+ n_frames = mel.shape[-1]
+ n_pad = 32 * ((n_frames - 1) // 32 + 1) - n_frames
+ if n_pad > 0:
+ mel = F.pad(mel, (0, n_pad), mode="constant")
+ if "privateuseone" in str(self.device):
+ onnx_input_name = self.model.get_inputs()[0].name
+ onnx_outputs_names = self.model.get_outputs()[0].name
+ hidden = self.model.run(
+ [onnx_outputs_names],
+ input_feed={onnx_input_name: mel.cpu().numpy()},
+ )[0]
+ else:
+ mel = mel.half() if self.is_half else mel.float()
+ hidden = self.model(mel)
+ return hidden[:, :n_frames]
+
+ def decode(self, hidden, thred=0.03):
+ cents_pred = self.to_local_average_cents(hidden, thred=thred)
+ f0 = 10 * (2 ** (cents_pred / 1200))
+ f0[f0 == 10] = 0
+ # f0 = np.array([10 * (2 ** (cent_pred / 1200)) if cent_pred else 0 for cent_pred in cents_pred])
+ return f0
+
+ def infer_from_audio(self, audio, thred=0.03):
+ # torch.cuda.synchronize()
+ # t0 = ttime()
+ if not torch.is_tensor(audio):
+ audio = torch.from_numpy(audio)
+ mel = self.mel_extractor(
+ audio.float().to(self.device).unsqueeze(0), center=True
+ )
+ # print(123123123,mel.device.type)
+ # torch.cuda.synchronize()
+ # t1 = ttime()
+ hidden = self.mel2hidden(mel)
+ # torch.cuda.synchronize()
+ # t2 = ttime()
+ # print(234234,hidden.device.type)
+ if "privateuseone" not in str(self.device):
+ hidden = hidden.squeeze(0).cpu().numpy()
+ else:
+ hidden = hidden[0]
+ if self.is_half == True:
+ hidden = hidden.astype("float32")
+
+ f0 = self.decode(hidden, thred=thred)
+ # torch.cuda.synchronize()
+ # t3 = ttime()
+ # print("hmvpe:%s\t%s\t%s\t%s"%(t1-t0,t2-t1,t3-t2,t3-t0))
+ return f0
+ def infer_from_audio_batch(self, audio, thred=0.03):
+ # torch.cuda.synchronize()
+ # t0 = ttime()
+ if not torch.is_tensor(audio):
+ audio = torch.from_numpy(audio)
+ mel = self.mel_extractor(
+ audio.float().to(self.device), center=True
+ )
+ # print(123123123,mel.device.type)
+ # torch.cuda.synchronize()
+ # t1 = ttime()
+ hidden = self.mel2hidden(mel)
+ # torch.cuda.synchronize()
+ # t2 = ttime()
+ # print(234234,hidden.device.type)
+ if "privateuseone" not in str(self.device):
+ hidden = hidden.cpu().numpy()
+ else:
+ pass
+ if self.is_half == True:
+ hidden = hidden.astype("float32")
+
+ f0s = []
+ for bib in range(hidden.shape[0]):
+ f0s.append(self.decode(hidden[bib], thred=thred))
+ f0s = np.stack(f0s)
+ f0s = torch.from_numpy(f0s).to(self.device)
+ # torch.cuda.synchronize()
+ # t3 = ttime()
+ # print("hmvpe:%s\t%s\t%s\t%s"%(t1-t0,t2-t1,t3-t2,t3-t0))
+ return f0s
+
+ def to_local_average_cents(self, salience, thred=0.05):
+ # t0 = ttime()
+ center = np.argmax(salience, axis=1) # 帧长#index
+ salience = np.pad(salience, ((0, 0), (4, 4))) # 帧长,368
+ # t1 = ttime()
+ center += 4
+ todo_salience = []
+ todo_cents_mapping = []
+ starts = center - 4
+ ends = center + 5
+ for idx in range(salience.shape[0]):
+ todo_salience.append(salience[:, starts[idx] : ends[idx]][idx])
+ todo_cents_mapping.append(self.cents_mapping[starts[idx] : ends[idx]])
+ # t2 = ttime()
+ todo_salience = np.array(todo_salience) # 帧长，9
+ todo_cents_mapping = np.array(todo_cents_mapping) # 帧长，9
+ product_sum = np.sum(todo_salience * todo_cents_mapping, 1)
+ weight_sum = np.sum(todo_salience, 1) # 帧长
+ devided = product_sum / weight_sum # 帧长
+ # t3 = ttime()
+ maxx = np.max(salience, axis=1) # 帧长
+ devided[maxx <= thred] = 0
+ # t4 = ttime()
+ # print("decode:%s\t%s\t%s\t%s" % (t1 - t0, t2 - t1, t3 - t2, t4 - t3))
+ return devided
diff --git a/modules/v2/__pycache__/ar.cpython-310.pyc b/modules/v2/__pycache__/ar.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1f4040e578018cb60779db0d90fe69fa71213122
Binary files /dev/null and b/modules/v2/__pycache__/ar.cpython-310.pyc differ
diff --git a/modules/v2/__pycache__/cfm.cpython-310.pyc b/modules/v2/__pycache__/cfm.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3c5f8bf0bdc3fb453ccee02ae8100b233cbad7f3
Binary files /dev/null and b/modules/v2/__pycache__/cfm.cpython-310.pyc differ
diff --git a/modules/v2/__pycache__/dit_model.cpython-310.pyc b/modules/v2/__pycache__/dit_model.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..087a9b9b4eb77693d2a898719b3dcb524bc29936
Binary files /dev/null and b/modules/v2/__pycache__/dit_model.cpython-310.pyc differ
diff --git a/modules/v2/__pycache__/dit_wrapper.cpython-310.pyc b/modules/v2/__pycache__/dit_wrapper.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5dc659c923bab9b1b759e45ca6707a6931df54b7
Binary files /dev/null and b/modules/v2/__pycache__/dit_wrapper.cpython-310.pyc differ
diff --git a/modules/v2/__pycache__/length_regulator.cpython-310.pyc b/modules/v2/__pycache__/length_regulator.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c0fc08b80f0a275c356f64325126b5c2ded4a781
Binary files /dev/null and b/modules/v2/__pycache__/length_regulator.cpython-310.pyc differ
diff --git a/modules/v2/__pycache__/model.cpython-310.pyc b/modules/v2/__pycache__/model.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..202b549ed99cfcf71724feae469c648534d13b30
Binary files /dev/null and b/modules/v2/__pycache__/model.cpython-310.pyc differ
diff --git a/modules/v2/__pycache__/vc_wrapper.cpython-310.pyc b/modules/v2/__pycache__/vc_wrapper.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5a0c6c759ae771550dd461dbf9de063ac8ab9914
Binary files /dev/null and b/modules/v2/__pycache__/vc_wrapper.cpython-310.pyc differ
diff --git a/modules/v2/ar.py b/modules/v2/ar.py
new file mode 100644
index 0000000000000000000000000000000000000000..1c38b99e14c0aa93bf4a908c231437904886ad8e
--- /dev/null
+++ b/modules/v2/ar.py
@@ -0,0 +1,763 @@
+import dataclasses
+import json
+import math
+from collections import OrderedDict
+from functools import partial, wraps
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Tuple, List
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch import Tensor
+from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
+
+
+def find_multiple(n: int, k: int) -> int:
+ if n % k == 0:
+ return n
+ return n + k - (n % k)
+
+def l2norm(t, groups = 1):
+ t = rearrange(t, '... (g d) -> ... g d', g = groups)
+ t = F.normalize(t, p = 2, dim = -1)
+ return rearrange(t, '... g d -> ... (g d)')
+
+@dataclass
+class BaseModelArgs:
+ model_type: str = "base"
+
+ vocab_size: int = 32000
+ n_layer: int = 32
+ n_head: int = 32
+ dim: int = 4096
+ intermediate_size: int = None
+ n_local_heads: int = -1
+ head_dim: int = 64
+ rope_base: float = 10000
+ norm_eps: float = 1e-5
+ max_seq_len: int = 4096
+ dropout: float = 0.0
+ tie_word_embeddings: bool = True
+ attention_qkv_bias: bool = False
+
+ # Gradient checkpointing
+ use_gradient_checkpointing: bool = False
+
+ # Initialize the model
+ initializer_range: float = 0.02
+
+ qk_norm: bool = False
+ layerscale: bool = False
+
+ def __post_init__(self):
+ if self.n_local_heads == -1:
+ self.n_local_heads = self.n_head
+ if self.intermediate_size is None:
+ hidden_dim = 4 * self.dim
+ n_hidden = int(2 * hidden_dim / 3)
+ self.intermediate_size = find_multiple(n_hidden, 256)
+ self.head_dim = self.dim // self.n_head
+
+ def save(self, path: str):
+ with open(path, "w") as f:
+ json.dump(self.__dict__, f, indent=4, sort_keys=True, ensure_ascii=False)
+
+
+@dataclass
+class NaiveModelArgs(BaseModelArgs):
+ model_type: str = "naive"
+
+
+class KVCache(nn.Module):
+ def __init__(
+ self, max_batch_size, max_seq_len, n_heads, head_dim, dtype=torch.bfloat16
+ ):
+ super().__init__()
+ cache_shape = (max_batch_size, n_heads, max_seq_len, head_dim)
+ self.register_buffer("k_cache", torch.zeros(cache_shape, dtype=dtype))
+ self.register_buffer("v_cache", torch.zeros(cache_shape, dtype=dtype))
+
+ def update(self, input_pos, k_val, v_val):
+ # input_pos: [S], k_val: [B, H, S, D]
+ assert input_pos.shape[0] == k_val.shape[2]
+
+ k_out = self.k_cache
+ v_out = self.v_cache
+ k_out[:, :, input_pos] = k_val
+ v_out[:, :, input_pos] = v_val
+
+ return k_out, v_out
+
+
+@dataclass
+class TransformerForwardResult:
+ token_logits: Tensor
+ token_targets: Tensor
+
+
+@dataclass
+class BaseTransformerForwardResult:
+ logits: Tensor
+ hidden_states: Tensor
+
+
+class BaseTransformer(nn.Module):
+ def __init__(
+ self,
+ config: BaseModelArgs,
+ init_weights: bool = True,
+ ) -> None:
+ super().__init__()
+ self.config = config
+
+ # Slow transformer
+ self.embeddings = nn.Embedding(
+ config.vocab_size,
+ config.dim,
+ )
+ self.layers = nn.ModuleList(
+ TransformerBlock(config, use_sdpa=True) for _ in range(config.n_layer)
+ )
+ self.norm = RMSNorm(config.dim, eps=config.norm_eps)
+
+ if self.config.tie_word_embeddings is False:
+ self.output = nn.Linear(
+ config.dim,
+ config.vocab_size,
+ bias=False,
+ )
+
+ self.register_buffer(
+ "freqs_cis",
+ precompute_freqs_cis(
+ config.max_seq_len,
+ config.dim // config.n_head,
+ config.rope_base,
+ ),
+ persistent=False,
+ )
+ self.register_buffer(
+ "causal_mask",
+ torch.tril(
+ torch.ones(
+ config.max_seq_len,
+ config.max_seq_len,
+ dtype=torch.bool,
+ )
+ ),
+ persistent=False,
+ )
+
+ self.output = nn.Linear(
+ config.dim,
+ config.vocab_size,
+ bias=False,
+ )
+
+ # For kv cache
+ self.max_batch_size = -1
+ self.max_seq_len = -1
+
+ if init_weights:
+ self.apply(self._init_weights)
+
+ def setup_caches(
+ self, max_batch_size: int, max_seq_len: int, dtype: torch.dtype = torch.bfloat16, device: torch.device = "cuda"
+ ):
+ if self.max_seq_len >= max_seq_len and self.max_batch_size >= max_batch_size:
+ return
+
+ head_dim = self.config.dim // self.config.n_head
+ max_seq_len = find_multiple(max_seq_len, 8)
+ self.max_seq_len = max_seq_len
+ self.max_batch_size = max_batch_size
+
+ for b in self.layers:
+ b.attention.kv_cache = KVCache(
+ max_batch_size,
+ max_seq_len,
+ self.config.n_local_heads,
+ head_dim,
+ dtype=dtype,
+ ).to(device)
+
+ def embed_base(self, x: Tensor, x_lens: Tensor) -> Tensor:
+ for bib in range(x.size(0)):
+ x[bib, x_lens[bib]:] = self.config.vocab_size - 1
+
+ x_emb = self.embeddings(x)
+ return x, x_emb
+
+ def forward(
+ self,
+ inp: Tensor,
+ key_padding_mask: Optional[Tensor] = None,
+ input_pos: Optional[Tensor] = None,
+ ) -> BaseTransformerForwardResult:
+ seq_len = inp.size(1)
+
+ # Here we want to merge the embeddings of the codebooks
+ # x = self.embed(inp)
+ x = inp.clone()
+
+ if input_pos is None:
+ freqs_cis = self.freqs_cis[:seq_len].repeat(inp.size(0), 1, 1, 1)
+ else:
+ freqs_cis = self.freqs_cis[input_pos]
+
+ # Not that the causal mask here follows the definition of scaled_dot_product_attention
+ # That is, FALSE means masked out
+ # To maintain consistency, key_padding_mask use TRUE to mask out
+ mask = None
+ if key_padding_mask is not None:
+ mask = self.causal_mask[None, None, :seq_len, :seq_len] # (B, N, Q, K)
+ mask = mask & key_padding_mask[:, None, None, :].logical_not()
+
+ for layer in self.layers:
+ if self.config.use_gradient_checkpointing and self.training:
+ x = checkpoint(layer, x, freqs_cis, mask, use_reentrant=True)
+ else:
+ x = layer(x, freqs_cis, mask)
+
+ # We got slow_out here
+ slow_out = self.norm(x)
+
+ if self.config.tie_word_embeddings:
+ token_logits = F.linear(slow_out, self.embeddings.weight)
+ else:
+ token_logits = self.output(slow_out)
+
+ return BaseTransformerForwardResult(
+ logits=token_logits,
+ hidden_states=x,
+ )
+
+ def forward_generate(
+ self,
+ inp: Tensor,
+ input_pos: Optional[Tensor] = None,
+ kv_pos: Optional[Tensor] = None,
+ return_all: bool = False,
+ ) -> BaseTransformerForwardResult:
+ # This is used for generation, optimized for torch compile
+
+ x = inp
+ max_seq_len = self.max_seq_len
+
+ mask = self.causal_mask[None, None, kv_pos, :max_seq_len] # (B, N, Q, K)
+ freqs_cis = self.freqs_cis[input_pos]
+
+ for layer in self.layers:
+ x = layer(x, freqs_cis, mask, input_pos=kv_pos)
+
+ x = x[:, -1:]
+
+ # We got slow_out here
+ slow_out = self.norm(x)
+
+ token_logits = self.output(slow_out)
+
+ return BaseTransformerForwardResult(
+ logits=token_logits,
+ hidden_states=x,
+ )
+
+ def _init_weights(self, module):
+ std = self.config.initializer_range
+ if isinstance(module, nn.Linear):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ elif isinstance(module, nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+
+class NaiveTransformer(BaseTransformer):
+ def __init__(self, config: NaiveModelArgs) -> None:
+ super().__init__(config, init_weights=False)
+ self.apply(self._init_weights)
+
+ def forward(
+ self,
+ inp: Tensor,
+ cond_lens: Tensor,
+ target: Tensor,
+ target_lens: Tensor,
+ key_padding_mask: Optional[Tensor] = None,
+ input_pos: Optional[Tensor] = None,
+ ) -> TransformerForwardResult:
+ parent_result = super().forward(
+ inp=inp,
+ key_padding_mask=key_padding_mask,
+ input_pos=input_pos,
+ )
+ token_logits = parent_result.logits
+
+ # construct targets for token_logits
+ token_targets = torch.zeros(token_logits.size(0), token_logits.size(1), dtype=torch.long,
+ device=target.device) - 100
+ for bib in range(token_targets.size(0)):
+ token_targets[bib, cond_lens[bib] + 1:cond_lens[bib] + target_lens[bib] + 1] = target[bib, :target_lens[bib]]
+ token_targets[bib, cond_lens[bib] + target_lens[bib] + 1] = self.config.vocab_size - 1
+ return TransformerForwardResult(
+ token_logits=token_logits,
+ token_targets=token_targets,
+ )
+
+ def infer_slow(self, inp: Tensor, input_pos: Optional[Tensor] = None):
+ # no kv cache used
+ parent_result = super().forward(inp, input_pos=input_pos)
+ latent = parent_result.hidden_states[:, -1]
+ base_logits = parent_result.logits[:, -1]
+ base_sampled, _ = topk_sampling(base_logits, top_k=-1, top_p=1.0)
+ return base_sampled
+
+ def forward_generate(
+ self,
+ x: Tensor,
+ input_pos: Optional[Tensor] = None,
+ kv_pos: Optional[Tensor] = None,
+ vq_masks: Optional[Tensor] = None,
+ ) -> TransformerForwardResult:
+ x = super().forward_generate(x, input_pos, kv_pos, vq_masks)
+ return x
+
+class NaiveWrapper(nn.Module):
+ def __init__(self, model: NaiveTransformer) -> None:
+ super().__init__()
+ self.model = model
+ self.sep_token_emb = nn.Parameter(torch.randn(model.config.dim))
+
+ def setup_caches(self, max_batch_size: int, max_seq_len: int, dtype: torch.dtype = torch.bfloat16, device: torch.device = "cuda"):
+ self.model.setup_caches(max_batch_size, max_seq_len, dtype, device)
+
+ def forward(self, cond: Tensor, cond_lens: Tensor, x: Tensor, x_lens: Tensor) -> torch.Tensor:
+ # style_emb = self.style_in(style).unsqueeze(1) # [B, 1, D]
+ sep_token_emb = self.sep_token_emb.expand(x.size(0), 1, -1)
+ _, x_emb = self.model.embed_base(x, x_lens)
+ emb_seq_list = []
+ for i in range(x.size(0)):
+ emb_seq = torch.cat([
+ sep_token_emb[i:i + 1],
+ cond[i:i+1, :cond_lens[i]],
+ sep_token_emb[i:i+1],
+ x_emb[i:i+1, :x_lens[i]]], dim=1)
+ emb_seq_list.append(emb_seq)
+ max_len = max([emb_seq.size(1) for emb_seq in emb_seq_list])
+ emb_seq = torch.cat([
+ F.pad(emb_seq, (0, 0, 0, max_len - emb_seq.size(1)), value=0)
+ for emb_seq in emb_seq_list
+ ], dim=0)
+ # input_pos = torch.arange(emb_seq.size(1), device=emb_seq.device).repeat(emb_seq.size(0), 1)
+ input_pos = torch.zeros(emb_seq.size(0), emb_seq.size(1), device=emb_seq.device, dtype=torch.long)
+ for i in range(x.size(0)):
+ input_pos[i, :cond_lens[i] + 1] = torch.arange(cond_lens[i] + 1, device=emb_seq.device)
+ input_pos[i, cond_lens[i] + 1: cond_lens[i] + x_lens[i] + 2] = torch.arange(x_lens[i] + 1, device=emb_seq.device)
+ out = self.model(emb_seq, cond_lens, x, x_lens, input_pos=input_pos)
+ loss = F.cross_entropy(out.token_logits.transpose(1, 2), out.token_targets.long(), ignore_index=-100)
+ return loss
+
+ @torch.no_grad()
+ def infer(self, cond: Tensor) -> torch.Tensor:
+ sep_token_emb = self.sep_token_emb.expand(1, 1, -1)
+ emb_seq = torch.cat([sep_token_emb, cond, sep_token_emb], dim=1)
+ pred_codes = []
+ input_pos = torch.arange(cond.size(1) + 1, device=cond.device)
+ for i in tqdm(range(4000)):
+ input_pos = torch.cat([input_pos, torch.LongTensor([i]).to(cond.device)], dim=0)
+ base = self.model.infer_slow(emb_seq, input_pos)
+ if base == self.model.config.vocab_size - 1:
+ break
+ new_emb = self.model.embed_base(base, torch.LongTensor([1]).to(base.device))[1]
+ emb_seq = torch.cat([emb_seq, new_emb], dim=1)
+ pred_codes.append(base)
+ return torch.cat(pred_codes, dim=-1)
+
+ @torch.no_grad()
+ def generate(
+ self,
+ prompt_text,
+ prompt_target,
+ compiled_decode_fn = None,
+ **sampling_kwargs,
+ ):
+ sep_token_emb = self.sep_token_emb.expand(1, 1, -1)
+ emb_seq = torch.cat([sep_token_emb, prompt_text, sep_token_emb], dim=1)
+ input_pos = torch.arange(prompt_text.size(1) + 1, device=emb_seq.device)
+ input_pos = torch.cat([input_pos, torch.LongTensor([0]).to(emb_seq.device)])
+ prompt_target_emb = self.model.embed_base(prompt_target,torch.LongTensor([prompt_target.size(1)]).to(prompt_target.device))[1]
+ emb_seq = torch.cat([emb_seq, prompt_target_emb], dim=1)
+ input_pos = torch.cat([input_pos, torch.arange(prompt_target_emb.size(1)).to(input_pos.device) + 1])
+
+ pred_codes = []
+ kv_pos = torch.arange(emb_seq.size(1), device=emb_seq.device)
+ next_tokens = self.decode_one_token_ar(emb_seq, input_pos, kv_pos, suppress_tokens=[self.model.config.vocab_size - 1], **sampling_kwargs)
+ pred_base = next_tokens[0]
+ pred_codes.append(pred_base)
+ new_emb = self.model.embed_base(pred_base.unsqueeze(0), torch.LongTensor([1]).to(pred_base.device))[1]
+ emb_seq = torch.cat([emb_seq, new_emb], dim=1)
+ for _ in tqdm(range(4000)):
+ suppress_eos = len(pred_codes) < 10
+ input_pos = input_pos[-1:] + 1
+ kv_pos = kv_pos[-1:] + 1
+ next_tokens = self.decode_one_token_ar(
+ emb_seq[:, -1:].reshape(1, 1, -1),
+ input_pos.reshape(1),
+ kv_pos.reshape(1),
+ previous_tokens=torch.cat(pred_codes),
+ suppress_tokens=[self.model.config.vocab_size - 1] if suppress_eos else None,
+ compiled_decode_fn=compiled_decode_fn,
+ **sampling_kwargs)
+ pred_base = next_tokens[0]
+ if pred_base == self.model.config.vocab_size - 1:
+ break
+ pred_codes.append(pred_base.clone())
+ new_emb = self.model.embed_base(pred_base.unsqueeze(0), torch.LongTensor([1]).to(pred_base.device))[1]
+ emb_seq = torch.cat([emb_seq, new_emb], dim=1)
+ return torch.stack(pred_codes, dim=-1)
+
+ def decode_one_token_ar(
+ self,
+ x: torch.Tensor,
+ input_pos: torch.Tensor,
+ kv_pos: torch.Tensor,
+ previous_tokens: torch.Tensor = None,
+ compiled_decode_fn = None,
+ **sampling_kwargs,
+ ) -> torch.Tensor:
+ if compiled_decode_fn is not None:
+ x = compiled_decode_fn(x, input_pos, kv_pos)
+ else:
+ x = self.model.forward_generate(x, input_pos, kv_pos)
+
+ sampling_kwargs_main = sampling_kwargs.copy()
+ codebooks = [
+ sample(
+ x.logits,
+ previous_tokens=(
+ previous_tokens[0] if previous_tokens is not None else None
+ ),
+ **sampling_kwargs_main,
+ )[0]
+ ]
+ codebooks = torch.stack(codebooks, dim=0)
+ return codebooks
+
+class TransformerBlock(nn.Module):
+ def __init__(self, config: BaseModelArgs, use_sdpa: bool = True) -> None:
+ super().__init__()
+ self.attention = Attention(config, use_sdpa=use_sdpa)
+ self.feed_forward = FeedForward(config)
+ self.ffn_norm = RMSNorm(config.dim, config.norm_eps)
+ self.attention_norm = RMSNorm(config.dim, config.norm_eps)
+
+ def forward(
+ self, x: Tensor, freqs_cis: Tensor, mask: Tensor, input_pos: Tensor = None
+ ) -> Tensor:
+ h = x + self.attention(self.attention_norm(x), freqs_cis, mask, input_pos)
+ out = h + self.feed_forward(self.ffn_norm(h))
+ return out
+
+
+class Attention(nn.Module):
+ def __init__(self, config: BaseModelArgs, use_sdpa: bool = True):
+ super().__init__()
+ assert config.dim % config.n_head == 0
+
+ total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim
+ # key, query, value projections for all heads, but in a batch
+ self.wqkv = nn.Linear(
+ config.dim, total_head_dim, bias=config.attention_qkv_bias
+ )
+ self.wo = nn.Linear(config.dim, config.dim, bias=False)
+ self.kv_cache = None
+
+ self.dropout = config.dropout
+ self.n_head = config.n_head
+ self.head_dim = config.head_dim
+ self.n_local_heads = config.n_local_heads
+ self.dim = config.dim
+ self.use_sdpa = use_sdpa
+ self._register_load_state_dict_pre_hook(self.load_hook)
+ self.qk_norm = config.qk_norm
+ self.qk_norm_groups = 1
+ self.qk_norm_scale = 10
+ self.qk_norm_dim_scale = False
+ self.qk_norm_q_scale = self.qk_norm_k_scale = 1
+
+ if self.qk_norm and self.qk_norm_dim_scale:
+ self.qk_norm_q_scale = nn.Parameter(torch.ones(self.n_head, 1, self.head_dim))
+ self.qk_norm_k_scale = nn.Parameter(torch.ones(self.n_head, 1, self.head_dim))
+ def load_hook(self, state_dict, prefix, *args):
+ if prefix + "wq.weight" in state_dict:
+ wq = state_dict.pop(prefix + "wq.weight")
+ wk = state_dict.pop(prefix + "wk.weight")
+ wv = state_dict.pop(prefix + "wv.weight")
+ state_dict[prefix + "wqkv.weight"] = torch.cat([wq, wk, wv])
+
+ def forward(
+ self,
+ x: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ input_pos: Optional[Tensor] = None,
+ ) -> Tensor:
+ bsz, seqlen, _ = x.shape
+
+ kv_size = self.n_local_heads * self.head_dim
+ q, k, v = self.wqkv(x).split([self.dim, kv_size, kv_size], dim=-1)
+
+ q = q.view(bsz, seqlen, self.n_head, self.head_dim)
+ k = k.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+ v = v.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+
+ if self.qk_norm:
+ qk_l2norm = partial(l2norm, groups = self.qk_norm_groups)
+ q, k = map(qk_l2norm, (q, k))
+ scale = self.qk_norm_scale
+
+ q = q * self.qk_norm_q_scale
+ k = k * self.qk_norm_k_scale
+
+ q = apply_rotary_emb(q, freqs_cis)
+ k = apply_rotary_emb(k, freqs_cis)
+
+ q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
+
+ if self.kv_cache is not None:
+ k, v = self.kv_cache.update(input_pos, k, v)
+
+ k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+
+ if self.use_sdpa:
+ if mask is None:
+ y = F.scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ dropout_p=self.dropout if self.training else 0.0,
+ is_causal=True,
+ # No third party attn_mask here to use flash_attention
+ )
+ else:
+ y = F.scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ attn_mask=mask,
+ dropout_p=self.dropout if self.training else 0.0,
+ )
+ else:
+ y = self.eq_scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ attn_mask=mask,
+ dropout_p=self.dropout if self.training else 0.0,
+ )
+
+ y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.dim)
+
+ return self.wo(y)
+
+ def eq_scaled_dot_product_attention(
+ self,
+ query,
+ key,
+ value,
+ attn_mask=None,
+ dropout_p=0.0,
+ ) -> torch.Tensor:
+ # This is a standard scaled dot product attention
+ # It's low efficient, but it doesn't raise cuda error
+
+ L, S = query.size(-2), key.size(-2)
+ scale_factor = 1 / math.sqrt(query.size(-1))
+ attn_bias = torch.zeros(1, 1, L, S, dtype=query.dtype, device=query.device)
+
+ if attn_mask is not None:
+ if attn_mask.dtype == torch.bool:
+ attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
+ else:
+ attn_bias += attn_mask
+
+ attn_weight = query @ key.transpose(-2, -1) * scale_factor
+ attn_weight += attn_bias
+ attn_weight = torch.softmax(attn_weight, dim=-1)
+ attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
+
+ return attn_weight @ value
+
+
+class FeedForward(nn.Module):
+ def __init__(self, config: BaseModelArgs) -> None:
+ super().__init__()
+ self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
+ self.dropout = nn.Dropout(p=config.dropout)
+
+ def forward(self, x: Tensor) -> Tensor:
+ return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x)))
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-5):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def _norm(self, x):
+ return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+ def forward(self, x: Tensor) -> Tensor:
+ output = self._norm(x.float()).type_as(x)
+ return output * self.weight
+
+
+def precompute_freqs_cis(seq_len: int, n_elem: int, base: int = 10000) -> Tensor:
+ freqs = 1.0 / (
+ base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem)
+ )
+ t = torch.arange(seq_len, device=freqs.device)
+ freqs = torch.outer(t, freqs)
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+ cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
+ return cache.to(dtype=torch.bfloat16)
+
+
+def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
+ xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+ freqs_cis = freqs_cis.view(x.size(0), xshaped.size(1), 1, xshaped.size(3), 2)
+ x_out2 = torch.stack(
+ [
+ xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],
+ xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],
+ ],
+ -1,
+ )
+
+ x_out2 = x_out2.flatten(3)
+ return x_out2.type_as(x)
+
+def top_k_top_p_filtering(
+ logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1
+):
+ """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+ Args:
+ logits: logits distribution shape (batch size, vocabulary size)
+ if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+ if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+ Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+ Make sure we keep at least min_tokens_to_keep per batch example in the output
+ From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+ """
+ if top_k > 0:
+ top_k = min(
+ max(top_k, min_tokens_to_keep), logits.size(-1)
+ ) # Safety check
+ # Remove all tokens with a probability less than the last token of the top-k
+ indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+ logits[indices_to_remove] = filter_value
+
+ if top_p < 1.0:
+ sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+ cumulative_probs = torch.cumsum(
+ F.softmax(sorted_logits, dim=-1), dim=-1
+ )
+
+ # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+ sorted_indices_to_remove = cumulative_probs > top_p
+ if min_tokens_to_keep > 1:
+ # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+ sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+ # Shift the indices to the right to keep also the first token above the threshold
+ sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
+ ..., :-1
+ ].clone()
+ sorted_indices_to_remove[..., 0] = 0
+
+ # scatter sorted tensors to original indexing
+ indices_to_remove = sorted_indices_to_remove.scatter(
+ 1, sorted_indices, sorted_indices_to_remove
+ )
+ logits[indices_to_remove] = filter_value
+ return logits
+
+def topk_sampling(logits, top_k=10, top_p=1.0, temperature=1.0):
+ # temperature: (`optional`) float
+ # The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+ # top_k: (`optional`) int
+ # The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.
+ # top_p: (`optional`) float
+ # The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
+
+ # Temperature (higher temperature => more likely to sample low probability tokens)
+ if temperature != 1.0:
+ logits = logits / temperature
+ # Top-p/top-k filtering
+ logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+ # Sample
+ token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+ logprobs = F.log_softmax(logits.float(), dim=-1)
+ current_logprobs = logprobs[torch.arange(logprobs.shape[0]), token.squeeze(1)]
+ return token, current_logprobs
+
+def sample(
+ logits,
+ previous_tokens: Optional[torch.Tensor] = None,
+ **sampling_kwargs,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ probs = logits_to_probs(
+ logits=logits[0, -1], previous_tokens=previous_tokens, **sampling_kwargs
+ )
+ idx_next = multinomial_sample_one_no_sync(probs)
+ return idx_next, probs
+
+def multinomial_sample_one_no_sync(
+ probs_sort,
+): # Does multinomial sampling without a cuda synchronization
+ q = torch.empty_like(probs_sort).exponential_(1)
+ return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.int)
+
+
+def logits_to_probs(
+ logits,
+ previous_tokens: Optional[torch.Tensor] = None,
+ suppress_tokens: Optional[List[int]] = None,
+ temperature: torch.Tensor = 0.7,
+ top_p: torch.Tensor = 0.7,
+ repetition_penalty: torch.Tensor = 1.5,
+) -> torch.Tensor:
+ # Apply repetition penalty
+ if previous_tokens is not None:
+ previous_tokens = previous_tokens.long()
+ score = torch.gather(logits, dim=0, index=previous_tokens)
+ score = torch.where(
+ score < 0, score * repetition_penalty, score / repetition_penalty
+ )
+ logits.scatter_(dim=0, index=previous_tokens, src=score)
+ if suppress_tokens is not None:
+ for token in suppress_tokens:
+ logits[token] = -float("Inf")
+
+ # Apply top-p sampling
+ sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+ cum_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
+ sorted_indices_to_remove = cum_probs > top_p
+ sorted_indices_to_remove[0] = False # keep at least one option
+ indices_to_remove = sorted_indices_to_remove.scatter(
+ dim=0, index=sorted_indices, src=sorted_indices_to_remove
+ )
+ logits = logits.masked_fill(indices_to_remove, -float("Inf"))
+
+ logits = logits / max(temperature, 1e-5)
+
+ probs = torch.nn.functional.softmax(logits, dim=-1)
+ return probs
diff --git a/modules/v2/cfm.py b/modules/v2/cfm.py
new file mode 100644
index 0000000000000000000000000000000000000000..b0ea58ef15f31c324bbcc061f15be8650824790e
--- /dev/null
+++ b/modules/v2/cfm.py
@@ -0,0 +1,173 @@
+import torch
+from tqdm import tqdm
+
+class CFM(torch.nn.Module):
+ def __init__(
+ self,
+ estimator: torch.nn.Module,
+ ):
+ super().__init__()
+ self.sigma_min = 1e-6
+ self.estimator = estimator
+ self.in_channels = estimator.in_channels
+ self.criterion = torch.nn.L1Loss()
+
+ @torch.inference_mode()
+ def inference(self,
+ mu: torch.Tensor,
+ x_lens: torch.Tensor,
+ prompt: torch.Tensor,
+ style: torch.Tensor,
+ n_timesteps=10,
+ temperature=1.0,
+ inference_cfg_rate=[0.5, 0.5],
+ random_voice=False,
+ ):
+ """Forward diffusion
+
+ Args:
+ mu (torch.Tensor): output of encoder
+ shape: (batch_size, n_feats, mel_timesteps)
+ x_lens (torch.Tensor): length of each mel-spectrogram
+ shape: (batch_size,)
+ prompt (torch.Tensor): prompt
+ shape: (batch_size, n_feats, prompt_len)
+ style (torch.Tensor): style
+ shape: (batch_size, style_dim)
+ n_timesteps (int): number of diffusion steps
+ temperature (float, optional): temperature for scaling noise. Defaults to 1.0.
+ inference_cfg_rate (float, optional): Classifier-Free Guidance inference introduced in VoiceBox. Defaults to 0.5.
+
+ Returns:
+ sample: generated mel-spectrogram
+ shape: (batch_size, n_feats, mel_timesteps)
+ """
+ B, T = mu.size(0), mu.size(1)
+ z = torch.randn([B, self.in_channels, T], device=mu.device) * temperature
+ t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device)
+ t_span = t_span + (-1) * (torch.cos(torch.pi / 2 * t_span) - 1 + t_span)
+ return self.solve_euler(z, x_lens, prompt, mu, style, t_span, inference_cfg_rate, random_voice)
+ def solve_euler(self, x, x_lens, prompt, mu, style, t_span, inference_cfg_rate=[0.5, 0.5], random_voice=False,):
+ """
+ Fixed euler solver for ODEs.
+ Args:
+ x (torch.Tensor): random noise
+ t_span (torch.Tensor): n_timesteps interpolated
+ shape: (n_timesteps + 1,)
+ mu (torch.Tensor): output of encoder
+ shape: (batch_size, n_feats, mel_timesteps)
+ x_lens (torch.Tensor): length of each mel-spectrogram
+ shape: (batch_size,)
+ prompt (torch.Tensor): prompt
+ shape: (batch_size, n_feats, prompt_len)
+ style (torch.Tensor): style
+ shape: (batch_size, style_dim)
+ inference_cfg_rate (float, optional): Classifier-Free Guidance inference introduced in VoiceBox. Defaults to 0.5.
+ sway_sampling (bool, optional): Sway sampling. Defaults to False.
+ amo_sampling (bool, optional): AMO sampling. Defaults to False.
+ """
+ t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]
+
+ # apply prompt
+ prompt_len = prompt.size(-1)
+ prompt_x = torch.zeros_like(x)
+ prompt_x[..., :prompt_len] = prompt[..., :prompt_len]
+ x[..., :prompt_len] = 0
+ for step in tqdm(range(1, len(t_span))):
+ if random_voice:
+ cfg_dphi_dt = self.estimator(
+ torch.cat([x, x], dim=0),
+ torch.cat([torch.zeros_like(prompt_x), torch.zeros_like(prompt_x)], dim=0),
+ torch.cat([x_lens, x_lens], dim=0),
+ torch.cat([t.unsqueeze(0), t.unsqueeze(0)], dim=0),
+ torch.cat([torch.zeros_like(style), torch.zeros_like(style)], dim=0),
+ torch.cat([mu, torch.zeros_like(mu)], dim=0),
+ )
+ cond_txt, uncond = cfg_dphi_dt[0:1], cfg_dphi_dt[1:2]
+ dphi_dt = ((1.0 + inference_cfg_rate[0]) * cond_txt - inference_cfg_rate[0] * uncond)
+ elif all(i == 0 for i in inference_cfg_rate):
+ dphi_dt = self.estimator(x, prompt_x, x_lens, t.unsqueeze(0), style, mu)
+ elif inference_cfg_rate[0] == 0:
+ # Classifier-Free Guidance inference introduced in VoiceBox
+ cfg_dphi_dt = self.estimator(
+ torch.cat([x, x], dim=0),
+ torch.cat([prompt_x, torch.zeros_like(prompt_x)], dim=0),
+ torch.cat([x_lens, x_lens], dim=0),
+ torch.cat([t.unsqueeze(0), t.unsqueeze(0)], dim=0),
+ torch.cat([style, torch.zeros_like(style)], dim=0),
+ torch.cat([mu, mu], dim=0),
+ )
+ cond_txt_spk, cond_txt = cfg_dphi_dt[0:1], cfg_dphi_dt[1:2]
+ dphi_dt = ((1.0 + inference_cfg_rate[1]) * cond_txt_spk - inference_cfg_rate[1] * cond_txt)
+ elif inference_cfg_rate[1] == 0:
+ cfg_dphi_dt = self.estimator(
+ torch.cat([x, x], dim=0),
+ torch.cat([prompt_x, torch.zeros_like(prompt_x)], dim=0),
+ torch.cat([x_lens, x_lens], dim=0),
+ torch.cat([t.unsqueeze(0), t.unsqueeze(0)], dim=0),
+ torch.cat([style, torch.zeros_like(style)], dim=0),
+ torch.cat([mu, torch.zeros_like(mu)], dim=0),
+ )
+ cond_txt_spk, uncond = cfg_dphi_dt[0:1], cfg_dphi_dt[1:2]
+ dphi_dt = ((1.0 + inference_cfg_rate[0]) * cond_txt_spk - inference_cfg_rate[0] * uncond)
+ else:
+ # Multi-condition Classifier-Free Guidance inference introduced in MegaTTS3
+ cfg_dphi_dt = self.estimator(
+ torch.cat([x, x, x], dim=0),
+ torch.cat([prompt_x, torch.zeros_like(prompt_x), torch.zeros_like(prompt_x)], dim=0),
+ torch.cat([x_lens, x_lens, x_lens], dim=0),
+ torch.cat([t.unsqueeze(0), t.unsqueeze(0), t.unsqueeze(0)], dim=0),
+ torch.cat([style, torch.zeros_like(style), torch.zeros_like(style)], dim=0),
+ torch.cat([mu, mu, torch.zeros_like(mu)], dim=0),
+ )
+ cond_txt_spk, cond_txt, uncond = cfg_dphi_dt[0:1], cfg_dphi_dt[1:2], cfg_dphi_dt[2:3]
+ dphi_dt = (1.0 + inference_cfg_rate[0] + inference_cfg_rate[1]) * cond_txt_spk - \
+ inference_cfg_rate[0] * uncond - inference_cfg_rate[1] * cond_txt
+ x = x + dt * dphi_dt
+ t = t + dt
+ if step < len(t_span) - 1:
+ dt = t_span[step + 1] - t
+ x[:, :, :prompt_len] = 0
+
+ return x
+
+ def forward(self, x1, x_lens, prompt_lens, mu, style):
+ """Computes diffusion loss
+
+ Args:
+ x1 (torch.Tensor): Target
+ shape: (batch_size, n_feats, mel_timesteps)
+ mask (torch.Tensor): target mask
+ shape: (batch_size, 1, mel_timesteps)
+ mu (torch.Tensor): output of encoder
+ shape: (batch_size, n_feats, mel_timesteps)
+ spks (torch.Tensor, optional): speaker embedding. Defaults to None.
+ shape: (batch_size, spk_emb_dim)
+
+ Returns:
+ loss: conditional flow matching loss
+ y: conditional flow
+ shape: (batch_size, n_feats, mel_timesteps)
+ """
+ b, _, t = x1.shape
+
+ # random timestep
+ t = torch.rand([b, 1, 1], device=mu.device, dtype=x1.dtype)
+ # sample noise p(x_0)
+ z = torch.randn_like(x1)
+
+ y = (1 - (1 - self.sigma_min) * t) * z + t * x1
+ u = x1 - (1 - self.sigma_min) * z
+ prompt = torch.zeros_like(x1)
+ for bib in range(b):
+ prompt[bib, :, :prompt_lens[bib]] = x1[bib, :, :prompt_lens[bib]]
+ # range covered by prompt are set to 0
+ y[bib, :, :prompt_lens[bib]] = 0
+
+ estimator_out = self.estimator(y, prompt, x_lens, t.squeeze(), style, mu)
+ loss = 0
+ for bib in range(b):
+ loss += self.criterion(estimator_out[bib, :, prompt_lens[bib]:x_lens[bib]], u[bib, :, prompt_lens[bib]:x_lens[bib]])
+ loss /= b
+
+ return loss
diff --git a/modules/v2/dit_model.py b/modules/v2/dit_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..4374ac86a4d4d0869788cdd16087115c4418ba5f
--- /dev/null
+++ b/modules/v2/dit_model.py
@@ -0,0 +1,250 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+from dataclasses import dataclass
+from typing import Optional, Union, Tuple, List
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import functional as F
+import time
+
+def find_multiple(n: int, k: int) -> int:
+ if n % k == 0:
+ return n
+ return n + k - (n % k)
+
+class AdaptiveLayerNorm(nn.Module):
+ r"""Adaptive Layer Normalization"""
+
+ def __init__(self, d_model, norm) -> None:
+ super(AdaptiveLayerNorm, self).__init__()
+ self.linear = nn.Linear(d_model, 6 * d_model)
+ self.act = nn.SiLU()
+ self.norm = norm
+ self.d_model = d_model
+ self.eps = self.norm.eps
+
+ def forward(self, x: Tensor, emb: Tensor) -> Tuple[Tensor]:
+ emb = self.linear(self.act(emb))
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=-1)
+
+ x = self.norm(x) * (1 + scale_msa) + shift_msa
+ return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+class AdaptiveLayerNormFinal(nn.Module):
+ r"""Adaptive Layer Normalization"""
+
+ def __init__(self, d_model, norm) -> None:
+ super(AdaptiveLayerNormFinal, self).__init__()
+ self.linear = nn.Linear(d_model, 2 * d_model)
+ self.act = nn.SiLU()
+ self.norm = norm
+ self.d_model = d_model
+ self.eps = self.norm.eps
+
+ def forward(self, x: Tensor, emb: Tensor) -> Tuple[Tensor]:
+ emb = self.linear(self.act(emb))
+ scale, shift = torch.chunk(emb, 2, dim=-1)
+
+ x = self.norm(x) * (1 + scale) + shift
+ return x
+
+@dataclass
+class ModelArgs:
+ block_size: int = 2048
+ vocab_size: int = 32000
+ n_layer: int = 32
+ n_head: int = 32
+ dim: int = 4096
+ intermediate_size: int = None
+ n_local_heads: int = -1
+ head_dim: int = 64
+ rope_base: float = 10000
+ norm_eps: float = 1e-5
+ uvit_skip_connection: bool = False
+ time_as_token: bool = False
+ dropout_rate: float = 0.1
+ attn_dropout_rate: float = 0.1
+
+ def __post_init__(self):
+ if self.n_local_heads == -1:
+ self.n_local_heads = self.n_head
+ if self.intermediate_size is None:
+ hidden_dim = 4 * self.dim
+ n_hidden = int(2 * hidden_dim / 3)
+ self.intermediate_size = find_multiple(n_hidden, 256)
+ # self.head_dim = self.dim // self.n_head
+
+class Transformer(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.config = config
+
+ self.layers = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer))
+ self.norm = AdaptiveLayerNormFinal(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+
+ self.max_batch_size = -1
+ self.max_seq_length = config.block_size
+
+ self.uvit_skip_connection = self.config.uvit_skip_connection
+ if self.uvit_skip_connection:
+ self.layers_emit_skip = [i for i in range(self.config.n_layer) if i < self.config.n_layer // 2]
+ self.layers_receive_skip = [i for i in range(self.config.n_layer) if i > self.config.n_layer // 2]
+ else:
+ self.layers_emit_skip = []
+ self.layers_receive_skip = []
+ freqs_cis = precompute_freqs_cis(self.config.block_size, self.config.head_dim,
+ self.config.rope_base)
+ self.register_buffer("freqs_cis", freqs_cis)
+
+ causal_mask = torch.tril(
+ torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool)
+ )
+ self.register_buffer("causal_mask", causal_mask)
+
+ def forward(self,
+ x: Tensor,
+ c: Tensor,
+ input_pos: Optional[Tensor] = None,
+ mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ mask = mask[..., input_pos]
+ freqs_cis = self.freqs_cis[input_pos]
+ for i, layer in enumerate(self.layers):
+ x = layer(x, c, freqs_cis, mask)
+ x = self.norm(x, c)
+ return x
+
+
+class TransformerBlock(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.attention = Attention(config)
+ self.feed_forward = FeedForward(config)
+ self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
+ self.attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+ def forward(self,
+ x: Tensor,
+ c: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ ) -> Tensor:
+ normed_x, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attention_norm(x, emb=c)
+ # attention
+ attn_output = self.attention(normed_x, freqs_cis, mask)
+ x = x + gate_msa * attn_output
+ normed_x = self.ffn_norm(x) * (1 + scale_mlp) + shift_mlp
+ ff_output = self.feed_forward(normed_x)
+ x = x + gate_mlp * ff_output
+ return x
+
+
+class Attention(nn.Module):
+ def __init__(self, config: ModelArgs, is_cross_attention: bool = False):
+ super().__init__()
+ assert config.dim % config.n_head == 0
+
+ total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim
+ # key, query, value projections for all heads, but in a batch
+ if is_cross_attention:
+ self.wq = nn.Linear(config.dim, config.n_head * config.head_dim, bias=False)
+ self.wkv = nn.Linear(config.context_dim, 2 * config.n_local_heads * config.head_dim, bias=False)
+ else:
+ self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False)
+ self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False)
+ self.kv_cache = None
+
+ self.n_head = config.n_head
+ self.head_dim = config.head_dim
+ self.n_local_heads = config.n_local_heads
+ self.dim = config.dim
+ self.attn_dropout_rate = config.attn_dropout_rate
+
+ def forward(self,
+ x: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ context: Optional[Tensor] = None,
+ context_freqs_cis: Optional[Tensor] = None,
+ ) -> Tensor:
+ bsz, seqlen, _ = x.shape
+
+ kv_size = self.n_local_heads * self.head_dim
+ q, k, v = self.wqkv(x).split([kv_size, kv_size, kv_size], dim=-1)
+ context_seqlen = seqlen
+
+ q = q.view(bsz, seqlen, self.n_head, self.head_dim)
+ k = k.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+ v = v.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+
+ q = apply_rotary_emb(q, freqs_cis)
+ k = apply_rotary_emb(k, freqs_cis)
+
+ q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
+
+ k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
+
+ y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.head_dim * self.n_head)
+
+ y = self.wo(y)
+ return y
+
+
+class FeedForward(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
+ self.dropout = nn.Dropout(config.dropout_rate)
+
+ def forward(self, x: Tensor) -> Tensor:
+ return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x)))
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-5):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def _norm(self, x):
+ return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+ def forward(self, x: Tensor) -> Tensor:
+ output = self._norm(x.float()).type_as(x)
+ return output * self.weight
+
+
+def precompute_freqs_cis(
+ seq_len: int, n_elem: int, base: int = 10000,
+ dtype: torch.dtype = torch.bfloat16
+) -> Tensor:
+ freqs = 1.0 / (base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem))
+ t = torch.arange(seq_len, device=freqs.device)
+ freqs = torch.outer(t, freqs)
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+ cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
+ return cache.to(dtype=dtype)
+
+
+def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
+ xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+ freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2)
+ x_out2 = torch.stack(
+ [
+ xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],
+ xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],
+ ],
+ -1,
+ )
+
+ x_out2 = x_out2.flatten(3)
+ return x_out2.type_as(x)
+
diff --git a/modules/v2/dit_wrapper.py b/modules/v2/dit_wrapper.py
new file mode 100644
index 0000000000000000000000000000000000000000..f653239de475fba17794d6b8d7e28a8edd1b65a0
--- /dev/null
+++ b/modules/v2/dit_wrapper.py
@@ -0,0 +1,152 @@
+import torch
+from torch import nn
+import math
+
+from modules.v2.dit_model import ModelArgs, Transformer
+from modules.commons import sequence_mask
+
+from torch.nn.utils import weight_norm
+
+def modulate(x, shift, scale):
+ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+#################################################################################
+# Embedding Layers for Timesteps and Class Labels #
+#################################################################################
+
+class TimestepEmbedder(nn.Module):
+ """
+ Embeds scalar timesteps into vector representations.
+ """
+ def __init__(self, hidden_size, frequency_embedding_size=256):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+ nn.SiLU(),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000, scale=1000):
+ """
+ Create sinusoidal timestep embeddings.
+ :param t: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ :param dim: the dimension of the output.
+ :param max_period: controls the minimum frequency of the embeddings.
+ :return: an (N, D) Tensor of positional embeddings.
+ """
+ # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+ ).to(device=t.device)
+ args = scale * t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t):
+ t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+ t_emb = self.mlp(t_freq)
+ return t_emb
+
+
+class DiT(torch.nn.Module):
+ def __init__(
+ self,
+ time_as_token,
+ style_as_token,
+ uvit_skip_connection,
+ block_size,
+ depth,
+ num_heads,
+ hidden_dim,
+ in_channels,
+ content_dim,
+ style_encoder_dim,
+ class_dropout_prob,
+ dropout_rate,
+ attn_dropout_rate,
+ ):
+ super(DiT, self).__init__()
+ self.time_as_token = time_as_token
+ self.style_as_token = style_as_token
+ self.uvit_skip_connection = uvit_skip_connection
+ model_args = ModelArgs(
+ block_size=block_size,
+ n_layer=depth,
+ n_head=num_heads,
+ dim=hidden_dim,
+ head_dim=hidden_dim // num_heads,
+ vocab_size=1, # we don't use this
+ uvit_skip_connection=self.uvit_skip_connection,
+ time_as_token=self.time_as_token,
+ dropout_rate=dropout_rate,
+ attn_dropout_rate=attn_dropout_rate,
+ )
+ self.transformer = Transformer(model_args)
+ self.in_channels = in_channels
+ self.out_channels = in_channels
+ self.num_heads = num_heads
+
+ self.x_embedder = weight_norm(nn.Linear(in_channels, hidden_dim, bias=True))
+
+ self.content_dim = content_dim # for continuous content
+ self.cond_projection = nn.Linear(content_dim, hidden_dim, bias=True) # continuous content
+
+ self.t_embedder = TimestepEmbedder(hidden_dim)
+
+ self.final_mlp = nn.Sequential(
+ nn.Linear(hidden_dim, hidden_dim),
+ nn.SiLU(),
+ nn.Linear(hidden_dim, in_channels),
+ )
+
+ self.class_dropout_prob = class_dropout_prob
+
+ self.cond_x_merge_linear = nn.Linear(hidden_dim + in_channels + in_channels, hidden_dim)
+ self.style_in = nn.Linear(style_encoder_dim, hidden_dim)
+
+ def forward(self, x, prompt_x, x_lens, t, style, cond):
+ class_dropout = False
+ content_dropout = False
+ if self.training and torch.rand(1) < self.class_dropout_prob:
+ class_dropout = True
+ if self.training and torch.rand(1) < 0.5:
+ content_dropout = True
+ cond_in_module = self.cond_projection
+
+ B, _, T = x.size()
+
+ t1 = self.t_embedder(t) # (N, D)
+ cond = cond_in_module(cond)
+
+ x = x.transpose(1, 2)
+ prompt_x = prompt_x.transpose(1, 2)
+
+ x_in = torch.cat([x, prompt_x, cond], dim=-1)
+ if class_dropout:
+ x_in[..., self.in_channels:self.in_channels*2] = 0
+ if content_dropout:
+ x_in[..., self.in_channels*2:] = 0
+ x_in = self.cond_x_merge_linear(x_in) # (N, T, D)
+
+ style = self.style_in(style)
+ style = torch.zeros_like(style) if class_dropout else style
+ if self.style_as_token:
+ x_in = torch.cat([style.unsqueeze(1), x_in], dim=1)
+ if self.time_as_token:
+ x_in = torch.cat([t1.unsqueeze(1), x_in], dim=1)
+ x_mask = sequence_mask(x_lens + self.style_as_token + self.time_as_token, max_length=x_in.size(1)).to(x.device).unsqueeze(1)
+ input_pos = torch.arange(x_in.size(1)).to(x.device)
+ x_mask_expanded = x_mask[:, None, :].repeat(1, 1, x_in.size(1), 1)
+ x_res = self.transformer(x_in, t1.unsqueeze(1), input_pos, x_mask_expanded)
+ x_res = x_res[:, 1:] if self.time_as_token else x_res
+ x_res = x_res[:, 1:] if self.style_as_token else x_res
+ x = self.final_mlp(x_res)
+ x = x.transpose(1, 2)
+ return x
diff --git a/modules/v2/length_regulator.py b/modules/v2/length_regulator.py
new file mode 100644
index 0000000000000000000000000000000000000000..7efe5a62bc5afba06a8abe5051aace9ad97dbf3e
--- /dev/null
+++ b/modules/v2/length_regulator.py
@@ -0,0 +1,105 @@
+from typing import Tuple
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from modules.commons import sequence_mask
+import numpy as np
+
+# f0_bin = 256
+f0_max = 1100.0
+f0_min = 50.0
+f0_mel_min = 1127 * np.log(1 + f0_min / 700)
+f0_mel_max = 1127 * np.log(1 + f0_max / 700)
+
+def f0_to_coarse(f0, f0_bin):
+ f0_mel = 1127 * (1 + f0 / 700).log()
+ a = (f0_bin - 2) / (f0_mel_max - f0_mel_min)
+ b = f0_mel_min * a - 1.
+ f0_mel = torch.where(f0_mel > 0, f0_mel * a - b, f0_mel)
+ # torch.clip_(f0_mel, min=1., max=float(f0_bin - 1))
+ f0_coarse = torch.round(f0_mel).long()
+ f0_coarse = f0_coarse * (f0_coarse > 0)
+ f0_coarse = f0_coarse + ((f0_coarse < 1) * 1)
+ f0_coarse = f0_coarse * (f0_coarse < f0_bin)
+ f0_coarse = f0_coarse + ((f0_coarse >= f0_bin) * (f0_bin - 1))
+ return f0_coarse
+
+class InterpolateRegulator(nn.Module):
+ def __init__(
+ self,
+ channels: int,
+ sampling_ratios: Tuple,
+ is_discrete: bool = False,
+ in_channels: int = None, # only applies to continuous input
+ codebook_size: int = 1024, # for discrete only
+ out_channels: int = None,
+ groups: int = 1,
+ f0_condition: bool = False,
+ n_f0_bins: int = 512,
+ ):
+ super().__init__()
+ self.sampling_ratios = sampling_ratios
+ out_channels = out_channels or channels
+ model = nn.ModuleList([])
+ if len(sampling_ratios) > 0:
+ self.interpolate = True
+ for _ in sampling_ratios:
+ module = nn.Conv1d(channels, channels, 3, 1, 1)
+ norm = nn.GroupNorm(groups, channels)
+ act = nn.Mish()
+ model.extend([module, norm, act])
+ else:
+ self.interpolate = False
+ model.append(
+ nn.Conv1d(channels, out_channels, 1, 1) if channels != out_channels else nn.Identity()
+ )
+ self.model = nn.Sequential(*model)
+ self.embedding = nn.Embedding(codebook_size, channels)
+ self.is_discrete = is_discrete
+
+ self.mask_token = nn.Parameter(torch.zeros(1, channels))
+
+ if f0_condition:
+ self.f0_embedding = nn.Embedding(n_f0_bins, channels)
+ self.f0_condition = f0_condition
+ self.n_f0_bins = n_f0_bins
+ self.f0_bins = torch.arange(2, 1024, 1024 // n_f0_bins)
+ self.f0_mask = nn.Parameter(torch.zeros(1, channels))
+ else:
+ self.f0_condition = False
+
+ if not is_discrete:
+ self.content_in_proj = nn.Linear(in_channels, channels)
+
+ def forward(self, x, ylens=None, f0=None):
+ if self.is_discrete:
+ if len(x.size()) == 2:
+ x = self.embedding(x)
+ else:
+ x = self.embedding(x[:, 0])
+ else:
+ x = self.content_in_proj(x)
+ # x in (B, T, D)
+
+ if self.interpolate:
+ mask = sequence_mask(ylens).unsqueeze(-1)
+ x = F.interpolate(x.transpose(1, 2).contiguous(), size=ylens.max(), mode='nearest')
+ else:
+ x = x.transpose(1, 2).contiguous()
+ mask = None
+ # mask = mask[:, :x.size(2), :]
+ # ylens = ylens.clamp(max=x.size(2)).long()
+ if self.f0_condition:
+ if f0 is None:
+ x = x + self.f0_mask.unsqueeze(-1)
+ else:
+ # quantized_f0 = torch.bucketize(f0, self.f0_bins.to(f0.device)) # (N, T)
+ quantized_f0 = f0_to_coarse(f0, self.n_f0_bins)
+ quantized_f0 = quantized_f0.clamp(0, self.n_f0_bins - 1).long()
+ f0_emb = self.f0_embedding(quantized_f0)
+ f0_emb = F.interpolate(f0_emb.transpose(1, 2).contiguous(), size=ylens.max(), mode='nearest')
+ x = x + f0_emb
+ out = self.model(x).transpose(1, 2).contiguous()
+ out = out * mask if mask is not None else out
+ olens = ylens
+ return out, olens
diff --git a/modules/v2/model.py b/modules/v2/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..a96dd0b6c58991ca3e203ca6c5247dc0413e48b4
--- /dev/null
+++ b/modules/v2/model.py
@@ -0,0 +1,302 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import functional as F
+
+
+def find_multiple(n: int, k: int) -> int:
+ if n % k == 0:
+ return n
+ return n + k - (n % k)
+
+class AdaptiveLayerNorm(nn.Module):
+ r"""Adaptive Layer Normalization"""
+
+ def __init__(self, d_model, norm) -> None:
+ super(AdaptiveLayerNorm, self).__init__()
+ self.project_layer = nn.Linear(d_model, 2 * d_model)
+ self.norm = norm
+ self.d_model = d_model
+ self.eps = self.norm.eps
+
+ def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor:
+ if embedding is None:
+ return self.norm(input)
+ weight, bias = torch.split(
+ self.project_layer(embedding),
+ split_size_or_sections=self.d_model,
+ dim=-1,
+ )
+ return weight * self.norm(input) + bias
+
+
+@dataclass
+class ModelArgs:
+ block_size: int = 2048
+ vocab_size: int = 32000
+ n_layer: int = 32
+ n_head: int = 32
+ dim: int = 4096
+ intermediate_size: int = None
+ n_local_heads: int = -1
+ head_dim: int = 64
+ rope_base: float = 10000
+ norm_eps: float = 1e-5
+ has_cross_attention: bool = False
+ context_dim: int = 0
+ uvit_skip_connection: bool = False
+ time_as_token: bool = False
+
+ def __post_init__(self):
+ if self.n_local_heads == -1:
+ self.n_local_heads = self.n_head
+ if self.intermediate_size is None:
+ hidden_dim = 4 * self.dim
+ n_hidden = int(2 * hidden_dim / 3)
+ self.intermediate_size = find_multiple(n_hidden, 256)
+ # self.head_dim = self.dim // self.n_head
+
+class Transformer(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.config = config
+
+ self.layers = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer))
+ self.norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+
+ self.freqs_cis: Optional[Tensor] = None
+ self.mask_cache: Optional[Tensor] = None
+ self.max_batch_size = -1
+ self.max_seq_length = -1
+
+ def setup_caches(self, max_batch_size, max_seq_length, use_kv_cache=False):
+ if self.max_seq_length >= max_seq_length and self.max_batch_size >= max_batch_size:
+ return
+ head_dim = self.config.dim // self.config.n_head
+ max_seq_length = find_multiple(max_seq_length, 8)
+ self.max_seq_length = max_seq_length
+ self.max_batch_size = max_batch_size
+ dtype = self.norm.project_layer.weight.dtype
+ device = self.norm.project_layer.weight.device
+
+ self.freqs_cis = precompute_freqs_cis(self.config.block_size, self.config.head_dim,
+ self.config.rope_base, dtype).to(device)
+ self.causal_mask = torch.tril(torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool)).to(device)
+ self.use_kv_cache = use_kv_cache
+ self.uvit_skip_connection = self.config.uvit_skip_connection
+ if self.uvit_skip_connection:
+ self.layers_emit_skip = [i for i in range(self.config.n_layer) if i < self.config.n_layer // 2]
+ self.layers_receive_skip = [i for i in range(self.config.n_layer) if i > self.config.n_layer // 2]
+ else:
+ self.layers_emit_skip = []
+ self.layers_receive_skip = []
+
+ def forward(self,
+ x: Tensor,
+ c: Tensor,
+ input_pos: Optional[Tensor] = None,
+ mask: Optional[Tensor] = None,
+ context: Optional[Tensor] = None,
+ context_input_pos: Optional[Tensor] = None,
+ cross_attention_mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ assert self.freqs_cis is not None, "Caches must be initialized first"
+ if mask is None: # in case of non-causal model
+ if not self.training and self.use_kv_cache:
+ mask = self.causal_mask[None, None, input_pos]
+ else:
+ mask = self.causal_mask[None, None, input_pos]
+ mask = mask[..., input_pos]
+ freqs_cis = self.freqs_cis[input_pos]
+ if context is not None:
+ context_freqs_cis = self.freqs_cis[context_input_pos]
+ else:
+ context_freqs_cis = None
+ skip_in_x_list = []
+ for i, layer in enumerate(self.layers):
+ if self.uvit_skip_connection and i in self.layers_receive_skip:
+ skip_in_x = skip_in_x_list.pop(-1)
+ else:
+ skip_in_x = None
+ x = layer(x, c, input_pos, freqs_cis, mask, context, context_freqs_cis, cross_attention_mask, skip_in_x)
+ if self.uvit_skip_connection and i in self.layers_emit_skip:
+ skip_in_x_list.append(x)
+ x = self.norm(x, c)
+ return x
+
+ @classmethod
+ def from_name(cls, name: str):
+ return cls(ModelArgs.from_name(name))
+
+
+class TransformerBlock(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.attention = Attention(config)
+ self.feed_forward = FeedForward(config)
+ self.ffn_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+ self.attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+
+ if config.has_cross_attention:
+ self.has_cross_attention = True
+ self.cross_attention = Attention(config, is_cross_attention=True)
+ self.cross_attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+ else:
+ self.has_cross_attention = False
+
+ if config.uvit_skip_connection:
+ self.skip_in_linear = nn.Linear(config.dim * 2, config.dim)
+ self.uvit_skip_connection = True
+ else:
+ self.uvit_skip_connection = False
+
+ self.time_as_token = config.time_as_token
+
+ def forward(self,
+ x: Tensor,
+ c: Tensor,
+ input_pos: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ context: Optional[Tensor] = None,
+ context_freqs_cis: Optional[Tensor] = None,
+ cross_attention_mask: Optional[Tensor] = None,
+ skip_in_x: Optional[Tensor] = None,
+ ) -> Tensor:
+ c = None if self.time_as_token else c
+ if self.uvit_skip_connection and skip_in_x is not None:
+ x = self.skip_in_linear(torch.cat([x, skip_in_x], dim=-1))
+ h = x + self.attention(self.attention_norm(x, c), freqs_cis, mask, input_pos)
+ if self.has_cross_attention:
+ h = h + self.cross_attention(self.cross_attention_norm(h, c), freqs_cis, cross_attention_mask, input_pos, context, context_freqs_cis)
+ out = h + self.feed_forward(self.ffn_norm(h, c))
+ return out
+
+
+class Attention(nn.Module):
+ def __init__(self, config: ModelArgs, is_cross_attention: bool = False):
+ super().__init__()
+ assert config.dim % config.n_head == 0
+
+ total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim
+ # key, query, value projections for all heads, but in a batch
+ if is_cross_attention:
+ self.wq = nn.Linear(config.dim, config.n_head * config.head_dim, bias=False)
+ self.wkv = nn.Linear(config.context_dim, 2 * config.n_local_heads * config.head_dim, bias=False)
+ else:
+ self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False)
+ self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False)
+ self.kv_cache = None
+
+ self.n_head = config.n_head
+ self.head_dim = config.head_dim
+ self.n_local_heads = config.n_local_heads
+ self.dim = config.dim
+ # self._register_load_state_dict_pre_hook(self.load_hook)
+
+ # def load_hook(self, state_dict, prefix, *args):
+ # if prefix + "wq.weight" in state_dict:
+ # wq = state_dict.pop(prefix + "wq.weight")
+ # wk = state_dict.pop(prefix + "wk.weight")
+ # wv = state_dict.pop(prefix + "wv.weight")
+ # state_dict[prefix + "wqkv.weight"] = torch.cat([wq, wk, wv])
+
+ def forward(self,
+ x: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ input_pos: Optional[Tensor] = None,
+ context: Optional[Tensor] = None,
+ context_freqs_cis: Optional[Tensor] = None,
+ ) -> Tensor:
+ bsz, seqlen, _ = x.shape
+
+ kv_size = self.n_local_heads * self.head_dim
+ if context is None:
+ q, k, v = self.wqkv(x).split([kv_size, kv_size, kv_size], dim=-1)
+ context_seqlen = seqlen
+ else:
+ q = self.wq(x)
+ k, v = self.wkv(context).split([kv_size, kv_size], dim=-1)
+ context_seqlen = context.shape[1]
+
+ q = q.view(bsz, seqlen, self.n_head, self.head_dim)
+ k = k.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+ v = v.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+
+ q = apply_rotary_emb(q, freqs_cis)
+ k = apply_rotary_emb(k, context_freqs_cis if context_freqs_cis is not None else freqs_cis)
+
+ q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
+
+ if self.kv_cache is not None:
+ k, v = self.kv_cache.update(input_pos, k, v)
+
+ k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
+
+ y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.head_dim * self.n_head)
+
+ y = self.wo(y)
+ return y
+
+
+class FeedForward(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
+
+ def forward(self, x: Tensor) -> Tensor:
+ return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-5):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def _norm(self, x):
+ return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+ def forward(self, x: Tensor) -> Tensor:
+ output = self._norm(x.float()).type_as(x)
+ return output * self.weight
+
+
+def precompute_freqs_cis(
+ seq_len: int, n_elem: int, base: int = 10000,
+ dtype: torch.dtype = torch.bfloat16
+) -> Tensor:
+ freqs = 1.0 / (base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem))
+ t = torch.arange(seq_len, device=freqs.device)
+ freqs = torch.outer(t, freqs)
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+ cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
+ return cache.to(dtype=dtype)
+
+
+def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
+ xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+ freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2)
+ x_out2 = torch.stack(
+ [
+ xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],
+ xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],
+ ],
+ -1,
+ )
+
+ x_out2 = x_out2.flatten(3)
+ return x_out2.type_as(x)
diff --git a/modules/v2/vc_wrapper.py b/modules/v2/vc_wrapper.py
new file mode 100644
index 0000000000000000000000000000000000000000..5b524a4c7f02540e2c81abcef08179872bace8aa
--- /dev/null
+++ b/modules/v2/vc_wrapper.py
@@ -0,0 +1,605 @@
+import spaces
+import torch
+import librosa
+import torchaudio
+import numpy as np
+from pydub import AudioSegment
+from hf_utils import load_custom_model_from_hf
+
+DEFAULT_REPO_ID = "Plachta/Seed-VC"
+DEFAULT_CFM_CHECKPOINT = "v2/cfm_small.pth"
+DEFAULT_AR_CHECKPOINT = "v2/ar_base.pth"
+
+DEFAULT_CE_REPO_ID = "Plachta/ASTRAL-quantization"
+DEFAULT_CE_NARROW_CHECKPOINT = "bsq32/bsq32_light.pth"
+DEFAULT_CE_WIDE_CHECKPOINT = "bsq2048/bsq2048_light.pth"
+
+DEFAULT_SE_REPO_ID = "funasr/campplus"
+DEFAULT_SE_CHECKPOINT = "campplus_cn_common.bin"
+
+class VoiceConversionWrapper(torch.nn.Module):
+ def __init__(
+ self,
+ sr: int,
+ hop_size: int,
+ mel_fn: callable,
+ cfm: torch.nn.Module,
+ cfm_length_regulator: torch.nn.Module,
+ content_extractor_narrow: torch.nn.Module,
+ content_extractor_wide: torch.nn.Module,
+ ar_length_regulator: torch.nn.Module,
+ ar: torch.nn.Module,
+ style_encoder: torch.nn.Module,
+ vocoder: torch.nn.Module,
+ ):
+ super(VoiceConversionWrapper, self).__init__()
+ self.sr = sr
+ self.hop_size = hop_size
+ self.mel_fn = mel_fn
+ self.cfm = cfm
+ self.cfm_length_regulator = cfm_length_regulator
+ self.content_extractor_narrow = content_extractor_narrow
+ self.content_extractor_wide = content_extractor_wide
+ self.vocoder = vocoder
+ self.ar_length_regulator = ar_length_regulator
+ self.ar = ar
+ self.style_encoder = style_encoder
+ # Set streaming parameters
+ self.overlap_frame_len = 16
+ self.bitrate = "320k"
+ self.compiled_decode_fn = None
+ self.dit_compiled = False
+ self.dit_max_context_len = 30 # in seconds
+ self.ar_max_content_len = 1500 # in num of narrow tokens
+ self.compile_len = 87 * self.dit_max_context_len
+
+ def compile_ar(self):
+ """
+ Compile the AR model for inference.
+ """
+ self.compiled_decode_fn = torch.compile(
+ self.ar.model.forward_generate,
+ fullgraph=True,
+ backend="inductor" if torch.cuda.is_available() else "aot_eager",
+ mode="reduce-overhead" if torch.cuda.is_available() else None,
+ )
+
+ def compile_cfm(self):
+ self.cfm.estimator.transformer = torch.compile(
+ self.cfm.estimator.transformer,
+ fullgraph=True,
+ backend="inductor" if torch.cuda.is_available() else "aot_eager",
+ mode="reduce-overhead" if torch.cuda.is_available() else None,
+ )
+ self.dit_compiled = True
+
+ @staticmethod
+ def strip_prefix(state_dict: dict, prefix: str = "module.") -> dict:
+ """
+ Strip the prefix from the state_dict keys.
+ """
+ new_state_dict = {}
+ for k, v in state_dict.items():
+ if k.startswith(prefix):
+ new_key = k[len(prefix):]
+ else:
+ new_key = k
+ new_state_dict[new_key] = v
+ return new_state_dict
+
+ @staticmethod
+ def duration_reduction_func(token_seq, n_gram=1):
+ """
+ Args:
+ token_seq: (T,)
+ Returns:
+ reduced_token_seq: (T')
+ reduced_token_seq_len: T'
+ """
+ n_gram_seq = token_seq.unfold(0, n_gram, 1)
+ mask = torch.all(n_gram_seq[1:] != n_gram_seq[:-1], dim=1)
+ reduced_token_seq = torch.cat(
+ (n_gram_seq[0, :n_gram], n_gram_seq[1:, -1][mask])
+ )
+ return reduced_token_seq, len(reduced_token_seq)
+
+ @staticmethod
+ def crossfade(chunk1, chunk2, overlap):
+ """Apply crossfade between two audio chunks."""
+ fade_out = np.cos(np.linspace(0, np.pi / 2, overlap)) ** 2
+ fade_in = np.cos(np.linspace(np.pi / 2, 0, overlap)) ** 2
+ if len(chunk2) < overlap:
+ chunk2[:overlap] = chunk2[:overlap] * fade_in[:len(chunk2)] + (chunk1[-overlap:] * fade_out)[:len(chunk2)]
+ else:
+ chunk2[:overlap] = chunk2[:overlap] * fade_in + chunk1[-overlap:] * fade_out
+ return chunk2
+
+ def _stream_wave_chunks(self, vc_wave, processed_frames, vc_mel, overlap_wave_len,
+ generated_wave_chunks, previous_chunk, is_last_chunk, stream_output):
+ """
+ Helper method to handle streaming wave chunks.
+
+ Args:
+ vc_wave: The current wave chunk
+ processed_frames: Number of frames processed so far
+ vc_mel: The mel spectrogram
+ overlap_wave_len: Length of overlap between chunks
+ generated_wave_chunks: List of generated wave chunks
+ previous_chunk: Previous wave chunk for crossfading
+ is_last_chunk: Whether this is the last chunk
+ stream_output: Whether to stream the output
+
+ Returns:
+ Tuple of (processed_frames, previous_chunk, should_break, mp3_bytes, full_audio)
+ where should_break indicates if processing should stop
+ mp3_bytes is the MP3 bytes if streaming, None otherwise
+ full_audio is the full audio if this is the last chunk, None otherwise
+ """
+ mp3_bytes = None
+ full_audio = None
+
+ if processed_frames == 0:
+ if is_last_chunk:
+ output_wave = vc_wave[0].cpu().numpy()
+ generated_wave_chunks.append(output_wave)
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=self.sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+ full_audio = (self.sr, np.concatenate(generated_wave_chunks))
+ else:
+ return processed_frames, previous_chunk, True, None, np.concatenate(generated_wave_chunks)
+
+ return processed_frames, previous_chunk, True, mp3_bytes, full_audio
+
+ output_wave = vc_wave[0, :-overlap_wave_len].cpu().numpy()
+ generated_wave_chunks.append(output_wave)
+ previous_chunk = vc_wave[0, -overlap_wave_len:]
+ processed_frames += vc_mel.size(2) - self.overlap_frame_len
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=self.sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+
+ elif is_last_chunk:
+ output_wave = self.crossfade(previous_chunk.cpu().numpy(), vc_wave[0].cpu().numpy(), overlap_wave_len)
+ generated_wave_chunks.append(output_wave)
+ processed_frames += vc_mel.size(2) - self.overlap_frame_len
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=self.sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+ full_audio = (self.sr, np.concatenate(generated_wave_chunks))
+ else:
+ return processed_frames, previous_chunk, True, None, np.concatenate(generated_wave_chunks)
+
+ return processed_frames, previous_chunk, True, mp3_bytes, full_audio
+
+ else:
+ output_wave = self.crossfade(previous_chunk.cpu().numpy(), vc_wave[0, :-overlap_wave_len].cpu().numpy(), overlap_wave_len)
+ generated_wave_chunks.append(output_wave)
+ previous_chunk = vc_wave[0, -overlap_wave_len:]
+ processed_frames += vc_mel.size(2) - self.overlap_frame_len
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=self.sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+
+ return processed_frames, previous_chunk, False, mp3_bytes, full_audio
+
+ def load_checkpoints(
+ self,
+ cfm_checkpoint_path = None,
+ ar_checkpoint_path = None,
+ ):
+ if cfm_checkpoint_path is None:
+ cfm_checkpoint_path = load_custom_model_from_hf(
+ repo_id=DEFAULT_REPO_ID,
+ model_filename=DEFAULT_CFM_CHECKPOINT,
+ )
+ if ar_checkpoint_path is None:
+ ar_checkpoint_path = load_custom_model_from_hf(
+ repo_id=DEFAULT_REPO_ID,
+ model_filename=DEFAULT_AR_CHECKPOINT,
+ )
+ # cfm
+ cfm_checkpoint = torch.load(cfm_checkpoint_path, map_location="cpu")
+ cfm_length_regulator_state_dict = self.strip_prefix(cfm_checkpoint["net"]['length_regulator'], "module.")
+ cfm_state_dict = self.strip_prefix(cfm_checkpoint["net"]['cfm'], "module.")
+ self.cfm.load_state_dict(cfm_state_dict, strict=False)
+ self.cfm_length_regulator.load_state_dict(cfm_length_regulator_state_dict, strict=False)
+
+ # ar
+ ar_checkpoint = torch.load(ar_checkpoint_path, map_location="cpu")
+ ar_length_regulator_state_dict = self.strip_prefix(ar_checkpoint["net"]['length_regulator'], "module.")
+ ar_state_dict = self.strip_prefix(ar_checkpoint["net"]['ar'], "module.")
+ self.ar.load_state_dict(ar_state_dict, strict=False)
+ self.ar_length_regulator.load_state_dict(ar_length_regulator_state_dict, strict=False)
+
+ # content extractor
+ content_extractor_narrow_checkpoint_path = load_custom_model_from_hf(
+ repo_id=DEFAULT_CE_REPO_ID,
+ model_filename=DEFAULT_CE_NARROW_CHECKPOINT,
+ )
+ content_extractor_narrow_checkpoint = torch.load(content_extractor_narrow_checkpoint_path, map_location="cpu")
+ self.content_extractor_narrow.load_state_dict(
+ content_extractor_narrow_checkpoint, strict=False
+ )
+
+ content_extractor_wide_checkpoint_path = load_custom_model_from_hf(
+ repo_id=DEFAULT_CE_REPO_ID,
+ model_filename=DEFAULT_CE_WIDE_CHECKPOINT,
+ )
+ content_extractor_wide_checkpoint = torch.load(content_extractor_wide_checkpoint_path, map_location="cpu")
+ self.content_extractor_wide.load_state_dict(
+ content_extractor_wide_checkpoint, strict=False
+ )
+
+ # style encoder
+ style_encoder_checkpoint_path = load_custom_model_from_hf(DEFAULT_SE_REPO_ID, DEFAULT_SE_CHECKPOINT, config_filename=None)
+ style_encoder_checkpoint = torch.load(style_encoder_checkpoint_path, map_location="cpu")
+ self.style_encoder.load_state_dict(style_encoder_checkpoint, strict=False)
+
+ def setup_ar_caches(self, max_batch_size=1, max_seq_len=4096, dtype=torch.float32, device=torch.device("cpu")):
+ self.ar.setup_caches(max_batch_size=max_batch_size, max_seq_len=max_seq_len, dtype=dtype, device=device)
+
+ def compute_style(self, waves_16k: torch.Tensor):
+ feat = torchaudio.compliance.kaldi.fbank(waves_16k,
+ num_mel_bins=80,
+ dither=0,
+ sample_frequency=16000)
+ feat = feat - feat.mean(dim=0, keepdim=True)
+ style = self.style_encoder(feat.unsqueeze(0))
+ return style
+
+ @torch.no_grad()
+ @torch.inference_mode()
+ def convert_timbre(
+ self,
+ source_audio_path: str,
+ target_audio_path: str,
+ diffusion_steps: int = 30,
+ length_adjust: float = 1.0,
+ inference_cfg_rate: float = 0.5,
+ use_sway_sampling: bool = False,
+ use_amo_sampling: bool = False,
+ device: torch.device = torch.device("cpu"),
+ dtype: torch.dtype = torch.float32,
+ ):
+ source_wave = librosa.load(source_audio_path, sr=self.sr)[0]
+ target_wave = librosa.load(target_audio_path, sr=self.sr)[0]
+ source_wave_tensor = torch.tensor(source_wave).unsqueeze(0).to(device)
+ target_wave_tensor = torch.tensor(target_wave).unsqueeze(0).to(device)
+
+ # get 16khz audio
+ source_wave_16k = librosa.resample(source_wave, orig_sr=self.sr, target_sr=16000)
+ target_wave_16k = librosa.resample(target_wave, orig_sr=self.sr, target_sr=16000)
+ source_wave_16k_tensor = torch.tensor(source_wave_16k).unsqueeze(0).to(device)
+ target_wave_16k_tensor = torch.tensor(target_wave_16k).unsqueeze(0).to(device)
+
+ # compute mel spectrogram
+ source_mel = self.mel_fn(source_wave_tensor)
+ target_mel = self.mel_fn(target_wave_tensor)
+ source_mel_len = source_mel.size(2)
+ target_mel_len = target_mel.size(2)
+
+ with torch.autocast(device_type=device.type, dtype=dtype):
+ # compute content features
+ _, source_content_indices, _ = self.content_extractor_wide(source_wave_16k_tensor, [source_wave_16k.size])
+ _, target_content_indices, _ = self.content_extractor_wide(target_wave_16k_tensor, [target_wave_16k.size])
+
+ # compute style features
+ target_style = self.compute_style(target_wave_16k_tensor)
+
+ # Length regulation
+ cond, _ = self.cfm_length_regulator(source_content_indices, ylens=torch.LongTensor([source_mel_len]).to(device))
+ prompt_condition, _, = self.cfm_length_regulator(target_content_indices, ylens=torch.LongTensor([target_mel_len]).to(device))
+
+ cat_condition = torch.cat([prompt_condition, cond], dim=1)
+ # generate mel spectrogram
+ vc_mel = self.cfm.inference(
+ cat_condition,
+ torch.LongTensor([cat_condition.size(1)]).to(device),
+ target_mel, target_style, diffusion_steps,
+ inference_cfg_rate=inference_cfg_rate,
+ sway_sampling=use_sway_sampling,
+ amo_sampling=use_amo_sampling,
+ )
+ vc_mel = vc_mel[:, :, target_mel_len:]
+ vc_wave = self.vocoder(vc_mel.float()).squeeze()[None]
+ return vc_wave.cpu().numpy()
+
+ @torch.no_grad()
+ @torch.inference_mode()
+ def convert_voice(
+ self,
+ source_audio_path: str,
+ target_audio_path: str,
+ diffusion_steps: int = 30,
+ length_adjust: float = 1.0,
+ inference_cfg_rate: float = 0.5,
+ top_p: float = 0.7,
+ temperature: float = 0.7,
+ repetition_penalty: float = 1.5,
+ use_sway_sampling: bool = False,
+ use_amo_sampling: bool = False,
+ device: torch.device = torch.device("cpu"),
+ dtype: torch.dtype = torch.float32,
+ ):
+ source_wave = librosa.load(source_audio_path, sr=self.sr)[0]
+ target_wave = librosa.load(target_audio_path, sr=self.sr)[0]
+ source_wave_tensor = torch.tensor(source_wave).unsqueeze(0).to(device)
+ target_wave_tensor = torch.tensor(target_wave).unsqueeze(0).to(device)
+
+ # get 16khz audio
+ source_wave_16k = librosa.resample(source_wave, orig_sr=self.sr, target_sr=16000)
+ target_wave_16k = librosa.resample(target_wave, orig_sr=self.sr, target_sr=16000)
+ source_wave_16k_tensor = torch.tensor(source_wave_16k).unsqueeze(0).to(device)
+ target_wave_16k_tensor = torch.tensor(target_wave_16k).unsqueeze(0).to(device)
+
+ # compute mel spectrogram
+ source_mel = self.mel_fn(source_wave_tensor)
+ target_mel = self.mel_fn(target_wave_tensor)
+ source_mel_len = source_mel.size(2)
+ target_mel_len = target_mel.size(2)
+
+ with torch.autocast(device_type=device.type, dtype=dtype):
+ # compute content features
+ _, source_content_indices, _ = self.content_extractor_wide(source_wave_16k_tensor, [source_wave_16k.size])
+ _, target_content_indices, _ = self.content_extractor_wide(target_wave_16k_tensor, [target_wave_16k.size])
+
+ _, source_narrow_indices, _ = self.content_extractor_narrow(source_wave_16k_tensor,
+ [source_wave_16k.size], ssl_model=self.content_extractor_wide.ssl_model)
+ _, target_narrow_indices, _ = self.content_extractor_narrow(target_wave_16k_tensor,
+ [target_wave_16k.size], ssl_model=self.content_extractor_wide.ssl_model)
+
+ src_narrow_reduced, src_narrow_len = self.duration_reduction_func(source_narrow_indices[0], 1)
+ tgt_narrow_reduced, tgt_narrow_len = self.duration_reduction_func(target_narrow_indices[0], 1)
+
+ ar_cond = self.ar_length_regulator(torch.cat([tgt_narrow_reduced, src_narrow_reduced], dim=0)[None])[0]
+
+ ar_out = self.ar.generate(ar_cond, target_content_indices, top_p=top_p, temperature=temperature, repetition_penalty=repetition_penalty)
+ ar_out_mel_len = torch.LongTensor([int(source_mel_len / source_content_indices.size(-1) * ar_out.size(-1) * length_adjust)]).to(device)
+ # compute style features
+ target_style = self.compute_style(target_wave_16k_tensor)
+
+ # Length regulation
+ cond, _ = self.cfm_length_regulator(ar_out, ylens=torch.LongTensor([ar_out_mel_len]).to(device))
+ prompt_condition, _, = self.cfm_length_regulator(target_content_indices, ylens=torch.LongTensor([target_mel_len]).to(device))
+
+ cat_condition = torch.cat([prompt_condition, cond], dim=1)
+ # generate mel spectrogram
+ vc_mel = self.cfm.inference(
+ cat_condition,
+ torch.LongTensor([cat_condition.size(1)]).to(device),
+ target_mel, target_style, diffusion_steps,
+ inference_cfg_rate=inference_cfg_rate,
+ sway_sampling=use_sway_sampling,
+ amo_sampling=use_amo_sampling,
+ )
+ vc_mel = vc_mel[:, :, target_mel_len:]
+ vc_wave = self.vocoder(vc_mel.float()).squeeze()[None]
+ return vc_wave.cpu().numpy()
+
+ def _process_content_features(self, audio_16k_tensor, is_narrow=False):
+ """Process audio through Whisper model to extract features."""
+ content_extractor_fn = self.content_extractor_narrow if is_narrow else self.content_extractor_wide
+ if audio_16k_tensor.size(-1) <= 16000 * 30:
+ # Compute content features
+ _, content_indices, _ = content_extractor_fn(audio_16k_tensor, [audio_16k_tensor.size(-1)], ssl_model=self.content_extractor_wide.ssl_model)
+ else:
+ # Process long audio in chunks
+ overlapping_time = 5 # 5 seconds
+ features_list = []
+ buffer = None
+ traversed_time = 0
+ while traversed_time < audio_16k_tensor.size(-1):
+ if buffer is None: # first chunk
+ chunk = audio_16k_tensor[:, traversed_time:traversed_time + 16000 * 30]
+ else:
+ chunk = torch.cat([
+ buffer,
+ audio_16k_tensor[:, traversed_time:traversed_time + 16000 * (30 - overlapping_time)]
+ ], dim=-1)
+ _, chunk_content_indices, _ = content_extractor_fn(chunk, [chunk.size(-1)], ssl_model=self.content_extractor_wide.ssl_model)
+ if traversed_time == 0:
+ features_list.append(chunk_content_indices)
+ else:
+ features_list.append(chunk_content_indices[:, 50 * overlapping_time:])
+ buffer = chunk[:, -16000 * overlapping_time:]
+ traversed_time += 30 * 16000 if traversed_time == 0 else chunk.size(-1) - 16000 * overlapping_time
+ content_indices = torch.cat(features_list, dim=1)
+
+ return content_indices
+
+ @spaces.GPU
+ @torch.no_grad()
+ @torch.inference_mode()
+ def convert_voice_with_streaming(
+ self,
+ source_audio_path: str,
+ target_audio_path: str,
+ diffusion_steps: int = 30,
+ length_adjust: float = 1.0,
+ intelligebility_cfg_rate: float = 0.7,
+ similarity_cfg_rate: float = 0.7,
+ top_p: float = 0.7,
+ temperature: float = 0.7,
+ repetition_penalty: float = 1.5,
+ convert_style: bool = False,
+ anonymization_only: bool = False,
+ device: torch.device = torch.device("cuda"),
+ dtype: torch.dtype = torch.float16,
+ stream_output: bool = True,
+ ):
+ """
+ Convert voice with streaming support for long audio files.
+
+ Args:
+ source_audio_path: Path to source audio file
+ target_audio_path: Path to target audio file
+ diffusion_steps: Number of diffusion steps (default: 30)
+ length_adjust: Length adjustment factor (default: 1.0)
+ intelligebility_cfg_rate: CFG rate for intelligibility (default: 0.7)
+ similarity_cfg_rate: CFG rate for similarity (default: 0.7)
+ top_p: Top-p sampling parameter (default: 0.7)
+ temperature: Temperature for sampling (default: 0.7)
+ repetition_penalty: Repetition penalty (default: 1.5)
+ device: Device to use (default: cpu)
+ dtype: Data type to use (default: float32)
+ stream_output: Whether to stream the output (default: True)
+
+ Returns:
+ If stream_output is True, yields (mp3_bytes, full_audio) tuples
+ If stream_output is False, returns the full audio as a numpy array
+ """
+ # Load audio
+ source_wave = librosa.load(source_audio_path, sr=self.sr)[0]
+ target_wave = librosa.load(target_audio_path, sr=self.sr)[0]
+
+ # Limit target audio to 25 seconds
+ target_wave = target_wave[:self.sr * (self.dit_max_context_len - 5)]
+
+ source_wave_tensor = torch.tensor(source_wave).unsqueeze(0).float().to(device)
+ target_wave_tensor = torch.tensor(target_wave).unsqueeze(0).float().to(device)
+
+ # Resample to 16kHz for feature extraction
+ source_wave_16k = librosa.resample(source_wave, orig_sr=self.sr, target_sr=16000)
+ target_wave_16k = librosa.resample(target_wave, orig_sr=self.sr, target_sr=16000)
+ source_wave_16k_tensor = torch.tensor(source_wave_16k).unsqueeze(0).to(device)
+ target_wave_16k_tensor = torch.tensor(target_wave_16k).unsqueeze(0).to(device)
+
+ # Compute mel spectrograms
+ source_mel = self.mel_fn(source_wave_tensor)
+ target_mel = self.mel_fn(target_wave_tensor)
+ source_mel_len = source_mel.size(2)
+ target_mel_len = target_mel.size(2)
+
+ # Set up chunk processing parameters
+ max_context_window = self.sr // self.hop_size * self.dit_max_context_len
+ overlap_wave_len = self.overlap_frame_len * self.hop_size
+
+ with torch.autocast(device_type=device.type, dtype=dtype):
+ # Compute content features
+ source_content_indices = self._process_content_features(source_wave_16k_tensor, is_narrow=False)
+ target_content_indices = self._process_content_features(target_wave_16k_tensor, is_narrow=False)
+ # Compute style features
+ target_style = self.compute_style(target_wave_16k_tensor)
+ prompt_condition, _, = self.cfm_length_regulator(target_content_indices,
+ ylens=torch.LongTensor([target_mel_len]).to(device))
+
+ # prepare for streaming
+ generated_wave_chunks = []
+ processed_frames = 0
+ previous_chunk = None
+ if convert_style:
+ with torch.autocast(device_type=device.type, dtype=dtype):
+ source_narrow_indices = self._process_content_features(source_wave_16k_tensor, is_narrow=True)
+ target_narrow_indices = self._process_content_features(target_wave_16k_tensor, is_narrow=True)
+ src_narrow_reduced, src_narrow_len = self.duration_reduction_func(source_narrow_indices[0], 1)
+ tgt_narrow_reduced, tgt_narrow_len = self.duration_reduction_func(target_narrow_indices[0], 1)
+ # Process src_narrow_reduced in chunks of max 1000 tokens
+ max_chunk_size = self.ar_max_content_len - tgt_narrow_len
+
+ # Process src_narrow_reduced in chunks
+ for i in range(0, len(src_narrow_reduced), max_chunk_size):
+ is_last_chunk = i + max_chunk_size >= len(src_narrow_reduced)
+ with torch.autocast(device_type=device.type, dtype=dtype):
+ chunk = src_narrow_reduced[i:i + max_chunk_size]
+ if anonymization_only:
+ chunk_ar_cond = self.ar_length_regulator(chunk[None])[0]
+ chunk_ar_out = self.ar.generate(chunk_ar_cond, torch.zeros([1, 0]).long().to(device),
+ compiled_decode_fn=self.compiled_decode_fn,
+ top_p=top_p, temperature=temperature,
+ repetition_penalty=repetition_penalty)
+ else:
+ # For each chunk, we need to include tgt_narrow_reduced as context
+ chunk_ar_cond = self.ar_length_regulator(torch.cat([tgt_narrow_reduced, chunk], dim=0)[None])[0]
+ chunk_ar_out = self.ar.generate(chunk_ar_cond, target_content_indices, compiled_decode_fn=self.compiled_decode_fn,
+ top_p=top_p, temperature=temperature,
+ repetition_penalty=repetition_penalty)
+ chunkar_out_mel_len = torch.LongTensor([int(source_mel_len / source_content_indices.size(
+ -1) * chunk_ar_out.size(-1) * length_adjust)]).to(device)
+ # Length regulation
+ chunk_cond, _ = self.cfm_length_regulator(chunk_ar_out, ylens=torch.LongTensor([chunkar_out_mel_len]).to(device))
+ cat_condition = torch.cat([prompt_condition, chunk_cond], dim=1)
+ original_len = cat_condition.size(1)
+ # pad cat_condition to compile_len
+ if self.dit_compiled:
+ cat_condition = torch.nn.functional.pad(cat_condition,
+ (0, 0, 0, self.compile_len - cat_condition.size(1),),
+ value=0)
+ # Voice Conversion
+ vc_mel = self.cfm.inference(
+ cat_condition,
+ torch.LongTensor([original_len]).to(device),
+ target_mel, target_style, diffusion_steps,
+ inference_cfg_rate=[intelligebility_cfg_rate, similarity_cfg_rate],
+ random_voice=anonymization_only,
+ )
+ vc_mel = vc_mel[:, :, target_mel_len:original_len]
+ vc_wave = self.vocoder(vc_mel).squeeze()[None]
+ processed_frames, previous_chunk, should_break, mp3_bytes, full_audio = self._stream_wave_chunks(
+ vc_wave, processed_frames, vc_mel, overlap_wave_len,
+ generated_wave_chunks, previous_chunk, is_last_chunk, stream_output
+ )
+
+ if stream_output and mp3_bytes is not None:
+ yield mp3_bytes, full_audio
+
+ if should_break:
+ if not stream_output:
+ return full_audio
+ break
+ else:
+ cond, _ = self.cfm_length_regulator(source_content_indices, ylens=torch.LongTensor([source_mel_len]).to(device))
+
+ # Process in chunks for streaming
+ max_source_window = max_context_window - target_mel.size(2)
+
+ # Generate chunk by chunk and stream the output
+ while processed_frames < cond.size(1):
+ chunk_cond = cond[:, processed_frames:processed_frames + max_source_window]
+ is_last_chunk = processed_frames + max_source_window >= cond.size(1)
+ cat_condition = torch.cat([prompt_condition, chunk_cond], dim=1)
+ original_len = cat_condition.size(1)
+ # pad cat_condition to compile_len
+ if self.dit_compiled:
+ cat_condition = torch.nn.functional.pad(cat_condition,
+ (0, 0, 0, self.compile_len - cat_condition.size(1),), value=0)
+ with torch.autocast(device_type=device.type, dtype=dtype):
+ # Voice Conversion
+ vc_mel = self.cfm.inference(
+ cat_condition,
+ torch.LongTensor([original_len]).to(device),
+ target_mel, target_style, diffusion_steps,
+ inference_cfg_rate=[intelligebility_cfg_rate, similarity_cfg_rate],
+ random_voice=anonymization_only,
+ )
+ vc_mel = vc_mel[:, :, target_mel_len:original_len]
+ vc_wave = self.vocoder(vc_mel).squeeze()[None]
+
+ processed_frames, previous_chunk, should_break, mp3_bytes, full_audio = self._stream_wave_chunks(
+ vc_wave, processed_frames, vc_mel, overlap_wave_len,
+ generated_wave_chunks, previous_chunk, is_last_chunk, stream_output
+ )
+
+ if stream_output and mp3_bytes is not None:
+ yield mp3_bytes, full_audio
+
+ if should_break:
+ if not stream_output:
+ return full_audio
+ break
\ No newline at end of file
diff --git a/modules/wavenet.py b/modules/wavenet.py
new file mode 100644
index 0000000000000000000000000000000000000000..6c81ffbe58c99fdfb30da9567d609cf193350d8d
--- /dev/null
+++ b/modules/wavenet.py
@@ -0,0 +1,174 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from modules.encodec import SConv1d
+
+from . import commons
+LRELU_SLOPE = 0.1
+
+class LayerNorm(nn.Module):
+ def __init__(self, channels, eps=1e-5):
+ super().__init__()
+ self.channels = channels
+ self.eps = eps
+
+ self.gamma = nn.Parameter(torch.ones(channels))
+ self.beta = nn.Parameter(torch.zeros(channels))
+
+ def forward(self, x):
+ x = x.transpose(1, -1)
+ x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+ return x.transpose(1, -1)
+
+
+class ConvReluNorm(nn.Module):
+ def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):
+ super().__init__()
+ self.in_channels = in_channels
+ self.hidden_channels = hidden_channels
+ self.out_channels = out_channels
+ self.kernel_size = kernel_size
+ self.n_layers = n_layers
+ self.p_dropout = p_dropout
+ assert n_layers > 1, "Number of layers should be larger than 0."
+
+ self.conv_layers = nn.ModuleList()
+ self.norm_layers = nn.ModuleList()
+ self.conv_layers.append(nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size // 2))
+ self.norm_layers.append(LayerNorm(hidden_channels))
+ self.relu_drop = nn.Sequential(
+ nn.ReLU(),
+ nn.Dropout(p_dropout))
+ for _ in range(n_layers - 1):
+ self.conv_layers.append(nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size // 2))
+ self.norm_layers.append(LayerNorm(hidden_channels))
+ self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+ self.proj.weight.data.zero_()
+ self.proj.bias.data.zero_()
+
+ def forward(self, x, x_mask):
+ x_org = x
+ for i in range(self.n_layers):
+ x = self.conv_layers[i](x * x_mask)
+ x = self.norm_layers[i](x)
+ x = self.relu_drop(x)
+ x = x_org + self.proj(x)
+ return x * x_mask
+
+
+class DDSConv(nn.Module):
+ """
+ Dialted and Depth-Separable Convolution
+ """
+
+ def __init__(self, channels, kernel_size, n_layers, p_dropout=0.):
+ super().__init__()
+ self.channels = channels
+ self.kernel_size = kernel_size
+ self.n_layers = n_layers
+ self.p_dropout = p_dropout
+
+ self.drop = nn.Dropout(p_dropout)
+ self.convs_sep = nn.ModuleList()
+ self.convs_1x1 = nn.ModuleList()
+ self.norms_1 = nn.ModuleList()
+ self.norms_2 = nn.ModuleList()
+ for i in range(n_layers):
+ dilation = kernel_size ** i
+ padding = (kernel_size * dilation - dilation) // 2
+ self.convs_sep.append(nn.Conv1d(channels, channels, kernel_size,
+ groups=channels, dilation=dilation, padding=padding
+ ))
+ self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
+ self.norms_1.append(LayerNorm(channels))
+ self.norms_2.append(LayerNorm(channels))
+
+ def forward(self, x, x_mask, g=None):
+ if g is not None:
+ x = x + g
+ for i in range(self.n_layers):
+ y = self.convs_sep[i](x * x_mask)
+ y = self.norms_1[i](y)
+ y = F.gelu(y)
+ y = self.convs_1x1[i](y)
+ y = self.norms_2[i](y)
+ y = F.gelu(y)
+ y = self.drop(y)
+ x = x + y
+ return x * x_mask
+
+
+class WN(torch.nn.Module):
+ def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0, causal=False):
+ super(WN, self).__init__()
+ conv1d_type = SConv1d
+ assert (kernel_size % 2 == 1)
+ self.hidden_channels = hidden_channels
+ self.kernel_size = kernel_size,
+ self.dilation_rate = dilation_rate
+ self.n_layers = n_layers
+ self.gin_channels = gin_channels
+ self.p_dropout = p_dropout
+
+ self.in_layers = torch.nn.ModuleList()
+ self.res_skip_layers = torch.nn.ModuleList()
+ self.drop = nn.Dropout(p_dropout)
+
+ if gin_channels != 0:
+ self.cond_layer = conv1d_type(gin_channels, 2 * hidden_channels * n_layers, 1, norm='weight_norm')
+
+ for i in range(n_layers):
+ dilation = dilation_rate ** i
+ padding = int((kernel_size * dilation - dilation) / 2)
+ in_layer = conv1d_type(hidden_channels, 2 * hidden_channels, kernel_size, dilation=dilation,
+ padding=padding, norm='weight_norm', causal=causal)
+ self.in_layers.append(in_layer)
+
+ # last one is not necessary
+ if i < n_layers - 1:
+ res_skip_channels = 2 * hidden_channels
+ else:
+ res_skip_channels = hidden_channels
+
+ res_skip_layer = conv1d_type(hidden_channels, res_skip_channels, 1, norm='weight_norm', causal=causal)
+ self.res_skip_layers.append(res_skip_layer)
+
+ def forward(self, x, x_mask, g=None, **kwargs):
+ output = torch.zeros_like(x)
+ n_channels_tensor = torch.IntTensor([self.hidden_channels])
+
+ if g is not None:
+ g = self.cond_layer(g)
+
+ for i in range(self.n_layers):
+ x_in = self.in_layers[i](x)
+ if g is not None:
+ cond_offset = i * 2 * self.hidden_channels
+ g_l = g[:, cond_offset:cond_offset + 2 * self.hidden_channels, :]
+ else:
+ g_l = torch.zeros_like(x_in)
+
+ acts = commons.fused_add_tanh_sigmoid_multiply(
+ x_in,
+ g_l,
+ n_channels_tensor)
+ acts = self.drop(acts)
+
+ res_skip_acts = self.res_skip_layers[i](acts)
+ if i < self.n_layers - 1:
+ res_acts = res_skip_acts[:, :self.hidden_channels, :]
+ x = (x + res_acts) * x_mask
+ output = output + res_skip_acts[:, self.hidden_channels:, :]
+ else:
+ output = output + res_skip_acts
+ return output * x_mask
+
+ def remove_weight_norm(self):
+ if self.gin_channels != 0:
+ torch.nn.utils.remove_weight_norm(self.cond_layer)
+ for l in self.in_layers:
+ torch.nn.utils.remove_weight_norm(l)
+ for l in self.res_skip_layers:
+ torch.nn.utils.remove_weight_norm(l)
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..074a97a518397e2f5a9aa24237f057e6a81f6a7f
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,20 @@
+torch --pre --index-url https://download.pytorch.org/whl/nightly/cu126
+torchvision --pre --index-url https://download.pytorch.org/whl/nightly/cu126
+torchaudio --pre --index-url https://download.pytorch.org/whl/nightly/cu126
+scipy==1.13.1
+librosa==0.10.2
+huggingface-hub
+munch==4.0.0
+einops==0.8.0
+pydub==0.25.1
+descript-audio-codec==1.0.0
+transformers==4.46.3
+soundfile==0.12.1
+sounddevice==0.5.0
+modelscope==1.18.1
+funasr==1.1.5
+numpy==1.26.4
+hydra-core==1.3.2
+pyyaml
+python-dotenv
+matplotlib
diff --git a/seed_vc_wrapper.py b/seed_vc_wrapper.py
new file mode 100644
index 0000000000000000000000000000000000000000..14caf0cfd3a18861e1a2b2c5fc3774aaa244e353
--- /dev/null
+++ b/seed_vc_wrapper.py
@@ -0,0 +1,463 @@
+import spaces
+import torch
+import torchaudio
+import librosa
+import numpy as np
+from pydub import AudioSegment
+import yaml
+from modules.commons import build_model, load_checkpoint, recursive_munch
+from hf_utils import load_custom_model_from_hf
+from modules.campplus.DTDNN import CAMPPlus
+from modules.bigvgan import bigvgan
+from modules.audio import mel_spectrogram
+from modules.rmvpe import RMVPE
+from transformers import AutoFeatureExtractor, WhisperModel
+
+class SeedVCWrapper:
+ def __init__(self, device=None):
+ """
+ Initialize the Seed-VC wrapper with all necessary models and configurations.
+
+ Args:
+ device: torch device to use. If None, will be automatically determined.
+ """
+ # Set device
+ if device is None:
+ if torch.cuda.is_available():
+ self.device = torch.device("cuda")
+ elif torch.backends.mps.is_available():
+ self.device = torch.device("mps")
+ else:
+ self.device = torch.device("cpu")
+ else:
+ self.device = device
+
+ # Load base model and configuration
+ self._load_base_model()
+
+ # Load F0 conditioned model
+ self._load_f0_model()
+
+ # Load additional modules
+ self._load_additional_modules()
+
+ # Set streaming parameters
+ self.overlap_frame_len = 16
+ self.bitrate = "320k"
+
+ def _load_base_model(self):
+ """Load the base DiT model for voice conversion."""
+ dit_checkpoint_path, dit_config_path = load_custom_model_from_hf(
+ "Plachta/Seed-VC",
+ "DiT_seed_v2_uvit_whisper_small_wavenet_bigvgan_pruned.pth",
+ "config_dit_mel_seed_uvit_whisper_small_wavenet.yml"
+ )
+ config = yaml.safe_load(open(dit_config_path, 'r'))
+ model_params = recursive_munch(config['model_params'])
+ self.model = build_model(model_params, stage='DiT')
+ self.hop_length = config['preprocess_params']['spect_params']['hop_length']
+ self.sr = config['preprocess_params']['sr']
+
+ # Load checkpoints
+ self.model, _, _, _ = load_checkpoint(
+ self.model, None, dit_checkpoint_path,
+ load_only_params=True, ignore_modules=[], is_distributed=False
+ )
+ for key in self.model:
+ self.model[key].eval()
+ self.model[key].to(self.device)
+ self.model.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192)
+
+ # Set up mel spectrogram function
+ mel_fn_args = {
+ "n_fft": config['preprocess_params']['spect_params']['n_fft'],
+ "win_size": config['preprocess_params']['spect_params']['win_length'],
+ "hop_size": config['preprocess_params']['spect_params']['hop_length'],
+ "num_mels": config['preprocess_params']['spect_params']['n_mels'],
+ "sampling_rate": self.sr,
+ "fmin": 0,
+ "fmax": None,
+ "center": False
+ }
+ self.to_mel = lambda x: mel_spectrogram(x, **mel_fn_args)
+
+ # Load whisper model
+ whisper_name = model_params.speech_tokenizer.whisper_name if hasattr(model_params.speech_tokenizer, 'whisper_name') else "openai/whisper-small"
+ self.whisper_model = WhisperModel.from_pretrained(whisper_name, torch_dtype=torch.float16).to(self.device)
+ del self.whisper_model.decoder
+ self.whisper_feature_extractor = AutoFeatureExtractor.from_pretrained(whisper_name)
+
+ def _load_f0_model(self):
+ """Load the F0 conditioned model for voice conversion."""
+ dit_checkpoint_path, dit_config_path = load_custom_model_from_hf(
+ "Plachta/Seed-VC",
+ "DiT_seed_v2_uvit_whisper_base_f0_44k_bigvgan_pruned_ft_ema.pth",
+ "config_dit_mel_seed_uvit_whisper_base_f0_44k.yml"
+ )
+ config = yaml.safe_load(open(dit_config_path, 'r'))
+ model_params = recursive_munch(config['model_params'])
+ self.model_f0 = build_model(model_params, stage='DiT')
+ self.hop_length_f0 = config['preprocess_params']['spect_params']['hop_length']
+ self.sr_f0 = config['preprocess_params']['sr']
+
+ # Load checkpoints
+ self.model_f0, _, _, _ = load_checkpoint(
+ self.model_f0, None, dit_checkpoint_path,
+ load_only_params=True, ignore_modules=[], is_distributed=False
+ )
+ for key in self.model_f0:
+ self.model_f0[key].eval()
+ self.model_f0[key].to(self.device)
+ self.model_f0.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192)
+
+ # Set up mel spectrogram function for F0 model
+ mel_fn_args_f0 = {
+ "n_fft": config['preprocess_params']['spect_params']['n_fft'],
+ "win_size": config['preprocess_params']['spect_params']['win_length'],
+ "hop_size": config['preprocess_params']['spect_params']['hop_length'],
+ "num_mels": config['preprocess_params']['spect_params']['n_mels'],
+ "sampling_rate": self.sr_f0,
+ "fmin": 0,
+ "fmax": None,
+ "center": False
+ }
+ self.to_mel_f0 = lambda x: mel_spectrogram(x, **mel_fn_args_f0)
+
+ def _load_additional_modules(self):
+ """Load additional modules like CAMPPlus, BigVGAN, and RMVPE."""
+ # Load CAMPPlus
+ campplus_ckpt_path = load_custom_model_from_hf("funasr/campplus", "campplus_cn_common.bin", config_filename=None)
+ self.campplus_model = CAMPPlus(feat_dim=80, embedding_size=192)
+ self.campplus_model.load_state_dict(torch.load(campplus_ckpt_path, map_location="cpu"))
+ self.campplus_model.eval()
+ self.campplus_model.to(self.device)
+
+ # Load BigVGAN models
+ self.bigvgan_model = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_22khz_80band_256x', use_cuda_kernel=False)
+ self.bigvgan_model.remove_weight_norm()
+ self.bigvgan_model = self.bigvgan_model.eval().to(self.device)
+
+ self.bigvgan_44k_model = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_44khz_128band_512x', use_cuda_kernel=False)
+ self.bigvgan_44k_model.remove_weight_norm()
+ self.bigvgan_44k_model = self.bigvgan_44k_model.eval().to(self.device)
+
+ # Load RMVPE for F0 extraction
+ model_path = load_custom_model_from_hf("lj1995/VoiceConversionWebUI", "rmvpe.pt", None)
+ self.rmvpe = RMVPE(model_path, is_half=False, device=self.device)
+
+ @staticmethod
+ def adjust_f0_semitones(f0_sequence, n_semitones):
+ """Adjust F0 values by a number of semitones."""
+ factor = 2 ** (n_semitones / 12)
+ return f0_sequence * factor
+
+ @staticmethod
+ def crossfade(chunk1, chunk2, overlap):
+ """Apply crossfade between two audio chunks."""
+ fade_out = np.cos(np.linspace(0, np.pi / 2, overlap)) ** 2
+ fade_in = np.cos(np.linspace(np.pi / 2, 0, overlap)) ** 2
+ if len(chunk2) < overlap:
+ chunk2[:overlap] = chunk2[:overlap] * fade_in[:len(chunk2)] + (chunk1[-overlap:] * fade_out)[:len(chunk2)]
+ else:
+ chunk2[:overlap] = chunk2[:overlap] * fade_in + chunk1[-overlap:] * fade_out
+ return chunk2
+
+ def _stream_wave_chunks(self, vc_wave, processed_frames, vc_target, overlap_wave_len,
+ generated_wave_chunks, previous_chunk, is_last_chunk, stream_output, sr):
+ """
+ Helper method to handle streaming wave chunks.
+
+ Args:
+ vc_wave: The current wave chunk
+ processed_frames: Number of frames processed so far
+ vc_target: The target mel spectrogram
+ overlap_wave_len: Length of overlap between chunks
+ generated_wave_chunks: List of generated wave chunks
+ previous_chunk: Previous wave chunk for crossfading
+ is_last_chunk: Whether this is the last chunk
+ stream_output: Whether to stream the output
+ sr: Sample rate
+
+ Returns:
+ Tuple of (processed_frames, previous_chunk, should_break, mp3_bytes, full_audio)
+ where should_break indicates if processing should stop
+ mp3_bytes is the MP3 bytes if streaming, None otherwise
+ full_audio is the full audio if this is the last chunk, None otherwise
+ """
+ mp3_bytes = None
+ full_audio = None
+
+ if processed_frames == 0:
+ if is_last_chunk:
+ output_wave = vc_wave[0].cpu().numpy()
+ generated_wave_chunks.append(output_wave)
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+ full_audio = (sr, np.concatenate(generated_wave_chunks))
+ else:
+ return processed_frames, previous_chunk, True, None, np.concatenate(generated_wave_chunks)
+
+ return processed_frames, previous_chunk, True, mp3_bytes, full_audio
+
+ output_wave = vc_wave[0, :-overlap_wave_len].cpu().numpy()
+ generated_wave_chunks.append(output_wave)
+ previous_chunk = vc_wave[0, -overlap_wave_len:]
+ processed_frames += vc_target.size(2) - self.overlap_frame_len
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+
+ elif is_last_chunk:
+ output_wave = self.crossfade(previous_chunk.cpu().numpy(), vc_wave[0].cpu().numpy(), overlap_wave_len)
+ generated_wave_chunks.append(output_wave)
+ processed_frames += vc_target.size(2) - self.overlap_frame_len
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+ full_audio = (sr, np.concatenate(generated_wave_chunks))
+ else:
+ return processed_frames, previous_chunk, True, None, np.concatenate(generated_wave_chunks)
+
+ return processed_frames, previous_chunk, True, mp3_bytes, full_audio
+
+ else:
+ output_wave = self.crossfade(previous_chunk.cpu().numpy(), vc_wave[0, :-overlap_wave_len].cpu().numpy(), overlap_wave_len)
+ generated_wave_chunks.append(output_wave)
+ previous_chunk = vc_wave[0, -overlap_wave_len:]
+ processed_frames += vc_target.size(2) - self.overlap_frame_len
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+
+ return processed_frames, previous_chunk, False, mp3_bytes, full_audio
+
+ def _process_whisper_features(self, audio_16k, is_source=True):
+ """Process audio through Whisper model to extract features."""
+ if audio_16k.size(-1) <= 16000 * 30:
+ # If audio is short enough, process in one go
+ inputs = self.whisper_feature_extractor(
+ [audio_16k.squeeze(0).cpu().numpy()],
+ return_tensors="pt",
+ return_attention_mask=True,
+ sampling_rate=16000
+ )
+ input_features = self.whisper_model._mask_input_features(
+ inputs.input_features, attention_mask=inputs.attention_mask
+ ).to(self.device)
+ outputs = self.whisper_model.encoder(
+ input_features.to(self.whisper_model.encoder.dtype),
+ head_mask=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ )
+ features = outputs.last_hidden_state.to(torch.float32)
+ features = features[:, :audio_16k.size(-1) // 320 + 1]
+ else:
+ # Process long audio in chunks
+ overlapping_time = 5 # 5 seconds
+ features_list = []
+ buffer = None
+ traversed_time = 0
+ while traversed_time < audio_16k.size(-1):
+ if buffer is None: # first chunk
+ chunk = audio_16k[:, traversed_time:traversed_time + 16000 * 30]
+ else:
+ chunk = torch.cat([
+ buffer,
+ audio_16k[:, traversed_time:traversed_time + 16000 * (30 - overlapping_time)]
+ ], dim=-1)
+ inputs = self.whisper_feature_extractor(
+ [chunk.squeeze(0).cpu().numpy()],
+ return_tensors="pt",
+ return_attention_mask=True,
+ sampling_rate=16000
+ )
+ input_features = self.whisper_model._mask_input_features(
+ inputs.input_features, attention_mask=inputs.attention_mask
+ ).to(self.device)
+ outputs = self.whisper_model.encoder(
+ input_features.to(self.whisper_model.encoder.dtype),
+ head_mask=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ )
+ chunk_features = outputs.last_hidden_state.to(torch.float32)
+ chunk_features = chunk_features[:, :chunk.size(-1) // 320 + 1]
+ if traversed_time == 0:
+ features_list.append(chunk_features)
+ else:
+ features_list.append(chunk_features[:, 50 * overlapping_time:])
+ buffer = chunk[:, -16000 * overlapping_time:]
+ traversed_time += 30 * 16000 if traversed_time == 0 else chunk.size(-1) - 16000 * overlapping_time
+ features = torch.cat(features_list, dim=1)
+
+ return features
+
+ @spaces.GPU
+ @torch.no_grad()
+ @torch.inference_mode()
+ def convert_voice(self, source, target, diffusion_steps=10, length_adjust=1.0,
+ inference_cfg_rate=0.7, f0_condition=False, auto_f0_adjust=True,
+ pitch_shift=0, stream_output=True):
+ """
+ Convert both timbre and voice from source to target.
+
+ Args:
+ source: Path to source audio file
+ target: Path to target audio file
+ diffusion_steps: Number of diffusion steps (default: 10)
+ length_adjust: Length adjustment factor (default: 1.0)
+ inference_cfg_rate: Inference CFG rate (default: 0.7)
+ f0_condition: Whether to use F0 conditioning (default: False)
+ auto_f0_adjust: Whether to automatically adjust F0 (default: True)
+ pitch_shift: Pitch shift in semitones (default: 0)
+ stream_output: Whether to stream the output (default: True)
+
+ Returns:
+ If stream_output is True, yields (mp3_bytes, full_audio) tuples
+ If stream_output is False, returns the full audio as a numpy array
+ """
+ # Select appropriate models based on F0 condition
+ inference_module = self.model if not f0_condition else self.model_f0
+ mel_fn = self.to_mel if not f0_condition else self.to_mel_f0
+ bigvgan_fn = self.bigvgan_model if not f0_condition else self.bigvgan_44k_model
+ sr = 22050 if not f0_condition else 44100
+ hop_length = 256 if not f0_condition else 512
+ max_context_window = sr // hop_length * 30
+ overlap_wave_len = self.overlap_frame_len * hop_length
+
+ # Load audio
+ source_audio = librosa.load(source, sr=sr)[0]
+ ref_audio = librosa.load(target, sr=sr)[0]
+
+ # Process audio
+ source_audio = torch.tensor(source_audio).unsqueeze(0).float().to(self.device)
+ ref_audio = torch.tensor(ref_audio[:sr * 25]).unsqueeze(0).float().to(self.device)
+
+ # Resample to 16kHz for feature extraction
+ ref_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000)
+ converted_waves_16k = torchaudio.functional.resample(source_audio, sr, 16000)
+
+ # Extract Whisper features
+ S_alt = self._process_whisper_features(converted_waves_16k, is_source=True)
+ S_ori = self._process_whisper_features(ref_waves_16k, is_source=False)
+
+ # Compute mel spectrograms
+ mel = mel_fn(source_audio.to(self.device).float())
+ mel2 = mel_fn(ref_audio.to(self.device).float())
+
+ # Set target lengths
+ target_lengths = torch.LongTensor([int(mel.size(2) * length_adjust)]).to(mel.device)
+ target2_lengths = torch.LongTensor([mel2.size(2)]).to(mel2.device)
+
+ # Compute style features
+ feat2 = torchaudio.compliance.kaldi.fbank(
+ ref_waves_16k,
+ num_mel_bins=80,
+ dither=0,
+ sample_frequency=16000
+ )
+ feat2 = feat2 - feat2.mean(dim=0, keepdim=True)
+ style2 = self.campplus_model(feat2.unsqueeze(0))
+
+ # Process F0 if needed
+ if f0_condition:
+ F0_ori = self.rmvpe.infer_from_audio(ref_waves_16k[0], thred=0.03)
+ F0_alt = self.rmvpe.infer_from_audio(converted_waves_16k[0], thred=0.03)
+
+ if self.device == "mps":
+ F0_ori = torch.from_numpy(F0_ori).float().to(self.device)[None]
+ F0_alt = torch.from_numpy(F0_alt).float().to(self.device)[None]
+ else:
+ F0_ori = torch.from_numpy(F0_ori).to(self.device)[None]
+ F0_alt = torch.from_numpy(F0_alt).to(self.device)[None]
+
+ voiced_F0_ori = F0_ori[F0_ori > 1]
+ voiced_F0_alt = F0_alt[F0_alt > 1]
+
+ log_f0_alt = torch.log(F0_alt + 1e-5)
+ voiced_log_f0_ori = torch.log(voiced_F0_ori + 1e-5)
+ voiced_log_f0_alt = torch.log(voiced_F0_alt + 1e-5)
+ median_log_f0_ori = torch.median(voiced_log_f0_ori)
+ median_log_f0_alt = torch.median(voiced_log_f0_alt)
+
+ # Shift alt log f0 level to ori log f0 level
+ shifted_log_f0_alt = log_f0_alt.clone()
+ if auto_f0_adjust:
+ shifted_log_f0_alt[F0_alt > 1] = log_f0_alt[F0_alt > 1] - median_log_f0_alt + median_log_f0_ori
+ shifted_f0_alt = torch.exp(shifted_log_f0_alt)
+ if pitch_shift != 0:
+ shifted_f0_alt[F0_alt > 1] = self.adjust_f0_semitones(shifted_f0_alt[F0_alt > 1], pitch_shift)
+ else:
+ F0_ori = None
+ F0_alt = None
+ shifted_f0_alt = None
+
+ # Length regulation
+ cond, _, codes, commitment_loss, codebook_loss = inference_module.length_regulator(
+ S_alt, ylens=target_lengths, n_quantizers=3, f0=shifted_f0_alt
+ )
+ prompt_condition, _, codes, commitment_loss, codebook_loss = inference_module.length_regulator(
+ S_ori, ylens=target2_lengths, n_quantizers=3, f0=F0_ori
+ )
+
+ # Process in chunks for streaming
+ max_source_window = max_context_window - mel2.size(2)
+ processed_frames = 0
+ generated_wave_chunks = []
+ previous_chunk = None
+
+ # Generate chunk by chunk and stream the output
+ while processed_frames < cond.size(1):
+ chunk_cond = cond[:, processed_frames:processed_frames + max_source_window]
+ is_last_chunk = processed_frames + max_source_window >= cond.size(1)
+ cat_condition = torch.cat([prompt_condition, chunk_cond], dim=1)
+
+ with torch.autocast(device_type=self.device.type, dtype=torch.float16):
+ # Voice Conversion
+ vc_target = inference_module.cfm.inference(
+ cat_condition,
+ torch.LongTensor([cat_condition.size(1)]).to(mel2.device),
+ mel2, style2, None, diffusion_steps,
+ inference_cfg_rate=inference_cfg_rate
+ )
+ vc_target = vc_target[:, :, mel2.size(-1):]
+
+ vc_wave = bigvgan_fn(vc_target.float())[0]
+
+ processed_frames, previous_chunk, should_break, mp3_bytes, full_audio = self._stream_wave_chunks(
+ vc_wave, processed_frames, vc_target, overlap_wave_len,
+ generated_wave_chunks, previous_chunk, is_last_chunk, stream_output, sr
+ )
+
+ if stream_output and mp3_bytes is not None:
+ yield mp3_bytes, full_audio
+
+ if should_break:
+ if not stream_output:
+ return full_audio
+ break
+
+ if not stream_output:
+ return np.concatenate(generated_wave_chunks)
+
+ return None, None
\ No newline at end of file