from typing import List, Tuple import gradio as gr from transformers import AutoTokenizer, AutoModelForCausalLM import note_seq from matplotlib.figure import Figure from numpy import ndarray import torch from constants import GM_INSTRUMENTS, SAMPLE_RATE from string_to_notes import token_sequence_to_note_sequence from model import get_model_and_tokenizer device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Load the tokenizer and the model model, tokenizer = get_model_and_tokenizer() def create_seed_string(genre: str = "OTHER") -> str: """ Creates a seed string for generating a new piece. Args: genre (str, optional): The genre of the piece. Defaults to "OTHER". Returns: str: The seed string. """ if genre == "RANDOM": seed_string = "PIECE_START" else: seed_string = f"PIECE_START GENRE={genre} TRACK_START" return seed_string def get_instruments(text_sequence: str) -> List[str]: """ Extracts the list of instruments from a text sequence. Args: text_sequence (str): The text sequence. Returns: List[str]: The list of instruments. """ instruments = [] parts = text_sequence.split() for part in parts: if part.startswith("INST="): if part[5:] == "DRUMS": instruments.append("Drums") else: index = int(part[5:]) instruments.append(GM_INSTRUMENTS[index]) return instruments def generate_new_instrument(seed: str, temp: float = 0.75) -> str: """ Generates a new instrument sequence from a given seed and temperature. Args: seed (str): The seed string for the generation. temp (float, optional): The temperature for the generation, which controls the randomness. Defaults to 0.75. Returns: str: The generated instrument sequence. """ seed_length = len(tokenizer.encode(seed)) while True: # Encode the conditioning tokens. input_ids = tokenizer.encode(seed, return_tensors="pt") # Move the input_ids tensor to the same device as the model input_ids = input_ids.to(model.device) # Generate more tokens. eos_token_id = tokenizer.encode("TRACK_END")[0] generated_ids = model.generate( input_ids, max_new_tokens=2048, do_sample=True, temperature=temp, eos_token_id=eos_token_id, ) generated_sequence = tokenizer.decode(generated_ids[0]) # Check if the generated sequence contains "NOTE_ON" beyond the seed new_generated_sequence = tokenizer.decode(generated_ids[0][seed_length:]) if "NOTE_ON" in new_generated_sequence: return generated_sequence def get_outputs_from_string( generated_sequence: str, qpm: int = 120 ) -> Tuple[ndarray, str, Figure, str, str]: """ Converts a generated sequence into various output formats including audio, MIDI, plot, etc. Args: generated_sequence (str): The generated sequence of tokens. qpm (int, optional): The quarter notes per minute. Defaults to 120. Returns: Tuple[ndarray, str, Figure, str, str]: The audio waveform, MIDI file name, plot figure, instruments string, and number of tokens string. """ instruments = get_instruments(generated_sequence) instruments_str = "\n".join(f"- {instrument}" for instrument in instruments) note_sequence = token_sequence_to_note_sequence(generated_sequence, qpm=qpm) synth = note_seq.fluidsynth array_of_floats = synth(note_sequence, sample_rate=SAMPLE_RATE) int16_data = note_seq.audio_io.float_samples_to_int16(array_of_floats) fig = note_seq.plot_sequence(note_sequence, show_figure=False) num_tokens = str(len(generated_sequence.split())) audio = gr.make_waveform((SAMPLE_RATE, int16_data)) note_seq.note_sequence_to_midi_file(note_sequence, "midi_ouput.mid") return audio, "midi_ouput.mid", fig, instruments_str, num_tokens def remove_last_instrument( text_sequence: str, qpm: int = 120 ) -> Tuple[ndarray, str, Figure, str, str, str]: """ Removes the last instrument from a song string and returns the various output formats. Args: text_sequence (str): The song string. qpm (int, optional): The quarter notes per minute. Defaults to 120. Returns: Tuple[ndarray, str, Figure, str, str, str]: The audio waveform, MIDI file name, plot figure, instruments string, new song string, and number of tokens string. """ # We split the song into tracks by splitting on 'TRACK_START' tracks = text_sequence.split("TRACK_START") # We keep all tracks except the last one modified_tracks = tracks[:-1] # We join the tracks back together, adding back the 'TRACK_START' that was removed by split new_song = "TRACK_START".join(modified_tracks) if len(tracks) == 2: # There is only one instrument, so start from scratch audio, midi_file, fig, instruments_str, new_song, num_tokens = generate_song( text_sequence=new_song ) elif len(tracks) == 1: # No instrument so start from empty sequence audio, midi_file, fig, instruments_str, new_song, num_tokens = generate_song( text_sequence="" ) else: audio, midi_file, fig, instruments_str, num_tokens = get_outputs_from_string( new_song, qpm ) return audio, midi_file, fig, instruments_str, new_song, num_tokens def regenerate_last_instrument( text_sequence: str, qpm: int = 120 ) -> Tuple[ndarray, str, Figure, str, str, str]: """ Regenerates the last instrument in a song string and returns the various output formats. Args: text_sequence (str): The song string. qpm (int, optional): The quarter notes per minute. Defaults to 120. Returns: Tuple[ndarray, str, Figure, str, str, str]: The audio waveform, MIDI file name, plot figure, instruments string, new song string, and number of tokens string. """ last_inst_index = text_sequence.rfind("INST=") if last_inst_index == -1: # No instrument so start from empty sequence audio, midi_file, fig, instruments_str, new_song, num_tokens = generate_song( text_sequence="", qpm=qpm ) else: # Take it from the last instrument and continue generation next_space_index = text_sequence.find(" ", last_inst_index) new_seed = text_sequence[:next_space_index] audio, midi_file, fig, instruments_str, new_song, num_tokens = generate_song( text_sequence=new_seed, qpm=qpm ) return audio, midi_file, fig, instruments_str, new_song, num_tokens def change_tempo( text_sequence: str, qpm: int ) -> Tuple[ndarray, str, Figure, str, str, str]: """ Changes the tempo of a song string and returns the various output formats. Args: text_sequence (str): The song string. qpm (int): The new quarter notes per minute. Returns: Tuple[ndarray, str, Figure, str, str, str]: The audio waveform, MIDI file name, plot figure, instruments string, text sequence, and number of tokens string. """ audio, midi_file, fig, instruments_str, num_tokens = get_outputs_from_string( text_sequence, qpm=qpm ) return audio, midi_file, fig, instruments_str, text_sequence, num_tokens def generate_song( genre: str = "OTHER", temp: float = 0.75, text_sequence: str = "", qpm: int = 120, ) -> Tuple[ndarray, str, Figure, str, str, str]: """ Generates a song given a genre, temperature, initial text sequence, and tempo. Args: model (AutoModelForCausalLM): The pretrained model used for generating the sequences. tokenizer (AutoTokenizer): The tokenizer used to encode and decode the sequences. genre (str, optional): The genre of the song. Defaults to "OTHER". temp (float, optional): The temperature for the generation, which controls the randomness. Defaults to 0.75. text_sequence (str, optional): The initial text sequence for the song. Defaults to "". qpm (int, optional): The quarter notes per minute. Defaults to 120. Returns: Tuple[ndarray, str, Figure, str, str, str]: The audio waveform, MIDI file name, plot figure, instruments string, generated song string, and number of tokens string. """ if text_sequence == "": seed_string = create_seed_string(genre) else: seed_string = text_sequence generated_sequence = generate_new_instrument(seed=seed_string, temp=temp) audio, midi_file, fig, instruments_str, num_tokens = get_outputs_from_string( generated_sequence, qpm ) return audio, midi_file, fig, instruments_str, generated_sequence, num_tokens