import asyncio import websockets import json import torch import numpy as np import base64 import soundfile as sf import os from audiocraft.models import MusicGen import time from BeatNet.BeatNet import BeatNet import librosa import pyrubberband as pyrb import faulthandler import random faulthandler.enable() beatnet = BeatNet(1, mode="offline", inference_model="DBN", plot=[], thread=False, device="cpu") next_job_id = 1 sample_rate = 0 class Messenger: def __init__(self, websocket): self.websocket = websocket async def progress(self, progress): payload = { "progress": progress } payload_serialized = json.dumps(payload) await self.websocket.send(payload_serialized) async def chunk(self, chunk): payload = { "chunk": chunk } payload_serialized = json.dumps(payload) await self.websocket.send(payload_serialized) async def success(self, result): payload = { "type": "audio", "success": result } payload_serialized = json.dumps(payload) await self.websocket.send(payload_serialized) await self.websocket.close() async def error(self, error): payload = { "error": error } payload_serialized = json.dumps(payload) await self.websocket.send(payload_serialized) await self.websocket.close() def estimate_beats(wav): # resample to BeatNet's sample rate beatnet_input = librosa.resample( wav, orig_sr=sample_rate, target_sr=beatnet.sample_rate, ) return beatnet.process(beatnet_input) # Gets audio prompt from given loop def get_audio_prompt(loop, beats_per_prompt): beats = estimate_beats(loop) bar_times = beats[:, 0][beats[:, 1] == 1] num_bars = len(bar_times) start_loop_time = beats[0][0] end_loop_time = beats[len(beats) - 1][0] num_beats = len(beats[(beats[:, 0] >= start_loop_time) & (beats[:, 0] < end_loop_time)]) print("Prev audio sample contains "+str(num_bars)+" bars") if num_beats < beats_per_prompt: raise ValueError("Not enough beats for audio prompt") end_time = bar_times[len(bar_times)-1] last_nbeats = beats[beats[:, 0] <= end_time][-beats_per_prompt:] audio_prompt_start_time = last_nbeats[0][0] audio_prompt_end_time = last_nbeats[-1][0] audio_prompt_start_sample = int(audio_prompt_start_time * sample_rate) audio_prompt_end_sample = int(audio_prompt_end_time * sample_rate) audio_prompt = torch.tensor(loop[audio_prompt_start_sample:audio_prompt_end_sample])[None] audio_prompt_duration = audio_prompt_end_sample - audio_prompt_start_sample audio_prompt_seconds = audio_prompt_end_time - audio_prompt_start_time return audio_prompt, audio_prompt_duration, audio_prompt_seconds def get_first_loop(prompt, model, max_duration, top_k, top_p, temperature, classifier_free_guidance): model.set_generation_params( duration=max_duration, top_k=top_k, top_p=top_p, temperature=temperature, cfg_coef=classifier_free_guidance, ) return model.generate([prompt], progress=True).cpu().numpy()[0, 0] def get_loop(audio_prompt, prompt, model, max_duration, audio_prompt_duration, audio_prompt_seconds, top_k, top_p, temperature, classifier_free_guidance): model.set_generation_params( duration=max_duration+audio_prompt_seconds, top_k=top_k, top_p=top_p, temperature=temperature, cfg_coef=classifier_free_guidance, ) return model.generate_continuation( prompt=audio_prompt, descriptions=[prompt], prompt_sample_rate=sample_rate, progress=True, ).cpu().numpy()[0, 0, audio_prompt_duration:] def arr_to_wav(numpy_array) -> str: numpy_array = numpy_array.astype('float32') tmp_filename = '/tmp/'+str(int(time.time()))+'.wav' sf.write(tmp_filename, numpy_array.T, sample_rate) with open(tmp_filename, 'rb') as wav_file: base64_string = base64.b64encode(wav_file.read()).decode() base64_string = "data:audio/wav;base64," + base64_string os.remove(tmp_filename) return base64_string # Websocket handler function, this is executed everytime new user connects async def handler(websocket): # Wait for the client to send a message message = await websocket.recv() # Parse the message as JSON message_data = json.loads(message) # Create messenger instance messenger = Messenger(websocket) # Check if the message contains a request to start the worker if message_data.get("started"): global next_job_id, beatnet, sample_rate # Assign & increment the current job ID current_job_id: int = next_job_id next_job_id += 1 print(f"Started worker with job id: {current_job_id}") # Get the prompt, and model from the message settings = message_data.get("settings") prompt = message_data.get("prompt") bpm = float(settings['bpm']) top_k = int(settings['top_k']) top_p = float(settings['top_p']) temperature = float(settings['temperature']) cfg_coef = float(settings['cfg_coef']) device = settings['device'] accuracy = settings['accuracy'] beats_per_chunk = int(settings['beats_per_chunk']) beats_per_prompt = int(settings['beats_per_prompt']) buffer = int(settings['buffer']) unguided = bool(settings['unguided']) if 'seed' in settings: seed = int(settings['seed']) random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) torch.backends.cudnn.deterministic = True max_duration = int(settings['max_duration']) if 'max_duration' in settings else np.floor((60 / bpm) * beats_per_chunk * 1.5) print("Target duration: "+str(max_duration)+"s") model = MusicGen.get_pretrained(accuracy, device=device) sample_rate = model.sample_rate print(message_data) loops = [] x = 0 prev_prompt = f"{prompt}, {bpm} bpm" if not unguided else prompt has_rewinded = False while True: is_prompt = x % buffer == 0 is_start = x == 0 if has_rewinded: has_rewinded = False elif is_prompt and not is_start: new_message = await websocket.recv() new_message_data = json.loads(new_message) print(new_message_data) if new_message_data.get('finished'): break prev_prompt = f"{new_message_data.get('prompt')}, {bpm} bpm" if not unguided else new_message_data.get('prompt') print(f"Generating chunk - {x}") print(f"Requesting new prompt - "+str(is_prompt and not is_start)) if not unguided: try: if is_start: new_loop = get_first_loop( prev_prompt, model, max_duration, top_k, top_p, temperature, cfg_coef ) else: new_audio_prompt, new_audio_prompt_duration, new_audio_prompt_seconds = get_audio_prompt(loops[x-1], beats_per_prompt) new_loop = get_loop( new_audio_prompt, prev_prompt, model, max_duration, new_audio_prompt_duration, new_audio_prompt_seconds, top_k, top_p, temperature, cfg_coef ) print("Processing loop") # Normalize new loops new_loop_normalized = new_loop / np.abs(new_loop).max() # Find beats and downbeats beats = estimate_beats(new_loop_normalized) downbeat_times = beats[:, 0][beats[:, 1] == 1] # Find best cutting points to beats i.e cut from downbeat to downbeat min_loop_time = downbeat_times[0] max_loop_time = downbeat_times[len(downbeat_times)-1] # Create new beats array always ends on downbeat-1 beat - 1, 2, 3, 4, 1, 2, 3, 4 trim_index_start = np.where(beats[:, 0] == min_loop_time)[0][0] trim_index_end = np.where(beats[:, 0] == max_loop_time)[0][0] processed_beats = beats[trim_index_start:trim_index_end] # Beat Sanity check if len(processed_beats) < beats_per_chunk: raise ValueError("Failed to produce loop with good beat count, retrying...") # Now select commanded number of beats from new array start_loop_time = processed_beats[0][0] end_loop_time = processed_beats[beats_per_chunk-1][0] # Finally truncate sample start_sample = int(start_loop_time * sample_rate) end_sample = int(end_loop_time * sample_rate) new_loop_truncated = new_loop_normalized[start_sample:end_sample] # Find bpm of adjusted beats array, time stretch sample accordingly to fit within given bpm num_beats = len(processed_beats[(processed_beats[:, 0] >= start_loop_time) & (processed_beats[:, 0] < end_loop_time)]) duration = end_loop_time - start_loop_time actual_bpm = num_beats / duration * 60 print("BPM count: " + str(actual_bpm)+" vs "+str(bpm)) if abs(actual_bpm - bpm) > 10: raise ValueError("Failed to produce loop in good bpm range, retrying...") loop_final = pyrb.time_stretch(new_loop_truncated, sample_rate, bpm / actual_bpm) except ValueError as e: print(getattr(e, 'message', repr(e))) has_rewinded = True continue else: if is_start: new_loop = get_first_loop( prev_prompt, model, max_duration, top_k, top_p, temperature, cfg_coef ) else: prev_loop = loops[x - 1] midpoint = len(prev_loop) // 2 new_audio_prompt = prev_loop[:midpoint] new_audio_prompt_seconds = len(new_audio_prompt) / sample_rate new_audio_prompt = torch.tensor(new_audio_prompt)[None] new_loop = get_loop( new_audio_prompt, prev_prompt, model, max_duration, midpoint, new_audio_prompt_seconds, top_k, top_p, temperature, cfg_coef ) loop_final = new_loop_truncated = new_loop / np.abs(new_loop).max() processed_loop_encoded = arr_to_wav(loop_final) loops.append(new_loop_truncated) x += 1 if x % buffer == 0: await messenger.progress(processed_loop_encoded) else: await messenger.chunk(processed_loop_encoded) print(f"Sequence complete for job id: {current_job_id}") # Send the result back to the client await messenger.success([]) if __name__ == "__main__": start_server = websockets.serve(handler, "0.0.0.0", 1337, ping_interval=None, max_size=1_000_000_000) asyncio.get_event_loop().run_until_complete(start_server) print("Server started...") try: asyncio.get_event_loop().run_forever() except KeyboardInterrupt: print("Keyboard Interrupt: Exiting server") except: print("Unknown error") finally: exit(1)