Generate music
from setuptools import setup, find_packages


def main():
    setup(
        name="tuneease",
        packages=find_packages(exclude=["tests*", "test_*", "*tests*"]),
        version="1.0.0",
        license="MIT",
        description="Generate music with AI",
        author="Pranit Shah",
        author_email="ppshah2023@gmail.com",
        url="https://www.pranitshah.cyou/musetune",
        download_url="https://github.com/Pshah2023/TuneEase/",
        keywords=["music", "AI", "getmusic"],
        install_requires=[
            "torch==2.0.1",
            "torchaudio==2.0.2",
            "torchvision==0.15.2",
            "absl-py==1.4.0",
            "blinker==1.6.2",
            "cachetools==5.3.1",
            "certifi==2023.7.22",
            "chardet==5.2.0",
            "charset-normalizer==3.2.0",
            "colorama==0.4.6",
            "contourpy==1.1.0",
            "cycler==0.11.0",
            "einops==0.6.1",
            "Flask==2.3.2",
            "Flask-Cors==4.0.0",
            "fonttools==4.42.0",
            "fsspec==2023.6.0",
            "google-auth==2.22.0",
            "google-auth-oauthlib==1.0.0",
            "grpcio==1.57.0",
            "huggingface-hub==0.16.4",
            "idna==3.4",
            "itsdangerous==2.1.2",
            "Jinja2==3.1.2",
            "joblib==1.3.2",
            "jsonpickle==3.0.2",
            "kiwisolver==1.4.4",
            "Markdown==3.4.4",
            "MarkupSafe==2.1.3",
            "matplotlib==3.7.2",
            "miditoolkit==0.1.16",
            "mido==1.3.0",
            "more-itertools==10.1.0",
            "music21==8.3.0",
            "numpy==1.23.5",  # this version is crucial for it to work out of the box
            "oauthlib==3.2.2",
            "packaging==23.1",
            "Pillow==10.0.0",
            "protobuf==4.24.0",
            "pyasn1==0.5.0",
            "pyasn1-modules==0.3.0",
            "pyparsing==3.0.9",
            "python-dateutil==2.8.2",
            "PyYAML==6.0.1",
            "regex==2023.8.8",
            "requests==2.31.0",
            "requests-oauthlib==1.3.1",
            "rsa==4.9",
            "safetensors==0.3.2",
            "scipy==1.11.1",
            "six==1.16.0",
            "tensorboard==2.14.0",
            "tensorboard-data-server==0.7.1",
            "tokenizers==0.13.3",
            "tqdm==4.66.1",
            "transformers==4.31.0",
            "typing_extensions==4.7.1",
            "urllib3==1.26.16",
            "webcolors==1.13",
            "Werkzeug==2.3.7",
        ],
        include_package_data=True,
        classifiers=[
            "Development Status :: 5 - Production/Stable",
            "Intended Audience :: Developers",
            "Topic :: Artistic Software",
            "License :: OSI Approved :: MIT License",
            "Programming Language :: Python :: 3.8",
            "Programming Language :: Python :: 3.9",
            "Programming Language :: Python :: 3.10",
            "Programming Language :: Python :: 3.11",
        ],
        entry_points={
            "console_scripts": [
                "tuneease=tuneease.__main__:main",
                "tuneease-generate=tuneease.tuneease:tuneease",
            ]
        },
    )


if __name__ == "__main__":
    main()

setup.py
Pshah2023/TuneEase
from .server import run_app


def main():
    run_app()


if __name__ == "__main__":
    main()

tuneease/__main__.py
Pshah2023/TuneEase
import os
import subprocess
from .logger import ServerLogger
from .pathutility import PathUtility


class Converter:
    """
    A class for converting music files to MusicXML format using MuseScore.

    Attributes:
        logger (ServerLogger): An instance of ServerLogger for logging messages.
        util (PathUtility): An instance of PathUtility for handling file paths.
        project_directory (str): The path to the project directory.
        museScore_path (str): The path to the MuseScore executable.

    Methods:
        __init__(self, log_file="tuneease.log", museScore_path=None):
        save_file(self, file, destination_folder="temp"):
        convert_to(self, filepath, output_extension=".xml"):
        convert(self, file, output_extension=".xml"):

    Raises:
        ValueError: If museScore_path is not provided during initialization.
    """

    def __init__(self, log_file="tuneease.log", museScore_path=None):
        self.logger = ServerLogger(log_file).get()
        self.util = PathUtility()
        self.project_directory = self.util.project_directory()
        if museScore_path == None:
            raise ValueError(
                "The value of the museScorePath cannot be null in this case. "
                "Provide input arguments to the path of museScore when running server.py"
            )
        self.museScore_path = museScore_path

    def save_file(self, file, destination_folder="temp"):
        """
        Save an uploaded file from the server to a temporary file storage folder.

        Args:
            file (FileStorage): The uploaded file object.
            destination_folder (str, optional): The name of the temporary folder. Defaults to "temp".

        Returns:
            str: The path to the saved file.

        Raises:
            AssertionError: If the uploaded file is empty or has no filename.
            Exception: If any other error occurs during file saving.

        Example:
            filepath = converter.save_file(uploaded_file, destination_folder="uploads")
        """
        try:
            assert file.filename != "", "File is empty"
            filepath = os.path.join(
                self.project_directory, destination_folder, file.filename
            )
            os.makedirs(os.path.dirname(filepath), exist_ok=True)
            file.save(filepath)
            return filepath
        except Exception as e:
            self.logger.error(
                f"Failed to save the file. Error message: {str(e)}", "error"
            )
            raise

    def convert_to(self, filepath, output_extension: str = ".xml"):
        """
        Convert some music file to the MusicXML format using MuseScore.

        Args:
            filepath (str): The path to the input music file.
            output_extension (str, optional): The file extension for the converted output. Defaults to ".xml".

        Returns:
            str: The path to the converted MusicXML file.

        Raises:
            subprocess.CalledProcessError: If the MuseScore command fails.
            Exception: If any other error occurs during conversion.

        Example:
            output_filepath = converter.convert_to(input_filepath, output_extension=".mxl")
        """
        output_filepath = os.path.splitext(filepath)[0] + output_extension
        try:
            subprocess.run(
                [self.museScore_path, filepath, "-o", output_filepath],
                check=True,
                cwd=self.project_directory,
            )
            return output_filepath
        except subprocess.CalledProcessError as e:
            self.logger.error(
                f"Failed to run MuseScore command. Error message: {str(e)}"
            )
        except Exception as e:
            self.logger.error(
                f"Failed to convert score to MusicXML. Error message: {str(e)}"
            )

    def convert(self, file, output_extension: str = ".xml"):
        """
        Save and convert an uploaded music file from the server to the MusicXML format.

        Args:
            file (FileStorage): The uploaded music file object.
            output_extension (str, optional): The file extension for the converted output. Defaults to ".xml".

        Returns:
            str: The path to the converted MusicXML file.

        Example:
            converted_filepath = converter.convert(uploaded_music_file, output_extension=".mxl")
        """
        filepath = self.save_file(file)
        output_filepath = self.convert_to(filepath, output_extension=output_extension)
        return output_filepath

tuneease/converter.py
Pshah2023/TuneEase
import random
import itertools as it
from ...getmusic.data.indexed_datasets import IndexedDataset
import numpy as np
import torch
from torch.utils.data import Dataset


class BigDataset(Dataset):
    def __init__(self, prefix, vocab_size, path=None):
        self.data_dir = path
        self.prefix = prefix
        self.ds_name = "train" if prefix == "train" else "valid"
        self.size = np.load(f"{self.data_dir}/{self.ds_name}_length.npy")
        self.empty_idx = vocab_size - 1
        self.indexed_ds = None

    def __len__(self):
        return self.size

    def _get_item(self, index):
        if self.indexed_ds is None:
            self.indexed_ds = IndexedDataset(f"{self.data_dir}/{self.ds_name}")
        return self.indexed_ds[index]

    def __getitem__(self, index):
        item = self._get_item(index)
        return item

    def collater(self, samples):
        if len(samples) == 0:
            assert 1 == 0
            return {}

        batch = {}
        batch["tempo"] = torch.LongTensor(np.array([s[0, -1] for s in samples]))
        batch["data"] = []
        batch["condition_pos"] = []
        batch["not_empty_pos"] = []

        for sample in samples:
            sample = sample[:, :-1]
            assert sample.shape == (14, 512)

            figure_size = 512

            track_not_empty_bool = torch.tensor(
                (sample == self.empty_idx).astype(float).sum(-1) != figure_size
            ).float()

            # because preprocessed music has at least 2 tracks
            # we have to randomly perform a single track generation
            if random.randint(0, 9) == 0:
                not_empty_track_index = torch.nonzero(
                    track_not_empty_bool[:-2]
                )  # can't only keep chord
                not_empty_track_index = [
                    i // 2
                    for i in not_empty_track_index.view(-1).tolist()
                    if i % 2 == 0
                ]
                single_track_index = random.choice(not_empty_track_index)
                sample[: 2 * single_track_index] = self.empty_idx
                sample[2 * single_track_index + 2 : -2] = self.empty_idx

            # which track is empty
            track_not_empty_bool = torch.tensor(
                (sample == self.empty_idx).astype(float).sum(-1) != figure_size
            ).float()

            track_not_empty_num = track_not_empty_bool.sum()

            content_num = random.randint(
                1, int(track_not_empty_num / 2) - 1
            )  # chord can not be a content track

            condition_bool = track_not_empty_bool.clone()

            # if content num != not empty num, conditional generation
            if content_num != track_not_empty_num:
                not_empty_track_index = torch.nonzero(track_not_empty_bool[:-2])
                not_empty_track_index = [
                    i // 2
                    for i in not_empty_track_index.view(-1).tolist()
                    if i % 2 == 0
                ]
                content_combination = it.combinations(
                    not_empty_track_index, content_num
                )
                content_combination = [i for i in content_combination]
                content_track_index = random.choice(content_combination)
                for c_idx in content_track_index:
                    condition_bool[c_idx * 2] = 0
                    condition_bool[c_idx * 2 + 1] = 0

            # randomly disable chord guidance
            if random.randint(0, 3) < 3:
                sample[-2:] = self.empty_idx
                condition_bool[-2:] = 0

            assert (track_not_empty_bool * condition_bool != track_not_empty_bool).any()
            condition_bool = condition_bool.view(14, 1).repeat(1, figure_size)

            track_not_empty_bool = (torch.tensor(sample) != self.empty_idx).float()

            batch["data"].append(sample)
            batch["condition_pos"].append(condition_bool)
            batch["not_empty_pos"].append(track_not_empty_bool)

        batch["data"] = torch.LongTensor(np.array(batch["data"]))
        batch["condition_pos"] = torch.stack(batch["condition_pos"], dim=0)
        batch["not_empty_pos"] = torch.stack(batch["not_empty_pos"], dim=0)

        return batch

tuneease/getmusic/data/bigdata.py
Pshah2023/TuneEase
import torch
from torch.utils.data import ConcatDataset
from ..utils.misc import instantiate_from_config


def build_dataloader(config, args=None, return_dataset=False):
    dataset_cfg = config["dataloader"]
    train_dataset = []
    for ds_cfg in dataset_cfg["train_datasets"]:
        ds = instantiate_from_config(ds_cfg)
        train_dataset.append(ds)
    if len(train_dataset) > 1:
        train_dataset = ConcatDataset(train_dataset)
    else:
        train_dataset = train_dataset[0]

    val_dataset = []
    for ds_cfg in dataset_cfg["validation_datasets"]:
        ds = instantiate_from_config(ds_cfg)
        val_dataset.append(ds)
    if len(val_dataset) > 1:
        val_dataset = ConcatDataset(val_dataset)
    else:
        val_dataset = val_dataset[0]

    if args is not None and args.distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(
            train_dataset, shuffle=True
        )
        val_sampler = torch.utils.data.distributed.DistributedSampler(
            val_dataset, shuffle=False
        )
        train_iters = len(train_sampler) // dataset_cfg["batch_size"]
        val_iters = len(val_sampler) // dataset_cfg["batch_size"]
    else:
        train_sampler = None
        val_sampler = None
        train_iters = len(train_dataset) // dataset_cfg["batch_size"]
        val_iters = len(val_dataset) // dataset_cfg["batch_size"]

    num_workers = dataset_cfg["num_workers"]
    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=dataset_cfg["batch_size"],
        shuffle=(train_sampler is None),
        num_workers=num_workers,
        pin_memory=True,
        sampler=train_sampler,
        drop_last=True,
        collate_fn=train_dataset.collater,
        persistent_workers=True,
    )

    val_loader = torch.utils.data.DataLoader(
        val_dataset,
        batch_size=dataset_cfg["batch_size"],
        # batch_size=1,
        shuffle=False,  # (val_sampler is None),
        num_workers=num_workers,
        pin_memory=True,
        sampler=val_sampler,
        drop_last=True,
        collate_fn=val_dataset.collater,
        persistent_workers=True,
    )

    dataload_info = {
        "train_loader": train_loader,
        "validation_loader": val_loader,
        "train_iterations": train_iters,
        "validation_iterations": val_iters,
    }

    if return_dataset:
        dataload_info["train_dataset"] = train_dataset
        dataload_info["validation_dataset"] = val_dataset

    return dataload_info

tuneease/getmusic/data/build.py
Pshah2023/TuneEase
import torch
from torch import distributed as dist
from torch import multiprocessing as mp
from . import distributed as dist_fn


def find_free_port():
    import socket

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    sock.bind(("", 0))
    port = sock.getsockname()[1]
    sock.close()

    return port


def launch(fn, n_gpu_per_machine, n_machine=1, machine_rank=0, dist_url=None, args=()):
    world_size = n_machine * n_gpu_per_machine

    if world_size > 1:
        if dist_url == "auto":
            if n_machine != 1:
                raise ValueError('dist_url="auto" not supported in multi-machine jobs')
            port = find_free_port()
            dist_url = f"tcp://127.0.0.1:{port}"
        print("dist_url ", dist_url)
        print("n_machine ", n_machine)
        print("args ", args)
        print("world_size ", world_size)
        print("machine_rank ", machine_rank)
        if n_machine > 1 and dist_url.startswith("file://"):
            raise ValueError(
                "file:// is not a reliable init method in multi-machine jobs. Prefer tcp://"
            )

        mp.spawn(
            distributed_worker,
            nprocs=n_gpu_per_machine,
            args=(fn, world_size, n_gpu_per_machine, machine_rank, dist_url, args),
            daemon=False,
        )
    else:
        local_rank = 0
        fn(local_rank, *args)


def distributed_worker(
    local_rank, fn, world_size, n_gpu_per_machine, machine_rank, dist_url, args
):
    if not torch.cuda.is_available():
        raise OSError("CUDA is not available. Please check your environments")

    global_rank = machine_rank * n_gpu_per_machine + local_rank
    print("local_rank ", local_rank)
    print("global_rank ", global_rank)
    try:
        dist.init_process_group(
            backend="NCCL",
            init_method=dist_url,
            world_size=world_size,
            rank=global_rank,
        )

    except Exception:
        raise OSError("failed to initialize NCCL groups")

    dist_fn.synchronize()

    if n_gpu_per_machine > torch.cuda.device_count():
        raise ValueError(
            f"specified n_gpu_per_machine larger than available device ({torch.cuda.device_count()})"
        )

    torch.cuda.set_device(local_rank)

    if dist_fn.LOCAL_PROCESS_GROUP is not None:
        raise ValueError("torch.distributed.LOCAL_PROCESS_GROUP is not None")

    n_machine = world_size // n_gpu_per_machine
    for i in range(n_machine):
        ranks_on_i = list(range(i * n_gpu_per_machine, (i + 1) * n_gpu_per_machine))
        pg = dist.new_group(ranks_on_i)

        if i == machine_rank:
            dist_fn.LOCAL_PROCESS_GROUP = pg

    fn(local_rank, *args)

tuneease/getmusic/distributed/launch.py
Pshah2023/TuneEase
class LinearDecayLRWithWarmup(object):
    """
    adjust lr:

    args:
        warmup_lr: float or None, the learning rate to be touched after warmup
        warmup: int, the number of steps to warmup
    """

    def __init__(
        self,
        optimizer,
        T_max,
        last_epoch=-1,
        verbose=False,
        min_lr=0,
        warmup_lr=None,
        warmup=0,
    ):
        self.optimizer = optimizer
        self.T_max = T_max
        self.last_epoch = last_epoch
        self.verbose = verbose
        self.warmup_lr = warmup_lr
        self.warmup = warmup

        if isinstance(min_lr, list) or isinstance(min_lr, tuple):
            if len(min_lr) != len(optimizer.param_groups):
                raise ValueError(
                    "expected {} min_lrs, got {}".format(
                        len(optimizer.param_groups), len(min_lr)
                    )
                )
            self.min_lrs = list(min_lr)
        else:
            self.min_lrs = [min_lr] * len(optimizer.param_groups)
        self.max_lrs = [lr for lr in self.min_lrs]

        self._prepare_for_warmup()

    def step(self):
        epoch = self.last_epoch + 1
        self.last_epoch = epoch

        if epoch <= self.warmup:
            self._increase_lr(epoch)
        else:
            self._reduce_lr(epoch)

    def _reduce_lr(self, epoch):
        for i, param_group in enumerate(self.optimizer.param_groups):
            progress = float(epoch - self.warmup) / float(
                max(1, self.T_max - self.warmup)
            )
            factor = max(0.0, 1 - progress)
            old_lr = float(param_group["lr"])
            new_lr = max(self.max_lrs[i] * factor, self.min_lrs[i])
            param_group["lr"] = new_lr
            if self.verbose:
                print(
                    "Epoch {:5d}: reducing learning rate"
                    " of group {} to {:.4e}.".format(epoch, i, new_lr)
                )

    def _increase_lr(self, epoch):
        for i, param_group in enumerate(self.optimizer.param_groups):
            old_lr = float(param_group["lr"])
            new_lr = old_lr + self.warmup_lr_steps[i]
            param_group["lr"] = new_lr
            self.max_lrs[i] = max(self.max_lrs[i], new_lr)
            if self.verbose:
                print(
                    "Epoch {:5d}: increasing learning rate"
                    " of group {} to {:.4e}.".format(epoch, i, new_lr)
                )

    def _prepare_for_warmup(self):
        if self.warmup_lr is not None:
            if isinstance(self.warmup_lr, (list, tuple)):
                if len(self.warmup_lr) != len(self.optimizer.param_groups):
                    raise ValueError(
                        "expected {} warmup_lrs, got {}".format(
                            len(self.optimizer.param_groups), len(self.warmup_lr)
                        )
                    )
                self.warmup_lrs = list(self.warmup_lr)
            else:
                self.warmup_lrs = [self.warmup_lr] * len(self.optimizer.param_groups)
        else:
            self.warmup_lrs = None
        if self.warmup > self.last_epoch:
            curr_lrs = [group["lr"] for group in self.optimizer.param_groups]
            self.warmup_lr_steps = [
                max(0, (self.warmup_lrs[i] - curr_lrs[i]) / float(self.warmup))
                for i in range(len(curr_lrs))
            ]
        else:
            self.warmup_lr_steps = None

    def state_dict(self):
        return {
            key: value for key, value in self.__dict__.items() if key != "optimizer"
        }

    def load_state_dict(self, state_dict):
        self.__dict__.update(state_dict)
        self._prepare_for_warmup()

tuneease/getmusic/engine/lr_scheduler.py
Pshah2023/TuneEase
import os
import time
import math
import torch
from ..utils.misc import (
    instantiate_from_config,
    format_seconds,
    get_model_parameters_info,
)
from ..distributed.distributed import reduce_dict
from ..distributed.distributed import is_primary
from ..engine.ema import EMA
from ..utils import midi_config as mc
from ...pipeline.presets import RootKinds
from ...pipeline.key_chord import TokenHelper

try:
    from torch.cuda.amp import autocast, GradScaler

    AMP = True
except:
    print("Warning: import torch.amp failed, so no amp will be used!")
    AMP = False


class Solver(object):
    def __init__(self, gpu, config, args, model, dataloader, logger, is_sample=False):
        self.gpu = gpu
        self.config = config
        self.args = args
        self.model = model
        self.dataloader = dataloader
        self.logger = logger
        self.t_h = TokenHelper()
        self.r_h = RootKinds()

        self.max_epochs = config["solver"]["max_epochs"]
        self.save_epochs = config["solver"]["save_epochs"]
        self.save_iterations = config["solver"].get("save_iterations", -1)
        self.validate_iterations = config["solver"]["validate_iterations"]

        self.validation_epochs = config["solver"]["validation_epochs"]  # 多少个epoch验证一次
        assert isinstance(self.save_epochs, (int, list))
        assert isinstance(self.validation_epochs, (int, list))

        self.last_epoch = -1
        self.last_iter = -1
        if not is_sample:
            self.ckpt_dir = os.path.join(args.save_dir, "checkpoint")
            self.oct_dir = os.path.join(args.save_dir, "oct")
            os.makedirs(self.ckpt_dir, exist_ok=True)
            os.makedirs(self.oct_dir, exist_ok=True)

        self.logger.debug(
            "Load dictionary: {} tokens.".format(len(self.t_h.ids_to_tokens))
        )

        beat_note_factor = mc.BEAT_NOTE_FACTOR
        max_notes_per_bar = mc.MAX_NOTES_PER_BAR
        pos_resolution = mc.POS_RESOLUTION
        bar_max = mc.BAR_MAX
        self.pos_in_bar = beat_note_factor * max_notes_per_bar * pos_resolution
        self.pad_index = mc.DURATION_MAX * mc.POS_RESOLUTION - 1
        self.figure_size = (
            mc.BAR_MAX * mc.BEAT_NOTE_FACTOR * mc.MAX_NOTES_PER_BAR * mc.POS_RESOLUTION
        )

        if "clip_grad_norm" in config["solver"]:
            self.clip_grad_norm = instantiate_from_config(
                config["solver"]["clip_grad_norm"]
            )
        else:
            self.clip_grad_norm = None

        adjust_lr = config["solver"].get("adjust_lr", "sqrt")
        base_lr = config["solver"].get("base_lr", 1.0e-4)

        if adjust_lr == "none":
            self.lr = base_lr
        elif adjust_lr == "sqrt":
            self.lr = base_lr * math.sqrt(
                args.world_size * config["dataloader"]["batch_size"]
            )
        elif adjust_lr == "linear":
            self.lr = base_lr * args.world_size * config["dataloader"]["batch_size"]
        else:
            raise NotImplementedError("Unknown type of adjust lr {}!".format(adjust_lr))
        self.logger.debug(
            "Get lr {} from base lr {} with {}".format(self.lr, base_lr, adjust_lr)
        )

        if hasattr(model, "get_optimizer_and_scheduler") and callable(
            getattr(model, "get_optimizer_and_scheduler")
        ):
            optimizer_and_scheduler = model.get_optimizer_and_scheduler(
                config["solver"]["optimizers_and_schedulers"]
            )
        else:
            optimizer_and_scheduler = self._get_optimizer_and_scheduler(
                config["solver"]["optimizers_and_schedulers"]
            )

        assert type(optimizer_and_scheduler) == type(
            {}
        ), "optimizer and schduler should be a dict!"
        self.optimizer_and_scheduler = optimizer_and_scheduler

        if "ema" in config["solver"] and args.local_rank == 0:
            ema_args = config["solver"]["ema"]
            ema_args["model"] = self.model
            self.ema = EMA(**ema_args)
        else:
            self.ema = None

        self.logger.debug(str(get_model_parameters_info(self.model)))

        if self.gpu:
            self.model.to(self.args.local_rank)

        self.device = self.args.local_rank

        if self.args.distributed:
            self.logger.debug("Distributed, begin DDP the model...")
            self.model = torch.nn.parallel.DistributedDataParallel(
                self.model,
                device_ids=[self.args.local_rank],
                find_unused_parameters=False,
            )  #
            self.logger.debug("Distributed, DDP model done!")

        self.args.amp = self.args.amp and AMP
        if self.args.amp:
            self.scaler = GradScaler()
            self.logger.debug("Using AMP for training!")
        self.batch_size = self.config["dataloader"]["batch_size"]
        self.logger.debug(
            "{}: global rank {}: prepare solver done!".format(
                self.args.name, self.args.global_rank
            )
        )

    def _get_optimizer_and_scheduler(self, op_sc_list):
        optimizer_and_scheduler = {}
        for op_sc_cfg in op_sc_list:
            op_sc = {
                "name": op_sc_cfg.get("name", "none"),
                "start_epoch": op_sc_cfg.get("start_epoch", 0),
                "end_epoch": op_sc_cfg.get("end_epoch", -1),
                "start_iteration": op_sc_cfg.get("start_iteration", 0),
                "end_iteration": op_sc_cfg.get("end_iteration", -1),
            }
            if op_sc["name"] == "none":
                parameters = filter(lambda p: p.requires_grad, self.model.parameters())
            else:
                parameters = self.model.parameters(name=op_sc["name"])

            op_cfg = op_sc_cfg.get(
                "optimizer", {"target": "torch.optim.SGD", "params": {}}
            )

            if "params" not in op_cfg:
                op_cfg["params"] = {}
            if "lr" not in op_cfg["params"]:
                op_cfg["params"]["lr"] = self.lr
            op_cfg["params"]["params"] = parameters
            optimizer = instantiate_from_config(op_cfg)

            op_sc["optimizer"] = {
                "module": optimizer,
                "step_iteration": op_cfg["step_iteration"],
            }

            assert isinstance(
                op_sc["optimizer"]["step_iteration"], int
            ), "optimizer steps should be a integer number of iterations"

            if "scheduler" in op_sc_cfg:
                sc_cfg = op_sc_cfg["scheduler"]
                sc_cfg["params"]["optimizer"] = optimizer
                if sc_cfg["target"].split(".")[-1] in [
                    "CosineAnnealingLRWithWarmup",
                    "CosineAnnealingLR",
                ]:
                    T_max = self.max_epochs * self.dataloader["train_iterations"]
                    sc_cfg["params"]["T_max"] = T_max
                scheduler = instantiate_from_config(sc_cfg)
                op_sc["scheduler"] = {
                    "module": scheduler,
                    "step_iteration": sc_cfg.get("step_iteration", 1),
                }
                if op_sc["scheduler"]["step_iteration"] == "epoch":
                    op_sc["scheduler"]["step_iteration"] = self.dataloader[
                        "train_iterations"
                    ]
            optimizer_and_scheduler[op_sc["name"]] = op_sc

        return optimizer_and_scheduler

    def _get_lr(self, return_type="str"):
        lrs = {}
        for op_sc_n, op_sc in self.optimizer_and_scheduler.items():
            lr = op_sc["optimizer"]["module"].state_dict()["param_groups"][0]["lr"]
            lrs[op_sc_n + "_lr"] = round(lr, 10)
        if return_type == "str":
            lrs = str(lrs)
            lrs = (
                lrs.replace("none", "lr")
                .replace("{", "")
                .replace("}", "")
                .replace("'", "")
            )
        elif return_type == "dict":
            pass
        else:
            raise ValueError("Unknow of return type: {}".format(return_type))
        return lrs

    def step(self, batch, phase="train"):
        loss = {}

        for op_sc_n, op_sc in self.optimizer_and_scheduler.items():
            if phase == "train":
                if op_sc["start_iteration"] > self.last_iter:
                    continue
                if (
                    op_sc["end_iteration"] > 0
                    and op_sc["end_iteration"] <= self.last_iter
                ):
                    continue
                if op_sc["start_epoch"] > self.last_epoch:
                    continue
                if op_sc["end_epoch"] > 0 and op_sc["end_epoch"] <= self.last_epoch:
                    continue

            input = {
                "batch": batch["data"].to(self.device),
                "tempo": batch["tempo"].to(self.device),
                "condition_pos": batch["condition_pos"].to(self.device),
                "not_empty_pos": batch["not_empty_pos"].to(self.device),
                "return_loss": True,
                "step": self.last_iter,
            }
            if op_sc_n != "none":
                input["name"] = op_sc_n

            if phase == "train":
                if self.args.amp:
                    with autocast():
                        output = self.model(**input)
                else:
                    output = self.model(**input)
            else:
                with torch.no_grad():
                    if self.args.amp:
                        with autocast():
                            output = self.model(**input)
                    else:
                        output = self.model(**input)

            if phase == "train":
                output["loss"] /= op_sc["optimizer"]["step_iteration"]
                output["loss"].backward()
                if self.clip_grad_norm is not None:
                    self.clip_grad_norm(self.model.parameters())
                if (
                    op_sc["optimizer"]["step_iteration"] > 0
                    and (self.last_iter + 1) % op_sc["optimizer"]["step_iteration"] == 0
                ):
                    op_sc["optimizer"]["module"].step()
                    op_sc["optimizer"]["module"].zero_grad(set_to_none=True)

                if "scheduler" in op_sc:
                    if (
                        op_sc["scheduler"]["step_iteration"] > 0
                        and (self.last_iter + 1) % op_sc["scheduler"]["step_iteration"]
                        == 0
                    ):
                        op_sc["scheduler"]["module"].step()

                if self.ema is not None:
                    self.ema.update(iteration=self.last_iter)

            loss[op_sc_n] = {
                k: v for k, v in output.items() if ("loss" in k or "acc" in k)
            }
        return loss

    def save(self, force=False):
        if is_primary():
            if self.save_iterations > 0:
                if (self.last_iter + 1) % self.save_iterations == 0:
                    save = True
                else:
                    save = False
            else:
                if isinstance(self.save_epochs, int):
                    save = (self.last_epoch + 1) % self.save_epochs == 0
                else:
                    save = (self.last_epoch + 1) in self.save_epochs

            if save or force:
                state_dict = {
                    "last_epoch": self.last_epoch,
                    "last_iter": self.last_iter,
                    "model": self.model.module.state_dict()
                    if isinstance(self.model, torch.nn.parallel.DistributedDataParallel)
                    else self.model.state_dict(),
                }
                if self.ema is not None:
                    state_dict["ema"] = self.ema.state_dict()
                if self.clip_grad_norm is not None:
                    state_dict["clip_grad_norm"] = self.clip_grad_norm.state_dict()

                optimizer_and_scheduler = {}
                for op_sc_n, op_sc in self.optimizer_and_scheduler.items():
                    state_ = {}
                    for k in op_sc:
                        if k in ["optimizer", "scheduler"]:
                            op_or_sc = {
                                kk: vv for kk, vv in op_sc[k].items() if kk != "module"
                            }
                            op_or_sc["module"] = op_sc[k]["module"].state_dict()
                            state_[k] = op_or_sc
                        else:
                            state_[k] = op_sc[k]
                    optimizer_and_scheduler[op_sc_n] = state_

                state_dict["optimizer_and_scheduler"] = optimizer_and_scheduler

                if save or force:
                    save_path = os.path.join(
                        self.ckpt_dir,
                        "{}e_{}iter.pth".format(
                            str(self.last_epoch).zfill(6), self.last_iter
                        ),
                    )
                    torch.save(state_dict, save_path)
                    self.logger.debug("saved in {}".format(save_path))

                save_path = os.path.join(self.ckpt_dir, "last.pth")
                torch.save(state_dict, save_path)
                self.logger.debug("saved in {}".format(save_path))

    def resume(self, path=None, load_optimizer_and_scheduler=True, load_others=True):
        if path is None:
            path = os.path.join(self.ckpt_dir, "last.pth")

        if os.path.exists(path):
            if self.gpu:
                state_dict = torch.load(
                    path, map_location="cuda:{}".format(self.args.local_rank)
                )
            else:
                state_dict = state_dict = torch.load(path, map_location="cpu")

            if load_others:
                self.last_epoch = state_dict["last_epoch"]
                self.last_iter = state_dict["last_iter"]

            if isinstance(self.model, torch.nn.parallel.DistributedDataParallel):
                try:
                    self.model.module.load_state_dict(state_dict["model"])
                except:
                    model_dict = self.model.module.state_dict()
                    temp_state_dict = {
                        k: v
                        for k, v in state_dict["model"].items()
                        if k in model_dict.keys()
                    }
                    model_dict.update(temp_state_dict)
                    self.model.module.load_state_dict(model_dict)
            else:
                self.model.load_state_dict(state_dict["model"])

            if "ema" in state_dict and self.ema is not None:
                try:
                    self.ema.load_state_dict(state_dict["ema"])
                except:
                    model_dict = self.ema.state_dict()
                    temp_state_dict = {
                        k: v
                        for k, v in state_dict["ema"].items()
                        if k in model_dict.keys()
                    }
                    model_dict.update(temp_state_dict)
                    self.ema.load_state_dict(model_dict)

            if "clip_grad_norm" in state_dict and self.clip_grad_norm is not None:
                self.clip_grad_norm.load_state_dict(state_dict["clip_grad_norm"])

            for op_sc_n, op_sc in state_dict["optimizer_and_scheduler"].items():
                for k in op_sc:
                    if k in ["optimizer", "scheduler"]:
                        for kk in op_sc[k]:
                            if kk == "module" and load_optimizer_and_scheduler:
                                self.optimizer_and_scheduler[op_sc_n][k][
                                    kk
                                ].load_state_dict(op_sc[k][kk])
                            elif load_others:  # such as step_iteration, ...
                                self.optimizer_and_scheduler[op_sc_n][k][kk] = op_sc[k][
                                    kk
                                ]
                    elif load_others:  # such as start_epoch, end_epoch, ....
                        self.optimizer_and_scheduler[op_sc_n][k] = op_sc[k]

            self.logger.debug(str("Resume from {}".format(path)))
        else:
            raise ValueError("checkpoint not found")

    def train_epoch(self):
        self.model.train()
        self.last_epoch += 1

        if self.args.distributed:
            self.dataloader["train_loader"].sampler.set_epoch(self.last_epoch)

        epoch_start = time.time()
        itr_start = time.time()
        itr = -1

        for itr, batch in enumerate(self.dataloader["train_loader"]):
            if itr == 0:
                print("time2 is " + str(time.time()))
            data_time = time.time() - itr_start
            step_start = time.time()
            self.last_iter += 1

            loss = self.step(batch, phase="train")

            if (
                self.logger is not None
                and self.last_iter % self.args.log_frequency == 0
            ):
                info = "{}: train".format(self.args.name)
                info = info + ": Epoch {}/{} iter {}/{}".format(
                    self.last_epoch,
                    self.max_epochs,
                    self.last_iter % self.dataloader["train_iterations"],
                    self.dataloader["train_iterations"],
                )
                for loss_n, loss_dict in loss.items():
                    info += " ||"
                    loss_dict = reduce_dict(loss_dict)
                    info += "" if loss_n == "none" else " {}".format(loss_n)

                    for k in loss_dict:
                        info += " | {}: {:.4f}".format(k, float(loss_dict[k]))
                        self.logger.debug(str("train/{}/{}".format(loss_n, k)))
                        self.logger.debug(str(float(loss_dict[k])))
                        self.logger.debug(str(self.last_iter))

                lrs = self._get_lr(return_type="dict")
                for k in lrs.keys():
                    lr = lrs[k]
                    self.logger.debug(str("train/{}_lr".format(k)))
                    self.logger.debug(str(lrs[k]))
                    self.logger.debug(str(self.last_iter))

                info += " || {}".format(self._get_lr())

                spend_time = time.time() - self.start_train_time
                itr_time_avg = spend_time / (self.last_iter + 1)
                info += " || data_time: {dt}s | fbward_time: {fbt}s | iter_time: {it}s | iter_avg_time: {ita}s | epoch_time: {et} | spend_time: {st} | left_time: {lt} ".format(
                    dt=round(data_time, 1),
                    it=round(time.time() - itr_start, 1),
                    fbt=round(time.time() - step_start, 1),
                    ita=round(itr_time_avg, 1),
                    et=format_seconds(time.time() - epoch_start),
                    st=format_seconds(spend_time),
                    lt=format_seconds(
                        itr_time_avg
                        * self.max_epochs
                        * self.dataloader["train_iterations"]
                        - spend_time
                    ),
                )

                if isinstance(self.model, torch.nn.parallel.DistributedDataParallel):
                    model = self.model.module
                else:
                    model = self.model

                self.logger.debug(str(info))

            itr_start = time.time()

            if (
                self.validate_iterations > 0
                and (self.last_iter + 1) % self.validate_iterations == 0
            ):
                if self.logger is not None:
                    self.logger.debug(
                        str("save model for iteration {}".format(self.last_iter))
                    )

                self.save(force=True)
                self.validate_epoch()
                self.model.train()

                if self.logger is not None:
                    self.logger.debug("save model done")

        assert itr >= 0, "The data is too less to form one iteration!"
        self.dataloader["train_iterations"] = itr + 1

    def validate_epoch(self):
        if self.logger is not None:
            self.logger.debug("Enter validate_epoch")

        if "validation_loader" not in self.dataloader:
            val = False
        else:
            if isinstance(self.validation_epochs, int):
                val = (self.last_epoch + 1) % self.validation_epochs == 0
            else:
                val = (self.last_epoch + 1) in self.validation_epochs

        if val:
            if self.args.distributed:
                self.dataloader["validation_loader"].sampler.set_epoch(self.last_epoch)
            self.model.eval()
            overall_loss = None
            epoch_start = time.time()
            itr_start = time.time()
            itr = -1
            for itr, batch in enumerate(self.dataloader["validation_loader"]):
                data_time = time.time() - itr_start
                step_start = time.time()
                loss = self.step(batch, phase="val")

                for loss_n, loss_dict in loss.items():
                    loss[loss_n] = reduce_dict(loss_dict)
                if overall_loss is None:
                    overall_loss = loss
                else:
                    for loss_n, loss_dict in loss.items():
                        for k, v in loss_dict.items():
                            overall_loss[loss_n][k] = (
                                overall_loss[loss_n][k] * itr + loss[loss_n][k]
                            ) / (itr + 1)

                if self.logger is not None and (itr + 1) % self.args.log_frequency == 0:
                    info = "{}: val".format(self.args.name)
                    info = info + ": Epoch {}/{} | iter {}/{}".format(
                        self.last_epoch,
                        self.max_epochs,
                        itr,
                        self.dataloader["validation_iterations"],
                    )
                    for loss_n, loss_dict in loss.items():
                        info += " ||"
                        info += "" if loss_n == "none" else " {}".format(loss_n)

                        for k in loss_dict:
                            info += " | {}: {:.4f}".format(k, float(loss_dict[k]))

                    itr_time_avg = (time.time() - epoch_start) / (itr + 1)
                    info += " || data_time: {dt}s | fbward_time: {fbt}s | iter_time: {it}s | epoch_time: {et} | left_time: {lt}".format(
                        dt=round(data_time, 1),
                        fbt=round(time.time() - step_start, 1),
                        it=round(time.time() - itr_start, 1),
                        et=format_seconds(time.time() - epoch_start),
                        lt=format_seconds(
                            itr_time_avg
                            * (self.dataloader["train_iterations"] - itr - 1)
                        ),
                    )

                    self.logger.debug(info)
                itr_start = time.time()

            assert itr >= 0, "The data is too less to form one iteration!"
            self.dataloader["validation_iterations"] = itr + 1

            if self.logger is not None:
                info = "{}: val".format(self.args.name)
                for loss_n, loss_dict in overall_loss.items():
                    info += "" if loss_n == "none" else " {}".format(loss_n)
                    info += ": Epoch {}/{}".format(self.last_epoch, self.max_epochs)
                    for k in loss_dict:
                        info += " | {}: {:.4f}".format(k, float(loss_dict[k]))
                        self.logger.debug(str("val/{}/{}".format(loss_n, k)))
                        self.logger.debug(str(float(loss_dict[k])))
                        self.logger.debug(str(self.last_epoch))
                self.logger.debug(info)

    def validate(self):
        self.validation_epoch()

    def train(self):
        start_epoch = self.last_epoch + 1
        self.start_train_time = time.time()
        self.logger.debug(
            str(
                "{}: global rank {}: start training...".format(
                    self.args.name, self.args.global_rank
                )
            )
        )

        for epoch in range(start_epoch, self.max_epochs):
            self.train_epoch()

    def infer_sample(
        self, x, tempo, not_empty_pos, condition_pos, use_ema=True, skip_step=0
    ):
        self.model.eval()
        tic = time.time()

        if use_ema and self.ema is not None:
            self.ema.modify_to_inference()

        if isinstance(self.model, torch.nn.parallel.DistributedDataParallel):
            model = self.model.module
        else:
            model = self.model

        with torch.no_grad():
            if self.gpu:
                x = x.cuda()
                tempo = tempo.cuda()
                not_empty_pos = not_empty_pos.cuda()
                condition_pos = condition_pos.cuda()
            assert x.size()[0] == 1
            ts = 6

            samples = model.infer_sample(
                x, tempo, not_empty_pos, condition_pos, skip_step=skip_step
            )

            print("sampling, the song has {} time units".format(samples.size()[-1]))

            datum = samples[0]

            encoding = []
            for t in range(samples.size()[-1]):
                bar = t // self.pos_in_bar
                pos = t % self.pos_in_bar
                main_pitch = datum[0][t].item()
                main_dur = datum[1][t].item()
                assert (main_dur <= self.pad_index) and (
                    main_pitch >= self.pad_index
                ), "pitch index is {} and dur index is {}".format(main_pitch, main_dur)
                if (main_pitch != self.pad_index) and (main_dur < self.pad_index):
                    p = self.t_h.ids_to_tokens[main_pitch]
                    if p[0] == "M":
                        p = p[1:]
                        encoding.append(
                            (
                                bar,
                                pos,
                                80,
                                p,
                                self.t_h.ids_to_tokens[main_dur][1:],
                                28,
                                ts,
                                tempo,
                            )
                        )
                    else:
                        print("out m")

                bass_pitch = datum[2][t].item()
                bass_dur = datum[3][t].item()
                assert (
                    bass_dur <= self.pad_index and bass_pitch >= self.pad_index
                ), "{}, {}".format(bass_dur, bass_pitch)
                if (bass_pitch != self.pad_index) and (bass_dur < self.pad_index):
                    pitch = self.t_h.ids_to_tokens[bass_pitch]
                    dur = self.t_h.ids_to_tokens[bass_dur]
                    if pitch[0] == "B":
                        pitch = pitch[1:].split(" ")
                        for p in pitch:
                            encoding.append((bar, pos, 32, p, dur[1:], 24, ts, tempo))
                    else:
                        print("out b")

                drums_pitch = datum[4][t].item()  # 128
                drums_dur = datum[5][t].item()
                assert (
                    drums_dur <= self.pad_index and drums_pitch >= self.pad_index
                ), "{}, {}".format(drums_dur, drums_pitch)
                if (drums_pitch != self.pad_index) and (drums_dur < self.pad_index):
                    pitch = self.t_h.ids_to_tokens[drums_pitch]
                    dur = self.t_h.ids_to_tokens[drums_dur]
                    if pitch[0] == "D":
                        pitch = pitch[1:].split(" ")
                        for p in pitch:
                            encoding.append((bar, pos, 128, p, dur[1:], 24, ts, tempo))
                    else:
                        print("out d")

                guitar_pitch = datum[6][t].item()  # 25
                guitar_dur = datum[7][t].item()
                assert (
                    guitar_dur <= self.pad_index and guitar_pitch >= self.pad_index
                ), "{}, {}".format(guitar_dur, guitar_pitch)
                if (guitar_pitch != self.pad_index) and (guitar_dur < self.pad_index):
                    pitch = self.t_h.ids_to_tokens[guitar_pitch]
                    dur = self.t_h.ids_to_tokens[guitar_dur]
                    if pitch[0] == "G":
                        pitch = pitch[1:].split(" ")
                        for p in pitch:
                            encoding.append((bar, pos, 25, p, dur[1:], 20, ts, tempo))
                    else:
                        print("out g")

                piano_pitch = datum[8][t].item()
                piano_dur = datum[9][t].item()
                assert (
                    piano_dur <= self.pad_index and piano_pitch >= self.pad_index
                ), "{}, {}".format(piano_dur, piano_pitch)
                if (piano_pitch != self.pad_index) and (piano_dur < self.pad_index):
                    pitch = self.t_h.ids_to_tokens[piano_pitch]
                    dur = self.t_h.ids_to_tokens[piano_dur]
                    if pitch[0] == "P":
                        pitch = pitch[1:].split(" ")
                        for p in pitch:
                            encoding.append((bar, pos, 0, p, dur[1:], 24, ts, tempo))
                    else:
                        print("out p")

                string_pitch = datum[10][t].item()  # 48
                string_dur = datum[11][t].item()
                assert (
                    string_dur <= self.pad_index and string_pitch >= self.pad_index
                ), "p:{},d:{}".format(string_pitch, string_dur)
                if (string_pitch != self.pad_index) and (string_dur < self.pad_index):
                    pitch = self.t_h.ids_to_tokens[string_pitch]
                    dur = self.t_h.ids_to_tokens[string_dur]
                    if pitch[0] == "S":
                        pitch = pitch[1:].split(" ")
                        for p in pitch:
                            encoding.append((bar, pos, 48, p, dur[1:], 12, ts, tempo))
                    else:
                        print("out s")

                # just for chord debug
                root_id = datum[12][t].item()
                kind_id = datum[13][t].item()
                if (
                    self.t_h.ids_to_tokens[root_id] in self.r_h.root_dict
                    and self.t_h.ids_to_tokens[kind_id] in self.r_h.kind_dict
                ):
                    root = self.r_h.root_dict[self.t_h.ids_to_tokens[root_id]]
                    kind = self.r_h.kind_dict[self.t_h.ids_to_tokens[kind_id]]
                    encoding.append((bar, pos, 129, root, kind, 1, ts, tempo))

        encoding.sort()
        oct_line = [
            "<0-{}> <1-{}> <2-{}> <3-{}> <4-{}> <5-{}> <6-{}> <7-{}>".format(
                e[0], e[1], e[2], e[3], e[4], e[5], e[6], e[7]
            )
            for e in encoding
        ]

        if use_ema and self.ema is not None:
            self.ema.modify_to_train()

        return " ".join(oct_line)

tuneease/getmusic/engine/solver.py
Pshah2023/TuneEase
from ..utils.misc import instantiate_from_config


def build_model(config, args=None):
    return instantiate_from_config(config["model"])

tuneease/getmusic/modeling/build.py
Pshah2023/TuneEase
import torch
from torch import nn
from ...utils.misc import instantiate_from_config


def disabled_train(self, mode=True):
    return self


class DFM(nn.Module):
    def __init__(self, *, diffusion_config):
        super().__init__()
        self.rfm = instantiate_from_config(diffusion_config)
        self.truncation_forward = False

    def forward(self, batch, name="none", **kwargs):
        output = self.rfm(batch, **kwargs)
        return output

    def parameters(self, recurse=True, name=None):
        if name is None or name == "none":
            return super().parameters(recurse=recurse)
        else:
            names = name.split("+")
            params = []
            for n in names:
                try:
                    params += getattr(self, name).parameters(recurse=recurse, name=name)
                except:
                    params += getattr(self, name).parameters(recurse=recurse)
            return params

    def device(self):
        return self.rfm.device()

    def get_ema_model(self):
        return self.rfm

    @torch.no_grad()
    def infer_sample(self, x, tempo, not_empty_pos, condition_pos, skip_step, **kwargs):
        self.eval()

        trans_out = self.rfm.sample(
            x=x,
            tempo=tempo,
            not_empty_pos=not_empty_pos,
            condition_pos=condition_pos,
            skip_step=skip_step,
            **kwargs
        )

        return trans_out

tuneease/getmusic/modeling/models/dfm.py
Pshah2023/TuneEase
import torch
from torch import nn
import torch.nn.functional as F
from tqdm import tqdm
import numpy as np
from torch.cuda.amp import autocast
from ...utils import midi_config as mc
from ...utils.misc import instantiate_from_config

eps = 1e-8


def sum_except_batch(x, num_dims=1):
    return x.reshape(*x.shape[:num_dims], -1).sum(-1)


def log_1_min_a(a):  # log(1-e_a)
    return torch.log(1 - a.exp() + 1e-40)


def log_add_exp(a, b):  # M + log(e_(a-M)+e_(b-M))
    maximum = torch.max(a, b)  # e(-70) is about 0
    return maximum + torch.log(torch.exp(a - maximum) + torch.exp(b - maximum))


def extract(a, t, x_shape):
    b, *_ = t.shape
    out = a.gather(-1, t)
    return out.reshape(b, *((1,) * (len(x_shape) - 1)))


def log_categorical(log_x_start, log_prob):
    return (log_x_start.exp() * log_prob).sum(dim=1)


def index_to_log_onehot(x, num_classes):
    assert x.max().item() < num_classes, f"Error: {x.max().item()} >= {num_classes}"
    x_onehot = F.one_hot(x, num_classes)
    permute_order = (0, -1) + tuple(range(1, len(x.size())))
    x_onehot = x_onehot.permute(permute_order)
    log_x = torch.log(x_onehot.float().clamp(min=1e-30))
    return log_x


def log_onehot_to_index(log_x):
    return log_x.argmax(1)


def alpha_schedule(
    time_step, N=100, att_1=0.99999, att_T=0.000009, ctt_1=0.000009, ctt_T=0.99999
):
    att = np.arange(0, time_step) / (time_step - 1) * (att_T - att_1) + att_1
    att = np.concatenate(([1], att))
    at = att[1:] / att[:-1]
    ctt = np.arange(0, time_step) / (time_step - 1) * (ctt_T - ctt_1) + ctt_1
    ctt = np.concatenate(([0], ctt))
    one_minus_ctt = 1 - ctt
    one_minus_ct = one_minus_ctt[1:] / one_minus_ctt[:-1]
    ct = 1 - one_minus_ct
    bt = (1 - at - ct) / N
    att = np.concatenate((att[1:], [1]))
    ctt = np.concatenate((ctt[1:], [0]))
    btt = (1 - att - ctt) / N
    return at, bt, ct, att, btt, ctt


class DiffusionRFM(nn.Module):
    def __init__(
        self,
        *,
        roformer_config=None,
        diffusion_step=100,
        alpha_init_type="cos",
        auxiliary_loss_weight=0,
        adaptive_auxiliary_loss=False,
        mask_weight=[1, 1],
    ):
        super().__init__()

        self.roformer = instantiate_from_config(roformer_config)
        self.amp = False
        self.num_classes = (
            self.roformer.vocab_size + 1
        )  # defined in vocabulary, add an additional mask
        self.cond_weight = self.roformer.cond_weight
        self.tracks = 14
        self.pad_index = mc.DURATION_MAX * mc.POS_RESOLUTION - 1
        self.figure_size = (
            mc.BAR_MAX * mc.BEAT_NOTE_FACTOR * mc.MAX_NOTES_PER_BAR * mc.POS_RESOLUTION
        )
        self.num_timesteps = diffusion_step
        self.parametrization = "x0"
        self.auxiliary_loss_weight = auxiliary_loss_weight
        self.adaptive_auxiliary_loss = adaptive_auxiliary_loss

        if alpha_init_type == "alpha1":
            at, bt, ct, att, btt, ctt = alpha_schedule(
                self.num_timesteps, N=self.num_classes - 1
            )
        else:
            print("alpha_init_type is Wrong !! ")

        at = torch.tensor(at.astype("float64"))
        bt = torch.tensor(bt.astype("float64"))
        ct = torch.tensor(ct.astype("float64"))
        log_at = torch.log(at)
        log_bt = torch.log(bt)
        log_ct = torch.log(ct)
        att = torch.tensor(att.astype("float64"))
        btt = torch.tensor(btt.astype("float64"))
        ctt = torch.tensor(ctt.astype("float64"))
        log_cumprod_at = torch.log(att)
        log_cumprod_bt = torch.log(btt)
        log_cumprod_ct = torch.log(ctt)

        log_1_min_ct = log_1_min_a(log_ct)  # log(1-e_a), log(1-ct)
        log_1_min_cumprod_ct = log_1_min_a(log_cumprod_ct)  # log(1-ctt)
        # M + log(e_(a-M)+e_(b-M))
        assert log_add_exp(log_ct, log_1_min_ct).abs().sum().item() < 1.0e-5
        assert (
            log_add_exp(log_cumprod_ct, log_1_min_cumprod_ct).abs().sum().item()
            < 1.0e-5
        )

        # Convert to float32 and register buffers.
        self.register_buffer("log_at", log_at.float())
        self.register_buffer("log_bt", log_bt.float())
        self.register_buffer("log_ct", log_ct.float())
        self.register_buffer("log_cumprod_at", log_cumprod_at.float())
        self.register_buffer("log_cumprod_bt", log_cumprod_bt.float())
        self.register_buffer("log_cumprod_ct", log_cumprod_ct.float())
        self.register_buffer("log_1_min_ct", log_1_min_ct.float())
        self.register_buffer("log_1_min_cumprod_ct", log_1_min_cumprod_ct.float())

        self.register_buffer("Lt_history", torch.zeros(self.num_timesteps))
        self.register_buffer("Lt_count", torch.zeros(self.num_timesteps))

        self.prior_ps = 1024  # max number to sample per step

    def multinomial_kl(self, log_prob1, log_prob2):  # compute KL loss on log_prob
        kl = (log_prob1.exp() * (log_prob1 - log_prob2)).sum(dim=1)
        return kl

    def q_pred_one_timestep(self, log_x_t, t):  # q(xt|xt_1)
        log_at = extract(self.log_at, t, log_x_t.shape)  # at
        log_bt = extract(self.log_bt, t, log_x_t.shape)  # bt
        log_ct = extract(self.log_ct, t, log_x_t.shape)  # ct
        log_1_min_ct = extract(self.log_1_min_ct, t, log_x_t.shape)  # 1-ct

        log_probs = torch.cat(
            [
                log_add_exp(log_x_t[:, :-1, :] + log_at, log_bt),
                log_add_exp(log_x_t[:, -1:, :] + log_1_min_ct, log_ct),
            ],
            dim=1,
        )

        return log_probs

    def q_pred(self, log_x_start, t):  # q(xt|x0)
        # log_x_start can be onehot or not
        t = (t + (self.num_timesteps + 1)) % (self.num_timesteps + 1)
        log_cumprod_at = extract(self.log_cumprod_at, t, log_x_start.shape)  # at~
        log_cumprod_bt = extract(self.log_cumprod_bt, t, log_x_start.shape)  # bt~
        log_cumprod_ct = extract(self.log_cumprod_ct, t, log_x_start.shape)  # ct~
        log_1_min_cumprod_ct = extract(
            self.log_1_min_cumprod_ct, t, log_x_start.shape
        )  # 1-ct~

        log_probs = torch.cat(
            [
                log_add_exp(log_x_start[:, :-1, :] + log_cumprod_at, log_cumprod_bt),
                log_add_exp(
                    log_x_start[:, -1:, :] + log_1_min_cumprod_ct, log_cumprod_ct
                ),
            ],
            dim=1,
        )

        return log_probs

    def log_sample_categorical(
        self, logits
    ):  # use gumbel to sample onehot vector from log probability
        uniform = torch.rand_like(logits)
        gumbel_noise = -torch.log(-torch.log(uniform + 1e-30) + 1e-30)
        sample = (gumbel_noise + logits).argmax(dim=1)
        log_sample = index_to_log_onehot(sample, self.num_classes)
        return log_sample

    def log_sample_categorical_infer(
        self, logits, figure_size
    ):  # use gumbel to sample onehot vector from log probability
        uniform = torch.rand_like(logits)
        gumbel_noise = -torch.log(-torch.log(uniform + 1e-30) + 1e-30)
        sample = gumbel_noise + logits
        for i in range(self.tracks - 2):  # do not decode chord
            track = sample[:, :, i * figure_size : (i + 1) * figure_size]
            if i % 2 == 1:  # duration
                track[
                    :, self.pad_index + 1 : -1, :
                ] = -70  # only decode duration tokens
            else:  # only decode pitch tokens in $i$-th track
                start = mc.TRACKS_START[i // 2]
                end = mc.TRACKS_END[i // 2]
                track[:, : self.pad_index, :] = -70
                track[:, self.pad_index + 1 : start, :] = -70
                track[:, end + 1 : -1, :] = -70
            sample[:, :, i * figure_size : (i + 1) * figure_size] = track
        sample = sample.argmax(dim=1)
        log_sample = index_to_log_onehot(sample, self.num_classes)
        return log_sample

    def q_sample(self, log_x_start, t):  # diffusion step, q(xt|x0) and sample xt
        log_EV_qxt_x0 = self.q_pred(log_x_start, t)  # log q(xt|x0)
        log_sample = self.log_sample_categorical(log_EV_qxt_x0)
        return log_sample

    def predict_start(self, log_x_t, t, condition_pos):  # p(x0|xt)
        x_t = log_onehot_to_index(log_x_t)
        if self.amp == True:
            with autocast():
                out = self.roformer(x_t, t, condition_pos)
        else:
            out = self.roformer(x_t, t, condition_pos)

        log_pred = F.log_softmax(out.double(), dim=2).float()
        batch_size = log_x_t.size()[0]
        zero_vector = torch.zeros(batch_size, out.size()[1], 2).type_as(log_x_t) - 70

        log_pred = torch.cat((log_pred, zero_vector), dim=2)
        log_pred = torch.clamp(log_pred, -70, 0)
        return log_pred.transpose(2, 1)

    def q_posterior(self, log_x_start, log_x_t, t):
        x_len = log_x_t.size()[-1]
        assert t.min().item() >= 0 and t.max().item() < self.num_timesteps
        batch_size = log_x_start.size()[0]
        onehot_x_t = log_onehot_to_index(log_x_t)  # get sample
        mask = (onehot_x_t == self.num_classes - 1).unsqueeze(1)  # select masked tokens
        log_one_vector = torch.zeros(batch_size, 1, 1).type_as(log_x_t)  # b,1,1
        log_zero_vector = torch.log(log_one_vector + 1.0e-30).expand(
            -1, -1, x_len
        )  # [B, 1, content_seq_len]

        log_qt = self.q_pred(log_x_t, t)  # q(xt|x0)
        log_qt = log_qt[:, :-1, :]
        log_cumprod_ct = extract(self.log_cumprod_ct, t, log_x_start.shape)  # ct~
        ct_cumprod_vector = log_cumprod_ct.expand(-1, self.num_classes - 1, -1)
        log_qt = (~mask) * log_qt + mask * ct_cumprod_vector

        # log(q(xt|xt_1,x0))
        log_qt_one_timestep = self.q_pred_one_timestep(log_x_t, t)  # q(xt|xt_1)
        log_qt_one_timestep = torch.cat(
            (log_qt_one_timestep[:, :-1, :], log_zero_vector), dim=1
        )
        log_ct = extract(self.log_ct, t, log_x_start.shape)  # ct
        ct_vector = log_ct.expand(-1, self.num_classes - 1, -1)
        ct_vector = torch.cat((ct_vector, log_one_vector), dim=1)
        log_qt_one_timestep = (~mask) * log_qt_one_timestep + mask * ct_vector

        q = log_x_start[:, :-1, :] - log_qt
        q = torch.cat((q, log_zero_vector), dim=1)
        q_log_sum_exp = torch.logsumexp(q, dim=1, keepdim=True)
        q = q - q_log_sum_exp

        log_EV_xtmin_given_xt_given_xstart = (
            self.q_pred(q, t - 1) + log_qt_one_timestep + q_log_sum_exp
        )
        return torch.clamp(log_EV_xtmin_given_xt_given_xstart, -70, 0)

    def p_pred(self, log_x, t, condition_pos):
        if self.parametrization == "x0":
            log_x_recon = self.predict_start(log_x, t, condition_pos)
            log_model_pred = self.q_posterior(
                log_x_start=log_x_recon, log_x_t=log_x, t=t
            )
        elif self.parametrization == "direct":
            log_model_pred = self.predict_start(log_x, t, condition_pos)
        else:
            raise ValueError
        return log_model_pred, log_x_recon

    @torch.no_grad()
    def p_sample(
        self,
        log_x,
        t,
        figure_size,
        condition_pos,
        not_empty_pos,
        sampled=None,
        to_sample=None,
    ):  # sample q(xt-1) for next step from  xt, actually is p(xt-1|xt)
        model_log_prob, log_x_recon = self.p_pred(log_x, t, condition_pos)

        max_sample_per_step = self.prior_ps  # max number to sample per step
        if t[0] > 0 and to_sample is not None:
            log_x_idx = log_onehot_to_index(log_x)

            # dim=1，vocaburay dimension
            score = torch.exp(log_x_recon).max(dim=1).values.softmax(dim=1)

            out = self.log_sample_categorical_infer(log_x_recon, figure_size)
            out_idx = log_onehot_to_index(out)

            out2_idx = log_x_idx.clone()
            _score = score.clone()
            if _score.sum() < 1e-6:
                _score += 1

            # only content has score
            _score = torch.where(
                ((1 - condition_pos) * not_empty_pos).type(torch.bool), _score, 0
            )
            # only mask has score
            _score[log_x_idx != self.num_classes - 1] = 0

            for j in range(6):  # do not decode chord
                __score = _score[0][2 * j * figure_size : (2 * j + 2) * figure_size]

                if __score.sum() == 0:
                    continue

                n_sample = min(to_sample[j] - sampled[j], max_sample_per_step)

                if to_sample[j] - sampled[j] - n_sample == 1:
                    n_sample = to_sample[j] - sampled[j]
                if n_sample <= 0:
                    continue

                sel = torch.multinomial(__score, n_sample.item())

                out2_idx[0][2 * j * figure_size : (2 * j + 2) * figure_size][
                    sel
                ] = out_idx[0][2 * j * figure_size : (2 * j + 2) * figure_size][sel]

                sampled[j] += (
                    (
                        out2_idx[0][2 * j * figure_size : (2 * j + 2) * figure_size]
                        != self.num_classes - 1
                    ).sum()
                    - (
                        log_x_idx[0][2 * j * figure_size : (2 * j + 2) * figure_size]
                        != self.num_classes - 1
                    ).sum()
                ).item()

            out = index_to_log_onehot(out2_idx, self.num_classes)
        else:
            # Gumbel sample
            out = self.log_sample_categorical_infer(model_log_prob, figure_size)
            sampled = None

        if to_sample is not None:
            return out, sampled
        else:
            return out

    def sample_time(self, b, device, method="uniform"):
        if method == "importance":
            if not (self.Lt_count > 10).all():
                return self.sample_time(b, device, method="uniform")

            Lt_sqrt = torch.sqrt(self.Lt_history + 1e-10) + 0.0001
            Lt_sqrt[0] = Lt_sqrt[1]
            pt_all = Lt_sqrt / Lt_sqrt.sum()

            t = torch.multinomial(pt_all, num_samples=b, replacement=True)

            pt = pt_all.gather(dim=0, index=t)
            return t, pt

        elif method == "uniform":
            t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
            pt = torch.ones_like(t).float() / self.num_timesteps
            return t, pt
        else:
            raise ValueError

    def _train_loss(self, x, condition_pos, not_empty_pos, is_train=True):
        assert x.size(2) == self.figure_size

        b, device = x.size(0), x.device

        condition_pos = condition_pos.view(b, -1)
        not_empty_pos = not_empty_pos.view(b, -1)

        x_start = x
        t, pt = self.sample_time(b, device, "importance")

        log_x_start = index_to_log_onehot(x_start.view(b, -1).long(), self.num_classes)

        log_xt = self.q_sample(log_x_start=log_x_start, t=t)

        log_empty = -70 * torch.ones_like(log_x_start)
        log_empty[:, -2, :] = 0

        # return condition to ground truth
        log_xt = torch.where(
            condition_pos.unsqueeze(1).type(torch.bool), log_x_start, log_xt
        )
        # empty
        log_xt = torch.where(
            (1 - not_empty_pos).unsqueeze(1).type(torch.bool), log_empty, log_xt
        )

        log_x0_recon = self.predict_start(log_xt, t=t, condition_pos=condition_pos)
        x0_recon = log_onehot_to_index(log_x0_recon)
        x0_real = x_start.view(b, -1).long()

        log_model_prob = self.q_posterior(
            log_x_start=log_x0_recon, log_x_t=log_xt, t=t
        )  # go through q(xt_1|xt,x0)
        xt_1_recon = log_onehot_to_index(log_model_prob)
        xt_recon = log_onehot_to_index(log_xt)

        position_weight = not_empty_pos * (
            (1 - condition_pos) + self.cond_weight * condition_pos
        )  # loss weight: condition/content = cond_weight / 1
        log_true_prob = self.q_posterior(log_x_start=log_x_start, log_x_t=log_xt, t=t)

        L_t_1 = self.multinomial_kl(log_true_prob, log_model_prob)
        L_t_1 = L_t_1 * position_weight
        L_t_1 = sum_except_batch(L_t_1)

        L_0 = -log_categorical(log_x_start, log_model_prob)
        L_0 *= position_weight
        L_0 = sum_except_batch(L_0)

        mask = (t == torch.zeros_like(t)).float()
        kl_loss = mask * L_0 + (1.0 - mask) * L_t_1

        Lt2 = kl_loss.pow(2)
        Lt2_prev = self.Lt_history.gather(dim=0, index=t)
        new_Lt_history = (0.1 * Lt2 + 0.9 * Lt2_prev).detach()
        self.Lt_history.scatter_(dim=0, index=t, src=new_Lt_history)
        self.Lt_count.scatter_add_(dim=0, index=t, src=torch.ones_like(Lt2))

        loss1 = kl_loss / pt
        vb_loss = loss1

        if self.auxiliary_loss_weight != 0 and is_train == True:
            kl_aux = self.multinomial_kl(
                log_x_start[:, :-2, :], log_x0_recon[:, :-2, :]
            )
            kl_aux = kl_aux * position_weight
            kl_aux = sum_except_batch(kl_aux)
            kl_aux_loss = mask * L_0 + (1.0 - mask) * kl_aux

            if self.adaptive_auxiliary_loss == True:
                addition_loss_weight = (1 - t / self.num_timesteps) + 1.0
            else:
                addition_loss_weight = 1.0

            loss2 = addition_loss_weight * self.auxiliary_loss_weight * kl_aux_loss / pt

            vb_loss += loss2

        scale = (position_weight != 0.0).sum(-1)

        scale = torch.where(scale == 0.0, 0.0, (1 / scale).double())

        return log_model_prob, vb_loss * scale, L_0 * scale

    def device(self):
        return self.roformer.input_transformers.encoder.embed_positions.weight.device

    def parameters(self, recurse=True, name=None):
        """
        Following minGPT:
        This long function is unfortunately doing something very simple and is being very defensive:
        We are separating out all parameters of the model into two buckets: those that will experience
        weight decay for regularization and those that won't (biases, and layernorm/embedding weights).
        We are then returning the PyTorch optimizer object.
        """
        # return super().parameters(recurse=True)
        if name is None or name == "none":
            return super().parameters(recurse=recurse)
        else:
            # separate out all parameters to those that will and won't experience regularizing weight decay
            print("GPTLikeunet: get parameters by the overwrite method!")
            decay = set()
            no_decay = set()
            whitelist_weight_modules = (torch.nn.Linear,)
            blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
            for mn, m in self.named_modules():
                for pn, p in m.named_parameters():
                    fpn = "%s.%s" % (mn, pn) if mn else pn  # full param name

                    if pn.endswith("bias"):
                        # all biases will not be decayed
                        no_decay.add(fpn)
                    elif pn.endswith("weight") and isinstance(
                        m, whitelist_weight_modules
                    ):
                        # weights of whitelist modules will be weight decayed
                        decay.add(fpn)
                    elif pn.endswith("weight") and isinstance(
                        m, blacklist_weight_modules
                    ):
                        # weights of blacklist modules will NOT be weight decayed
                        no_decay.add(fpn)
            # special case the position embedding parameter as not decayed
            module_name = ["condition_emb", "content_emb"]
            pos_emb_name = [
                "pos_emb",
                "width_emb",
                "height_emb",
                "pad_emb",
                "token_type_emb",
            ]
            for mn in module_name:
                if hasattr(self, mn) and getattr(self, mn) is not None:
                    for pn in pos_emb_name:
                        if hasattr(getattr(self, mn), pn):
                            if isinstance(
                                getattr(getattr(self, mn), pn), torch.nn.Parameter
                            ):
                                no_decay.add("{}.{}".format(mn, pn))

            # validate that we considered every parameter
            param_dict = {
                pn: p for pn, p in self.roformer.named_parameters()
            }  # if p.requires_grad}
            inter_params = decay & no_decay
            union_params = decay | no_decay
            assert (
                len(inter_params) == 0
            ), "parameters %s made it into both decay/no_decay sets!" % (
                str(inter_params),
            )
            assert (
                len(param_dict.keys() - union_params) == 0
            ), "parameters %s were not separated into either decay/no_decay set!" % (
                str(param_dict.keys() - union_params),
            )

            # create the pytorch optimizer object
            optim_groups = [
                {
                    "params": [param_dict[pn] for pn in sorted(list(decay))],
                    "weight_decay": 0.01,
                },
                {
                    "params": [param_dict[pn] for pn in sorted(list(no_decay))],
                    "weight_decay": 0.0,
                },
            ]
            return optim_groups

    def forward(
        self,
        x,
        tempo,
        condition_pos,
        not_empty_pos,
        return_loss=False,
        return_logits=True,
        return_att_weight=False,
        is_train=True,
        **kwargs,
    ):
        if kwargs.get("autocast") == True:
            self.amp = True

        if is_train == True:
            log_model_prob, loss, decoder_nll = self._train_loss(
                x, condition_pos, not_empty_pos
            )
            loss = loss.sum()
            decoder_nll = decoder_nll.sum()

        out = {}
        if return_logits:
            out["logits"] = torch.exp(log_model_prob)

        if return_loss:
            out["loss"] = loss
            out["decoder_nll"] = decoder_nll

        self.amp = False
        return out

    def sample(self, x, tempo, not_empty_pos, condition_pos, skip_step=0, **kwargs):
        batch_size = x.size()[0]
        assert batch_size == 1

        num_to_be_generated_per_track = (x == self.num_classes - 1).sum(-1)[0]
        num_to_be_generated_per_track = (
            num_to_be_generated_per_track[0::2] + num_to_be_generated_per_track[1::2]
        )

        x = x.view(batch_size, -1).long()

        condition_pos = condition_pos.view(batch_size, -1)
        not_empty_pos = not_empty_pos.view(batch_size, -1)

        sample_len = x.size()[1]
        assert sample_len % self.tracks == 0
        figure_size = sample_len // self.tracks

        device = self.log_at.device
        start_step = self.num_timesteps

        log_x_start = index_to_log_onehot(x, self.num_classes)

        log_pad = torch.ones_like(log_x_start) * -70
        log_pad[:, self.pad_index, :] = 0
        log_empty = -70 * torch.ones_like(log_x_start)
        log_empty[:, -2, :] = 0

        log_p_x_on_y = -70 * torch.ones(
            (batch_size, self.num_classes, sample_len), device=device
        )
        log_p_x_on_y[:, -1, :] = 0

        start_step = self.num_timesteps

        to_be_sampled_per_step = torch.floor(
            num_to_be_generated_per_track / (self.num_timesteps - 1)
        )

        self.n_sample = to_be_sampled_per_step.unsqueeze(0).repeat(
            self.num_timesteps - 1, 1
        )

        compensate = num_to_be_generated_per_track - self.n_sample.sum(0)

        for i in range(6):
            while compensate[i] > 0:
                compensate_freq = (
                    torch.ceil(self.num_timesteps / compensate[i]).long().item()
                )
                if compensate_freq > 1:
                    self.n_sample[0::compensate_freq, i] += 1
                compensate[i] = (
                    num_to_be_generated_per_track[i] - self.n_sample.sum(0)[i]
                )

        self.n_sample = torch.cat(
            [torch.ones(1, 7, device=self.n_sample.device), self.n_sample], dim=0
        )
        self.n_sample = torch.where(self.n_sample == 0, -1, self.n_sample).long()

        with torch.no_grad():
            for diffusion_index in tqdm(range(start_step - 1, -1, -1)):
                t = torch.full(
                    (batch_size,), diffusion_index, device=device, dtype=torch.long
                )
                sampled = torch.zeros(7, dtype=torch.long, device=self.n_sample.device)

                while (sampled < self.n_sample[diffusion_index]).any():
                    log_p_x_on_y = torch.where(
                        condition_pos.unsqueeze(1).type(torch.bool),
                        log_x_start,
                        log_p_x_on_y,
                    )
                    log_p_x_on_y = torch.where(
                        (1 - not_empty_pos).unsqueeze(1).type(torch.bool),
                        log_empty,
                        log_p_x_on_y,
                    )
                    log_p_x_on_y, sampled = self.p_sample(
                        log_p_x_on_y,
                        t,
                        figure_size,
                        condition_pos,
                        not_empty_pos,
                        sampled,
                        self.n_sample[diffusion_index],
                    )  # log_z is log_onehot
                    if sampled is None:
                        assert t[0] == 0
                        break

            log_p_x_on_y = torch.where(
                condition_pos.unsqueeze(1).type(torch.bool), log_x_start, log_p_x_on_y
            )
            log_p_x_on_y = torch.where(
                (1 - not_empty_pos).unsqueeze(1).type(torch.bool), log_pad, log_p_x_on_y
            )
            content_token = log_onehot_to_index(log_p_x_on_y)

        return content_token.view(batch_size, self.tracks, -1)

tuneease/getmusic/modeling/roformer/diffusion_roformer.py
Pshah2023/TuneEase
import math
from typing import Optional, Tuple, Union
import torch
import torch.utils.checkpoint
from torch import nn
from transformers import RoFormerPreTrainedModel, PretrainedConfig
from transformers import RoFormerModel
from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions
from einops import rearrange


class RoFormerConfig(PretrainedConfig):
    model_type = "roformer"

    def __init__(
        self,
        vocab_size=None,
        embedding_size=96,
        tracks=14,
        hidden_size=768,
        num_hidden_layers=12,
        num_attention_heads=12,
        intermediate_size=3072,
        hidden_act="gelu",
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        max_position_embeddings=1536,
        type_vocab_size=2,
        initializer_range=0.02,
        layer_norm_eps=1e-12,
        pad_token_id=None,
        rotary_value=False,
        use_cache=True,
        **kwargs,
    ):
        super().__init__(pad_token_id=pad_token_id, **kwargs)

        self.vocab_size = vocab_size
        self.tracks = tracks
        self.embedding_size = hidden_size if embedding_size is None else embedding_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.hidden_act = hidden_act
        self.intermediate_size = intermediate_size
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.max_position_embeddings = max_position_embeddings
        self.type_vocab_size = type_vocab_size
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps
        self.rotary_value = rotary_value
        self.use_cache = use_cache


def timestep_embedding(timesteps, dim, max_period=10000):
    """
    Create sinusoidal timestep embeddings.
    :param timesteps: a 1-D Tensor of N indices, one per batch element.
                    These may be fractional.
    :param dim: the dimension of the output.
    :param max_period: controls the minimum frequency of the embeddings.
    :return: an [N x dim] Tensor of positional embeddings.
    """
    half = dim // 2
    freqs = torch.exp(
        -math.log(max_period)
        * torch.arange(start=0, end=half, dtype=torch.float32)
        / half
    ).to(device=timesteps.device)
    args = timesteps[:, None].float() * freqs[None]
    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
    if dim % 2:
        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
    return embedding


class SiLU(nn.Module):
    def forward(self, x):
        return x * torch.sigmoid(x)


class DiffusionRoFormerEmbeddings(nn.Module):
    """Construct the embeddings from word and token_type embeddings."""

    def __init__(self, config):
        super().__init__()

        self.word_embeddings = nn.Embedding(
            config.vocab_size, config.embedding_size, padding_idx=config.pad_token_id
        )
        self.token_type_embeddings = nn.Embedding(
            config.type_vocab_size, config.embedding_size
        )
        self.time_embeddings = nn.Sequential(
            nn.Linear(128, config.embedding_size),
            SiLU(),
            nn.Linear(config.embedding_size, config.embedding_size),
        )

        self.LayerNorm = nn.LayerNorm(
            config.embedding_size * config.tracks, eps=config.layer_norm_eps
        )
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, input_ids=None, token_type_ids=None, timesteps=None):
        if input_ids is not None:
            input_shape = input_ids.size()
        else:
            input_shape = inputs_embeds.size()[:-1]

        inputs_embeds = self.word_embeddings(input_ids)  # b, t, l, d

        if token_type_ids is None:
            token_type_ids = torch.zeros(
                input_shape, dtype=torch.long, device=inputs_embeds.device
            )

        time_embeddings = (
            self.time_embeddings(timestep_embedding(timesteps, 128))
            .unsqueeze(1)
            .unsqueeze(1)
        )  # b, 1, 1, d
        token_type_embeddings = self.token_type_embeddings(
            token_type_ids.long()
        )  # b, t, l, d

        embeddings = inputs_embeds + token_type_embeddings + time_embeddings
        embeddings = rearrange(embeddings, "b t l d -> b l (t d)")
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)

        return embeddings


class DiffusionRoFormerModel(RoFormerPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.embeddings = DiffusionRoFormerEmbeddings(config)

        self.embeddings_project = nn.Sequential(
            nn.Linear(config.embedding_size * config.tracks, config.hidden_size),
            SiLU(),
            nn.Linear(config.hidden_size, config.hidden_size),
        )

        self.outputs_project = nn.Sequential(
            nn.Linear(config.hidden_size, config.embedding_size * config.tracks),
            SiLU(),
            nn.Linear(
                config.embedding_size * config.tracks,
                config.embedding_size * config.tracks,
            ),
        )

        self.encoder = RoFormerModel(config).encoder

        self.lm_head = nn.Linear(config.embedding_size, config.vocab_size)
        with torch.no_grad():
            self.lm_head.weight = self.embeddings.word_embeddings.weight

    def get_extended_attention_mask(self, attention_mask, input_shape, device):
        if attention_mask.dim() == 3:
            extended_attention_mask = attention_mask[:, None, :, :]
        elif attention_mask.dim() == 2:
            if self.config.is_decoder:
                batch_size, seq_length = input_shape
                seq_ids = torch.arange(seq_length, device=device)
                causal_mask = (
                    seq_ids[None, None, :].repeat(batch_size, seq_length, 1)
                    <= seq_ids[None, :, None]
                )
                causal_mask = causal_mask.to(attention_mask.dtype)

                if causal_mask.shape[1] < attention_mask.shape[1]:
                    prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]
                    causal_mask = torch.cat(
                        [
                            torch.ones(
                                (batch_size, seq_length, prefix_seq_len),
                                device=device,
                                dtype=causal_mask.dtype,
                            ),
                            causal_mask,
                        ],
                        axis=-1,
                    )

                extended_attention_mask = (
                    causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
                )
            else:
                extended_attention_mask = attention_mask[:, None, None, :]
        else:
            raise ValueError(
                "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
                    input_shape, attention_mask.shape
                )
            )

        extended_attention_mask = extended_attention_mask.to(
            dtype=self.dtype
        )  # fp16 compatibility
        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0

        return extended_attention_mask

    def forward(
        self,
        input_ids,
        timesteps,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[BaseModelOutputWithPastAndCrossAttentions, Tuple[torch.Tensor]]:
        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        use_cache = False

        if input_ids is not None:
            input_shape = input_ids.size()
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        batch_size, track_num, seq_length = input_shape
        # inputs_embeds = self.word_embeddings(input_ids) # b, t, l, d
        device = input_ids.device if input_ids is not None else inputs_embeds.device

        past_key_values_length = (
            past_key_values[0][0].shape[2] if past_key_values is not None else 0
        )

        if attention_mask is None:
            attention_mask = torch.ones(
                ((batch_size, seq_length + past_key_values_length)), device=device
            )
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(
            attention_mask, input_shape, device=input_ids.device
        )

        encoder_extended_attention_mask = None

        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        embedding_output = self.embeddings(
            input_ids=input_ids, token_type_ids=token_type_ids, timesteps=timesteps
        )

        embedding_output = self.embeddings_project(embedding_output)

        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=extended_attention_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]

        sequence_output = self.outputs_project(sequence_output)
        sequence_output = rearrange(
            sequence_output, "b l (t d) -> b (t l) d", d=self.config.embedding_size
        )

        return self.lm_head(sequence_output)[:, :, :-2]

tuneease/getmusic/modeling/roformer/roformer.py
Pshah2023/TuneEase
import torch.nn as nn
import torch
from .roformer import RoFormerConfig, DiffusionRoFormerModel


class DiffusionRoformerModel(nn.Module):
    def __init__(
        self,
        vocab_size=None,
        cond_weight=None,
    ):
        super().__init__()

        self.vocab_size = vocab_size

        config = RoFormerConfig(vocab_size=vocab_size + 1, pad_token_id=vocab_size - 1)

        self.cond_weight = cond_weight

        self.input_transformers = DiffusionRoFormerModel(config)

    def forward(self, x, timesteps, condition_pos):
        b = x.size()[0]
        x_seq_len = x.size()[1]
        figure_size = x_seq_len // 14
        x = x.reshape(b, 14, figure_size)
        condition_pos = condition_pos.reshape(b, 14, figure_size)

        attention_mask = None
        # extrapolation
        if not self.training and figure_size > 512:
            e = torch.ones(b, figure_size, figure_size, device=x.device)
            attention_mask = torch.tril(e, 255) * torch.triu(e, -255)

        outputs = self.input_transformers(
            input_ids=x,
            token_type_ids=condition_pos,
            timesteps=timesteps,
            attention_mask=attention_mask,
        )

        return outputs

tuneease/getmusic/modeling/roformer/roformer_utils.py
Pshah2023/TuneEase
import sys
import yaml
import torch
import json


def save_config_to_yaml(config, path):
    assert path.endswith(".yaml")
    with open(path, "w") as f:
        f.write(yaml.dump(config))
        f.close()


def save_dict_to_json(d, path, indent=None):
    json.dump(d, open(path, "w"), indent=indent)


def load_dict_from_json(path):
    return json.load(open(path, "r"))


def write_args(args, path):
    args_dict = dict(
        (name, getattr(args, name)) for name in dir(args) if not name.startswith("_")
    )
    with open(path, "a") as args_file:
        args_file.write("==> torch version: {}\n".format(torch.__version__))
        args_file.write(
            "==> cudnn version: {}\n".format(torch.backends.cudnn.version())
        )
        args_file.write("==> Cmd:\n")
        args_file.write(str(sys.argv))
        args_file.write("\n==> args:\n")
        for k, v in sorted(args_dict.items()):
            args_file.write("  %s: %s\n" % (str(k), str(v)))
        args_file.close()

tuneease/getmusic/utils/io.py
Pshah2023/TuneEase
# Copyright 2021 The Magenta Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#!PS Probably modified by the authors of getmusic and now me for making the class.
#!PS It was previously not usable because of time of import issues.
#!PS I attribute everything to the original authors, whoever they are, for the beautiful algorithms.

"""Chord inference for NoteSequences."""
import bisect
import itertools
import numpy as np


# Names of pitch classes to use (mostly ignoring spelling).
class Magenta:
    _PITCH_CLASS_NAMES = None
    _KEY_PITCHES = None
    _CHORD_KIND_PITCHES = None
    _CHORD_KINDS = None
    NO_CHORD = None
    _CHORDS = None
    _KEY_CHORDS = None
    _MAX_NUM_CHORDS = None
    UNPITCHED_PROGRAMS = None
    _DEFAULT_TIME_SIGNATURE_CHORDS_PER_BAR = None

    def __init__(self) -> None:
        self._PITCH_CLASS_NAMES = [
            "C",
            "C#",
            "D",
            "Eb",
            "E",
            "F",
            "F#",
            "G",
            "Ab",
            "A",
            "Bb",
            "B",
        ]

        # Pitch classes in a key (rooted at zero).
        self._KEY_PITCHES = [0, 2, 4, 5, 7, 9, 11]

        # Pitch classes in each chord kind (rooted at zero).
        self._CHORD_KIND_PITCHES = {
            "": [0, 4, 7],
            "m": [0, 3, 7],
            "+": [0, 4, 8],
            "dim": [0, 3, 6],
            "7": [0, 4, 7, 10],
            "maj7": [0, 4, 7, 11],
            "m7": [0, 3, 7, 10],
            "m7b5": [0, 3, 6, 10],
        }

        self._CHORD_KINDS = self._CHORD_KIND_PITCHES.keys()
        self.NO_CHORD = "N.C."
        # All usable chords, including no-chord.
        self._CHORDS = [self.NO_CHORD] + list(
            itertools.product(range(12), self._CHORD_KINDS)
        )

        # All key-chord pairs.
        self._KEY_CHORDS = list(itertools.product(range(12), self._CHORDS))

        # Maximum length of chord sequence to infer.
        self._MAX_NUM_CHORDS = 1000

        # MIDI programs that typically sound unpitched.
        self.UNPITCHED_PROGRAMS = (
            list(range(96, 104)) + list(range(112, 120)) + list(range(120, 128))
        )

        # Mapping from time signature to number of chords to infer per bar.
        self._DEFAULT_TIME_SIGNATURE_CHORDS_PER_BAR = {
            (2, 2): 1,
            (2, 4): 1,
            (3, 4): 1,
            (4, 4): 2,
            (6, 8): 2,
        }

    def _key_chord_distribution(self, chord_pitch_out_of_key_prob):
        """Probability distribution over chords for each key."""
        num_pitches_in_key = np.zeros([12, len(self._CHORDS)], dtype=np.int32)
        num_pitches_out_of_key = np.zeros([12, len(self._CHORDS)], dtype=np.int32)

        # For each key and chord, compute the number of chord notes in the key and the
        # number of chord notes outside the key.
        for key in range(12):
            key_pitches = set((key + offset) % 12 for offset in self._KEY_PITCHES)
            for i, chord in enumerate(self._CHORDS[1:]):
                root, kind = chord
                chord_pitches = set(
                    (root + offset) % 12 for offset in self._CHORD_KIND_PITCHES[kind]
                )
                num_pitches_in_key[key, i + 1] = len(chord_pitches & key_pitches)
                num_pitches_out_of_key[key, i + 1] = len(chord_pitches - key_pitches)

        # Compute the probability of each chord under each key, normalizing to sum to
        # one for each key.
        mat = (
            (1 - chord_pitch_out_of_key_prob) ** num_pitches_in_key
            * chord_pitch_out_of_key_prob**num_pitches_out_of_key
        )
        mat /= mat.sum(axis=1)[:, np.newaxis]
        return mat

    def _key_chord_transition_distribution(
        self, key_chord_distribution, key_change_prob, chord_change_prob
    ):
        """Transition distribution between key-chord pairs."""
        mat = np.zeros([len(self._KEY_CHORDS), len(self._KEY_CHORDS)])

        for i, key_chord_1 in enumerate(self._KEY_CHORDS):
            key_1, chord_1 = key_chord_1
            chord_index_1 = i % len(self._CHORDS)

            for j, key_chord_2 in enumerate(self._KEY_CHORDS):
                key_2, chord_2 = key_chord_2
                chord_index_2 = j % len(self._CHORDS)

                if key_1 != key_2:
                    # Key change. Chord probability depends only on key and not previous
                    # chord.
                    mat[i, j] = key_change_prob / 11
                    mat[i, j] *= key_chord_distribution[key_2, chord_index_2]

                else:
                    # No key change.
                    mat[i, j] = 1 - key_change_prob
                    if chord_1 != chord_2:
                        # Chord probability depends on key, but we have to redistribute the
                        # probability mass on the previous chord since we know the chord
                        # changed.
                        mat[i, j] *= chord_change_prob * (
                            key_chord_distribution[key_2, chord_index_2]
                            + key_chord_distribution[key_2, chord_index_1]
                            / (len(self._CHORDS) - 1)
                        )
                    else:
                        # No chord change.
                        mat[i, j] *= 1 - chord_change_prob

        return mat

    def _chord_pitch_vectors(self):
        """Unit vectors over pitch classes for all chords."""
        x = np.zeros([len(self._CHORDS), 12])
        for i, chord in enumerate(self._CHORDS[1:]):
            root, kind = chord
            for offset in self._CHORD_KIND_PITCHES[kind]:
                x[i + 1, (root + offset) % 12] = 1
        x[1:, :] /= np.linalg.norm(x[1:, :], axis=1)[:, np.newaxis]
        return x

    def sequence_note_pitch_vectors(self, sequence, seconds_per_frame):
        """Compute pitch class vectors for temporal frames across a sequence.
        Args:
        sequence: The NoteSequence for which to compute pitch class vectors.
        seconds_per_frame: The size of the frame corresponding to each pitch class
            vector, in seconds. Alternatively, a list of frame boundary times in
            seconds (not including initial start time and final end time).
        Returns:
        A numpy array with shape `[num_frames, 12]` where each row is a unit-
        normalized pitch class vector for the corresponding frame in `sequence`.
        """
        frame_boundaries = sorted(seconds_per_frame)
        num_frames = len(frame_boundaries) + 1
        x = np.zeros([num_frames, 12])

        for note in sequence:
            # if note.is_drum:
            #     continue
            # if note.program in UNPITCHED_PROGRAMS:
            #     continue

            start_frame = bisect.bisect_right(frame_boundaries, note.start)
            end_frame = bisect.bisect_left(frame_boundaries, note.end)
            pitch_class = note.pitch % 12

            if start_frame >= end_frame:
                x[start_frame, pitch_class] += note.end - note.start
            else:
                x[start_frame, pitch_class] += (
                    frame_boundaries[start_frame] - note.start
                )
                for frame in range(start_frame + 1, end_frame):
                    x[frame, pitch_class] += (
                        frame_boundaries[frame] - frame_boundaries[frame - 1]
                    )
                x[end_frame, pitch_class] += note.end - frame_boundaries[end_frame - 1]

        x_norm = np.linalg.norm(x, axis=1)
        nonzero_frames = x_norm > 0
        x[nonzero_frames, :] /= x_norm[nonzero_frames, np.newaxis]

        return x

    def _chord_frame_log_likelihood(self, note_pitch_vectors, chord_note_concentration):
        """Log-likelihood of observing each frame of note pitches under each chord."""
        return chord_note_concentration * np.dot(
            note_pitch_vectors, self._chord_pitch_vectors().T
        )

    def _key_chord_viterbi(
        self, chord_frame_loglik, key_chord_loglik, key_chord_transition_loglik
    ):
        """Use the Viterbi algorithm to infer a sequence of key-chord pairs."""
        num_frames, num_chords = chord_frame_loglik.shape
        num_key_chords = len(key_chord_transition_loglik)

        loglik_matrix = np.zeros([num_frames, num_key_chords])
        path_matrix = np.zeros([num_frames, num_key_chords], dtype=np.int32)

        # Initialize with a uniform distribution over keys.
        for i, key_chord in enumerate(self._KEY_CHORDS):
            key, unused_chord = key_chord
            chord_index = i % len(self._CHORDS)
            loglik_matrix[0, i] = (
                -np.log(12)
                + key_chord_loglik[key, chord_index]
                + chord_frame_loglik[0, chord_index]
            )

        for frame in range(1, num_frames):
            # At each frame, store the log-likelihood of the best sequence ending in
            # each key-chord pair, along with the index of the parent key-chord pair
            # from the previous frame.
            mat = (
                np.tile(loglik_matrix[frame - 1][:, np.newaxis], [1, num_key_chords])
                + key_chord_transition_loglik
            )
            path_matrix[frame, :] = mat.argmax(axis=0)
            loglik_matrix[frame, :] = mat[
                path_matrix[frame, :], range(num_key_chords)
            ] + np.tile(chord_frame_loglik[frame], 12)

        # Reconstruct the most likely sequence of key-chord pairs.
        path = [np.argmax(loglik_matrix[-1])]
        for frame in range(num_frames, 1, -1):
            path.append(path_matrix[frame - 1, path[-1]])

        return [
            (index // num_chords, self._CHORDS[index % num_chords])
            for index in path[::-1]
        ]

    def infer_chords_for_sequence(
        self,
        sequence,
        pos_per_chord,
        max_chords,
        key_chord_loglik=None,
        key_chord_transition_loglik=None,
        key_change_prob=0.001,
        chord_change_prob=0.5,
        chord_pitch_out_of_key_prob=0.01,
        chord_note_concentration=100.0,
        add_key_signatures=False,
    ):
        """Infer chords for a NoteSequence using the Viterbi algorithm.
        This uses some heuristics to infer chords for a quantized NoteSequence. At
        each chord position a key and chord will be inferred, and the chords will be
        added (as text annotations) to the sequence.
        If the sequence is quantized relative to meter, a fixed number of chords per
        bar will be inferred. Otherwise, the sequence is expected to have beat
        annotations and one chord will be inferred per beat.
        Args:
            sequence: The NoteSequence for which to infer chords. This NoteSequence will
                be modified in place.

            key_change_prob: Probability of a key change between two adjacent frames.
            chord_change_prob: Probability of a chord change between two adjacent
                frames.
            chord_pitch_out_of_key_prob: Probability of a pitch in a chord not belonging
                to the current key.
            chord_note_concentration: Concentration parameter for the distribution of
                observed pitches played over a chord. At zero, all pitches are equally
                likely. As concentration increases, observed pitches must match the
                chord pitches more closely.
            add_key_signatures: If True, also add inferred key signatures to
                `quantized_sequence` (and remove any existing key signatures).
        Raises:
            SequenceAlreadyHasChordsError: If `sequence` already has chords.
            QuantizationStatusError: If `sequence` is not quantized relative to
                meter but `chords_per_bar` is specified or no beat annotations are
                present.
            UncommonTimeSignatureError: If `chords_per_bar` is not specified and
                `sequence` is quantized and has an uncommon time signature.
            NonIntegerStepsPerChordError: If the number of quantized steps per chord
                is not an integer.
            EmptySequenceError: If `sequence` is empty.
            SequenceTooLongError: If the number of chords to be inferred is too
                large.
        """
        beats = [pos_per_chord * i for i in range(max_chords)]
        if len(beats) == 0:
            raise Exception("max chords should > 0")
        num_chords = len(beats)
        if num_chords > self._MAX_NUM_CHORDS:
            raise Exception(
                "NoteSequence too long for chord inference: %d frames" % num_chords
            )

        # Compute pitch vectors for each chord frame, then compute log-likelihood of
        # observing those pitch vectors under each possible chord.
        note_pitch_vectors = self.sequence_note_pitch_vectors(sequence, beats)
        chord_frame_loglik = self._chord_frame_log_likelihood(
            note_pitch_vectors, chord_note_concentration
        )

        # Compute distribution over chords for each key, and transition distribution
        # between key-chord pairs.
        if key_chord_loglik is None:
            key_chord_distribution = self._key_chord_distribution(
                chord_pitch_out_of_key_prob=chord_pitch_out_of_key_prob
            )
            key_chord_loglik = np.log(key_chord_distribution)

        if key_chord_transition_loglik is None:
            key_chord_transition_distribution = self._key_chord_transition_distribution(
                key_chord_distribution,
                key_change_prob=key_change_prob,
                chord_change_prob=chord_change_prob,
            )
            key_chord_transition_loglik = np.log(key_chord_transition_distribution)

        key_chords = self._key_chord_viterbi(
            chord_frame_loglik, key_chord_loglik, key_chord_transition_loglik
        )

        # if add_key_signatures:
        #     del sequence.key_signatures[:]

        # Add the inferred chord changes to the sequence, optionally adding key
        # signature(s) as well.
        # current_key_name = None
        # current_chord_name = None
        chords = []
        for frame, (key, chord) in enumerate(key_chords):
            # time = beats[frame]

            # if _PITCH_CLASS_NAMES[key] != current_key_name:
            #     # A key change was inferred.
            #     if add_key_signatures:
            #         ks = sequence.key_signatures.add()
            #         ks.time = time
            #         ks.key = key
            #     else:
            #         if current_key_name is not None:
            #             logging.info(
            #                 'Sequence has key change from %s to %s at %f seconds.',
            #                 current_key_name, _PITCH_CLASS_NAMES[key], time)
            #
            # current_key_name = _PITCH_CLASS_NAMES[key]
            if chord == self.NO_CHORD:
                figure = self.NO_CHORD
            else:
                root, kind = chord
                figure = "%s:%s" % (self._PITCH_CLASS_NAMES[root], kind)
            chords.append(figure)
        return chords
        # if figure != current_chord_name:
        #     ta = sequence.text_annotations.add()
        #     ta.time = time
        #     ta.quantized_step = 0 if frame == 0 else sorted_beat_steps[frame - 1]
        #     ta.text = figure
        # current_chord_name = figure


#!PS Not exactly sure what the bottom portion is doing except ruining my pylint score apparently?


class ChordInferenceError(Exception):  # pylint:disable=g-bad-exception-name
    pass


class SequenceAlreadyHasChordsError(ChordInferenceError):
    pass


class UncommonTimeSignatureError(ChordInferenceError):
    pass


class NonIntegerStepsPerChordError(ChordInferenceError):
    pass


class EmptySequenceError(ChordInferenceError):
    pass


class SequenceTooLongError(ChordInferenceError):
    pass

tuneease/getmusic/utils/magenta_chord_recognition.py
Pshah2023/TuneEase
import importlib
import random
import warnings
import os
import numpy as np
import torch


def seed_everything(seed, cudnn_deterministic=False):
    """
    Function that sets seed for pseudo-random number generators in:
    pytorch, numpy, python.random

    Args:
        seed: the integer value seed for global random state
    """
    if seed is not None:
        random.seed(seed)
        os.environ["PYTHONHASHSEED"] = str(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)

    if cudnn_deterministic:
        torch.backends.cudnn.deterministic = True
        warnings.warn(
            "You have chosen to seed training. "
            "This will turn on the CUDNN deterministic setting, "
            "which can slow down your training considerably! "
            "You may see unexpected behavior when restarting "
            "from checkpoints."
        )


def merge_opts_to_config(config, opts):
    def modify_dict(c, nl, v):
        if len(nl) == 1:
            c[nl[0]] = type(c[nl[0]])(v)
        else:
            # print(nl)
            c[nl[0]] = modify_dict(c[nl[0]], nl[1:], v)
        return c

    if opts is not None and len(opts) > 0:
        assert (
            len(opts) % 2 == 0
        ), "each opts should be given by the name and values! The length shall be even number!"
        for i in range(len(opts) // 2):
            name = opts[2 * i]
            value = opts[2 * i + 1]
            config = modify_dict(config, name.split("."), value)
    return config


def modify_config_for_debug(config):
    config["dataloader"]["num_workers"] = 0
    config["dataloader"]["batch_size"] = 1
    return config


def get_model_parameters_info(model):
    # for mn, m in model.named_modules():
    parameters = {"overall": {"trainable": 0, "non_trainable": 0, "total": 0}}
    for child_name, child_module in model.named_children():
        parameters[child_name] = {"trainable": 0, "non_trainable": 0}
        for pn, p in child_module.named_parameters():
            if p.requires_grad:
                parameters[child_name]["trainable"] += p.numel()
            else:
                parameters[child_name]["non_trainable"] += p.numel()
        parameters[child_name]["total"] = (
            parameters[child_name]["trainable"]
            + parameters[child_name]["non_trainable"]
        )

        parameters["overall"]["trainable"] += parameters[child_name]["trainable"]
        parameters["overall"]["non_trainable"] += parameters[child_name][
            "non_trainable"
        ]
        parameters["overall"]["total"] += parameters[child_name]["total"]

    # format the numbers
    def format_number(num):
        K = 2**10
        M = 2**20
        G = 2**30
        if num > G:  # K
            uint = "G"
            num = round(float(num) / G, 2)
        elif num > M:
            uint = "M"
            num = round(float(num) / M, 2)
        elif num > K:
            uint = "K"
            num = round(float(num) / K, 2)
        else:
            uint = ""

        return "{}{}".format(num, uint)

    def format_dict(d):
        for k, v in d.items():
            if isinstance(v, dict):
                format_dict(v)
            else:
                d[k] = format_number(v)

    format_dict(parameters)
    return parameters


def format_seconds(seconds):
    h = int(seconds // 3600)
    m = int(seconds // 60 - h * 60)
    s = int(seconds % 60)

    d = int(h // 24)
    h = h - d * 24

    if d == 0:
        if h == 0:
            if m == 0:
                ft = "{:02d}s".format(s)
            else:
                ft = "{:02d}m:{:02d}s".format(m, s)
        else:
            ft = "{:02d}h:{:02d}m:{:02d}s".format(h, m, s)

    else:
        ft = "{:d}d:{:02d}h:{:02d}m:{:02d}s".format(d, h, m, s)

    return ft


def instantiate_from_config(config):
    if config is None:
        return None
    if not "target" in config:
        raise KeyError("Expected key `target` to instantiate.")
    module, cls = config["target"].rsplit(".", 1)
    cls = getattr(importlib.import_module(module, package=None), cls)
    return cls(**config.get("params", dict()))


def class_from_string(class_name):
    module, cls = class_name.rsplit(".", 1)
    cls = getattr(importlib.import_module(module, package=None), cls)
    return cls


def get_all_file(dir, end_with=".h5"):
    if isinstance(end_with, str):
        end_with = [end_with]
    filenames = []
    for root, dirs, files in os.walk(dir):
        for f in files:
            for ew in end_with:
                if f.endswith(ew):
                    filenames.append(os.path.join(root, f))
                    break
    return filenames


def get_sub_dirs(dir, abs=True):
    sub_dirs = os.listdir(dir)
    if abs:
        sub_dirs = [os.path.join(dir, s) for s in sub_dirs]
    return sub_dirs


def get_model_buffer(model):
    state_dict = model.state_dict()
    buffers_ = {}
    params_ = {n: p for n, p in model.named_parameters()}

    for k in state_dict:
        if k not in params_:
            buffers_[k] = state_dict[k]
    return buffers_

tuneease/getmusic/utils/misc.py
Pshah2023/TuneEase
import logging


class ServerLogger:
    """
    A utility class that configures and then can provide a logger instance for server-side logging.

    Args:
        name (str): The name of the logger, typically the name of the log file.

    Attributes:
        logger (logging.Logger): The configured logger instance.

    Methods:
        __init__(self, name):
        get(self):

    Example:
        >>> ServerLogger("tuneease.log").get()
    """

    logged_messages = set()

    def __init__(self, name="tuneease.log"):
        """
        Initialize a ServerLogger with the given name.

        Args:
            name (str): The name of the logger, typically the name of the log file.

        Notes:
            Sets the loggger's log level to DEBUG, configures the log message format, and setups logs to ouput to a file and the console.
        """
        self.name = name.replace(".log", "")
        self.logger = logging.getLogger(name)

    def get(self):
        """
        Get the configured logger instance.

        Returns:
            logging.Logger: The configured logger instance for server-side logging.

        Example:
            >>> ServerLogger("tuneease.log").get()
        """
        self.logger.setLevel(logging.DEBUG)
        formatter = logging.Formatter(
            "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
            datefmt="%m/%d/%Y %I:%M:%S %p",
        )
        filehandler = logging.FileHandler(self.name + ".log")
        filehandler.setLevel(logging.DEBUG)
        filehandler.setFormatter(formatter)
        self.logger.addHandler(filehandler)
        streamhandler = logging.StreamHandler()
        streamhandler.setLevel(logging.INFO)
        streamhandler.setFormatter(formatter)
        streamhandler.addFilter(self.filter)
        self.logger.addHandler(streamhandler)
        return self.logger

    def filter(self, record):
        msg = record.getMessage()
        if msg in self.logged_messages:
            return False
        self.logged_messages.add(msg)
        return True

tuneease/logger.py
Pshah2023/TuneEase
import os
import sys
import shutil
import urllib.request
from .logger import ServerLogger
from argparse import ArgumentParser


def get_args():
    parser = ArgumentParser()
    parser.add_argument("--musescore_path", default=None)
    parser.add_argument("--checkpoint_path", default=None)
    args = parser.parse_args()
    logger = ServerLogger("tuneease.log").get()
    logger.debug("checkpoint_path from args: " + str(args.checkpoint_path))
    logger.debug("musescore_path from args: " + str(args.musescore_path))
    return args


class PathUtility:
    """
    A utility class for handling file paths and locating the MuseScore executable's path.

    Attributes:
        current_file_path (str): The absolute path of the current file.

    Methods:
        __init__(self):
        project_directory(self):
        museScore_path(self):

    Example:
        >>> path_util = PathUtility()
        >>> path_util.current_file_path
        /path/to/the/file
    """

    logger = ServerLogger("tuneease.log").get()

    def __init__(self):
        self.current_file_path = os.path.abspath(__file__)
        args = get_args()
        self.logger = ServerLogger("tuneease.log").get()
        if args.checkpoint_path:
            self.checkpoint_path = lambda: os.path.abspath(args.checkpoint_path)
        else:
            self.checkpoint_path(printsuccess=True)
        if args.musescore_path:
            self.musescore_path = lambda: os.path.abspath(args.musescore_path)
        else:
            self.musescore_path(printsuccess=True)

    def project_directory(self):
        """
        Get the path to the project directory.

        Args:
            None

        Returns:
            str: The path to the project directory.

        Example:
            >>> path_util = PathUtility()
            >>> project_dir = path_util.project_directory()
        """
        return os.path.dirname(os.path.dirname(self.current_file_path))

    def tuneease_dir(self):
        """
        Get the path to the python code directory.

        Args:
            None

        Returns:
            str: The path to the python code directory.

        Example:
            >>> path_util = PathUtility()
            >>> tuneease_dir = path_util.tuneease_dir()
        """
        return os.path.dirname(self.current_file_path)

    def musescore_path(self, printsuccess=False):
        """
        Locate and return the path to the MuseScore executable.

        Args:
            None

        Returns:
            str: The path to the MuseScore executable.

        Raises:
            FileNotFoundError: If MuseScore installation is not found.

        Example:
            >>> path_util = PathUtility()
            >>> muse_score_path = path_util.museScore_path()

        Notes:
            - On Linux, `<project directory>/temp/MuseScore.AppImage`
            - On macOS, the standard installation path and the user's 'bin' directory.
            - On Windows, it searches in common installation paths for 32 and 64-bit versions.
        """
        mscore_mac = "mscore" if sys.platform != "win32" else "MuseScore"
        musescore_windows = "MuseScore" if sys.platform != "win32" else "MuseScore4.exe"
        musescore_linux = (
            "MuseScore.AppImage" if sys.platform.startswith("linux") else None
        )
        path = shutil.which(mscore_mac)
        if path:
            return path
        path = shutil.which(musescore_windows)
        if path:
            return path
        directories = list()
        if sys.platform.startswith("linux"):
            directories += [os.path.join(self.project_directory(), "temp")]
        elif sys.platform == "darwin":
            directories += [
                "/Applications/MuseScore 4.app/Contents/Resources/bin",
                os.path.join(os.getenv("HOME"), "bin"),
            ]
        elif sys.platform == "win32":
            directories += [
                r"C:\Program Files\MuseScore 4\bin",
                r"C:\Program Files (x86)\MuseScore 4\bin",
            ]
        for directory in directories:
            if musescore_windows:
                path = os.path.join(directory, musescore_windows)
                if os.path.exists(path):
                    if printsuccess:
                        toprint = f"Found museScore {path}"
                        self.logger.info(toprint)
                    return path
            if mscore_mac:
                path = os.path.join(directory, mscore_mac)
                if os.path.exists(path):
                    if printsuccess:
                        toprint = f"Found museScore {path}"
                        self.logger.info(toprint)
                    return path
            if musescore_linux:
                path = os.path.join(directory, musescore_linux)
                if os.path.exists(path):
                    if printsuccess:
                        toprint = f"Found museScore {path}"
                        self.logger.info(toprint)
                    return path
                if sys.platform.startswith("linux"):
                    self.logger.info(
                        "Sometimes terminal downloads are too slow. Download with your browser:"
                    )
                    fileurl = "https://cdn.jsdelivr.net/musescore/v4.1.1/MuseScore-4.1.1.232071203-x86_64.AppImage"
                    self.logger.info(fileurl, "\nThen, move it this spot:")
                    filename = os.path.join(
                        self.project_directory(), "temp", "MuseScore.AppImage"
                    )
                    self.logger.info(filename)
                    os.makedirs(
                        os.path.join(self.project_directory(), "temp"), exist_ok=True
                    )
                    urllib.request.urlretrieve(
                        fileurl, filename, reporthook=self.download_progress_bar
                    )
                    toprint = f"Downloaded {filename}."
                    self.logger.info(toprint)
                    toprint = f"Please perform `chmod +x {filename}`"
                    self.logger.info(toprint)
                    return filename
        self.logger.info("On Windows and macOS, install MuseScore through the website:")
        self.logger.info("[https://musescore.org/en](https://musescore.org/en).")
        self.logger.info(
            "Attempts are made to find it automatically. If necessary, use --musescore_path <path>"
        )
        return None

    def checkpoint_path(self, printsuccess=False):
        path = os.path.join(self.project_directory(), "checkpoint.pth")
        if not os.path.exists(path):
            self.logger.info(
                "Install the below at the following path or use --checkpoint_path <path>"
            )
            fileurl = "https://1drv.ms/u/s!ArHNvccy1VzPkWGKXZDQY5k-kDi4?e=fFxcEq"
            self.logger.info(fileurl)
            self.logger.info(path)
            return None
        else:
            if printsuccess:
                toprint = f"Found museScore {path}"
                self.logger.info(toprint)
            return path

    def download_progress_bar(self, blocknum, blocksize, totalsize):
        downloaded = blocknum * blocksize
        percent_complete = (downloaded / totalsize) * 100
        self.logger.info(f"Downloaded {percent_complete:.2f}% from {totalsize}")

tuneease/pathutility.py
Pshah2023/TuneEase
import os
import datetime
import argparse
import numpy as np
from .compute import Node
from ..getmusic.utils.misc import seed_everything
from ..pathutility import PathUtility


def validate_binary_string(arg):
    if not arg.isdigit() or len(arg) != 7:
        raise argparse.ArgumentTypeError(
            "Binary string must be 7 characters long and consist of 0s and 1s only."
        )
    return arg


def get_args(input_args):
    node = Node()
    parser = argparse.ArgumentParser(description="PyTorch Training script")
    pathutil = PathUtility()
    if not pathutil.checkpoint_path():
        raise FileNotFoundError("Could not find the checkpoint.pth")
    # Contains additional args as necessary
    common_args = [
        ("--config_file", str, "configs/train.yaml", "path of config file"),
        (
            "--name",
            str,
            "inference_cache",
            "the name of this experiment, if not provided, set to the name of config file",
        ),
        ("--output", str, "cache", "directory to save the results"),
        ("--log_frequency", int, 10, "print frequency"),
        (
            "--load_path",
            str,
            pathutil.checkpoint_path(),
            "path to model that need to be loaded used for loading pretrained model",
        ),
        ("--resume_name", str, None, "resume one experiment with the given name"),
        ("--auto_resume", bool, False, "resume the training"),
        ("--num_node", int, node.NUM_NODE, "number of nodes for distributed training"),
        ("--ngpus_per_node", int, 8, "number of gpu on one node"),
        ("--node_rank", int, node.NODE_RANK, "node rank for distributed training"),
        ("--dist_url", str, node.DIST_URL, "url used to set up distributed training"),
        (
            "--gpu",
            int,
            0,
            "GPU id to use. If given, only the specific gpu will be used, and ddp will be disabled",
        ),
        ("--local_rank", int, -1, "node rank for distributed training"),
        ("--sync_bn", bool, False, "use sync BN layer"),
        ("--tensorboard", bool, False, "use tensorboard for logging"),
        ("--timestamp", bool, False, "use tensorboard for logging"),
        ("--seed", int, 0, "seed for initializing training"),
        ("--cudnn_deterministic", bool, False, "set cudnn.deterministic True"),
        ("--amp", bool, False, "automatic mixture of precesion"),
        (
            "--conditional_name",
            str,
            None,
            "I'm not sure what this does. Send a pull request please!",
        ),
        (
            "--content_name",
            str,
            None,
            "I'm not sure what this does. Send a pull request please!",
        ),
        (
            "--conditional_tracks",
            validate_binary_string,
            "0000000",
            "the content tracks to base the generation on",
        ),
        (
            "--content_tracks",
            validate_binary_string,
            "0000110",
            "the tracks to generate",
        ),
        ("--debug", bool, False, "set as debug mode"),
        (
            "--do_sample",
            bool,
            True,
            "I'm not sure what this does. Send a pull request please!",
        ),
        ("--file_path", str, None, "the output filepath"),
        (
            "--skip_step",
            int,
            0,
            "I'm not sure what this does. Send a pull request please!",
        ),
        (
            "--decode_chord",
            bool,
            False,
            "I'm not sure what this does. Send a pull request please!",
        ),
        (
            "--chord_from_single",
            bool,
            False,
            "I'm not sure what this does. Send a pull request please!",
        ),
        (
            "--no_ema",
            bool,
            True,
            "I'm not sure what this does. Send a pull request please!",
        ),
        (
            "--no_load_optimizer_and_scheduler",
            bool,
            True,
            "I'm not sure what this does. Send a pull request please!",
        ),
        (
            "--no_load_others",
            bool,
            True,
            "I'm not sure what this does. Send a pull request please!",
        ),
        (
            "--training_model",
            bool,
            False,
            "If you are training the model, this changes the output directory to 'OUTPUT.'",
        ),
        ("--use_gpu", int, 0, "Whether to use GPU or not. False for CPU."),
    ]
    for arg, arg_type, default_val, help in common_args:
        parser.add_argument(arg, type=arg_type, default=default_val, help=help)
    parser.add_argument(
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )
    if input_args == None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(input_args)
    args.cwd = os.path.abspath(os.path.dirname(__file__))
    # Common to all Args
    seed_everything(args.seed, args.cudnn_deterministic)
    # Modified logic to work with all args
    if args.resume_name is not None:
        args.name = args.resume_name
        args.config_file = os.path.join(
            args.output, args.resume_name, "configs", "config.yaml"
        )
        args.auto_resume = True
    else:
        if args.name == "":
            args.name = os.path.basename(args.config_file).replace(".yaml", "")
        if args.timestamp:
            assert (
                not args.auto_resume
            ), "for timstamp, auto resume is hard to find the save directory"
            time_str = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M")
            args.name = time_str + "-" + args.name
    if args.debug:
        args.name = "debug"
        if args.gpu is None:
            args.gpu = 0
    if args.training_model:
        args.output = "OUTPUT"
    # Common to all args
    random_seconds_shift = datetime.timedelta(seconds=np.random.randint(60))
    now = (datetime.datetime.now() - random_seconds_shift).strftime("%Y-%m-%dT%H-%M-%S")
    args.save_dir = os.path.join(args.output, args.name, now)
    args.use_gpu = bool(int(args.use_gpu))
    return args

tuneease/pipeline/args.py
Pshah2023/TuneEase
import os


class Node:
    NODE_RANK = None
    MASTER_ADDR = None
    MASTER_PORT = None
    DIST_URL = None
    NUM_NODE = None

    def __init__(self) -> None:
        self.NODE_RANK = int(os.environ["INDEX"] if "INDEX" in os.environ else 0)

        self.MASTER_ADDR, self.MASTER_PORT = (
            (os.environ["CHIEF_IP"], 22275)
            if "CHIEF_IP" in os.environ
            else ("127.0.0.1", 29500)
        )
        self.DIST_URL = f"tcp://{self.MASTER_ADDR}:%{self.MASTER_PORT}"

        self.NUM_NODE = os.environ["HOST_NUM"] if "HOST_NUM" in os.environ else 1

tuneease/pipeline/compute.py
Pshah2023/TuneEase
from os import path


class Config:
    config = None

    def __init__(self) -> None:
        self.config = {
            "model": {
                "target": "tuneease.getmusic.modeling.models.dfm.DFM",
                "params": {
                    "diffusion_config": {
                        "target": "tuneease.getmusic.modeling.roformer.diffusion_roformer.DiffusionRFM",
                        "params": {
                            "diffusion_step": 100,
                            "alpha_init_type": "alpha1",
                            "auxiliary_loss_weight": 0.001,
                            "adaptive_auxiliary_loss": True,
                            "roformer_config": {
                                "target": "tuneease.getmusic.modeling.roformer.roformer_utils.DiffusionRoformerModel",
                                "params": {"vocab_size": 11880, "cond_weight": 0.5},
                            },
                        },
                    }
                },
            },
            "solver": {
                "base_lr": 3e-06,
                "adjust_lr": "none",
                "max_epochs": 50,
                "save_epochs": 10,
                "validation_epochs": 1,
                "sample_iterations": "epoch",
                "validate_iterations": 1000,
                "vocab_path": path.join(
                    path.dirname(path.dirname(path.abspath(__file__))),
                    "getmusic",
                    "utils",
                    "dict.txt",
                ),
                "print_specific_things": True,
                "ema": {"decay": 0.9, "update_interval": 100, "device": "cpu"},
                "clip_grad_norm": {
                    "target": "tuneease.getmusic.engine.clip_grad_norm.ClipGradNorm",
                    "params": {
                        "start_iteration": 0,
                        "end_iteration": 5000,
                        "max_norm": 0.5,
                    },
                },
                "optimizers_and_schedulers": [
                    {
                        "name": "none",
                        "optimizer": {
                            "target": "torch.optim.AdamW",
                            "step_iteration": 1,
                            "params": {"betas": (0.9, 0.999), "weight_decay": 0.01},
                        },
                        "scheduler": {
                            "step_iteration": 1,
                            "target": "tuneease.getmusic.engine.lr_scheduler.LinearDecayLRWithWarmup",
                            "params": {
                                "min_lr": 1e-06,
                                "warmup_lr": 0.0001,
                                "warmup": 1000,
                                "T_max": 300000,
                            },
                        },
                    }
                ],
            },
            "dataloader": {
                "batch_size": 3,
                "num_workers": 28,
                # 'train_datasets': [
                #     {
                #         'target':
                #             'tuneease.getmusic.data.bigdata.BigDataset',
                #             'params': {
                #                 'prefix': 'train',
                #                 'path': '/your-data-path',
                #                 'vocab_size': 11880
                #                 }
                #             }
                #     ], 'validation_datasets': [
                #         {
                #             'target': 'tuneease.getmusic.data.bigdata.BigDataset',
                #             'params': {
                #                 'prefix': 'valid',
                #                 'path': '/your-data-path',
                #                 'vocab_size': 11880
                #                 }
                #             }
                #         ]
            },
        }

tuneease/pipeline/config.py
Pshah2023/TuneEase
import miditoolkit
import math
from ..getmusic.utils.midi_config import (
    POS_RESOLUTION,
    BEAT_NOTE_FACTOR,
    TRUNC_POS,
    MAX_PITCH,
    MAX_INST,
    FILTER_SYMBOLIC,
    FILTER_SYMBOLIC_PPL,
)
from .encoding_helpers import (
    t2e,
    e2t,
    d2e,
    e2d,
    time_signature_reduce,
    b2e,
    e2b,
    v2e,
    e2v,
)
from .presets import RootKinds
from .item import Item
from ..getmusic.utils.magenta_chord_recognition import Magenta
from .processing import normalize_to_c_major
from .key_chord import KeyChordDetails


def MIDI_to_encoding(
    filename, midi_obj, with_chord=None, condition_inst=None, chord_from_single=None
):
    magenta = Magenta()
    k_c_d = KeyChordDetails()
    r_h = RootKinds()
    if filename == "track_generation.py":

        def time_to_pos(t):
            return round(t * POS_RESOLUTION / midi_obj.ticks_per_beat)

        notes_start_pos = [
            time_to_pos(j.start) for i in midi_obj.instruments for j in i.notes
        ]
        if len(notes_start_pos) == 0:
            return list()
        max_pos = max(notes_start_pos) + 1
        pos_to_info = [[None for _ in range(4)] for _ in range(max_pos)]
        tsc = midi_obj.time_signature_changes
        tpc = midi_obj.tempo_changes
        for i in range(len(tsc)):
            for j in range(
                time_to_pos(tsc[i].time),
                time_to_pos(tsc[i + 1].time) if i < len(tsc) - 1 else max_pos,
            ):
                if j < len(pos_to_info):
                    pos_to_info[j][1] = t2e(
                        time_signature_reduce(tsc[i].numerator, tsc[i].denominator)
                    )
        for i in range(len(tpc)):
            for j in range(
                time_to_pos(tpc[i].time),
                time_to_pos(tpc[i + 1].time) if i < len(tpc) - 1 else max_pos,
            ):
                if j < len(pos_to_info):
                    pos_to_info[j][3] = b2e(tpc[i].tempo)
        for j in range(len(pos_to_info)):
            if pos_to_info[j][1] is None:
                # MIDI default time signature
                pos_to_info[j][1] = t2e(time_signature_reduce(4, 4))
            if pos_to_info[j][3] is None:
                pos_to_info[j][3] = b2e(120.0)  # MIDI default tempo (BPM)
        cnt = 0
        bar = 0
        measure_length = None
        for j in range(len(pos_to_info)):
            timesignature = e2t(pos_to_info[j][1])
            if cnt == 0:
                measure_length = (
                    timesignature[0]
                    * BEAT_NOTE_FACTOR
                    * POS_RESOLUTION
                    // timesignature[1]
                )
            pos_to_info[j][0] = bar
            pos_to_info[j][2] = cnt
            cnt += 1
            if cnt >= measure_length:
                assert (
                    cnt == measure_length
                ), "invalid time signature change: pos = {}".format(j)
                cnt -= measure_length
                bar += 1
        encoding = []

        for inst in midi_obj.instruments:
            for note in inst.notes:
                if time_to_pos(note.start) >= TRUNC_POS:
                    continue
                info = pos_to_info[time_to_pos(note.start)]
                duration = d2e(time_to_pos(note.end) - time_to_pos(note.start))
                encoding.append(
                    [
                        info[0],
                        info[2],
                        MAX_INST + 1 if inst.is_drum else inst.program,
                        note.pitch + MAX_PITCH + 1 if inst.is_drum else note.pitch,
                        duration,
                        v2e(note.velocity),
                        info[1],
                        info[3],
                    ]
                )
        if len(encoding) == 0:
            return list()
        encoding.sort()
        encoding, is_major, pitch_shift = normalize_to_c_major(filename, encoding)
        # extract chords
        if with_chord:
            max_pos = 0
            note_items = []
            for note in encoding:
                if (0 < note[3] < 128) and (note[2] in [0, 25, 32, 48, 80]):
                    if chord_from_single and (str(note[2]) not in condition_inst):
                        continue
                    timesignature = e2t(note[6])
                    measure_length = (
                        timesignature[0]
                        * BEAT_NOTE_FACTOR
                        * POS_RESOLUTION
                        // timesignature[1]
                    )
                    max_pos = max(
                        max_pos, measure_length * note[0] + note[1] + e2d(note[4])
                    )
                    note_items.append(
                        Item(
                            name="On",
                            start=measure_length * note[0] + note[1],
                            end=measure_length * note[0] + note[1] + e2d(note[4]),
                            vel=e2v(note[5]),
                            pitch=note[3],
                            track=0,
                        )
                    )
            note_items.sort(key=lambda x: (x.start, -x.end))
            pos_per_chord = measure_length
            max_chords = round(max_pos // pos_per_chord + 0.5)
            if max_chords > 0:
                chords = magenta.infer_chords_for_sequence(
                    note_items,
                    pos_per_chord=pos_per_chord,
                    max_chords=max_chords,
                    key_chord_loglik=k_c_d.key_chord_loglik,
                    key_chord_transition_loglik=k_c_d.key_chord_transition_loglik,
                )
            else:
                chords = []
            bar_idx = 0
            for chord in chords:
                if chord == "N.C.":
                    bar_idx += 1
                    continue
                r, k = chord.split(":")
                if k == "":
                    k = "null"
                elif k == "7":
                    k = "seven"
                encoding.append(
                    (
                        bar_idx,
                        0,
                        129,
                        r_h.root_dict[r],
                        r_h.kind_dict[k],
                        0,
                        t2e(time_signature_reduce(4, 4)),
                        0,
                    )
                )
                bar_idx += 1
            encoding.sort()
        return encoding, pitch_shift, tpc
    elif filename == "position_generation.py":

        def time_to_pos(t):
            return round(t * POS_RESOLUTION / midi_obj.ticks_per_beat)

        notes_start_pos = [
            time_to_pos(j.start) for i in midi_obj.instruments for j in i.notes
        ]
        if len(notes_start_pos) == 0:
            return list()
        max_pos = max(notes_start_pos) + 1
        pos_to_info = [
            [None for _ in range(4)] for _ in range(max_pos)
        ]  # (Measure, TimeSig, Pos, Tempo)
        tsc = (
            midi_obj.time_signature_changes
        )  # [TimeSignature(numerator=4, denominator=4, time=0)]
        tpc = midi_obj.tempo_changes  # [TempoChange(tempo=120.0, time=0)]
        for i in range(len(tsc)):
            for j in range(
                time_to_pos(tsc[i].time),
                time_to_pos(tsc[i + 1].time) if i < len(tsc) - 1 else max_pos,
            ):
                if j < len(pos_to_info):
                    pos_to_info[j][1] = t2e(
                        time_signature_reduce(tsc[i].numerator, tsc[i].denominator)
                    )
        for i in range(len(tpc)):
            for j in range(
                time_to_pos(tpc[i].time),
                time_to_pos(tpc[i + 1].time) if i < len(tpc) - 1 else max_pos,
            ):
                if j < len(pos_to_info):
                    pos_to_info[j][3] = b2e(tpc[i].tempo)
        for j in range(len(pos_to_info)):
            if pos_to_info[j][1] is None:
                # MIDI default time signature
                pos_to_info[j][1] = t2e(time_signature_reduce(4, 4))
            if pos_to_info[j][3] is None:
                pos_to_info[j][3] = b2e(120.0)  # MIDI default tempo (BPM)
        cnt = 0
        bar = 0
        measure_length = None
        for j in range(len(pos_to_info)):  # 它这里是不管这个位置有没有音符，都占个位
            timesignature = e2t(pos_to_info[j][1])
            if cnt == 0:
                measure_length = (
                    timesignature[0]
                    * BEAT_NOTE_FACTOR
                    * POS_RESOLUTION
                    // timesignature[1]
                )  # 比如一个3/4的ts，一个4/4的小节有16pos，所以3/4一小节就有12
            pos_to_info[j][0] = bar
            pos_to_info[j][2] = cnt
            cnt += 1
            if cnt >= measure_length:  # 如果cnt>了measure长度，就是下一个小节了，cnt清零，bar index加一
                assert (
                    cnt == measure_length
                ), "invalid time signature change: pos = {}".format(j)
                cnt -= measure_length
                bar += 1
        encoding = []
        for inst in midi_obj.instruments:
            for note in inst.notes:
                if time_to_pos(note.start) >= TRUNC_POS:
                    continue
                info = pos_to_info[time_to_pos(note.start)]
                duration = d2e(time_to_pos(note.end) - time_to_pos(note.start))
                encoding.append(
                    [
                        info[0],
                        info[2],
                        MAX_INST + 1 if inst.is_drum else inst.program,
                        note.pitch + MAX_PITCH + 1 if inst.is_drum else note.pitch,
                        duration,
                        v2e(note.velocity),
                        info[1],
                        info[3],
                    ]
                )
        if len(encoding) == 0:
            return list()
        encoding.sort()
        encoding, is_major, pitch_shift = normalize_to_c_major(filename, encoding)
        # extract chords
        if with_chord:
            max_pos = 0
            note_items = []
            for note in encoding:
                if 0 < note[3] < 128:  # and str(note[2]) in condition_inst:
                    timesignature = e2t(note[6])
                    measure_length = (
                        timesignature[0]
                        * BEAT_NOTE_FACTOR
                        * POS_RESOLUTION
                        // timesignature[1]
                    )
                    max_pos = max(
                        max_pos, measure_length * note[0] + note[1] + e2d(note[4])
                    )
                    note_items.append(
                        Item(
                            name="On",
                            start=measure_length * note[0] + note[1],
                            end=measure_length * note[0] + note[1] + e2d(note[4]),
                            vel=e2v(note[5]),
                            pitch=note[3],
                            track=0,
                        )
                    )
            note_items.sort(key=lambda x: (x.start, -x.end))
            pos_per_chord = measure_length
            max_chords = round(max_pos // pos_per_chord + 0.5)
            if max_chords > 0:
                chords = magenta.infer_chords_for_sequence(
                    note_items,
                    pos_per_chord=pos_per_chord,
                    max_chords=max_chords,
                    key_chord_loglik=k_c_d.key_chord_loglik,
                    key_chord_transition_loglik=k_c_d.key_chord_transition_loglik,
                )
            else:
                chords = []
            bar_idx = 0
            for chord in chords:
                if chord == "N.C.":
                    bar_idx += 1
                    continue
                r, k = chord.split(":")
                if k == "":
                    k = "null"
                elif k == "7":
                    k = "seven"
                encoding.append(
                    (
                        bar_idx,
                        0,
                        129,
                        r_h.root_dict[r],
                        r_h.kind_dict[k],
                        0,
                        t2e(time_signature_reduce(4, 4)),
                        0,
                    )
                )
                bar_idx += 1
            encoding.sort()
        return encoding, pitch_shift, tpc
    elif filename == "to_oct.py":

        def time_to_pos(t):
            return round(t * POS_RESOLUTION / midi_obj.ticks_per_beat)

        notes_start_pos = [
            time_to_pos(j.start) for i in midi_obj.instruments for j in i.notes
        ]
        if len(notes_start_pos) == 0:
            return list()
        max_pos = min(max(notes_start_pos) + 1, TRUNC_POS)
        pos_to_info = [
            [None for _ in range(4)] for _ in range(max_pos)
        ]  # (Measure, TimeSig, Pos, Tempo)
        tsc = (
            midi_obj.time_signature_changes
        )  # [TimeSignature(numerator=4, denominator=4, time=0)]
        tpc = midi_obj.tempo_changes  # [TempoChange(tempo=120.0, time=0)]
        # filter tempo and ts change
        if len(tsc) > 1 or len(tpc) > 1:
            return ["welcome use my code"]
        for i in range(len(tsc)):
            for j in range(
                time_to_pos(tsc[i].time),
                time_to_pos(tsc[i + 1].time) if i < len(tsc) - 1 else max_pos,
            ):
                if j < len(pos_to_info):
                    pos_to_info[j][1] = t2e(
                        time_signature_reduce(tsc[i].numerator, tsc[i].denominator)
                    )
        for i in range(len(tpc)):
            for j in range(
                time_to_pos(tpc[i].time),
                time_to_pos(tpc[i + 1].time) if i < len(tpc) - 1 else max_pos,
            ):
                if j < len(pos_to_info):
                    pos_to_info[j][3] = b2e(tpc[i].tempo)
        for j in range(len(pos_to_info)):
            if pos_to_info[j][1] is None:
                # MIDI default time signature
                pos_to_info[j][1] = t2e(time_signature_reduce(4, 4))
            if pos_to_info[j][3] is None:
                pos_to_info[j][3] = b2e(120.0)  # MIDI default tempo (BPM)
        cnt = 0
        bar = 0
        measure_length = None
        for j in range(len(pos_to_info)):
            timesignature = e2t(pos_to_info[j][1])
            if cnt == 0:
                measure_length = (
                    timesignature[0]
                    * BEAT_NOTE_FACTOR
                    * POS_RESOLUTION
                    // timesignature[1]
                )
            pos_to_info[j][0] = bar
            pos_to_info[j][2] = cnt
            cnt += 1
            if cnt >= measure_length:
                assert (
                    cnt == measure_length
                ), "invalid time signature change: pos = {}".format(j)
                cnt -= measure_length
                bar += 1
        encoding = []
        start_distribution = [0] * POS_RESOLUTION
        for inst in midi_obj.instruments:
            for note in inst.notes:
                if time_to_pos(note.start) >= TRUNC_POS:
                    continue
                start_distribution[time_to_pos(note.start) % POS_RESOLUTION] += 1
                info = pos_to_info[time_to_pos(note.start)]
                duration = d2e(time_to_pos(note.end) - time_to_pos(note.start))
                encoding.append(
                    (
                        info[0],
                        info[2],
                        MAX_INST + 1 if inst.is_drum else inst.program,
                        note.pitch + MAX_PITCH + 1 if inst.is_drum else note.pitch,
                        duration,
                        v2e(note.velocity),
                        info[1],
                        info[3],
                    )
                )
        if len(encoding) == 0:
            return list()
        tot = sum(start_distribution)
        start_ppl = 2 ** sum(
            (
                0 if x == 0 else -(x / tot) * math.log2((x / tot))
                for x in start_distribution
            )
        )
        # filter unaligned music
        if FILTER_SYMBOLIC:
            assert (
                start_ppl <= FILTER_SYMBOLIC_PPL
            ), "filtered out by the symbolic filter: ppl = {:.2f}".format(start_ppl)
        # normalize
        encoding.sort()
        encoding, is_major = normalize_to_c_major(filename, encoding)
        # extract chords
        max_pos = 0
        note_items = []
        for note in encoding:
            if 0 <= note[3] < 128:
                timesignature = e2t(note[6])
                measure_length = (
                    timesignature[0]
                    * BEAT_NOTE_FACTOR
                    * POS_RESOLUTION
                    // timesignature[1]
                )
                max_pos = max(
                    max_pos, measure_length * note[0] + note[1] + e2d(note[4])
                )
                note_items.append(
                    Item(
                        name="On",
                        start=measure_length * note[0] + note[1],
                        end=measure_length * note[0] + note[1] + e2d(note[4]),
                        vel=e2v(note[5]),
                        pitch=note[3],
                        track=0,
                    )
                )
        note_items.sort(key=lambda x: (x.start, -x.end))
        pos_per_chord = measure_length
        max_chords = round(max_pos // pos_per_chord + 0.5)
        chords = magenta.infer_chords_for_sequence(
            note_items,
            pos_per_chord=pos_per_chord,
            max_chords=max_chords,
            key_chord_loglik=k_c_d.key_chord_loglik,
            key_chord_transition_loglik=k_c_d.key_chord_transition_loglik,
        )
        bar_idx = 0
        for chord in chords:
            r, k = chord.split(":")
            if k == "":
                k = "null"
            elif k == "7":
                k = "seven"
            encoding.append(
                (
                    bar_idx,
                    0,
                    129,
                    r_h.root_dict[r],
                    r_h.kind_dict[k],
                    0,
                    t2e(time_signature_reduce(4, 4)),
                    0,
                )
            )
            bar_idx += 1
        encoding.sort()
        return encoding


def encoding_to_MIDI(filename, encoding, tpc=None, decode_chord=None):
    magenta = Magenta()
    r_h = RootKinds()
    if filename == "track_generation.py":
        tmp = encoding.strip().split("<0-")[1:]
        encoding = []
        for item in tmp:
            tmp2 = item.strip()[:-1].split("> <")
            encoding.append([int(tmp2[0])] + [int(i[2:]) for i in tmp2[1:]])
        del tmp
        bar_to_timesig = [list() for _ in range(max(map(lambda x: x[0], encoding)) + 1)]
        for i in encoding:
            bar_to_timesig[i[0]].append(i[6])
        bar_to_timesig = [
            max(set(i), key=i.count) if len(i) > 0 else None for i in bar_to_timesig
        ]
        for i in range(len(bar_to_timesig)):
            if bar_to_timesig[i] is None:
                bar_to_timesig[i] = (
                    t2e(time_signature_reduce(4, 4))
                    if i == 0
                    else bar_to_timesig[i - 1]
                )
        bar_to_pos = [None] * len(bar_to_timesig)
        cur_pos = 0
        for i in range(len(bar_to_pos)):
            bar_to_pos[i] = cur_pos
            ts = e2t(bar_to_timesig[i])
            measure_length = ts[0] * BEAT_NOTE_FACTOR * POS_RESOLUTION // ts[1]
            cur_pos += measure_length
        pos_to_tempo = [
            list() for _ in range(cur_pos + max(map(lambda x: x[1], encoding)))
        ]
        for i in encoding:
            pos_to_tempo[bar_to_pos[i[0]] + i[1]].append(i[7])
        pos_to_tempo = [
            round(sum(i) / len(i)) if len(i) > 0 else None for i in pos_to_tempo
        ]
        for i in range(len(pos_to_tempo)):
            if pos_to_tempo[i] is None:
                pos_to_tempo[i] = b2e(120.0) if i == 0 else pos_to_tempo[i - 1]
        midi_obj = miditoolkit.midi.parser.MidiFile()
        midi_obj.tempo_changes = tpc

        def get_tick(bar, pos):
            return (bar_to_pos[bar] + pos) * midi_obj.ticks_per_beat // POS_RESOLUTION

        midi_obj.instruments = [
            miditoolkit.containers.Instrument(
                program=(0 if i == 128 else i), is_drum=(i == 128), name=str(i)
            )
            for i in range(128 + 1)
        ]
        for i in encoding:
            start = get_tick(i[0], i[1])
            program = i[2]
            if program == 129 and decode_chord:
                root_name = r_h.root_list[i[3]]
                kind_name = r_h.kind_list[i[4]]
                root_pitch_shift = r_h.root_dict[root_name]
                end = start + get_tick(0, e2d(1))
                for kind_shift in magenta._CHORD_KIND_PITCHES[kind_name]:
                    pitch = 36 + root_pitch_shift + kind_shift
                    midi_obj.instruments[1].notes.append(
                        miditoolkit.containers.Note(
                            start=start, end=end, pitch=pitch, velocity=e2v(20)
                        )
                    )
            elif program != 129:
                pitch = i[3] - 128 if program == 128 else i[3]
                if pitch < 0:
                    continue
                duration = get_tick(0, e2d(i[4]))
                if duration == 0:
                    duration = 1
                end = start + duration
                velocity = e2v(i[5])

                midi_obj.instruments[program].notes.append(
                    miditoolkit.containers.Note(
                        start=start, end=end, pitch=pitch, velocity=velocity
                    )
                )
        midi_obj.instruments = [i for i in midi_obj.instruments if len(i.notes) > 0]
        cur_ts = None
        for i in range(len(bar_to_timesig)):
            new_ts = bar_to_timesig[i]
            if new_ts != cur_ts:
                numerator, denominator = e2t(new_ts)
                midi_obj.time_signature_changes.append(
                    miditoolkit.containers.TimeSignature(
                        numerator=numerator,
                        denominator=denominator,
                        time=get_tick(i, 0),
                    )
                )
                cur_ts = new_ts
        cur_tp = None
        for i in range(len(pos_to_tempo)):
            new_tp = pos_to_tempo[i]
            if new_tp != cur_tp:
                tempo = e2b(new_tp)
                midi_obj.tempo_changes.append(
                    miditoolkit.containers.TempoChange(tempo=tempo, time=get_tick(0, i))
                )
                cur_tp = new_tp
        return midi_obj
    elif filename == "to_oct.py":
        # TODO: filter out non-valid notes and error handling
        bar_to_timesig = [list() for _ in range(max(map(lambda x: x[0], encoding)) + 1)]
        for i in encoding:
            bar_to_timesig[i[0]].append(i[6])
        bar_to_timesig = [
            max(set(i), key=i.count) if len(i) > 0 else None for i in bar_to_timesig
        ]
        for i in range(len(bar_to_timesig)):
            if bar_to_timesig[i] is None:
                bar_to_timesig[i] = (
                    t2e(time_signature_reduce(4, 4))
                    if i == 0
                    else bar_to_timesig[i - 1]
                )
        bar_to_pos = [None] * len(bar_to_timesig)
        cur_pos = 0
        for i in range(len(bar_to_pos)):
            bar_to_pos[i] = cur_pos
            ts = e2t(bar_to_timesig[i])
            measure_length = ts[0] * BEAT_NOTE_FACTOR * POS_RESOLUTION // ts[1]
            cur_pos += measure_length
        pos_to_tempo = [
            list() for _ in range(cur_pos + max(map(lambda x: x[1], encoding)))
        ]
        for i in encoding:
            pos_to_tempo[bar_to_pos[i[0]] + i[1]].append(i[7])
        pos_to_tempo = [
            round(sum(i) / len(i)) if len(i) > 0 else None for i in pos_to_tempo
        ]
        for i in range(len(pos_to_tempo)):
            if pos_to_tempo[i] is None:
                pos_to_tempo[i] = b2e(120.0) if i == 0 else pos_to_tempo[i - 1]
        midi_obj = miditoolkit.midi.parser.MidiFile()

        def get_tick(bar, pos):
            return (bar_to_pos[bar] + pos) * midi_obj.ticks_per_beat // POS_RESOLUTION

        midi_obj.instruments = [
            miditoolkit.containers.Instrument(
                program=(0 if i == 128 else i), is_drum=(i == 128), name=str(i)
            )
            for i in range(128 + 1)
        ]
        for i in encoding:
            start = get_tick(i[0], i[1])
            program = i[2]
            pitch = i[3] - 128 if program == 128 else i[3]
            duration = get_tick(0, e2d(i[4]))
            if duration == 0:
                duration = 1
            end = start + duration
            velocity = e2v(i[5])
            midi_obj.instruments[program].notes.append(
                miditoolkit.containers.Note(
                    start=start, end=end, pitch=pitch, velocity=velocity
                )
            )
        midi_obj.instruments = [i for i in midi_obj.instruments if len(i.notes) > 0]
        cur_ts = None
        for i in range(len(bar_to_timesig)):
            new_ts = bar_to_timesig[i]
            if new_ts != cur_ts:
                numerator, denominator = e2t(new_ts)
                midi_obj.time_signature_changes.append(
                    miditoolkit.containers.TimeSignature(
                        numerator=numerator,
                        denominator=denominator,
                        time=get_tick(i, 0),
                    )
                )
                cur_ts = new_ts
        cur_tp = None
        for i in range(len(pos_to_tempo)):
            new_tp = pos_to_tempo[i]
            if new_tp != cur_tp:
                tempo = e2b(new_tp)
                midi_obj.tempo_changes.append(
                    miditoolkit.containers.TempoChange(tempo=tempo, time=get_tick(0, i))
                )
                cur_tp = new_tp
        return midi_obj
    elif filename == "position_generation.py":
        tmp = encoding.strip().split("<0-")[1:]
        encoding = []
        for item in tmp:
            tmp2 = item.strip()[:-1].split("> <")
            encoding.append([int(tmp2[0])] + [int(i[2:]) for i in tmp2[1:]])
        del tmp

        bar_to_timesig = [list() for _ in range(max(map(lambda x: x[0], encoding)) + 1)]
        for i in encoding:
            bar_to_timesig[i[0]].append(i[6])
        bar_to_timesig = [
            max(set(i), key=i.count) if len(i) > 0 else None for i in bar_to_timesig
        ]
        for i in range(len(bar_to_timesig)):
            if bar_to_timesig[i] is None:
                bar_to_timesig[i] = (
                    t2e(time_signature_reduce(4, 4))
                    if i == 0
                    else bar_to_timesig[i - 1]
                )
        bar_to_pos = [None] * len(bar_to_timesig)
        cur_pos = 0
        for i in range(len(bar_to_pos)):
            bar_to_pos[i] = cur_pos
            ts = e2t(bar_to_timesig[i])
            measure_length = ts[0] * BEAT_NOTE_FACTOR * POS_RESOLUTION // ts[1]
            cur_pos += measure_length
        pos_to_tempo = [
            list() for _ in range(cur_pos + max(map(lambda x: x[1], encoding)))
        ]
        for i in encoding:
            pos_to_tempo[bar_to_pos[i[0]] + i[1]].append(i[7])
        pos_to_tempo = [
            round(sum(i) / len(i)) if len(i) > 0 else None for i in pos_to_tempo
        ]
        for i in range(len(pos_to_tempo)):
            if pos_to_tempo[i] is None:
                pos_to_tempo[i] = b2e(120.0) if i == 0 else pos_to_tempo[i - 1]
        midi_obj = miditoolkit.midi.parser.MidiFile()
        midi_obj.tempo_changes = tpc

        def get_tick(bar, pos):
            return (bar_to_pos[bar] + pos) * midi_obj.ticks_per_beat // POS_RESOLUTION

        midi_obj.instruments = [
            miditoolkit.containers.Instrument(
                program=(0 if i == 128 else i), is_drum=(i == 128), name=str(i)
            )
            for i in range(128 + 1)
        ]
        for i in encoding:
            start = get_tick(i[0], i[1])
            program = i[2]
            if program == 129 and decode_chord:
                root_name = r_h.root_list[i[3]]
                kind_name = r_h.kind_list[i[4]]
                root_pitch_shift = r_h.root_dict[root_name]
                end = start + get_tick(0, e2d(1))
                for kind_shift in magenta._CHORD_KIND_PITCHES[kind_name]:
                    pitch = 36 + root_pitch_shift + kind_shift
                    midi_obj.instruments[1].notes.append(
                        miditoolkit.containers.Note(
                            start=start, end=end, pitch=pitch, velocity=e2v(20)
                        )
                    )
            elif program != 129:
                pitch = i[3] - 128 if program == 128 else i[3]
                if pitch < 0:
                    continue
                duration = get_tick(0, e2d(i[4]))
                if duration == 0:
                    duration = 1
                end = start + duration
                velocity = e2v(i[5])
                midi_obj.instruments[program].notes.append(
                    miditoolkit.containers.Note(
                        start=start, end=end, pitch=pitch, velocity=velocity
                    )
                )
        midi_obj.instruments = [i for i in midi_obj.instruments if len(i.notes) > 0]
        cur_ts = None
        for i in range(len(bar_to_timesig)):
            new_ts = bar_to_timesig[i]
            if new_ts != cur_ts:
                numerator, denominator = e2t(new_ts)
                midi_obj.time_signature_changes.append(
                    miditoolkit.containers.TimeSignature(
                        numerator=numerator,
                        denominator=denominator,
                        time=get_tick(i, 0),
                    )
                )
                cur_ts = new_ts
        cur_tp = None
        for i in range(len(pos_to_tempo)):
            new_tp = pos_to_tempo[i]
            if new_tp != cur_tp:
                tempo = e2b(new_tp)
                midi_obj.tempo_changes.append(
                    miditoolkit.containers.TempoChange(tempo=tempo, time=get_tick(0, i))
                )
                cur_tp = new_tp
        return midi_obj

tuneease/pipeline/encoding.py
Pshah2023/TuneEase
from ..getmusic.utils.midi_config import (
    MAX_TS_DENOMINATOR,
    MAX_NOTES_PER_BAR,
    DURATION_MAX,
    POS_RESOLUTION,
    MIN_TEMPO,
    MAX_TEMPO,
    TEMPO_QUANT,
    VELOCITY_QUANT,
)
import math


class TS:
    ts_dict = dict()
    ts_list = list()

    def __init__(self):
        for i in range(0, MAX_TS_DENOMINATOR + 1):  # 1 ~ 64
            for j in range(1, ((2**i) * MAX_NOTES_PER_BAR) + 1):
                self.ts_dict[(j, 2**i)] = len(self.ts_dict)
                self.ts_list.append((j, 2**i))


def t2e(x):
    ts = TS()
    assert x in ts.ts_dict, "unsupported time signature: " + str(x)
    return ts.ts_dict[x]


def e2t(x):
    ts = TS()
    return ts.ts_list[x]


class Dur:
    dur_enc = list()
    dur_dec = list()

    def __init__(self) -> None:
        for i in range(DURATION_MAX):
            for j in range(POS_RESOLUTION):
                self.dur_dec.append(len(self.dur_enc))
                for k in range(2**i):
                    self.dur_enc.append(len(self.dur_dec) - 1)


def d2e(x):
    dur = Dur()
    return dur.dur_enc[x] if x < len(dur.dur_enc) else dur.dur_enc[-1]


def e2d(x):
    dur = Dur()
    return dur.dur_dec[x] if x < len(dur.dur_dec) else dur.dur_dec[-1]


def time_signature_reduce(numerator, denominator):
    # reduction (when denominator is too large)
    global MAX_TS_DENOMINATOR
    global MAX_NOTES_PER_BAR
    while (
        denominator > 2**MAX_TS_DENOMINATOR
        and denominator % 2 == 0
        and numerator % 2 == 0
    ):
        denominator //= 2
        numerator //= 2
    # decomposition (when length of a bar exceed max_notes_per_bar)
    while numerator > MAX_NOTES_PER_BAR * denominator:
        for i in range(2, numerator + 1):
            if numerator % i == 0:
                numerator //= i
                break
    return numerator, denominator


def v2e(x):
    return x // VELOCITY_QUANT


def e2v(x):
    return (x * VELOCITY_QUANT) + (VELOCITY_QUANT // 2)


def b2e(x):
    x = max(x, MIN_TEMPO)
    x = min(x, MAX_TEMPO)
    x = x / MIN_TEMPO
    e = round(math.log2(x) * TEMPO_QUANT)
    return e


def e2b(x):
    return 2 ** (x / TEMPO_QUANT) * MIN_TEMPO

tuneease/pipeline/encoding_helpers.py
Pshah2023/TuneEase
import miditoolkit
import torch
import numpy as np
from .encoding import MIDI_to_encoding
from .key_chord import TokenHelper, KeyChordDetails
from .encoding_helpers import b2e, t2e, time_signature_reduce
from .presets import inst_to_row, prog_to_abrv, RootKinds
from ..getmusic.utils.midi_config import (
    BAR_MAX,
    DEDUPLICATE,
    SAMPLE_LEN_MAX,
    SAMPLE_OVERLAP_RATE,
)
from .file_helpers import get_midi_dict, timeout, get_hash, lock_set, lock_write, writer

track_name = ["lead", "bass", "drum", "guitar", "piano", "string"]


def create_pos_from_str(str_cmd, pos):
    if str_cmd == "-":
        return pos
    track_cmds = str_cmd.split(";")
    for track_cmd in track_cmds:
        track_id, start_pos, end_pos = track_cmd.split(",")
        if end_pos != "":
            pos[int(track_id) * 2][int(start_pos) : int(end_pos)] = 1
            pos[int(track_id) * 2 + 1][int(start_pos) : int(end_pos)] = 1
        else:
            pos[int(track_id) * 2][int(start_pos) :] = 1
            pos[int(track_id) * 2 + 1][int(start_pos) :] = 1
    return pos


def F(
    filename,
    file_name,
    conditional_tracks=None,
    content_tracks=None,
    condition_inst=None,
    chord_from_single=None,
):
    t_h = TokenHelper()
    k_c_d = KeyChordDetails()
    r_h = RootKinds()
    if filename == "track_generation.py":
        empty_tracks = ~conditional_tracks & ~content_tracks
        conditional_tracks &= ~empty_tracks  # emptied tracks can not be condition
        conditional_tracks = torch.tensor(conditional_tracks).float()
        conditional_tracks = conditional_tracks.view(7, 1).repeat(1, 2).reshape(14, 1)
        empty_tracks = torch.tensor(empty_tracks).float()
        empty_tracks = empty_tracks.view(7, 1).repeat(1, 2).reshape(14, 1)
        midi_obj = miditoolkit.midi.parser.MidiFile(file_name)
        if conditional_tracks[-1]:
            with_chord = True
        else:
            with_chord = False
        encoding, pitch_shift, tpc = MIDI_to_encoding(
            "track_generation.py",
            midi_obj,
            with_chord,
            condition_inst,
            chord_from_single,
        )

        if len(encoding) == 0:
            print("ERROR(BLANK): " + file_name + "\n", end="")
            return None, 0
        bar_index_offset = 0
        figure_size = encoding[-1][0] * k_c_d.pos_in_bar + encoding[-1][1]
        pad_length = 1  # (512 - figure_size % 512)
        figure_size += pad_length
        conditional_bool = conditional_tracks.repeat(1, figure_size)
        empty_pos = empty_tracks.repeat(1, figure_size).type(torch.bool)
        datum = t_h.pad_index * torch.ones(14, figure_size, dtype=float)
        oov = 0
        inv = 0
        chord_list = []
        tempo = b2e(67)
        lead_start = 0
        idx = 0
        while idx != len(encoding) - 1:
            e = encoding[idx]
            bar = e[0]
            pos = e[1]
            inst = e[2]
            pitch = e[3]
            if inst == 80:
                tempo = e[7]
                assert tempo != 0, "bad tempo"
            if e[2] == 129:
                row = inst_to_row[str(inst)]
                r = r_h.root_list[e[3]]
                k = r_h.kind_list[e[4]]
                datum[2 * row][
                    k_c_d.pos_in_bar * bar + pos : k_c_d.pos_in_bar * (bar + 1) + pos
                ] = t_h.tokens_to_ids[r]
                datum[2 * row + 1][
                    k_c_d.pos_in_bar * bar + pos : k_c_d.pos_in_bar * (bar + 1) + pos
                ] = t_h.tokens_to_ids[k]
                idx += 1
                continue
            chord_list = [str(e[3])]
            for f_idx in range(idx + 1, len(encoding)):
                if (
                    (encoding[f_idx][0] == bar)
                    and (encoding[f_idx][1] == pos)
                    and (encoding[f_idx][2] == inst)
                ):
                    if encoding[f_idx][3] != pitch:
                        chord_list.append(str(encoding[f_idx][3]))
                        pitch = encoding[f_idx][3]
                else:
                    break
            idx = max(idx + 1, f_idx)
            dur = e[4]
            if dur == 0:
                continue
            if not (str(inst) in inst_to_row):
                continue
            row = inst_to_row[str(inst)]
            dur = t_h.tokens_to_ids["T" + str(e[4])]  # duration
            chord_string = " ".join(chord_list)
            token = prog_to_abrv[str(inst)] + chord_string
            if token in t_h.tokens_to_ids:
                pitch = t_h.tokens_to_ids[token]
                assert (dur < t_h.pad_index) and (
                    pitch > t_h.pad_index
                ), "pitch index is {} and dur index is {}".format(pitch, dur)
                datum[2 * row][k_c_d.pos_in_bar * bar + pos] = pitch
                datum[2 * row + 1][k_c_d.pos_in_bar * bar + pos] = dur
                inv += 1
            else:
                oov += 1
        datum = torch.where(empty_pos, t_h.empty_index, datum)
        datum = torch.where(
            ((datum != t_h.empty_index).float() * (1 - conditional_bool)).type(
                torch.bool
            ),
            t_h.empty_index + 1,
            datum,
        )
        conditional_bool = conditional_bool[:, :512]
        not_empty_pos = (torch.tensor(np.array(datum)) != t_h.empty_index).float()
        have_cond = True
        for i in range(14):
            if (
                with_chord
                and conditional_tracks[i] == 1
                and (
                    (datum[i] == t_h.pad_index).sum()
                    + (datum[i] == t_h.empty_index).sum()
                )
                == min(512, figure_size)
            ):
                have_cond = False
                break
        return (
            datum.unsqueeze(0),
            torch.tensor(tempo),
            not_empty_pos,
            conditional_bool,
            pitch_shift,
            tpc,
            have_cond,
        )
    elif filename == "position_generation.py":
        midi_obj = miditoolkit.midi.parser.MidiFile(file_name)
        encoding, pitch_shift, tpc = MIDI_to_encoding(
            "position_generation.py", midi_obj
        )
        if len(encoding) == 0:
            print("ERROR(BLANK): " + file_name + "\n", end="")
            return None, 0
        bar_index_offset = 0
        figure_size = max(encoding[-1][0] * k_c_d.pos_in_bar + encoding[-1][1], 512)
        pad_length = 1  # (512 - figure_size % 512)
        figure_size += pad_length
        datum = t_h.pad_index * torch.ones(14, figure_size, dtype=float)
        oov = 0
        inv = 0
        chord_list = []
        tempo = b2e(67)
        lead_start = 0
        idx = 0
        track_set = set()
        while idx != len(encoding) - 1:
            e = encoding[idx]
            bar = e[0]
            pos = e[1]
            inst = e[2]
            pitch = e[3]
            if inst == 80:
                tempo = e[7]
                assert tempo != 0, "bad tempo"
            # assert e[6] == 6
            if e[2] == 129:
                row = inst_to_row[str(inst)]
                r = r_h.root_list[e[3]]
                k = r_h.kind_list[e[4]]
                datum[2 * row][
                    k_c_d.pos_in_bar * bar + pos : k_c_d.pos_in_bar * (bar + 1) + pos
                ] = t_h.tokens_to_ids[r]
                datum[2 * row + 1][
                    k_c_d.pos_in_bar * bar + pos : k_c_d.pos_in_bar * (bar + 1) + pos
                ] = t_h.tokens_to_ids[k]
                idx += 1
                continue
            chord_list = [str(e[3])]
            for f_idx in range(idx + 1, len(encoding)):
                if (
                    (encoding[f_idx][0] == bar)
                    and (encoding[f_idx][1] == pos)
                    and (encoding[f_idx][2] == inst)
                ):
                    if encoding[f_idx][3] != pitch:
                        chord_list.append(str(encoding[f_idx][3]))
                        pitch = encoding[f_idx][3]
                else:
                    break
            idx = max(idx + 1, f_idx)
            dur = e[4]
            if dur == 0:
                continue
            if not (str(inst) in inst_to_row):
                continue
            row = inst_to_row[str(inst)]
            dur = t_h.tokens_to_ids["T" + str(e[4])]  # duration
            chord_string = " ".join(chord_list)
            token = prog_to_abrv[str(inst)] + chord_string
            track_set.add(track_name[prog_to_abrv[str(inst)]])
            if token in t_h.tokens_to_ids:
                pitch = t_h.tokens_to_ids[token]
                assert (dur < t_h.pad_index) and (
                    pitch > t_h.pad_index
                ), "pitch index is {} and dur index is {}".format(pitch, dur)
                datum[2 * row][k_c_d.pos_in_bar * bar + pos] = pitch
                datum[2 * row + 1][k_c_d.pos_in_bar * bar + pos] = dur
                inv += 1
            else:
                oov += 1
        datum[:, -pad_length:] = t_h.empty_index
        print(
            "The music has {} tracks, with {} positions".format(
                track_set, datum.size()[1]
            )
        )
        print("Representation Visualization:")
        print(
            "\t0,1,2,3,4,5,6,7,8,...\n(0)lead\n(1)bass\n(2)drum\n(3)guitar\n(4)piano\n(5)string\n(6)chord"
        )
        print(
            "Example: condition on 100 to 200 position of lead, 300 to 400 position of piano, write command like this:'0,100,200;4,300,400"
        )
        condition_str = input("Input positions you want to condition on:")
        empty_str = input("Input positions you want to empty:")
        empty_pos = torch.zeros_like(datum)
        condition_pos = torch.zeros_like(datum)
        empty_pos = create_pos_from_str(empty_str, empty_pos)
        condition_pos = create_pos_from_str(condition_str, condition_pos)
        datum = torch.where(empty_pos.type(torch.bool), t_h.empty_index, datum)
        datum = torch.where(
            ((datum != t_h.empty_index).float() * (1 - condition_pos)).type(torch.bool),
            t_h.empty_index + 1,
            datum,
        )
        not_empty_pos = (torch.tensor(np.array(datum)) != t_h.empty_index).float()
        return (
            datum.unsqueeze(0),
            torch.tensor(tempo),
            not_empty_pos,
            condition_pos,
            pitch_shift,
            tpc,
        )
    elif filename == "to_oct.py":
        try:
            with timeout(seconds=600):
                midi_obj = miditoolkit.midi.parser.MidiFile(file_name)
            # check abnormal values in parse result
            assert all(
                0 <= j.start < 2**31 and 0 <= j.end < 2**31
                for i in midi_obj.instruments
                for j in i.notes
            ), "bad note time"
            assert all(
                0 < j.numerator < 2**31 and 0 < j.denominator < 2**31
                for j in midi_obj.time_signature_changes
            ), "bad time signature value"
            assert 0 < midi_obj.ticks_per_beat < 2**31, "bad ticks per beat"
        except BaseException as e:
            print("ERROR(PARSE): " + file_name + " " + str(e) + "\n", end="")
            return None
        midi_notes_count = sum(len(inst.notes) for inst in midi_obj.instruments)
        if midi_notes_count == 0:
            print("ERROR(BLANK): " + file_name + "\n", end="")
            return None
        no_empty_tracks = {"80": 0, "32": 0, "128": 0, "25": 0, "0": 0, "48": 0}
        for inst in midi_obj.instruments:
            no_empty_tracks[str(inst.program)] = 1
        if no_empty_tracks["80"] == 0 or sum(no_empty_tracks.values()) <= 1:
            print("ERROR(BAD TRACKS): " + file_name + "\n", end="")
            return False
        try:
            e = MIDI_to_encoding("to_oct.py", midi_obj)
            if len(e) == 0:
                print("ERROR(BLANK): " + file_name + "\n", end="")
                return None
            if len(e) == 1:
                print("ERROR(TEMPO CHANGE): " + file_name + "\n", end="")
                return False
            # if ts_filter:
            allowed_ts = t2e(time_signature_reduce(4, 4))
            if not all(i[6] == allowed_ts for i in e):
                print("ERROR(TSFILT): " + file_name + "\n", end="")
                return None
            if DEDUPLICATE:
                duplicated = False
                dup_file_name = ""
                midi_hash = "0" * 32
                try:
                    midi_hash = get_hash(e)
                except BaseException as e:
                    pass
                lock_set.acquire()
                midi_dict = get_midi_dict()
                if midi_hash in midi_dict:
                    dup_file_name = midi_dict[midi_hash]
                    duplicated = True
                else:
                    midi_dict[midi_hash] = file_name
                lock_set.release()
                if duplicated:
                    print(
                        "ERROR(DUPLICATED): "
                        + midi_hash
                        + " "
                        + file_name
                        + " == "
                        + dup_file_name
                        + "\n",
                        end="",
                    )
                    return None
            output_str_list = []
            sample_step = max(round(SAMPLE_LEN_MAX / SAMPLE_OVERLAP_RATE), 1)
            for p in range(0, len(e), sample_step):
                L = p
                R = min(p + SAMPLE_LEN_MAX, len(e)) - 1
                bar_index_list = [
                    e[i][0] for i in range(L, R + 1) if e[i][0] is not None
                ]
                bar_index_min = 0
                bar_index_max = 0
                if len(bar_index_list) > 0:
                    bar_index_min = min(bar_index_list)
                    bar_index_max = max(bar_index_list)
                # to make bar index start from 0
                bar_index_offset = -bar_index_min
                e_segment = []
                for i in e[L : R + 1]:
                    if i[0] is None or i[0] + bar_index_offset < BAR_MAX:
                        e_segment.append(i)
                    else:
                        break
                tokens_per_note = 8
                output_words = [
                    (
                        "<{}-{}>".format(j, k if j > 0 else k + bar_index_offset)
                        if k is not None
                        else "<unk>"
                    )
                    for i in e_segment
                    for j, k in enumerate(i)
                ]  # tokens_per_note - 1 for append_eos functionality of binarizer in fairseq
                output_str_list.append(" ".join(output_words))
            # no empty
            if not all(
                len(i.split()) > tokens_per_note * 2 - 1 for i in output_str_list
            ):
                print("ERROR(ENCODE): " + file_name + "\n", end="")
                return False
            try:
                lock_write.acquire()
                writer(file_name, output_str_list)
            except BaseException as e:
                print("ERROR(WRITE): " + file_name + " " + str(e) + "\n", end="")
                return False
            finally:
                lock_write.release()
            print("SUCCESS: " + file_name + "\n", end="")
            return True
        except BaseException as e:
            print("ERROR(PROCESS): " + file_name + " " + str(e) + "\n", end="")
            return False
        print("ERROR(GENERAL): " + file_name + "\n", end="")
        return False

tuneease/pipeline/file.py
Pshah2023/TuneEase
import signal
import hashlib
from multiprocessing import Lock, Manager

lock_file = Lock()
lock_write = Lock()
lock_set = Lock()


def get_midi_dict():
    manager = Manager()
    midi_dict = manager.dict()
    return midi_dict


output_file = None


class timeout:
    def __init__(self, seconds=1, error_message="Timeout"):
        self.seconds = seconds
        self.error_message = error_message

    def handle_timeout(self, signum, frame):
        raise TimeoutError(self.error_message)

    def __enter__(self):
        signal.signal(signal.SIGALRM, self.handle_timeout)
        signal.alarm(self.seconds)

    def __exit__(self, exc_type, value, traceback):
        signal.alarm(0)


def get_hash(encoding):
    # add i[4] and i[5] for stricter match
    midi_tuple = tuple((i[2], i[3]) for i in encoding)
    midi_hash = hashlib.md5(str(midi_tuple).encode("ascii")).hexdigest()
    return midi_hash


def writer(file_name, output_str_list):
    # note: parameter "file_name" is reserved for patching
    with open(output_file, "a") as f:
        for output_str in output_str_list:
            if (
                "<2-80>" in output_str
                and "<2-129>" in output_str
                and (
                    ("<2-0>" in output_str)
                    or ("<2-25>" in output_str)
                    or ("<2-32>" in output_str)
                    or ("<2-48>" in output_str)
                    or ("<2-128>" in output_str)
                )
            ):
                f.write(output_str + "\n")

tuneease/pipeline/file_helpers.py
Pshah2023/TuneEase
class Item(object):
    def __init__(self, name, start, end, vel=0, pitch=0, track=0, value=""):
        self.name = name
        self.start = start  # start step
        self.end = end  # end step
        self.vel = vel
        self.pitch = pitch
        self.track = track
        self.value = value

    def __repr__(self):
        return (
            f"Item(name={self.name:>10s}, start={self.start:>4d}, end={self.end:>4d}, "
            f"vel={self.vel:>3d}, pitch={self.pitch:>3d}, track={self.track:>2d}, "
            f"value={self.value:>10s})\n"
        )

    def __eq__(self, other):
        return (
            self.name == other.name
            and self.start == other.start
            and self.pitch == other.pitch
            and self.track == other.track
        )

tuneease/pipeline/item.py
Pshah2023/TuneEase
import numpy as np
import os
import pickle
from ..getmusic.utils.magenta_chord_recognition import Magenta
from ..pipeline.config import Config


class TokenHelper:
    """
    Initialize a TokenHelper instance.

    This constructor method updates the class-level instance variables based on a vocabulary file specified config class.

    Args:
        None

    Returns:
        None
    """

    tokens_to_ids = {}
    ids_to_tokens = []
    pad_index = None
    empty_index = None

    def __init__(self) -> None:
        self.updateClass(TokenHelper)

    @staticmethod
    def updateClass(cls):
        """
        Update class-level instance variables with token information.

        Args:
            cls (class): The class instance to be updated.

        Returns:
            None

        Example:
            >>> token_helper = TokenHelper()
            >>> token_helper.tokens_to_ids  # Updated with token mappings
            {'<token1>': 0, '<token2>': 1, ...}
        """
        if not bool(len(cls.tokens_to_ids)):
            config = Config().config
            with open(config["solver"]["vocab_path"], "r") as f:
                tokens = f.readlines()

                for id, token in enumerate(tokens):
                    token, freq = token.strip().split("\t")
                    cls.tokens_to_ids[token] = id
                cls.ids_to_tokens = list(cls.tokens_to_ids.keys())
                cls.pad_index = cls.tokens_to_ids["<pad>"]
                cls.empty_index = len(cls.ids_to_tokens)


class KeyChordDetails:
    pos_in_bar = None
    key_profile_path = None
    key_profile = None
    chord_pitch_out_of_key_prob = None
    chord_change_prob = None
    key_change_prob = None
    key_chord_distribution = None
    key_chord_loglik = None
    key_chord_transition_distribution = None
    beat_note_factor = 4  # In MIDI format a note is always 4 beats
    max_notes_per_bar = 1  # 1/64 ... 128/64 #
    pos_resolution = 4  # 16  # per beat (quarter note)
    bar_max = 32
    figure_size = None

    def __init__(self) -> None:
        if not bool(self.pos_in_bar):
            magenta = Magenta()
            self.pos_in_bar = (
                self.beat_note_factor * self.max_notes_per_bar * self.pos_resolution
            )
            self.figure_size = self.bar_max * self.pos_in_bar
            self.key_profile_path = os.path.join(
                os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                "getmusic",
                "utils",
                "key_profile.pickle",
            )
            self.key_profile = pickle.load(open(self.key_profile_path, "rb"))

            self.chord_pitch_out_of_key_prob = 0.01
            self.key_change_prob = 0.001
            self.chord_change_prob = 0.5
            self.key_chord_distribution = magenta._key_chord_distribution(
                chord_pitch_out_of_key_prob=self.chord_pitch_out_of_key_prob
            )
            self.key_chord_loglik = np.log(self.key_chord_distribution)
            self.key_chord_transition_distribution = (
                magenta._key_chord_transition_distribution(
                    self.key_chord_distribution,
                    key_change_prob=self.key_change_prob,
                    chord_change_prob=self.chord_change_prob,
                )
            )
            self.key_chord_transition_loglik = np.log(
                self.key_chord_transition_distribution
            )

tuneease/pipeline/key_chord.py
Pshah2023/TuneEase
prog_to_abrv = {
    "80": "M",
    "32": "B",
    "128": "D",
    "25": "G",
    "0": "P",
    "48": "S",
}
inst_to_row = {"80": 0, "32": 1, "128": 2, "25": 3, "0": 4, "48": 5, "129": 6}
track_name = {
    "M": "lead",
    "B": "bass",
    "D": "drum",
    "G": "guitar",
    "P": "piano",
    "S": "string",
}
prog_to_name = {
    "32": "Bass",
    "128": "Drums",
    "25": "Guitar",
    "0": "Grand_Piano",
    "48": "Strings",
}


class RootKinds:
    root_dict = None
    kind_dict = None
    root_list = None
    kind_list = None

    def __init__(self) -> None:
        self.root_dict = {
            "C": 0,
            "C#": 1,
            "D": 2,
            "Eb": 3,
            "E": 4,
            "F": 5,
            "F#": 6,
            "G": 7,
            "Ab": 8,
            "A": 9,
            "Bb": 10,
            "B": 11,
        }
        self.kind_dict = {
            "null": 0,
            "m": 1,
            "+": 2,
            "dim": 3,
            "seven": 4,
            "maj7": 5,
            "m7": 6,
            "m7b5": 7,
        }
        self.root_list = list(self.root_dict.keys())
        self.kind_list = list(self.kind_dict.keys())

tuneease/pipeline/presets.py
Pshah2023/TuneEase
import numpy as np
from .key_chord import KeyChordDetails


def normalize_to_c_major(filename, e):
    k_c_d = KeyChordDetails()
    if filename == "track_generation.py":

        def get_pitch_class_histogram(
            notes, use_duration=True, use_velocity=True, normalize=True
        ):
            weights = np.ones(len(notes))
            if use_duration:
                weights *= [note[4] for note in notes]
            if use_velocity:
                weights *= [note[5] for note in notes]
            histogram, _ = np.histogram(
                [note[3] % 12 for note in notes],
                bins=np.arange(13),
                weights=weights,
                density=normalize,
            )
            if normalize:
                histogram /= histogram.sum() + (histogram.sum() == 0)
            return histogram

        pitch_histogram = [i for i in e if i[2] < 128]
        if len(pitch_histogram) == 0:
            return e, True, 0

        histogram = get_pitch_class_histogram(pitch_histogram)
        key_candidate = np.dot(k_c_d.key_profile, histogram)
        key_temp = np.where(key_candidate == max(key_candidate))
        major_index = key_temp[0][0]
        minor_index = key_temp[0][1]
        major_count = histogram[major_index]
        minor_count = histogram[minor_index % 12]
        key_number = 0
        if major_count < minor_count:
            key_number = minor_index
            is_major = False
        else:
            key_number = major_index
            is_major = True
        real_key = key_number
        # transposite to C major or A minor
        if real_key <= 11:
            trans = 0 - real_key
        else:
            trans = 21 - real_key
        pitch_shift = trans

        e = [
            tuple(
                k + pitch_shift if j == 3 and i[2] != 128 else k
                for j, k in enumerate(i)
            )
            for i in e
        ]
        return e, is_major, pitch_shift
    elif filename == "position_generation.py":

        def get_pitch_class_histogram(
            notes, use_duration=True, use_velocity=True, normalize=True
        ):
            weights = np.ones(len(notes))
            # Assumes that duration and velocity have equal weight
            if use_duration:
                weights *= [note[4] for note in notes]  # duration
            if use_velocity:
                weights *= [note[5] for note in notes]  # velocity
            histogram, _ = np.histogram(
                [note[3] % 12 for note in notes],
                bins=np.arange(13),
                weights=weights,
                density=normalize,
            )
            if normalize:
                histogram /= histogram.sum() + (histogram.sum() == 0)
            return histogram

        histogram = get_pitch_class_histogram([i for i in e if i[2] < 128])
        key_candidate = np.dot(k_c_d.key_profile, histogram)
        key_temp = np.where(key_candidate == max(key_candidate))
        major_index = key_temp[0][0]
        minor_index = key_temp[0][1]
        major_count = histogram[major_index]
        minor_count = histogram[minor_index % 12]
        key_number = 0
        if major_count < minor_count:
            key_number = minor_index
            is_major = False
        else:
            key_number = major_index
            is_major = True
        real_key = key_number
        # transposite to C major or A minor
        if real_key <= 11:
            trans = 0 - real_key
        else:
            trans = 21 - real_key
        pitch_shift = trans

        e = [
            tuple(
                k + pitch_shift if j == 3 and i[2] != 128 else k
                for j, k in enumerate(i)
            )
            for i in e
        ]
        return e, is_major, pitch_shift
    elif filename == "to_oct.py":

        def get_pitch_class_histogram(
            notes, use_duration=True, use_velocity=True, normalize=True
        ):
            weights = np.ones(len(notes))
            # Assumes that duration and velocity have equal weight
            if use_duration:
                weights *= [note[4] for note in notes]  # duration
            if use_velocity:
                weights *= [note[5] for note in notes]  # velocity
            histogram, _ = np.histogram(
                [note[3] % 12 for note in notes],
                bins=np.arange(13),
                weights=weights,
                density=normalize,
            )
            if normalize:
                histogram /= histogram.sum() + (histogram.sum() == 0)
            return histogram

        histogram = get_pitch_class_histogram([i for i in e if i[2] < 128])
        key_candidate = np.dot(k_c_d.key_profile, histogram)
        key_temp = np.where(key_candidate == max(key_candidate))
        major_index = key_temp[0][0]
        minor_index = key_temp[0][1]
        major_count = histogram[major_index]
        minor_count = histogram[minor_index % 12]
        key_number = 0
        if major_count < minor_count:
            key_number = minor_index
            is_major = False
        else:
            key_number = major_index
            is_major = True
        real_key = key_number
        # transposite to C major or A minor
        if real_key <= 11:
            trans = 0 - real_key
        else:
            trans = 21 - real_key
        pitch_shift = trans

        # _e = []
        # for i in e:
        #     for j, k in enumerate(i):
        #         if i[2] == 128:
        #             _e.append(i)
        #         else:

        e = [
            tuple(
                k + pitch_shift if j == 3 and i[2] != 128 else k
                for j, k in enumerate(i)
            )
            for i in e
        ]
        return e, is_major

tuneease/pipeline/processing.py
Pshah2023/TuneEase
import os
import warnings
import torch
from .getmusic.modeling.build import build_model
from .getmusic.utils.misc import merge_opts_to_config, modify_config_for_debug
from .pipeline.config import Config
from .logger import ServerLogger
from getmusic.engine.solver import Solver
from getmusic.distributed.launch import launch
from .pipeline.args import get_args
from .pipeline.encoding import encoding_to_MIDI
from .pipeline.file import F


def main_worker(local_rank, args):
    args.local_rank = local_rank

    args.global_rank = args.local_rank + args.node_rank * args.ngpus_per_node
    args.distributed = args.world_size > 1

    config = Config()
    config = config.config
    config = merge_opts_to_config(config, args.opts)
    if args.debug:
        config = modify_config_for_debug(config)

    logger = ServerLogger("tuneease.log").get()
    logger.debug(str(args))

    model = build_model(config, args)

    if args.sync_bn:
        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)

    dataloader_info = None

    solver = Solver(
        config=config,
        args=args,
        model=model,
        dataloader=dataloader_info,
        logger=logger,
        is_sample=True,
    )

    assert args.load_path is not None
    solver.resume(path=args.load_path)

    file_list = [
        os.path.join(args.file_path, n)
        for n in os.listdir(args.file_path)
        if (n[-4:].lower() == ".mid" or n[-5:].lower() == ".midi")
        and ("iter" not in n.lower())
    ]
    file_list.sort()

    for file_name in file_list:
        print(file_name)

        if ".pth" in file_name:
            continue
        y = input("skip?")
        if "y" in y:
            continue

        x, tempo, not_empty_pos, condition_pos, pitch_shift, tpc = F(
            "position_generation.py", file_name
        )

        oct_line = solver.infer_sample(
            x,
            tempo,
            not_empty_pos,
            condition_pos,
            use_ema=args.no_ema,
            skip_step=args.skip_step,
        )

        data = oct_line.split(" ")

        oct_final_list = []
        for start in range(3, len(data), 8):
            if "pad" not in data[start] and "pad" not in data[start + 1]:
                pitch = int(data[start][:-1].split("-")[1])
                if data[start - 1] != "<2-129>" and data[start - 1] != "<2-128>":
                    pitch -= pitch_shift
                data[start] = "<3-{}>".format(pitch)  # re-normalize
                oct_final_list.append(" ".join(data[start - 3 : start + 5]))

        oct_final = " ".join(oct_final_list)

        midi_obj = encoding_to_MIDI(
            "position_generation.py", oct_final, tpc, args.decode_chord
        )

        save_path = os.path.join(
            args.file_path, "position-{}".format(file_name.split("/")[-1])
        )

        midi_obj.dump(save_path)


def main():
    args = get_args()

    if args.gpu is not None:
        warnings.warn(
            "You have chosen a specific GPU. This will completely disable ddp."
        )
        torch.cuda.set_device(args.gpu)
        args.ngpus_per_node = 1
        args.world_size = 1
    else:
        print("args.num_node ", args.num_node)
        if args.num_node == 1:
            args.dist_url == "auto"
        else:
            assert args.num_node > 1

        args.ngpus_per_node = torch.cuda.device_count()
        args.world_size = args.ngpus_per_node * args.num_node  #

    launch(
        main_worker,
        args.ngpus_per_node,
        args.num_node,
        args.node_rank,
        args.dist_url,
        args=(args,),
    )


if __name__ == "__main__":
    main()

tuneease/position_generation.py
Pshah2023/TuneEase
from ..getmusic.data.indexed_datasets import IndexedDatasetBuilder
from ..pipeline.presets import prog_to_abrv, inst_to_row, RootKinds
from ..pipeline.key_chord import TokenHelper, KeyChordDetails

import multiprocessing as mp
from tqdm import tqdm
import numpy as np
import sys
import os

sys.path.append("/".join(os.path.abspath(__file__).split("/")[:-2]))


def oct_to_rep(line):
    t_h = TokenHelper()
    k_c_d = KeyChordDetails()
    r_h = RootKinds()

    no_empty_tracks = {"80": 0, "32": 0, "128": 0, "25": 0, "0": 0, "48": 0}
    main_num = line.count("<2-80>")
    string_num = line.count("<2-48>")
    drum_num = line.count("<2-128>")

    oov = 0
    inv = 0

    if "<2-80>" not in line:
        return None

    f = (
        ("<2-0>" in line)
        + ("<2-25>" in line)
        + ("<2-32>" in line)
        + ("<2-48>" in line)
        + ("<2-128>" in line)
    )

    if f == 0:
        return None

    tmp = line.strip().split("<0-")[1:]
    encoding = []
    for item in tmp:
        tmp2 = item.strip()[:-1].split("> <")
        encoding.append([int(tmp2[0])] + [int(i[2:]) for i in tmp2[1:]])

    chord_list = []

    datum = t_h.pad_index * np.ones([14, 1 + k_c_d.figure_size], dtype=float)

    idx = 0
    while idx != len(encoding) - 1:
        e = encoding[idx]

        bar = e[0]
        pos = e[1]
        inst = e[2]
        pitch = e[3]
        tempo = e[7]

        assert e[6] == 6  # ts = 4/4

        if e[2] == 129:
            row = inst_to_row[str(inst)]
            r = r_h.root_list[e[3]]
            k = r_h.kind_list[e[4]]
            datum[2 * row][
                k_c_d.pos_in_bar * bar + pos : k_c_d.pos_in_bar * (bar + 1) + pos
            ] = tokens_to_ids[r]
            datum[2 * row + 1][
                k_c_d.pos_in_bar * bar + pos : k_c_d.pos_in_bar * (bar + 1) + pos
            ] = tokens_to_ids[k]
            idx += 1
            continue

        if tempo == 0:
            return None

        dur = e[4]
        if e[4] == 0:
            e[4] = 1

        chord_list = [str(e[3])]

        for f_idx in range(idx + 1, len(encoding)):
            if (
                (encoding[f_idx][0] == bar)
                and (encoding[f_idx][1] == pos)
                and (encoding[f_idx][2] == inst)
            ):
                if encoding[f_idx][3] != pitch:
                    chord_list.append(str(encoding[f_idx][3]))
                    pitch = encoding[f_idx][3]
            else:
                break

        idx = max(idx + 1, f_idx)

        row = inst_to_row[str(inst)]
        dur = tokens_to_ids["T" + str(e[4])]

        chord_string = " ".join(chord_list)
        token = prog_to_abrv[str(inst)] + chord_string

        if token in tokens_to_ids:
            pitch = tokens_to_ids[token]
            no_empty_tracks[str(e[2])] = 1
            assert (dur < t_h.pad_index) and (
                pitch > t_h.pad_index
            ), "pitch index is {} and dur index is {}".format(pitch, dur)
            datum[2 * row][k_c_d.pos_in_bar * bar + pos] = pitch
            datum[2 * row + 1][k_c_d.pos_in_bar * bar + pos] = dur
            inv += 1
        else:
            oov += 1

    if sum(no_empty_tracks.values()) > 1 and no_empty_tracks["80"] == 1:  # 伴奏 + 主旋律
        assert (datum == t_h.pad_index).sum() != 14 * (
            k_c_d.pos_in_bar * k_c_d.bar_max + 1
        )

        pad_in_a_row = (datum == t_h.pad_index).sum(-1)
        for row_id in range(6):
            if pad_in_a_row[2 * row_id] == k_c_d.figure_size + 1:
                datum[2 * row_id] = t_h.empty_index
                datum[2 * row_id + 1] = t_h.empty_index

        datum[:, k_c_d.pos_in_bar * bar + pos + 1 :] = t_h.empty_index
        datum[:, -1] = tempo

        return datum
    else:
        return None


if __name__ == "__main__":
    tokens_to_ids = {}

    oct_path = sys.argv[2]
    binary_data_dir = sys.argv[3]
    os.makedirs(binary_data_dir, exist_ok=True)

    num_valid_split = 10

    with open(oct_path, "r") as f:
        all_midi_fns = f.readlines()

    for split in ["valid", "train"]:
        p = mp.Pool(int(os.getenv("N_PROC", os.cpu_count())))
        futures = []
        midi_fns = (
            all_midi_fns[num_valid_split:]
            if split == "train"
            else all_midi_fns[:num_valid_split]
        )
        print(f"| #{split} set: {len(midi_fns)}")
        ds_builder = IndexedDatasetBuilder(f"{binary_data_dir}/{split}")
        for midi_fn in midi_fns:
            futures.append(p.apply_async(oct_to_rep, args=[midi_fn]))
        p.close()

        num = 0
        for f in tqdm(futures):
            item = f.get()
            if item is None:
                continue
            num += 1
            ds_builder.add_item(item)
        ds_builder.finalize()
        p.join()
        np.save(f"{binary_data_dir}/{split}_length.npy".format(split), num)
        print("{} set has {} reps".format(split, num))

tuneease/preprocess/binarize.py
Pshah2023/TuneEase
#!PS Commenting this out for now. This also just wouldn't work on Windows.

# from tqdm import tqdm
# import sys
# import os
# from ..getmusic.utils.midi_config import duration_max, pos_resolution
# from ..pipeline.presets import prog_to_name, RootKinds
# sys.path.append('/'.join(os.path.abspath(__file__).split('/')[:-2]))

# dur_vocab_size = duration_max * pos_resolution
# # 9 celesta
# # 25 acoustic guitar(steel)
# # 32 acoustic bass
# # 48 string ensemble
# # 80 synth
# # 128 drum

# dpath = sys.argv[1]
# output_txt_path = dpath
# num_threshold = int(sys.argv[2])
# oct_path = os.path.join(dpath,'oct.txt')

# with open(oct_path,'r') as f:
#     lines = f.readlines()
#     token_dict = {'32':{},'128':{},'25':{},'0':{},'48':{}}
#     for line in tqdm(lines):
#         if '<2-80>' not in line:
#             continue

#         track_numbers = ('<2-0>' in line) + ('<2-25>' in line) + ('<2-32>' in line) + ('<2-48>' in line) + ('<2-128>' in line)

#         if track_numbers == 0:
#             continue

#         tmp = line.strip().split('<0-')[1:]
#         encoding = []
#         for item in tmp:
#             tmp2 = item.strip()[:-1].split('> <')
#             encoding.append([int(tmp2[0])] + [int(i[2:]) for i in tmp2[1:]])

#         idx = 0
#         while idx != len(encoding) - 1:
#             e = encoding[idx]

#             if e[2] == 80 or e[2] == 129 or e[2] not in [0,25,32,48,80,128]:
#                 idx += 1
#                 continue

#             bar = e[0]
#             pos = e[1]
#             inst = e[2]
#             pitch = e[3]
#             dur = e[4]

#             chord_list = [str(e[3])]
#             dur_list = [e[4]]

#             for f_idx in range(idx + 1, len(encoding)):
#                 if (encoding[f_idx][0] == bar) and (encoding[f_idx][1] == pos) and (encoding[f_idx][2] == inst):
#                     if encoding[f_idx][3] != pitch:
#                         chord_list.append(str(encoding[f_idx][3]))
#                         dur_list.append(encoding[f_idx][4])
#                         pitch = encoding[f_idx][3]
#                 else:
#                     break

#             if len(chord_list) > 1:
#                 chord_str = ' '.join(chord_list)
#                 if chord_str in token_dict[str(inst)]:
#                     token_dict[str(inst)][chord_str] += 1
#                 else:
#                     token_dict[str(inst)][chord_str] = 1

#             idx = max(idx + 1, f_idx)

# output_txt_path = os.path.join(output_txt_path, 'pitch_dict.txt')
# w = open(output_txt_path,'w')

# for j in range(1, dur_vocab_size):
#     w.write('T'+str(j)+'\t'+'0\n') # duration
# w.write('<pad>\t0\n')

# for j in range(0,128):
#     w.write('M'+str(j)+'\t'+'0\n') # pitch

# start_ls = [dur_vocab_size]
# end_ls = [dur_vocab_size + 127]

# start_idx = dur_vocab_size + 128
# for idx, inst_program in enumerate(token_dict):
#     start_ls.append(start_idx)
#     threshold = sum(token_dict[inst_program].values()) * 0.9
#     inst_list_of_tuples = [(k, v) for k, v in token_dict[inst_program].items()]

#     inst_list_of_tuples = sorted(inst_list_of_tuples, key=lambda tup: tup[1], reverse=True)

#     if inst_program == '128':
#         for j in range(128,256):
#             w.write('D'+str(j)+'\t'+'0\n') # pitch
#     elif inst_program != '0':
#         for j in range(0,128):
#             w.write(prog_to_name[inst_program][0]+str(j)+'\t'+'0\n') # pitch
#     else:
#         for j in range(0,128):
#             w.write('P'+str(j)+'\t'+'0\n') # pitch

#     chord_sum = 0
#     for idx,j in enumerate(inst_list_of_tuples):
#         chord_sum += j[1]
#         if chord_sum < threshold and j[1] > num_threshold:
#             if inst_program != '0':
#                 w.write(prog_to_name[inst_program][0] + j[0] + '\t' + str(j[1]) + '\n')
#             else:
#                 w.write('P' + j[0] + '\t' + str(j[1]) + '\n')
#         else:
#             print(prog_to_name[inst_program])
#             print("total tokens:",int(1.1 * threshold))
#             print('we keep {}, which is {} ({}%) tokens'.format(chord_sum, idx, chord_sum / (1.1 * threshold)))
#             break
#     start_idx += 128 + idx - 1
#     end_ls.append(start_idx)
#     start_idx += 1

# print(start_ls)
# print(end_ls)
# for r in root_dict:
#     w.write(r+'\t'+'0\n')

# for k in kind_dict:
#     w.write(k+'\t'+'0\n')

# w.close()

tuneease/preprocess/make_dict.py
Pshah2023/TuneEase
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
#

from multiprocessing import Pool
import os
import sys
from ..getmusic.utils.midi_config import (
    BAR_MAX,
    MAX_NOTES_PER_BAR,
    BEAT_NOTE_FACTOR,
    POS_RESOLUTION,
    MAX_INST,
    MAX_PITCH,
    DURATION_MAX,
    MAX_TEMPO,
    MAX_VELOCITY,
    SAMPLE_LEN_MAX,
    POOL_NUM,
)
from ..pipeline.encoding_helpers import TS, v2e, b2e
from ..pipeline.file import F

sys.path.append("/".join(os.path.abspath(__file__).split("/")[:-2]))

data_zip = None


# (0 Measure, 1 Pos, 2 Program, 3 Pitch, 4 Duration, 5 Velocity, 6 TimeSig, 7 Tempo)
# (Measure, TimeSig)
# (Pos, Tempo)
# Percussion: Program=128 Pitch=[128,255]
def gen_dictionary(file_name):
    ts = TS()
    num = 0
    with open(file_name, "w") as f:
        for j in range(BAR_MAX):
            print("<0-{}>".format(j), num, file=f)
        for j in range(BEAT_NOTE_FACTOR * MAX_NOTES_PER_BAR * POS_RESOLUTION):
            print("<1-{}>".format(j), num, file=f)
        for j in range(MAX_INST + 1 + 1):
            # max_inst + 1 for percussion
            print("<2-{}>".format(j), num, file=f)
        for j in range(2 * MAX_PITCH + 1 + 1):
            # max_pitch + 1 ~ 2 * max_pitch + 1 for percussion
            print("<3-{}>".format(j), num, file=f)
        for j in range(DURATION_MAX * POS_RESOLUTION):
            print("<4-{}>".format(j), num, file=f)
        for j in range(v2e(MAX_VELOCITY) + 1):
            print("<5-{}>".format(j), num, file=f)
        for j in range(len(ts.ts_list)):
            print("<6-{}>".format(j), num, file=f)
        for j in range(b2e(MAX_TEMPO) + 1):
            print("<7-{}>".format(j), num, file=f)


def G(file_name):
    try:
        return F("to_oct.py", file_name)
    except BaseException as e:
        print("ERROR(UNCAUGHT): " + file_name + "\n", end="")
        return False


def str_to_encoding(s):
    encoding = [int(i[3:-1]) for i in s.split() if "s" not in i]
    tokens_per_note = 8
    assert len(encoding) % tokens_per_note == 0
    encoding = [
        tuple(encoding[i + j] for j in range(tokens_per_note))
        for i in range(0, len(encoding), tokens_per_note)
    ]
    return encoding


def encoding_to_str(e):
    bar_index_offset = 0
    p = 0
    tokens_per_note = 8
    return " ".join(
        (
            [
                "<{}-{}>".format(j, k if j > 0 else k + bar_index_offset)
                for i in e[p : p + SAMPLE_LEN_MAX]
                if i[0] + bar_index_offset < BAR_MAX
                for j, k in enumerate(i)
            ]
        )
    )  # 8 - 1 for append_eos functionality of binarizer in fairseq


if __name__ == "__main__":
    data_path = sys.argv[1]
    prefix = sys.argv[2]
    if os.path.exists(prefix):
        print(
            "Output path {} already exists! Please delete existing files".format(prefix)
        )
        sys.exit(0)
    os.system("mkdir -p {}".format(prefix))
    file_list = []
    for dirpath, dirnames, filenames in os.walk(data_path):
        for filename in filenames:
            if filename.endswith(".mid"):
                file_list.append(os.path.join(dirpath, filename))
    file_list.sort()
    # random.shuffle(file_list)
    # gen_dictionary('{}/dict.txt'.format(prefix))
    ok_cnt = 0
    all_cnt = 0
    total_file_cnt = len(file_list)
    file_list_split = file_list
    output_file = "{}/oct.txt".format(prefix)
    with Pool(POOL_NUM) as p:
        result = list(p.imap_unordered(G, file_list_split))
        all_cnt += sum((1 if i is not None else 0 for i in result))
        ok_cnt += sum((1 if i is True else 0 for i in result))
    output_file = None
    print(
        "{}/{} ({:.2f}%) MIDI files successfully processed".format(
            ok_cnt, all_cnt, ok_cnt / all_cnt * 100
        )
    )

tuneease/preprocess/to_oct.py
Pshah2023/TuneEase
import os
from flask import Flask, request, send_file, send_from_directory
from flask_cors import CORS
import music21
import random
from .logger import ServerLogger
from .pathutility import PathUtility
from .converter import Converter
from .track_generation import main as music_generator

util = PathUtility()
logger = ServerLogger().get()
static_path = os.path.join(util.tuneease_dir(), "build", "static")
logger.debug(static_path)
app = Flask(__name__, static_url_path="/static", static_folder=static_path)
app.config["museScore_path"] = util.musescore_path()
CORS(app)


@app.route("/")
def index():
    util = PathUtility()
    build_dir = os.path.join(util.tuneease_dir(), "build")
    return send_from_directory(build_dir, "index.html")


@app.route("/<path:filename>")
def serve_file(filename):
    util = PathUtility()
    build_dir = os.path.join(util.tuneease_dir(), "build")
    logger = ServerLogger("tuneease.log").get()
    logger.info("Received a request " + build_dir)
    return send_from_directory(build_dir, filename)


@app.route("/convert", methods=["POST"])
def convert():
    logger = ServerLogger("tuneease.log").get()
    logger.info("Received a /convert request")
    if app.config["museScore_path"]:
        converter = Converter(museScore_path=app.config["museScore_path"])
        try:
            file = request.files["file"]
            output_filepath = converter.convert(file)
            return send_file(
                output_filepath, as_attachment=True, download_name=output_filepath
            )
        except Exception as e:
            return str(e), 500
    else:
        file = request.files["file"]
        assert file.filename != "", "File is empty"
        filetail = os.path.split(file.filename)[-1]
        filepath = os.path.join(util.project_directory(), "temp", filetail)
        file.save(filepath)
        fileext = os.path.splitext(filepath)[0]
        musicxml_file = music21.converter.parse(filepath)
        try:
            musicxml_file.write("musicxml", fp=fileext + ".xml")
        except Exception as E:
            logger.error(
                "Had an error converting the generated content to xml (likely some error with the input file)"
                + str(E)
            )
            return 500
        return send_file(
            fileext + ".xml", as_attachment=True, download_name=fileext + ".xml"
        )


@app.route("/split", methods=["POST"])
def split():
    logger = ServerLogger("tuneease.log").get()
    logger.info(
        "Received a /split request " + str(request.files) + " with " + str(request.form)
    )
    try:
        file = request.files["file"]
        assert file.filename != "", "File is empty"
        assert os.path.splitext(file.filename)[1] == ".xml", "File is not .xml"
        start_number = int(request.form["start"])
        if (end_number := request.form.get("end")) is not None:
            end_number = int(end_number)
        else:
            end_number = ""
    except Exception as e:
        logger.error(f"Failed to start /split. Error message: {str(e)}")
        return "Possible Clientside error", 500
    util = PathUtility()
    filepath = os.path.join(util.project_directory(), "temp", file.filename)
    measure_filepath = os.path.join(
        util.project_directory(),
        "temp",
        os.path.splitext(file.filename)[0]
        + "-measure-"
        + str(start_number)
        + "-"
        + str(end_number),
    )
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    file.save(filepath)
    try:
        parsed_song = music21.converter.parse(filepath)
        measures = []
        for part in parsed_song.parts:
            measures.extend(part.getElementsByClass("Measure"))
        try:
            if end_number == "":
                stream = measures[start_number].write(
                    "musicxml", measure_filepath + ".xml"
                )
            else:
                measures_to_combine = measures[start_number - 1 : end_number]
                stream = music21.stream.Stream()
                for measure in measures_to_combine:
                    stream.append(measure)
        except Exception as E:
            logger.error(
                f"Your input numbers were outside of the measures of the part (probably): {str(e)}"
            )
            return (
                f"Your input numbers were outside of the measures of the part (probably)",
                500,
            )
        stream.write("musicxml", measure_filepath + ".xml")
        return send_file(
            measure_filepath + ".xml",
            as_attachment=True,
            download_name=measure_filepath + ".xml",
        )
    except Exception as e:
        logger.error(f"Error: {str(e)}")
        return f"Error", 500


@app.route("/number", methods=["POST"])
def number():
    logger = ServerLogger("tuneease.log").get()
    logger.info("Received a /number request " + str(request.files))
    try:
        file = request.files["file"]
        assert file.filename != "", "File is empty"
        assert os.path.splitext(file.filename)[1] == ".xml", "File is not .musixml"
    except Exception as e:
        logger.error(f"Failed to start /number. Error message: {str(e)}")
        return "Possible Clientside error", 500

    util = PathUtility()
    filepath = os.path.join(util.project_directory(), "temp", file.filename)
    output_filepath = os.path.join(
        util.project_directory(),
        "temp",
        os.path.splitext(file.filename)[0] + "-number.xml",
    )
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    file.save(filepath)

    try:
        score = music21.converter.parse(filepath)
        for part in score.parts:
            for i, measure in enumerate(
                part.getElementsByClass(music21.stream.Measure)
            ):
                text_expression = music21.expressions.TextExpression(f"Measure {i + 1}")
                text_expression.style.absoluteY = 20  # Adjust for positioning
                text_expression.style.justify = "left"
                measure.insert(0, text_expression)
        score.write("musicxml", output_filepath)
        return send_file(
            output_filepath, as_attachment=True, download_name=output_filepath
        )
    except Exception as e:
        logger.error(f"Error: {str(e)}")
        return f"Error", 500


@app.route("/random", methods=["GET"])
def random_score():
    logger = ServerLogger("tuneease.log").get()
    logger.info("Received a /random request")
    try:
        time_signature = request.form.get("time_signature", "4/4")
        quarter_length = 1.0 / (int(time_signature.split("/")[1]) / 4)
        number_measures = int(request.form.get("number_measures", "8")) * int(
            time_signature.split("/")[0]
        )
        s = music21.stream.Stream()
        ts = music21.meter.TimeSignature(time_signature)
        s.append(ts)
        md = music21.metadata.Metadata()
        md.title = "Random"
        md.composer = "MuseTune"
        s.insert(0, md)
        notes = ["A4", "B4", "C4", "D4", "E4", "F4", "G4"]
        lyrics = ["✔", "x"]
        for i in range(number_measures):
            note_name = random.choice(notes)
            n = music21.note.Note(note_name, quarterLength=quarter_length)
            n.lyric = random.choice(lyrics)
            s.append(n)

        util = PathUtility()
        filepath = os.path.join(util.project_directory(), "temp", "random.xml")
        os.makedirs(os.path.dirname(filepath), exist_ok=True)

        s.write("musicxml", filepath)
        return send_file(filepath, as_attachment=True, download_name=filepath)
    except Exception as e:
        logger.error(f"Error: {str(e)}")
        return f"Error", 500


@app.route("/generate", methods=["POST", "GET"])
def generate():
    """Do note the lowercase in the boolean values!"""
    logger = ServerLogger("tuneease.log").get()
    logger.info("Received a /generate request with " + str(request.form))
    conditional_tracks = "".join(
        [
            str(int(request.form.get("condition-lead", "false") == "true")),
            str(int(request.form.get("condition-bass", "false") == "true")),
            str(int(request.form.get("condition-drums", "false") == "true")),
            str(int(request.form.get("condition-guitar", "false") == "true")),
            str(int(request.form.get("condition-piano", "false") == "true")),
            str(int(request.form.get("condition-strings", "false") == "true")),
            str(int(True)),
        ]
    )
    content_tracks = "".join(
        [
            str(int(request.form.get("content-lead", "false") == "true")),
            str(int(request.form.get("content-bass", "false") == "true")),
            str(int(request.form.get("content-drums", "false") == "true")),
            str(int(request.form.get("content-guitar", "false") == "true")),
            str(int(request.form.get("content-piano", "true") == "true")),
            str(int(request.form.get("content-strings", "false") == "true")),
            str(int(False)),
        ]
    )
    seed = request.form.get("seed", "0")

    if "file" in request.files:
        logger.info("Performing with input file")
        file = request.files["file"]
        assert file.filename != "", "File is empty"
        assert (
            os.path.splitext(file.filename)[1] == ".xml"
            or os.path.splitext(file.filename)[1] == ".mid"
        ), "File is not .xml or .mid"
        util = PathUtility()
        filepath = os.path.join(util.project_directory(), "temp", file.filename)
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        file.save(filepath)
    else:
        logger.info("Performing with template file")
        time_signature = request.form.get("time_signature", "4/4")
        quarter_length = 1.0 / (int(time_signature.split("/")[1]) / 4)
        number_measures = int(request.form.get("number_measures", "10")) * int(
            time_signature.split("/")[0]
        )
        score = music21.stream.Stream()
        time_signature = music21.meter.TimeSignature(time_signature)
        score.append(time_signature)
        metadata = music21.metadata.Metadata()
        metadata.title = "Random"
        metadata.composer = "MuseTune"
        score.insert(0, metadata)
        note_names = ["A4", "B4", "C4", "D4", "E4", "F4", "G4"]
        for _ in range(number_measures):
            note_name = random.choice(note_names)
            note = music21.note.Note(note_name, quarterLength=quarter_length)
            score.append(note)
        util = PathUtility()
        filepath = os.path.join(util.project_directory(), "temp", "template.mid")
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        score.write("midi", filepath)
    if os.path.splitext(filepath)[1] == ".xml":
        output_filepath = os.path.splitext(filepath)[0] + ".mid"
        musicxml_file = music21.converter.parse(filepath)
        try:
            musicxml_file.write("midi", fp=output_filepath)
        except Exception as exception:
            logger.error(
                "Had an error converting to midi with {} - {}".format(
                    filepath, str(exception)
                )
            )
            raise exception
    elif os.path.splitext(filepath)[1] == ".mid":
        output_filepath = filepath
    cmd = [
        "--file_path",
        output_filepath,
        "--seed",
        str(seed),
        "--conditional_tracks",
        conditional_tracks,
        "--content_tracks",
        content_tracks,
    ]
    save_path = os.path.splitext(music_generator(cmd))[0]
    musicxml_file = music21.converter.parse(save_path + ".mid")
    try:
        musicxml_file.write("musicxml", fp=save_path + ".xml")
    except Exception as E:
        logger.error(
            "Had an error converting the generated content to xml (likely some error with the input file)"
            + str(E)
        )
        return 500
    return send_file(
        save_path + ".xml", as_attachment=True, download_name=save_path + ".xml"
    )


def run_app():
    app.run(host="0.0.0.0", port=8080, debug=True)


if __name__ == "__main__":
    run_app()

tuneease/server.py
Pshah2023/TuneEase
# pip install memory_profiler
# python -m memory_profiler os.path.abspath(__file__) --checkpoint_path <required path>
# from memory_profiler import profile
# from tuneease.tuneease import tuneease

# @profile
# def main():
#     tuneease()

# if __name__ == "__main__":
#     main()

tuneease/tests/memory.py
Pshah2023/TuneEase
import os
import pytest
from werkzeug.datastructures import FileStorage
from ..converter import Converter
from ..pathutility import PathUtility
from ..tuneease import TuneEase


@pytest.fixture
def get_converter():
    path_util = PathUtility()
    museScore_path = path_util.musescore_path()
    return Converter(log_file="tuneease.log", museScore_path=museScore_path)


@pytest.fixture
def get_music_file():
    tuneease = TuneEase()
    random_score_file = tuneease.random_score()
    assert os.path.exists(random_score_file)
    return random_score_file


def test_converter_init(get_converter):
    converter = get_converter
    assert isinstance(converter, Converter)


def test_converter_save_file(get_converter, get_music_file):
    converter = get_converter
    music_file = get_music_file
    destination_folder = "temp"
    filepath = converter.save_file(
        FileStorage(
            stream=open(music_file),
            filename=music_file,
            name=os.path.split(music_file)[-1],
        ),
        destination_folder,
    )
    assert os.path.exists(filepath)


@pytest.mark.benchmark(min_rounds=10)
def test_benchmark_convert(benchmark, get_converter, get_music_file):
    converter = get_converter
    music_file = get_music_file
    result = benchmark(converter.convert_to, music_file, output_extension=".xml")
    assert os.path.exists(result)

tuneease/tests/test_converter.py
Pshah2023/TuneEase
import pytest
import logging
from ..logger import ServerLogger


def test_ServerLogger_init():
    server_logger = ServerLogger("test.log")
    assert isinstance(server_logger, ServerLogger)
    assert isinstance(server_logger.logger, logging.Logger)


def test_ServerLogger_get():
    server_logger = ServerLogger("test.log")
    logger = server_logger.get()
    assert logger is server_logger.logger
    assert isinstance(logger, logging.Logger)

tuneease/tests/test_logger.py
Pshah2023/TuneEase
import os
from ..pathutility import PathUtility


def test_path_utility_init():
    path_util = PathUtility()
    assert isinstance(path_util, PathUtility)
    assert os.path.exists(os.path.abspath(__file__))


def test_path_utility_project_directory():
    path_util = PathUtility()
    project_dir = path_util.project_directory()
    assert os.path.isdir(project_dir)


def test_path_utility_museScore_path():
    path_util = PathUtility()
    musescore_path = path_util.musescore_path()
    assert os.path.exists(musescore_path)

tuneease/tests/test_pathutility.py
Pshah2023/TuneEase
import os
import pytest
from ..tuneease import TuneEase
from ..pathutility import PathUtility
from ..converter import Converter


@pytest.fixture
def get_tuneease():
    path_util = PathUtility()
    museScore_path = path_util.musescore_path()
    return TuneEase(museScore_path=museScore_path)


@pytest.fixture
def get_music_file():
    tuneease = TuneEase()
    random_score_file = tuneease.random_score()
    assert os.path.exists(random_score_file)
    return random_score_file


@pytest.fixture
def get_converter():
    path_util = PathUtility()
    museScore_path = path_util.musescore_path()
    return Converter(log_file="tuneease.log", museScore_path=museScore_path)


def test_tuneease_init(get_tuneease):
    tuneease = get_tuneease
    assert isinstance(tuneease, TuneEase)


def test_convert(get_tuneease, get_music_file, get_converter):
    converter = get_converter
    music_file = get_music_file
    midi_path = converter.convert_to(music_file, output_extension=".midi")
    assert os.path.exists(midi_path)
    tuner = get_tuneease
    output_filepath = tuner.convert(midi_path)
    assert os.path.exists(output_filepath)


def test_split(get_tuneease, get_music_file):
    music_file = get_music_file
    tuner = get_tuneease
    output_filepath = tuner.split(music_file, 1, 2)
    assert os.path.exists(output_filepath)


def test_number(get_tuneease, get_music_file):
    music_file = get_music_file
    tuner = get_tuneease
    output_filepath = tuner.number(music_file)
    assert os.path.exists(output_filepath)


def test_random_score(get_tuneease):
    tuneease = get_tuneease
    random_score_file = tuneease.random_score()
    assert os.path.exists(random_score_file)


def test_generate(get_tuneease):
    tuneease = get_tuneease
    test_param = {
        "condition-lead": "True",
        "content-bass": "True",
    }
    output_filepath = tuneease.generate(test_param)
    assert os.path.exists(output_filepath)

tuneease/tests/test_tuneease.py
Pshah2023/TuneEase
import os
import torch
import numpy as np
from .getmusic.modeling.build import build_model
from .getmusic.utils.misc import merge_opts_to_config
from .logger import ServerLogger
from .getmusic.engine.solver import Solver
from .pipeline.args import get_args
from .pipeline.encoding import encoding_to_MIDI
from .pipeline.file import F
from .pipeline.config import Config


def main(input_args=None):
    args = get_args(input_args)
    logger = ServerLogger("tuneease.log").get()
    logger.debug(str(args))
    gpu = args.use_gpu
    logger.info("Use GPU set to " + str(gpu))
    if gpu:
        torch.cuda.set_device(0)
    args.ngpus_per_node = 1
    args.world_size = 1
    args.local_rank = 0
    args.global_rank = args.local_rank + args.node_rank * args.ngpus_per_node
    args.distributed = args.world_size > 1
    config = Config().config
    config = merge_opts_to_config(config, args.opts)
    model = build_model(config, args)
    conditional_track = np.array([False, False, False, False, False, False, True])
    conditional_name = args.conditional_tracks
    condition_inst = []
    if conditional_name[0] == "1":
        conditional_track[0] = True
        condition_inst.append("80")
    if conditional_name[1] == "1":
        conditional_track[1] = True
        condition_inst.append("32")
    if conditional_name[2] == "1":
        conditional_track[2] = True
    if conditional_name[3] == "1":
        conditional_track[3] = True
        condition_inst.append("25")
    if conditional_name[4] == "1":
        conditional_track[4] = True
        condition_inst.append("0")
    if conditional_name[5] == "1":
        conditional_track[5] = True
        condition_inst.append("48")
    # if 'c' in conditional_name:
    #     conditional_track[6] = True
    if all(conditional_track):
        print(
            "You can't set all tracks as condition. We conduct uncontional generation based on selected content tracks. If you skip content tracks, this song is skipped."
        )
    content_track = np.array([False, False, False, False, False, False, False])
    content_name = args.content_tracks
    if content_name[0] == "1":
        content_track[0] = True
    if content_name[1] == "1":
        content_track[1] = True
    if content_name[2] == "1":
        content_track[2] = True
    if content_name[3] == "1":
        content_track[3] = True
    if content_name[4] == "1":
        content_track[4] = True
    if content_name[5] == "1":
        content_track[5] = True
    x, tempo, not_empty_pos, condition_pos, pitch_shift, tpc, have_cond = F(
        "track_generation.py",
        args.file_path,
        conditional_track,
        content_track,
        condition_inst,
        args.chord_from_single,
    )
    if not have_cond:
        print("chord error")
        return ""
    dataloader_info = None
    solver = Solver(
        gpu=gpu,
        config=config,
        args=args,
        model=model,
        dataloader=dataloader_info,
        logger=logger,
        is_sample=True,
    )
    assert args.load_path is not None
    solver.resume(path=args.load_path)
    oct_line = solver.infer_sample(
        x, tempo, not_empty_pos, condition_pos, use_ema=args.no_ema
    )
    data = oct_line.split(" ")
    oct_final_list = []
    for start in range(3, len(data), 8):
        if "pad" not in data[start] and "pad" not in data[start + 1]:
            pitch = int(data[start][:-1].split("-")[1])
            if data[start - 1] != "<2-129>" and data[start - 1] != "<2-128>":
                pitch -= pitch_shift
            data[start] = "<3-{}>".format(pitch)  # re-normalize
            oct_final_list.append(" ".join(data[start - 3 : start + 5]))
    oct_final = " ".join(oct_final_list)
    midi_obj = encoding_to_MIDI(
        "track_generation.py", oct_final, tpc, args.decode_chord
    )
    save_path = os.path.join(
        os.path.splitext(args.file_path)[0]
        + "{}2{}".format(conditional_name, content_name)
        + ".mid"
    )
    midi_obj.dump(save_path)
    return save_path


if __name__ == "__main__":
    main()

tuneease/track_generation.py
Pshah2023/TuneEase
import warnings
import torch
from .getmusic.modeling.build import build_model
from .getmusic.data.build import build_dataloader
from .getmusic.utils.misc import merge_opts_to_config, modify_config_for_debug
from .logger import ServerLogger
from .getmusic.engine.solver import Solver
from .getmusic.distributed.launch import launch
from .pipeline.args import get_args
from .pipeline.config import Config


def main():
    args = get_args()

    if args.gpu is not None:
        warnings.warn(
            "You have chosen a specific GPU. This will completely disable ddp."
        )
        torch.cuda.set_device(args.gpu)
        args.ngpus_per_node = 1
        args.world_size = 1
    else:
        print("args.num_node ", args.num_node)
        if args.num_node == 1:
            args.dist_url == "auto"
        else:
            assert args.num_node > 1

        args.ngpus_per_node = torch.cuda.device_count()
        args.world_size = args.ngpus_per_node * args.num_node  #

    launch(
        main_worker,
        args.ngpus_per_node,
        args.num_node,
        args.node_rank,
        args.dist_url,
        args=(args,),
    )


def main_worker(local_rank, args):
    args.local_rank = local_rank
    args.global_rank = args.local_rank + args.node_rank * args.ngpus_per_node
    args.distributed = args.world_size > 1
    print(args)
    config = Config().config
    config = merge_opts_to_config(config, args.opts)
    if args.debug:
        config = modify_config_for_debug(config)

    logger = ServerLogger("tuneease.log").get()
    logger.debug(str(args))

    model = build_model(config, args)

    if args.sync_bn:
        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)

    dataloader_info = build_dataloader(config, args)

    solver = Solver(
        config=config, args=args, model=model, dataloader=dataloader_info, logger=logger
    )

    if args.load_path is not None:
        solver.resume(
            path=args.load_path,
            load_optimizer_and_scheduler=args.no_load_optimizer_and_scheduler,
            load_others=args.no_load_others,
        )
    if args.auto_resume:
        solver.resume()
    solver.train()


if __name__ == "__main__":
    main()

tuneease/train.py
Pshah2023/TuneEase
import os
from click import FileError
import music21
import random
from .logger import ServerLogger
from .pathutility import PathUtility
from .converter import Converter
from .track_generation import main as music_generator


class TuneEase:
    """
    A class for generating, converting, and processing music using TuneEase.

    Attributes:
        util (PathUtility): An instance of PathUtility for handling file paths.
        museScore_path (str): The path to the MuseScore executable.
        logger (ServerLogger): An instance of ServerLogger for logging messages.
        converter (Converter): An instance of Converter for file conversion.

    Methods:
        __init__(self, museScore_path=""):
        convert(self, filepath):
        split(self, filepath, start, end=""):
        number(self, filepath):
        random_score(self, time_signature='4/4', number_measures='8'):
        generate(self, form, file=None):

    Example:
        >>> tuneease = TuneEase()
        >>> output_path = tuneease.convert(input_path)
    """

    util = PathUtility()
    museScore_path = util.musescore_path()
    logger = ServerLogger("tuneease.log").get()

    def __init__(self, museScore_path=""):
        if museScore_path:
            self.museScore_path = museScore_path
        self.converter = Converter(museScore_path=self.museScore_path)

    def convert(self, filepath):
        """
        Converts a music file to the MusicXML format.

        Args:
            filepath (str): The path to the input music file.

        Returns:
            str or Exception: The path to the converted MusicXML file if successful, or an exception if an error occurs.

        Example:
            >>> tuneease = TuneEase()
            >>> converted_file = tuneease.convert("input_music.mid")
        """
        self.logger.info("Running convert")
        try:
            output_filepath = self.converter.convert_to(filepath)
            return output_filepath
        except Exception as e:
            return e

    def split(self, filepath, start, end=""):
        """
        Get specific measures from a music file and save the result as MusicXML.

        Args:
            filepath (str): The path to the input music file.
            start (int): The starting measure number.
            end (int, optional): The ending measure number. If not provided, only the start measure is extracted.

        Returns:
            str: The path to the split MusicXML file.

        Example:
            >>> tuneease = TuneEase()
            >>> split_filepath = tuneease.split('inputfile.xml', 2, 5)
        """
        self.logger.info(
            "Running split" + filepath + " with " + str(start) + "-" + str(end)
        )
        filepath = os.path.join(self.util.project_directory(), "temp", filepath)
        measure_filepath = os.path.join(
            self.util.project_directory(),
            "temp",
            os.path.splitext(filepath)[0]
            + "-measure-"
            + str(start)
            + "-"
            + str(end)
            + ".xml",
        )
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        try:
            parsed_song = music21.converter.parse(filepath)
            measures = []
            for part in parsed_song.parts:
                measures.extend(part.getElementsByClass("Measure"))
            if end == "":
                stream = measures[start]
            else:
                measures_to_combine = measures[start - 1 : end]
                stream = music21.stream.Stream()
                for measure in measures_to_combine:
                    stream.append(measure)
            stream.write("musicxml", measure_filepath)
        finally:
            return measure_filepath

    def number(self, filepath):
        """
        Add measure numbers to a music score file.

        Args:
            filepath (str): The path to the input music score file.

        Returns:
            str: The path to the output music score file with measure numbers.

        Example:
            >>> tuneease = TuneEase()
            >>> output_file = tuneease.number("input.xml")
        """
        self.logger.info("Running number on {}".format(filepath))
        filepath = os.path.join(self.util.project_directory(), "temp", filepath)
        output_filepath = os.path.join(
            self.util.project_directory(),
            "temp",
            os.path.splitext(filepath)[0] + "-number.xml",
        )
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        score = music21.converter.parse(filepath)
        for part in score.parts:
            for i, measure in enumerate(
                part.getElementsByClass(music21.stream.Measure)
            ):
                text_expression = music21.expressions.TextExpression(
                    "Measure {}".format(i + 1)
                )
                text_expression.style.absoluteY = 20
                text_expression.style.justify = "left"
                measure.insert(0, text_expression)
        score.write("musicxml", output_filepath)
        return output_filepath

    def random_score(self, time_signature="4/4", number_measures="8"):
        """
        Generate a random music score.

        Args:
            time_signature (str): The time signature of the generated score.
            number_measures (str): The number of measures in the generated score.

        Returns:
            str: The path to the generated random music score file.

        Example:
            >>> tuneease = TuneEase()
            >>> random_score_file = tuneease.random_score(time_signature='4/4', number_measures='8')

        Notes:
            - This method generates a random music score with random notes along with random correct or incorrect signs above each image.
        """
        self.logger.info("Running random score")
        quarter_length = 1.0 / (int(time_signature.split("/")[1]) / 4)
        number_measures = int(number_measures) * int(time_signature.split("/")[0])
        s = music21.stream.Stream()
        ts = music21.meter.TimeSignature(time_signature)
        s.append(ts)
        md = music21.metadata.Metadata()
        md.title = "Random"
        md.composer = "MuseTune"
        s.insert(0, md)
        notes = ["A4", "B4", "C4", "D4", "E4", "F4", "G4"]
        lyrics = ["✔", "x"]
        for i in range(number_measures):
            note_name = random.choice(notes)
            n = music21.note.Note(note_name, quarterLength=quarter_length)
            n.lyric = random.choice(lyrics)
            s.append(n)
        idx = 0
        while os.path.exists(
            os.path.join(self.util.project_directory(), "temp", f"random{idx}.xml")
        ):
            idx += 1
        filepath = os.path.join(
            self.util.project_directory(), "temp", f"random{idx}.xml"
        )
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        s.write("musicxml", filepath)
        return filepath

    def generate(self, input_dict=dict(), file=None, gpu=False):
        """
        Generate music based on user input or a template file.

        Args:
            form (dict): A dictionary containing user input parameters.
            file (str): Optional. A file to base the generated content on.

        Returns:
            str: The path to the generated music XML file.

        Example:
            >>> tuneease = TuneEase()
            >>> user_input = {
            ...     'condition-lead': 'True',
            ...     'condition-bass': 'False',
            ... }
            >>> generated_file = tuneease.generate()

        Notes:
            - Defaults follow below: condition is what is in the input, content is what is to be outputed
            conditional_tracks = "".join([
                str(int(form.get("condition-lead", False) == True)),
                str(int(form.get("condition-bass", False) == True)),
                str(int(form.get("condition-drums", False) == True)),
                str(int(form.get("condition-guitar", False) == True)),
                str(int(form.get("condition-piano", False) == True)),
                str(int(form.get("condition-strings", False) == True)),
                str(int(True))
            ])
            content_tracks = "".join([
                str(int(form.get("content-lead", False) == True)),
                str(int(form.get("content-bass", False) == True)),
                str(int(form.get("content-drums", False) == True)),
                str(int(form.get("content-guitar", False) == True)),
                str(int(form.get("content-piano", True) == True)),
                str(int(form.get("content-strings", False) == True)),
                str(int(False))
            ])
            seed = input_dict.get("seed", 0)
        """
        self.logger.info("Received a generate request with {}".format(str(input_dict)))
        conditional_tracks = "".join(
            [
                str(int(input_dict.get("condition-lead", False))),
                str(int(input_dict.get("condition-bass", False))),
                str(int(input_dict.get("condition-drums", False))),
                str(int(input_dict.get("condition-guitar", False))),
                str(int(input_dict.get("condition-piano", False))),
                str(int(input_dict.get("condition-strings", False))),
                str(int(True)),
            ]
        )
        content_tracks = "".join(
            [
                str(int(input_dict.get("content-lead", False))),
                str(int(input_dict.get("content-bass", False))),
                str(int(input_dict.get("content-drums", False))),
                str(int(input_dict.get("content-guitar", False))),
                str(int(input_dict.get("content-piano", True))),
                str(int(input_dict.get("content-strings", True))),
                str(int(False)),
            ]
        )
        seed = input_dict.get("seed", 0)
        if file:
            assert os.path.exists(file)
            filepath = file
        else:
            self.logger.info(
                "Performing with template file with "
                + str(input_dict.get("number_measures", "5") + " measures")
            )
            time_signature = input_dict.get("time_signature", "4/4")
            quarter_length = 1.0 / (int(time_signature.split("/")[1]) / 4)
            number_measures = int(input_dict.get("number_measures", "5")) * int(
                time_signature.split("/")[0]
            )
            stream = music21.stream.Stream()
            time_signature = music21.meter.TimeSignature(time_signature)
            stream.append(time_signature)
            metadata = music21.metadata.Metadata()
            metadata.title = "Random"
            metadata.composer = "MuseTune"
            stream.insert(0, metadata)
            notes = ["A4", "B4", "C4", "D4", "E4", "F4", "G4"]
            for _ in range(number_measures):
                note_name = random.choice(notes)
                n = music21.note.Note(note_name, quarterLength=quarter_length)
                stream.append(n)
            filepath = os.path.join(
                self.util.project_directory(), "temp", "template.xml"
            )
            os.makedirs(os.path.dirname(filepath), exist_ok=True)
            stream.write("musicxml", filepath)

        if os.path.splitext(filepath)[1] == ".xml":
            output_filepath = os.path.splitext(filepath)[0] + ".mid"
            musicxml_file = music21.converter.parse(filepath)
            try:
                musicxml_file.write("midi", fp=output_filepath)
            except Exception as e:
                self.logger.error(
                    "Had an error converting to midi with {} - {}".format(
                        filepath, str(e)
                    )
                )
                raise e
        elif os.path.splitext(filepath)[1] == ".mid":
            output_filepath = filepath
        else:
            raise FileError("Only .mid or .xml accepted atm")
        cmd = [
            "--file_path",
            output_filepath,
            "--seed",
            str(seed),
            "--conditional_tracks",
            conditional_tracks,
            "--content_tracks",
            content_tracks,
            "--use_gpu",
            int(gpu),
        ]
        save_path = os.path.splitext(music_generator(cmd))[0]
        musicxml_file = music21.converter.parse(save_path + ".mid")
        try:
            musicxml_file.write("musicxml", fp=save_path + ".xml")
        except Exception as e:
            self.logger.error(
                "Had an error converting the generated content to xml (probably some error with input file) - {}".format(
                    str(e)
                )
            )
            raise e
        return save_path + ".xml"


def tuneease():
    tuneease = TuneEase()
    print(tuneease.generate())


if __name__ == "__main__":
    tuneease()

tuneease/tuneease.py
Pshah2023/TuneEase