token_telephone/tt.py

from ttutil import hsv_to_rgb, dbg, log, set_debug, pow2db, db2pow
from dataclasses import dataclass, field
import os
from pathlib import Path
import random
import time
from threading import Thread
import gc
gc.disable()

import sounddevice as sd

from blessed import Terminal

import numpy as np
import torch
from einops import rearrange

PROFILE = False
DEBUG = False
DEBUG_NO_VAMPNET = False
set_debug(DEBUG)
# if DEBUG:
# import gc
# # log when gc start and stops
# gc.set_debug(gc.DEBUG_STATS)

@dataclass 
class LoadState:
 t0: float = None
 loaded: bool = False

load_state = LoadState() 

def on_random_color():
 def random_rgb_bg():
 return np.random.randint(0, 255), np.random.randint(0, 255), np.random.randint(0, 255)
 return term.on_color_rgb(*random_rgb_bg())

# draw the intro screen before slow imports
def color_tokenize_txt(text: str):
 # apply a random bg color to each letter
 return "".join(on_random_color()(letter) for letter in text)

def color_tokenize_words(text: str):
 return " ".join(on_random_color()(word) for word in text.split(" "))

def draw_intro_screen():
 global load_state
 load_state.t0 = time.time()
 avg_time = 20 # average loading time

 while not load_state.loaded: 
 print(term.clear)
 print(term.move_xy(0, 1) + term.center(color_tokenize_words("hugo flores garcía")))
 print(term.move_xy(0, 3) + term.center(color_tokenize_words("and")))
 print(term.move_xy(0, 5) + term.center(color_tokenize_words("stephan moore")))
 print(term.move_xy(0, 7) + term.center(color_tokenize_words("present")))
 print(term.move_xy(0, 9) + term.center(term.bold(color_tokenize_txt("token telephone"))))

 # print(term.move_xy(0, 10) + term.center(color_tokenize_txt("loading ")), end="")
 # make a little loading bar
 elapsed = time.time() - load_state.t0
 num_dots = int((elapsed / avg_time) * 20)
 num_spaces = 20 - num_dots
 print(term.move_xy(0, 12) + term.center(color_tokenize_words("loading")))
 print(term.move_xy(0, 13) + term.center(color_tokenize_txt(f"[{'.' * num_dots}") + f"{' ' * num_spaces}]"))
 time.sleep(0.3)
 
 log(f"loading took {time.time() - load_state.t0} seconds")
 return
 
# the program
term = Terminal()

# draw the intro screen on a background thread
Thread(target=draw_intro_screen).start()

# disable garbage collection
from audiotools import AudioSignal
from vamp_helper import load_interface, ez_variation

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 
# ~~~~~~ configs! ~~~~~~~~
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

MAX_LOUDNESS = -20
MIN_LOUDNESS = -40
COLS = 40
ROWS = 13

device = 'Scarlett 4i4 4th Gen'
sample_rate = 48000
num_channels = 4
blocksize = 16384


# TODO: 
# still some quirks to work around recording time: 
# do we wanna stop recording and wait a full cycle before letting people record again? 
# how do we wanna balance the volume of a new input vs what's currently gonig on?
# should people have to take turns in between new loops? 
# otherwise, we're doing great i think
# we also need to add a crossfade. This means maybe cutting off the last 0.1 seconds of the loop, and the beginning 0.1
# and use that to crossfade.

# TODO: do I wanna train a diff model to swap every 2hrs or something? 
# how lond does model swapping take? how can I make it faster?

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# ~~~~~~ looper ~~~~~~~~
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@dataclass 
class State:
 # looper state
 feedback: float = 0.25
 duration: float = 5.0
 record_channel: int = 0

 loopbuf: np.ndarray = None # the main loop buffer. the token telephone audio is here
 looper_in: np.ndarray = None # a buffer that stores the audio that's being recorded

 buf_in: np.ndarray = None # the input block with audio samples in the audio callbac
 lookback_buf: np.ndarray = None # stores some lookback audio for when the threshold is passed, to propery capture transients

 recording: bool = False
 playing: bool = False

 # ramps
 record_ramp_in: bool = False
 record_ramp_out: bool = False

 # n_record_layers: int = 2 # number of times we'll record over before clearing
 # cur_rec_layer: int = 0
 recording_locked: bool = False

 rec_time: float = 0
 cur_hold_time: float = None
 pos: int = 0
 rms_db: float = float("-inf")

 trig_threshold_db = -25 # a more sane default is -20
 hold_seconds = 1.0
 rel_threshold_db = -40 # a more sane default is -30

 status: str = field(default=None)

 # token telephone configs
 z_buf: torch.Tensor = None
 input_ready = False
 input_channel = 0
 token_telephone_processing: bool = False
 num_telephone_chans = 4
 tt_cur_ch = 0

 def __post_init__(self):
 self.loopbuf = np.zeros((num_channels, int(self.duration * sample_rate)))
 self.looper_in = np.zeros((1, int(self.duration * sample_rate)))

 # hold 200ms of lookback to account for rising attacks. 
 num_lookback_samples = max(int(sample_rate * 0.2), int(blocksize))
 log(f"num_lookback_samples {num_lookback_samples} ({num_lookback_samples / sample_rate} seconds)")
 self.lookback_buf = np.zeros((1, num_lookback_samples)) 
 
 self.buf_in = np.zeros((num_channels, blocksize))



def check_if_record(st: State, ain: np.ndarray, on_release_callback=None):
 # get our rms value
 rms = pow2db(np.sqrt(np.mean(ain**2)))
 st.rms_db = rms

 # determine if we should ater the looper state
 # if we werent recording and we cross the trigger threshold
 # start recording
 # if not st.recording and rms > st.trig_threshold_db and not st.recording_locked:
 if not st.recording and rms > st.trig_threshold_db and not st.recording_locked:
 st.recording = True
 st.record_ramp_in = True

 # if we were recording and we cross the release threshold
 # begin the hold period
 if (st.recording and rms < st.rel_threshold_db) or st.rec_time > (st.duration-st.hold_seconds):
 # if we dont have a hold time, set it
 if st.cur_hold_time is None:
 st.cur_hold_time = time.time()

 # release if we have a hold time and we've held for the required time, 
 if (time.time() - st.cur_hold_time) > st.hold_seconds:
 st.record_ramp_out = True
 st.rec_time = 0
 if on_release_callback is not None:
 st.input_ready = True
 on_release_callback(st)
 st.cur_hold_time = None
 else:
 pass
 else:
 st.cur_hold_time = None


def launch_token_telephone(st: State):
 if interface is None:
 log("no interface loaded, can't do token telephone!")
 time.sleep(10)
 return

 # if we're already processing, do nothing
 if st.token_telephone_processing:
 return 
 else:
 log("starting token telephone!")
 Thread(target=do_token_telephone, args=(st,)).start()


def do_token_telephone(st: State,):
 st.token_telephone_processing = True
 while True:
 lrc = st.record_channel
 t0 = time.time()
 cur_ch = st.tt_cur_ch

 # if there was input ready, start back from the top. 
 if st.input_ready:
 log(f"there was input ready, processing!")
 # NOTE: hugo, trying something new here. what happens if
 # we don't reset the channel when input is ready, 
 # and instead let it come in anywhere in the cycle?
 # st.tt_cur_ch = 0 # uncomment to go back to reality

 # clear the lrc, reset for next record. 
 st.input_ready = False

 # reocrd the channel that we'll be processing in and lock recording
 st.input_channel = cur_ch
 st.recording_locked = True

 # first, let's preprocess looper in
 sig_looper_in = AudioSignal(
 torch.from_numpy(st.looper_in).unsqueeze(0),
 sample_rate=sample_rate
 )
 sig_loopbuf_curch = AudioSignal(
 torch.from_numpy(st.loopbuf[cur_ch:cur_ch+1]).unsqueeze(0),
 sample_rate=sample_rate
 )
 # make sure looperin matches the midpoint in loudness
 ldns_mid = max(sig_loopbuf_curch.loudness(), sig_looper_in.loudness())
 sig_looper_in = sig_looper_in.normalize(ldns_mid)
 st.looper_in = sig_looper_in.samples.cpu().numpy().squeeze(0)

 st.loopbuf[cur_ch:cur_ch + 1] = (
 st.looper_in + st.loopbuf[cur_ch:cur_ch+1] * st.feedback
 )
 # also lower the volumes of the other channels
 for i in range(4):
 if i != cur_ch:
 st.loopbuf[i:i+1] = st.loopbuf[i:i+1] * 0.5 # -3dB

 st.looper_in = np.zeros_like(st.looper_in)

 loop_input = st.loopbuf[cur_ch:cur_ch+1]

 # ~~~ VAMPNET STUFF ~~~~
 sig = AudioSignal(
 torch.from_numpy(loop_input).unsqueeze(0),
 sample_rate=sample_rate
 )
 input_loudness = sig.loudness()
 log(f"INPUT loudness {input_loudness}")
 if input_loudness > MAX_LOUDNESS:
 log(f"input loudness {input_loudness} is over {MAX_LOUDNESS}!")
 sig = sig.normalize(MAX_LOUDNESS)
 elif input_loudness < MIN_LOUDNESS:
 log(f"input loudness {input_loudness} is under {MIN_LOUDNESS}!")
 sig = sig.normalize(MIN_LOUDNESS)

 sig = ez_variation(interface, sig)
 sig = sig.resample(sample_rate)

 # notify if we've gone over the loudness
 sig = sig.normalize(input_loudness)
 outloudness = sig.loudness()
 if outloudness > MAX_LOUDNESS:
 log(f"out loudness {sig.loudness()} is over {MAX_LOUDNESS}!")
 sig = sig.normalize(MAX_LOUDNESS)
 elif outloudness < MIN_LOUDNESS:
 log(f"out loudness {sig.loudness()} is under {MIN_LOUDNESS}!")
 sig = sig.normalize(MIN_LOUDNESS)

 # put it back in the loopbuf
 # write to the next channel
 # (TODO: instead of trimming to loopbuf.shape[1], maybe we can just have the loopbuf be the right size from init time.)
 cur_ch = (cur_ch + 1) % st.num_telephone_chans
 st.tt_cur_ch = cur_ch
 if False: # HUGO: is there a time where we want feedback? 
 st.loopbuf[cur_ch:cur_ch+1] = (
 sig.samples.cpu().numpy().squeeze(0)[:, :st.loopbuf.shape[1]]
 + st.feedback * st.loopbuf[cur_ch:cur_ch+1]
 )
 else:
 st.loopbuf[cur_ch:cur_ch+1] = (
 sig.samples.cpu().numpy().squeeze(0)[:, :st.loopbuf.shape[1]]
 )

 log(f"output loudness {sig.loudness()}")
 log(f"telephone loop took {time.time() - t0} seconds... next channel {cur_ch}\n\n")

 # if we've made it back to the input channel, we can unlock the recording
 log(f"cur_ch {cur_ch} input_channel {st.input_channel}")
 if cur_ch == st.input_channel:
 st.recording_locked = False
 log(f"recording unlocked!")


 # unlock the recording if we've successfully written to all channels
 # if st.recording_locked and cur_ch == 0:
 # st.recording_locked = False
 # log(f"recording locked {st.recording_locked}")

 st.token_telephone_processing = False
 return

# TODO: since we're using this really high threshold
# we always need to record about 100ms in advance, to catch the beginning of the attacks. 

def looper_process_block(st, block: np.ndarray):
 lrc = st.record_channel

 # treat the lookback buffer as a circular buffer
 st.lookback_buf = np.roll(st.lookback_buf, block.shape[1], axis=1)
 st.lookback_buf[:, -block.shape[1]:] = block[lrc:lrc+1, :]


 # check if we need to record.
 if st.recording:
 start_i = (st.pos + block.shape[1]) - st.lookback_buf.shape[1]
 end_i = st.pos + st.lookback_buf.shape[1]

 indices = np.take(
 np.arange(st.loopbuf.shape[1]),
 np.arange(start_i, end_i),
 mode="wrap"
 )
 _audio_in = st.lookback_buf[:, :]
 # ramp in if we need to
 if st.record_ramp_in:
 _audio_in = _audio_in * np.linspace(0, 1, _audio_in.shape[1])
 st.record_ramp_in=False
 
 if st.record_ramp_out:
 _audio_in = _audio_in * np.linspace(1, 0, _audio_in.shape[1])
 st.record_ramp_out=False
 st.recording = False

 st.looper_in[:, indices] = (
 0.9 * st.looper_in[:, indices] + _audio_in
 )

 # incremement the recording time
 st.rec_time += st.lookback_buf.shape[1] / sample_rate
 
 # check if we need to play
 crossfade_samples = int(0.1 * sample_rate)
 if st.playing:
 play_pos = (st.pos + block.shape[1]) % st.loopbuf.shape[1] # read one buffer ahead
 indices = np.arange(play_pos, play_pos + block.shape[1]) 
 block = st.loopbuf.take(indices, axis=1, mode="wrap")[:, :] # this doesn't have any crossfading. # TODO: this is still not working! 

 # if we've recorded more than the loop size
 if st.rec_time > st.duration and st.recording:
 # play the loop
 play_pos = st.pos + block.shape[1] # read one buffer ahead
 indices = np.arange(play_pos, play_pos + block.shape[1])

 block[lrc:lrc] = st.looper_in.take(indices, axis=1, mode="wrap")[:, :]
 
 # advance looper state
 st.pos = (st.pos + block.shape[1]) % st.loopbuf.shape[1] 
 
 return block

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# ~~~~~~ drawing ~~~~~~~~
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def draw_rms_bar(st, x, y, width, height):
 rms_min = -50
 rms_max = -10
 rms = st.rms_db
 rms = max(rms, rms_min)
 threshold = st.trig_threshold_db
 rel_threshold = st.rel_threshold_db

 rms_block = int((rms - rms_min) / (rms_max - rms_min) * height)
 threshold_block = (threshold - rms_min) / (rms_max - rms_min) * height
 rel_threshold_block = (rel_threshold - rms_min) / (rms_max - rms_min) * height

 # draw the rms curve
 for i in range(rms_block, height+4):
 with term.location(x+4, y+height-i):
 print(term.clear_bol)
 for i in range(rms_block):
 rms_val = i * (rms_max - rms_min) / height + rms_min
 with term.location(x, y+height-2-i):
 if i < threshold_block:
 print(" " + term.on_green(f"*"))
 else:
 print(" " + term.on_red(f"*"))

 # at the very bottom of the bar, draw the rms value
 with term.location(x, y+height-1):
 print(f"{rms:.1f}dB")
 # print(f" rms")


def draw_looper(st):
 x = 0
 y = 0
 width = COLS
 height = ROWS

 tt_refresh_every = 0.3
 if not hasattr(draw_looper, "last_draw"):
 draw_looper.last_draw = 0
 should_draw = True
 else:
 should_draw = (time.time() - draw_looper.last_draw) > tt_refresh_every
 if should_draw:
 draw_looper.last_draw = time.time()


 draw_rms_bar(st, x, y, width - 10, height)
 
 if should_draw:
 with term.location(width // 2-4, 1):
 for i, letter in enumerate("token telephone"):
 print(on_random_color()(letter), end="")

 # with term.location(ROWS-2, COLS // 2):
 # print(f"status {st.status}!!!")


 # if we're recording, draw a red unlderlined "rec" sign on the bottom right
 # with term.location(width-8, height-1):
 # if st.recording:
 # print(term.on_red("rec"))
 # else:
 # print(term.on_gray50("rec"))
 
 # # if we're playing draw a green underline "play" sign on the bottom right
 # with term.location(width-4, height-1):
 # if st.playing:
 # print(term.on_green("play"))
 # else:
 # print(term.on_gray50("play"))


 # draw the timeline at the bottom using ---
 with term.location(6, height):
 timeline = ["-"] * (width - 12)
 playhead = int((st.pos / st.loopbuf.shape[1]) * (width - 12))
 timeline[playhead] = "v"
 print("|"+"".join(timeline) + "|")

 
 # draw the main message at the very center:
 msg_loc = (width // 2, height // 2+1)
 _x, _y = msg_loc
 if not st.recording:
 if not st.recording_locked:
 print(term.move_xy(0, _y-1) + term.center("make a sound", width=width+5))
 print(term.move_xy(0, _y+0) + term.center("to", width=width+5))
 print(term.move_xy(0, _y+1) + term.center("record", width=width+5))
 else:
 # how many seconds left until we can record again?
 # how many more chs do we need to go through before we can record again?
 if st.tt_cur_ch < st.input_channel:
 chs_remaining = st.input_channel - st.tt_cur_ch
 else:
 chs_remaining = 4-st.tt_cur_ch + st.input_channel
 locked_time_remaining = chs_remaining * st.duration + st.duration - (st.pos / sample_rate)
 print(term.move_xy(0, _y-1) + term.center("please wait", width=width+5))
 print(term.move_xy(0, _y+0) + term.center(term.on_green(f"{locked_time_remaining:.1f}s"), width=width+5))
 print(term.move_xy(0, _y+1) + term.center("for your turn :)", width=width+5))
 else: 
 print(term.move_xy(0, _y-1) + term.center(term.on_red("recording"), width=width+5))
 print(term.move_xy(0, _y+0) + term.center(f"{(st.duration) - st.rec_time:.1f}s left", width=width+5))
 print(term.move_xy(0, _y+1) + term.center("", width=width+5))


 # we'll draw channel 0 (1) on the bottom right corner
 # channel 1 (2) on the top right corner
 # channel 2 (3) on the top left corner
 # channel 3 (4) on the bottom left corner
 my = 3 # margin
 mx = 10
 locations = {
 1: (width - mx, height - my),
 2: (width - mx, 1+my),
 3: (mx, 1+my),
 4: (mx, height - my),
 }
 for i in range(1, 5):
 if should_draw:
 if st.tt_cur_ch == i - 1 and st.token_telephone_processing:
 x, y = locations[i]
 on_random_colors = lambda n: "".join(on_random_color()(" ") for _ in range(n))
 print(term.move_xy(x, y-1) + on_random_colors(5))
 print(term.move_xy(x, y) + on_random_color()(" ") + f" {i} " + on_random_color()(" "))
 print(term.move_xy(x, y+1) + on_random_colors(5))
 else:
 # same thing, but a gray instead of random colors
 x, y = locations[i]
 on_gray_colors = lambda n: "".join(term.on_gray50(" ") for _ in range(n))
 print(term.move_xy(x, y-1) + on_gray_colors(5))
 print(term.move_xy(x, y) + term.on_gray50(" ") + f" {i} " + term.on_gray50(" "))
 print(term.move_xy(x, y+1) + on_gray_colors(5))


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# ~~~~~~ live audio ~~~~~~~~
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def audio_init():
 sd.default.samplerate = sample_rate
 sd.default.device = device
 
# ~~~~~~ the main audio callback ~~~~~~~~~
def callback(st, indata, outdata, frames, _time, status):
 t0 = time.time()
 lrc = st.record_channel

 if status:
 log(f"status is {status}")
 st.status = status
 
 # log dtype, status, frames, time, max min
 # log(f"indata {indata.dtype} max {indata.max()} min {indata.min()} {status} {frames} {_time}")


 ain = rearrange(indata, 't n -> n t', n=num_channels)
 
 # convert audio to from int32 to float32
 ain = ain.astype(np.float32) / np.iinfo(np.int16).max
 buf_in = ain

 # if it's all zeros, we're not recording
 # so we can just pass it through
 if np.all(buf_in == 0):
 st.status = st.status + "no input"
 return 

 st.buf_in = buf_in
 check_if_record(
 st, buf_in, 
 on_release_callback=launch_token_telephone
 )
 buf_in = looper_process_block(st, buf_in)

 # pass our st.loopbuf to the output
 ain = buf_in

 # convert back to int32
 ain = (ain * np.iinfo(np.int16).max).astype(np.int16)

 outdata[:] = rearrange(ain, 'n t -> t n')

 # log(f"outdata {outdata.dtype} max {outdata.max()} min {outdata.min()} --- took {time.time() - t0} seconds")



if DEBUG_NO_VAMPNET:
 interface=None
else:
 interface = load_interface(model_choice="opera")

load_state.loaded = True

def main():
 if PROFILE:
 import yappi
 yappi.start()

 try:
 audio_init()
 st = State()
 st.playing = True

 from functools import partial
 cb = partial(callback, st)

 with term.fullscreen(), term.cbreak():
 with sd.Stream(channels=num_channels, callback=cb, blocksize=blocksize, prime_output_buffers_using_stream_callback=True, dtype=np.int16):
 while True:
 with term.hidden_cursor():
 if DEBUG: 
 time.sleep(100)
 else:
 draw_looper(st)

 except KeyboardInterrupt:
 print(term.clear)
 if PROFILE: 
 yappi.stop()

 # retrieve thread stats by their thread id (given by yappi)
 threads = yappi.get_thread_stats()
 for thread in threads:
 print(
 "Function stats for (%s) (%d)" % (thread.name, thread.id)
 ) # it is the Thread.__class__.__name__
 yappi.get_func_stats(ctx_id=thread.id).print_all()

main()