Introduce code from Rust monorepo. This is a mess

2026-10-24 08:36:10 +00:00 · 2022-02-28 05:43:54 -05:00
parent 921fb49f9a
commit ba326bfcfd
1058 changed files with 192192 additions and 0 deletions
@@ -0,0 +1,335 @@
+import sys
+sys.path.append('waveglow/')
+import shutil
+import json
+
+import numpy as np
+import torch
+
+from tacotron2.hparams import create_hparams
+from tacotron2.model import Tacotron2
+from tacotron2.layers import TacotronSTFT
+from tacotron2.audio_processing import griffin_lim
+from tacotron2.text import text_to_sequence
+
+from waveglow.denoiser import Denoiser
+
+from scipy.io.wavfile import write as write_wav
+
+# Hifi-Gan
+from hifigan.env import AttrDict
+from hifigan.models import Generator
+from hifigan.denoiser import Denoiser as HifiDenoiser
+from hifigan.meldataset import mel_spectrogram
+import resampy
+import scipy
+
+# For metadata
+import subprocess
+import magic
+import os
+
+def print_gpu_info():
+    print('========================================')
+    print('Python interpreter', sys.executable)
+    print('PyTorch version', torch.__version__)
+    print('CUDA Available?', torch.cuda.is_available())
+    print('CUDA Device count', torch.cuda.device_count())
+    #try:
+    #    from tensorflow.python.client import device_lib
+    #    print('local devices', str(device_lib.list_local_devices()).replace("\n", "\n  "))
+    #except ImportError:
+    #    print('no tensorflow - cannot list devices')
+    #    pass
+    print('========================================', flush=True)
+
+def load_tacotron_model(checkpoint_path):
+    hparams = create_hparams()
+    hparams.sampling_rate = 22050 # Don't change this
+    hparams.max_decoder_steps = 1000 # How long the audio will be before it cuts off (1000 is about 11 seconds)
+    hparams.gate_threshold = 0.1 # Model must be 90% sure the clip is over before ending generation (the higher this number is, the more likely that the AI will keep generating until it reaches the Max Decoder Steps)
+    model = Tacotron2(hparams)
+    print('Loading synthesizer at path: {}'.format(checkpoint_path))
+    model.load_state_dict(torch.load(checkpoint_path)['state_dict'])
+    _ = model.cuda().eval().half()
+    # NB(bt, 2021-05-31): Trying to get everything on the same device
+    device = "cuda:0"
+    model = model.to(device)
+    return model
+
+def load_waveglow_model(checkpoint_path):
+    waveglow = torch.load(checkpoint_path)['model']
+    waveglow.cuda().eval().half()
+    for k in waveglow.convinv:
+        k.float()
+    denoiser = Denoiser(waveglow)
+    return waveglow
+
+def load_hifigan_model(checkpoint_path, super_resolution = False):
+    conf_name = 'config_32k' if super_resolution else 'config_v1'
+    conf = os.path.join("hifigan", conf_name + ".json")
+    with open(conf) as f:
+        json_config = json.loads(f.read())
+    h = AttrDict(json_config)
+    torch.manual_seed(h.seed)
+    hifigan = Generator(h).to(torch.device("cuda"))
+    state_dict_g = torch.load(checkpoint_path, map_location=torch.device("cuda"))
+    hifigan.load_state_dict(state_dict_g["generator"])
+    hifigan.eval()
+    hifigan.remove_weight_norm()
+    denoiser = HifiDenoiser(hifigan, mode="normal")
+    return hifigan, h, denoiser
+
+def preprocess_text(text, raw_input=True):
+    # TODO: Handle ARPAbet
+    pre_sequence = ""
+    for line in text.split("\n"):
+        if len(line) < 1:
+            continue
+        if raw_input:
+            if line[-1] != ";":
+                line = line + ";"
+        else:
+            #line = ARPA(line)
+            pass
+        pre_sequence += line
+
+    with torch.no_grad():
+        # save VRAM by not including gradients
+        sequence = np.array(text_to_sequence(pre_sequence, ['english_cleaners']))[None, :]
+        sequence = torch.autograd.Variable(torch.from_numpy(sequence)).cuda().long()
+    return sequence
+
+class NumpyEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return json.JSONEncoder.default(self, obj)
+
+def save_spectorgram_json_file(mel_outputs_postnet, output_spectrogram_filename):
+    # squeeze(0) -> 3D to 2D (by removing the singular 1-length "wrapper" dimension)
+    # transpose() -> originally 80x{N}, we turn to {N}x80
+    mel_for_json = mel_outputs_postnet.cpu().numpy().squeeze(0).transpose()
+    mel_for_scaling = mel_for_json.copy()
+
+    max_value = np.amax(mel_for_scaling)
+    min_value = np.amin(mel_for_scaling)
+    mel_range = max_value - min_value
+
+    mel_for_scaling -= min_value
+    mel_for_scaling *= (255.0 / mel_range)
+    mel_for_scaling = mel_for_scaling.astype('int32')
+
+    json_data = {
+        'mel': mel_for_json,
+        'mel_scaled': mel_for_scaling,
+    }
+
+    with open(output_spectrogram_filename, 'w') as outfile:
+        json.dump(json_data, outfile, cls=NumpyEncoder)
+
+def save_wav_audio_file(audio, output_audio_filename, sampling_rate=22050):
+    output_audio = audio[0].data.cpu().numpy().astype(np.float32)
+    write_wav(output_audio_filename, sampling_rate, output_audio)
+
+def generate_metadata_file(output_audio_filename, output_metadata_filename):
+    # ==== METADATA (1) ====
+    command = [
+        "ffprobe",
+        "-loglevel",  "quiet",
+        "-print_format", "json",
+        "-show_format",
+        "-show_streams",
+        output_audio_filename
+    ]
+
+    pipe = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    out, err = pipe.communicate()
+    ffmpeg_metadata = json.loads(out)
+
+    print(ffmpeg_metadata)
+
+    # ==== METADATA (2) ====
+    mime_type = magic.from_file(output_audio_filename, mime=True)
+    file_size_bytes = os.path.getsize(output_audio_filename)
+
+    duration_millis = int(float(ffmpeg_metadata['format']['duration']) * 1000)
+
+    metadata = {
+        'duration_millis': duration_millis,
+        'mimetype': mime_type,
+        'file_size_bytes': file_size_bytes,
+    }
+
+    with open(output_metadata_filename, 'w') as json_file:
+        json.dump(metadata, json_file)
+
+
+class TacotronWaveglowPipeline:
+
+    def __init__(self):
+        # Cache of "path" => loaded model
+        self.tacotron_model_cache = {}
+        self.waveglow = None
+        self.waveglow_filename = None
+
+        self.hifigan = None
+        self.hifigan_denoiser = None
+        self.hifigan_h = None
+        self.hifigan_filename = None
+        self.hifigan_super_resolution = None
+        self.hifigan_super_resolution_denoiser = None
+        self.hifigan_super_resolution_h = None
+        self.hifigan_super_resolution_filename = None
+
+    def maybe_load_hifigan(self, checkpoint_path):
+        if self.hifigan_filename == checkpoint_path:
+            print('Hifi-gan model already loaded into memory: {}'.format(checkpoint_path))
+            return
+        print('Dumping previous Hifi-gan model from memory: {}'.format(self.hifigan_filename))
+        print('Loading Hifi-gan model into memory: {}'.format(checkpoint_path), flush=True)
+        hifigan, h, denoiser = load_hifigan_model(checkpoint_path, False)
+        self.hifigan = hifigan
+        self.hifigan_denoiser = denoiser
+        self.hifigan_h = h
+        self.hifigan_filename = checkpoint_path
+
+    def maybe_load_hifigan_super_resolution(self, checkpoint_path):
+        if self.hifigan_super_resolution_filename == checkpoint_path:
+            print('Hifi-gan (SuperRes) model already loaded into memory: {}'.format(checkpoint_path))
+            return
+        print('Dumping previous Hifi-gan (SuperRes) model from memory: {}'.format(self.hifigan_super_resolution_filename))
+        print('Loading Hifi-gan (SuperRes) model into memory: {}'.format(checkpoint_path), flush=True)
+        hifigan, h, denoiser = load_hifigan_model(checkpoint_path, True)
+        self.hifigan_super_resolution = hifigan
+        self.hifigan_super_resolution_denoiser = denoiser
+        self.hifigan_super_resolution_h = h
+        self.hifigan_super_resolution_filename = checkpoint_path
+
+    def maybe_load_waveglow_model(self, checkpoint_path):
+        if self.waveglow_filename == checkpoint_path:
+            print('Waveglow model already loaded into memory: {}'.format(checkpoint_path))
+            return
+        print('Dumping previous waveglow model from memory: {}'.format(self.waveglow_filename))
+        print('Loading waveglow model into memory: {}'.format(checkpoint_path), flush=True)
+        self.waveglow = load_waveglow_model(checkpoint_path)
+        self.waveglow_filename = checkpoint_path
+
+    def maybe_load_tacotron_model_to_cache(self, checkpoint_path):
+        if checkpoint_path in self.tacotron_model_cache:
+            print('Tacotron model already loaded into memory: {}'.format(checkpoint_path))
+            return
+        print('Loading tacotron model into memory: {}'.format(checkpoint_path), flush=True)
+        model = load_tacotron_model(checkpoint_path)
+        self.tacotron_model_cache[checkpoint_path] = model
+        print('Now {} synthesizer models in memory'.format(len(self.tacotron_model_cache)))
+
+    def uncache_tacotron_model(self, checkpoint_path):
+        print('Currently {} synthesizer models in memory'.format(
+            len(self.tacotron_model_cache)))
+        print('Clearing synthesizer from memory: {}'.format(checkpoint_path))
+        self.tacotron_model_cache.pop(checkpoint_path, None)
+        print('Now reduced to {} synthesizer models in memory'.format(
+            len(self.tacotron_model_cache)))
+
+    def infer(self, args):
+        assert('synthesizer_checkpoint_path' in args)
+        assert('vocoder_type' in args)
+        assert('raw_text' in args)
+        assert('output_audio_filename' in args)
+        assert('output_spectrogram_filename' in args)
+        assert('output_metadata_filename' in args)
+
+        print('Raw text: {}'.format(args['raw_text']))
+
+        sequence = preprocess_text(args['raw_text'])
+
+        tacotron = self.tacotron_model_cache[args['synthesizer_checkpoint_path']]
+
+        with torch.no_grad():
+            mel_outputs, mel_outputs_postnet, _, alignments = tacotron.inference(sequence)
+
+        if args['vocoder_type'] == 'hifigan-superres':
+            self.vocode_hifigan_superres(args, mel_outputs_postnet)
+        elif args['vocoder_type'] == 'waveglow':
+            self.vocode_waveglow(args, mel_outputs_postnet)
+        else:
+            raise Exception('Wrong vocoder type: {}'.format(args['vocoder_type']))
+
+        print('Saving spectrogram as JSON...')
+        save_spectorgram_json_file(mel_outputs_postnet, args['output_spectrogram_filename'])
+
+        print('Generating metadata file...')
+        generate_metadata_file(args['output_audio_filename'], args['output_metadata_filename'])
+
+    def vocode_waveglow(self, args, mel_outputs_postnet):
+        print('Running waveglow...')
+        with torch.no_grad():
+            sigma = 0.8
+            audio = self.waveglow.infer(mel_outputs_postnet, sigma=sigma)
+
+        print('Encoding and saving audio...')
+        save_wav_audio_file(audio, args['output_audio_filename'])
+
+    def vocode_hifigan_superres(self, args, mel_outputs_postnet):
+        print('Running hifigan...')
+        with torch.no_grad():
+            y_g_hat = self.hifigan(mel_outputs_postnet.float())
+
+            MAX_WAV_VALUE = 32768.0
+            audio = y_g_hat.squeeze()
+            audio = audio * MAX_WAV_VALUE
+            audio_denoised = self.hifigan_denoiser(audio.view(1, -1), strength=35)[:, 0]
+
+            # Resample to 32k
+            audio_denoised = audio_denoised.cpu().numpy().reshape(-1)
+
+            normalize = (MAX_WAV_VALUE / np.max(np.abs(audio_denoised))) ** 0.9
+            audio_denoised = audio_denoised * normalize
+            wave = resampy.resample(
+                audio_denoised,
+                self.hifigan_h.sampling_rate,
+                self.hifigan_super_resolution_h.sampling_rate,
+                filter="sinc_window",
+                window=scipy.signal.windows.hann,
+                num_zeros=8,
+            )
+            wave_out = wave.astype(np.int16)
+
+            # HiFi-GAN super-resolution
+            print('Running hifigan super resolution...')
+            wave = wave / MAX_WAV_VALUE
+            wave = torch.FloatTensor(wave).to(torch.device("cuda"))
+            new_mel = mel_spectrogram(
+                wave.unsqueeze(0),
+                self.hifigan_super_resolution_h.n_fft,
+                self.hifigan_super_resolution_h.num_mels,
+                self.hifigan_super_resolution_h.sampling_rate,
+                self.hifigan_super_resolution_h.hop_size,
+                self.hifigan_super_resolution_h.win_size,
+                self.hifigan_super_resolution_h.fmin,
+                self.hifigan_super_resolution_h.fmax,
+            )
+            y_g_hat2 = self.hifigan_super_resolution(new_mel)
+            audio2 = y_g_hat2.squeeze()
+            audio2 = audio2 * MAX_WAV_VALUE
+            audio2_denoised = self.hifigan_denoiser(audio2.view(1, -1), strength=35)[:, 0]
+
+            # High-pass filter, mixing and denormalizing
+            audio2_denoised = audio2_denoised.cpu().numpy().reshape(-1)
+            b = scipy.signal.firwin(
+                101, cutoff=10500, fs=self.hifigan_super_resolution_h.sampling_rate, pass_zero=False
+            )
+            y = scipy.signal.lfilter(b, [1.0], audio2_denoised)
+            superres_strength = 4.0
+            y *= superres_strength
+            y_out = y.astype(np.int16)
+            y_padded = np.zeros(wave_out.shape)
+            y_padded[: y_out.shape[0]] = y_out
+            sr_mix = wave_out + y_padded
+            sr_mix = sr_mix / normalize
+
+            print('Encoding and saving audio...')
+            rate = self.hifigan_super_resolution_h.sampling_rate
+            output_audio = sr_mix.astype(np.int16)
+            write_wav(args['output_audio_filename'], rate, output_audio)