# -*- coding: utf-8 -*-
"""softvc_pynb.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1FzpL3Ze3X9ESt_nB4w9rTgnfGJxunha0
"""

!nvidia-smi

#@title Download code
!git clone https://github.com/bshall/acoustic-model
!git clone https://github.com/bshall/hubert
!git clone https://github.com/bshall/hifigan

#@title Download and apply patches
!pip install gdown
import gdown
gdown.download(id="1tcaBRORA627ZGoyMqPEGbmW5Hr2PBDXK", output="/content/acoustic-model-train.patch")
!patch -p2 < /content/acoustic-model-train.patch

#@title Create directories and setup the colab. Grab the models that are needed for this.
from pathlib import Path
import os
import requests


def download_url(url:str,path:Path):
    url_file_name = url.split("/")[-1]
    response = requests.get(url)
    with open(path / Path(f"{url_file_name}"), 'wb') as f:
        f.write(response.content)

dataset = Path("./dataset")

if dataset.exists() == False:
    os.mkdir(dataset)

wavs = Path("./dataset/wavs")
if wavs.exists() == False:
    os.mkdir(wavs)

mels = Path("./dataset/mels")
if mels.exists() == False:
    os.mkdir(mels)

train = Path("./dataset/wavs/train")
if train.exists() == False:
    os.mkdir(train)

dev = Path("./dataset/wavs/dev")
if dev.exists() == False:
    os.mkdir(dev)


check_point = Path("./checkpoint/")
if check_point.exists() == False:
    os.mkdir(check_point)

batch_convert_wavs = Path("./batch_convert/")
if check_point.exists() == False:
    os.mkdir(batch_convert_wavs)


download_url("https://github.com/bshall/acoustic-model/releases/download/v0.1/hubert-soft-0321fd7e.pt",check_point)
download_url("https://github.com/bshall/acoustic-model/releases/download/v0.1/hubert-discrete-d49e1c77.pt",check_point)

print("Models and Directories Setup")

#@title Mount Google Drive
from google.colab import drive
drive.mount('/content/drive')

#@title Unzip dataset
#@markdown the inside of the dataset zip should be structured like this:
#@markdown ```
#@markdown /wavs/
#@markdown ------/train
#@markdown -------------/1.wav
#@markdown -------------/2.wav
#@markdown -------------/etc.wav
#@markdown ------/dev
#@markdown ------------/3.wav
#@markdown ------------/4.wav
#@markdown ------------/etc.wav
#@markdown ```
#@markdown - dev is the validation dataset and should be roughly 1/10th of the training data
#@markdown - the wav names don't matter
#@markdown - the training data should be at least half an hour long for decent results
    
import zipfile
path_to_zipfile = "/content/drive/MyDrive/tts/datasets/bshall-JGSpeech-1.7.zip"#@param { type: "string"}
with zipfile.ZipFile(path_to_zipfile, 'r') as zip_ref:
    zip_ref.extractall("/content/dataset")

!ls /content/dataset/wavs/dev | wc -l
!ls /content/dataset/wavs/train | wc -l
!du -sh /content/dataset/wavs/dev
!du -sh /content/dataset/wavs/train

#@title Convert to mels
!python ./acoustic-model/mels.py ./dataset/wavs ./dataset/mels/ 
!python ./hubert/encode.py --extension .wav soft ./dataset/wavs ./dataset/soft



#@title Start training
#@markdown - Edit hyperparameters in /content/acoustic-model/train.py before you run this cell
#@markdown - to resume from a previous session, change resume_path to your previous checkpoint
#@markdown - If this cell hangs with no output, restart the runtime

import os, sys, argparse
from pathlib import Path
import torch
import torch.distributed as dist

sys.path.append("/content/acoustic-model/")
from train import train

checkpoint_save_path = "/content/drive/MyDrive/tts/softvc-checkpoints/jeff_v2/acoustic" #@param {type: "string"}
_checkpoint_save_path = Path(checkpoint_save_path)
if _checkpoint_save_path.exists() == False:
    os.mkdir(_checkpoint_save_path)

resume_path = "/content/checkpoint/hubert-soft-0321fd7e.pt" #@param {type: "string"}

world_size = torch.cuda.device_count()


class DictObj:
  def __init__(self, in_dict:dict):
    assert isinstance(in_dict, dict)
    for key, val in in_dict.items():
      if isinstance(val, (list, tuple)):
        setattr(self, key, [DictObj(x) if isinstance(x, dict) else x for x in val])
      else:
        setattr(self, key, DictObj(val) if isinstance(val, dict) else val)

args=DictObj({"checkpoint_dir": _checkpoint_save_path, "dataset_dir": Path("/content/dataset"), "resume":Path(resume_path), "discrete": False})
try:
  dist.destroy_process_group()
except:
  pass
train(0, world_size, args)

"""# Inference"""

#@title Load the models 
#@markdown leave hifigan blank to use defaults
import sys
import torch
from typing import Dict, Any
from collections import OrderedDict
from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present

def modify_model(state_dict:Dict[str,Any]):
    new_state_dict = OrderedDict()
    for k in state_dict['acoustic-model'].keys():
        value = state_dict['acoustic-model'][k]
        name = k[7:] # remove `module.` because they used ddp training to save
        new_state_dict[name] = value
    return new_state_dict

sys.path.append("/content/acoustic-model/")
from acoustic import AcousticModel

hubert = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda()

check_point_path = '/content/drive/MyDrive/tts/softvc-checkpoints/tay/acoustic/taylor_swift-acoustic.pt' #@param {type:"string"}

# Load the acoustic model (either hubert_soft or hubert_discrete) 
acoustic_dict = torch.load(check_point_path) # this used ddp
acoustic = AcousticModel().cuda()
acoustic.eval()

import importlib.util
import sys

spec = importlib.util.spec_from_file_location("hifigan.utils", "/content/hifigan/hifigan/utils.py")
u = importlib.util.module_from_spec(spec)
sys.modules["hifigan.utils"] = u
spec.loader.exec_module(u)

spec = importlib.util.spec_from_file_location("hifigan.generator", "/content/hifigan/hifigan/generator.py")
g = importlib.util.module_from_spec(spec)
sys.modules["hifigan.generator"] = g
spec.loader.exec_module(g)

new_state_dict = modify_model(state_dict=acoustic_dict)
acoustic.load_state_dict(new_state_dict)
hifigan_checkpoint_path = "/content/drive/MyDrive/tts/softvc-checkpoints/tay/hifigan/taylor_swift-hifigan.pt" #@param {type: "string"}
if hifigan_checkpoint_path:
  checkpoint = torch.load(hifigan_checkpoint_path)["generator"]["model"]
  hifigan = g.HifiganGenerator()
  consume_prefix_in_state_dict_if_present(checkpoint, "module.")
  hifigan.load_state_dict(checkpoint)
  hifigan.remove_weight_norm()
  hifigan.eval()
 
else:
  hifigan = torch.hub.load("bshall/hifigan:main", "hifigan_hubert_soft")
hifigan.cuda()

#@title This code selects the source file
import IPython.display as display
import torch, torchaudio
import sys  
sys.path.insert(0, './acoustic-model/acoustic/')
from model import AcousticModel
import soundfile as sf
from google.colab import files
from time import time
sys.path.append("/content/hifigan/")
from hifigan.generator import HifiganGenerator
from hifigan.generator import consume_prefix_in_state_dict_if_present
#uploaded = files.upload()
file_name = "/content/drive/MyDrive/20220914110418-9dd00c7845-godzilla_mdxKaraoke_vocals.[mvsep.com].wav" #@param {type: "string"}

#for name in uploaded.keys():
    #file_name = name    
source, sr = torchaudio.load(f"{file_name}")
source = torchaudio.functional.resample(source, sr, 16000)
source = torch.mean(source, dim=0).unsqueeze(0)
source = source.unsqueeze(0).cuda()
display.Audio(source.squeeze().cpu(), rate=16000)

#@title Run the inference and display the result
# Convert to the target speaker
with torch.inference_mode():
    # Extract speech units
    start = time()
    units = hubert.units(source)
    end = time()
    # Generate target spectrogram
    print(f"time for feature embedding:{end-start}")

    start = time()
    mel = acoustic.generate(units).transpose(1, 2)
    end = time()
    print(f"time for inference:{end-start}")
    # Generate audio waveform
    start = time()
    target = hifigan(mel)
    end = time()
    print(f"time for wavform {end-start}")

display.Audio(target.squeeze().cpu(), rate=16000)

#@title Save the output to file
filename = "bruce.wav" #@param{type: "string"}
sf.write(filename, target.squeeze().cpu(), samplerate=16000)

#@title Batch convert
def batch_convert_wavs(folder:Path):
    # Load the content encoder (either hubert_soft or hubert_discrete)
    hubert = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda()

    # Load the acoustic model (either hubert_soft or hubert_discrete) 
    acoustic_dict = torch.load(check_point_path) # this used ddp
    acoustic = AcousticModel().cuda()
    acoustic.eval()

    new_state_dict = modify_model(state_dict=acoustic_dict)
    acoustic.load_state_dict(new_state_dict)

    # Load the vocoder (either hifigan_hubert_soft or hifigan_hubert_discrete)
    hifigan = torch.hub.load("bshall/hifigan:main", "hifigan_hubert_soft").cuda()

    
    # Convert to the target speaker
    with torch.inference_mode():
        # Extract speech units
        units = hubert.units(source)
        # Generate target spectrogram
        mel = acoustic.generate(units).transpose(1, 2)
        # Generate audio waveform
        target = hifigan(mel)

    # TODO later.