44 lines
1.1 KiB
Python
44 lines
1.1 KiB
Python
import numpy as np
|
|
import random
|
|
import time
|
|
from tinygrad import Tensor, nn
|
|
from tinygrad.nn.state import safe_load, load_state_dict
|
|
import librosa
|
|
import sounddevice as sd
|
|
from model import gen
|
|
from data import spec_to_audio
|
|
|
|
SAMPLE_RATE = 22050
|
|
|
|
def load_model(filepath="model.safetensors"):
|
|
"""Loads the model structure and weights."""
|
|
model = gen()
|
|
state_dict = safe_load(filepath)
|
|
load_state_dict(model, state_dict)
|
|
return model
|
|
|
|
def load_data(filepath="data.npz"):
|
|
"""Loads the pre-processed spectrogram data."""
|
|
print(f"Loading data from {filepath}...")
|
|
data = np.load(filepath)
|
|
x = data["arr_0"]
|
|
return x
|
|
|
|
def play_spec(spec,i):
|
|
"""Converts a spectrogram numpy array to audio and plays it."""
|
|
audio = spec_to_audio(spec)
|
|
sd.wait()
|
|
print(f"chunk:{i}")
|
|
sd.play(audio, samplerate=SAMPLE_RATE)
|
|
|
|
def run_prediction_loop(model, data_x):
|
|
current_spect = data_x[0:1]
|
|
for i in range(10):
|
|
play_spec(current_spect[0][0],i)
|
|
current_spect = model(Tensor(current_spect)).numpy()
|
|
|
|
if __name__ == "__main__":
|
|
model = load_model()
|
|
data_x = load_data()
|
|
run_prediction_loop(model, data_x)
|