# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '1'

# +
import sys

SOURCE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__name__)))
sys.path.insert(0, SOURCE_DIR)

# +
import pickle

metadata = pickle.load(open('metadata.pkl', "rb"))
# -

import numpy as np
from math import ceil
import tensorflow as tf
from malaya_speech.train.model import fastspeech, fastvc
import malaya_speech

sbmt_i = metadata[0]
x_org = sbmt_i[2]
c = np.random.normal(size = (2, 512)).astype(np.float32)

mel = tf.placeholder(tf.float32, (None, None, 80))
ori_vector = tf.placeholder(tf.float32, (None, 512))
target_vector = tf.placeholder(tf.float32, (None, 512))
mel_lengths = tf.placeholder(tf.int32, (None))

dim_neck = 32
dim_speaker = 256
config = malaya_speech.config.fastspeech_config
config['encoder_hidden_size'] = dim_speaker + 80
config['decoder_hidden_size'] = dim_speaker + dim_neck
config = fastspeech.Config(vocab_size = 1, **config)

model = fastvc.model.Model(dim_neck, config, dim_speaker = dim_speaker)

encoder_outputs, mel_before, mel_after, codes = model(mel, ori_vector, target_vector, mel_lengths)

codes_ = model.call_second(mel_after, ori_vector, mel_lengths)

encoder_outputs, mel_before, mel_after, codes, codes_

loss_f = tf.losses.mean_squared_error
max_length = tf.cast(tf.reduce_max(mel_lengths), tf.int32)
mask = tf.sequence_mask(
    lengths = mel_lengths, maxlen = max_length, dtype = tf.float32
)
mask = tf.expand_dims(mask, axis = -1)
mask

# +
from functools import partial
from malaya_speech.train.loss import calculate_2d_loss, calculate_3d_loss

mse_mel = partial(
    loss_f,
    weights = mask
)
# -

mel_loss_before = calculate_3d_loss(mel, mel_before, mse_mel)
mel_loss_before

mel_loss_after = calculate_3d_loss(mel, mel_after, mse_mel)
mel_loss_after

g_loss_cd = tf.losses.absolute_difference(codes, codes_)

loss = mel_loss_before + mel_loss_after + g_loss_cd

sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())

o = sess.run([encoder_outputs, mel_before, mel_after], feed_dict = {mel: [x_org] * 2,
                                                                    ori_vector: c,
                                                                    target_vector: c,
                                                                    mel_lengths: [len(x_org)] * 2})

o[0].shape, o[1].shape, o[2].shape

o = sess.run([mel_loss_before, mel_loss_after, g_loss_cd], feed_dict = {mel: [x_org] * 2,
                                                                    ori_vector: c,
                                                                    target_vector: c,
                                                                    mel_lengths: [len(x_org)] * 2})

o

tf.trainable_variables()

saver = tf.train.Saver(tf.trainable_variables())
saver.save(sess, "test/model.ckpt")

# !ls -lh test

# !rm -rf test

StarcoderdataJupyterScriptsDedupFiltered