Initial commit

app.py

@@ -0,0 +1,245 @@
+import streamlit as st
+from source.languagemodel import LanguageModel
+from source.utilities import (
+    convert_tokens_to_songdata,
+    convert_songdata_to_notesequence,
+    convert_songdata_to_pianoroll,
+    convert_notesequence_to_wave,
+    convert_notesequence_to_midi
+)
+
+# Define the MIDI instruments.
+midi_instruments = {
+    "Harpsichord": 6,
+    "Church Organ": 19,
+    "Piano": 0,
+}
+
+# Load the model once and cache it.
+@st.cache_resource
+def load_model():
+    model = LanguageModel("TristanBehrens/bach-garland-mambaplus")
+    return model
+model = load_model()
+
+
+# Initialize token_sequence in session state if it doesn't exist
+if "token_sequence" not in st.session_state:
+    st.session_state.token_sequence = "GARLAND_START"
+    st.session_state.song_data = None
+    st.session_state.piano_roll = None
+    st.session_state.wave = None
+    st.session_state.note_sequence = None
+    st.session_state.midi_file_content = None
+    st.session_state.temperature = 0.1
+    st.session_state.bpm = 100
+    st.session_state.instrument = "Piano"
+
+
+# Define the main function.
+def main():
+
+    columns = st.columns([0.7, 0.3])
+
+    # Set up the Streamlit application
+    column = columns.pop(0)
+    with column:
+
+        # Change the colors of the a-tag to (255, 75, 75).
+        st.markdown("<style>a:link { color: #FF4B4B; } a:visited { color: #FF4B4B; }</style>", unsafe_allow_html=True)
+
+        # Add a title.
+        st.title("Garland Composer")
+        linkedin_url = "https://huggingface.co/TristanBehrens/bach-garland-mambaplus/"
+        x_url = "https://huggingface.co/TristanBehrens/bach-garland-mambaplus/"
+        st.write(f"By Dr. Tristan Behrens. Find me on [LinkedIn]({linkedin_url}) and [X]({x_url}).")
+        hf_url = "https://huggingface.co/TristanBehrens/bach-garland-mambaplus/"
+        st.write(f"Model available on [Hugging Face]({hf_url}).")
+
+    # Add a picture.
+    column = columns.pop(0)
+    with column:
+        st.write(" ")
+        st.write(" ")
+        st.write(" ")
+        st.image("garland.jpg", use_column_width=True)
+
+    # Add a horizontal line.
+    st.markdown("---")
+
+    # Create two columns.
+    columns = st.columns(3)
+
+    # Add a slider to control the temperature.
+    state_temperature = st.session_state.temperature
+    with columns.pop(0):
+        temperature = st.slider("Temperature", 0.0, 1.0, state_temperature)
+    st.session_state.temperature = temperature
+
+    # Add a slider to control the bpm.
+    state_bpm = st.session_state.bpm
+    with columns.pop(0):
+        bpm = st.slider("BPM", 80, 120, state_bpm, 5)
+    st.session_state.bpm = bpm
+
+    # Dropdown for the instrument.
+    state_instrument = st.session_state.instrument
+    with columns.pop(0):
+        instrument = st.selectbox("Instrument", list(midi_instruments.keys()), index=list(midi_instruments.keys()).index(state_instrument))
+    st.session_state.instrument = instrument
+    
+    # Get the token sequence from the session state.
+    token_sequence = st.session_state.token_sequence
+
+    # Columns for the buttons.
+    columns = st.columns(5)
+
+    # Add a button to generate the next bar.
+    column = columns.pop(0)
+    with column:
+        if st.button("Add a bar", use_container_width=True):
+            token_sequence = extend_sequence(model, token_sequence, temperature)
+            refresh(token_sequence, bpm, instrument)
+
+    # Add a button to compose long.
+    column = columns.pop(0)
+    with column:
+        if st.button("Auto compose", use_container_width=True):
+            token_sequence = auto_compose(model, token_sequence, temperature)
+            refresh(token_sequence, bpm, instrument)
+
+    # Add a button to remove the last bar.
+    column = columns.pop(0)
+    with column:
+        if st.button("Remove last", use_container_width=True):
+            token_sequence = shortened_sequence(token_sequence)
+            refresh(token_sequence, bpm, instrument)
+
+    # Add a button to reset the sequence.
+    column = columns.pop(0)
+    if token_sequence != "GARLAND_START":
+        with column:
+            if st.button("Reset", use_container_width=True):
+                with columns.pop(0):
+                    token_sequence = "GARLAND_START"
+                    refresh(token_sequence, bpm, instrument)
+
+    # Provide a download button for the MIDI file.
+    column = columns.pop(0)
+    if "midi_file_content" in st.session_state and st.session_state.midi_file_content is not None:
+        with column:
+            midi_file_content = st.session_state.midi_file_content
+            if st.download_button(
+                label="Download MIDI",
+                data=midi_file_content,
+                file_name="music.mid",
+                mime="audio/midi",
+                use_container_width=True
+            ):
+                pass
+
+    # Add a horizontal line.
+    st.markdown("---")
+
+    # Display the piano roll.
+    if "piano_roll" in st.session_state and st.session_state.piano_roll is not None:
+        st.image(st.session_state.piano_roll)
+
+    # Display an audio player.
+    if "wave" in st.session_state and st.session_state.wave is not None:
+        st.audio(st.session_state.wave, format="audio/wav", sample_rate=44100, autoplay=True)
+
+    # Add a horizontal line.
+    st.markdown("---")
+
+    # Set the text color to (255, 31, 75).
+    if token_sequence.endswith("GARLAND_END"):
+        st.write("The AI believes that the music is finished.")
+    else:
+        st.write("The AI believes that the music is not finished.")
+
+
+def auto_compose(model, token_sequence, temperature):
+
+    max_iterations = 100
+    for _ in range(max_iterations):
+        token_sequence = extend_sequence(model, token_sequence, temperature)
+        if token_sequence.endswith("GARLAND_END"):
+            break
+    return token_sequence
+
+
+def extend_sequence(model, token_sequence, temperature):
+
+    # Replace the last GARLAND_END token with NEXT.
+    if token_sequence.endswith("GARLAND_END"):
+        token_sequence = token_sequence.replace("GARLAND_END", "NEXT")
+
+    # The maximum length of the generated music.
+    max_length = 16_384
+
+    # When to stop the generation.
+    end_tokens = ["NEXT", "GARLAND_END"]
+
+    # Compose the music iterativelybar by bar.
+    output_dict = model.generate(
+        prompt=token_sequence,
+        temperature=temperature,
+        max_length=max_length,
+        end_tokens=end_tokens,
+        forbidden_tokens=["[PAD]", "[EOS]"],
+        return_structured_output=True
+    )
+    output = output_dict["output"]
+    return output
+
+
+def shortened_sequence(token_sequence):
+
+    # Find the position of the next to last NEXT token.
+    next_tokens = token_sequence.split()
+    next_positions = [i for i, x in enumerate(next_tokens) if x == "NEXT" or x == "GARLAND_END"]
+    if len(next_positions) <= 1:
+        token_sequence = "GARLAND_START"
+    else:
+        next_position = next_positions[-2]
+        token_sequence = " ".join(next_tokens[:next_position + 1])
+    return token_sequence
+
+
+def refresh(token_sequence="GARLAND_START", bpm=120, instrument="Piano"):
+
+    # Get the token sequence into the session state.
+    st.session_state.token_sequence = token_sequence
+
+    # Convert to song data.
+    song_data = convert_tokens_to_songdata(token_sequence)
+    song_data["bpm"] = bpm
+    st.session_state.song_data = song_data
+
+    # Set the instrument.
+    for track in song_data["tracks"]:
+        track["instrument"] = midi_instruments[instrument]
+
+    # Convert to piano roll.
+    piano_roll = convert_songdata_to_pianoroll(song_data)
+    st.session_state.piano_roll = piano_roll
+
+    # Convert to note sequence.
+    note_sequence = convert_songdata_to_notesequence(song_data)
+    st.session_state.note_sequence = note_sequence
+
+    # Play the note sequence.
+    wave = convert_notesequence_to_wave(note_sequence)
+    st.session_state.wave = wave
+
+    # Get the MIDI file content.
+    midi_file_content = convert_notesequence_to_midi(note_sequence)
+    st.session_state.midi_file_content = midi_file_content
+
+    # Rerun the app.
+    st.rerun()
+
+
+if __name__ == "__main__":
+    main()

assets/asciilogo.txt

@@ -0,0 +1,11 @@
+   ▄█    █▄       ▄████████  ▄█        ▄█  ▀█████████▄     ▄████████ ███    █▄  ███▄▄▄▄   ███▄▄▄▄      ▄████████ 
+  ███    ███     ███    ███ ███       ███    ███    ███   ███    ███ ███    ███ ███▀▀▀██▄ ███▀▀▀██▄   ███    ███ 
+  ███    ███     ███    █▀  ███       ███▌   ███    ███   ███    ███ ███    ███ ███   ███ ███   ███   ███    ███ 
+ ▄███▄▄▄▄███▄▄  ▄███▄▄▄     ███       ███▌  ▄███▄▄▄██▀   ▄███▄▄▄▄██▀ ███    ███ ███   ███ ███   ███   ███    ███ 
+▀▀███▀▀▀▀███▀  ▀▀███▀▀▀     ███       ███▌ ▀▀███▀▀▀██▄  ▀▀███▀▀▀▀▀   ███    ███ ███   ███ ███   ███ ▀███████████ 
+  ███    ███     ███    █▄  ███       ███    ███    ██▄ ▀███████████ ███    ███ ███   ███ ███   ███   ███    ███ 
+  ███    ███     ███    ███ ███▌    ▄ ███    ███    ███   ███    ███ ███    ███ ███   ███ ███   ███   ███    ███ 
+  ███    █▀      ██████████ █████▄▄██ █▀   ▄█████████▀    ███    ███ ████████▀   ▀█   █▀   ▀█   █▀    ███    █▀  
+                            ▀                             ███    ███                                             
+
+By Dr. Tristan Behrens

requirements.txt

@@ -0,0 +1,11 @@
+dacite==1.8.1
+colorama==0.4.6
+omegaconf==2.3.0
+streamlit==1.38.0
+note_seq==0.0.5
+pyfluidsynth==1.3.2
+torch==2.2.0
+transformers==4.44.0
+mamba-ssm==2.2.2
+einops==0.8.0
+mambapy==1.2.0

source/languagemodel.py

@@ -0,0 +1,288 @@
+# Helibrunna - A HuggingFace compatible xLSTM trainer.
+# Copyright (c) 2024 Dr. Tristan Behrens
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+import os
+import glob
+from omegaconf import OmegaConf
+from transformers import PreTrainedTokenizerFast
+import torch
+from safetensors.torch import load_file
+import time
+from .utilities import display_logo, model_from_config
+
+
+class LanguageModel:
+
+    def __init__(self, model_path_or_repo, config_overrides={}, mask_special_tokens=True, device="auto"):
+        """
+        Initializes the LanguageModel object.
+        Args:
+            model_path_or_repo (str): The path to the model or the repository ID.
+        Raises:
+            ValueError: If the model checkpoint, tokenizer, config, or weights are not found.
+            Exception: If failed to download the model.
+        Returns:
+            None
+        """
+
+        # Set the maskt_special_tokens flag.
+        self.mask_special_tokens = mask_special_tokens
+
+        # Set the device. CPU is default.
+        if device != "auto":
+            
+            # Check if CUDA is available.
+            if not torch.cuda.is_available() and device == "cuda":
+                raise ValueError("CUDA is not available on this system.")
+            
+            # Check if MPS is available.
+            if not torch.backends.mps.is_available() and device == "mps":
+                raise ValueError("MPS is not available on this system.")
+
+            # Set the device.
+            self.device = device
+
+        # Set the device to auto.
+        else:     
+
+            # Set the device to CPU if auto is selected.
+            self.device = "cpu" if device == "auto" else device
+
+            # Check if CUDA is available.
+            if torch.cuda.is_available() and device == "auto":
+                self.device = "cuda"
+
+            # See if MPS is available.
+            # Note: This is disabled for now. It's not working as expected. It is very slow.
+            #if torch.backends.mps.is_available():
+            #    self.device = "mps"
+
+        # Display the logo.
+        display_logo()
+
+        # Download the model if it doesn't exist. Or at least try to.
+        if not os.path.exists(model_path_or_repo):
+            from huggingface_hub import snapshot_download
+            try:
+                model_path=snapshot_download(repo_id=model_path_or_repo)
+                tokenizer_path=model_path
+            except Exception as e:
+                raise f"Failed to download the model: {e}"
+        
+        # Use a local model.
+        else:
+            # Set the model path and tokenizer path.
+            model_path = None
+            tokenizer_path = model_path_or_repo
+
+            # Find all the checkpoint folders, folders that start with "checkpoint-". Then find the last one.
+            checkpoint_folders = glob.glob(os.path.join(model_path_or_repo, "checkpoint-*"))
+            for checkpoint_folder in checkpoint_folders:
+                if checkpoint_folder.endswith("-last"):
+                    model_path = checkpoint_folder
+                    break
+            if model_path is None:
+                raise ValueError("No model checkpoint found.")
+
+            # Find the tokenizer folder.
+            if os.path.exists(os.path.join(model_path_or_repo, "tokenizer.json")):
+                tokenizer_path = model_path_or_repo
+            if not os.path.exists(tokenizer_path):
+                raise ValueError("Tokenizer not found.")
+
+        # Load the config.
+        config_path = os.path.join(model_path, "config.yaml")
+        if not os.path.exists(config_path):
+            raise ValueError(f"Config not found at {config_path}")
+        model_config = OmegaConf.load(config_path)
+
+        # Override the config.
+        if config_overrides != {} and config_overrides is not None:
+            model_config = OmegaConf.merge(model_config, config_overrides)
+        import json
+        print(json.dumps(OmegaConf.to_container(model_config), indent=4))
+
+        # Create the model from the config.
+        model = model_from_config(model_config, device=self.device)
+        model.to(self.device)
+        self.config = model_config
+
+        # Load the weights from the checkpoint.
+        weights_path = os.path.join(model_path, "model.safetensors")
+        if not os.path.exists(weights_path):
+            raise ValueError(f"Weights not found at {weights_path}")
+        state_dict = load_file(weights_path)
+
+        # TODO: Permute the last two dimensions of these parameters: xlstm_block_stack.blocks.2.xlstm.slstm_cell._recurrent_kernel_:
+        # Check if we have an xLSTM model and if CUDA is not available.
+        if not torch.cuda.is_available() and model_config.get("type", "xLSTMLMModel") == "xLSTMLMModel":
+            print(state_dict.keys())
+            endings = ["xlstm.slstm_cell._recurrent_kernel_"]
+            for key, values in state_dict.items():
+                for ending in endings:
+                    if key.endswith(ending):
+                        print(key)
+                        print(values.shape)
+                        
+                        # Option: Permute the last two dimensions.
+                        values = values.permute(0, 2, 1)
+                        
+                        # Option: View the tensor.
+                        #new_shape = (values.shape[0], values.shape[2], values.shape[1])
+                        #values = values.view(new_shape)
+                        
+                        print(values.shape)
+                        state_dict[key] = values
+                        break
+
+        # Load the weights into the model.
+        model.load_state_dict(state_dict)
+        self.model = model
+
+        # Load the tokenizer.
+        tokenizer_path = os.path.join(tokenizer_path, "tokenizer.json")
+        if not os.path.exists(tokenizer_path):
+            raise ValueError(f"Tokenizer not found at {tokenizer_path}")
+        tokenizer = PreTrainedTokenizerFast.from_pretrained(tokenizer_path)
+        self.tokenizer = tokenizer
+
+
+    def generate(
+        self,
+        prompt: str,
+        temperature: float = 1.0,
+        max_length: int = 100,
+        end_tokens: list[str] = [],
+        forbidden_tokens: list[str] = [],
+        return_structured_output: bool = False
+    ):
+        """
+        Generates a continuation for a given prompt using the language model.
+        Args:
+            prompt (str): The prompt to generate a continuation for.
+            temperature (float, optional): The temperature value for controlling the randomness of the generated output. 
+                Higher values (e.g., 1.0) make the output more random, while lower values (e.g., 0.5) make it more deterministic. 
+                Defaults to 1.0.
+            max_length (int, optional): The maximum length of the generated output. Defaults to 100.
+            end_tokens (list[str], optional): A list of end tokens that, if encountered, will stop the generation process. 
+                Defaults to an empty list.
+            return_structured_output (bool, optional): If True, returns a dictionary with the generated output, elapsed time, 
+                and tokens per second. If False, returns only the generated output as a string. Defaults to False.
+        Returns:
+            str or dict: The generated output as a string if return_structured_output is False. 
+                A dictionary with the generated output, elapsed time, and tokens per second if return_structured_output is True.
+        """    
+
+        # Tokenize the prompt.
+        inputs = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
+        assert inputs.shape[0] == 1
+
+        # Determine the end tokens ids.
+        end_token_ids = []
+        for end_token in end_tokens:
+            assert end_token in self.tokenizer.vocab
+            end_token_ids.append(self.tokenizer(end_token).input_ids[0])
+
+        # Initialize the ids to mask.
+        ids_to_mask = []
+
+        # Mask the forbidden tokens.
+        for forbidden_token in forbidden_tokens:
+            assert forbidden_token in self.tokenizer.vocab
+            ids_to_mask.extend(self.tokenizer(forbidden_token).input_ids)
+
+        # Generate the continuation.
+        start_time = time.time()
+        tokens_count = 0
+        while inputs.shape[1] < max_length:
+
+            # Stop if the maximum context length is reached.
+            if inputs.shape[1] >= self.config.context_length:
+                print("Warning: The maximum context length has been reached.")
+                break
+
+            # Generate the continuation.
+            outputs = self.model(inputs.to(device=self.device))
+            assert outputs.shape[0] == 1
+
+            # Mask the tokens.
+            outputs[:, :, self.tokenizer.all_special_ids] = float("-inf")
+
+            # Use the temperature to sample from the distribution.
+            outputs = outputs / temperature
+            outputs = torch.nn.functional.softmax(outputs, dim=-1)
+            outputs = torch.multinomial(outputs[0, -1], num_samples=1)
+
+            # Add to the inputs.
+            inputs = torch.cat([inputs, outputs.unsqueeze(0)], dim=1)
+            
+            # Increment the tokens count.
+            tokens_count += 1
+
+            # Check if the end token is reached.
+            if outputs[0] in end_token_ids:
+                break
+
+        # Print the elapsed time and tokens per second.
+        elapsed_time = time.time() - start_time
+        tokens_per_second = tokens_count / elapsed_time
+
+        # Decode the output.
+        output = self.tokenizer.decode(inputs[0].tolist())
+
+        # Return the output.
+        if not return_structured_output:
+            return output
+        
+        # Return the structured output.
+        else:
+            return {
+                "output": output,
+                "elapsed_time": elapsed_time,
+                "tokens_per_second": tokens_per_second
+            }
+        
+    def summary(self):
+        """
+        Prints a summary of the model. Makes the model architecture readable. Includes the number of parameters.
+        """
+
+        # Print the model.
+        print(self.model)
+
+        # Get the number of parameters.
+        number_of_parameters = sum(p.numel() for p in self.model.parameters())
+        print(f"Number of parameters: {number_of_parameters:_}")
+        sizes = ["", "K", "M", "B", "T"]
+        size_index = 0
+        while number_of_parameters > 1000:
+            number_of_parameters /= 1000
+            size_index += 1
+        print(f"Number of parameters: {number_of_parameters:.2f}{sizes[size_index]}")
+        
+        # Size of the model.
+        # Get the total size of all the markdown files. And make it human readable.
+        number_of_parameters = sum(p.numel() for p in self.model.parameters())
+        total_size = number_of_parameters * 4
+        sizes = ["B", "KB", "MB", "GB", "TB"]
+        size_index = 0
+        while total_size > 1024:
+            total_size /= 1024
+            size_index += 1
+        print(f"Total size of the model: {total_size:.2f}{sizes[size_index]} for precision 32-bit floats.")
+
+        # Print on which device the model is running. 
+        print(f"Device: {self.device}")

source/utilities.py

@@ -0,0 +1,331 @@
+import copy
+import note_seq
+from PIL import Image
+import tempfile
+import os
+import colorama
+from omegaconf import DictConfig, OmegaConf
+import torch
+from typing import List, Tuple, Dict
+from dacite import from_dict
+from collections.abc import MutableMapping
+import sys
+
+
+# NOTE: Imported from helibrunna.
+def display_logo():
+    """
+    Display the logo by printing it line by line with a cyberpunk color scheme.
+
+    Raises:
+        FileNotFoundError: If the logo file is missing.
+    """
+
+    # Get the path of this script and use it to find the logo.
+    script_path = os.path.dirname(os.path.realpath(__file__))
+    search_path = os.path.dirname(script_path)
+
+    # Load the logo.
+    logo_path = os.path.join(search_path, "assets", "asciilogo.txt")
+    if not os.path.exists(logo_path):
+        raise FileNotFoundError("The logo file is missing.")
+    with open(logo_path, "r") as f:
+        logo = f.read()
+
+    # Print the logo line by line. Use colorama to colorize the output. Use a cyberpunk color scheme.
+    for line_index, line in enumerate(logo.split("\n")):
+        color = colorama.Fore.GREEN
+        style = colorama.Style.BRIGHT if line_index % 2 == 0 else colorama.Style.NORMAL
+        print(color + style + line)
+    print(colorama.Style.RESET_ALL)
+
+
+# NOTE: Imported from helibrunna.
+def model_from_config(model_config: DictConfig, device:str) -> torch.nn.Module:
+    """
+    Create a model based on the provided model configuration.
+
+    Args:
+        model_config (DictConfig): The configuration for the model.
+
+    Returns:
+        The created model.
+
+    Raises:
+        ValueError: If the model type is unknown.
+    """
+    
+    # Get the model type from the configuration.
+    model_type = model_config.get("type", "xLSTMLMModel")
+    
+    # Create the xLSTMLMModel.
+    if model_type == "xLSTMLMModel":
+        print("Creating xLSTMLMModel...")
+        from xlstm.xlstm_lm_model import xLSTMLMModel, xLSTMLMModelConfig
+        
+        # If there is no GPU, use the vanilla backend.
+        if not torch.cuda.is_available():
+            #model_config.backend = "vanilla"
+            model_config.slstm_block.slstm.backend = "vanilla"
+            model_config.mlstm_block.mlstm.backend = "vanilla"
+        model_config_object = from_dict(xLSTMLMModelConfig, OmegaConf.to_container(model_config))
+        
+        # Create the model.
+        model = xLSTMLMModel(model_config_object)
+        model.reset_parameters()
+    
+    # Create the GPT2LMModel.
+    elif model_type == "gpt2":
+        print("Creating GPT2LMModel...")
+        from .models.gpttwo import GPT2LMModel, GPT2LMModelConfig
+        model_config_object = from_dict(GPT2LMModelConfig, OmegaConf.to_container(model_config))
+        model = GPT2LMModel(model_config_object)
+    
+    # Create the MambaLM.
+    elif model_type == "mamba":
+        print("Creating Mamba LM...")
+        from mambapy.lm import LM, MambaConfig
+        model_config_object = from_dict(MambaConfig, OmegaConf.to_container(model_config))
+        model = LM(model_config_object, model_config.vocab_size)
+    
+    # Create the Transformer.
+    elif model_type == "transformer":
+        from .models.transformer import TransformerConfig, Transformer
+        model_config_object = from_dict(TransformerConfig, OmegaConf.to_container(model_config))
+        model = Transformer(model_config_object)
+    
+    # Create a Pharia instance.
+    elif model_type == "pharia":
+        from .models.pharia import PhariaConfig, PhariaModel
+        model_config_object = from_dict(PhariaConfig, OmegaConf.to_container(model_config))
+        model = PhariaModel(model_config_object)
+    
+    # Create a TransformerXL instance.
+    else:
+        raise ValueError(f"Unknown model type: {model_type}")
+    
+    # Move the model to the device.
+    model.to(device)
+    return model
+
+
+def convert_tokens_to_songdata(tokens):
+
+    if isinstance(tokens, str):
+        tokens = tokens.split()
+
+    song_data = {}
+
+    song_data["tracks"] = []
+
+    current_track_index = 0
+    current_timestep = 0
+    for token in tokens:
+        if token == "GARLAND_START":
+            pass
+        elif token == "BAR_START":
+            if current_track_index == len(song_data["tracks"]):
+                song_data["tracks"] += [{"bars": [], "instrument": "0"}]
+            bar_data = {"notes": []}
+            song_data["tracks"][current_track_index]["bars"] += [bar_data]
+            current_timestep = 0
+        elif token.startswith("INST="):
+            instrument = token.split("=")[1]
+            song_data["tracks"][current_track_index]["instrument"] = instrument
+        elif token.startswith("DENSITY="):
+            pass
+        elif token.startswith("NOTE_ON="):
+            note_pitch = int(token.split("=")[1])
+            note_data = {
+                "note": note_pitch,
+                "start": current_timestep,
+                "end": current_timestep,
+                "veloctiy": 80
+            }
+            song_data["tracks"][current_track_index]["bars"][-1]["notes"] += [note_data]
+            pass
+        elif token.startswith("TIME_DELTA="):
+            current_timestep += int(token.split("=")[1])
+        elif token.startswith("NOTE_OFF="):
+            note_pitch = int(token.split("=")[1])
+            for note_data in song_data["tracks"][current_track_index]["bars"][-1]["notes"]:
+                if note_data["note"] == note_pitch and note_data["start"] == note_data["end"]:
+                    note_data["end"] = current_timestep
+                    break
+            pass
+        elif token == "BAR_END":
+            current_track_index += 1
+        elif token == "NEXT":
+            current_track_index = 0
+        elif token == "GARLAND_END":
+            pass
+        elif token == "[PAD]":
+            pass
+        elif token == "[EOS]":
+            pass
+        else:
+            raise Exception(f"Unknown token: {token}")
+    
+    assert isinstance(song_data, dict)
+    return song_data
+
+
+def convert_songdata_to_notesequence(song_data:dict, quantize_steps_per_quarter=8, remove_disabled_tracks=True):
+
+    assert isinstance(song_data, dict), f"Invalid song data type: {type(song_data)}"
+
+    # Clone the song data.
+    song_data = copy.deepcopy(song_data)
+
+    # Sort the tracks by instrument.
+    assert "tracks" in song_data, f"Invalid song data: {song_data.keys()}"
+    tracks = sorted(song_data["tracks"], key=lambda t: t["instrument"])
+    song_data["tracks"] = tracks
+
+    # Remove tracks that are not enabled.
+    if remove_disabled_tracks:
+        song_data["tracks"] = [t for t in song_data["tracks"] if t.get("enabled", True)]
+
+    # Create an empy note sequence.
+    note_sequence = note_seq.protobuf.music_pb2.NoteSequence()
+
+    # Add the tempo.
+    bpm = song_data["bpm"] if "bpm" in song_data else 120
+    note_sequence.tempos.add().qpm = bpm
+
+    # Compute some lengths.
+    step_length_seconds = 60.0 / bpm / quantize_steps_per_quarter
+    bar_length_seconds = 4 * step_length_seconds * quantize_steps_per_quarter
+
+    # Get the instruments.
+    instruments = list(set([t["instrument"] for t in song_data["tracks"]]))
+
+    # Add the tracks.
+    for track_index, track_data in enumerate(song_data["tracks"]):
+        instrument = track_data["instrument"]
+        for bar_index, bar_data in enumerate(track_data["bars"]):
+            bar_start_time = bar_index * bar_length_seconds
+            for note_data in bar_data["notes"]:
+                assert "note" in note_data
+                assert "start" in note_data
+                assert "end" in note_data
+                note = note_sequence.notes.add()
+                #note.instrument = instrument TODO
+                note.pitch = note_data["note"]
+                note.start_time = note_data["start"] * step_length_seconds + bar_start_time
+                note.end_time = note_data["end"] * step_length_seconds + bar_start_time
+                if "velocity" in note_data:
+                    note.velocity = note_data["velocity"]
+                else:
+                    note.velocity = 80
+                note.instrument = track_index
+                if instrument == "drums":
+                    note.is_drum = True
+                else:
+                    note.is_drum = False
+                    note.program = int(instrument)
+
+    return note_sequence
+
+
+def convert_songdata_to_pianoroll(song_data):
+
+    # The bars are 4/4 and the quantization is 8 steps per quarter, aka 32 steps per bar.
+    # We will render a grid. The height is 64 pixels. The width is 32 pixels per bar
+
+    # Create a new image.
+    lengths = [len(track["bars"]) for track in song_data["tracks"]]
+    if lengths == []:
+        return None
+    assert len(set(lengths)) == 1, f"Unequal number of bars: {lengths}"
+    num_bars = lengths[0]
+
+    # Get the note extremes.
+    min_note = 128
+    max_note = 0
+    for track_data in song_data["tracks"]:
+        for bar_data in track_data["bars"]:
+            for note_data in bar_data["notes"]:
+                min_note = min(min_note, note_data["note"])
+                max_note = max(max_note, note_data["note"])
+
+    # The width depends on the bars.
+    width = 32 * num_bars
+    
+    # The width depends on the notes.
+    height = 1 + max_note - min_note
+
+    # Create the image.
+    image = Image.new("RGB", (width, height), (14, 17, 23))
+
+    # Define some colors.
+    base_color = (255, 75, 75)
+    adjustments = [1.2, 1.0, 0.8, 0.6]
+    colors = []
+    for adjustment in adjustments:
+        import colorsys
+        rgb = base_color
+        rgb = [float(c) / 255.0 for c in rgb]
+        hsv = colorsys.rgb_to_hsv(*rgb)
+        # Rotate the hue.
+        offset = (adjustment - 1.0) * 0.1
+        hsv = (hsv[0] + offset, hsv[1], hsv[2])
+        rgb = colorsys.hsv_to_rgb(*hsv)
+        rgb = tuple([int(255.0 * c) for c in rgb])
+        colors += [rgb]
+        print("")
+
+    for color in colors:
+        print(color)
+        
+
+
+    # Draw the grid.
+    for track_index, track_data in enumerate(song_data["tracks"]):
+        color = colors[track_index % len(colors)]
+        for bar_index, bar_data in enumerate(track_data["bars"]):
+            x = bar_index * 32
+            
+            for note_data in bar_data["notes"]:
+                y = max_note - note_data["note"]
+                assert y >= 0 and y < height, f"Invalid y: {y}, note {note_data['note']} min_note: {min_note}, max_note: {max_note}, difference: {max_note - min_note}, height: {height}"
+                for i in range(note_data["start"], note_data["end"]):
+                    image.putpixel((x + i, y), color)
+
+    # Resize the image. Use nearest neighbor for pixel art.
+    factor = 4
+    image = image.resize((width * factor, height * factor), Image.NEAREST)
+
+    return image
+
+
+def convert_notesequence_to_wave(note_sequence):
+
+    if len(note_sequence.notes) == 0:
+        return None
+
+    try:
+        synthesizer = note_seq.fluidsynth
+        wave = synthesizer(note_sequence, sample_rate=44100)
+        return wave
+    except Exception as e:
+        synthesizer = note_seq.synthesize
+        wave = synthesizer(note_sequence)
+        return wave
+    
+
+def convert_notesequence_to_midi(note_sequence, filename="output.mid"):
+
+    if len(note_sequence.notes) == 0:
+        return None
+
+    # Returns the file content of the midi file.
+    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
+        filename = temp_file.name
+        note_seq.sequence_proto_to_midi_file(note_sequence, filename)
+        with open(filename, "rb") as file:
+            content = file.read()
+    return content
+
+
+