# 互換性のために forward を extract_features として定義 
cmodel.forward = cmodel.extract_features

OUTPUT_DIR = Path("output") 
BASE_MODEL_NAME = "cmodel" 

OUTPUT_DIR.mkdir(exist_ok=True) 

ir_cmodel_path = Path(OUTPUT_DIR / (BASE_MODEL_NAME + "_ir")).with_suffix(".xml") 

length = 32000 

dummy_input = torch.randn(1, length)

core = ov.Core() 

class ModelWrapper(torch.nn.Module): 
    def __init__(self, model): 
        super().__init__() 
        self.model = model 

    def forward(self, input): 
        return self.model(input)[0] 

if not ir_cmodel_path.exists(): 
    ir_cmodel = ov.convert_model(ModelWrapper(cmodel), example_input=dummy_input) 
    ov.save_model(ir_cmodel, ir_cmodel_path) 
else: 
    ir_cmodel = core.read_model(ir_cmodel_path)

import ipywidgets as widgets 

device = widgets.Dropdown( 
    options=core.available_devices + ["AUTO"], 
    value="AUTO", 
    description="Device:", 
    disabled=False, 
) 

device

compiled_cmodel = core.compile_model(ir_cmodel, device.value)

OUTPUT_DIR = Path("output") 
BASE_MODEL_NAME = "smodel" 

OUTPUT_DIR.mkdir(exist_ok=True) 

ir_smodel_path = Path(OUTPUT_DIR / (BASE_MODEL_NAME + "ir")).with_suffix(".xml") 

length = 32000 

dummy_input = torch.randn(1, length, 40) 

if not ir_smodel_path.exists(): 
    ir_smodel = ov.convert_model(smodel, example_input=dummy_input) ov.save_model(ir_smodel, ir_smodel_path) 
else: 
    ir_smodel = core.read_model(ir_smodel_path)

from speaker_encoder.hparams import sampling_rate, mel_window_step, partials_n_frames 
from speaker_encoder import audio 

def compute_partial_slices(n_samples: int, rate, min_coverage): 
    """ 
    Computes where to split an utterance waveform and its corresponding mel spectrogram to 
    obtain partial utterances of <partials_n_frames> each.Both the waveform and the 
    mel spectrogram slices are returned, so as to make each partial utterance waveform 
    correspond to its spectrogram.

    The returned ranges may be indexing further than the length of the waveform. It is
    recommended that you pad the waveform with zeros up to wav_slices[-1].stop.
 
    :param n_samples: the number of samples in the waveform 
    :param rate: how many partial utterances should occur per second.Partial utterances must 
    cover the span of the entire utterance, thus the rate should not be lower than the inverse 
    of the duration of a partial utterance.By default, partial utterances are 1.6s long and 
    the minimum rate is thus 0.625.
    :param min_coverage: when reaching the last partial utterance, it may or may not have 
    enough frames.If at least <min_pad_coverage> of <partials_n_frames> are present, 
    then the last partial utterance will be considered by zero-padding the audio.Otherwise, 
    it will be discarded.If there aren't enough frames for one partial utterance, 
    this parameter is ignored so that the function always returns at least one slice. 
    :return: the waveform slices and mel spectrogram slices as lists of array slices.Index 
    respectively the waveform and the mel spectrogram with these slices to obtain the partial 
    utterances.
    """ 
    assert 0 < min_coverage <= 1 

    # 2 つの部分的な発話の間に何フレームあるかを計算 
    samples_per_frame = int((sampling_rate * mel_window_step / 1000)) 
    n_frames = int(np.ceil((n_samples + 1) / samples_per_frame)) 
    frame_step = int(np.round((sampling_rate / rate) / samples_per_frame)) 
    assert 0 < frame_step, "The rate is too high" 
    assert frame_step <= partials_n_frames, "The rate is too low, it should be %f at least" % (sampling_rate / (samples_per_frame * partials_n_frames)) 

    # スライスを計算 
    wav_slices, mel_slices = [], [] 
    steps = max(1, n_frames - partials_n_frames + frame_step + 1) 
    for i in range(0, steps, frame_step): 
        mel_range = np.array([i, i + partials_n_frames]) 
        wav_range = mel_range * samples_per_frame 
        mel_slices.append(slice(*mel_range)) 
        wav_slices.append(slice(*wav_range)) 

    # 追加のパディングが必要であるか評価 
    last_wav_range = wav_slices[-1] 
    coverage = (n_samples - last_wav_range.start) / (last_wav_range.stop - last_wav_range.start) 
    if coverage < min_coverage and len(mel_slices) > 1: 
        mel_slices = mel_slices[:-1] 
        wav_slices = wav_slices[:-1] 

    return wav_slices, mel_slices 

def embed_utterance( 
    wav: np.ndarray, 
    smodel: ov.CompiledModel, 
    return_partials=False, rate=1.3, 
    min_coverage=0.75, 
): 
    """ 
    Computes an embedding for a single utterance.The utterance is divided in partial 
    utterances and an embedding is computed for each.The complete utterance embedding is the 
    L2-normed average embedding of the partial utterances.

    :param wav: a preprocessed utterance waveform as a numpy array of float32 
    :param smodel: compiled speaker encoder model.
    :param return_partials: if True, the partial embeddings will also be returned along with 
    the wav slices corresponding to each partial utterance. 
    :param rate: how many partial utterances should occur per second.Partial utterances must 
    cover the span of the entire utterance, thus the rate should not be lower than the inverse 
    of the duration of a partial utterance.By default, partial utterances are 1.6s long and 
    the minimum rate is thus 0.625.
    :param min_coverage: when reaching the last partial utterance, it may or may not have 
    enough frames. If at least <min_pad_coverage> of <partials_n_frames> are present, 
    then the last partial utterance will be considered by zero-padding the audio.Otherwise, 
    it will be discarded.If there aren't enough frames for one partial utterance, 
    this parameter is ignored so that the function always returns at least one slice. 
    :return: the embedding as a numpy array of float32 of shape (model_embedding_size,).If 
    <return_partials> is True, the partial utterances as a numpy array of float32 of shape 
    (n_partials, model_embedding_size) and the wav partials as a list of slices will also be 
    returned.
    """ 
    # 発話をどこで部分的に分割するかを計算し、部分的な発話が 
    # より広い範囲をカバーする場合は波形をゼロで埋める 
    wav_slices, mel_slices = compute_partial_slices(len(wav), rate, min_coverage) 
    max_wave_length = wav_slices[-1].stop 
    if max_wave_length >= len(wav): 
        wav = np.pad(wav, (0, max_wave_length - len(wav)), "constant") 

    # 発話を部分的に分割してモデルに転送 
    mel = audio.wav_to_mel_spectrogram(wav) 
    mels = np.array([mel[s] for s in mel_slices]) 
    with torch.no_grad(): 
        mels = torch.from_numpy(mels).to(torch.device("cpu")) 
        output_layer = smodel.output(0) 
        partial_embeds = smodel(mels)[output_layer] 

    # 部分埋め込みから発話埋め込みを計算 
    raw_embed = np.mean(partial_embeds, axis=0) 
    embed = raw_embed / np.linalg.norm(raw_embed, 2) 

    if return_partials: 
        return embed, partial_embeds, wav_slices 
    return embed

device

compiled_smodel = core.compile_model(ir_smodel, device.value)

OUTPUT_DIR = Path("output") 
BASE_MODEL_NAME = "net_g" 
onnx_net_g_path = Path(OUTPUT_DIR / (BASE_MODEL_NAME + "_fp32")).with_suffix(".onnx") 
ir_net_g_path = Path(OUTPUT_DIR / (BASE_MODEL_NAME + "ir")).with_suffix(".xml") 

dummy_input_1 = torch.randn(1, 1024, 81) 
dummy_input_2 = torch.randn(1, 256) 

# 前方を推測として定義
net_g.forward = net_g.infer 

if not ir_net_g_path.exists(): 
    ir_net_g_model = ov.convert_model(net_g, example_input=(dummy_input_1, dummy_input_2)) 
    ov.save_model(ir_net_g_model, ir_net_g_path) 
else: 
    ir_net_g_model = core.read_model(ir_net_g_path)

compiled_ir_net_g_model = core.compile_model(ir_net_g_model, device.value)

def synthesize_audio(src, tgt): 
    wav_tgt, _ = librosa.load(tgt, sr=hps.data.sampling_rate) 
    wav_tgt, _ = librosa.effects.trim(wav_tgt, top_db=20) 

    g_tgt = embed_utterance(wav_tgt, compiled_smodel) 
    g_tgt = torch.from_numpy(g_tgt).unsqueeze(0) 

    # src 
    wav_src, _ = librosa.load(src, sr=hps.data.sampling_rate) 
    wav_src = np.expand_dims(wav_src, axis=0) 

    output_layer = compiled_cmodel.output(0) 
    c = compiled_cmodel(wav_src)[output_layer] 
    c = c.transpose((0, 2, 1)) 

    output_layer = compiled_ir_net_g_model.output(0) 
    tgt_audio = compiled_ir_net_g_model((c, g_tgt))[output_layer] 
    tgt_audio = tgt_audio[0][0] 

    return tgt_audio

result_wav_names = [] 

with torch.no_grad(): 
    for line in tqdm(zip(srcs, tgts)): 
        src, tgt = line 

        output_audio = synthesize_audio(src, tgt) 

        timestamp = time.strftime("%m-%d_%H-%M", time.localtime()) 
        result_name = f"{timestamp}.wav" 
        result_wav_names.append(result_name) 
        write( 
            os.path.join("outputs/freevc", result_name), 
            hps.data.sampling_rate, 
            output_audio, 
        )

import IPython.display as ipd 

ipd.Audio(srcs[0])

ipd.Audio(tgts[0])

ipd.Audio(f"outputs/freevc/{result_wav_names[0]}")

import gradio as gr 

audio1 = gr.Audio(label="Source Audio", type="filepath") 
audio2 = gr.Audio(label="Reference Audio", type="filepath") 
outputs = gr.Audio(label="Output Audio", type="filepath") 
examples = [[audio1_name, audio2_name]] 

title = "FreeVC with Gradio" 
description = 'Gradio Demo for FreeVC and OpenVINO™.Upload a source audio and a reference audio, then click the "Submit" button to inference.' 

def infer(src, tgt): 
    output_audio = synthesize_audio(src, tgt) 

    timestamp = time.strftime("%m-%d_%H-%M", time.localtime()) 
    result_name = f"{timestamp}.wav" 
    write(result_name, hps.data.sampling_rate, output_audio) 

    return result_name 

iface = gr.Interface( 
    infer, 
    [audio1, audio2], 
    outputs, 
    title=title, 
    description=description, 
    examples=examples, 
) 
iface.launch() 
# リモートで起動する場合は、server_name と server_port を指定 
# iface.launch(server_name='your server name', server_port='server port in int') 
# プラットフォーム上で起動する際に問題がある場合は、起動メソッドに share=True を渡すことができます: 
# iface.launch(share=True) 
# インターフェイスの公開共有可能なリンクを作成。詳細はドキュメントをご覧ください: https://gradio.app/docs/ 

iface.close()