#%% Setup.
print('Importing.')
import signal
import sys

import numpy as np
import scipy.io.wavfile

from keras.utils.visualize_util import plot
from keras.callbacks import TensorBoard, ModelCheckpoint

from tqdm import tqdm

from eva.models.wavenet import Wavenet, compute_receptive_field

from eva.util.mutil import sparse_labels
from eva.util.autil import to_pcm8

#%% Generation Config.
print('Preparing generation.')
RATE = 4000
UNITS = 356415

#%% Model Config.
print('Preparing the model.')
MODEL = Wavenet
FILTERS = 32
DEPTH = 9
STACKS = 2
BINS = 256
LAST = RATE
LENGTH = LAST + compute_receptive_field(RATE, DEPTH, STACKS)[0]


#%% Model.
print('Loading the model.')
INPUT = (LENGTH, BINS)
ARGS = (INPUT, FILTERS, DEPTH, STACKS)

M = MODEL(*ARGS)
M.load_weights('model.h5')

M.summary()

plot(M)

#%% Generate.
print('Generating.')
def save():
    print('Saving.')
    np.save('samples.npy', samples)
    np.save('audio.npy', audio)

    pcm8_64bit_wide = to_pcm8(audio)

    for i in tqdm(range(audio.shape[0])):
        scipy.io.wavfile.write('audio{}.wav'.format(i), RATE, pcm8_64bit_wide[i])

def save_gracefully(signal, frame):
    save()
    sys.exit(0)

signal.signal(signal.SIGINT, save_gracefully)

BATCH_SIZE = 3

samples = np.zeros(shape=(BATCH_SIZE, UNITS, BINS))
audio = np.zeros(shape=(BATCH_SIZE, UNITS))
for i in tqdm(range(UNITS+LENGTH-1)):
    if i >= UNITS-1:
        break
    x = i+LENGTH-1
    samples[:, x] = M.predict(samples[:, i:i+LENGTH], batch_size=1)[0]
    samples[:,x,np.argmax(samples[:,x], axis=-1)] += 1-np.sum(samples[:, x], axis=-1)
    audio[:, x] = np.array([np.random.choice(256, p=p) for p in samples[:, x]])
    if i % (RATE//2) == 0:
        print(str(i/(RATE)) + " Seconds generated!")

save()
@@ -532,9 +532,10 @@
 H = 
-9
+10
 %0ASTA
@@ -540,17 +540,17 @@
 TACKS = 
-2
+5
 %0ABINS = 
@@ -1224,17 +1224,17 @@
 _SIZE = 
-3
+1
 %0A%0Asample
@@ -1458,24 +1458,19 @@
 es%5B:, i:
-i+LENGTH
+x+1
 %5D, batch
@@ -1478,17 +1478,21 @@
 size=1)%5B
-0
+:, -1
 %5D%0A    sa

Fix and update inference args