tunesformer / app.py
sander-wood's picture
Update app.py
3198fcb
raw
history blame
No virus
20.4 kB
import gradio as gr
import torch
import random
from unidecode import unidecode
import re
import os
from tqdm import tqdm
import requests
from samplings import top_p_sampling, top_k_sampling, temperature_sampling
from transformers import GPT2Config, GPT2Model, GPT2LMHeadModel, PreTrainedModel
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
PATCH_LENGTH = 128 # Patch Length
PATCH_SIZE = 32 # Patch Size
PATCH_NUM_LAYERS = 9 # Number of layers in the encoder
CHAR_NUM_LAYERS = 3 # Number of layers in the decoder
NUM_EPOCHS = 32 # Number of epochs to train for (if early stopping doesn't intervene)
LEARNING_RATE = 5e-5 # Learning rate for the optimizer
PATCH_SAMPLING_BATCH_SIZE = 0 # Batch size for patch during training, 0 for full context
LOAD_FROM_CHECKPOINT = False # Whether to load weights from a checkpoint
SHARE_WEIGHTS = False # Whether to share weights between the encoder and decoder
description = """
<div>
<a style="display:inline-block" href='https://github.com/sander-wood/tunesformer'><img src='https://img.shields.io/github/stars/sander-wood/tunesformer?style=social' /></a>
<a style="display:inline-block" href="https://huggingface.co/datasets/sander-wood/irishman"><img src="https://img.shields.io/badge/huggingface-dataset-ffcc66.svg"></a>
<a style="display:inline-block" href="https://arxiv.org/pdf/2301.02884.pdf"><img src="https://img.shields.io/badge/arXiv-2301.02884-b31b1b.svg"></a>
<a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/sander-wood/tunesformer?duplicate=true"><img style="margin-top:0;margin-bottom:0" src="https://huggingface.co/datasets/huggingface/badges/raw/main/duplicate-this-space-md-dark.svg" alt="Duplicate Space"></a>
</div>
## ℹ️ How to use this demo?
1. Enter the control codes to set the musical form of the generated music. For details, please refer to the below "Control Codes" section. (optional)
2. Enter the prompt of the generated music. You can set the control codes to set the musical form, ABC header (i.e., note length, tempo, meter, and key) and the motif of the melody. (optional)
2. You can set the parameters (i.e., number of tunes, maximum patch, top-p, top-k, temperature and random seed) for the generation. (optional)
3. Click "Submit" and wait for the result.
4. The generated ABC notation can be played or edited using [EasyABC]([ABC Sheet Music Editor - EasyABC](https://easyabc.sourceforge.net/), you can also use this [Online ABC Player](https://abc.rectanglered.com/) to render the tune.
## 📝 Control Codes
Inspired by [CTRL](https://huggingface.co/ctrl), we incorporate control codes into TunesFormer to represent musical forms. These codes, positioned ahead of the ABC notation, enable users to specify the structures of the generated tunes. The following control codes are introduced:
- **S:number of sections**: determines the number of sections in the entire melody. It counts on several symbols that can be used to represent section boundaries: `[|`, `||`, `|]`, `|:`, `::`, and `:|`. In our dataset, the range is 1 to 8 (e.g., `S:1` for a single-section melody, and `S:8` for a melody with eight sections).
- **B:number of bars**: specifies the desired number of bars within a section. It counts on the bar symbol `|`. In our dataset, the range is 1 to 32 (e.g., `B:1` for a one-bar section, and `B:32` for a section with 32 bars).
- **E:edit distance similarity**: controls the similarity level between the current section $c$ and a previous section $p$ in the melody. It is based on the Levenshtein distance $lev(c,p)$ , quantifying the difference between sections for creating variations or contrasts. Mathematically, it can be expressed as:
```
eds(c,p) = 1 - lev(c,p) / max(|c|,|p|)
```
where $|c|$ and $|p|$ are the string lengths of the two sections. It is discretized into 11 levels, ranging from no match at all to an exact match (e.g., `E:0` for no similarity, and `E:10` for an exact match).
## ❕Caution
ABC notation is a specialized notation of representing sheet music, and it follows a specific standard format. When interacting with TunesFormer, all trained ABC notation adheres to these standard formats.
If you are unfamiliar with the details of ABC notation, we strongly recommend against manually entering ABC notation. Otherwise, the model may not recognize and generate the music correctly. Inputting incorrect formats might lead to unpredictable outputs or other issues.
A general recommendation is to adjust the desired musical structure and form through control codes and ABC header, rather than directly editing the ABC notation itself.
Please make sure to operate according to the provided formats and guidelines to fully leverage the capabilities of TunesFormer and achieve a satisfying music generation experience.
"""
class Patchilizer:
"""
A class for converting music bars to patches and vice versa.
"""
def __init__(self):
self.delimiters = ["|:", "::", ":|", "[|", "||", "|]", "|"]
self.regexPattern = '(' + '|'.join(map(re.escape, self.delimiters)) + ')'
self.pad_token_id = 0
self.bos_token_id = 1
self.eos_token_id = 2
def split_bars(self, body):
"""
Split a body of music into individual bars.
"""
bars = re.split(self.regexPattern, ''.join(body))
bars = list(filter(None, bars)) # remove empty strings
if bars[0] in self.delimiters:
bars[1] = bars[0] + bars[1]
bars = bars[1:]
bars = [bars[i * 2] + bars[i * 2 + 1] for i in range(len(bars) // 2)]
return bars
def bar2patch(self, bar, patch_size=PATCH_SIZE):
"""
Convert a bar into a patch of specified length.
"""
patch = [self.bos_token_id] + [ord(c) for c in bar] + [self.eos_token_id]
patch = patch[:patch_size]
patch += [self.pad_token_id] * (patch_size - len(patch))
return patch
def patch2bar(self, patch):
"""
Convert a patch into a bar.
"""
return ''.join(chr(idx) if idx > self.eos_token_id else '' for idx in patch if idx != self.eos_token_id)
def encode(self, abc_code, patch_length=PATCH_LENGTH, patch_size=PATCH_SIZE, add_special_patches=False):
"""
Encode music into patches of specified length.
"""
lines = unidecode(abc_code).split('\n')
lines = list(filter(None, lines)) # remove empty lines
body = ""
patches = []
for line in lines:
if len(line) > 1 and ((line[0].isalpha() and line[1] == ':') or line.startswith('%%score')):
if body:
bars = self.split_bars(body)
patches.extend(self.bar2patch(bar + '\n' if idx == len(bars) - 1 else bar, patch_size)
for idx, bar in enumerate(bars))
body = ""
patches.append(self.bar2patch(line + '\n', patch_size))
else:
body += line + '\n'
if body:
patches.extend(self.bar2patch(bar, patch_size) for bar in self.split_bars(body))
if add_special_patches:
bos_patch = [self.bos_token_id] * (patch_size-1) + [self.eos_token_id]
eos_patch = [self.bos_token_id] + [self.eos_token_id] * (patch_size-1)
patches = [bos_patch] + patches + [eos_patch]
return patches[:patch_length]
def decode(self, patches):
"""
Decode patches into music.
"""
return ''.join(self.patch2bar(patch) for patch in patches)
class PatchLevelDecoder(PreTrainedModel):
"""
An Patch-level Decoder model for generating patch features in an auto-regressive manner.
It inherits PreTrainedModel from transformers.
"""
def __init__(self, config):
super().__init__(config)
self.patch_embedding = torch.nn.Linear(PATCH_SIZE * 128, config.n_embd)
torch.nn.init.normal_(self.patch_embedding.weight, std=0.02)
self.base = GPT2Model(config)
def forward(self, patches: torch.Tensor) -> torch.Tensor:
"""
The forward pass of the patch-level decoder model.
:param patches: the patches to be encoded
:return: the encoded patches
"""
patches = torch.nn.functional.one_hot(patches, num_classes=128).float()
patches = patches.reshape(len(patches), -1, PATCH_SIZE * 128)
patches = self.patch_embedding(patches.to(self.device))
return self.base(inputs_embeds=patches)
class CharLevelDecoder(PreTrainedModel):
"""
A Char-level Decoder model for generating the characters within each bar patch sequentially.
It inherits PreTrainedModel from transformers.
"""
def __init__(self, config):
super().__init__(config)
self.pad_token_id = 0
self.bos_token_id = 1
self.eos_token_id = 2
self.base = GPT2LMHeadModel(config)
def forward(self, encoded_patches: torch.Tensor, target_patches: torch.Tensor, patch_sampling_batch_size: int):
"""
The forward pass of the char-level decoder model.
:param encoded_patches: the encoded patches
:param target_patches: the target patches
:return: the decoded patches
"""
# preparing the labels for model training
target_masks = target_patches == self.pad_token_id
labels = target_patches.clone().masked_fill_(target_masks, -100)
# masking the labels for model training
target_masks = torch.ones_like(labels)
target_masks = target_masks.masked_fill_(labels == -100, 0)
# select patches
if patch_sampling_batch_size!=0 and patch_sampling_batch_size<target_patches.shape[0]:
indices = list(range(len(target_patches)))
random.shuffle(indices)
selected_indices = sorted(indices[:patch_sampling_batch_size])
target_patches = target_patches[selected_indices,:]
target_masks = target_masks[selected_indices,:]
encoded_patches = encoded_patches[selected_indices,:]
labels = labels[selected_indices,:]
# get input embeddings
inputs_embeds = torch.nn.functional.embedding(target_patches, self.base.transformer.wte.weight)
# concatenate the encoded patches with the input embeddings
inputs_embeds = torch.cat((encoded_patches.unsqueeze(1), inputs_embeds[:,1:,:]), dim=1)
return self.base(inputs_embeds=inputs_embeds,
attention_mask=target_masks,
labels=labels)
def generate(self, encoded_patch: torch.Tensor, tokens: torch.Tensor):
"""
The generate function for generating a patch based on the encoded patch and already generated tokens.
:param encoded_patch: the encoded patch
:param tokens: already generated tokens in the patch
:return: the probability distribution of next token
"""
encoded_patch = encoded_patch.reshape(1, 1, -1)
tokens = tokens.reshape(1, -1)
# Get input embeddings
tokens = torch.nn.functional.embedding(tokens, self.base.transformer.wte.weight)
# Concatenate the encoded patch with the input embeddings
tokens = torch.cat((encoded_patch, tokens[:,1:,:]), dim=1)
# Get output from model
outputs = self.base(inputs_embeds=tokens)
# Get probabilities of next token
probs = torch.nn.functional.softmax(outputs.logits.squeeze(0)[-1], dim=-1)
return probs
class TunesFormer(PreTrainedModel):
"""
TunesFormer is a hierarchical music generation model based on bar patching.
It includes a patch-level decoder and a character-level decoder.
It inherits PreTrainedModel from transformers.
"""
def __init__(self, encoder_config, decoder_config, share_weights=False):
super().__init__(encoder_config)
self.pad_token_id = 0
self.bos_token_id = 1
self.eos_token_id = 2
if share_weights:
max_layers = max(encoder_config.num_hidden_layers, decoder_config.num_hidden_layers)
max_context_size = max(encoder_config.max_length, decoder_config.max_length)
max_position_embeddings = max(encoder_config.max_position_embeddings, decoder_config.max_position_embeddings)
encoder_config.num_hidden_layers = max_layers
encoder_config.max_length = max_context_size
encoder_config.max_position_embeddings = max_position_embeddings
decoder_config.num_hidden_layers = max_layers
decoder_config.max_length = max_context_size
decoder_config.max_position_embeddings = max_position_embeddings
self.patch_level_decoder = PatchLevelDecoder(encoder_config)
self.char_level_decoder = CharLevelDecoder(decoder_config)
if share_weights:
self.patch_level_decoder.base = self.char_level_decoder.base.transformer
def forward(self, patches: torch.Tensor, patch_sampling_batch_size: int=PATCH_SAMPLING_BATCH_SIZE):
"""
The forward pass of the TunesFormer model.
:param patches: the patches to be both encoded and decoded
:return: the decoded patches
"""
patches = patches.reshape(len(patches), -1, PATCH_SIZE)
encoded_patches = self.patch_level_decoder(patches)["last_hidden_state"]
return self.char_level_decoder(encoded_patches.squeeze(0)[:-1, :], patches.squeeze(0)[1:, :], patch_sampling_batch_size)
def generate(self, patches: torch.Tensor,
tokens: torch.Tensor,
top_p: float=1,
top_k: int=0,
temperature: float=1,
seed: int=None):
"""
The generate function for generating patches based on patches.
:param patches: the patches to be encoded
:return: the generated patches
"""
patches = patches.reshape(len(patches), -1, PATCH_SIZE)
encoded_patches = self.patch_level_decoder(patches)["last_hidden_state"]
if tokens==None:
tokens = torch.tensor([self.bos_token_id], device=self.device)
generated_patch = []
random.seed(seed)
while True:
if seed!=None:
n_seed = random.randint(0, 1000000)
random.seed(n_seed)
else:
n_seed = None
prob = self.char_level_decoder.generate(encoded_patches[0][-1], tokens).cpu().detach().numpy()
prob = top_p_sampling(prob, top_p=top_p, return_probs=True)
prob = top_k_sampling(prob, top_k=top_k, return_probs=True)
token = temperature_sampling(prob, temperature=temperature, seed=n_seed)
generated_patch.append(token)
if token == self.eos_token_id or len(tokens) >= PATCH_SIZE - 1:
break
else:
tokens = torch.cat((tokens, torch.tensor([token], device=self.device)), dim=0)
return generated_patch, n_seed
def generate_abc(prompt,
num_tunes,
max_patch,
top_p,
top_k,
temperature,
seed,
show_control_code=True):
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
patchilizer = Patchilizer()
patch_config = GPT2Config(num_hidden_layers=PATCH_NUM_LAYERS,
max_length=PATCH_LENGTH,
max_position_embeddings=PATCH_LENGTH,
vocab_size=1)
char_config = GPT2Config(num_hidden_layers=CHAR_NUM_LAYERS,
max_length=PATCH_SIZE,
max_position_embeddings=PATCH_SIZE,
vocab_size=128)
model = TunesFormer(patch_config, char_config, share_weights=SHARE_WEIGHTS)
filename = "weights.pth"
if os.path.exists(filename):
print(f"Weights already exist at '{filename}'. Loading...")
else:
print(f"Downloading weights to '{filename}' from huggingface.co...")
try:
url = 'https://huggingface.co/sander-wood/tunesformer/resolve/main/weights.pth'
response = requests.get(url, stream=True)
total_size = int(response.headers.get('content-length', 0))
chunk_size = 1024
with open(filename, 'wb') as file, tqdm(
desc=filename,
total=total_size,
unit='B',
unit_scale=True,
unit_divisor=1024,
) as bar:
for data in response.iter_content(chunk_size=chunk_size):
size = file.write(data)
bar.update(size)
except Exception as e:
print(f"Error: {e}")
exit()
checkpoint = torch.load('weights.pth')
model.load_state_dict(checkpoint['model'])
model = model.to(device)
model.eval()
tunes = ""
print("\n"+" OUTPUT TUNES ".center(60, "#"))
for i in range(num_tunes):
tune = "X:"+str(i+1) + "\n" + prompt
lines = re.split(r'(\n)', tune)
tune = ""
skip = False
for line in lines:
if show_control_code or line[:2] not in ["S:", "B:", "E:"]:
if not skip:
print(line, end="")
tune += line
skip = False
else:
skip = True
input_patches = torch.tensor([patchilizer.encode(prompt, add_special_patches=True)[:-1]], device=device)
if tune=="":
tokens = None
else:
prefix = patchilizer.decode(input_patches[0])
remaining_tokens = prompt[len(prefix):]
tokens = torch.tensor([patchilizer.bos_token_id]+[ord(c) for c in remaining_tokens], device=device)
while input_patches.shape[1]<max_patch:
predicted_patch, seed = model.generate(input_patches,
tokens,
top_p=top_p,
top_k=top_k,
temperature=temperature,
seed=seed)
tokens = None
if predicted_patch[0]!=patchilizer.eos_token_id:
next_bar = patchilizer.decode([predicted_patch])
if show_control_code or next_bar[:2] not in ["S:", "B:", "E:"]:
print(next_bar, end="")
tune += next_bar
if next_bar=="":
break
next_bar = remaining_tokens+next_bar
remaining_tokens = ""
predicted_patch = torch.tensor(patchilizer.bar2patch(next_bar), device=device).unsqueeze(0)
input_patches = torch.cat([input_patches, predicted_patch.unsqueeze(0)], dim=1)
else:
break
tunes += tune+"\n\n"
return tunes
default_prompt = """S:2
B:9
E:4
B:9
L:1/8
M:3/4
K:D
de |"D" """
input_prompt = gr.inputs.Textbox(lines=5, label="Prompt", default=default_prompt)
input_num_tunes = gr.inputs.Slider(minimum=1, maximum=10, step=1, default=1, label="Number of Tunes")
input_max_patch = gr.inputs.Slider(minimum=10, maximum=128, step=1, default=128, label="Max Patch")
input_top_p = gr.inputs.Slider(minimum=0.0, maximum=1.0, step=0.05, default=0.8, label="Top P")
input_top_k = gr.inputs.Slider(minimum=1, maximum=20, step=1, default=8, label="Top K")
input_temperature = gr.inputs.Slider(minimum=0.0, maximum=2.0, step=0.05, default=1.2, label="Temperature")
input_seed = gr.inputs.Textbox(lines=1, label="Seed (int)", default="None")
output_abc = gr.outputs.Textbox(label="Generated Tunes")
gr.Interface(generate_abc,
[input_prompt, input_num_tunes, input_max_patch, input_top_p, input_top_k, input_temperature, input_seed],
output_abc,
title="TunesFormer: Forming Irish Tunes with Control Codes by Bar Patching",
description=description).launch()