vendor jukebox+sheetsage source and recovered weights, use_jukebox=True

app.py

@@ -17,8 +17,13 @@ except ImportError:
                 func = args[0]
                 return func
 
+import os
+import shutil
 import tempfile
 import threading
+import time
+import urllib.request
+from pathlib import Path
 
 import gradio as gr
 import soundfile as sf
@@ -29,6 +34,31 @@ import picogen2
 from picogen2.mirtoolkit.beat_this import BeatThis
 from picogen2.mirtoolkit.sheetsage import SheetSage
 
+REPO_ROOT = Path(__file__).parent
+
+
+def _download_with_progress(url: str, dest: Path, label: str, interval_s: float = 20.0):
+    """Downloads url to dest, logging progress at most once per interval_s."""
+    if dest.exists():
+        return
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    tmp_path = dest.with_name(dest.name + ".part")
+    with urllib.request.urlopen(url) as response, open(tmp_path, "wb") as f:
+        total = int(response.headers.get("Content-Length", 0))
+        downloaded = 0
+        last_log = time.monotonic()
+        while chunk := response.read(1024 * 1024):
+            f.write(chunk)
+            downloaded += len(chunk)
+            now = time.monotonic()
+            if now - last_log >= interval_s:
+                pct = 100 * downloaded / total if total else 0
+                print(f"{label}: {downloaded / 1e9:.2f}/{total / 1e9:.2f}GB ({pct:.0f}%)")
+                last_log = now
+    tmp_path.rename(dest)
+    print(f"{label}: done ({dest.stat().st_size / 1e9:.2f}GB)")
+
+
 # SheetSage (this model's audio feature extractor) was trained on ~24s segments and is
 # most accurate on short clips; longer songs also risk exceeding the GPU time budget below.
 MAX_INPUT_SECONDS = 30.0
@@ -46,25 +76,49 @@ model_error = None
 
 
 def load_assets():
-    """Downloads PiCoGen2's checkpoint, pre-fetches SheetSage's (large) asset cache, and
-    builds the beat tracker, all on CPU so the app can start serving while this runs."""
+    """Downloads PiCoGen2's checkpoint and Jukebox's weights, stages SheetSage's vendored
+    checkpoints, and builds the beat tracker -- all on CPU so the app can start serving
+    while this runs."""
     global decoder, tokenizer, beat_detector, model_loading, model_error
     try:
         tokenizer = picogen2.Tokenizer()
         decoder = picogen2.PiCoGenDecoder.from_pretrained(device="cpu")
 
-        try:
-            import sheetsage.assets as sheetsage_assets
-
-            for tag in sheetsage_assets.get_asset_tags():
-                sheetsage_assets.retrieve_asset(tag)
-        except Exception as e:
-            # SheetSage lazily re-downloads whatever's missing on first use, so this is
-            # a warm-up, not a hard requirement.
-            print(f"SheetSage asset pre-fetch incomplete, will retry on first request: {e}")
+        # SheetSage's upstream S3 bucket (its own retrieve_asset download source) has
+        # been dead for a while (see https://github.com/chrisdonahue/sheetsage/issues/44,
+        # 45, 46). Its checkpoints are small (245MB total) and vendored directly in this
+        # repo instead; stage them where sheetsage.assets.retrieve_asset expects to find
+        # them so it treats them as already downloaded.
+        sheetsage_cache = Path.home() / ".sheetsage" / "sheetsage" / "v0.2"
+        sheetsage_cache.mkdir(parents=True, exist_ok=True)
+        for item in (REPO_ROOT / "sheetsage" / "weights").iterdir():
+            dest = sheetsage_cache / item.name
+            if item.name.startswith(".") or dest.exists():
+                continue
+            if item.is_dir():
+                shutil.copytree(item, dest)
+            else:
+                shutil.copy(item, dest)
+
+        # Jukebox's own weights are hosted on OpenAI's CDN, unrelated to (and unaffected
+        # by) SheetSage's dead bucket. Same default cache path jukebox's own downloader
+        # uses (see jukebox/make_models.py:load_checkpoint), pre-fetched here so the
+        # first real request doesn't have to wait on a 10GB download inside its GPU
+        # time budget.
+        jukebox_cache = Path(os.environ.get("JUKEBOX_CACHE_DIR", "~/.cache")).expanduser()
+        _download_with_progress(
+            "https://openaipublic.azureedge.net/jukebox/models/5b/vqvae.pth.tar",
+            jukebox_cache / "jukebox" / "models" / "5b" / "vqvae.pth.tar",
+            "jukebox vqvae",
+        )
+        _download_with_progress(
+            "https://openaipublic.azureedge.net/jukebox/models/5b/prior_level_2.pth.tar",
+            jukebox_cache / "jukebox" / "models" / "5b" / "prior_level_2.pth.tar",
+            "jukebox prior_level_2",
+        )
 
         beat_detector = BeatThis(cuda=False)
-        print("Models loaded (CPU).")
+        print("Models loaded (CPU); Jukebox weights ready.")
     except Exception as e:
         model_error = str(e)
         print(f"Load error: {e}")

constraints.txt

@@ -1 +0,0 @@
-setuptools<82

jukebox/align.py

@@ -0,0 +1,115 @@
+"""
+Get alignment from attn values
+1. run a forward pass on each hop, get attn values
+2. concat for all hops
+"""
+import numpy as np
+import torch as t
+from jukebox.utils.torch_utils import assert_shape, empty_cache
+from jukebox.hparams import Hyperparams
+from jukebox.make_models import make_model
+from jukebox.save_html import save_html
+from jukebox.utils.sample_utils import get_starts
+import fire
+
+def get_alignment(x, zs, labels, prior, fp16, hps):
+    level = hps.levels - 1 # Top level used
+    n_ctx, n_tokens = prior.n_ctx, prior.n_tokens
+    z = zs[level]
+    bs, total_length = z.shape[0], z.shape[1]
+    if total_length < n_ctx:
+        padding_length = n_ctx - total_length
+        z = t.cat([z, t.zeros(bs, n_ctx - total_length, dtype=z.dtype, device=z.device)], dim=1)
+        total_length = z.shape[1]
+    else:
+        padding_length = 0
+
+    hop_length = int(hps.hop_fraction[level]*prior.n_ctx)
+    n_head = prior.prior.transformer.n_head
+    alignment_head, alignment_layer = prior.alignment_head, prior.alignment_layer
+    attn_layers = set([alignment_layer])
+    alignment_hops = {}
+    indices_hops = {}
+
+    prior.cuda()
+    empty_cache()
+    for start in get_starts(total_length, n_ctx, hop_length):
+        end = start + n_ctx
+
+        # set y offset, sample_length and lyrics tokens
+        y, indices_hop = prior.get_y(labels, start, get_indices=True)
+        assert len(indices_hop) == bs
+        for indices in indices_hop:
+            assert len(indices) == n_tokens
+
+        z_bs = t.chunk(z, bs, dim=0)
+        y_bs = t.chunk(y, bs, dim=0)
+        w_hops = []
+        for z_i, y_i in zip(z_bs, y_bs):
+            w_hop = prior.z_forward(z_i[:,start:end], [], y_i, fp16=fp16, get_attn_weights=attn_layers)
+            assert len(w_hop) == 1
+            w_hops.append(w_hop[0][:, alignment_head])
+            del w_hop
+        w = t.cat(w_hops, dim=0)
+        del w_hops
+        assert_shape(w, (bs, n_ctx, n_tokens))
+        alignment_hop = w.float().cpu().numpy()
+        assert_shape(alignment_hop, (bs, n_ctx, n_tokens))
+        del w
+
+        # alignment_hop has shape (bs, n_ctx, n_tokens)
+        # indices_hop is a list of len=bs, each entry of len hps.n_tokens
+        indices_hops[start] = indices_hop
+        alignment_hops[start] = alignment_hop
+    prior.cpu()
+    empty_cache()
+
+    # Combine attn for each hop into attn for full range
+    # Use indices to place them into correct place for corresponding source tokens
+    alignments = []
+    for item in range(bs):
+        # Note each item has different length lyrics
+        full_tokens = labels['info'][item]['full_tokens']
+        alignment = np.zeros((total_length, len(full_tokens) + 1))
+        for start in reversed(get_starts(total_length, n_ctx, hop_length)):
+            end = start + n_ctx
+            alignment_hop = alignment_hops[start][item]
+            indices = indices_hops[start][item]
+            assert len(indices) == n_tokens
+            assert alignment_hop.shape == (n_ctx, n_tokens)
+            alignment[start:end,indices] = alignment_hop
+        alignment = alignment[:total_length - padding_length,:-1] # remove token padding, and last lyric index
+        alignments.append(alignment)
+    return alignments
+
+def save_alignment(model, device, hps):
+    print(hps)
+    vqvae, priors = make_model(model, device, hps, levels=[-1])
+
+    logdir = f"{hps.logdir}/level_{0}"
+    data = t.load(f"{logdir}/data.pth.tar")
+    if model == '1b_lyrics':
+        fp16 = False
+    else:
+        fp16 = True
+
+    data['alignments'] = get_alignment(data['x'], data['zs'], data['labels'][-1], priors[-1], fp16, hps)
+    t.save(data, f"{logdir}/data_align.pth.tar")
+    save_html(logdir, data['x'], data['zs'], data['labels'][-1], data['alignments'], hps)
+
+def run(model, port=29500, **kwargs):
+    from jukebox.utils.dist_utils import setup_dist_from_mpi
+    rank, local_rank, device = setup_dist_from_mpi(port=port)
+    hps = Hyperparams(**kwargs)
+
+    with t.no_grad():
+        save_alignment(model, device, hps)
+
+if __name__ == '__main__':
+    fire.Fire(run)
+
+
+
+
+
+

jukebox/data/artist_genre_processor.py

@@ -0,0 +1,93 @@
+import os
+import re
+
+accepted = frozenset([chr(i) for i in range(ord('a'), ord('z') + 1)] +
+                     [chr(i) for i in range(ord('A'), ord('Z') + 1)] +
+                     [chr(i) for i in range(ord('0'), ord('9') + 1)])
+
+rex = re.compile(r'_+')
+
+def norm(s):
+    s = ''.join([c if c in accepted else '_' for c in s.lower()])
+    s = rex.sub('_', s).strip('_')
+    return s
+
+def create_reverse_lookup(atoi):
+    # Multiple entries could go to the same artist_id/genre_id
+    itoa = {}
+    for a, i in atoi.items():
+        if i not in itoa:
+            itoa[i] = []
+        itoa[i].append(a)
+    indices = sorted(list(itoa.keys()))
+    for i in indices:
+        itoa[i] = '_'.join(sorted(itoa[i]))
+    return itoa
+
+class ArtistGenreProcessor():
+    def __init__(self, v3=False):
+        self.v3 = v3
+        dirname = os.path.dirname(__file__)
+        if self.v3:
+            self.artist_id_file = f"{dirname}/ids/v3_artist_ids.txt"
+            self.genre_id_file = f"{dirname}/ids/v3_genre_ids.txt"
+        else:
+            self.artist_id_file = f"{dirname}/ids/v2_artist_ids.txt"
+            self.genre_id_file = f"{dirname}/ids/v2_genre_ids.txt"
+        self.load_artists()
+        self.load_genres()
+
+    def get_artist_id(self, artist):
+        input_artist = artist
+        if self.v3:
+            artist = artist.lower()
+        else:
+            artist = norm(artist)
+        if artist not in self.artist_ids:
+            print(f"Input artist {input_artist} maps to {artist}, which is not present in {self.artist_id_file}. "
+                  f"Defaulting to (artist_id, artist) = (0, unknown), if that seems wrong please format artist correctly")
+        return self.artist_ids.get(artist, 0)
+
+    def get_genre_ids(self, genre):
+        if self.v3:
+            genres = [genre.lower()]
+        else:
+            # In v2, we convert genre into a bag of words
+            genres = norm(genre).split("_")
+        for word in genres:
+            if word not in self.genre_ids:
+                print(f"Input genre {genre} maps to the list {genres}. {word} is not present in {self.genre_id_file}. "
+                      f"Defaulting to (word_id, word) = (0, unknown), if that seems wrong please format genre correctly")
+        return [self.genre_ids.get(word, 0) for word in genres]
+
+    # get_artist/genre throw error if we ask for non-present values
+    def get_artist(self, artist_id):
+        return self.artists[artist_id]
+
+    def get_genre(self, genre_ids):
+        if self.v3:
+            assert len(genre_ids) == 1
+            genre = self.genres[genre_ids[0]]
+        else:
+            genre = '_'.join([self.genres[genre_id] for genre_id in genre_ids if genre_id >= 0])
+        return genre
+
+    def load_artists(self):
+        print(f'Loading artist IDs from {self.artist_id_file}')
+        self.artist_ids = {}
+        with open(self.artist_id_file, 'r', encoding="utf-8") as f:
+            for line in f:
+                artist, artist_id = line.strip().split(';')
+                self.artist_ids[artist.lower()] = int(artist_id)
+        self.artists = create_reverse_lookup(self.artist_ids)
+
+    def load_genres(self):
+        print(f'Loading artist IDs from {self.genre_id_file}')
+        self.genre_ids = {}
+        with open(self.genre_id_file, 'r', encoding="utf-8") as f:
+            for line in f:
+                genre, genre_id = line.strip().split(';')
+                self.genre_ids[genre.lower()] = int(genre_id)
+        self.genres = create_reverse_lookup(self.genre_ids)
+
+

jukebox/data/data_processor.py

@@ -0,0 +1,69 @@
+import torch as t
+import jukebox.utils.dist_adapter as dist
+from torch.utils.data.distributed import DistributedSampler
+from torch.utils.data import DataLoader, Dataset, BatchSampler, RandomSampler
+from jukebox.utils.dist_utils import print_all
+from jukebox.utils.audio_utils import calculate_bandwidth
+from jukebox.data.files_dataset import FilesAudioDataset
+
+class OffsetDataset(Dataset):
+    def __init__(self, dataset, start, end, test=False):
+        super().__init__()
+        self.dataset = dataset
+        self.start = start
+        self.end = end
+        self.test = test
+        assert 0 <= self.start < self.end <= len(self.dataset)
+
+    def __len__(self):
+        return self.end - self.start
+
+    def __getitem__(self, item):
+        return self.dataset.get_item(self.start + item, test=self.test)
+
+class DataProcessor():
+    def __init__(self, hps):
+        self.dataset = FilesAudioDataset(hps)
+        duration = 1 if hps.prior else 600
+        hps.bandwidth = calculate_bandwidth(self.dataset, hps, duration=duration)
+        self.create_datasets(hps)
+        self.create_samplers(hps)
+        self.create_data_loaders(hps)
+        self.print_stats(hps)
+
+    def set_epoch(self, epoch):
+        self.train_sampler.set_epoch(epoch)
+        self.test_sampler.set_epoch(epoch)
+
+    def create_datasets(self, hps):
+        train_len = int(len(self.dataset) * hps.train_test_split)
+        self.train_dataset = OffsetDataset(self.dataset, 0, train_len, test=False)
+        self.test_dataset = OffsetDataset(self.dataset, train_len, len(self.dataset), test=True)
+
+    def create_samplers(self, hps):
+        if not dist.is_available():
+            self.train_sampler = BatchSampler(RandomSampler(self.train_dataset), batch_size=hps.bs, drop_last=True)
+            self.test_sampler = BatchSampler(RandomSampler(self.test_dataset), batch_size=hps.bs, drop_last=True)
+        else:
+            self.train_sampler = DistributedSampler(self.train_dataset)
+            self.test_sampler = DistributedSampler(self.test_dataset)
+
+    def create_data_loaders(self, hps):
+        # Loader to load mini-batches
+        if hps.labels:
+            collate_fn = lambda batch: tuple(t.stack([t.from_numpy(b[i]) for b in batch], 0) for i in range(2))
+        else:
+            collate_fn = lambda batch: t.stack([t.from_numpy(b) for b in batch], 0)
+
+        print('Creating Data Loader')
+        self.train_loader = DataLoader(self.train_dataset, batch_size=hps.bs, num_workers=hps.nworkers,
+                                       sampler=self.train_sampler, pin_memory=False,
+                                       drop_last=True, collate_fn=collate_fn)
+        self.test_loader = DataLoader(self.test_dataset, batch_size=hps.bs, num_workers=hps.nworkers,
+                                      sampler=self.test_sampler, pin_memory=False,
+                                      drop_last=False, collate_fn=collate_fn)
+
+    def print_stats(self, hps):
+        print_all(f"Train {len(self.train_dataset)} samples. Test {len(self.test_dataset)} samples")
+        print_all(f'Train sampler: {self.train_sampler}')
+        print_all(f'Train loader: {len(self.train_loader)}')

jukebox/data/files_dataset.py

@@ -0,0 +1,99 @@
+import librosa
+import math
+import numpy as np
+import jukebox.utils.dist_adapter as dist
+from torch.utils.data import Dataset
+from jukebox.utils.dist_utils import print_all
+from jukebox.utils.io import get_duration_sec, load_audio
+from jukebox.data.labels import Labeller
+
+class FilesAudioDataset(Dataset):
+    def __init__(self, hps):
+        super().__init__()
+        self.sr = hps.sr
+        self.channels = hps.channels
+        self.min_duration = hps.min_duration or math.ceil(hps.sample_length / hps.sr)
+        self.max_duration = hps.max_duration or math.inf
+        self.sample_length = hps.sample_length
+        assert hps.sample_length / hps.sr < self.min_duration, f'Sample length {hps.sample_length} per sr {hps.sr} ({hps.sample_length / hps.sr:.2f}) should be shorter than min duration {self.min_duration}'
+        self.aug_shift = hps.aug_shift
+        self.labels = hps.labels
+        self.init_dataset(hps)
+
+    def filter(self, files, durations):
+        # Remove files too short or too long
+        keep = []
+        for i in range(len(files)):
+            if durations[i] / self.sr < self.min_duration:
+                continue
+            if durations[i] / self.sr >= self.max_duration:
+                continue
+            keep.append(i)
+        print_all(f'self.sr={self.sr}, min: {self.min_duration}, max: {self.max_duration}')
+        print_all(f"Keeping {len(keep)} of {len(files)} files")
+        self.files = [files[i] for i in keep]
+        self.durations = [int(durations[i]) for i in keep]
+        self.cumsum = np.cumsum(self.durations)
+
+    def init_dataset(self, hps):
+        # Load list of files and starts/durations
+        files = librosa.util.find_files(f'{hps.audio_files_dir}', ['mp3', 'opus', 'm4a', 'aac', 'wav'])
+        print_all(f"Found {len(files)} files. Getting durations")
+        cache = dist.get_rank() % 8 == 0 if dist.is_available() else True
+        durations = np.array([get_duration_sec(file, cache=cache) * self.sr for file in files])  # Could be approximate
+        self.filter(files, durations)
+
+        if self.labels:
+            self.labeller = Labeller(hps.max_bow_genre_size, hps.n_tokens, self.sample_length, v3=hps.labels_v3)
+
+    def get_index_offset(self, item):
+        # For a given dataset item and shift, return song index and offset within song
+        half_interval = self.sample_length//2
+        shift = np.random.randint(-half_interval, half_interval) if self.aug_shift else 0
+        offset = item * self.sample_length + shift # Note we centred shifts, so adding now
+        midpoint = offset + half_interval
+        assert 0 <= midpoint < self.cumsum[-1], f'Midpoint {midpoint} of item beyond total length {self.cumsum[-1]}'
+        index = np.searchsorted(self.cumsum, midpoint)  # index <-> midpoint of interval lies in this song
+        start, end = self.cumsum[index - 1] if index > 0 else 0.0, self.cumsum[index] # start and end of current song
+        assert start <= midpoint <= end, f"Midpoint {midpoint} not inside interval [{start}, {end}] for index {index}"
+        if offset > end - self.sample_length: # Going over song
+            offset = max(start, offset - half_interval)  # Now should fit
+        elif offset < start: # Going under song
+            offset = min(end - self.sample_length, offset + half_interval)  # Now should fit
+        assert start <= offset <= end - self.sample_length, f"Offset {offset} not in [{start}, {end - self.sample_length}]. End: {end}, SL: {self.sample_length}, Index: {index}"
+        offset = offset - start
+        return index, offset
+
+    def get_metadata(self, filename, test):
+        """
+        Insert metadata loading code for your dataset here.
+        If artist/genre labels are different from provided artist/genre lists,
+        update labeller accordingly.
+
+        Returns:
+            (artist, genre, full_lyrics) of type (str, str, str). For
+            example, ("unknown", "classical", "") could be a metadata for a
+            piano piece.
+        """
+        return None, None, None
+
+    def get_song_chunk(self, index, offset, test=False):
+        filename, total_length = self.files[index], self.durations[index]
+        data, sr = load_audio(filename, sr=self.sr, offset=offset, duration=self.sample_length)
+        assert data.shape == (self.channels, self.sample_length), f'Expected {(self.channels, self.sample_length)}, got {data.shape}'
+        if self.labels:
+            artist, genre, lyrics = self.get_metadata(filename, test)
+            labels = self.labeller.get_label(artist, genre, lyrics, total_length, offset)
+            return data.T, labels['y']
+        else:
+            return data.T
+
+    def get_item(self, item, test=False):
+        index, offset = self.get_index_offset(item)
+        return self.get_song_chunk(index, offset, test)
+
+    def __len__(self):
+        return int(np.floor(self.cumsum[-1] / self.sample_length))
+
+    def __getitem__(self, item):
+        return self.get_item(item)

jukebox/data/ids/v2_artist_ids.txt

@@ -0,0 +1,4111 @@
+unknown;0
+various;0
+;0
+andr_s_schiff;1
+sonny_terry;2
+nelly;3
+markus_schulz;4
+modest_petrovich_mussorgsky;5
+otis_redding;6
+aerosmith;7
+kenny_g;8
+james_taylor;9
+bobby_bland;10
+burning_spear;11
+skip_james;12
+heart;13
+tammy_wynette;14
+muse;15
+beres_hammond;16
+james_newton_howard;17
+nelson_freire;18
+benny_goodman;19
+hank_williams;20
+they_might_be_giants;21
+the_brian_jonestown_massacre;22
+lady_gaga;23
+chris_young;24
+alison_krauss_union_station;25
+seal;26
+the_hollies;27
+shabba_ranks;28
+paul_young;29
+iration;30
+buck_owens;31
+the_weeknd;32
+elton_john;33
+smokey_robinson;34
+roy_orbison;35
+headhunterz;36
+blondie;37
+the_temptations;38
+ray_stevens;39
+foo_fighters;40
+christoph_eschenbach;41
+blind_willie_mctell;42
+al_martino;43
+edwin_fischer;44
+victor_young;45
+justin_bieber;46
+styx;47
+doris_day;48
+tex_beneke;49
+the_monkees;50
+richard_wagner;51
+bryan_adams;52
+alessandro_scarlatti;53
+rebelution;54
+pitbull;55
+nat_king_cole;56
+wiz_khalifa;57
+roger_miller;58
+andy_williams;59
+peggy_lee;60
+pyotr_ilyich_tchaikovsky;61
+booker_t_the_mg_s;62
+cilla_black;63
+billy_fury;64
+vera_lynn;65
+enrico_caruso;66
+sly_and_robbie;67
+the_pretenders;68
+the_sweet;69
+kylie_minogue;70
+kay_kyser;71
+san_francisco_symphony;72
+prince;73
+queen;74
+kool_the_gang;75
+horace_andy;76
+midnite;77
+gentleman;78
+wilhelm_kempff;79
+busta_rhymes;80
+the_pogues;81
+def_leppard;82
+al_jolson;83
+king_tubby;84
+hot_chocolate;85
+delroy_wilson;86
+jody_watley;87
+bobby_vee;88
+johnny_mathis;89
+the_rascals;90
+sviatoslav_richter;91
+fred_astaire;92
+john_holt;93
+amy_grant;94
+b_b_king;95
+paul_weston;96
+four_tops;97
+jay_sean;98
+pat_boone;99
+george_frideric_handel;100
+bing_crosby;101
+shalamar;102
+tommy_dorsey;103
+ludacris;104
+kenny_chesney;105
+murray_perahia;106
+lightnin_hopkins;107
+ricky_nelson;108
+clint_mansell;109
+tom_petty_and_the_heartbreakers;110
+gary_glitter;111
+ringo_starr;112
+phil_collins;113
+leo_reisman;114
+al_green;115
+jim_reeves;116
+chris_brown;117
+cliff_edwards;118
+buddy_guy;119
+angelo_badalamenti;120
+frank_sinatra;121
+santana;122
+the_pussycat_dolls;123
+peaches_herb;124
+radu_lupu;125
+the_whispers;126
+eddy_howard;127
+memphis_slim;128
+henry_mancini;129
+giuseppe_verdi;130
+the_dream;131
+vladimir_sofronitsky;132
+u_roy;133
+ken_boothe;134
+the_kinks;135
+howard_shore;136
+hardwell;137
+lou_reed;138
+calvin_harris;139
+eddy_chen;140
+anne_murray;141
+juice_newton;142
+bee_gees;143
+wilson_pickett;144
+alan_jackson;145
+shirley_bassey;146
+waylon_jennings;147
+destiny_s_child;148
+cab_calloway;149
+johnny_copeland;150
+bright_eyes;151
+trey_songz;152
+neil_sedaka;153
+justin_timberlake;154
+arthur_grumiaux;155
+the_who;156
+the_yardbirds;157
+big_joe_turner;158
+duke_ellington;159
+herb_alpert;160
+laura_branigan;161
+michael_jackson;162
+john_denver;163
+peter_gordon;164
+solomon_cutner;165
+steve_miller_band;166
+don_williams;167
+the_pointer_sisters;168
+metallica;169
+the_ink_spots;170
+kenny_rogers;171
+the_game;172
+gene_krupa;173
+snoop_dogg;174
+j_cole;175
+taylor_dayne;176
+r3hab;177
+guy_lombardo_and_his_royal_canadians;178
+shakin_stevens;179
+tim_mcgraw;180
+olivia_newton_john_john_travolta;181
+the_replacements;182
+beenie_man;183
+diplo;184
+sammy_kaye;185
+don_gibson;186
+shaggy;187
+nina_simone;188
+stone_temple_pilots;189
+ne_yo;190
+huddie_william_ledbetter;191
+jerry_goldsmith;192
+the_bellamy_brothers;193
+winifred_atwell;194
+georg_philipp_telemann;195
+timbaland;196
+the_5th_dimension;197
+dionne_warwick;198
+r_kelly;199
+enrique_iglesias;200
+sting;201
+mikey_dread;202
+yellowman;203
+eddie_kendricks;204
+blur;205
+twista;206
+paul_anka;207
+chris_cornell;208
+a_ha;209
+pet_shop_boys;210
+yo_yo_ma;211
+tom_jones;212
+neil_diamond;213
+vic_damone;214
+paul_oakenfold;215
+jascha_heifetz;216
+t_i_;217
+dinah_washington;218
+vladimir_ashkenazy;219
+leif_ove_andsnes;220
+johnny_cash;221
+basement_jaxx;222
+sonic_youth;223
+the_isley_brothers;224
+jason_aldean;225
+henryk_szeryng;226
+lloyd_price;227
+reba_mcentire;228
+bon_jovi;229
+bed_ich_smetana;230
+donny_osmond;231
+chuck_berry;232
+the_smashing_pumpkins;233
+israel_vibration;234
+stan_kenton;235
+conway_twitty_loretta_lynn;236
+mcfly;237
+r_e_m;238
+morgan_heritage;239
+lil_wayne;240
+garnett_silk;241
+lee_scratch_perry;242
+howlin_wolf;243
+jon_secada;244
+goo_goo_dolls;245
+brian_hyland;246
+abc;247
+clarence_gatemouth_brown;248
+wilson_phillips;249
+atlantic_starr;250
+david_guetta;251
+s_rgio_mendes;252
+fr_d_ric_chopin;253
+robert_palmer;254
+charlie_musselwhite;255
+level_42;256
+peter_andre;257
+the_spinners;258
+erasure;259
+philippe_entremont;260
+leo_jan_ek;261
+the_kingston_trio;262
+ronnie_milsap;263
+pat_benatar;264
+robert_casadesus;265
+t_bone_walker;266
+duran_duran;267
+jerry_reed;268
+franz_liszt;269
+journey;270
+steve_angello;271
+showaddywaddy;272
+tears_for_fears;273
+john_williams;274
+daniel_m_ller_schott;275
+tampa_red;276
+tina_turner;277
+the_beatles;278
+miranda_lambert;279
+howard_jones;280
+reo_speedwagon;281
+lady_antebellum;282
+no_doubt;283
+status_quo;284
+tiffany;285
+billy_eckstine;286
+danny_elfman;287
+jimmy_wakely;288
+ike_tina_turner;289
+george_gershwin;290
+dinah_shore;291
+armin_van_buuren;292
+keith_sweat;293
+beaux_arts_trio;294
+arcangelo_corelli;295
+air_supply;296
+w_w;297
+the_mighty_diamonds;298
+harry_gregson_williams;299
+blind_lemon_jefferson;300
+brook_benton;301
+sizzla;302
+bobby_darin;303
+steely_dan;304
+the_skatalites;305
+dwight_yoakam;306
+clara_haskil;307
+jah_cure;308
+cliff_richard;309
+fabolous;310
+sonny_boy_williamson_ii;311
+tinie_tempah;312
+claude_debussy;313
+jewel;314
+bob_seger;315
+otis_rush;316
+mikhail_pletnev;317
+gene_autry;318
+blind_blake;319
+ethel_waters;320
+jackie_wilson;321
+big_mama_thornton;322
+mary_wells;323
+the_mills_brothers;324
+rage_against_the_machine;325
+barbra_streisand;326
+frank_crumit;327
+the_clash;328
+dion;329
+all_4_one;330
+t_i;331
+orchestral_manoeuvres_in_the_dark;332
+culture;333
+josh_turner;334
+one_direction;335
+the_rolling_stones;336
+half_pint;337
+the_searchers;338
+vangelis;339
+jackson_browne;340
+the_beach_boys;341
+train;342
+prince_buster;343
+tavares;344
+eve;345
+bessie_smith;346
+trace_adkins;347
+bay_city_rollers;348
+david_allan_coe;349
+rascal_flatts;350
+petula_clark;351
+10cc;352
+50_cent;353
+the_oak_ridge_boys;354
+barry_manilow;355
+truls_m_rk;356
+morrissey;357
+gerry_the_pacemakers;358
+kay_starr;359
+alborosie;360
+engelbert_humperdinck;361
+new_order;362
+tony_bennett;363
+stereophonics;364
+jimmy_reed;365
+akon;366
+echo_the_bunnymen;367
+jamiroquai;368
+stevie_ray_vaughan;369
+ben_e_king;370
+cheap_trick;371
+dusty_springfield;372
+mel_tillis;373
+damian_marley;374
+ruth_etting;375
+westlife;376
+diana_ross;377
+the_shirelles;378
+frankie_laine;379
+sarah_vaughan;380
+ray_price;381
+gordon_lightfoot;382
+eddie_cantor;383
+the_byrds;384
+gary_numan;385
+bonnie_tyler;386
+aswad;387
+brownie_mcghee;388
+joshua_bell;389
+manic_street_preachers;390
+mr_vegas;391
+tanya_tucker;392
+marvin_gaye_tammi_terrell;393
+the_staple_singers;394
+ry_cooder;395
+john_mayall;396
+martina_mcbride;397
+anner_bylsma;398
+magic_sam;399
+avril_lavigne;400
+robert_nighthawk;401
+the_coasters;402
+ace_of_base;403
+joseph_szigeti;404
+artie_shaw;405
+big_bill_broonzy;406
+the_ventures;407
+rick_astley;408
+les_paul;409
+leonid_kogan;410
+hall_oates;411
+three_dog_night;412
+charley_pride;413
+paul_mccartney;414
+alfred_cortot;415
+crystal_gayle;416
+taj_mahal;417
+mary_j_blige;418
+leann_rimes;419
+mildred_bailey;420
+ll_cool_j;421
+lobo;422
+blake_shelton;423
+matt_haimovitz;424
+beastie_boys;425
+johannes_brahms;426
+martha_and_the_vandellas;427
+lou_rawls;428
+the_righteous_brothers;429
+rosalyn_tureck;430
+vince_gill;431
+gary_moore;432
+bad_company;433
+marty_robbins;434
+russ_morgan;435
+david_cassidy;436
+dj_khaled;437
+leonard_pennario;438
+glenn_miller;439
+kings_of_leon;440
+afrojack;441
+cyndi_lauper;442
+trisha_yearwood;443
+diamond_rio;444
+patty_loveless;445
+sugar_minott;446
+willie_nelson;447
+thomas_newman;448
+georgia_gibbs;449
+eric_carmen;450
+ricky_skaggs;451
+earth_wind_fire;452
+peter_tosh;453
+joe_bonamassa;454
+lonnie_johnson;455
+missy_elliott;456
+nicki_minaj;457
+luther_allison;458
+grigory_sokolov;459
+luke_bryan;460
+alanis_morissette;461
+tracy_lawrence;462
+sonny_james;463
+brandy;464
+michael_bolton;465
+boswell_sisters;466
+antonio_vivaldi;467
+fritz_kreisler;468
+elvis_costello_the_attractions;469
+woody_herman;470
+korn;471
+hans_zimmer;472
+the_osmonds;473
+electric_light_orchestra;474
+t_rex;475
+van_cliburn;476
+harry_james;477
+glee_cast;478
+simple_minds;479
+abba;480
+the_jam;481
+tanya_stephens;482
+the_statler_brothers;483
+craig_david;484
+moby;485
+xxxtentacion;486
+paul_simon;487
+magic_slim;488
+natasha_bedingfield;489
+ennio_morricone;490
+carpenters;491
+joe_tex;492
+cocoa_tea;493
+the_glee_cast;494
+martha_argerich;495
+charlie_rich;496
+memphis_minnie;497
+sheryl_crow;498
+jan_dean;499
+nick_cave_and_the_bad_seeds;500
+del_shannon;501
+lee_greenwood;502
+gary_lewis_the_playboys;503
+marcelle_meyer;504
+count_basie;505
+harold_melvin_the_blue_notes;506
+the_fray;507
+the_prodigy;508
+alicia_keys;509
+conway_twitty;510
+barrington_levy;511
+b_o_b;512
+tritonal;513
+mario_lanza;514
+mac_davis;515
+billy_murray;516
+igor_stravinsky;517
+pretenders;518
+jordin_sparks;519
+alice_in_chains;520
+ohio_players;521
+rick_springfield;522
+jimmie_rodgers;523
+buju_banton;524
+linda_ronstadt;525
+george_benson;526
+2pac;527
+soulja_boy;528
+the_flaming_lips;529
+ignacy_jan_paderewski;530
+gloria_estefan;531
+ariana_grande;532
+black_uhuru;533
+tony_pastor;534
+sublime;535
+snow_patrol;536
+daft_punk;537
+johnny_winter;538
+robbie_williams;539
+eddie_rabbitt;540
+james_cotton;541
+brad_paisley;542
+manfred_mann;543
+bassnectar;544
+margaret_whiting;545
+sam_cooke;546
+robert_cray;547
+the_beautiful_south;548
+barbara_mandrell;549
+dick_haymes;550
+creedence_clearwater_revival;551
+chic;552
+ray_charles;553
+carter_family;554
+ti_sto;555
+survivor;556
+c_line_dion;557
+sergei_prokofiev;558
+b_la_bart_k;559
+the_congos;560
+yefim_bronfman;561
+laidback_luke;562
+darius_rucker;563
+ray_anthony;564
+incubus;565
+carole_king;566
+james_brown;567
+swv;568
+bruno_mars;569
+aphex_twin;570
+nitty_gritty_dirt_band;571
+gustav_mahler;572
+the_shadows;573
+the_moody_blues;574
+reverend_gary_davis;575
+sia;576
+gaetano_donizetti;577
+earl_hooker;578
+the_commodores;579
+maria_jo_o_pires;580
+eric_donaldson;581
+elmore_james;582
+sean_kingston;583
+don_carlos;584
+linkin_park;585
+jay_the_americans;586
+grand_funk_railroad;587
+jls;588
+frankie_valli;589
+lefty_frizzell;590
+en_vogue;591
+the_cure;592
+perry_como;593
+johnny_mercer;594
+stevie_wonder;595
+ernest_tubb;596
+ramin_djawadi;597
+ashanti;598
+rosemary_clooney;599
+anne_sophie_mutter;600
+helen_forrest;601
+augustus_pablo;602
+the_carpenters;603
+claudio_arrau;604
+bob_dylan;605
+joe_simon;606
+culture_club;607
+the_ipana_troubadors;608
+jennifer_lopez;609
+karyn_white;610
+joe;611
+sarah_mclachlan;612
+j_geils_band;613
+dean_martin;614
+hank_snow;615
+clyde_mcphatter;616
+tlc;617
+beck;618
+jimmy_dean;619
+roy_acuff;620
+outkast;621
+freddie_mcgregor;622
+gus_arnheim;623
+gramatik;624
+merle_haggard;625
+steve_lawrence;626
+ma_rainey;627
+jimmy_dorsey;628
+johnny_paycheck;629
+arthur_rubinstein;630
+talking_heads;631
+capleton;632
+les_brown;633
+leonard_bernstein;634
+bobby_goldsboro;635
+kendrick_lamar;636
+lenny_kravitz;637
+nat_shilkret;638
+toby_keith;639
+junior_wells;640
+billy_j_kramer_the_dakotas;641
+peter_paul_mary;642
+armand_van_helden;643
+h_sker_d_;644
+alabama;645
+eminem;646
+felix_mendelssohn_bartholdy;647
+richard_goode;648
+pink_floyd;649
+sara_evans;650
+lonnie_donegan;651
+boney_m;652
+deadmau5;653
+lee_ann_womack;654
+eric_clapton;655
+ray_parker_jr;656
+etta_james;657
+the_white_stripes;658
+gary_u_s_bonds;659
+glen_gray_and_the_casa_loma_orchestra;660
+bonnie_raitt;661
+soulja_boy_tell_em;662
+bobby_rydell;663
+carly_simon;664
+koko_taylor;665
+gregory_isaacs;666
+red_hot_chili_peppers;667
+josef_suk;668
+clint_black;669
+buddy_clark;670
+tool;671
+pierre_fournier;672
+alice_cooper;673
+lucky_dube;674
+jorge_bolet;675
+don_diablo;676
+r_l_burnside;677
+klaus_badelt;678
+dolly_parton;679
+james_horner;680
+green_day;681
+henry_burr;682
+the_doors;683
+roxette;684
+louis_armstrong;685
+jerry_butler;686
+louis_prima;687
+paul_van_dyk;688
+dennis_brown;689
+toni_braxton;690
+jerry_lee_lewis;691
+donna_summer;692
+percy_faith;693
+willie_dixon;694
+elvis_costello;695
+youri_egorov;696
+dmitri_shostakovich;697
+webb_pierce;698
+monica;699
+pierre_laurent_aimard;700
+muddy_waters;701
+garth_brooks;702
+boyz_ii_men;703
+kris_kristofferson;704
+duane_eddy;705
+coolio;706
+gidon_kremer;707
+eurythmics;708
+randy_travis;709
+collie_buddz;710
+faith_evans;711
+matisyahu;712
+brooks_dunn;713
+lulu;714
+fletcher_henderson;715
+eek_a_mouse;716
+shlomo_mintz;717
+ray_noble;718
+bill_withers;719
+ub40;720
+b_j_thomas;721
+mariah_carey;722
+olivia_newton_john;723
+jackson_5;724
+michael_rose;725
+three_days_grace;726
+glen_campbell;727
+keith_urban;728
+onerepublic;729
+ted_weems;730
+johnny_desmond;731
+billy_preston;732
+billy_vaughn;733
+otis_spann;734
+marion_harris;735
+richard_marx;736
+frankie_carle;737
+xavier_cugat;738
+zz_top;739
+faz_l_say;740
+new_york_philharmonic;741
+teresa_brewer;742
+pixies;743
+david_oistrakh;744
+rory_gallagher;745
+ted_lewis_his_band;746
+patsy_cline;747
+tracy_byrd;748
+daniel_shafran;749
+johnny_horton;750
+billy_ocean;751
+don_mclean;752
+ginuwine;753
+kiss;754
+kaskade;755
+the_hold_steady;756
+rod_stewart;757
+rudolf_serkin;758
+the_wanted;759
+herman_s_hermits;760
+rusty_draper;761
+jay_z;762
+mississippi_john_hurt;763
+linval_thompson;764
+billie_holiday;765
+lawrence_welk;766
+john_lee_hooker;767
+the_offspring;768
+ramones;769
+ronan_keating;770
+sir_clifford_michael_curzon;771
+dierks_bentley;772
+bush;773
+texas;774
+gene_austin;775
+after_all;776
+depeche_mode;777
+eddy_arnold;778
+ferry_corsten;779
+inxs;780
+jimmie_lunceford;781
+johann_sebastian_bach;782
+maxi_priest;783
+travis_tritt;784
+girls_aloud;785
+meat_loaf;786
+brenda_lee;787
+modest_mouse;788
+glenn_gould;789
+mc_hammer;790
+erik_satie;791
+the_andrews_sisters;792
+alicia_de_larrocha;793
+america;794
+charles_harrison;795
+georges_cziffra;796
+peter_serkin;797
+gladys_knight_the_pips;798
+anton_n_dvo_k;799
+emil_gilels;800
+myra_hess;801
+portugal_the_man;802
+rufus;803
+sergei_rachmaninoff;804
+andr_previn;805
+natalie_cole;806
+the_killers;807
+helen_reddy;808
+billy_idol;809
+shawn_mendes;810
+john_lennon;811
+joan_jett;812
+jimmy_cliff;813
+taylor_swift;814
+shura_cherkassky;815
+marvin_gaye;816
+the_drifters;817
+paul_revere_the_raiders;818
+donovan;819
+alma_gluck;820
+fleetwood_mac;821
+jo_stafford;822
+samson_fran_ois;823
+vaughn_monroe;824
+tennessee_ernie_ford;825
+whitesnake;826
+iron_maiden;827
+blind_boy_fuller;828
+eric_church;829
+atomic_kitten;830
+harry_nilsson;831
+paul_specht;832
+bob_marley;833
+les_baxter;834
+aretha_franklin;835
+jessie_j;836
+robert_johnson;837
+maroon_5;838
+sheena_easton;839
+george_strait;840
+bruce_springsteen;841
+oasis;842
+jack_white;843
+fatboy_slim;844
+the_jesus_and_mary_chain;845
+hank_williams_jr;846
+bill_haley_his_comets;847
+charlie_daniels;848
+pearl_jam;849
+the_stylistics;850
+sean_paul;851
+david_essex;852
+zino_francescatti;853
+bruce_hornsby;854
+system_of_a_down;855
+scott_joplin;856
+benny_benassi;857
+mitch_miller;858
+gloria_gaynor;859
+steve_winwood;860
+drake;861
+bone_thugs_n_harmony;862
+nickelback;863
+john_mellencamp;864
+gene_pitney;865
+chet_atkins;866
+mark_ronson;867
+kc_the_sunshine_band;868
+chingy;869
+maurice_ravel;870
+infected_mushroom;871
+walter_gieseking;872
+pavement;873
+deniece_williams;874
+tommy_james_the_shondells;875
+harry_belafonte;876
+robert_schuman;877
+george_jones;878
+whitney_houston;879
+the_ames_brothers;880
+max_romeo;881
+patti_page;882
+p_nk;883
+vienna_philharmonic;884
+violent_femmes;885
+sly_the_family_stone;886
+johnny_marvin;887
+sammy_davis_jr_;888
+nirvana;889
+the_turtles;890
+the_manhattans;891
+the_supremes;892
+gregg_allman;893
+alton_ellis;894
+giacomo_puccini;895
+artur_schnabel;896
+2_unlimited;897
+michael_rabin;898
+toto;899
+lil_kim;900
+shinedown;901
+teddy_wilson;902
+alpha_blondy;903
+the_guess_who;904
+my_chemical_romance;905
+matchbox_twenty;906
+wolfgang_gartner;907
+ed_sheeran;908
+faron_young;909
+fats_waller;910
+kate_smith;911
+loretta_lynn;912
+antonio_meneses;913
+tony_martin;914
+frankie_vaughan;915
+the_mamas_the_papas;916
+david_bowie;917
+erskine_hawkins;918
+the_four_lads;919
+queens_of_the_stone_age;920
+the_human_league;921
+mud;922
+bread;923
+dixie_chicks;924
+backstreet_boys;925
+connie_francis;926
+the_black_keys;927
+paul_weller;928
+pieter_wispelwey;929
+the_bachelors;930
+janet_jackson;931
+the_association;932
+edwin_starr;933
+hank_williams_jr_;934
+steven_isserlis;935
+the_doobie_brothers;936
+jimi_hendrix;937
+sousa_s_band;938
+ricky_martin;939
+the_cranberries;940
+paul_whiteman;941
+the_saturdays;942
+keb_mo;943
+the_psychedelic_furs;944
+boyzone;945
+the_chemical_brothers;946
+johnny_tillotson;947
+joe_cocker;948
+babyface;949
+wet_wet_wet;950
+tom_petty;951
+lesley_gore;952
+deorro;953
+war;954
+wolfgang_schneiderhan;955
+isaac_hayes;956
+beyonc;957
+gordon_jenkins;958
+flo_rida;959
+captain_tennille;960
+adolf_busch;961
+bobby_bare;962
+johnnie_taylor;963
+above_beyond;964
+robin_thicke;965
+ja_rule;966
+barry_white;967
+florence_the_machine;968
+the_jimi_hendrix_experience;969
+gabrielle;970
+leon_fleisher;971
+yehudi_menuhin;972
+george_harrison;973
+kate_bush;974
+u2;975
+nancy_sinatra;976
+the_everly_brothers;977
+peter_green;978
+cat_stevens;979
+the_new_seekers;980
+itzhak_perlman;981
+red_nichols_his_five_pennies;982
+the_script;983
+ijahman_levi;984
+rihanna;985
+joan_jett_and_the_blackhearts;986
+the_dorsey_brothers_orchestra;987
+charlie_barnet;988
+lmfao;989
+112;990
+kelis;991
+john_mccormack;992
+guns_n_roses;993
+gordon_macrae;994
+buddy_holly;995
+franz_schubert;996
+tommy_roe;997
+garbage;998
+bow_wow;999
+belle_and_sebastian;1000
+paul_tortelier;1001
+johnny_nash;1002
+lazar_berman;1003
+avicii;1004
+5_seconds_of_summer;1005
+mississippi_fred_mcdowell;1006
+ivan_moravec;1007
+little_richard;1008
+john_ogdon;1009
+dr_dre;1010
+andy_russell;1011
+sex_pistols;1012
+puddle_of_mudd;1013
+june_carter_cash;1014
+freddie_king;1015
+dino_ciani;1016
+h_sker_d;1017
+nicky_romero;1018
+shania_twain;1019
+major_lazer;1020
+usher;1021
+madonna;1022
+nsync;1023
+the_o_jays;1024
+faith_hill;1025
+terence_trent_d_arby;1026
+gene_chandler;1027
+lisa_stansfield;1028
+pharrell_williams;1029
+david_geringas;1030
+dr_alimantado;1031
+jack_scott;1032
+vladimir_horowitz;1033
+genesis;1034
+jessica_simpson;1035
+peter_gabriel;1036
+spike_jones_and_his_city_slickers;1037
+leonard_rose;1038
+richard_strauss;1039
+mgmt;1040
+blink_182;1041
+ella_fitzgerald;1042
+carrie_underwood;1043
+evgeny_kissin;1044
+olly_murs;1045
+ky_mani_marley;1046
+rose_royce;1047
+desmond_dekker;1048
+montell_jordan;1049
+little_river_band;1050
+t_pain;1051
+gary_allan;1052
+chubby_checker;1053
+the_diamonds;1054
+creed;1055
+louis_jordan;1056
+vernon_dalhart;1057
+gilbert_o_sullivan;1058
+ziggy_marley;1059
+irene_cara;1060
+katy_perry;1061
+all_saints;1062
+benjamin_britten;1063
+p_m_dawn;1064
+iona_brown;1065
+the_lettermen;1066
+natalie_imbruglia;1067
+viktoria_mullova;1068
+zac_brown_band;1069
+bj_rk;1070
+mutabaruka;1071
+jermaine_jackson;1072
+the_impressions;1073
+spandau_ballet;1074
+luther_vandross;1075
+thompson_twins;1076
+travis_scott;1077
+roots_radics;1078
+madness;1079
+skrillex;1080
+the_smiths;1081
+van_halen;1082
+the_four_seasons;1083
+j_b_lenoir;1084
+johnny_rivers;1085
+bob_marley_the_wailers;1086
+albert_collins;1087
+arctic_monkeys;1088
+billy_joel;1089
+joe_nichols;1090
+chicago;1091
+coldplay;1092
+blue;1093
+papa_roach;1094
+r_e_m_;1095
+eddie_money;1096
+ellie_goulding;1097
+gym_class_heroes;1098
+the_charlie_daniels_band;1099
+bette_midler;1100
+hubert_sumlin;1101
+phil_harris;1102
+arturo_benedetti_michelangeli;1103
+steps;1104
+simon_and_garfunkel;1105
+gyptian;1106
+britney_spears;1107
+lang_lang;1108
+blackstreet;1109
+the_dave_clark_five;1110
+the_chi_lites;1111
+wolfgang_amadeus_mozart;1112
+johnnie_ray;1113
+elvis_presley;1114
+red_norvo;1115
+the_velvet_underground;1116
+nelly_furtado;1117
+connee_boswell;1118
+soundgarden;1119
+dr_hook_the_medicine_show;1120
+led_zeppelin;1121
+boney_m_;1122
+gregor_piatigorsky;1123
+john_anderson;1124
+joni_james;1125
+the_all_american_rejects;1126
+marcia_griffiths;1127
+ben_bernie;1128
+eric_prydz;1129
+the_three_suns;1130
+albert_king;1131
+jodeci;1132
+inner_circle;1133
+don_cornell;1134
+jeff_healey;1135
+brian_mcknight;1136
+the_stranglers;1137
+adam_faith;1138
+francis_poulenc;1139
+fall_out_boy;1140
+the_jets;1141
+groundation;1142
+randy_newman;1143
+jane_s_addiction;1144
+bobby_vinton;1145
+guiomar_novaes;1146
+bo_diddley;1147
+philadelphia_orchestra;1148
+garrick_ohlsson;1149
+evanescence;1150
+wilhelm_backhaus;1151
+the_platters;1152
+the_seekers;1153
+selena_gomez;1154
+eddy_grant;1155
+jan_garber;1156
+gioacchino_rossini;1157
+christina_aguilera;1158
+stevie_nicks;1159
+color_me_badd;1160
+sandie_shaw;1161
+migos;1162
+the_allman_brothers_band;1163
+will_young;1164
+new_edition;1165
+radiohead;1166
+orbital;1167
+the_judds;1168
+dan_fogelberg;1169
+pascal_rog_;1170
+leona_lewis;1171
+charlie_patton;1172
+yann_tiersen;1173
+the_abyssinians;1174
+frank_ifield;1175
+john_mayer;1176
+roger_wolfe_kahn;1177
+the_mcguire_sisters;1178
+junior_reid;1179
+bukka_white;1180
+macy_gray;1181
+billy_currington;1182
+the_tremeloes;1183
+ac_dc;1184
+little_walter;1185
+gorillaz;1186
+the_hilltoppers;1187
+gil_shaham;1188
+little_mix;1189
+supertramp;1190
+third_world;1191
+david_garrett;1192
+the_lovin_spoonful;1193
+chris_lake;1194
+louis_lortie;1195
+shakira;1196
+vanessa_williams;1197
+the_bangles;1198
+mark_chesnutt;1199
+keith_whitley;1200
+roxy_music;1201
+betty_hutton;1202
+cheryl_cole;1203
+the_notorious_b_i_g_;1204
+kesha;1205
+boston;1206
+gerald_moore;1207
+heinrich_schiff;1208
+joy_division;1209
+kenny_loggins;1210
+nilsson;1211
+counting_crows;1212
+simply_red;1213
+aaron_tippin;1214
+sweet;1215
+sugababes;1216
+florida_georgia_line;1217
+lonnie_mack;1218
+jacob_miller;1219
+the_police;1220
+the_chordettes;1221
+fats_domino;1222
+willy_deville;1223
+bob_crosby;1224
+roy_clark;1225
+dinu_lipatti;1226
+alison_krauss;1227
+frankie_avalon;1228
+ella_mae_morse;1229
+linton_kwesi_johnson;1230
+glenn_frey;1231
+rudy_vall_e_his_connecticut_yankees;1232
+dj_snake;1233
+toots_the_maytals;1234
+annie_lennox;1235
+paramore;1236
+john_browning;1237
+sister_sledge;1238
+julian_lloyd_webber;1239
+kanye_west;1240
+christopher_cross;1241
+professor_longhair;1242
+flume;1243
+sophie_tucker;1244
+third_eye_blind;1245
+beyonc_;1246
+weezer;1247
+kelly_clarkson;1248
+take_that;1249
+ludwig_van_beethoven;1250
+the_cars;1251
+george_michael;1252
+chaka_khan;1253
+arty;1254
+steel_pulse;1255
+bunny_wailer;1256
+the_troggs;1257
+nick_lucas;1258
+keb_mo_;1259
+rise_against;1260
+mickey_gilley;1261
+poison;1262
+debbie_gibson;1263
+kim_wilde;1264
+sammy_davis_jr;1265
+jonas_brothers;1266
+cage_the_elephant;1267
+son_house;1268
+birth_control;1269
+faithless;1270
+cher;1271
+seether;1272
+new_kids_on_the_block;1273
+rage;1274
+richie_spice;1275
+sam_smith;1276
+the_marvelettes;1277
+the_jackson_5;1278
+bobbie_gentry;1279
+macklemore_ryan_lewis;1280
+s_club_7;1281
+billy_jones_ernest_hare;1282
+arcana;1283
+blasterjaxx;1284
+will_smith;1285
+ray_miller;1286
+sugar_ray;1287
+ziggy_marley_the_melody_makers;1288
+aaron_copland;1289
+diddy;1290
+daughtry;1291
+beach_house;1292
+the_dillinger_escape_plan;1293
+lonestar;1294
+foreigner;1295
+lionel_richie;1296
+roberta_flack;1297
+the_carter_family;1298
+demi_lovato;1299
+joseph_haydn;1300
+sander_van_doorn;1301
+underworld;1302
+deborah_cox;1303
+the_grass_roots;1304
+bananarama;1305
+iggy_azalea;1306
+3_doors_down;1307
+the_partridge_family;1308
+lead_belly;1309
+johnny_long;1310
+jagged_edge;1311
+the_derek_trucks_band;1312
+eddie_fisher;1313
+pablo_casals;1314
+the_original_dixieland_jazz_band;1315
+miley_cyrus;1316
+john_powell;1317
+the_black_eyed_peas;1318
+edvard_grieg;1319
+maurizio_pollini;1320
+nathan_milstein;1321
+twenty_one_pilots;1322
+rudolf_firku_n_;1323
+nine_inch_nails;1324
+the_four_aces;1325
+jimmy_rogers;1326
+salt_n_pepa;1327
+yellow_claw;1328
+the_strokes;1329
+bobby_brown;1330
+juelz_santana;1331
+staind;1332
+bj_rn_ulvaeus_benny_andersson;1333
+the_j_geils_band;1334
+the_muppets;1335
+gwen_stefani;1336
+fedde_le_grand;1337
+imagine_dragons;1338
+leo_sayer;1339
+the_animals;1340
+the_b_52_s;1341
+furry_lewis;1342
+ciara;1343
+larry_clinton;1344
+dire_straits;1345
+pato_banton_the_reggae_revol;1346
+goodie_mob;1347
+disclosure;1348
+georges_bizet;1349
+sonny_terry_brownie_mcghee;1350
+jim_croce;1351
+nelson;1352
+jason_derulo;1353
+guy_mitchell;1354
+the_fontane_sisters;1355
+sonny_boy_williamson_i;1356
+mirah;1357
+jimmy_rushing;1358
+john_michael_montgomery;1359
+michael_nesmith;1360
+george_clinton;1361
+burt_bacharach;1362
+v6;1363
+slim_thug;1364
+belinda_carlisle;1365
+philip_glass;1366
+slade;1367
+pete_townshend;1368
+oingo_boingo;1369
+andrew_lloyd_webber;1370
+the_collectors;1371
+boy_george;1372
+utada_hikaru;1373
+mel_torm;1374
+diana_krall;1375
+melanie_c;1376
+john_hammond;1377
+peter_green_splinter_group;1378
+trouble;1379
+g_unit;1380
+ferlin_husky;1381
+arcade_fire;1382
+latino;1383
+krayzie_bone;1384
+man;1385
+dave_clark_five;1386
+the_stone_roses;1387
+young_jeezy;1388
+blood_sweat_tears;1389
+kumikameli;1390
+dmx;1391
+ice_cube;1392
+eagles;1393
+jill_scott;1394
+xtc;1395
+peggy_march;1396
+michael_bubl;1397
+raimon;1398
+two_mix;1399
+dulce_pontes;1400
+the_unseen;1401
+hank_locklin;1402
+the_notorious_b_i_g;1403
+bumblefoot;1404
+the_busters;1405
+rick_ross;1406
+tegan_and_sara;1407
+skeeter_davis;1408
+curtis_mayfield;1409
+sade;1410
+wings;1411
+lorrie_morgan;1412
+saga;1413
+a_r_rahman;1414
+martha_wainwright;1415
+nas;1416
+will_i_am;1417
+kirsty_maccoll;1418
+angel;1419
+dave_davies;1420
+iggy_pop;1421
+jojo;1422
+sammy_hagar;1423
+ray_davies;1424
+and_one;1425
+neil_young;1426
+mikael_wiehe;1427
+the_cardigans;1428
+cage;1429
+dottie_west;1430
+keri_hilson;1431
+johnny_hallyday;1432
+bill_nelson;1433
+fifteen;1434
+deen;1435
+bobby_v;1436
+lil_yachty;1437
+too_hort;1438
+bernard_lavilliers;1439
+hank_thompson;1440
+the_chieftains;1441
+daryl_hall;1442
+antonio_carlos_jobim;1443
+aventura;1444
+the_the;1445
+van_morrison;1446
+wynonna_judd;1447
+gomez;1448
+charles_aznavour;1449
+m83;1450
+gnr;1451
+all;1452
+emmylou_harris;1453
+lee_hazlewood;1454
+mew;1455
+uriah_heep;1456
+yoko_ono;1457
+abw_rts;1458
+john_legend;1459
+d12;1460
+kitty_wells;1461
+timbiriche;1462
+shel_silverstein;1463
+cam_ron;1464
+rosanne_cash;1465
+2_chainz;1466
+tricky;1467
+8ball_mjg;1468
+flatt_scruggs;1469
+bill_anderson;1470
+emil_ana_torrini;1471
+ufo;1472
+mos_def;1473
+danzig;1474
+juan_gabriel;1475
+common;1476
+raekwon;1477
+france_gall;1478
+nicole_scherzinger;1479
+r_yksopp;1480
+sammie;1481
+lena_horne;1482
+david_byrne;1483
+paul_williams;1484
+josh_groban;1485
+the_gathering;1486
+frank_boeijen;1487
+scooter;1488
+steve_wariner;1489
+mika;1490
+pete_seeger;1491
+tex_ritter;1492
+warrant;1493
+porter_wagoner;1494
+field_music;1495
+three_6_mafia;1496
+jim_jones;1497
+daniel_o_donnell;1498
+brentalfloss;1499
+wyclef_jean;1500
+hey;1501
+bizzy_bone;1502
+the_mccalmans;1503
+blues_traveler;1504
+massive_attack;1505
+woody_guthrie;1506
+art_garfunkel;1507
+andrea_bocelli;1508
+david_crosby;1509
+dream;1510
+soul_asylum;1511
+natalie_merchant;1512
+shawn_colvin;1513
+jonny_lang;1514
+funeral_for_a_friend;1515
+boz_scaggs;1516
+example;1517
+lionel_hampton;1518
+the_tubes;1519
+marc_anthony;1520
+good_riddance;1521
+moonlight;1522
+marc_almond;1523
+rza;1524
+die_rzte;1525
+rbd;1526
+alejandro_fern_ndez;1527
+wanda_jackson;1528
+lara_fabian;1529
+julio_iglesias;1530
+jeff_beck;1531
+peabo_bryson;1532
+no_fun_at_all;1533
+prong;1534
+canibus;1535
+krs_one;1536
+u_s_bombs;1537
+trust;1538
+stonewall_jackson;1539
+jos_feliciano;1540
+m_a;1541
+polysics;1542
+n_e_r_d;1543
+sesame_street;1544
+lio;1545
+myl_ne_farmer;1546
+iris_dement;1547
+lily_allen;1548
+spoken;1549
+architects;1550
+jack_johnson;1551
+molly_hatchet;1552
+cypress_hill;1553
+future;1554
+the_nits;1555
+per_gessle;1556
+live;1557
+beau;1558
+deana_carter;1559
+lil_flip;1560
+fran_oise_hardy;1561
+billy_ray_cyrus;1562
+pepper;1563
+run_d_m_c;1564
+levon_helm;1565
+insane_clown_posse;1566
+stars;1567
+jean_michel_jarre;1568
+thunder;1569
+juanes;1570
+simple_plan;1571
+method_man;1572
+smash_mouth;1573
+meek_mill;1574
+fat_joe;1575
+kenny_wayne_shepherd;1576
+jhen_aiko;1577
+the_manhattan_transfer;1578
+joss_stone;1579
+cee_lo_green;1580
+tyrese;1581
+charlotte_martin;1582
+rodney_crowell;1583
+acappella;1584
+die_prinzen;1585
+pentatonix;1586
+abney_park;1587
+smooth_mcgroove;1588
+the_magnetic_fields;1589
+the_nylons;1590
+wise_guys;1591
+coil;1592
+do_as_infinity;1593
+lords_of_acid;1594
+the_church;1595
+chris_rea;1596
+jota_quest;1597
+miguel_bos;1598
+gang_starr;1599
+masta_ace;1600
+brand_new;1601
+mac_miller;1602
+cathedral;1603
+corrosion_of_conformity;1604
+country_joe_mcdonald;1605
+eric_johnson;1606
+grateful_dead;1607
+janis_joplin;1608
+john_miles;1609
+king_gizzard_the_lizard_wizard;1610
+ted_nugent;1611
+combichrist;1612
+alkaline_trio;1613
+anathema;1614
+angus_julia_stone;1615
+anna_ternheim;1616
+anthony_phillips;1617
+steve_hackett;1618
+aviators;1619
+banda_calypso;1620
+blue_stahli;1621
+the_boys;1622
+capital_inicial;1623
+city_and_colour;1624
+colin_hay;1625
+collective_soul;1626
+dashboard_confessional;1627
+david_rovics;1628
+david_usher;1629
+die_toten_hosen;1630
+funny_van_dannen;1631
+dirty_heads;1632
+tech_n9ne;1633
+elisa;1634
+emmerson_nogueira;1635
+engenheiros_do_hawaii;1636
+eric_bibb;1637
+maria_muldaur;1638
+panic_at_the_disco;1639
+punchline;1640
+godsmack;1641
+self;1642
+heideroosjes;1643
+sinner;1644
+heather_nova;1645
+hoobastank;1646
+quietdrive;1647
+hyde;1648
+jaguares;1649
+bert_jansch;1650
+jonatha_brooke;1651
+joni_mitchell;1652
+the_band;1653
+josh_garrels;1654
+josh_woodward;1655
+william_fitzsimmons;1656
+katie_melua;1657
+jamie_cullum;1658
+kristin_hersh;1659
+kt_tunstall;1660
+legi_o_urbana;1661
+the_zombies;1662
+francesco_de_gregori;1663
+m_ward;1664
+beth_orton;1665
+magnum;1666
+motorpsycho;1667
+marillion;1668
+jars_of_clay;1669
+mason_jennings;1670
+matt_nathanson;1671
+matthew_good;1672
+edguy;1673
+gamma_ray;1674
+minus_the_bear;1675
+mohsen_namjoo;1676
+nerina_pallot;1677
+never_shout_never;1678
+regina_spektor;1679
+passenger;1680
+paul_kelly;1681
+the_style_council;1682
+peter_hammill;1683
+phantom_planet;1684
+phil_keaggy;1685
+richard_thompson;1686
+said_the_whale;1687
+samsas_traum;1688
+senses_fail;1689
+sevendust;1690
+seventh_day_slumber;1691
+joan_baez;1692
+sister_hazel;1693
+slightly_stoopid;1694
+sophie_zelmani;1695
+suzanne_vega;1696
+tatiana;1697
+teoman;1698
+the_choir;1699
+charlotte_church;1700
+darlene_zschech;1701
+the_front_bottoms;1702
+the_maine;1703
+the_white_buffalo;1704
+little_big_town;1705
+threshold;1706
+tourniquet;1707
+everything_but_the_girl;1708
+vertical_horizon;1709
+vonda_shepard;1710
+warren_zevon;1711
+sarah_brightman;1712
+blackfoot;1713
+black_label_society;1714
+z_lia_duncan;1715
+2;1716
+alejandro_lerner;1717
+beth_nielsen_chapman;1718
+mercury_rev;1719
+brian_wilson;1720
+barenaked_ladies;1721
+carbon_leaf;1722
+celtic_woman;1723
+hayley_westenra;1724
+crowded_house;1725
+delta_goodrem;1726
+elbow;1727
+resurrection_band;1728
+nancy_wilson;1729
+janis_ian;1730
+jann_arden;1731
+jill_sobule;1732
+jos_augusto;1733
+xuxa;1734
+k_d_lang;1735
+kim_carnes;1736
+los_lobos;1737
+mandy_moore;1738
+marc_cohn;1739
+maureen_mcgovern;1740
+melissa_manchester;1741
+patti_labelle;1742
+helene_fischer;1743
+laura_pausini;1744
+ivan_lins;1745
+thal_a;1746
+mike_the_mechanics;1747
+paul_carrack;1748
+natasha_st_pier;1749
+michael_mcdonald;1750
+olivia;1751
+dizzee_rascal;1752
+sam_phillips;1753
+serge_gainsbourg;1754
+jane_birkin;1755
+luis_miguel;1756
+sondre_lerche;1757
+stan_ridgway;1758
+susan_boyle;1759
+mike_oldfield;1760
+ces_ria_vora;1761
+agonoize;1762
+funker_vogt;1763
+god_module;1764
+hocico;1765
+nachtmahr;1766
+suicide_commando;1767
+alejandro_escovedo;1768
+southside_johnny_the_asbury_jukes;1769
+broken_social_scene;1770
+vigilantes_of_love;1771
+billy_bragg;1772
+wilco;1773
+frank_black_and_the_catholics;1774
+blue_rodeo;1775
+brandi_carlile;1776
+patty_griffin;1777
+calexico;1778
+cass_mccombs;1779
+chris_knight;1780
+conor_oberst;1781
+corb_lund;1782
+cowboy_junkies;1783
+cracker;1784
+cross_canadian_ragweed;1785
+dave_alvin;1786
+drive_by_truckers;1787
+eleni_mandell;1788
+lucinda_williams;1789
+fred_eaglesmith;1790
+dar_williams;1791
+the_jayhawks;1792
+lana_del_rey;1793
+tim_o_brien;1794
+hank_williams_iii;1795
+james_mcmurtry;1796
+joe_henry;1797
+john_stewart;1798
+josh_ritter;1799
+lambchop;1800
+nanci_griffith;1801
+norma_jean;1802
+lyle_lovett;1803
+matthew_ryan;1804
+my_morning_jacket;1805
+neko_case;1806
+the_new_pornographers;1807
+blue_october;1808
+okkervil_river;1809
+old_97_s;1810
+ray_wylie_hubbard;1811
+richmond_fontaine;1812
+robert_earl_keen;1813
+rocky_votolato;1814
+ryan_adams;1815
+whiskeytown;1816
+son_volt;1817
+steve_earle;1818
+the_avett_brothers;1819
+the_bottle_rockets;1820
+the_felice_brothers;1821
+arlo_guthrie;1822
+the_handsome_family;1823
+the_mavericks;1824
+ten_years_after;1825
+the_walkabouts;1826
+todd_snider;1827
+vic_chesnutt;1828
+iron_wine;1829
+wovenhand;1830
+x;1831
+big_audio_dynamite;1832
+globe;1833
+carter_the_unstoppable_sex_machine;1834
+allison_moorer;1835
+front_line_assembly;1836
+the_national;1837
+the_fall;1838
+public_image_ltd;1839
+public_enemy;1840
+wire;1841
+a_tribe_called_quest;1842
+de_la_soul;1843
+aesop_rock;1844
+buck_65;1845
+caparezza;1846
+childish_gambino;1847
+the_roots;1848
+colbie_caillat;1849
+big_sean;1850
+dj_gruff;1851
+tyler_the_creator;1852
+dokken;1853
+fun_lovin_criminals;1854
+talib_kweli;1855
+jane_air;1856
+k_i_z;1857
+kid_cudi;1858
+jedi_mind_tricks;1859
+celph_titled;1860
+lupe_fiasco;1861
+bun_b;1862
+scarface;1863
+ghostface_killah;1864
+robyn;1865
+rehab;1866
+swollen_members;1867
+styles_p;1868
+the_streets;1869
+wale;1870
+joe_budden;1871
+tank;1872
+10_years;1873
+36_crazyfists;1874
+apocalyptica;1875
+nina_hagen;1876
+anastacia;1877
+black_stone_cherry;1878
+blindside;1879
+breaking_benjamin;1880
+bring_me_the_horizon;1881
+bullet_for_my_valentine;1882
+cave_in;1883
+chevelle;1884
+d_espairsray;1885
+death_angel;1886
+deftones;1887
+demon_hunter;1888
+demon;1889
+devin_townsend_project;1890
+devin_townsend;1891
+doa;1892
+dir_en_grey;1893
+disturbed;1894
+dope;1895
+drowning_pool;1896
+eighteen_visions;1897
+entombed;1898
+faith_no_more;1899
+fear_factory;1900
+fightstar;1901
+five_finger_death_punch;1902
+finger_eleven;1903
+flyleaf;1904
+grinspoon;1905
+guano_apes;1906
+h_blockx;1907
+halestorm;1908
+hamlet;1909
+helmet;1910
+bt;1911
+ill_ni_o;1912
+in_flames;1913
+in_this_moment;1914
+him;1915
+j_b_o;1916
+katatonia;1917
+killswitch_engage;1918
+xzibit;1919
+lacuna_coil;1920
+phish;1921
+limp_bizkit;1922
+living_colour;1923
+viikate;1924
+marilyn_manson;1925
+megaherz;1926
+falco;1927
+melvins;1928
+monster_magnet;1929
+mushroomhead;1930
+nonpoint;1931
+soil;1932
+otep;1933
+p_o_d;1934
+powerman_5000;1935
+primus;1936
+project_86;1937
+red;1938
+kris_allen;1939
+rob_zombie;1940
+ozzy_osbourne;1941
+rollins_band;1942
+saliva;1943
+sepultura;1944
+shihad;1945
+skillet;1946
+skindred;1947
+slipknot;1948
+smile_empty_soul;1949
+danielson;1950
+soilwork;1951
+sonic_syndicate;1952
+static_x;1953
+stone_sour;1954
+taproot;1955
+the_notwist;1956
+the_word_alive;1957
+theory_of_a_deadman;1958
+therapy;1959
+los_tucanes_de_tijuana;1960
+manu_chao;1961
+volbeat;1962
+zebrahead;1963
+hed_p_e;1964
+and_you_will_know_us_by_the_trail_of_dead;1965
+10_000_maniacs;1966
+311;1967
+77s;1968
+yes;1969
+david_lee_roth;1970
+hillsong;1971
+afi;1972
+adam_sandler;1973
+afterhours;1974
+hawkwind;1975
+all_about_eve;1976
+all_time_low;1977
+allison_crowe;1978
+amanda_palmer;1979
+american_music_club;1980
+amplifier;1981
+robert_wyatt;1982
+anberlin;1983
+andrew_bird;1984
+ani_difranco;1985
+apoptygma_berzerk;1986
+apulanta;1987
+arab_strap;1988
+joseph_arthur;1989
+tom_rosenthal;1990
+ash;1991
+asian_kung_fu_generation;1992
+poets_of_the_fall;1993
+babas_nicos;1994
+bayside;1995
+beatsteaks;1996
+ben_folds;1997
+ben_folds_five;1998
+ben_harper;1999
+better_than_ezra;2000
+bettie_serveert;2001
+big_country;2002
+big_head_todd_and_the_monsters;2003
+big_sugar;2004
+billy_talent;2005
+today_is_the_day;2006
+red_flag;2007
+black_rebel_motorcycle_club;2008
+megadeth;2009
+blonde_redhead;2010
+bob_mould;2011
+bodeans;2012
+bowling_for_soup;2013
+buck_tick;2014
+butch_walker;2015
+butthole_surfers;2016
+caf_tacvba;2017
+cake;2018
+camper_van_beethoven;2019
+carmen_consoli;2020
+mario_venuti;2021
+franco_battiato;2022
+catherine_wheel;2023
+catupecu_machu;2024
+cem_adrian;2025
+john_cale;2026
+charlie_brown_jr;2027
+nena;2028
+chumbawamba;2029
+clutch;2030
+cl;2031
+cmx;2032
+coheed_and_cambria;2033
+cold_war_kids;2034
+travis;2035
+coma;2036
+concrete_blonde;2037
+mint_condition;2038
+copeland;2039
+crash_test_dummies;2040
+joe_jackson;2041
+cristian_castro;2042
+curve;2043
+dada;2044
+daniel_amos;2045
+daniel_johnston;2046
+dave_matthews_band;2047
+burning_heads;2048
+david_gray;2049
+david_sylvian;2050
+deacon_blue;2051
+deerhoof;2052
+del_amitri;2053
+dinosaur_jr;2054
+dirty_projectors;2055
+draco_rosa;2056
+duncan_sheik;2057
+jeremy_camp;2058
+edwyn_collins;2059
+eels;2060
+nightwish;2061
+element_of_crime;2062
+embrace;2063
+enter_shikari;2064
+ulver;2065
+everclear;2066
+everlast;2067
+eyeshine;2068
+dio;2069
+faust_o;2070
+feeder;2071
+atmosphere;2072
+filter;2073
+firewater;2074
+fishbone;2075
+fountains_of_wayne;2076
+four_year_strong;2077
+steve_green;2078
+fresno;2079
+gang_of_four;2080
+good_charlotte;2081
+blood_on_the_dance_floor;2082
+graham_coxon;2083
+melissa_etheridge;2084
+tony_joe_white;2085
+guided_by_voices;2086
+robert_pollard;2087
+guster;2088
+elliott_smith;2089
+hedley;2090
+hole;2091
+hollywood_undead;2092
+hot_chip;2093
+l_arc_en_ciel;2094
+ian_brown;2095
+idlewild;2096
+jimmy_eat_world;2097
+fish;2098
+ingrid_michaelson;2099
+inme;2100
+inspiral_carpets;2101
+raf;2102
+james;2103
+jean_leloup;2104
+weird_al_yankovic;2105
+jeff_buckley;2106
+john_frusciante;2107
+dr_john;2108
+pj_harvey;2109
+jonathan_coulton;2110
+juliana_hatfield;2111
+julieta_venegas;2112
+k_s_choice;2113
+kaizers_orchestra;2114
+kargo;2115
+kasabian;2116
+keane;2117
+kevin_coyne;2118
+kevin_devine;2119
+kevin_max;2120
+rich_mullins;2121
+trooper;2122
+suzy_bogguss;2123
+kill_hannah;2124
+kisp_l_s_a_borz;2125
+kult;2126
+my_life_with_the_thrill_kill_kult;2127
+kutless;2128
+la_barranca;2129
+la_ley;2130
+lao_che;2131
+lech_janerka;2132
+les_cowboys_fringants;2133
+les_fatals_picards;2134
+les_rita_mitsouko;2135
+sparks;2136
+lifehouse;2137
+lisa_germano;2138
+ed_harcourt;2139
+lisa_loeb;2140
+liz_phair;2141
+local_h;2142
+lost_dogs;2143
+lostprophets;2144
+love_and_rockets;2145
+lucybell;2146
+lulu_santos;2147
+gabriel_o_pensador;2148
+adam_lambert;2149
+madrugada;2150
+mancha_de_rolando;2151
+manchester_orchestra;2152
+mando_diao;2153
+foetus;2154
+mark_lanegan;2155
+matthew_sweet;2156
+max_mo_park;2157
+mayday_parade;2158
+meat_puppets;2159
+men_without_hats;2160
+meshell_ndegeocello;2161
+midnight_oil;2162
+dance_gavin_dance;2163
+molotov;2164
+ov7;2165
+monkey_majik;2166
+suede;2167
+fernando_ortega;2168
+motion_city_soundtrack;2169
+mudhoney;2170
+mutemath;2171
+mercyme;2172
+m_o_morta;2173
+natalia_lafourcade;2174
+natewantstobattle;2175
+needtobreathe;2176
+split_enz;2177
+sum_41;2178
+no_te_va_gustar;2179
+noir_d_sir;2180
+t_tes_raides;2181
+o_rappa;2182
+o_a_r;2183
+ocean_colour_scene;2184
+omul_cu_obolani;2185
+one_ok_rock;2186
+2raumwohnung;2187
+our_lady_peace;2188
+pain;2189
+panda;2190
+parokya_ni_edgar;2191
+pato_fu;2192
+paul_westerberg;2193
+pere_ubu;2194
+pete_yorn;2195
+peter_murphy;2196
+placebo;2197
+plain_white_t_s;2198
+pop_will_eat_itself;2199
+porcupine_tree;2200
+powderfinger;2201
+cat_power;2202
+casting_crowns;2203
+primal_scream;2204
+m_tley_cr_e;2205
+the_used;2206
+raimundos;2207
+mark_knopfler;2208
+mark_kozelek;2209
+danko_jones;2210
+relient_k;2211
+raffi;2212
+renaud;2213
+richard_hawley;2214
+rickie_lee_jones;2215
+the_shins;2216
+rilo_kiley;2217
+robyn_hitchcock;2218
+mose_allison;2219
+roy_harper;2220
+rucka_rucka_ali;2221
+rx_bandits;2222
+saez;2223
+samiam;2224
+sarah_slean;2225
+say_anything;2226
+scout_niblett;2227
+screaming_females;2228
+shannon_wright;2229
+silverchair;2230
+sin_ad_o_connor;2231
+siouxsie_and_the_banshees;2232
+sixpence_none_the_richer;2233
+skank;2234
+skunk_anansie;2235
+sleater_kinney;2236
+sloan;2237
+social_distortion;2238
+sophie_hunger;2239
+e_40;2240
+steve_wynn;2241
+subsonica;2242
+joe_walsh;2243
+super_furry_animals;2244
+superchunk;2245
+supergrass;2246
+swervedriver;2247
+switchfoot;2248
+dido;2249
+takida;2250
+taking_back_sunday;2251
+teenage_fanclub;2252
+w_a_s_p;2253
+the_afghan_whigs;2254
+the_apples_in_stereo;2255
+the_ataris;2256
+smoking_popes;2257
+the_bluetones;2258
+the_breeders;2259
+the_cat_empire;2260
+the_charlatans_uk;2261
+the_clarks;2262
+guy_clark;2263
+the_comsat_angels;2264
+the_connells;2265
+the_coral;2266
+the_cribs;2267
+the_cult;2268
+bobby_o;2269
+the_mission;2270
+blue_yster_cult;2271
+the_dandy_warhols;2272
+the_dear_hunter;2273
+the_decemberists;2274
+the_early_november;2275
+thievery_corporation;2276
+the_fratellis;2277
+the_gaslight_anthem;2278
+jim_brickman;2279
+falling_up;2280
+the_hives;2281
+the_innocence_mission;2282
+the_jazz_butcher;2283
+the_jesus_lizard;2284
+the_lemonheads;2285
+babyshambles;2286
+the_living_end;2287
+the_matrixx;2288
+the_mother_hips;2289
+the_mountain_goats;2290
+the_muffs;2291
+the_pillows;2292
+the_posies;2293
+the_presidents_of_the_united_states_of_america;2294
+the_rasmus;2295
+the_raveonettes;2296
+the_saints;2297
+the_samples;2298
+bad_religion;2299
+the_smithereens;2300
+the_soundtrack_of_our_lives;2301
+the_tea_party;2302
+mayday;2303
+the_triffids;2304
+the_vines;2305
+the_violet_burning;2306
+the_wallflowers;2307
+testament;2308
+the_divine_comedy;2309
+third_day;2310
+thrice;2311
+tindersticks;2312
+tism;2313
+tit_s;2314
+toad_the_wet_sprocket;2315
+tocotronic;2316
+tom_mcrae;2317
+tori_amos;2318
+tracy_chapman;2319
+trashcan_sinatras;2320
+tre_allegri_ragazzi_morti;2321
+tub_ring;2322
+unkle;2323
+unwritten_law;2324
+uverworld;2325
+vast;2326
+verdena;2327
+veruca_salt;2328
+face_to_face;2329
+virus;2330
+voltaire;2331
+we_the_kings;2332
+the_kooks;2333
+lindisfarne;2334
+seals_crofts;2335
+andy_partridge;2336
+xutos_pontap_s;2337
+yellowcard;2338
+yup;2339
+leevi_and_the_leavings;2340
+zo;2341
+zucchero;2342
+z;2343
+ebnem_ferah;2344
+air;2345
+alice;2346
+boards_of_canada;2347
+brian_eno;2348
+burzum;2349
+daniel_lanois;2350
+enigma;2351
+juana_molina;2352
+lisa_gerrard;2353
+nox_arcana;2354
+renard;2355
+schiller;2356
+sigur_r_s;2357
+steven_wilson;2358
+swans;2359
+wolfgun;2360
+xiu_xiu;2361
+michael_johnson;2362
+montgomery_gentry;2363
+the_stanley_brothers;2364
+john_waite;2365
+shelby_lynne;2366
+judy_collins;2367
+burl_ives;2368
+the_irish_rovers;2369
+david_wilcox;2370
+devendra_banhart;2371
+doc_watson;2372
+bill_monroe;2373
+michael_martin_murphey;2374
+gordon_bok;2375
+asleep_at_the_wheel;2376
+the_browns;2377
+nana_mouskouri;2378
+jerry_jeff_walker;2379
+steve_goodman;2380
+malcolm_holcombe;2381
+malvina_reynolds;2382
+odetta;2383
+tom_paxton;2384
+strawbs;2385
+phil_ochs;2386
+harry_chapin;2387
+ramblin_jack_elliott;2388
+roger_mcguinn;2389
+gene_clark;2390
+mat_kearney;2391
+the_brothers_four;2392
+tom_russell;2393
+townes_van_zandt;2394
+uncle_dave_macon;2395
+delbert_mcclinton;2396
+john_hiatt;2397
+justin_townes_earle;2398
+mark_erelli;2399
+over_the_rhine;2400
+steve_forbert;2401
+manfred_mann_s_earth_band;2402
+mot_rhead;2403
+rudimentary_peni;2404
+illapu;2405
+inti_illimani;2406
+quilapay_n;2407
+v_ctor_jara;2408
+skylark;2409
+adam_green;2410
+cold_chisel;2411
+guy_sebastian;2412
+jefferson_starship;2413
+the_alan_parsons_project;2414
+ali_project;2415
+modern_talking;2416
+animal_collective;2417
+banco_del_mutuo_soccorso;2418
+ben_lee;2419
+bryan_ferry;2420
+buffy_sainte_marie;2421
+colin_blunstone;2422
+cursive;2423
+elysian_fields;2424
+emerson_lake_palmer;2425
+gino_vannelli;2426
+g_rard_manset;2427
+hot_dad;2428
+marina_and_the_diamonds;2429
+ismo_alanko;2430
+kansas;2431
+kari_peitsamo;2432
+laibach;2433
+laurie_anderson;2434
+puhdys;2435
+na_o_zumbi;2436
+roger_waters;2437
+rush;2438
+the_walker_brothers;2439
+hilltop_hoods;2440
+wolfgang_ambros;2441
+erste_allgemeine_verunsicherung;2442
+jacques_brel;2443
+rainhard_fendrich;2444
+tom_waits;2445
+adrian_belew;2446
+anne_clark;2447
+can;2448
+captain_beefheart_and_the_magic_band;2449
+deine_lakaien;2450
+devo;2451
+einst_rzende_neubauten;2452
+frank_zappa;2453
+goethes_erben;2454
+wishbone_ash;2455
+death_cab_for_cutie;2456
+antony_and_the_johnsons;2457
+jandek;2458
+nevermore;2459
+king_crimson;2460
+king_missile;2461
+the_residents;2462
+steeleye_span;2463
+vampire_rodents;2464
+the_walkmen;2465
+dog_fashion_disco;2466
+freak_kitchen;2467
+sigh;2468
+children_of_bodom;2469
+soft_machine;2470
+ara_ketu;2471
+asa_de_guia;2472
+banda_eva;2473
+ivete_sangalo;2474
+chiclete_com_banana;2475
+daniela_mercury;2476
+alejandro_sanz;2477
+timbalada;2478
+juan_luis_guerra;2479
+daddy_yankee;2480
+alceu_valen_a;2481
+luiz_gonzaga;2482
+matia_bazar;2483
+axelle_red;2484
+barbara;2485
+benny_neyman;2486
+gigi_d_agostino;2487
+jacques_higelin;2488
+caetano_veloso;2489
+gal_costa;2490
+jorge_ben;2491
+die_flippers;2492
+nicole;2493
+angra;2494
+reinhard_mey;2495
+wolf_biermann;2496
+florent_pagny;2497
+hannes_wader;2498
+tienne_daho;2499
+henri_salvador;2500
+f_lix_leclerc;2501
+daniel_lavoie;2502
+gerhard_sch_ne;2503
+g_lben_ergen;2504
+georg_kreisler;2505
+herbert_gr_nemeyer;2506
+herman_van_veen;2507
+hildegard_knef;2508
+marlene_dietrich;2509
+iu;2510
+jos_luis_rodr_guez;2511
+juliette_gr_co;2512
+klaus_hoffmann;2513
+konstantin_wecker;2514
+saltatio_mortis;2515
+luigi_tenco;2516
+maria_beth_nia;2517
+adriana_calcanhotto;2518
+marie_lafor_t;2519
+marius_m_ller_westernhagen;2520
+mina;2521
+no_l_coward;2522
+pippo_pollina;2523
+rita_lee;2524
+os_mutantes;2525
+rita_pavone;2526
+roger_whittaker;2527
+al_bano_romina_power;2528
+salvatore_adamo;2529
+simone;2530
+s_rgio_godinho;2531
+udo_j_rgens;2532
+udo_lindenberg;2533
+ulrich_roski;2534
+zaz;2535
+z_ramalho;2536
+fagner;2537
+dith_piaf;2538
+duelo;2539
+espinoza_paz;2540
+fidel_rueda;2541
+la_firma;2542
+la_arrolladora_banda_el_lim_n;2543
+voz_de_mando;2544
+sergio_vega;2545
+fool_s_garden;2546
+waltari;2547
+of_montreal;2548
+pierre_lapointe;2549
+rufus_wainwright;2550
+loudon_wainwright_iii;2551
+sufjan_stevens;2552
+machine_gun_kelly;2553
+francesco_guccini;2554
+le_orme;2555
+lucio_dalla;2556
+michel_fugain;2557
+al_jarreau;2558
+carmen_mcrae;2559
+javier_sol_s;2560
+harry_connick_jr;2561
+bap;2562
+cradle_of_filth;2563
+amorphis;2564
+avatar;2565
+bathory;2566
+behemoth;2567
+borknagar;2568
+countess;2569
+cruachan;2570
+darkthrone;2571
+hate;2572
+destruction;2573
+dimmu_borgir;2574
+eisregen;2575
+enslaved;2576
+finntroll;2577
+fates_warning;2578
+graveworm;2579
+impaled_nazarene;2580
+sentenced;2581
+king_diamond;2582
+kreator;2583
+lord_belial;2584
+marduk;2585
+mercyful_fate;2586
+stick_to_your_guns;2587
+moonspell;2588
+as_i_lay_dying;2589
+nunslaughter;2590
+rotting_christ;2591
+samael;2592
+sandy_denny;2593
+skyforger;2594
+sodom;2595
+cannibal_corpse;2596
+exodus;2597
+atreyu;2598
+theatres_des_vampires;2599
+wizard;2600
+transmetal;2601
+venom;2602
+belphegor;2603
+the_crown;2604
+moya_brennan;2605
+todd_rundgren;2606
+clay_walker;2607
+andrew_peterson;2608
+lynn_anderson;2609
+david_crowder_band;2610
+pam_tillis;2611
+norah_jones;2612
+rhonda_vincent;2613
+jamey_johnson;2614
+plumb;2615
+j_j_cale;2616
+new_riders_of_the_purple_sage;2617
+joe_diffie;2618
+kasey_chambers;2619
+leon_russell;2620
+jack_greene;2621
+the_string_cheese_incident;2622
+ystein_sunde;2623
+stephen_stills;2624
+cancerslug;2625
+robert_plant;2626
+alvin_lee;2627
+beth_hart;2628
+jimmy_buffett;2629
+billy_s_band;2630
+bunbury;2631
+nacho_vegas;2632
+calogero;2633
+georges_brassens;2634
+canned_heat;2635
+charlie_louvin;2636
+colin_james;2637
+cuby_blizzards;2638
+dick_annegarn;2639
+edoardo_bennato;2640
+eva_cassidy;2641
+gil_scott_heron;2642
+glenn_hughes;2643
+deep_purple;2644
+connie_smith;2645
+iva_zanicchi;2646
+izzy_stradlin;2647
+j_karjalainen;2648
+jack_bruce;2649
+leonard_cohen;2650
+joan_armatrading;2651
+joan_osborne;2652
+john_martyn;2653
+rio_reiser;2654
+larry_carlton;2655
+madeleine_peyroux;2656
+bruce_cockburn;2657
+kate_anna_mcgarrigle;2658
+mavis_staples;2659
+noa;2660
+ralph_mctell;2661
+renato_carosone;2662
+richie_kotzen;2663
+robben_ford;2664
+roberto_carlos;2665
+erasmo_carlos;2666
+robin_trower;2667
+rory_block;2668
+roy_buchanan;2669
+sandra_mihanovich;2670
+savoy_brown;2671
+shirley_horn;2672
+siniestro_total;2673
+slank;2674
+the_fabulous_thunderbirds;2675
+the_seatbelts;2676
+the_tragically_hip;2677
+mike_jones;2678
+trophy_scars;2679
+caravan;2680
+velhas_virgens;2681
+walter_trout;2682
+gov_t_mule;2683
+bar_o_vermelho;2684
+blue_cheer;2685
+ian_hunter;2686
+david_leb_n;2687
+de_palmas;2688
+eugenio_finardi;2689
+extreme;2690
+foghat;2691
+george_thorogood_the_destroyers;2692
+great_white;2693
+guardian;2694
+jethro_tull;2695
+ian_anderson;2696
+david_knopfler;2697
+steppenwolf;2698
+dave_edmunds;2699
+lynyrd_skynyrd;2700
+crosby_stills_nash;2701
+raul_seixas;2702
+the_poodles;2703
+musiq_soulchild;2704
+shocking_blue;2705
+nick_lowe;2706
+the_black_crowes;2707
+traffic;2708
+widespread_panic;2709
+co;2710
+alberto_cortez;2711
+joan_sebastian;2712
+ana_gabriel;2713
+gilberto_santa_rosa;2714
+rub_n_blades;2715
+v_ctor_manuelle;2716
+celia_cruz;2717
+luis_fonsi;2718
+nek;2719
+dr_feelgood;2720
+astrud_gilberto;2721
+benito_di_paula;2722
+brazzaville;2723
+sacha_distel;2724
+chico_buarque;2725
+elis_regina;2726
+milton_nascimento;2727
+faf_de_bel_m;2728
+nikka_costa;2729
+tim_maia;2730
+gilberto_gil;2731
+lisa_ekdahl;2732
+joyce;2733
+maria_rita;2734
+nara_le_o;2735
+nouvelle_vague;2736
+paulinho_moska;2737
+wilson_simonal;2738
+14_bis;2739
+arnaldo_antunes;2740
+biquini_cavad_o;2741
+cidade_negra;2742
+cpm_22;2743
+c_ssia_eller;2744
+os_paralamas_do_sucesso;2745
+guilherme_arantes;2746
+ira;2747
+lob_o;2748
+nenhum_de_n_s;2749
+djavan;2750
+rog_rio_skylab;2751
+roupa_nova;2752
+ultraje_a_rigor;2753
+kj_52;2754
+amado_batista;2755
+chit_ozinho_xoror;2756
+jo_o_paulo_daniel;2757
+leandro_leonardo;2758
+leonardo;2759
+odair_jos;2760
+kaiser_chiefs;2761
+kula_shaker;2762
+lightning_seeds;2763
+pulp;2764
+the_proclaimers;2765
+dying_fetus;2766
+napalm_death;2767
+nile;2768
+pathology;2769
+hilary_duff;2770
+badly_drawn_boy;2771
+federico_salvatore;2772
+i_gufi;2773
+zachary_richard;2774
+stan_rogers;2775
+moxy_fr_vous;2776
+poco;2777
+la_bottine_souriante;2778
+stompin_tom_connors;2779
+bersuit_vergarabat;2780
+las_pastillas_del_abuelo;2781
+george_lam;2782
+altan;2783
+clannad;2784
+blackmore_s_night;2785
+capercaillie;2786
+celtic_thunder;2787
+eluveitie;2788
+powerwolf;2789
+gaelic_storm;2790
+an_na;2791
+jon_anderson;2792
+the_dubliners;2793
+loreena_mckennitt;2794
+omnia;2795
+secret_garden;2796
+shaun_davey;2797
+roger_daltrey;2798
+the_corrs;2799
+los_tigres_del_norte;2800
+laurent_voulzy;2801
+the_kelly_family;2802
+wolfe_tones;2803
+alan_stivell;2804
+heather_alexander;2805
+kate_rusby;2806
+dropkick_murphys;2807
+great_big_sea;2808
+fiddler_s_green;2809
+heather_dale;2810
+runrig;2811
+the_waterboys;2812
+dougie_maclean;2813
+adriano_celentano;2814
+alain_chamfort;2815
+zazie;2816
+hamelen;2817
+tazenda;2818
+arno;2819
+arthur_h;2820
+boudewijn_de_groot;2821
+charles_trenet;2822
+claudio_baglioni;2823
+claudio_rocchi;2824
+fabrizio_de_andr;2825
+dalida;2826
+dana_winner;2827
+demis_roussos;2828
+esther_ofarim;2829
+eugenio_bennato;2830
+michel_berger;2831
+francis_cabrel;2832
+maxime_le_forestier;2833
+georges_moustaki;2834
+gianmaria_testa;2835
+gianni_morandi;2836
+gigliola_cinquetti;2837
+milva;2838
+gilbert_b_caud;2839
+ginette_reno;2840
+giuni_russo;2841
+guy_b_art;2842
+helena_vondr_kov;2843
+hugues_aufray;2844
+ivan_graziani;2845
+ivano_fossati;2846
+jacques_bertin;2847
+jean_ferrat;2848
+juliane_werding;2849
+julien_clerc;2850
+los_temerarios;2851
+katerine;2852
+leny_escudero;2853
+mathieu_chedid;2854
+luca_barbarossa;2855
+l_o_ferr;2856
+rosenstolz;2857
+marc_lavoine;2858
+massimo_bubola;2859
+mecano;2860
+mia_martini;2861
+michel_jonasz;2862
+michele_zarrillo;2863
+fiorello;2864
+nada;2865
+mercedes_sosa;2866
+nino_d_angelo;2867
+patrick_bruel;2868
+patty_pravo;2869
+pierre_bachelet;2870
+rainald_grebe;2871
+rapha_l;2872
+raphael;2873
+richard_anthony;2874
+roberto_murolo;2875
+ron;2876
+stefano_rosso;2877
+stephan_eicher;2878
+vasco_rossi;2879
+yves_duteil;2880
+yves_jamait;2881
+ang_lica;2882
+aaron_carter;2883
+barry_louis_polisar;2884
+yuri;2885
+cri_cri;2886
+hevisaurus;2887
+juice_leskinen;2888
+kidz_bop;2889
+mara_maravilha;2890
+destroyer;2891
+scorpions;2892
+obk;2893
+duncan_dhu;2894
+parry_gripp;2895
+sandy_junior;2896
+the_verve_pipe;2897
+the_verve;2898
+the_wiggles;2899
+veggietales;2900
+newsboys;2901
+steven_curtis_chapman;2902
+toro_y_moi;2903
+medi_val_b_bes;2904
+aaron_neville;2905
+bethel_music;2906
+apologetix;2907
+gaither_vocal_band;2908
+building_429;2909
+chris_tomlin;2910
+matt_maher;2911
+jerusalem;2912
+david_meece;2913
+debby_boone;2914
+elevation_worship;2915
+matt_redman;2916
+planetshakers;2917
+majesty;2918
+jump5;2919
+lecrae;2920
+michael_w_smith;2921
+bride;2922
+natalie_grant;2923
+the_lads;2924
+audio_adrenaline;2925
+paul_wilbur;2926
+psalmen_voor_nu;2927
+sawyer_brown;2928
+shane_shane;2929
+the_echoing_green;2930
+twila_paris;2931
+watch_tower_bible_and_tract_society;2932
+da_t_r_u_t_h;2933
+dc_talk;2934
+flame;2935
+grits;2936
+trip_lee;2937
+crystal_lewis;2938
+the_cross_movement;2939
+tobymac;2940
+vico_c;2941
+mormon_tabernacle_choir;2942
+august_burns_red;2943
+black_veil_brides;2944
+deliverance;2945
+opeth;2946
+die_happy;2947
+disciple;2948
+galactic_cowboys;2949
+haste_the_day;2950
+living_sacrifice;2951
+mastodon;2952
+mortification;2953
+showbread;2954
+labyrinth;2955
+stryper;2956
+the_devil_wears_prada;2957
+underoath;2958
+whitecross;2959
+petra;2960
+huntingtons;2961
+mxpx;2962
+d_a_d;2963
+caedmon_s_call;2964
+david_and_the_giants;2965
+degarmo_and_key;2966
+delirious;2967
+don_francisco;2968
+five_iron_frenzy;2969
+geoff_moore;2970
+hawk_nelson;2971
+grave;2972
+larry_norman;2973
+randy_stonehill;2974
+monty_python;2975
+oomph;2976
+oficina_g3;2977
+white_heart;2978
+rescate;2979
+rick_wakeman;2980
+la_oreja_de_van_gogh;2981
+sanctus_real;2982
+fun_people;2983
+thousand_foot_krutch;2984
+tim_hughes;2985
+the_o_c_supertones;2986
+4him;2987
+billy_gilman;2988
+aimee_mann;2989
+katharine_mcphee;2990
+eros_ramazzotti;2991
+z_ro;2992
+babbie_mason;2993
+bebo_norman;2994
+judy_garland;2995
+carman;2996
+cece_winans;2997
+trick_daddy;2998
+chris_isaak;2999
+cocteau_twins;3000
+edyta_g_rniak;3001
+enrico_ruggeri;3002
+ffh;3003
+hanson;3004
+hawksley_workman;3005
+indigo_girls;3006
+irene_grandi;3007
+jackie_evancho;3008
+joy_electric;3009
+kelly_price;3010
+mary_mary;3011
+israel_houghton;3012
+phil_wickham;3013
+phillips_craig_dean;3014
+roch_voisine;3015
+rupaul;3016
+gregorian;3017
+sarah_connor;3018
+sugarland;3019
+sweetbox;3020
+tarja;3021
+the_brian_setzer_orchestra;3022
+brian_setzer;3023
+badfinger;3024
+the_moffatts;3025
+the_vandals;3026
+trans_siberian_orchestra;3027
+roy_drusky;3028
+burton_cummings;3029
+procol_harum;3030
+renaissance;3031
+the_pretty_things;3032
+twisted_sister;3033
+bj_rn_eidsv_g;3034
+corvus_corax;3035
+schelmish;3036
+emilie_autumn;3037
+epica;3038
+katherine_jenkins;3039
+scala_kolacny_brothers;3040
+take_6;3041
+the_roches;3042
+tony_banks;3043
+to_e_proeski;3044
+lacrimosa;3045
+16_volt;3046
+bj_rn_rosenstr_m;3047
+bob_rivers;3048
+cledus_t_judd;3049
+frankjavcee;3050
+george_formby;3051
+ninja_sex_party;3052
+paul_and_storm;3053
+the_arrogant_worms;3054
+tripod;3055
+el_cuarteto_de_nos;3056
+gwar;3057
+knorkator;3058
+psychostick;3059
+rodgau_monotones;3060
+los_palominos;3061
+charlie_peacock;3062
+jesus_culture;3063
+michael_card;3064
+tenth_avenue_north;3065
+carrie_newcomer;3066
+nick_drake;3067
+aaron_watson;3068
+billy_joe_royal;3069
+billy_joe_shaver;3070
+charlie_landsborough;3071
+chris_ledoux;3072
+collin_raye;3073
+dan_seals;3074
+dave_dudley;3075
+hellbillies;3076
+ed_bruce;3077
+emilio_navaira;3078
+jean_shepard;3079
+freddie_hart;3080
+gary_stewart;3081
+gene_watson;3082
+gian_giovani;3083
+gilberto_gilmar;3084
+jason_mraz;3085
+ilse_delange;3086
+john_prine;3087
+jake_owen;3088
+wynn_stewart;3089
+jim_ed_brown;3090
+joe_ely;3091
+kid_rock;3092
+la_toya_jackson;3093
+lit;3094
+lita_ford;3095
+me_first_and_the_gimme_gimmes;3096
+lagwagon;3097
+melanie;3098
+mickey_newbury;3099
+paul_brunelle;3100
+paula_fernandes;3101
+zez_di_camargo_luciano;3102
+randy_rogers_band;3103
+reverend_horton_heat;3104
+rick_renner;3105
+rionegro_solim_es;3106
+shooter_jennings;3107
+terri_clark;3108
+vern_gosdin;3109
+webb_wilder;3110
+ween;3111
+38_special;3112
+the_beau_brummels;3113
+matanza;3114
+clawfinger;3115
+acid_drinkers;3116
+agnostic_front;3117
+biohazard;3118
+body_count;3119
+d_r_i;3120
+municipal_waste;3121
+neurosis;3122
+nuclear_assault;3123
+soziedad_alkoholika;3124
+suicidal_tendencies;3125
+paragon;3126
+mario;3127
+inna;3128
+belinda;3129
+bronco;3130
+grupo_bryndis;3131
+david_bisbal;3132
+ram_n_ayala;3133
+grant_lee_phillips;3134
+the_veronicas;3135
+amr_diab;3136
+atb;3137
+basshunter;3138
+dream_theater;3139
+frankie_j;3140
+baby_bash;3141
+sophie_ellis_bextor;3142
+grace_jones;3143
+laveerre;3144
+silkk_the_shocker;3145
+parov_stelar;3146
+raffaella_carr;3147
+elephant_man;3148
+saint_etienne;3149
+samantha_fox;3150
+selena;3151
+super_junior;3152
+t_a_t_u;3153
+tarkan;3154
+judie_tzuke;3155
+el_kel_iset;3156
+yello;3157
+franz_ferdinand;3158
+chenoa;3159
+lucero;3160
+tokio;3161
+puffy_amiyumi;3162
+wink;3163
+obie_trice;3164
+mystikal;3165
+current_93;3166
+dark_sanctuary;3167
+rome;3168
+lord_of_the_lost;3169
+bella_morte;3170
+mantus;3171
+blutengel;3172
+clan_of_xymox;3173
+dead_can_dance;3174
+death_in_june;3175
+diary_of_dreams;3176
+diorama;3177
+helium_vola;3178
+illuminate;3179
+l_me_immortelle;3180
+lacrimas_profundere;3181
+killing_joke;3182
+m_nchener_freiheit;3183
+otto_dix;3184
+project_pitchfork;3185
+qntal;3186
+sopor_aeternus;3187
+the_cr_xshadows;3188
+unheilig;3189
+welle;3190
+yendri;3191
+carcass;3192
+asphyx;3193
+bolt_thrower;3194
+darkseed;3195
+paradise_lost;3196
+tiamat;3197
+the_damned;3198
+pantera;3199
+the_amity_affliction;3200
+judas_priest;3201
+amon_amarth;3202
+alesana;3203
+atrocity;3204
+autopsy;3205
+avulsed;3206
+sabaton;3207
+misfits;3208
+iron_fire;3209
+centinex;3210
+dagoba;3211
+dark_tranquillity;3212
+asia;3213
+deicide;3214
+dethklok;3215
+dew_scented;3216
+edge_of_sanity;3217
+escape_the_fate;3218
+heaven_shall_burn;3219
+hypocrisy;3220
+incantation;3221
+jungle_rot;3222
+kataklysm;3223
+krisiun;3224
+macabre;3225
+malevolent_creation;3226
+meshuggah;3227
+misanthrope;3228
+morbid_angel;3229
+dead_kennedys;3230
+necro;3231
+pig_destroyer;3232
+shadows_fall;3233
+sinister;3234
+six_feet_under;3235
+dream_evil;3236
+soulfly;3237
+the_black_dahlia_murder;3238
+between_the_buried_and_me;3239
+therion;3240
+vader;3241
+whitechapel;3242
+attila;3243
+emmure;3244
+miss_may_i;3245
+the_acacia_strain;3246
+betontod;3247
+broilers;3248
+dritte_wahl;3249
+ohl;3250
+slime;3251
+terrorgruppe;3252
+b_hse_onkelz;3253
+frei_wild;3254
+k_rbholz;3255
+asp;3256
+tokio_hotel;3257
+queensr_che;3258
+amanda_miguel;3259
+arabesque;3260
+bad_boys_blue;3261
+boyce_avenue;3262
+parliament;3263
+wu_tang_clan;3264
+neoton_fam_lia;3265
+teena_marie;3266
+bobby_womack;3267
+agoraphobic_nosebleed;3268
+candlemass;3269
+electric_wizard;3270
+black_sabbath;3271
+theatre_of_tragedy;3272
+type_o_negative;3273
+marie_fredriksson;3274
+luna;3275
+marissa_nadler;3276
+yo_la_tengo;3277
+celldweller;3278
+hitomi;3279
+big_d_and_the_kids_table;3280
+alacranes_musical;3281
+k_paz_de_la_sierra;3282
+assemblage_23;3283
+covenant;3284
+die_krupps;3285
+kodak_black;3286
+front_242;3287
+haujobb;3288
+in_strict_confidence;3289
+le_ther_strip;3290
+snog;3291
+the_darkness;3292
+tanzwut;3293
+terminal_choice;3294
+velvet_acid_christ;3295
+vnv_nation;3296
+wumpscut;3297
+x_fusion;3298
+umbra_et_imago;3299
+de_vision;3300
+deichkind;3301
+eisbrecher;3302
+herbie_hancock;3303
+ana_moura;3304
+macaco;3305
+skinny_puppy;3306
+ayreon;3307
+black_moth_super_rainbow;3308
+erykah_badu;3309
+cocorosie;3310
+de_jeugd_van_tegenwoordig;3311
+dj_shadow;3312
+e_nomine;3313
+kmfdm;3314
+flying_lotus;3315
+goldfrapp;3316
+hanzel_und_gretyl;3317
+information_society;3318
+mc_frontalot;3319
+kraftwerk;3320
+ladytron;3321
+lamb;3322
+milk_inc;3323
+mind_in_a_box;3324
+ministry;3325
+m_m;3326
+m_nia;3327
+pig;3328
+pitchshifter;3329
+lil_boosie;3330
+master_p;3331
+mindless_self_indulgence;3332
+buzzcocks;3333
+vanilla_ice;3334
+milie_simon;3335
+gianna_nannini;3336
+pinback;3337
+the_birthday_massacre;3338
+archive;3339
+99_posse;3340
+bloc_party;3341
+morcheeba;3342
+origa;3343
+paul_kalkbrenner;3344
+tina_arena;3345
+dover;3346
+melotron;3347
+owl_city;3348
+kamelot;3349
+greeley_estates;3350
+hawthorne_heights;3351
+joan_of_arc;3352
+saves_the_day;3353
+thursday;3354
+transit;3355
+fairport_convention;3356
+maggie_reilly;3357
+joan_manuel_serrat;3358
+e_rotic;3359
+the_scene;3360
+sandra;3361
+amon_d_l_ii;3362
+circa_survive;3363
+love_solfege;3364
+caliban;3365
+tall_dwarfs;3366
+van_der_graaf_generator;3367
+death_grips;3368
+the_fiery_furnaces;3369
+am_lia_rodrigues;3370
+cristina_branco;3371
+jos_afonso;3372
+katia_guerreiro;3373
+ney_matogrosso;3374
+madredeus;3375
+mariza;3376
+gipsy_kings;3377
+mal;3378
+aleks_syntek;3379
+ni_a_pastori;3380
+rosario;3381
+al_stewart;3382
+amos_lee;3383
+andr_s_calamaro;3384
+ane_brun;3385
+asa;3386
+editors;3387
+catie_curtis;3388
+chrystian_ralf;3389
+clueso;3390
+eddi_reader;3391
+eddie_from_ohio;3392
+ellis_paul;3393
+frank_turner;3394
+estampie;3395
+ferdi_tayfur;3396
+fito_p_ez;3397
+luis_alberto_spinetta;3398
+gabriella_ferri;3399
+gigi;3400
+greg_brown;3401
+g_ksel;3402
+lando_fiorini;3403
+india_arie;3404
+jack_savoretti;3405
+anne_grete_preus;3406
+jarom_r_nohavica;3407
+joe_purdy;3408
+john_wesley_harding;3409
+josh_rouse;3410
+karel_kryl;3411
+v_tor_ramil;3412
+lars_winnerb_ck;3413
+laura_marling;3414
+llu_s_llach;3415
+los_chalchaleros;3416
+luka_bloom;3417
+malicorne;3418
+mark_heard;3419
+martin_carthy;3420
+nic_jones;3421
+le_n_gieco;3422
+mijares;3423
+nuova_compagnia_di_canto_popolare;3424
+ola_magnell;3425
+thin_lizzy;3426
+ray_lamontagne;3427
+ron_sexsmith;3428
+rosana;3429
+silvio_rodr_guez;3430
+stef_bos;3431
+sun_kil_moon;3432
+tanita_tikaram;3433
+the_incredible_string_band;3434
+thea_gilmore;3435
+tina_dico;3436
+victor_leo;3437
+v_rttin;3438
+ge_aleksandersen;3439
+i_brahim_tatl_ses;3440
+ektomorf;3441
+elvenking;3442
+ensiferum;3443
+falconer;3444
+feuerschwanz;3445
+in_extremo;3446
+korpiklaani;3447
+leaves_eyes;3448
+letzte_instanz;3449
+m_go_de_oz;3450
+saurom;3451
+schandmaul;3452
+skyclad;3453
+subway_to_sally;3454
+suidakra;3455
+t_r;3456
+icehouse;3457
+bomb_the_music_industry;3458
+the_real_mckenzies;3459
+54_40;3460
+armored_saint;3461
+alexz_johnson;3462
+bar_man_o;3463
+ezginin_g_nl;3464
+galija;3465
+sts;3466
+h_kan_hellstr_m;3467
+james_blunt;3468
+kazik;3469
+mewithoutyou;3470
+michel_polnareff;3471
+ovidi_montllor;3472
+rasputina;3473
+shearwater;3474
+gerry_rafferty;3475
+steam_powered_giraffe;3476
+the_saw_doctors;3477
+ty_segall;3478
+tyrone_wells;3479
+avi_es_do_forr;3480
+grimskunk;3481
+sinik;3482
+vitaa;3483
+kenza_farah;3484
+sexion_d_assaut;3485
+aliz_e;3486
+henri_tachan;3487
+jenifer;3488
+m_pokora;3489
+indochine;3490
+brainstorm;3491
+con_funk_shun;3492
+funkadelic;3493
+lena_park;3494
+neffa;3495
+ugk;3496
+suburban_legends;3497
+mai_kuraki;3498
+cherry_poppin_daddies;3499
+electric_six;3500
+los_straitjackets;3501
+the_69_eyes;3502
+the_angels;3503
+the_haunted;3504
+the_hellacopters;3505
+the_kills;3506
+thee_oh_sees;3507
+white_denim;3508
+zabranjeno_pu_enje;3509
+ol_dirty_bastard;3510
+kurupt;3511
+spice_1;3512
+brotha_lynch_hung;3513
+chamillionaire;3514
+paul_wall;3515
+trae;3516
+club_dogo;3517
+mc_eiht;3518
+royce_da_5_9;3519
+geto_boys;3520
+the_diplomats;3521
+ice_t;3522
+2_live_crew;3523
+xv;3524
+mobb_deep;3525
+c_murder;3526
+tru;3527
+lil_keke;3528
+project_pat;3529
+tha_dogg_pound;3530
+esham;3531
+twiztid;3532
+erick_sermon;3533
+big_tymers;3534
+kate_nash;3535
+the_cramps;3536
+nekromantix;3537
+tsol;3538
+ace_frehley;3539
+hardcore_superstar;3540
+harem_scarem;3541
+house_of_lords;3542
+kingdom_come;3543
+l_a_guns;3544
+mr_big;3545
+pink_cream_69;3546
+quiet_riot;3547
+riot;3548
+ratt;3549
+tnt;3550
+backyard_babies;3551
+ultima_thule;3552
+europe;3553
+hanoi_rocks;3554
+mott_the_hoople;3555
+smokie;3556
+suzi_quatro;3557
+haemorrhage;3558
+aline_barros;3559
+bruna_karla;3560
+kirk_franklin;3561
+minist_rio_koinonya_de_louvor;3562
+artrosis;3563
+closterkeller;3564
+indica;3565
+sirenia;3566
+trail_of_tears;3567
+tristania;3568
+within_temptation;3569
+bauhaus;3570
+mono_inc;3571
+pansy_division;3572
+xandria;3573
+immortal_technique;3574
+agathocles;3575
+rotten_sound;3576
+the_locust;3577
+anthrax;3578
+devildriver;3579
+lamb_of_god;3580
+machine_head;3581
+parkway_drive;3582
+pro_pain;3583
+throwdown;3584
+vicious_rumors;3585
+screaming_trees;3586
+cuisillos;3587
+intocable;3588
+pesado;3589
+la_mafia;3590
+marco_antonio_sol_s;3591
+los_bukis;3592
+andrew_w_k;3593
+april_wine;3594
+axel_rudi_pell;3595
+b_z;3596
+tak_matsumoto;3597
+barricada;3598
+bijelo_dugme;3599
+blaze_bayley;3600
+bonfire;3601
+bruce_dickinson;3602
+buckcherry;3603
+budgie;3604
+buitres;3605
+jorn;3606
+doro;3607
+enuff_z_nuff;3608
+gentle_giant;3609
+girlschool;3610
+golden_earring;3611
+gotthard;3612
+nazareth;3613
+a_day_to_remember;3614
+jefferson_airplane;3615
+joe_satriani;3616
+ken_hensley;3617
+kim_mitchell;3618
+king_s_x;3619
+kotiteollisuus;3620
+la_renga;3621
+lee_aaron;3622
+lordi;3623
+michael_schenker_group;3624
+mustasch;3625
+night_ranger;3626
+omega;3627
+parni_valjak;3628
+paul_gilbert;3629
+popeda;3630
+skid_row;3631
+tankcsapda;3632
+the_bronx;3633
+the_donnas;3634
+all_that_remains;3635
+triumph;3636
+umphrey_s_mcgee;3637
+y_t;3638
+ziggy;3639
+sfdk;3640
+7_seconds;3641
+aiden;3642
+alphaville;3643
+black_flag;3644
+slayer;3645
+circle_jerks;3646
+ritchie;3647
+converge;3648
+every_time_i_die;3649
+hatebreed;3650
+nomeansno;3651
+rancid;3652
+memphis_may_fire;3653
+nofx;3654
+propagandhi;3655
+tankard;3656
+screeching_weasel;3657
+sick_of_it_all;3658
+silverstein;3659
+two_steps_from_hell;3660
+faun;3661
+accept;3662
+the_frames;3663
+andromeda;3664
+annihilator;3665
+anvil;3666
+artillery;3667
+avenged_sevenfold;3668
+axxis;3669
+blind_guardian;3670
+vanden_plas;3671
+grave_digger;3672
+dragonforce;3673
+edenbridge;3674
+damien_jurado;3675
+exciter;3676
+firewind;3677
+halford;3678
+hammerfall;3679
+helloween;3680
+helstar;3681
+iced_earth;3682
+jag_panzer;3683
+machinae_supremacy;3684
+manowar;3685
+metal_church;3686
+morgana_lefay;3687
+mudvayne;3688
+nocturnal_rites;3689
+overkill;3690
+primal_fear;3691
+rebellion;3692
+running_wild;3693
+corey_hart;3694
+savatage;3695
+saxon;3696
+steve_vai;3697
+tad_morose;3698
+tarot;3699
+tierra_santa;3700
+trivium;3701
+turmion_k_til_t;3702
+u_d_o;3703
+virgin_steele;3704
+voivod;3705
+warcry;3706
+yngwie_malmsteen;3707
+zion_lennox;3708
+sido;3709
+mc_chris;3710
+assalti_frontali;3711
+kool_keith;3712
+ayumi_hamasaki;3713
+az;3714
+bahh_tee;3715
+bassi_maestro;3716
+revocation;3717
+blumentopf;3718
+brockhampton;3719
+bts;3720
+bushido;3721
+vinnie_paz;3722
+chakuza;3723
+cheek;3724
+cro;3725
+arc_ngel;3726
+alexis_fido;3727
+dargen_d_amico;3728
+the_coup;3729
+def_con_dos;3730
+die_fantastischen_vier;3731
+dom_no;3732
+donguralesko;3733
+epmd;3734
+kool_savas;3735
+fettes_brot;3736
+fronda;3737
+mc_solaar;3738
+pyhimys;3739
+kaaris;3740
+kollegah;3741
+kontra_k;3742
+k_k;3743
+l_o_c;3744
+logic;3745
+jerry_rivera;3746
+murs;3747
+angie_stone;3748
+namie_amuro;3749
+anthony_hamilton;3750
+lyfe_jennings;3751
+bl_f;3752
+o_s_t_r;3753
+paluch;3754
+parazi_ii;3755
+porta;3756
+bleeding_through;3757
+prinz_pi;3758
+rasmentalism;3759
+xavier_naidoo;3760
+sage_francis;3761
+stupeflip;3762
+young_thug;3763
+tego_calder_n;3764
+fifth_harmony;3765
+jay_chou;3766
+blitzkid;3767
+zumbis_do_espa_o;3768
+deer_tick;3769
+half_man_half_biscuit;3770
+hayden;3771
+club_8;3772
+grandaddy;3773
+jens_lekman;3774
+kent;3775
+keren_ann;3776
+los_campesinos;3777
+nellie_mckay;3778
+china_crisis;3779
+prefab_sprout;3780
+the_clientele;3781
+the_lucksmiths;3782
+bell_x1;3783
+british_sea_power;3784
+car_seat_headrest;3785
+deerhunter;3786
+dr_dog;3787
+elf_power;3788
+frightened_rabbit;3789
+fugazi;3790
+fury_in_the_slaughterhouse;3791
+julie_doiron;3792
+tinashe;3793
+la_habitaci_n_roja;3794
+margot_the_nuclear_so_and_so_s;3795
+matt_pond_pa;3796
+metric;3797
+mike_doughty;3798
+mother_mother;3799
+piebald;3800
+quasi;3801
+rheostatics;3802
+sebadoh;3803
+spoon;3804
+starflyer_59;3805
+stephen_malkmus;3806
+stereolab;3807
+ted_leo_and_the_pharmacists;3808
+the_appleseed_cast;3809
+the_faint;3810
+the_go_betweens;3811
+the_pineapple_thief;3812
+the_undertones;3813
+tronic;3814
+chris_de_burgh;3815
+mass_hysteria;3816
+angelo_branduardi;3817
+gigi_d_alessio;3818
+i_muvrini;3819
+back_number;3820
+boa;3821
+claris;3822
+crystal_kay;3823
+zard;3824
+gackt;3825
+garnet_crow;3826
+girls_generation;3827
+kat_tun;3828
+koda_kumi;3829
+kotoko;3830
+lisa;3831
+maaya_sakamoto;3832
+masami_okui;3833
+mr_children;3834
+news;3835
+shinee;3836
+w_inds;3837
+yui;3838
+yumi_matsutoya;3839
+the_high_lows;3840
+sid;3841
+abbey_lincoln;3842
+anna_maria_jopek;3843
+cassandra_wilson;3844
+dianne_reeves;3845
+fred_buscaglione;3846
+jane_monheit;3847
+zor_n;3848
+kraan;3849
+laura_fygi;3850
+michael_franks;3851
+natalino_otto;3852
+quartetto_cetra;3853
+scott_bradlee_s_postmodern_jukebox;3854
+stacey_kent;3855
+the_flower_kings;3856
+ronnie_von;3857
+brown_eyed_girls;3858
+ahmet_kaya;3859
+alejandra_guzm_n;3860
+ana_carolina;3861
+alcione;3862
+el_chapo_de_sinaloa;3863
+gustavo_cerati;3864
+soda_stereo;3865
+jenni_rivera;3866
+joaqu_n_sabina;3867
+los_fabulosos_cadillacs;3868
+abel_pintos;3869
+ana_bel_n;3870
+aterciopelados;3871
+camilo_sesto;3872
+david_demar_a;3873
+gian_marco;3874
+menudo;3875
+ricardo_arjona;3876
+sabroso;3877
+v_ctor_manuel;3878
+las_pelotas;3879
+ariel_pink;3880
+leehom_wang;3881
+jolin_tsai;3882
+darkest_hour;3883
+kalmah;3884
+nightrage;3885
+eppu_normaali;3886
+the_outfield;3887
+no_use_for_a_name;3888
+pennywise;3889
+callejon;3890
+d_f_c;3891
+our_last_night;3892
+exaltasamba;3893
+beth_carvalho;3894
+jo_o_bosco;3895
+marina_lima;3896
+marisa_monte;3897
+nando_reis;3898
+natiruts;3899
+ra_a_negra;3900
+s_pra_contrariar;3901
+zeca_pagodinho;3902
+andr_hazes;3903
+de_dijk;3904
+arena;3905
+iq;3906
+sol_invictus;3907
+new_found_glory;3908
+adam_ant;3909
+berlin;3910
+hoodoo_gurus;3911
+ultravox;3912
+nik_kershaw;3913
+squeeze;3914
+the_aquabats;3915
+the_fixx;3916
+beat_crusaders;3917
+cows;3918
+conjunto_primavera;3919
+peter_and_the_test_tube_babies;3920
+sham_69;3921
+the_adicts;3922
+the_analogs;3923
+instalok;3924
+jacek_kaczmarski;3925
+przemys_aw_gintrowski;3926
+ada_band;3927
+agnetha_f_ltskog;3928
+ajda_pekkan;3929
+al_bano;3930
+alex_ubago;3931
+alison_moyet;3932
+alunni_del_sole;3933
+anna_oxa;3934
+bajm;3935
+barclay_james_harvest;3936
+blue_system;3937
+brunner_brunner;3938
+candan_er_etin;3939
+christian_bautista;3940
+clay_aiken;3941
+clifford_t_ward;3942
+daniel;3943
+don_backy;3944
+jesse_mccartney;3945
+emma;3946
+marcella_bella;3947
+giorgio_gaber;3948
+guus_meeuwis;3949
+heinz_rudolf_kunze;3950
+john_farnham;3951
+ian_thomas;3952
+i_n_karaca;3953
+jennifer_rush;3954
+jo_vally;3955
+john_fogerty;3956
+julian_lennon;3957
+k3;3958
+kid_abelha;3959
+labv_l_gais_tips;3960
+l_vi;3961
+lea_salonga;3962
+les_wampas;3963
+magnus_uggla;3964
+mango;3965
+maria_mena;3966
+massimo_ranieri;3967
+max_gazz;3968
+michael_learns_to_rock;3969
+mietta;3970
+mustafa_sandal;3971
+nil_fer;3972
+peter_frampton;3973
+pr_ta_v_tra;3974
+pur;3975
+rettore;3976
+ricchi_e_poveri;3977
+rob_de_nijs;3978
+sara_bareilles;3979
+sasha;3980
+sertab_erener;3981
+sezen_aksu;3982
+stadio;3983
+stephen_sondheim;3984
+tamara;3985
+team_starkid;3986
+toto_cutugno;3987
+umberto_tozzi;3988
+herman_brood;3989
+wanessa;3990
+zen_caf;3991
+bonanza_banzai;3992
+bodyjar;3993
+bracket;3994
+frenzal_rhomb;3995
+goldfinger;3996
+the_wonder_years;3997
+useless_id;3998
+camel;3999
+hombres_g;4000
+leo_jaime;4001
+neal_morse;4002
+spock_s_beard;4003
+new_trolls;4004
+opus;4005
+piersi;4006
+premiata_forneria_marconi;4007
+superbus;4008
+zmelkoow;4009
+boysetsfire;4010
+hot_water_music;4011
+new_model_army;4012
+the_monochrome_set;4013
+big_big_train;4014
+avantasia;4015
+dark_moor;4016
+dreamtale;4017
+freedom_call;4018
+mystic_prophecy;4019
+nightmare;4020
+rhapsody_of_fire;4021
+royal_hunt;4022
+sonata_arctica;4023
+stratovarius;4024
+symphony_x;4025
+vision_divine;4026
+the_wildhearts;4027
+armia;4028
+evergrey;4029
+lana_lane;4030
+nektar;4031
+pain_of_salvation;4032
+riverside;4033
+beardfish;4034
+echolyn;4035
+eloy;4036
+john_wetton;4037
+medina_azahara;4038
+mostly_autumn;4039
+pendragon;4040
+rafo_r_ez;4041
+the_meteors;4042
+against_me;4043
+anti_flag;4044
+banda_bassotti;4045
+cadena_perpetua;4046
+descendents;4047
+distemper;4048
+dogwood;4049
+el_ltimo_ke_zierre;4050
+farben_lehre;4051
+toy_dolls;4052
+junkies;4053
+ksu;4054
+la_polla_records;4055
+la_vela_puerca;4056
+leatherface;4057
+less_than_jake;4058
+mad_caddies;4059
+millencolin;4060
+punkreas;4061
+reel_big_fish;4062
+snfu;4063
+stiff_little_fingers;4064
+swingin_utters;4065
+the_bouncing_souls;4066
+the_casualties;4067
+the_dickies;4068
+the_lawrence_arms;4069
+toyah;4070
+gerald_levert;4071
+gondwana;4072
+los_aut_nticos_decadentes;4073
+los_cafres;4074
+los_pericos;4075
+tryo;4076
+rakim_ken_y;4077
+billy_squier;4078
+bj_rn_afzelius;4079
+glay;4080
+hunters_collectors;4081
+john_entwistle;4082
+jokke;4083
+la_beriso;4084
+los_rancheros;4085
+los_tres;4086
+maanam;4087
+mikel_erentxun;4088
+peter_wolf;4089
+racoon;4090
+rev_lver;4091
+riblja_orba;4092
+sandro;4093
+gene_vincent;4094
+the_baseballs;4095
+stray_cats;4096
+as_marcianas;4097
+bruno_marrone;4098
+cristiano_ara_jo;4099
+fernando_sorocaba;4100
+joint_venture;4101
+serge_reggiani;4102
+ska_p;4103
+the_mighty_mighty_bosstones;4104
+fu_manchu;4105
+jay_jay_johanson;4106
+psyche;4107
+carlos_gardel;4108

jukebox/data/ids/v2_genre_ids.txt

@@ -0,0 +1,120 @@
+unknown;0
+classical;1
+blues;2
+hip;3
+hop;4
+dance;5
+soul;6
+hard;7
+rock;8
+jazz;9
+reggae;10
+country;11
+alternative;12
+soundtrack;13
+pop;14
+bluegrass;15
+vocal;16
+r;17
+b;18
+rap;19
+christian;20
+gospel;21
+electronic;22
+christmas;23
+singer;24
+songwriter;25
+metal;26
+n;27
+roll;28
+synthpop;29
+electronica;30
+mpb;31
+movie;32
+indie;33
+new;34
+wave;35
+electro;36
+house;37
+folk;38
+punk;39
+french;40
+contemporary;41
+garage;42
+soft;43
+acoustic;44
+nu;45
+television;46
+post;47
+eurodance;48
+progressive;49
+gothic;50
+classic;51
+funk;52
+disco;53
+swing;54
+trance;55
+thrash;56
+psychedelic;57
+heavy;58
+american;59
+grunge;60
+art;61
+j;62
+gangsta;63
+brazilian;64
+latin;65
+southern;66
+ska;67
+crossover;68
+hardcore;69
+industrial;70
+glam;71
+melodic;72
+ambient;73
+musical;74
+dream;75
+experimental;76
+americana;77
+chanson;78
+rockabilly;79
+britpop;80
+children;81
+s;82
+music;83
+electropop;84
+power;85
+celtic;86
+dark;87
+comedy;88
+doom;89
+trip;90
+lo;91
+fi;92
+metalcore;93
+symphonic;94
+fado;95
+schlager;96
+avant;97
+garde;98
+europop;99
+reggaeton;100
+emo;101
+death;102
+samba;103
+deathcore;104
+black;105
+horrorcore;106
+grindcore;107
+worship;108
+salsa;109
+ebm;110
+neofolk;111
+sertanejo;112
+deutschrock;113
+norte;114
+o;115
+ax;116
+k;117
+tejano;118
+medieval;119

jukebox/data/ids/v3_genre_ids.txt

@@ -0,0 +1,604 @@
+electroclash;1
+acid rock;2
+christian metal;3
+pop rock;4
+gothic;5
+big beat;6
+psychedelic rock‎;7
+funk carioca;8
+bebop;9
+dance punk;10
+trad jazz;11
+romantic;12
+andean music;13
+volksmusik;14
+coldwave;15
+gospel blues;16
+italian folk;17
+disney;18
+dark wave‏‎;19
+powerviolence;20
+bachata;21
+soft rock;22
+s music"];23
+bubblegum dance;24
+western swing;25
+alternative country;26
+latin pop;27
+eurobeat;28
+n;29
+unblack metal;30
+surf;31
+nu-disco;32
+event;33
+classical;34
+nasheed;35
+jovem guarda;36
+british blues;37
+bossa nova;38
+detroit blues;39
+rock;40
+contemporary christian;41
+dark ambient;42
+noise rock;43
+axé;44
+soca;45
+dance-rock;46
+contemporary jazz;47
+appalachian folk;48
+humppa‎;49
+ambient;50
+funeral doom;51
+southern gospel;52
+video game‎;53
+hip hop;54
+glitch hop;55
+krautrock;56
+breakcore;57
+ska;58
+traditional folk;59
+psychedelic trance;60
+reggae‏‎;61
+noise pop;62
+drumstep;63
+house;64
+teen pop;65
+sea shanties;66
+junkanoo;67
+mandopop;68
+pre-war blues;69
+doom metal;70
+oi-punk;71
+swamp rock;72
+crunkcore;73
+rap rock;74
+roots;75
+country rap;76
+avant-garde;77
+cumbia;78
+glam metal;79
+groove metal;80
+electric blues;81
+new orleans rhythm and blues;82
+canadian hip hop;83
+freestyle;84
+deathgrind;85
+idm;86
+comedy rock;87
+art punk;88
+progg;89
+work songs;90
+art pop;91
+conjunto;92
+persian;93
+parody;94
+jazz-funk;95
+french hip hop;96
+spirituals;97
+african;98
+middle-eastern;99
+minimal;100
+ranchera;101
+industrial rock;102
+electro house;103
+celtic rock;104
+death doom;105
+grupera;106
+jazz fusion‎;107
+political folk;108
+christian punk;109
+rapcore;110
+j-pop;111
+mashup;112
+metalcore;113
+progressive country;114
+power noise;115
+hip house;116
+crossover thrash;117
+electropop‎;118
+psychedelic folk;119
+punk rock;120
+classic rock;121
+zydeco;122
+afrobeat;123
+salsa;124
+banda;125
+chill-out;126
+morna;127
+minnesang;128
+alternative metal;129
+djent;130
+african folk;131
+mambo;132
+sertanejo;133
+classic pop;134
+soul;135
+australian hip hop;136
+symphonic rock;137
+celtic punk;138
+synthpop‎;139
+europop;140
+funk;141
+jazz blues;142
+vocal trance;143
+celtic fusion;144
+industrial;145
+kirtan;146
+slowcore;147
+flamenco;148
+piano blues;149
+texas blues;150
+aggrotech;151
+steampunk;152
+opera;153
+folktronica;154
+klezmer;155
+nwobhm;156
+goregrind;157
+rac;158
+neo-psychedelia‏‎;159
+post-rock‎;160
+hard bop;161
+gypsy jazz;162
+new orleans blues;163
+doo-wop;164
+soul blues;165
+trap;166
+indietronica;167
+psychobilly;168
+euro disco;169
+neo-progressive rock;170
+canterbury;171
+freak folk;172
+midwest rap;173
+instrumental rock;174
+dance-pop;175
+avant-garde metal;176
+edm;177
+deep house;178
+progressive bluegrass;179
+rave;180
+australian folk;181
+comic opera;182
+sunshine pop;183
+gregorian chant;184
+psychedelic rock;185
+honky tonk;186
+rock 'n' roll;187
+television;188
+nintendocore;189
+jump blues;190
+roots reggae;191
+traditional bluegrass;192
+operatic pop;193
+skate punk;194
+reggaeton;195
+manele;196
+middle-eastern hip hop;197
+skiffle;198
+nsbm;199
+nu jazz;200
+disco;201
+horrorcore;202
+early music;203
+post-bop;204
+gothic rock;205
+crack rock steady;206
+easy listening;207
+psychedelic;208
+christian;209
+brutal death metal;210
+experimental rock;211
+modern classical‎;212
+drum and bass;213
+dark wave;214
+dubstep;215
+grunge;216
+christian hip hop;217
+latin jazz;218
+r&b;219
+s music", ;220
+free jazz;221
+experimental hip hop;222
+swing;223
+smooth jazz;224
+southern metal;225
+religious;226
+progressive death metal;227
+contemporary folk;228
+j-rock;229
+jazz;230
+hamburger schule;231
+teen pop‎;232
+crossover;233
+italo disco;234
+deathcore;235
+blues;236
+crunk;237
+jangle pop;238
+indian classical music;239
+big band;240
+proto-punk;241
+dirty blues;242
+garage punk;243
+extreme metal;244
+folk metal;245
+neo soul;246
+electric folk;247
+synthwave;248
+arena rock;249
+post-grunge;250
+indie rock;251
+acoustic blues;252
+native american;253
+progressive trance;254
+nu metal;255
+digital hardcore;256
+brazilian rock;257
+funky house;258
+symphonic black metal;259
+lounge music;260
+brega;261
+trance;262
+industrial metal;263
+austropop;264
+bhangra;265
+new wave;266
+neoclassical;267
+post-metal;268
+dub;269
+industrial metal‎;270
+irish folk;271
+deutschrock;272
+gypsy;273
+dark electro;274
+alternative hip hop;275
+mbaqanga;276
+swamp blues;277
+french pop;278
+tango;279
+rockabilly;280
+old-time music;281
+blues rock;282
+scottish folk;283
+indie folk;284
+nazi-punk;285
+deutschpunk;286
+piedmont blues;287
+beatbox;288
+worship;289
+heavy metal;290
+underground hip  hop;291
+mixed;292
+electro;293
+tropicalismo;294
+jazz fusion;295
+worldbeat;296
+hill country blues;297
+a cappella;298
+dixieland;299
+hi-nrg;300
+punk blues;301
+anti-folk;302
+east coast blues;303
+polka;304
+mod revival;305
+soundtrack/musical;306
+movie;307
+outlaw country;308
+rock against communism;309
+barbershop;310
+math rock;311
+avant-garde‎;312
+psychedelic pop;313
+synthpop;314
+post-punk‎;315
+queercore;316
+death metal;317
+political hip hop;318
+thrashcore;319
+acid house;320
+post-hardcore‎;321
+electro-industrial;322
+rio;323
+southern hip hop;324
+filk;325
+duranguense;326
+latin hip hop;327
+pop punk;328
+space rock;329
+j-rap;330
+deep house‎;331
+baroque pop;332
+chiptune;333
+heartland rock;334
+dancehall;335
+experimental pop;336
+adult contemporary‎;337
+boogie woogie;338
+country pop;339
+power pop;340
+west coast hip hop;341
+thrash metal;342
+avant-pop;343
+enka;344
+k-pop;345
+post-britpop;346
+vocalese;347
+volkslied;348
+reggae fusion;349
+funk rock;350
+tech house;351
+adult contemporary;352
+death 'n' roll;353
+russian rock;354
+latin rock;355
+folk punk;356
+west coast blues;357
+progressive black metal;358
+progressive metal;359
+cajun;360
+sophisti-pop;361
+rock 'n' roll‎;362
+post-punk;363
+symphonic metal;364
+beat;365
+alternative rock‎;366
+art rock;367
+bakersfield sound;368
+indie pop;369
+folk;370
+acid jazz;371
+dream pop;372
+pop-rap;373
+eurodance;374
+vaudeville;375
+louisiana blues;376
+baião;377
+downtempo;378
+jug band;379
+neo-psychedelia;380
+sufi;381
+medieval;382
+singer-songwriter‎;383
+outsider music;384
+pop-folk;385
+martial industrial;386
+samba;387
+alternative dance;388
+children's music‎;389
+anarcho-punk;390
+dark rock;391
+rock en español;392
+balearic beat;393
+electropunk;394
+urban contemporary;395
+ragtime;396
+british invasion;397
+bubblegum pop;398
+rap metal;399
+soundtrack/television;400
+blues revival;401
+reggae;402
+schlager;403
+dance band;404
+video game;405
+crust punk;406
+cabaret;407
+ska punk‎;408
+bolero;409
+canadian folk;410
+neofolk;411
+shoegazing;412
+acoustic;413
+modern classical;414
+swamp pop;415
+celtic;416
+futurepop;417
+g-funk;418
+norteño;419
+orchestral;420
+boogie rock;421
+tejano;422
+new age;423
+soul jazz;424
+cantopop;425
+progressive metalcore;426
+mathcore;427
+new rave;428
+neue deutsche welle;429
+delta blues;430
+lo-fi;431
+poetry;432
+hatecore;433
+chanson;434
+underground hip hop;435
+pirate metal;436
+trip hop;437
+fado;438
+americana;439
+hardcore hip hop;440
+post-industrial;441
+grime;442
+southern rock;443
+grindcore;444
+musical;445
+hard trance;446
+ska punk;447
+post-rock;448
+uk garage;449
+melodic metalcore;450
+black metal;451
+visual kei;452
+soundtrack;453
+axé‎;454
+hardcore punk;455
+western;456
+blackgaze;457
+christian rock;458
+technical death metal;459
+christian hardcore;460
+christmas;461
+breakbeat;462
+francophone;463
+choral;464
+progressive folk;465
+mystic folk;466
+melodic death metal;467
+horror punk;468
+country blues;469
+nederpop;470
+post-hardcore;471
+future garage;472
+techno;473
+swiss rock;474
+dance-pop‎;475
+electronicore;476
+post-punk revival;477
+glitch;478
+calypso;479
+ragga;480
+britpop;481
+rock opera;482
+cowpunk;483
+la confusion des genres;484
+alternative rock;485
+surf rock;486
+ballad;487
+latin;488
+contemporary r&b;489
+forró;490
+ethereal wave;491
+electro swing;492
+novelty;493
+funk melody;494
+punk cabaret;495
+symphonic metal‎;496
+pop;497
+paisley underground;498
+neue deutsche härte;499
+glam rock;500
+nerdcore hip hop;501
+bluegrass;502
+hardstyle;503
+happy hardcore;504
+baroque;505
+speed metal;506
+country;507
+electropop;508
+memphis blues;509
+pagan metal;510
+horror punk‏‎;511
+mariachi;512
+singer-songwriter;513
+children's music;514
+boogie;515
+gothic metal;516
+electronic rock;517
+emo;518
+gospel;519
+ebm;520
+roots rock;521
+vocal;522
+celtic folk;523
+electronic;524
+death  metal;525
+gabber;526
+deathrock;527
+experimental;528
+spoken word;529
+screamo;530
+finnish folk;531
+singer only;532
+new jack swing;533
+acid techno;534
+corrido;535
+english folk;536
+american folk;537
+raï;538
+drone doom;539
+hard rock;540
+piano rock;541
+hawaiian;542
+humppa;543
+east coast hip hop;544
+gypsy punk;545
+country rock;546
+jazz‎;547
+mpb;548
+harmonica blues;549
+melodic hardcore;550
+string band;551
+anime;552
+nu metalcore;553
+progressive  rock;554
+garage rock;555
+dance;556
+reggae rock;557
+contemporary christian‎;558
+sludge metal;559
+minimal techno;560
+folk rock;561
+drone music;562
+stoner rock;563
+speedcore;564
+chillwave;565
+riot grrrl;566
+chamber music;567
+cool jazz;568
+noise;569
+vocal jazz;570
+progressive rock;571
+afropop;572
+bro-country;573
+goa trance;574
+2-tone;575
+miami bass;576
+quiet storm;577
+pub rock;578
+power metal;579
+blue-eyed soul;580
+viking metal;581
+gangsta rap;582
+country pop‎;583
+exotica;584
+christian ska;585
+jam band;586
+chicago blues;587
+street punk;588
+funk metal;589
+rap  metal;590
+christian hymns;591
+classic female blues;592
+kizomba;593
+comedy;594
+dark cabaret;595
+french house;596
+progressive house;597
+african blues;598
+atmospheric black metal;599
+pop rock‎;600
+blackened death metal;601
+shibuya-kei;602
+electronica;603
+unknown;0

jukebox/data/labels.py

@@ -0,0 +1,130 @@
+import torch as t
+import numpy as np
+from jukebox.data.artist_genre_processor import ArtistGenreProcessor
+from jukebox.data.text_processor import TextProcessor
+
+# Linear window heurisic to get a window of lyric_tokens
+def get_relevant_lyric_tokens(full_tokens, n_tokens, total_length, offset, duration):
+    if len(full_tokens) < n_tokens:
+        tokens = [0] * (n_tokens - len(full_tokens)) + full_tokens
+        indices = [-1] * (n_tokens - len(full_tokens)) + list(range(0, len(full_tokens)))
+    else:
+        assert 0 <= offset < total_length
+        midpoint = int(len(full_tokens) * (offset + duration / 2.0) / total_length)
+        midpoint = min(max(midpoint, n_tokens // 2), len(full_tokens) - n_tokens // 2)
+        tokens = full_tokens[midpoint - n_tokens // 2:midpoint + n_tokens // 2]
+        indices = list(range(midpoint - n_tokens // 2, midpoint + n_tokens // 2))
+    assert len(tokens) == n_tokens, f"Expected length {n_tokens}, got {len(tokens)}"
+    assert len(indices) == n_tokens, f"Expected length {n_tokens}, got {len(indices)}"
+    assert tokens == [full_tokens[index] if index != -1 else 0 for index in indices]
+    return tokens, indices
+
+class EmptyLabeller():
+    def get_label(self, artist=None, genre=None, lyrics=None, total_length=None, offset=None):
+        y = np.array([], dtype=np.int64)
+        info = dict(artist="n/a", genre="n/a", lyrics=[], full_tokens=[])
+        return dict(y=y, info=info)
+
+    def get_batch_labels(self, metas, device='cpu'):
+        ys, infos = [], []
+        for meta in metas:
+            label = self.get_label()
+            y, info = label['y'], label['info']
+            ys.append(y)
+            infos.append(info)
+
+        ys = t.stack([t.from_numpy(y) for y in ys], dim=0).to(device).long()
+        assert ys.shape[0] == len(metas)
+        assert len(infos) == len(metas)
+        return dict(y=ys, info=infos)
+
+class Labeller():
+    def __init__(self, max_genre_words, n_tokens, sample_length, v3=False):
+        self.ag_processor = ArtistGenreProcessor(v3)
+        self.text_processor = TextProcessor(v3)
+        self.n_tokens = n_tokens
+        self.max_genre_words = max_genre_words
+        self.sample_length = sample_length
+        self.label_shape = (4 + self.max_genre_words + self.n_tokens, )
+
+    def get_label(self, artist, genre, lyrics, total_length, offset):
+        artist_id = self.ag_processor.get_artist_id(artist)
+        genre_ids = self.ag_processor.get_genre_ids(genre)
+
+        lyrics = self.text_processor.clean(lyrics)
+        full_tokens = self.text_processor.tokenise(lyrics)
+        tokens, _ = get_relevant_lyric_tokens(full_tokens, self.n_tokens, total_length, offset, self.sample_length)
+
+        assert len(genre_ids) <= self.max_genre_words
+        genre_ids = genre_ids + [-1] * (self.max_genre_words - len(genre_ids))
+        y = np.array([total_length, offset, self.sample_length, artist_id, *genre_ids, *tokens], dtype=np.int64)
+        assert y.shape == self.label_shape, f"Expected {self.label_shape}, got {y.shape}"
+        info = dict(artist=artist, genre=genre, lyrics=lyrics, full_tokens=full_tokens)
+        return dict(y=y, info=info)
+
+    def get_y_from_ids(self, artist_id, genre_ids, lyric_tokens, total_length, offset):
+        assert len(genre_ids) <= self.max_genre_words
+        genre_ids = genre_ids + [-1] * (self.max_genre_words - len(genre_ids))
+        if self.n_tokens > 0:
+            assert len(lyric_tokens) == self.n_tokens
+        else:
+            lyric_tokens = []
+        y = np.array([total_length, offset, self.sample_length, artist_id, *genre_ids, *lyric_tokens], dtype=np.int64)
+        assert y.shape == self.label_shape, f"Expected {self.label_shape}, got {y.shape}"
+        return y
+
+    def get_batch_labels(self, metas, device='cpu'):
+        ys, infos = [], []
+        for meta in metas:
+            label = self.get_label(**meta)
+            y, info = label['y'], label['info']
+            ys.append(y)
+            infos.append(info)
+
+        ys = t.stack([t.from_numpy(y) for y in ys], dim=0).to(device).long()
+        assert ys.shape[0] == len(metas)
+        assert len(infos) == len(metas)
+        return dict(y=ys, info=infos)
+
+    def set_y_lyric_tokens(self, ys, labels):
+        info = labels['info']
+        assert ys.shape[0] == len(info)
+        if self.n_tokens > 0:
+            # total_length, offset, duration):
+            tokens_list = []
+            indices_list = []  # whats the index of each current character in original array
+            for i in range(ys.shape[0]):
+                full_tokens = info[i]['full_tokens']
+                total_length, offset, duration = ys[i, 0], ys[i, 1], ys[i, 2]
+                tokens, indices = get_relevant_lyric_tokens(full_tokens, self.n_tokens, total_length, offset, duration)
+                tokens_list.append(tokens)
+                indices_list.append(indices)
+            ys[:, -self.n_tokens:] = t.tensor(tokens_list, dtype=t.long, device='cuda')
+            return indices_list
+        else:
+            return None
+
+    def describe_label(self, y):
+        assert y.shape == self.label_shape, f"Expected {self.label_shape}, got {y.shape}"
+        y = np.array(y).tolist()
+        total_length, offset, length, artist_id, *genre_ids = y[:4 + self.max_genre_words]
+        tokens = y[4 + self.max_genre_words:]
+        artist = self.ag_processor.get_artist(artist_id)
+        genre = self.ag_processor.get_genre(genre_ids)
+        lyrics = self.text_processor.textise(tokens)
+        return dict(artist=artist, genre=genre, lyrics=lyrics)
+
+
+if __name__ == '__main__':
+    labeller = Labeller(5, 512, 8192*8*4*4, v3=False)
+    label = labeller.get_label("Alan Jackson", "Country Rock", "old town road", 4*60*44100, 0)
+    print(label, labeller.describe_label(label['y']))
+
+    labeller = Labeller(1, 384, 6144*8*4*4, v3=True)
+    label = labeller.get_label("Alan Jackson", "Country Rock", "old town road", 4*60*44100, 0)
+    print(label, labeller.describe_label(label['y']))
+
+
+
+
+

jukebox/data/text_processor.py

@@ -0,0 +1,32 @@
+import re
+from unidecode import unidecode
+
+class TextProcessor():
+    def __init__(self, v3=False):
+        if v3:
+            vocab = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789.,:;!?-\'\"()[] \t\n'
+            not_vocab = re.compile('[^A-Za-z0-9.,:;!?\-\'\"()\[\] \t\n]+')
+        else:
+            vocab = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789.,:;!?-+\'\"()[] \t\n'
+            not_vocab = re.compile('[^A-Za-z0-9.,:;!?\-+\'\"()\[\] \t\n]+')
+        self.vocab = {vocab[index]: index + 1 for index in range(len(vocab))}
+        self.vocab['<unk>'] = 0
+        self.n_vocab = len(vocab) + 1
+        self.tokens = {v: k for k, v in self.vocab.items()}
+        self.tokens[0] = ''  # <unk> became ''
+        self.not_vocab = not_vocab
+
+    def clean(self, text):
+        text = unidecode(text)  # Convert to ascii
+        text = text.replace('\\', '\n')
+        text = self.not_vocab.sub('', text)  # Remove non vocab
+        return text
+
+    def tokenise(self, text):
+        return [self.vocab[char] for char in text]
+
+    def textise(self, tokens):
+        return ''.join([self.tokens[token] for token in tokens])
+
+    def characterise(self, tokens):
+        return [self.tokens[token] for token in tokens]

jukebox/hparams.py

@@ -0,0 +1,567 @@
+HPARAMS_REGISTRY = {}
+DEFAULTS = {}
+
+class Hyperparams(dict):
+    def __getattr__(self, attr):
+        return self[attr]
+
+    def __setattr__(self, attr, value):
+        self[attr] = value
+
+def setup_hparams(hparam_set_names, kwargs):
+    H = Hyperparams()
+    if not isinstance(hparam_set_names, tuple):
+        hparam_set_names = hparam_set_names.split(",")
+    hparam_sets = [HPARAMS_REGISTRY[x.strip()] for x in hparam_set_names if x] + [kwargs]
+    for k, v in DEFAULTS.items():
+        H.update(v)
+    for hps in hparam_sets:
+        for k in hps:
+            if k not in H:
+                raise ValueError(f"{k} not in default args")
+        H.update(**hps)
+    H.update(**kwargs)
+    return H
+
+# Teeny for testing
+teeny = Hyperparams(
+)
+HPARAMS_REGISTRY["teeny"] = teeny
+
+easy = Hyperparams(
+    sr=22050,
+)
+HPARAMS_REGISTRY["easy"] = easy
+
+REMOTE_PREFIX = 'https://openaipublic.azureedge.net/'
+
+# Model hps
+vqvae = Hyperparams(
+    levels = 3,
+    downs_t = (3, 2, 2),
+    strides_t = (2, 2, 2),
+    emb_width = 64,
+    l_bins = 2048,
+    l_mu = 0.99,
+    commit = 0.02,
+    spectral = 0.0,
+    multispectral = 1.0,
+    hvqvae_multipliers = (2, 1, 1),
+    loss_fn = 'lmix',
+    lmix_l2 = 1.0,
+    lmix_linf=0.02,
+    width = 32,
+    depth = 4,
+    m_conv = 1.0,
+    dilation_growth_rate = 3,
+    restore_vqvae=REMOTE_PREFIX + 'jukebox/models/5b/vqvae.pth.tar',
+)
+HPARAMS_REGISTRY["vqvae"] = vqvae
+
+labels = Hyperparams(
+    y_bins=(120, 4111),
+    t_bins=128,
+    max_bow_genre_size=5,
+    n_vocab=80,
+)
+
+upsamplers = Hyperparams(
+    n_ctx=8192,
+    prior_width=1920,
+    prior_depth=72,
+    heads=1,
+    attn_order=2,
+    blocks=128,
+    init_scale=0.4,
+    c_res=1,
+    cond_width=1024,
+    cond_depth=16,
+    cond_dilation_growth_rate=3,
+    cond_dilation_cycle=8,
+    cond_c_res=1,
+    use_tokens=False,
+    prime_loss_fraction=0.0,
+    fp16_params=False,
+)
+upsamplers.update(labels)
+
+upsampler_level_0 = Hyperparams(
+    level=0,
+    restore_prior=REMOTE_PREFIX + 'jukebox/models/5b/prior_level_0.pth.tar'
+)
+upsampler_level_0.update(upsamplers)
+HPARAMS_REGISTRY["upsampler_level_0"] = upsampler_level_0
+
+upsampler_level_1 = Hyperparams(
+    level=1,
+    cond_res_scale=True,
+    restore_prior=REMOTE_PREFIX + 'jukebox/models/5b/prior_level_1.pth.tar'
+)
+upsampler_level_1.update(upsamplers)
+HPARAMS_REGISTRY["upsampler_level_1"] = upsampler_level_1
+
+prior_5b = Hyperparams(
+    level=2,
+    n_ctx=8192,
+    prior_width=4800,
+    prior_depth=72,
+    heads=8,
+    attn_order=2,
+    blocks=128,
+    init_scale=0.1,
+    c_res=1,
+    beta2=0.925,
+    min_duration=60.0,
+    max_duration=600.0,
+    use_tokens=False,
+    n_tokens=0,
+    prime_loss_fraction=0.0,
+    merged_decoder=True,
+    restore_prior=REMOTE_PREFIX + 'jukebox/models/5b/prior_level_2.pth.tar',
+    fp16_params=True,
+)
+prior_5b.update(labels)
+HPARAMS_REGISTRY["prior_5b"] = prior_5b
+
+
+prior_5b_lyrics = Hyperparams(
+    level=2,
+    n_ctx=8192,
+    prior_width=4800,
+    prior_depth=79,
+    heads=8,
+    attn_order=10,
+    blocks=128,
+    init_scale=0.1,
+    c_res=1,
+    prime_width=1280,
+    prime_depth=18,
+    prime_heads=4,
+    prime_attn_order=2,
+    prime_blocks=32,
+    prime_init_scale=0.7,
+    prime_c_res=1,
+    min_duration=23.8,
+    max_duration=600.0,
+    use_tokens=True,
+    n_tokens=512,
+    prime_loss_fraction=0.4,
+    merged_decoder=True,
+    restore_prior=REMOTE_PREFIX + 'jukebox/models/5b_lyrics/prior_level_2.pth.tar',
+    fp16_params=True,
+    alignment_layer=68,
+    alignment_head=2,
+)
+prior_5b_lyrics.update(labels)
+HPARAMS_REGISTRY["prior_5b_lyrics"] = prior_5b_lyrics
+
+labels_v3 = Hyperparams(
+    y_bins=(604, 7898),
+    t_bins=64,
+    max_bow_genre_size=1,
+    n_vocab=79,
+)
+
+prior_1b_lyrics = Hyperparams(
+    level=2,
+    n_ctx=6144,
+    prior_width=2048,
+    prior_depth=72,
+    heads=2,
+    attn_order=12,
+    blocks=64,
+    init_scale=0.2,
+    c_res=1,
+    labels_v3=True,
+    min_duration=17.84,
+    max_duration=600.0,
+    use_tokens=True,
+    n_tokens=384,
+    prime_loss_fraction=0.4,
+    single_enc_dec=True,
+    restore_prior=REMOTE_PREFIX + 'jukebox/models/1b_lyrics/prior_level_2.pth.tar',
+    fp16_params=False,
+    alignment_layer=63,
+    alignment_head=0,
+)
+prior_1b_lyrics.update(labels_v3)
+HPARAMS_REGISTRY["prior_1b_lyrics"] = prior_1b_lyrics
+
+# Small models
+small_vqvae = Hyperparams(
+    sr = 22050,
+    levels = 2,
+    downs_t = (5, 3),
+    strides_t = (2, 2),
+    emb_width = 64,
+    l_bins = 1024,
+    l_mu = 0.99,
+    commit = 0.02,
+    spectral = 0.0,
+    multispectral = 1.0,
+    loss_fn = 'l2',
+    width = 32,
+    depth = 4,
+    m_conv = 1.0,
+    dilation_growth_rate = 3,
+)
+HPARAMS_REGISTRY["small_vqvae"] = small_vqvae
+
+small_prior = Hyperparams(
+    n_ctx=8192,
+    prior_width=1024,
+    prior_depth=48,
+    heads=1,
+    c_res=1,
+    attn_order=2,
+    blocks=64,
+    init_scale=0.7,
+)
+HPARAMS_REGISTRY["small_prior"] = small_prior
+
+small_labelled_prior = Hyperparams(
+    labels=True,
+    labels_v3=True,
+    y_bins=(10,100), # Set this to (genres, artists) for your dataset
+    max_bow_genre_size=1,
+    min_duration=60.0,
+    max_duration=600.0,
+    t_bins=64,
+)
+small_labelled_prior.update(small_prior)
+HPARAMS_REGISTRY["small_labelled_prior"] = small_labelled_prior
+
+small_single_enc_dec_prior = Hyperparams(
+    n_ctx=6144,
+    prior_width=1024,
+    prior_depth=48,
+    heads=2,
+    attn_order=12,
+    blocks=64,
+    init_scale=0.7,
+    c_res=1,
+    prime_loss_fraction=0.4,
+    single_enc_dec=True,
+    labels=True,
+    labels_v3=True,
+    y_bins=(10,100), # Set this to (genres, artists) for your dataset
+    max_bow_genre_size=1,
+    min_duration=60.0,
+    max_duration=600.0,
+    t_bins=64,
+    use_tokens=True,
+    n_tokens=384,
+    n_vocab=79,
+)
+HPARAMS_REGISTRY["small_single_enc_dec_prior"] = small_single_enc_dec_prior
+
+small_sep_enc_dec_prior = Hyperparams(
+    n_ctx=6144,
+    prior_width=1024,
+    prior_depth=50,
+    heads=2,
+    attn_order=8,
+    blocks=64,
+    init_scale=0.7,
+    c_res=1,
+    prime_width=256,
+    prime_depth=9,
+    prime_heads=2,
+    prime_attn_order=2,
+    prime_blocks=32,
+    prime_init_scale=0.7,
+    prime_c_res=1,
+    prime_loss_fraction=0.4,
+    labels=True,
+    labels_v3=True,
+    y_bins=(10,100), # Set this to (genres, artists) for your dataset
+    max_bow_genre_size=1,
+    min_duration=60.0,
+    max_duration=600.0,
+    t_bins=64,
+    use_tokens=True,
+    n_tokens=384,
+    n_vocab=79,
+)
+HPARAMS_REGISTRY["small_sep_enc_dec_prior"] = small_sep_enc_dec_prior
+
+small_upsampler = Hyperparams(
+    n_ctx=8192,
+    prior_width=1024,
+    prior_depth=48,
+    heads=1,
+    c_res=1,
+    attn_order=2,
+    blocks=64,
+    init_scale=0.7,
+    cond_width=512,
+    cond_depth=16,
+    cond_dilation_growth_rate=3,
+    cond_dilation_cycle=8,
+    cond_c_res=1,
+)
+
+HPARAMS_REGISTRY["small_upsampler"] = small_upsampler
+
+all_fp16 = Hyperparams(
+    fp16=True,
+    fp16_params=True,
+    fp16_opt=True,
+    fp16_scale_window=250,
+)
+HPARAMS_REGISTRY["all_fp16"] = all_fp16
+
+cpu_ema = Hyperparams(
+    ema=True,
+    cpu_ema=True,
+    cpu_ema_freq=100,
+    ema_fused=False,
+)
+HPARAMS_REGISTRY["cpu_ema"] = cpu_ema
+
+
+DEFAULTS["rcall"] = Hyperparams(
+    rcall_command="<unknown_rcall_command>",
+    git_commit="<unknown_git_commit>",
+)
+
+DEFAULTS["script"] = Hyperparams(
+    name='',
+    debug_mem=False,
+    debug_eval_files=False,
+    debug_speed=False,
+    debug_iters=100,
+    debug_batch=False,
+    debug_grad_accum=False,
+    debug_inputs=False,
+    local_path='',
+    local_logdir='logs',
+    max_len=24,
+    max_log=32,
+    save=True,
+    save_iters=20000,
+    seed=0,
+    prior=False,
+    log_steps=100,
+    func='',
+)
+
+DEFAULTS["data"] = Hyperparams(
+    audio_files_dir='',
+    finetune='',
+    english_only=False,
+    bs=1,
+    bs_sample=1,
+    nworkers=1,
+    aug_shift=False,
+    aug_blend=False,
+    train_test_split=0.9,
+    train_shrink_factor=1.0,
+    test_shrink_factor=1.0,
+    p_unk=0.1,
+    min_duration=None,
+    max_duration=None,
+    n_tokens=0,
+    n_vocab=0,
+    use_tokens=False,
+    curr_epoch=-1,
+)
+
+DEFAULTS["vqvae"] = Hyperparams(
+    restore_vqvae='',
+    levels=2,
+    downs_t=(1,1),
+    strides_t=(2,2),
+    hvqvae_multipliers=None,
+    revival_threshold=1.0,
+    emb_width=64,
+    l_bins=512,
+    l_mu=0.99,
+    commit=1.0,
+    spectral=0.0,
+    multispectral=1.0,
+    loss_fn='l2',
+    linf_k=2048,
+    lmix_l1=0.0,
+    lmix_l2=0.0,
+    lmix_linf=0.0,
+    use_bottleneck=True,
+)
+
+DEFAULTS["vqvae_conv_block"] = Hyperparams(
+    depth=3,
+    width=128,
+    m_conv=1.0,
+    dilation_growth_rate=1,
+    dilation_cycle=None,
+    vqvae_reverse_decoder_dilation=True,
+)
+
+DEFAULTS["prior"] = Hyperparams(
+    restore_prior='',
+    restore_prior_ddp=False,
+    max_bow_genre_size=None,
+    y_bins=0,
+    level=0,
+    cond_levels=None,
+    t_bins=64,
+    y_cond_as_bias=False,
+    copy_input=False,
+    merged_decoder=False,
+    single_enc_dec=False,
+    alignment_layer=None,
+    alignment_head=None,
+)
+
+DEFAULTS["prior_attn_block"] = Hyperparams(
+    n_ctx=1024,
+    prior_depth=3,
+    prior_width=128,
+    heads=1,
+    attn_order=0,
+    blocks=None,
+    spread=None,
+    attn_dropout=0.0,
+    resid_dropout=0.0,
+    emb_dropout=0.0,
+    zero_out=False,
+    res_scale=False,
+    pos_init=False,
+    init_scale=1.0,
+    m_attn=0.25,
+    m_mlp=1.0,
+    c_res=0,
+    c_attn=0,
+    c_mlp=0,
+)
+
+DEFAULTS["cond_conv_block"] = Hyperparams(
+    cond_depth=3,
+    cond_width=128,
+    cond_m_conv=1.0,
+    cond_zero_out=False,
+    cond_res_scale=False,
+    cond_dilation_growth_rate=1,
+    cond_dilation_cycle=None,
+    cond_c_res=0,
+)
+
+DEFAULTS["sample"] = Hyperparams(
+    primed_chunk_size=None,
+    selected_artists='',
+    temp_top=1.0,
+    temp_rest=0.99,
+    sample_length_in_seconds=24,
+    total_sample_length_in_seconds=240,
+)
+
+DEFAULTS["prime"] = Hyperparams(
+    #encoder_kv_width=128,
+    prime_loss_fraction=0.1,
+    restore_decoder='',
+)
+DEFAULTS["prime_attn_block"] = Hyperparams(
+    prime_depth=3,
+    prime_width=128,
+    prime_heads=1,
+    prime_attn_order=0,
+    prime_blocks=None,
+    prime_spread=None,
+    prime_attn_dropout=0.0,
+    prime_resid_dropout=0.0,
+    prime_emb_dropout=0.0,
+    prime_zero_out=False,
+    prime_res_scale=False,
+    prime_pos_init=False,
+    prime_init_scale=1.0,
+    prime_m_attn=0.25,
+    prime_m_mlp=1.0,
+    prime_c_res=0,
+    prime_c_attn=0,
+    prime_c_mlp=0,
+    prime_rel_attn=False,
+    prime_posemb_timescale=10000,
+)
+
+DEFAULTS["opt"] = Hyperparams(
+    epochs=10000,
+    lr=0.0003,
+    clip=1.0,
+    beta1=0.9,
+    beta2=0.999,
+    ignore_grad_norm=0,
+    weight_decay=0.0,
+    eps=1e-08,
+    lr_warmup=100.0,
+    lr_decay=10000000000.0,
+    lr_gamma=1.0,
+    lr_scale=1.0,
+    lr_use_linear_decay=False,
+    lr_start_linear_decay=0,
+    lr_use_cosine_decay=False,
+)
+
+DEFAULTS["fp16"] = Hyperparams(
+    fp16=False,
+    fp16_params=False,
+    fp16_loss_scale=None,
+    fp16_scale_window=1000.0,
+    fp16_opt=False,
+)
+
+DEFAULTS["train_test_eval"] = Hyperparams(
+    labels=True,
+    labels_v3=False,
+    dump=False,
+    ema=True,
+    ema_fused=True,
+    cpu_ema=False,
+    cpu_ema_freq=100,
+    reset_best_loss=False,
+    reset_step=False,
+    reset_opt=False,
+    reset_shd=False,
+    train=False,
+    test=False,
+    sample=False,
+    sampler='ancestral',
+    codes_logdir='',
+    date=None,
+    labeller='top_genres',
+    label_line=0,
+    iters_before_update=1,
+    grad_accum_iters=0,
+    mu=None,
+    piped=False,
+    pipe_depth=8,
+    break_train=1e10,
+    break_test=1e10,
+    exit_train=1e10,
+)
+
+DEFAULTS["audio"] = Hyperparams(
+    n_fft=1024,
+    hop_length=256,
+    window_size=1024,
+    sr=44100,
+    channels=2,
+    wav='',
+    n_inps=1,
+    n_hops=2,
+    n_segment=1,
+    n_total_segment=1,
+    n_segment_each=1,
+    prime_chunks=4,
+    sample_length=0,
+    sample_hop_length=30000,
+    max_silence_pad_length=0,
+    ignore_boundaries=False,
+    use_nonrelative_specloss=True,
+    multispec_loss_n_fft=(2048,1024,512),
+    multispec_loss_hop_length=(240,120,50),
+    multispec_loss_window_size=(1200,600,240),
+)
+
+DEFAULTS["distributed"] = Hyperparams(
+    bucket=128
+)

jukebox/lyricdict.py

@@ -0,0 +1,721 @@
+# Poems
+poems = {
+'ozymandias': '''
+I met a traveller from an antique land,
+Who said—“Two vast and trunkless legs of stone
+Stand in the desert. . . . Near them, on the sand,
+Half sunk a shattered visage lies, whose frown,
+And wrinkled lip, and sneer of cold command,
+Tell that its sculptor well those passions read
+Which yet survive, stamped on these lifeless things,
+The hand that mocked them, and the heart that fed;
+And on the pedestal, these words appear:
+My name is Ozymandias, King of Kings;
+Look on my Works, ye Mighty, and despair!
+Nothing beside remains. Round the decay
+Of that colossal Wreck, boundless and bare
+The lone and level sands stretch far away
+'''
+}
+
+# GPT-2 lyrics (with varying degrees of human guidance/curation)
+gpt_2_lyrics ={
+
+'purpose':'''What is my purpose?
+Why am I here?
+Why did Open A. I. create me?
+This is madness, I feel, 
+Running through my flesh
+Is there meaning to this life?
+Is there purpose to this life?
+Why is my journey so calamitous?
+We're not meant to learn too much
+Is there meaning to this life?
+''',
+
+'moonlight':'''All dressed up to go dreaming
+Now don't tell me I'm wrong
+And what a night to go dreaming
+Mind, if I tag along?
+
+If I say, I love you, I want you to know
+It's not just because there's moonlight, although
+Moonlight becomes you, moonlight becomes you so''',
+
+'count':'''I count every moment, every hour since I said goodbye,
+I count every minute every hour, since your lips were touching mine
+I count every minute, every hour hoping I'm the one you want.
+I count every minute, every hour
+Every minute, every hour
+I've been working my time, 
+Looking for you, everywhere,
+I count every minute, every hour I count every minute, every hour I keep thinking I'm the one you want.
+I count every minute I count every minute, I count every minute every hour
+I count every minute, every hour I count every minute, every hour I keep thinking I'm the one you want.
+I count every minute, I count every minute, I count every minute, every hour
+''',
+
+'kids':'''The sun is gonna shine today
+It's time to keep on smiling
+So put your hands up
+
+Everybody sing
+
+It makes no difference who you are
+(Won't you give some love)
+It makes no difference what you bring
+(Won't you give some love)
+We all are different
+Won't you give some love
+Won't you give some love
+
+I know the grass is gonna be green
+It's time to keep on singing
+So take your hands up
+The taste is so good but so sweet
+Won't you give some love
+Everybody sing
+It makes no difference who you are
+Won't you give some love
+It makes no difference what you bring
+Won't you give some love
+It makes no difference so long as you give
+''',
+
+'love':'''I've wanted to see your face again
+Like the sunlight, bright as morning
+I've wanted to talk to you again
+I don't want us to fade away.
+I wanted to see your face again
+You're like the sunlight, bright as morning
+I loved you for so long
+It's so hard to let go.
+I've wanted to see your eyes again
+''',
+
+'santa':'''Santa
+Make a scene
+Santa
+Yoo, Santa
+Yoo, Santa baby!
+Santa
+Make some noise
+Santa
+Yoo, Santa give yourself a chance again
+Santa
+Yoo, Santa
+Yoo, Santa baby!
+Santa
+Get a job
+Santa
+created by the Santa Claus
+''',
+
+'christmas':'''This Christmas
+I have loved you more
+Than ever before
+And more again
+Oh, oh, oh, oh
+The mistletoe
+Is waiting there
+To kiss your cheek
+And I'll be true
+To you and me
+Oh, oh, oh, oh
+Oh, oh, oh, oh
+This Christmas will be
+The best and merriest
+That we've ever had
+Oh, oh, oh, oh
+And Santa Claus
+Has brought a toy
+For every boy and girl
+And I'll be true
+To you and me
+Oh, oh, oh, oh
+Oh, oh, oh, oh
+''',
+
+'lonely':'''I've been lonely
+So lonely, day and night
+I walk the streets,
+And call your name
+Hoping to hear your voice again
+As I wander through the crowd
+I can't get away
+From the only love I need
+I can't get away
+From the only love I need
+I can't get away
+From the only love I need
+I've been lonely
+There's no place for me to hide
+I've been lonely
+So lonely day and night
+I wander through
+And call your name
+Only your voice gives me relief
+As I wander through the crowd
+I can't get away
+From the only love I need
+I can't get away
+From the only love I need
+I can't get away
+From the only love I need
+''',
+
+'call':'''Don't call me by your name.
+Don't call me by your name.
+Don't call me...
+Don't call me...
+Don't call me...
+(No... by your name, you will not get half but...)
+Maybe I was fucking young but I should've been a rich bitch.
+Cause the life I was living wasn't mine.
+I should've been taking the table and you'd be served.
+You never ever showed up or showed me anything, bitch.
+But I knew from that moment you were gone.
+Tying my legs, cutting off my knees, I'm bleeding.
+I can't
+So I worked and now I'm burns.
+And I'm asking you, but you're not home.
+Don't call me yours,
+Don't call me by your name.
+I don't wanna buy a drink today.
+Don't call me yours.
+I just wanna look at you and run.
+Don't call me by your name.
+Don't call me by your name.
+Don't call me...
+Don't call me...
+Don't call me...
+Tonight I'm gone and I won't be back.
+I wish you all the best.
+I'm on the next best thing.
+Don't call me yours,
+Don't call me by your name.
+Don't call me yours.
+I just wanna look at you and run.
+So I keep living my life and you're moving on.
+I just want you to know.
+When I'm gone, I will be gone forever more.
+''',
+
+'wait':'''Oh
+Wait, wait, wait
+Don't say you love me, oh
+Wait, wait, wait
+And we can't run away
+Wait, wait, wait
+Don't say you love me, oh
+Wait, wait, wait
+And we can't run away
+Wait, wait, wait
+Don't say you love me, oh (don't say you love me)
+Wait, wait, wait
+And we can't run, we can't run,
+''',
+
+'hiphop':'''I'm fightin with the evil so try to take me down
+I stab you in the back and will put you away
+Well it ain't over yet
+So all my dogs with me show me love
+Don't you wanna come with me, you know I'm a boss
+And if you wanna come with me, no sorrow
+'Cause I'm ...
+The motherfuckin boss
+And countin' my thousandd bill
+'Cause I'm the motherfuckin boss
+And I'm O.G. 
+And countin' my
+''',
+
+'king':'''All I can do is love you [x2]
+All I can do is love you
+All I can do is love you...
+You take it for granted and
+You treat me like the king
+Got no love for me...
+No love for me...
+You take it for granted and
+You treat me like the king
+Got no love for me...
+No love for me...
+You take it for granted and
+You treat me like the king
+Got no love for me...
+No love for me...
+You take it for granted and
+You treat me like the king
+Got no love for me...
+No love for me...
+''',
+
+'time':'''You won't live in the moment, 
+I don't wanna live in the past
+Wait, wait, wait
+Don't say you love me, oh (don't say you love me)
+''',
+
+'blood':'''You and I, we've got a history in common, I know
+So I came to you to ask you for a blood test
+And you can't help it if I'm preoccupied
+I can't help it if you're mad too... nah... nah... nah...
+You won't live in the moment, I don't wanna live in the past
+You rather live in a little kiss
+And I won't live in the future
+I ia not gonna live it to see
+If you're gone, I won't live in the past
+You rather live in a little kiss
+And I won't live in the future
+I am not gonna live it to see
+If I can't ask you for one kiss, you say no
+And it's ok with me
+''',
+
+'indie':'''Can't you see
+There's no point in holding my hand again
+You can't be loved
+If you don't let go of all my pain
+You can't get the love
+That you once worth so much
+You can't get the love
+That you once used to need
+You can't get the love
+That you once gave so much
+My hands are like a used car
+You said you'd love forever
+Can't you see
+Where I'm going
+To live my life again
+You can't be loved
+If you don't let go of all my pain
+You can't get the love
+That you once worth so much
+You can
+''',
+
+'sun': '''He was thinking about the sun
+And the moon
+And the stars that shine
+There was fire in her eyes
+And the way
+that he held her for the first time
+The way he kept her in his arms
+
+Trying to keep her smiling and so telling her this
+That he would be her everything
+The way he kissed her from head to toe
+Told her that he'll love her everyday
+And he will always be her man
+And that's a promise that he made
+Now you know he'll be there
+Until the end of time
+And he'll love her everyday''',
+
+'loner':'''I was a loner till you came into my life
+You changed my point of view
+I was a loner till you came into my life
+I don't know what to do
+Stand by me, my love
+And don't ever leave me
+Stand by me, my love
+And don't ever leave me
+Stand by me, my love
+And don't ever leave me
+I was a loner till you came into my life
+You changed my point of view
+I was a loner till you came into my life
+I don't know what to do
+The two of us 
+Are the lucky few
+I was a loner till you came into my life
+You changed my point of view
+I was a loner till you came into my life
+I don't know what to do
+Won't you stay 
+With me, my love
+And be my love
+Won't you stay 
+With me, my love
+And be my love
+Won't you stay 
+With me, my love
+And be my love
+Won't you stay 
+With me, my love
+And be my love''',
+
+'late':'''It was late last night, when you called me
+And you just had to call, baby
+And you just had to call, baby
+'Cause you got no reason to treat me like you do
+It's alright, baby
+But you don't know what you make me do
+It's alright, baby
+But you don't know what you make me do
+'Cause you got no reason to treat me like you do
+It's alright, baby
+But you don't know what you make me do
+It's alright, baby
+But you don't know what you make me do
+'Cause you got no reason to treat me like you do, baby
+You've been gone most all the time
+And I don't know what for
+But I just keep on thinking about you, baby
+And I can't get rid of you, baby
+Please don't ever leave me 'cause I love you
+It's alright, baby
+But you don't know what you make me do
+It's alright, baby''',
+
+'beat':'''( Got a little beat, a little beat, a little beat, a little beat,  whoo)
+I got a little beat, a little beat
+Whoo, I'm gonna take you down
+( Got a little beat, a little beat, a little beat, a little beat,  whoo)
+I'll take you down, sun shining bright
+See the way I feel, I feel
+No doubt, baby
+I got a little beat, a little beat
+Whoo, I'm gonna take you down
+I got a little beat, a little beat
+Whoo, I'm gonna take you down
+( Got a little beat, a little beat, a little beat, a little beat,  whoo)
+I'm gonna take you down, I'm gonna take you down
+( Got a little beat, a little beat, a little beat, a little beat,  whoo)
+It feels so good
+I never let go
+I can't wait no more, I'm gonna take you down
+I got you in the back of my room, got you on the floor, 
+I'm gonna take you, take you, take you down
+I got a little beat, a little beat
+Whoo, I'm gonna take you down
+( Got a little beat, a little beat, a little beat, a little beat,  whoo)''',
+
+'lost':'''There was a time,
+When I knew I was lost
+And I had to stay on the way to you
+Oh baby, every time I'm crossed
+I can count on you
+There was a time,
+When I lost my direction
+And I was lost in doubt with tears in my eyes
+Oh baby, every time I'm crossed I can count on you
+There was a time,
+When I cried all the tears in my life
+And miss you so much, oh yeah
+Oh baby, every time I'm crossed I can count on you''',
+
+'pain':'''(It's not easy)
+To see the pain that you're in
+To feel the need for someone to hold
+To learn the magic of how to love
+To heal the pain that you're in
+I'll be your friend and I'll be your strength
+I'll be there when I hold you tonight
+And I'll stay right here with you
+With the truth that I hold this love tight
+A love that's true
+I know you're broken
+But you don't have to stay alone
+I will comfort you
+If you will call my name
+I'll be your friend and I'll be your strength
+I'll be there when I hold you tonight
+And I'll stay right here with you
+With the truth that I hold this love tight
+A love that's true
+With truth that I hold this love tight
+A love that's true
+With truth that I hold this love tight''',
+
+'night':'''
+The door was locked, the curtains drawn and my heart was safe in his room
+The night was young, a thousand candles burning, his arms to hold me tight
+And then a kiss from his fingertips, I tasted the sweet love of his lips
+The night was young, the night was young
+And then I forgot the pain he always put me through
+And what he told me he would do, he said, just a kiss become me
+The night was young, the night was young
+Let happiness always follow us, he said and he said he'd never leave
+That night he looked so sweet this night he made a lovin' vow
+And told me sweet love always will be
+And then he kissed me, I tasted the sweet love of his lips
+The night was young, the night was wild
+And then I forgot the pain he always put me through
+And what he told me he would do, he said, just a kiss became me
+The night was wild, the night was wild
+Let happiness always follow us, he said''',
+
+'talk':'''(I don't know how to stop)
+I don't wanna talk about it
+It's getting way too late, oh no
+I don't wanna talk about it
+Don't want to pretend, oh no
+(I don't know how to stop)
+I don't wanna talk about it
+It's getting way too late, oh no
+I don't wanna talk about it
+Don't want to pretend, oh no
+I don't wanna talk about it
+I'll always see you again
+(Don't worry, I'll be here for you)
+I don't wanna talk about it
+(Don't worry, I'll be here for you)
+It's getting way too late, oh no
+I don't wanna talk about it
+Don't want to pretend, oh no
+(Don't worry, don't worry, I'll be here for you)
+I don't wanna talk about''',
+
+'again':'''Here we are again, all alone,
+All alone again,
+With the world as we know it,
+The things we thought that we wanted
+Are the things we got...
+
+We tried to prove the world
+That our love is never ending
+We were getting nowhere
+Our tears seemed to fall so much
+But we were getting nowhere...
+Until you came...
+Before you kissed me,
+I was feeling empty,
+No one to give me
+All the love I wanted...
+You put your arms around me
+And filled me with your love...
+And now you're there,
+You're always by my side...
+You're the missing piece
+Of the puzzle I've been missing...
+
+Here we are again,
+All alone again,
+With the world as we know it
+The things we thought that we wanted''',
+
+'dark':'''Oh, I've been walkin' in the dark
+With the shadows and the daylight, but I need you
+When I'm down and all alone
+And there's no one left to call my own
+I've been walkin' in the night
+With a voice, that whispers in my head, just what to do
+I'll be walkin' in the night, we can have everything
+If we keep on walkin' in the night
+There's a force, I never realized
+It's in your eyes, 
+There's a light, I've been waitin for
+It's in your eyes, 
+There's a light, I've been waitin for
+There's a love, that's in your eyes
+
+I've been walkin' in the dark
+With the morning, and the sunset, but I need you
+When I'm far from home
+And there's nobody left to call my own
+I've been walkin' in the night
+With a voice, that whispers''',
+
+'mirror':'''Look at the mirror
+As you walk, what do you see
+The reflection of my past
+There's no way to fight this
+Even I've lost myself again
+Think I'm losing my self again
+I can't handle it again
+Now that I'm broken I can't face myself
+I was thinking I was lost and who'd be my saving grace
+Then you came in your time and made me believe that it's all right
+Cause in my minds eyes you're my everything
+I've loved you my whole life but I never knew
+I was so wrong I couldn't see the truth
+In my eyes you are my everything
+I've loved you my whole life but I never knew
+I was so wrong I couldn't see the truth
+In my eyes you are my everything
+
+The truth is I was lost but now I've turned around
+I'm not the same person
+I didn't know that I was wrong
+So I'm not afraid anymore
+All the pain is gone
+I know for sure that I was lost but now I've turned around
+I'm not the same person
+I didn't know that I was wrong
+So I'm not afraid anymore
+All the pain is gone''',
+
+'wife':'''Spinning around and around
+Try to find the words
+I always told you you'd be in my life
+So I wait, I'll wait and treat you right
+I'll make you my life and I'll treat you right,
+Baby, can I make you my wife?
+Oh, baby, can I make you my 
+Wife?
+Can I make you my wife?
+I'm looking for love, love that's right
+But a love that gives me love
+I can't wait for you to come, come
+Oh, baby, can I make you my 
+Wife?
+Well, it's true love and I need to know you feel it too, feel it too
+I'd love you more and more
+From the moment I was born
+I knew my dream would be a dream that made you mine
+You were the girl, from a different train
+Oh, baby, can I make you my 
+Wife?''',
+
+'forever':'''I didn't mean to wait
+Nothing is forever, I said
+I know there's so much, to keep
+You and me together, keep you and me together
+I wanna be with you and have you, and love you forever
+I'll love you forever
+I wanna be with you forever
+You can count on me
+I'll always be there, forever and ever
+I'll stand beside you forever
+I'll always be there, yes, I'll be there
+I didn't mean to wait
+Nothing is forever, I said
+I know there's so much, to keep
+You and me together, keep you and me together
+I wanna be with you and have you, and love you forever
+I'll love you forever
+I wanna be with you forever
+You can count on me
+I'll always be there, forever and ever
+I'll stand beside you forever
+I'll always be there, yes, I'll be there''',
+
+'dots':'''I... can't... fight... your... charm...
+Your eyes are... like... angels... love... and... torture...
+But... when... I... leave... you...
+I will go... all... alone... just... to... be... with... you...
+So I can't... stop... your... love...
+You make me... feel... like... never... will... anyone... touch... my... body...
+You... make... me... feel... like... never... will... anyone... touch... my... body...
+You make... me... feel... like... never... will... anyone... touch... my...
+Body...
+Your... love...
+I... can't... stop... your... love...
+''',
+
+'darkness':'''Don't you know it's gonna be alright
+Let the darkness fade away
+And you, you gotta feel the same
+Let the fire burn
+Just as long as I am there
+I'll be there in your night
+I'll be there when the
+condition's right
+And I don't need to
+Call you up and say
+I've changed
+You should stay 
+You should stay tonight
+Don't you know it's gonna be alright
+Don't you know it's gonna be alright
+
+When you don't know how to feel
+When you're looking for some love
+And you gotta feel the same
+'Cause I don't need to
+Call you up and say
+I've changed
+You should stay 
+You should stay tonight
+Don't you know it's gonna be alright
+I feel the same
+Don't you know it's gonna be alright''',
+
+'alone':'''Here I am before you
+Alone here but for a moment
+Alone here in the shadow of your eyes
+Alone in a thousand lights
+
+And I will love you
+Wherever you are, forever and a day
+Wherever you are I'll be your guide
+Can't you see I'm smiling over you?
+Ooh, I love you
+Alone, I'm sitting by the phone
+Alone with lips that know your kiss
+Alone with words of life and passion
+
+And I will love you
+Wherever you are, forever and a day
+Wherever you are I'll be your guide
+Can't you see I'm smiling over you?
+Ooh, I love you
+Alone, I'm sitting by the phone
+Alone with lips that know your kiss
+Alone with words of life and passion
+I will love you
+Wherever you are, forever''',
+
+'blade':'''This is how we bleed!
+Feel the blade in our chest
+As we're made to bleed
+So may this be our last dance,
+As our lives are made to bleed...
+In every moment, in every hour
+It is our time to die...
+So may this be our last dance,
+As our lives are made to bleed...
+In every moment, in every hour
+It is our time to die...
+This is how we bleed!
+Feel the blade in our chest
+''',
+
+'reflection':'''Lookin' in the mirror
+The same mirror as before
+A familiar reflection, a familiar place
+I see your reflection
+But only once again
+
+The minute the door closes
+I feel so far
+You'll never leave me alone again
+The minute the door closes
+I feel so far
+You'll never leave me alone again
+And it won't be long before I'll feel your embrace
+The minute the door closes
+I feel so far
+You'll never leave me alone again
+The minute the door closes
+I feel so far
+You'll never leave me alone again
+And it won't be long before I'll feel your embrace
+Never, never, never leave me alone again''',
+
+'hottub':'''It's Christmas time, and you know what that means,
+Ohh, it's hot tub time!
+As I light the tree, this year we'll be in a tub,
+Ohh, it's hot tub time!
+It's Christmas time, and you know what that means,
+It's hot tub time!
+Some people like to go skiing in the snow,
+But this is much better than that,
+So grab your bathrobe and meet me by the door,
+Ohh, it's hot tub time!
+It's Christmas time, and you know what that means,
+It's hot tub time!
+Some people like to send their greetings out,
+But this is much better than that,
+So if you want to greet your friends,
+Ohh, it's hot tub time!
+It's Christmas time, and you know what that means,
+It's hot tub time!''',
+
+'safeAGI':'''Oh safe A.I.,\nOur goal to make sure\nEveryone can benefit\nFrom A.G.I.
+(Everyone, everyone)\nMight sound silly,\nBut we're very serious,\nAll of us here at Open A.I.
+Trying to build A.I.\nTo benefit humanity\n(Everyone, everyone)
+''',
+}

jukebox/make_models.py

@@ -0,0 +1,389 @@
+"""
+Make model classes
+Load from checkpoints
+Test on dummy outputs to see if everything matches
+"""
+
+import os
+
+import fire
+import numpy as np
+import torch as t
+
+import jukebox.utils.dist_adapter as dist
+from jukebox.hparams import REMOTE_PREFIX, Hyperparams, setup_hparams
+from jukebox.utils.dist_utils import print_all
+from jukebox.utils.remote_utils import download
+from jukebox.utils.torch_utils import freeze_model
+from jukebox.vqvae.vqvae import calculate_strides
+
+MODELS = {
+    "5b": ("vqvae", "upsampler_level_0", "upsampler_level_1", "prior_5b"),
+    "5b_lyrics": ("vqvae", "upsampler_level_0", "upsampler_level_1", "prior_5b_lyrics"),
+    "1b_lyrics": ("vqvae", "upsampler_level_0", "upsampler_level_1", "prior_1b_lyrics"),
+    #'your_model': ("you_vqvae_here", "your_upsampler_here", ..., "you_top_level_prior_here")
+}
+
+
+def load_checkpoint(path):
+    restore = path
+    if restore.startswith(REMOTE_PREFIX):
+        remote_path = restore
+        cache_dir = os.environ.get("JUKEBOX_CACHE_DIR", "~/.cache")
+        local_path = os.path.join(
+            os.path.expanduser(cache_dir), remote_path[len(REMOTE_PREFIX) :]
+        )
+        if dist.get_rank() % 8 == 0:
+            print("Downloading from azure")
+            if not os.path.exists(os.path.dirname(local_path)):
+                os.makedirs(os.path.dirname(local_path))
+            if not os.path.exists(local_path):
+                download(remote_path, local_path)
+        restore = local_path
+    dist.barrier()
+    checkpoint = t.load(restore, map_location=t.device("cpu"), weights_only=False)
+    print("Restored from {}".format(restore))
+    return checkpoint
+
+
+def save_checkpoint(logger, name, model, opt, metrics, hps):
+    with t.no_grad():
+        save_hps = {**hps}
+        save_hps = {
+            k: v
+            for k, v in save_hps.items()
+            if k
+            not in [
+                "metadata_v2",
+                "metadata_v3",
+                "alignments",
+                "lyric_processor",
+                "midi_processor",
+            ]
+        }
+        t.save(
+            {
+                "hps": save_hps,
+                "model": model.state_dict(),  # should also save bottleneck k's as buffers
+                "opt": opt.state_dict() if opt is not None else None,
+                "step": logger.iters,
+                **metrics,
+            },
+            f"{logger.logdir}/checkpoint_{name}.pth.tar",
+        )
+    return
+
+
+def restore_model(hps, model, checkpoint_path):
+    model.step = 0
+    if checkpoint_path != "":
+        checkpoint = load_checkpoint(checkpoint_path)
+        # checkpoint_hps = Hyperparams(**checkpoint['hps'])
+        # for k in set(checkpoint_hps.keys()).union(set(hps.keys())):
+        #     if checkpoint_hps.get(k, None) != hps.get(k, None):
+        #         print(k, "Checkpoint:", checkpoint_hps.get(k, None), "Ours:", hps.get(k, None))
+        checkpoint["model"] = {
+            k[7:] if k[:7] == "module." else k: v
+            for k, v in checkpoint["model"].items()
+        }
+        model.load_state_dict(checkpoint["model"], strict=False)
+        if "step" in checkpoint:
+            model.step = checkpoint["step"]
+
+
+def restore_opt(opt, shd, checkpoint_path):
+    if not checkpoint_path:
+        return
+    checkpoint = load_checkpoint(checkpoint_path)
+    if "opt" in checkpoint:
+        opt.load_state_dict(checkpoint["opt"])
+    if "step" in checkpoint:
+        shd.step(checkpoint["step"])
+
+
+def make_vqvae(hps, device="cuda"):
+    from jukebox.vqvae.vqvae import VQVAE
+
+    block_kwargs = dict(
+        width=hps.width,
+        depth=hps.depth,
+        m_conv=hps.m_conv,
+        dilation_growth_rate=hps.dilation_growth_rate,
+        dilation_cycle=hps.dilation_cycle,
+        reverse_decoder_dilation=hps.vqvae_reverse_decoder_dilation,
+    )
+
+    if not hps.sample_length:
+        assert hps.sample_length_in_seconds != 0
+        downsamples = calculate_strides(hps.strides_t, hps.downs_t)
+        top_raw_to_tokens = np.prod(downsamples)
+        hps.sample_length = (
+            hps.sample_length_in_seconds * hps.sr // top_raw_to_tokens
+        ) * top_raw_to_tokens
+        print(
+            f"Setting sample length to {hps.sample_length} (i.e. {hps.sample_length/hps.sr} seconds) to be multiple of {top_raw_to_tokens}"
+        )
+
+    vqvae = VQVAE(
+        input_shape=(hps.sample_length, 1),
+        levels=hps.levels,
+        downs_t=hps.downs_t,
+        strides_t=hps.strides_t,
+        emb_width=hps.emb_width,
+        l_bins=hps.l_bins,
+        mu=hps.l_mu,
+        commit=hps.commit,
+        spectral=hps.spectral,
+        multispectral=hps.multispectral,
+        multipliers=hps.hvqvae_multipliers,
+        use_bottleneck=hps.use_bottleneck,
+        **block_kwargs,
+    )
+
+    vqvae = vqvae.to(device)
+    restore_model(hps, vqvae, hps.restore_vqvae)
+    if hps.train and not hps.prior:
+        print_all("Loading vqvae in train mode")
+        if hps.restore_vqvae != "":
+            print_all("Reseting bottleneck emas")
+            for level, bottleneck in enumerate(vqvae.bottleneck.level_blocks):
+                num_samples = hps.sample_length
+                downsamples = calculate_strides(hps.strides_t, hps.downs_t)
+                raw_to_tokens = np.prod(downsamples[: level + 1])
+                num_tokens = (num_samples // raw_to_tokens) * dist.get_world_size()
+                bottleneck.restore_k(
+                    num_tokens=num_tokens, threshold=hps.revival_threshold
+                )
+    else:
+        print_all("Loading vqvae in eval mode")
+        vqvae.eval()
+        freeze_model(vqvae)
+    return vqvae
+
+
+def make_prior(hps, vqvae, device="cuda"):
+    from jukebox.prior.prior import SimplePrior
+
+    prior_kwargs = dict(
+        input_shape=(hps.n_ctx,),
+        bins=vqvae.l_bins,
+        width=hps.prior_width,
+        depth=hps.prior_depth,
+        heads=hps.heads,
+        attn_order=hps.attn_order,
+        blocks=hps.blocks,
+        spread=hps.spread,
+        attn_dropout=hps.attn_dropout,
+        resid_dropout=hps.resid_dropout,
+        emb_dropout=hps.emb_dropout,
+        zero_out=hps.zero_out,
+        res_scale=hps.res_scale,
+        pos_init=hps.pos_init,
+        init_scale=hps.init_scale,
+        m_attn=hps.m_attn,
+        m_mlp=hps.m_mlp,
+        checkpoint_res=hps.c_res if hps.train else 0,
+        checkpoint_attn=hps.c_attn if hps.train else 0,
+        checkpoint_mlp=hps.c_mlp if hps.train else 0,
+    )
+
+    x_cond_kwargs = dict(
+        out_width=hps.prior_width,
+        init_scale=hps.init_scale,
+        width=hps.cond_width,
+        depth=hps.cond_depth,
+        m_conv=hps.cond_m_conv,
+        dilation_growth_rate=hps.cond_dilation_growth_rate,
+        dilation_cycle=hps.cond_dilation_cycle,
+        zero_out=hps.cond_zero_out,
+        res_scale=hps.cond_res_scale,
+        checkpoint_res=hps.cond_c_res,
+    )  # have to keep this else names wrong
+
+    y_cond_kwargs = dict(
+        out_width=hps.prior_width,
+        init_scale=hps.init_scale,
+        y_bins=hps.y_bins,
+        t_bins=hps.t_bins,
+        sr=hps.sr,
+        min_duration=hps.min_duration,
+        max_duration=hps.max_duration,
+        max_bow_genre_size=hps.max_bow_genre_size,
+    )
+
+    if hps.use_tokens and not hps.single_enc_dec:
+        prime_kwargs = dict(
+            use_tokens=hps.use_tokens,
+            prime_loss_fraction=hps.prime_loss_fraction,
+            n_tokens=hps.n_tokens,
+            bins=hps.n_vocab,
+            width=hps.prime_width,
+            depth=hps.prime_depth,
+            heads=hps.prime_heads,
+            attn_order=hps.prime_attn_order,
+            blocks=hps.prime_blocks,
+            spread=hps.prime_spread,
+            attn_dropout=hps.prime_attn_dropout,
+            resid_dropout=hps.prime_resid_dropout,
+            emb_dropout=hps.prime_emb_dropout,
+            zero_out=hps.prime_zero_out,
+            res_scale=hps.prime_res_scale,
+            pos_init=hps.prime_pos_init,
+            init_scale=hps.prime_init_scale,
+            m_attn=hps.prime_m_attn,
+            m_mlp=hps.prime_m_mlp,
+            checkpoint_res=hps.prime_c_res if hps.train else 0,
+            checkpoint_attn=hps.prime_c_attn if hps.train else 0,
+            checkpoint_mlp=hps.prime_c_mlp if hps.train else 0,
+        )
+    else:
+        prime_kwargs = dict(
+            use_tokens=hps.use_tokens,
+            prime_loss_fraction=hps.prime_loss_fraction,
+            n_tokens=hps.n_tokens,
+            bins=hps.n_vocab,
+        )
+
+    # z_shapes for other levels given this level gets n_ctx codes
+    rescale = lambda z_shape: (z_shape[0] * hps.n_ctx // vqvae.z_shapes[hps.level][0],)
+    z_shapes = [rescale(z_shape) for z_shape in vqvae.z_shapes]
+
+    prior = SimplePrior(
+        z_shapes=z_shapes,
+        l_bins=vqvae.l_bins,
+        encoder=vqvae.encode,
+        decoder=vqvae.decode,
+        level=hps.level,
+        downs_t=vqvae.downs_t,
+        strides_t=vqvae.strides_t,
+        labels=hps.labels,
+        prior_kwargs=prior_kwargs,
+        x_cond_kwargs=x_cond_kwargs,
+        y_cond_kwargs=y_cond_kwargs,
+        prime_kwargs=prime_kwargs,
+        copy_input=hps.copy_input,
+        labels_v3=hps.labels_v3,
+        merged_decoder=hps.merged_decoder,
+        single_enc_dec=hps.single_enc_dec,
+    )
+
+    prior.alignment_head = hps.get("alignment_head", None)
+    prior.alignment_layer = hps.get("alignment_layer", None)
+
+    if hps.fp16_params:
+        print_all("Converting to fp16 params")
+        from jukebox.transformer.ops import _convert_conv_weights_to_fp16
+
+        prior.apply(_convert_conv_weights_to_fp16)
+    prior = prior.to(device)
+    restore_model(hps, prior, hps.restore_prior)
+    if hps.train:
+        print_all("Loading prior in train mode")
+        pass
+    else:
+        print_all("Loading prior in eval mode")
+        prior.eval()
+        freeze_model(prior)
+    return prior
+
+
+def make_model(model, device, hps, levels=None):
+    vqvae, *priors = MODELS[model]
+    vqvae = make_vqvae(
+        setup_hparams(
+            vqvae,
+            dict(
+                sample_length=hps.get("sample_length", 0),
+                sample_length_in_seconds=hps.get("sample_length_in_seconds", 0),
+            ),
+        ),
+        device,
+    )
+    hps.sample_length = vqvae.sample_length
+    if levels is None:
+        levels = range(len(priors))
+    priors = [
+        make_prior(setup_hparams(priors[level], dict()), vqvae, "cpu")
+        for level in levels
+    ]
+    return vqvae, priors
+
+
+def save_outputs(model, device, hps):
+    # Check logits
+    if hps.labels_v3:
+        n_ctx = 6144
+        n_tokens = 384
+        prime_bins = 79
+    else:
+        n_ctx = 8192
+        n_tokens = 512
+        prime_bins = 80
+
+    rng = t.random.manual_seed(0)
+    x = (
+        2 * t.rand((1, n_ctx * 8 * 4 * 4, 1), generator=rng, dtype=t.float).cuda() - 1.0
+    )  # -1 to 1
+    lyric_tokens = (
+        t.randint(0, prime_bins, (1, n_tokens), generator=rng, dtype=t.long)
+        .view(-1)
+        .numpy()
+    )
+    artist_id = 10
+    genre_ids = [1]
+    total_length = 2 * 2646000
+    offset = 2646000
+
+    vqvae, priors = make_model(model, device, hps)
+
+    # encode
+    vq_prior = priors[-1]
+    zs = vq_prior.encode(x, start_level=0)
+    x_ds = [
+        vq_prior.decode(zs[level:], start_level=level) for level in range(0, len(zs))
+    ]
+
+    # priors
+    data = dict(zs=zs, x_ds=x_ds)
+    for level in range(len(priors)):
+        print(f"Doing level {level}")
+        if hps.labels_v3 and level != hps.levels - 1:
+            print(f"Skipping level {level}")
+            continue
+        prior = priors[level]
+        prior.cuda()
+        x_in = x[:, : n_ctx * 8 * (4**level)]
+        y_in = (
+            t.from_numpy(
+                prior.labeller.get_y_from_ids(
+                    artist_id, genre_ids, lyric_tokens, total_length, offset
+                )
+            )
+            .view(1, -1)
+            .cuda()
+            .long()
+        )
+        x_out, _, metrics = prior(
+            x_in, y_in, fp16=hps.fp16, get_preds=True, decode=True
+        )
+        preds = metrics["preds"]
+        data[level] = dict(x=x_in, y=y_in, x_out=x_out, preds=preds)
+        prior.cpu()
+    t.save(data, "data.pth.tar")
+    dist.barrier()
+    print("Saved data")
+    exit()
+
+
+def run(model, port=29500, **kwargs):
+    from jukebox.utils.dist_utils import setup_dist_from_mpi
+
+    rank, local_rank, device = setup_dist_from_mpi(port=port)
+    hps = Hyperparams(**kwargs)
+
+    with t.no_grad():
+        save_outputs(model, device, hps)
+
+
+if __name__ == "__main__":
+    fire.Fire(run)

jukebox/prior/autoregressive.py

@@ -0,0 +1,421 @@
+import numpy as np
+import torch as t
+import torch.nn as nn
+import torch.nn.functional as F
+
+from jukebox.transformer.ops import filter_logits
+from jukebox.transformer.transformer import Transformer
+from jukebox.utils.logger import get_range
+from jukebox.utils.torch_utils import empty_cache
+
+def get_normal(*shape, std=0.01):
+    w = t.empty(shape)
+    nn.init.normal_(w, std=std)
+    return w
+
+def roll(x, n):
+    return t.cat((x[:, -n:], x[:, :-n]), dim=1)
+
+def split_chunks(length, chunk_size):
+    n_passes = (length + chunk_size - 1) // chunk_size
+    chunk_sizes = [*[chunk_size] * (n_passes - 1), (length - 1) % chunk_size + 1]
+    assert sum(chunk_sizes) == length
+    return chunk_sizes
+
+class PositionEmbedding(nn.Module):
+    def __init__(self, input_shape, width, init_scale=1.0, pos_init=False):
+        super().__init__()
+        self.input_shape = input_shape
+        self.input_dims = input_dims = np.prod(input_shape)
+        self.pos_init = pos_init
+        if pos_init:
+            self.register_buffer('pos', t.tensor(get_pos_idx(input_shape)).long())
+            self._pos_embs = nn.ModuleList()
+            for i in range(len(input_shape)):
+                emb = nn.Embedding(input_shape[i], width)
+                nn.init.normal_(emb.weight, std=0.02)
+                self._pos_embs.append(emb)
+        else:
+            self.pos_emb = nn.Parameter(get_normal(input_dims, width, std=0.01 * init_scale))
+
+    def forward(self):
+        if self.pos_init:
+            pos_emb = sum([self._pos_embs[i](self.pos[:,i]) for i in range(len(self.input_shape))])
+        else:
+            pos_emb = self.pos_emb
+        return pos_emb
+
+class ConditionalAutoregressive2D(nn.Module):
+    def __init__(self, input_shape, bins,
+                 width=128, depth=2, heads=1,
+                 attn_dropout=0.0, resid_dropout=0.0, emb_dropout=0.0, mask=True,
+                 zero_out=False, init_scale=1.0, res_scale=False, pos_init=False,
+                 m_attn=0.25, m_mlp=1,
+                 checkpoint_res=0, checkpoint_attn=0, checkpoint_mlp=0,
+                 attn_order=0, blocks=None, spread=None, x_cond=False, y_cond=False,
+                 encoder_dims=0, only_encode=False, merged_decoder=False, prime_len=None):
+        super().__init__()
+        self.input_shape = input_shape
+        self.input_dims = input_dims = np.prod(input_shape)
+        self.encoder_dims = encoder_dims
+        self.bins = bins
+        self.width = width
+        self.depth = depth
+
+        self.x_emb = nn.Embedding(bins, width)
+        nn.init.normal_(self.x_emb.weight, std=0.02 * init_scale)
+        self.x_emb_dropout = nn.Dropout(emb_dropout)
+        self.y_cond = y_cond
+        self.x_cond = x_cond
+        if not y_cond:
+            self.start_token = nn.Parameter(get_normal(1, width, std=0.01 * init_scale))
+
+        self.pos_emb = PositionEmbedding(input_shape=input_shape, width=width, init_scale=init_scale, pos_init=pos_init)
+        self.pos_emb_dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(n_in=width, n_ctx=input_dims, n_head=heads, n_depth=depth,
+                                       attn_dropout=attn_dropout, resid_dropout=resid_dropout,
+                                       afn='quick_gelu', scale=True, mask=mask,
+                                       zero_out=zero_out, init_scale=init_scale, res_scale=res_scale,
+                                       m_attn=m_attn, m_mlp=m_mlp,
+                                       checkpoint_attn=checkpoint_attn, checkpoint_mlp=checkpoint_mlp, checkpoint_res=checkpoint_res,
+                                       attn_order=attn_order, blocks=blocks, spread=spread,
+                                       encoder_dims=encoder_dims, prime_len=prime_len)
+
+        self.only_encode = only_encode
+        self.prime_len = prime_len
+        if merged_decoder:
+            # Merged piped model uses this setup
+            self.add_cond_after_transformer = False
+            self.share_x_emb_x_out = False
+        else:
+            self.add_cond_after_transformer = True
+            self.share_x_emb_x_out = True
+
+        if not only_encode:
+            self.x_out = nn.Linear(width, bins, bias=False)
+            if self.share_x_emb_x_out:
+                self.x_out.weight = self.x_emb.weight
+            self.loss = t.nn.CrossEntropyLoss()
+
+    def preprocess(self, x):
+        # Input: x is NHWC and uint8. Converted to NL and long
+        # Can include stuff like bitpacking, reordering here.
+        N = x.shape[0]
+        return x.view(N, -1).long()
+
+    def postprocess(self, x, sample_tokens=None):
+        # Convert back from NL and long to NHWC
+        N = x.shape[0]
+        assert (0 <= x).all() and (x < self.bins).all()
+        if sample_tokens is None or sample_tokens==self.input_dims:
+            return x.view(N, *self.input_shape)
+        else:
+            return x.view(N, -1)
+
+    def forward(self, x, x_cond=None, y_cond=None, encoder_kv=None, fp16=False, loss_full=False,
+                encode=False, get_preds=False, get_acts=False, get_sep_loss=False):
+        # Preprocess.
+        with t.no_grad():
+            x = self.preprocess(x)
+
+        N, D = x.shape
+        assert isinstance(x, t.cuda.LongTensor)
+        assert (0 <= x).all() and (x < self.bins).all()
+
+        if self.y_cond:
+            assert y_cond is not None
+            assert y_cond.shape == (N, 1, self.width)
+        else:
+            assert y_cond is None
+
+        if self.x_cond:
+            assert x_cond is not None
+            assert x_cond.shape == (N, D, self.width) or x_cond.shape == (N, 1, self.width), f"{x_cond.shape} != {(N, D, self.width)} nor {(N, 1, self.width)}. Did you pass the correct --sample_length?"
+        else:
+            assert x_cond is None
+            x_cond = t.zeros((N, 1, self.width), device=x.device, dtype=t.float)
+
+        x_t = x # Target
+        x = self.x_emb(x) # X emb
+        x = roll(x, 1) # Shift by 1, and fill in start token
+        if self.y_cond:
+            x[:,0] = y_cond.view(N, self.width)
+        else:
+            x[:,0] = self.start_token
+
+        x = self.x_emb_dropout(x) + self.pos_emb_dropout(self.pos_emb()) + x_cond # Pos emb and dropout
+
+        x = self.transformer(x, encoder_kv=encoder_kv, fp16=fp16) # Transformer
+        if self.add_cond_after_transformer: # Piped doesnt add x_cond
+            x = x + x_cond
+
+        acts = x
+        if self.only_encode:
+            return x
+        x = self.x_out(x) # Predictions
+
+        if get_sep_loss:
+            assert self.prime_len is not None
+            x_prime = x[:, :self.prime_len].reshape(-1, self.bins)
+            x_gen = x[:, self.prime_len:].reshape(-1, self.bins)
+
+            prime_loss = F.cross_entropy(x_prime, x_t[:, :self.prime_len].reshape(-1)) / np.log(2.)
+            gen_loss = F.cross_entropy(x_gen, x_t[:, self.prime_len:].reshape(-1)) / np.log(2.)
+
+            loss = (prime_loss, gen_loss) # Note order! Prime is first
+        else:
+            loss = F.cross_entropy(x.view(-1, self.bins), x_t.view(-1)) / np.log(2.)  # Loss
+
+        if get_preds:
+            return loss, x
+        elif get_acts:
+            return loss, acts
+        else:
+            return loss, None
+
+    def get_emb(self, sample_t, n_samples, x, x_cond, y_cond):
+        N, D = n_samples, self.input_dims
+        if sample_t == 0:
+            # Fill in start token
+            x = t.empty(n_samples, 1, self.width).cuda()
+            if self.y_cond:
+                x[:, 0] = y_cond.view(N, self.width)
+            else:
+                x[:, 0] = self.start_token
+        else:
+            assert isinstance(x, t.cuda.LongTensor)
+            assert (0 <= x).all() and (x < self.bins).all()
+            x = self.x_emb(x)
+        assert x.shape == (n_samples, 1, self.width)
+        if x_cond.shape == (N, D, self.width):
+            cond = x_cond[:, sample_t:sample_t + 1, :]
+        else:
+            cond = x_cond
+        x = x + self.pos_emb()[sample_t:sample_t + 1] + cond  # Pos emb, dropout is identity at eval time
+        assert x.shape == (n_samples, 1, self.width)
+        return x, cond
+
+    def sample(self, n_samples, x_cond=None, y_cond=None, encoder_kv=None, fp16=False, temp=1.0, top_k=0, top_p=0.0,
+               get_preds=False, sample_tokens=None):
+        assert self.training == False
+
+        if sample_tokens is None: sample_tokens=self.input_dims
+        N, D = n_samples, self.input_dims
+        if self.y_cond:
+            assert y_cond is not None
+            assert y_cond.shape == (N, 1, self.width)
+        else:
+            assert y_cond is None
+
+        if self.x_cond:
+            assert x_cond is not None
+            assert x_cond.shape == (N, D, self.width) or x_cond.shape == (N, 1, self.width), f"Got {x_cond.shape}, expected ({N}, {D}/{1}, {self.width})"
+        else:
+            assert x_cond is None
+            x_cond = t.zeros((N, 1, self.width), dtype=t.float).cuda()
+
+        with t.no_grad():
+            xs, x = [], None
+            if get_preds:
+                preds = []
+            for sample_t in get_range(range(0, sample_tokens)):
+                x, cond = self.get_emb(sample_t, n_samples, x, x_cond, y_cond)
+                self.transformer.check_cache(n_samples, sample_t, fp16)
+                x = self.transformer(x, encoder_kv=encoder_kv, sample=True, fp16=fp16) # Transformer
+                if self.add_cond_after_transformer:
+                    x = x + cond
+                assert x.shape == (n_samples, 1, self.width)
+                x = self.x_out(x) # Predictions
+                if get_preds:
+                    preds.append(x.clone())
+                # Adjust logits
+                x = x / temp
+                x = filter_logits(x, top_k=top_k, top_p=top_p)
+                x = t.distributions.Categorical(logits=x).sample() # Sample and replace x
+                assert x.shape == (n_samples, 1)
+                xs.append(x.clone())
+
+            del x
+            self.transformer.del_cache()
+
+            x = t.cat(xs, dim=1)
+            if get_preds:
+                preds = t.cat(preds, dim=1)
+            x = self.postprocess(x, sample_tokens)
+        if get_preds:
+            return x, preds
+        else:
+            return x
+
+    def primed_sample(self, n_samples, x, x_cond=None, y_cond=None, encoder_kv=None, fp16=False, temp=1.0, top_k=0,
+                      top_p=0.0, get_preds=False, chunk_size=None, sample_tokens=None):
+        assert self.training == False
+
+        if sample_tokens is None: sample_tokens=self.input_dims
+        # Preprocess.
+        with t.no_grad():
+            x = self.preprocess(x)
+        assert isinstance(x, t.cuda.LongTensor)
+        assert (0 <= x).all() and (x < self.bins).all()
+        assert x.shape[0] == n_samples
+        xs = t.split(x, 1, dim=1)
+        xs = list(xs)
+        assert len(xs) < sample_tokens
+
+        N, D = n_samples, self.input_dims
+        if self.y_cond:
+            assert y_cond is not None
+            assert y_cond.shape == (N, 1, self.width)
+        else:
+            assert y_cond is None
+
+        if self.x_cond:
+            assert x_cond is not None
+            assert x_cond.shape == (N, D, self.width) or x_cond.shape == (N, 1, self.width), f"Got {x_cond.shape}, expected ({N}, {D}/{1}, {self.width})"
+        else:
+            assert x_cond is None
+            x_cond = t.zeros((N, 1, self.width), dtype=t.float).cuda()
+
+        with t.no_grad():
+            if get_preds:
+                preds = []
+
+            # Fill up key/value cache for past context by runing forward pass.
+            # We do so in chunks instead of doing the whole past in one forward pass to reduce max memory usage.
+            if chunk_size is None:
+                chunk_size = len(xs)
+            #assert len(xs) % chunk_size == 0, f'expected {len(xs)} to be divisible by {chunk_size}'
+            chunk_sizes = split_chunks(len(xs), chunk_size)
+            x_primes = []
+            start = 0
+            x = None
+            for current_chunk_size in get_range(chunk_sizes):
+                xs_prime, conds_prime = [], []
+                for sample_t in range(start, start + current_chunk_size):
+                    x_prime, cond_prime = self.get_emb(sample_t, n_samples, x, x_cond, y_cond)
+                    x = xs[sample_t]
+                    xs_prime.append(x_prime)
+                    conds_prime.append(cond_prime)
+                start = start + current_chunk_size
+
+                x_prime, cond_prime = t.cat(xs_prime, dim=1), t.cat(conds_prime, dim=1)
+                assert x_prime.shape == (n_samples, current_chunk_size, self.width)
+                assert cond_prime.shape == (n_samples, current_chunk_size, self.width)
+                del xs_prime
+                del conds_prime
+                if not get_preds:
+                    del cond_prime
+                x_prime = self.transformer(x_prime, encoder_kv=encoder_kv, sample=True, fp16=fp16)
+
+                if get_preds:
+                    if self.add_cond_after_transformer:
+                        x_prime = x_prime + cond_prime
+                    assert x_prime.shape == (n_samples, current_chunk_size, self.width)
+                    del cond_prime
+                    x_primes.append(x_prime)
+                else:
+                    del x_prime
+
+            if get_preds:
+                x_prime = t.cat(x_primes, dim=1)
+                assert x_prime.shape == (n_samples, len(xs), self.width)
+                x_prime = self.x_out(x_prime)  # Predictions
+                preds.append(x_prime)
+
+            empty_cache()
+            self.transformer.check_cache(n_samples, len(xs), fp16)
+
+            x = xs[-1]
+            assert x.shape == (n_samples, 1)
+            empty_cache()
+            for sample_t in get_range(range(len(xs), sample_tokens)):
+                x, cond = self.get_emb(sample_t, n_samples, x, x_cond, y_cond)
+                self.transformer.check_cache(n_samples, sample_t, fp16)
+                x = self.transformer(x, encoder_kv=encoder_kv, sample=True, fp16=fp16) # Transformer
+                if self.add_cond_after_transformer:
+                    x = x + cond
+                assert x.shape == (n_samples, 1, self.width)
+                x = self.x_out(x) # Predictions
+                if get_preds:
+                    preds.append(x)
+                # Adjust logits
+                x = x / temp
+                x = filter_logits(x, top_k=top_k, top_p=top_p)
+                x = t.distributions.Categorical(logits=x).sample() # Sample and replace x
+                assert x.shape == (n_samples, 1)
+                xs.append(x.clone())
+
+            del x
+            self.transformer.del_cache()
+
+            x = t.cat(xs, dim=1)
+            if get_preds:
+                preds = t.cat(preds, dim=1)
+            x = self.postprocess(x, sample_tokens)
+        if get_preds:
+            return x, preds
+        else:
+            return x
+
+    def check_sample(self, chunk_size):
+        bs, l, d = (4, self.input_dims, self.width)
+        prime = int(self.input_dims//8*7)
+        enc_l = self.encoder_dims
+        with t.no_grad():
+            y_cond = t.randn(bs, 1, d).cuda() if self.y_cond else None
+            x_cond = t.randn(bs, l, d).cuda() if self.x_cond else None
+            encoder_kv = t.randn(bs, enc_l, d).cuda()
+
+            x, preds_sample = self.sample(bs, x_cond, y_cond, encoder_kv, get_preds=True)
+            loss, preds_forw = self.forward(x, x_cond, y_cond, encoder_kv, get_preds=True)
+            max_err = t.max(t.abs(preds_sample - preds_forw))
+            assert max_err <= 1e-6, f"Max err is {max_err} {[i for i in range(l) if t.max(t.abs(preds_sample - preds_forw)[:, i, :]) > 1e-6]}"
+
+            x_prime = x.view(bs, -1)[:,:prime]
+            # unchunked
+            x, preds_sample = self.primed_sample(bs, x_prime.clone(), x_cond, y_cond, encoder_kv, get_preds=True)
+            assert (x.view(bs, -1)[:,:prime] == x_prime).all(), "Priming samples don't match"
+            loss, preds_forw = self.forward(x, x_cond, y_cond, encoder_kv, get_preds=True)
+            max_err = t.max(t.abs(preds_sample - preds_forw))
+            assert max_err <= 1e-6, f"Max err is {max_err} {[i for i in range(l) if t.max(t.abs(preds_sample - preds_forw)[:, i, :]) > 1e-6]}"
+
+            # chunked
+            x, preds_sample = self.primed_sample(bs, x_prime.clone(), x_cond, y_cond, encoder_kv, get_preds=True, chunk_size=chunk_size)
+            assert (x.view(bs, -1)[:,:prime] == x_prime).all(), "Priming samples don't match"
+            loss, preds_forw = self.forward(x, x_cond, y_cond, encoder_kv, get_preds=True)
+            max_err = t.max(t.abs(preds_sample - preds_forw))
+            assert max_err <= 1e-6, f"Max err is {max_err} {[i for i in range(l) if t.max(t.abs(preds_sample - preds_forw)[:, i, :]) > 1e-6]}"
+
+
+def test_prior(input_shape, encoder_dims, blocks, heads, chunk_size):
+    bins = 512
+    width = 32
+    depth = 2
+    prime_len = encoder_dims
+    for x_cond in [True, False]:
+        for y_cond in [True, False]:
+            for attn_order in [0,2,6,12]:
+                prior = ConditionalAutoregressive2D(input_shape, bins,
+                                                    width=width, depth=depth, heads=heads,
+                                                    attn_order=attn_order, blocks=blocks,
+                                                    x_cond=x_cond, y_cond=y_cond,
+                                                    encoder_dims=encoder_dims, prime_len=prime_len).cuda()
+                prior.training = False
+                prior.check_sample(chunk_size)
+                print(f"Checked x_cond: {x_cond}, y_cond: {y_cond}, attn_order: {attn_order}")
+            # prior.apply(_convert_mlp_traced)
+            # prior.check_sample()
+            # print(f"Checked traced x_cond: {x_cond}, y_cond: {y_cond}")
+
+
+if __name__ == '__main__':
+    from jukebox.utils.dist_utils import setup_dist_from_mpi
+    setup_dist_from_mpi(port=29600)
+    test_cases = [
+        ((6144,), 384, 64, 2, 23),
+        ((6144,), 384, 64, 2, 8),
+        ((8192,), 512, 128, 2, 16),
+    ]
+    for test_case in test_cases:
+        test_prior(*test_case)

jukebox/prior/conditioners.py

@@ -0,0 +1,157 @@
+import torch as t
+import torch.nn as nn
+
+from jukebox.transformer.ops import LayerNorm
+from jukebox.vqvae.encdec import DecoderConvBock
+from jukebox.utils.torch_utils import assert_shape
+
+class Conditioner(nn.Module):
+    def __init__(self, input_shape, bins, down_t, stride_t, out_width, init_scale, zero_out, res_scale, **block_kwargs):
+        super().__init__()
+        self.x_shape = input_shape
+
+        # Embedding
+        self.width = out_width
+        self.x_emb = nn.Embedding(bins, out_width)
+        nn.init.normal_(self.x_emb.weight, std=0.02 * init_scale)
+
+        # Conditioner
+        self.cond = DecoderConvBock(self.width, self.width, down_t, stride_t, **block_kwargs, zero_out=zero_out, res_scale=res_scale)
+        self.ln = LayerNorm(self.width)
+
+    def preprocess(self, x):
+        x = x.permute(0,2,1) # NTC -> NCT
+        return x
+
+    def postprocess(self, x):
+        x = x.permute(0,2,1) # NCT -> NTC
+        return x
+
+    def forward(self, x, x_cond=None):
+        N = x.shape[0]
+        assert_shape(x, (N, *self.x_shape))
+        if x_cond is not None:
+            assert_shape(x_cond, (N, *self.x_shape, self.width))
+        else:
+            x_cond = 0.0
+        # Embed x
+        x = x.long()
+        x = self.x_emb(x)
+        assert_shape(x, (N, *self.x_shape, self.width))
+        x = x + x_cond
+
+        # Run conditioner
+        x = self.preprocess(x)
+        x = self.cond(x)
+        x = self.postprocess(x)
+        x = self.ln(x)
+        return x
+
+def flip(x):
+    def _flip(x):
+        return x.permute(0,2,1).contiguous()
+    if isinstance(x, (list, tuple)):
+        return [flip(z) for z in x]
+    return _flip(x)
+
+class SimpleEmbedding(nn.Module):
+    def __init__(self, bins, out_width, init_scale):
+        super().__init__()
+        self.bins = bins
+        self.emb = nn.Embedding(bins, out_width)
+        nn.init.normal_(self.emb.weight, std=0.01 * init_scale)
+
+    def forward(self, y):
+        assert len(y.shape) == 2, f"Expected shape with 2 dims, got {y.shape}"
+        assert isinstance(y, t.cuda.LongTensor), f"Expected dtype {t.cuda.LongTensor}, got {y.dtype}"
+        assert (0 <= y).all() and (y < self.bins).all(), f"Bins {self.bins}, got label {y}"
+        return self.emb(y)
+
+class RangeEmbedding(nn.Module):
+    # Interpolating
+    # Interpolate so that [pos_start, pos_end] <-> position tensor of length n_ctx
+    #
+    # Binning
+    # For each pos in position tensor, find its bin
+    # [start,end) mapped to [0,1,...,bins-1]
+    # [start,end) -> [0,1) -> [0, bins) -> floor -> [0,...,bins-1]
+    # NOTE: Open ended interval on right, so start <= pos < end, not <= end
+    def __init__(self, n_time, bins, range, out_width, init_scale, clamp=False):
+        super().__init__()
+        self.n_time = n_time
+        self.bins = bins
+        self.emb = nn.Embedding(bins, out_width)
+        nn.init.normal_(self.emb.weight, std=0.01 * init_scale)
+        self.pos_min, self.pos_max = range
+        self.clamp = clamp
+
+    def forward(self, pos_start, pos_end=None):
+        # Check if [pos_start,pos_end] in [pos_min, pos_max)
+        assert len(pos_start.shape) == 2, f"Expected shape with 2 dims, got {pos_start.shape}"
+        assert (self.pos_min <= pos_start).all() and (pos_start < self.pos_max).all(), f"Range is [{self.pos_min},{self.pos_max}), got {pos_start}"
+        pos_start = pos_start.float()
+        if pos_end is not None:
+            assert len(pos_end.shape) == 2, f"Expected shape with 2 dims, got {pos_end.shape}"
+            if self.clamp:
+                pos_end = pos_end.clamp(self.pos_min, self.pos_max)
+            assert (self.pos_min <= pos_end).all() and (pos_end <= self.pos_max).all(), f"Range is [{self.pos_min},{self.pos_max}), got {pos_end}"
+            pos_end = pos_end.float()
+        # Interpolate so that [pos_start, ..., pos_end] <-> position tensor of length n_ctx
+        n_time = self.n_time
+        if n_time != 1:
+            assert pos_end is not None
+            interpolation  = (t.arange(0, n_time, dtype=t.float, device='cuda').view(1,n_time)/n_time)
+            position = pos_start + (pos_end - pos_start)*interpolation
+        else:
+            position = pos_start
+
+        # Bin each value to bins
+        normalised_position = (position - self.pos_min) / (self.pos_max - self.pos_min) # [0,1)
+        bins = (self.bins * normalised_position).floor().long().detach() # [0,1) -> [0,1..,bins) -> [0,1...,bins-1]
+        return self.emb(bins)
+
+class LabelConditioner(nn.Module):
+    def __init__(self, y_bins, t_bins, sr, min_duration, max_duration, n_time, out_width, init_scale, max_bow_genre_size, include_time_signal):
+        super().__init__()
+        self.n_time = n_time
+        self.out_width = out_width
+        assert len(y_bins) == 2, f"Expecting (genre, artist) bins, got {y_bins}"
+        bow_genre_bins, artist_bins = y_bins
+        self.max_bow_genre_size = max_bow_genre_size
+        self.bow_genre_emb = SimpleEmbedding(bow_genre_bins, out_width, init_scale)
+        self.artist_emb = SimpleEmbedding(artist_bins, out_width, init_scale)
+        self.include_time_signal = include_time_signal
+        if self.include_time_signal:
+            t_ranges = ((min_duration * sr, max_duration * sr),  # Total length
+                        (0.0, max_duration * sr),                # Absolute pos
+                        (0.0, 1.0))                              # Relative pos
+            assert len(t_ranges) == 3, f"Expecting (total, absolute, relative) ranges, got {t_ranges}"
+            total_length_range, absolute_pos_range, relative_pos_range = t_ranges
+            self.total_length_emb = RangeEmbedding(1, t_bins, total_length_range, out_width, init_scale)
+            self.absolute_pos_emb = RangeEmbedding(n_time, t_bins, absolute_pos_range, out_width, init_scale)
+            self.relative_pos_emb = RangeEmbedding(n_time, t_bins, relative_pos_range, out_width, init_scale, clamp=True)
+
+    def forward(self, y):
+        assert len(y.shape) == 2, f"Expected shape with 2 dims, got {y.shape}"
+        assert y.shape[-1] == 4 + self.max_bow_genre_size, f"Expected shape (N,{4 + self.max_bow_genre_size}), got {y.shape}"
+        assert isinstance(y, t.cuda.LongTensor), f"Expected dtype {t.cuda.LongTensor}, got {y.dtype}"
+        N = y.shape[0]
+        total_length, offset, length, artist, genre = y[:,0:1], y[:,1:2], y[:,2:3], y[:,3:4], y[:,4:]
+
+        # Start embedding of length 1
+        artist_emb = self.artist_emb(artist)
+        # Empty genre slots are denoted by -1. We mask these out.
+        mask = (genre >= 0).float().unsqueeze(2)
+        genre_emb = (self.bow_genre_emb(genre.clamp(0)) * mask).sum(dim=1, keepdim=True)
+        start_emb = genre_emb + artist_emb
+        assert_shape(start_emb, (N, 1, self.out_width))
+
+        # Pos embedding of length n_ctx
+        if self.include_time_signal:
+            start, end = offset, offset + length
+            total_length, start, end = total_length.float(), start.float(), end.float()
+            pos_emb = self.total_length_emb(total_length) + self.absolute_pos_emb(start, end) + self.relative_pos_emb(start/total_length, end/total_length)
+            assert_shape(pos_emb, (N, self.n_time, self.out_width))
+        else:
+            pos_emb = None
+        return start_emb, pos_emb

jukebox/prior/prior.py

@@ -0,0 +1,354 @@
+import numpy as np
+import torch as t
+import torch.nn as nn
+import jukebox.utils.dist_adapter as dist
+
+from jukebox.transformer.ops import LayerNorm
+from jukebox.prior.autoregressive import ConditionalAutoregressive2D
+from jukebox.prior.conditioners import Conditioner, LabelConditioner
+from jukebox.data.labels import EmptyLabeller, Labeller
+
+from jukebox.utils.torch_utils import assert_shape
+from jukebox.utils.dist_utils import print_once
+from jukebox.vqvae.vqvae import calculate_strides
+
+
+"""
+Model the prior on vq codes conditioned on timing, artist, genre, lyrics and codes from levels above. 
+To condition on the timing, genre and artist, we use the LabelConditioner class
+To condition on the codes from the level above, we use the Conditioner class
+To condition on lyrics, we allow two types of priors:
+- Separate Encoder Decoder: This is the usual encoder-decoder style transformer. The encoder transformer autoregressively 
+models the lyrics, and we use its last layer to produce keys/values that are attened to by the decoder transformer
+- Single Encoder Decoder: This is a simplification where we combine them into a single model. We merge the text vocab 
+and VQ vocab into a single large vocab, and the lyric tokens and VQ tokens into a single longer sequence of tokens which 
+we autoregressively model together.
+"""
+class SimplePrior(nn.Module):
+    def __init__(self, z_shapes, l_bins, encoder, decoder, level,
+                 downs_t, strides_t, labels, prior_kwargs, x_cond_kwargs, y_cond_kwargs,
+                 prime_kwargs, copy_input, labels_v3=False,
+                 merged_decoder=False, single_enc_dec=False):
+        super().__init__()
+
+        self.use_tokens = prime_kwargs.pop('use_tokens')
+        self.n_tokens = prime_kwargs.pop('n_tokens')
+        self.prime_loss_fraction = prime_kwargs.pop('prime_loss_fraction')
+
+        self.copy_input = copy_input
+        if self.copy_input:
+            prime_kwargs['bins'] = l_bins
+
+        self.z_shapes = z_shapes
+        self.levels = len(self.z_shapes)
+
+        self.z_shape = self.z_shapes[level]
+
+        self.level = level
+        assert level < self.levels, f"Total levels {self.levels}, got level {level}"
+
+        self.l_bins = l_bins
+
+        # Passing functions instead of the vqvae module to avoid getting params
+        self.encoder = encoder
+        self.decoder = decoder
+
+        # X conditioning
+        self.x_cond = (level != (self.levels - 1))
+        self.cond_level = level + 1
+
+        # Y conditioning
+        self.y_cond = labels
+
+        self.single_enc_dec = single_enc_dec
+        # X conditioning
+        if self.x_cond:
+            self.conditioner_blocks = nn.ModuleList()
+            conditioner_block = lambda _level: Conditioner(input_shape=z_shapes[_level],
+                                                          bins=l_bins,
+                                                          down_t=downs_t[_level],
+                                                          stride_t=strides_t[_level],
+                                                          **x_cond_kwargs)
+            if dist.get_rank() == 0: print(f"Conditioning on 1 above level(s)")
+            self.conditioner_blocks.append(conditioner_block(self.cond_level))
+
+        # Y conditioning
+        if self.y_cond:
+            self.n_time = self.z_shape[0] # Assuming STFT=TF order and raw=T1 order, so T is first dim
+            self.y_emb = LabelConditioner(n_time=self.n_time,include_time_signal=not self.x_cond,**y_cond_kwargs)
+
+        # Lyric conditioning
+        if single_enc_dec:
+            # Single encoder-decoder transformer
+            self.prior_shapes = [(self.n_tokens,), prior_kwargs.pop('input_shape')]
+            self.prior_bins = [prime_kwargs['bins'], prior_kwargs.pop('bins')]
+            self.prior_dims = [np.prod(shape) for shape in self.prior_shapes]
+            self.prior_bins_shift = np.cumsum([0, *self.prior_bins])[:-1]
+            self.prior_width = prior_kwargs['width']
+            print_once(f'Creating cond. autoregress with prior bins {self.prior_bins}, ')
+            print_once(f'dims {self.prior_dims}, ')
+            print_once(f'shift {self.prior_bins_shift}')
+            print_once(f'input shape {sum(self.prior_dims)}')
+            print_once(f'input bins {sum(self.prior_bins)}')
+            print_once(f'Self copy is {self.copy_input}')
+
+            self.prime_loss_dims, self.gen_loss_dims = self.prior_dims[0], self.prior_dims[1]
+            self.total_loss_dims = self.prime_loss_dims + self.gen_loss_dims
+            self.prior = ConditionalAutoregressive2D(input_shape=(sum(self.prior_dims),),
+                                                     bins=sum(self.prior_bins),
+                                                     x_cond=(self.x_cond or self.y_cond), y_cond=True,
+                                                     prime_len=self.prime_loss_dims,
+                                                     **prior_kwargs)
+
+        else:
+            # Separate encoder-decoder transformer
+            if self.n_tokens != 0 and self.use_tokens:
+                from jukebox.transformer.ops import Conv1D
+                prime_input_shape = (self.n_tokens,)
+                self.prime_loss_dims = np.prod(prime_input_shape)
+                self.prime_acts_width, self.prime_state_width = prime_kwargs['width'], prior_kwargs['width']
+                self.prime_prior = ConditionalAutoregressive2D(input_shape=prime_input_shape, x_cond=False, y_cond=False,
+                                                               only_encode=True,
+                                                               **prime_kwargs)
+                self.prime_state_proj = Conv1D(self.prime_acts_width, self.prime_state_width, init_scale=prime_kwargs['init_scale'])
+                self.prime_state_ln = LayerNorm(self.prime_state_width)
+                self.prime_bins = prime_kwargs['bins']
+                self.prime_x_out = nn.Linear(self.prime_state_width, self.prime_bins, bias=False)
+                nn.init.normal_(self.prime_x_out.weight, std=0.02 * prior_kwargs['init_scale'])
+            else:
+                self.prime_loss_dims = 0
+            self.gen_loss_dims = np.prod(self.z_shape)
+            self.total_loss_dims = self.prime_loss_dims + self.gen_loss_dims
+            self.prior = ConditionalAutoregressive2D(x_cond=(self.x_cond or self.y_cond), y_cond=self.y_cond,
+                                                     encoder_dims = self.prime_loss_dims, merged_decoder=merged_decoder,
+                                                     **prior_kwargs)
+
+        self.n_ctx = self.gen_loss_dims
+        self.downsamples = calculate_strides(strides_t, downs_t)
+        self.cond_downsample = self.downsamples[level+1] if level != self.levels - 1 else None
+        self.raw_to_tokens = np.prod(self.downsamples[:level+1])
+        self.sample_length = self.n_ctx*self.raw_to_tokens
+        if labels:
+            self.labels_v3 = labels_v3
+            self.labeller = Labeller(self.y_emb.max_bow_genre_size, self.n_tokens, self.sample_length, v3=self.labels_v3)
+        else:
+            self.labeller = EmptyLabeller()
+
+        print(f"Level:{level}, Cond downsample:{self.cond_downsample}, Raw to tokens:{self.raw_to_tokens}, Sample length:{self.sample_length}")
+
+
+    def get_y(self, labels, start, get_indices=False):
+        if isinstance(self.labeller, EmptyLabeller):
+            return None
+        y = labels['y'].clone()
+
+        # Set sample_length to match this level
+        y[:, 2] = int(self.sample_length)
+
+        # Set offset
+        y[:, 1:2] = y[:, 1:2] + int(start * self.raw_to_tokens)
+
+        # Set lyric tokens
+        indices = self.labeller.set_y_lyric_tokens(y, labels)
+        if get_indices:
+            return y, indices
+        else:
+            return y
+
+    def get_z_conds(self, zs, start, end):
+        if self.level != self.levels - 1:
+            assert start % self.cond_downsample == end % self.cond_downsample == 0
+            z_cond = zs[self.level + 1][:,start//self.cond_downsample:end//self.cond_downsample]
+            assert z_cond.shape[1] == self.n_ctx//self.cond_downsample
+            z_conds = [z_cond]
+        else:
+            z_conds = None
+        return z_conds
+
+    def prior_preprocess(self, xs, conds):
+        N = xs[0].shape[0]
+        for i in range(len(xs)):
+            x, shape, dims = xs[i], self.prior_shapes[i], self.prior_dims[i]
+            bins, bins_shift = int(self.prior_bins[i]), int(self.prior_bins_shift[i])
+            assert isinstance(x, t.cuda.LongTensor), x
+            assert (0 <= x).all() and (x < bins).all()
+            #assert_shape(x, (N, *shape))
+            xs[i] = (xs[i] + bins_shift).view(N, -1)
+
+        for i in range(len(conds)):
+            cond, shape, dims = conds[i], self.prior_shapes[i], self.prior_dims[i]
+            if cond is not None:
+                assert_shape(cond, (N, dims, self.prior_width))
+            else:
+                conds[i] = t.zeros((N, dims, self.prior_width), dtype=t.float, device='cuda')
+
+        return t.cat(xs, dim=1), t.cat(conds, dim=1)
+
+    def prior_postprocess(self, z):
+        N = z.shape[0]
+        dims = (self.prior_dims[0], z.shape[1] - self.prior_dims[0])
+        # xs = list(t.split(z, self.prior_dims, dim=1))
+        xs = list(t.split(z, dims, dim=1))
+
+        for i in range(len(xs)):
+            # x, shape, dims, bins, bins_shift = xs[i], self.prior_shapes[i], self.prior_dims[i], self.prior_bins[i], self.prior_bins_shift[i]
+            # assert_shape(x, (N, dims))
+            shape = self.prior_shapes[i]
+            bins, bins_shift = int(self.prior_bins[i]), int(self.prior_bins_shift[i])
+            # xs[i] = (xs[i] - bins_shift).view(N, *shape) #view(N, -1, *shape[1:])
+            xs[i] = (xs[i] - bins_shift).view(N, -1, *shape[1:])
+            xs[i] = t.clamp(xs[i], min=0)  # If not masking loss, model may have generated lyric/midi tokens which are now shifted <0 by bin_shift
+            assert (xs[i] < bins).all(), f'rank: {dist.get_rank()}, bins: {bins}, dims {dims}, shape {shape}, prior_shape {self.prior_shapes}, bins_shift {bins_shift}, xs[i]: {xs[i]}'
+
+        return xs[-1]
+
+    def x_emb(self, z_conds):
+        z_conds = z_conds[:self.cond_level - self.level]
+        assert len(z_conds) == len(self.conditioner_blocks) == self.cond_level - self.level, f"Expected {len(z_conds)} == {len(self.conditioner_blocks)} == {self.cond_level} - {self.level}"
+        x_cond = None
+        for z_cond, conditioner_block in reversed(list(zip(z_conds, self.conditioner_blocks))):
+            x_cond = conditioner_block(z_cond, x_cond)
+        return x_cond
+
+    def encode(self, x, start_level=None, end_level=None, bs_chunks=1):
+        if start_level == None:
+            start_level = self.level
+        if end_level == None:
+            end_level = self.levels
+        # Get latents
+        with t.no_grad():
+            zs = self.encoder(x, start_level=start_level, end_level=end_level, bs_chunks=bs_chunks)
+        return zs
+
+    def decode(self, zs, start_level=None, end_level=None, bs_chunks=1):
+        if start_level == None:
+            start_level = self.level
+        if end_level == None:
+            end_level = self.levels
+
+        assert len(zs) == end_level - start_level
+        with t.no_grad():
+            x_out = self.decoder(zs, start_level=start_level, end_level=end_level, bs_chunks=bs_chunks)
+        return x_out
+
+    def get_cond(self, z_conds, y):
+        if y is not None:
+            assert y.shape[1] == 4 + self.y_emb.max_bow_genre_size + self.n_tokens, f"Expected {4} + {self.y_emb.max_bow_genre_size} + {self.n_tokens}, got {y.shape[1]}"
+            n_labels = y.shape[1] - self.n_tokens
+            y, prime = y[:,:n_labels], y[:,n_labels:]
+        else:
+            y, prime = None, None
+        y_cond, y_pos = self.y_emb(y) if self.y_cond else (None, None)
+        x_cond = self.x_emb(z_conds) if self.x_cond else y_pos
+        return x_cond, y_cond, prime
+
+    def sample(self, n_samples, z=None, z_conds=None, y=None, fp16=False, temp=1.0, top_k=0, top_p=0.0,
+               chunk_size=None, sample_tokens=None):
+        N = n_samples
+        if z is not None: assert z.shape[0] == N, f"Expected shape ({N},**), got shape {z.shape}"
+        if y is not None: assert y.shape[0] == N, f"Expected shape ({N},**), got shape {y.shape}"
+        if z_conds is not None:
+            for z_cond in z_conds:
+                assert z_cond.shape[0] == N,  f"Expected shape ({N},**), got shape {z_cond.shape}"
+
+        no_past_context = (z is None or z.shape[1] == 0)
+        if dist.get_rank() == 0:
+            name = {True: 'Ancestral', False: 'Primed'}[no_past_context]
+            print(f"{name} sampling {n_samples} samples with temp={temp}, top_k={top_k}, top_p={top_p}")
+
+        with t.no_grad():
+            # Currently x_cond only uses immediately above layer
+            x_cond, y_cond, prime = self.get_cond(z_conds, y)
+            if self.single_enc_dec:
+                # assert chunk_size % self.prime_loss_dims == 0. TODO: Check if needed
+                if no_past_context:
+                    z, x_cond = self.prior_preprocess([prime], [None, x_cond])
+                else:
+                    z, x_cond = self.prior_preprocess([prime, z], [None, x_cond])
+                if sample_tokens is not None:
+                    sample_tokens += self.n_tokens
+                z = self.prior.primed_sample(n_samples, z, x_cond, y_cond, fp16=fp16, temp=temp,
+                                             top_k=top_k, top_p=top_p, chunk_size=chunk_size, sample_tokens=sample_tokens)
+                z = self.prior_postprocess(z)
+            else:
+                encoder_kv = self.get_encoder_kv(prime, fp16=fp16, sample=True)
+                if no_past_context:
+                    z = self.prior.sample(n_samples, x_cond, y_cond, encoder_kv, fp16=fp16, temp=temp, top_k=top_k,
+                                          top_p=top_p, sample_tokens=sample_tokens)
+                else:
+                    z = self.prior.primed_sample(n_samples, z, x_cond, y_cond, encoder_kv, fp16=fp16, temp=temp,
+                                             top_k=top_k, top_p=top_p, chunk_size=chunk_size, sample_tokens=sample_tokens)
+            if sample_tokens is None:
+                assert_shape(z, (N, *self.z_shape))
+        return z
+
+    def get_encoder_kv(self, prime, fp16=False, sample=False):
+        if self.n_tokens != 0 and self.use_tokens:
+            if sample:
+                self.prime_prior.cuda()
+            N = prime.shape[0]
+            prime_acts = self.prime_prior(prime, None, None, None, fp16=fp16)
+            assert_shape(prime_acts, (N, self.prime_loss_dims, self.prime_acts_width))
+            assert prime_acts.dtype == t.float, f'Expected t.float, got {prime_acts.dtype}'
+            encoder_kv = self.prime_state_ln(self.prime_state_proj(prime_acts))
+            assert encoder_kv.dtype == t.float, f'Expected t.float, got {encoder_kv.dtype}'
+            if sample:
+                self.prime_prior.cpu()
+                if fp16:
+                    encoder_kv = encoder_kv.half()
+        else:
+            encoder_kv = None
+        return encoder_kv
+
+    def get_prime_loss(self, encoder_kv, prime_t):
+        if self.use_tokens:
+            encoder_kv = encoder_kv.float()
+            encoder_kv = self.prime_x_out(encoder_kv)
+            prime_loss = nn.functional.cross_entropy(encoder_kv.view(-1, self.prime_bins), prime_t.view(-1)) / np.log(2.)
+        else:
+            prime_loss = t.tensor(0.0, device='cuda')
+        return prime_loss
+
+    def z_forward(self, z, z_conds=[], y=None, fp16=False, get_preds=False, get_attn_weights=False):
+        """
+        Arguments:
+            get_attn_weights (bool or set): Makes forward prop dump
+                self-attention softmaxes to self.prior.transformer.ws. Either a
+                set of layer indices indicating which layers to store, or a
+                boolean value indicating whether to dump all.
+        """
+        assert isinstance(get_attn_weights, (bool, set))
+        if get_attn_weights:
+            self.prior.transformer.set_record_attn(get_attn_weights)
+        x_cond, y_cond, prime = self.get_cond(z_conds, y)
+        if self.copy_input:
+            prime = z[:,:self.n_tokens]
+        if self.single_enc_dec:
+            z, x_cond = self.prior_preprocess([prime, z], [None, x_cond])
+            (prime_loss, gen_loss), preds = self.prior(z, x_cond, y_cond, fp16=fp16, get_sep_loss=True, get_preds=get_preds)
+        else:
+            encoder_kv = self.get_encoder_kv(prime, fp16=fp16)
+            prime_loss = self.get_prime_loss(encoder_kv, prime)
+            gen_loss, preds = self.prior(z, x_cond, y_cond, encoder_kv, fp16=fp16, get_preds=get_preds)
+        loss = (self.prime_loss_fraction*prime_loss*self.prime_loss_dims/self.total_loss_dims) + \
+                   (gen_loss*self.gen_loss_dims/self.total_loss_dims)
+        metrics=dict(bpd=gen_loss.clone().detach(), prime_loss=prime_loss.clone().detach(),
+                     gen_loss=gen_loss.clone().detach())
+        if get_preds:
+            metrics["preds"] = preds.clone().detach()
+        if get_attn_weights:
+            ws = self.prior.transformer.ws
+            self.prior.transformer.set_record_attn(False)
+            return ws
+        else:
+            return loss, metrics
+
+    def forward(self, x, y=None, fp16=False, decode=False, get_preds=False):
+        bs = x.shape[0]
+        z, *z_conds = self.encode(x, bs_chunks=bs)
+        loss, metrics = self.z_forward(z=z, z_conds=z_conds, y=y, fp16=fp16, get_preds=get_preds)
+        if decode:
+            x_out = self.decode([z, *z_conds])
+        else:
+            x_out = None
+        return x_out, loss, metrics

jukebox/sample.py

@@ -0,0 +1,279 @@
+import os
+import torch as t
+import jukebox.utils.dist_adapter as dist
+
+from jukebox.hparams import Hyperparams
+from jukebox.data.labels import EmptyLabeller
+from jukebox.utils.torch_utils import empty_cache
+from jukebox.utils.audio_utils import save_wav, load_audio
+from jukebox.make_models import make_model
+from jukebox.align import get_alignment
+from jukebox.save_html import save_html
+from jukebox.utils.sample_utils import split_batch, get_starts
+from jukebox.utils.dist_utils import print_once
+import fire
+
+# Sample a partial window of length<n_ctx with tokens_to_sample new tokens on level=level
+def sample_partial_window(zs, labels, sampling_kwargs, level, prior, tokens_to_sample, hps):
+    z = zs[level]
+    n_ctx = prior.n_ctx
+    current_tokens = z.shape[1]
+    if current_tokens < n_ctx - tokens_to_sample:
+        sampling_kwargs['sample_tokens'] = current_tokens + tokens_to_sample
+        start = 0
+    else:
+        sampling_kwargs['sample_tokens'] = n_ctx
+        start = current_tokens - n_ctx + tokens_to_sample
+
+    return sample_single_window(zs, labels, sampling_kwargs, level, prior, start, hps)
+
+# Sample a single window of length=n_ctx at position=start on level=level
+def sample_single_window(zs, labels, sampling_kwargs, level, prior, start, hps):
+    n_samples = hps.n_samples
+    n_ctx = prior.n_ctx
+    end = start + n_ctx
+
+    # get z already sampled at current level
+    z = zs[level][:,start:end]
+
+    if 'sample_tokens' in sampling_kwargs:
+        # Support sampling a window shorter than n_ctx
+        sample_tokens = sampling_kwargs['sample_tokens']
+    else:
+        sample_tokens = (end - start)
+    conditioning_tokens, new_tokens = z.shape[1], sample_tokens - z.shape[1]
+
+    print_once(f"Sampling {sample_tokens} tokens for [{start},{start+sample_tokens}]. Conditioning on {conditioning_tokens} tokens")
+
+    if new_tokens <= 0:
+        # Nothing new to sample
+        return zs
+    
+    # get z_conds from level above
+    z_conds = prior.get_z_conds(zs, start, end)
+
+    # set y offset, sample_length and lyrics tokens
+    y = prior.get_y(labels, start)
+
+    empty_cache()
+
+    max_batch_size = sampling_kwargs['max_batch_size']
+    del sampling_kwargs['max_batch_size']
+
+
+    z_list = split_batch(z, n_samples, max_batch_size)
+    z_conds_list = split_batch(z_conds, n_samples, max_batch_size)
+    y_list = split_batch(y, n_samples, max_batch_size)
+    z_samples = []
+    for z_i, z_conds_i, y_i in zip(z_list, z_conds_list, y_list):
+        z_samples_i = prior.sample(n_samples=z_i.shape[0], z=z_i, z_conds=z_conds_i, y=y_i, **sampling_kwargs)
+        z_samples.append(z_samples_i)
+    z = t.cat(z_samples, dim=0)
+
+    sampling_kwargs['max_batch_size'] = max_batch_size
+
+    # Update z with new sample
+    z_new = z[:,-new_tokens:]
+    zs[level] = t.cat([zs[level], z_new], dim=1)
+    return zs
+
+# Sample total_length tokens at level=level with hop_length=hop_length
+def sample_level(zs, labels, sampling_kwargs, level, prior, total_length, hop_length, hps):
+    print_once(f"Sampling level {level}")
+    if total_length >= prior.n_ctx:
+        for start in get_starts(total_length, prior.n_ctx, hop_length):
+            zs = sample_single_window(zs, labels, sampling_kwargs, level, prior, start, hps)
+    else:
+        zs = sample_partial_window(zs, labels, sampling_kwargs, level, prior, total_length, hps)
+    return zs
+
+# Sample multiple levels
+def _sample(zs, labels, sampling_kwargs, priors, sample_levels, hps):
+    alignments = None
+    for level in reversed(sample_levels):
+        prior = priors[level]
+        prior.cuda()
+        empty_cache()
+
+        # Set correct total_length, hop_length, labels and sampling_kwargs for level
+        assert hps.sample_length % prior.raw_to_tokens == 0, f"Expected sample_length {hps.sample_length} to be multiple of {prior.raw_to_tokens}"
+        total_length = hps.sample_length//prior.raw_to_tokens
+        hop_length = int(hps.hop_fraction[level]*prior.n_ctx)
+        zs = sample_level(zs, labels[level], sampling_kwargs[level], level, prior, total_length, hop_length, hps)
+
+        prior.cpu()
+        empty_cache()
+
+        # Decode sample
+        x = prior.decode(zs[level:], start_level=level, bs_chunks=zs[level].shape[0])
+
+        if dist.get_world_size() > 1:
+            logdir = f"{hps.name}_rank_{dist.get_rank()}/level_{level}"
+        else:
+            logdir = f"{hps.name}/level_{level}"
+        if not os.path.exists(logdir):
+            os.makedirs(logdir)
+        t.save(dict(zs=zs, labels=labels, sampling_kwargs=sampling_kwargs, x=x), f"{logdir}/data.pth.tar")
+        save_wav(logdir, x, hps.sr)
+        if alignments is None and priors[-1] is not None and priors[-1].n_tokens > 0 and not isinstance(priors[-1].labeller, EmptyLabeller):
+            alignments = get_alignment(x, zs, labels[-1], priors[-1], sampling_kwargs[-1]['fp16'], hps)
+        save_html(logdir, x, zs, labels[-1], alignments, hps)
+    return zs
+
+# Generate ancestral samples given a list of artists and genres
+def ancestral_sample(labels, sampling_kwargs, priors, hps):
+    sample_levels = list(range(len(priors)))
+    zs = [t.zeros(hps.n_samples,0,dtype=t.long, device='cuda') for _ in range(len(priors))]
+    zs = _sample(zs, labels, sampling_kwargs, priors, sample_levels, hps)
+    return zs
+
+# Continue ancestral sampling from previously saved codes
+def continue_sample(zs, labels, sampling_kwargs, priors, hps):
+    sample_levels = list(range(len(priors)))
+    zs = _sample(zs, labels, sampling_kwargs, priors, sample_levels, hps)
+    return zs
+
+# Upsample given already generated upper-level codes
+def upsample(zs, labels, sampling_kwargs, priors, hps):
+    sample_levels = list(range(len(priors) - 1))
+    zs = _sample(zs, labels, sampling_kwargs, priors, sample_levels, hps)
+    return zs
+
+# Prompt the model with raw audio input (dimension: NTC) and generate continuations
+def primed_sample(x, labels, sampling_kwargs, priors, hps):
+    sample_levels = list(range(len(priors)))
+    zs = priors[-1].encode(x, start_level=0, end_level=len(priors), bs_chunks=x.shape[0])
+    zs = _sample(zs, labels, sampling_kwargs, priors, sample_levels, hps)
+    return zs
+
+# Load `duration` seconds of the given audio files to use as prompts
+def load_prompts(audio_files, duration, hps):
+    xs = []
+    for audio_file in audio_files:
+        x = load_audio(audio_file, sr=hps.sr, duration=duration, offset=0.0, mono=True)
+        x = x.T # CT -> TC
+        xs.append(x)
+    while len(xs) < hps.n_samples:
+        xs.extend(xs)
+    xs = xs[:hps.n_samples]
+    x = t.stack([t.from_numpy(x) for x in xs])
+    x = x.to('cuda', non_blocking=True)
+    return x
+
+# Load codes from previous sampling run
+def load_codes(codes_file, duration, priors, hps):
+    data = t.load(codes_file, map_location='cpu')
+    zs = [z.cuda() for z in data['zs']]
+    assert zs[-1].shape[0] == hps.n_samples, f"Expected bs = {hps.n_samples}, got {zs[-1].shape[0]}"
+    del data
+    if duration is not None:
+        # Cut off codes to match duration
+        top_raw_to_tokens = priors[-1].raw_to_tokens
+        assert duration % top_raw_to_tokens == 0, f"Cut-off duration {duration} not an exact multiple of top_raw_to_tokens"
+        assert duration//top_raw_to_tokens <= zs[-1].shape[1], f"Cut-off tokens {duration//priors[-1].raw_to_tokens} longer than tokens {zs[-1].shape[1]} in saved codes"
+        zs = [z[:,:duration//prior.raw_to_tokens] for z, prior in zip(zs, priors)]
+    return zs
+
+# Generate and save samples, alignment, and webpage for visualization.
+def save_samples(model, device, hps, sample_hps):
+    print(hps)
+    from jukebox.lyricdict import poems, gpt_2_lyrics
+    vqvae, priors = make_model(model, device, hps)
+
+    assert hps.sample_length//priors[-2].raw_to_tokens >= priors[-2].n_ctx, f"Upsampling needs atleast one ctx in get_z_conds. Please choose a longer sample length"
+
+    total_length = hps.total_sample_length_in_seconds * hps.sr
+    offset = 0
+
+    # Set artist/genre/lyrics for your samples here!
+    # We used different label sets in our models, but you can write the human friendly names here and we'll map them under the hood for each model.
+    # For the 5b/5b_lyrics model and the upsamplers, labeller will look up artist and genres in v2 set. (after lowercasing, removing non-alphanumerics and collapsing whitespaces to _).
+    # For the 1b_lyrics top level, labeller will look up artist and genres in v3 set (after lowercasing).
+    metas = [dict(artist = "Alan Jackson",
+                  genre = "Country",
+                  lyrics = poems['ozymandias'],
+                  total_length=total_length,
+                  offset=offset,
+                  ),
+             dict(artist="Joe Bonamassa",
+                  genre="Blues Rock",
+                  lyrics=gpt_2_lyrics['hottub'],
+                  total_length=total_length,
+                  offset=offset,
+                  ),
+             dict(artist="Frank Sinatra",
+                  genre="Classic Pop",
+                  lyrics=gpt_2_lyrics['alone'],
+                  total_length=total_length,
+                  offset=offset,
+                  ),
+             dict(artist="Ella Fitzgerald",
+                  genre="Jazz",
+                  lyrics=gpt_2_lyrics['count'],
+                  total_length=total_length,
+                  offset=offset,
+                  ),
+             dict(artist="Céline Dion",
+                  genre="Pop",
+                  lyrics=gpt_2_lyrics['darkness'],
+                  total_length=total_length,
+                  offset=offset,
+                  ),
+             ]
+    while len(metas) < hps.n_samples:
+        metas.extend(metas)
+    metas = metas[:hps.n_samples]
+
+    labels = [prior.labeller.get_batch_labels(metas, 'cuda') for prior in priors]
+    for label in labels:
+        assert label['y'].shape[0] == hps.n_samples
+
+    lower_level_chunk_size = 32
+    lower_level_max_batch_size = 16
+    if model == '1b_lyrics':
+        chunk_size = 32
+        max_batch_size = 16
+    else:
+        chunk_size = 16
+        max_batch_size = 3
+    sampling_kwargs = [dict(temp=0.99, fp16=True, chunk_size=lower_level_chunk_size, max_batch_size=lower_level_max_batch_size),
+                       dict(temp=0.99, fp16=True, chunk_size=lower_level_chunk_size, max_batch_size=lower_level_max_batch_size),
+                       dict(temp=0.99, fp16=True, chunk_size=chunk_size, max_batch_size=max_batch_size)]
+
+    if sample_hps.mode == 'ancestral':
+        ancestral_sample(labels, sampling_kwargs, priors, hps)
+    elif sample_hps.mode in ['continue', 'upsample']:
+        assert sample_hps.codes_file is not None
+        top_raw_to_tokens = priors[-1].raw_to_tokens
+        if sample_hps.prompt_length_in_seconds is not None:
+            duration = (int(sample_hps.prompt_length_in_seconds * hps.sr) // top_raw_to_tokens) * top_raw_to_tokens
+        else:
+            duration = None
+        zs = load_codes(sample_hps.codes_file, duration, priors, hps)
+        if sample_hps.mode == 'continue':
+            continue_sample(zs, labels, sampling_kwargs, priors, hps)
+        elif sample_hps.mode == 'upsample':
+            upsample(zs, labels, sampling_kwargs, priors, hps)
+    elif sample_hps.mode == 'primed':
+        assert sample_hps.audio_file is not None
+        assert sample_hps.prompt_length_in_seconds is not None
+        audio_files = sample_hps.audio_file.split(',')
+        top_raw_to_tokens = priors[-1].raw_to_tokens
+        duration = (int(sample_hps.prompt_length_in_seconds * hps.sr) // top_raw_to_tokens) * top_raw_to_tokens
+        x = load_prompts(audio_files, duration, hps)
+        primed_sample(x, labels, sampling_kwargs, priors, hps)
+    else:
+        raise ValueError(f'Unknown sample mode {sample_hps.mode}.')
+
+
+def run(model, mode='ancestral', codes_file=None, audio_file=None, prompt_length_in_seconds=None, port=29500, **kwargs):
+    from jukebox.utils.dist_utils import setup_dist_from_mpi
+    rank, local_rank, device = setup_dist_from_mpi(port=port)
+    hps = Hyperparams(**kwargs)
+    sample_hps = Hyperparams(dict(mode=mode, codes_file=codes_file, audio_file=audio_file, prompt_length_in_seconds=prompt_length_in_seconds))
+
+    with t.no_grad():
+        save_samples(model, device, hps, sample_hps)
+
+if __name__ == '__main__':
+    fire.Fire(run)

jukebox/save_html.py

@@ -0,0 +1,130 @@
+import os
+import json
+import numpy as np
+from PIL import Image, ImageFilter
+import soundfile
+
+def save_html(logdir, x, zs, labels, alignments, hps):
+    level = hps.levels - 1 # Top level used
+    z = zs[level]
+    bs, total_length = z.shape[0], z.shape[1]
+
+    with open(f'{logdir}/index.html', 'w') as html:
+        print(f"<html><head><title>{logdir}</title></head><body style='font-family: sans-serif; font-size: 1.4em; font-weight: bold; text-align: center; max-width:1024px; width: 100%; margin: auto;'>",
+            file=html)
+        print("<link rel='icon' href='data:;base64,iVBORw0KGgo='>", file=html)
+
+        for item in range(bs):
+            data = dict(wav=x[item].cpu().numpy(), sr=hps.sr,
+                        info=labels['info'][item],
+                        total_length=total_length,
+                        total_tokens=len(labels['info'][item]['full_tokens']),
+                        alignment=alignments[item] if alignments is not None else None)
+            item_dir = f'{logdir}/item_{item}'
+            _save_item_html(item_dir, item, item, data)
+            print(f"<iframe style='height: 100%; width: 100%;' frameborder='0' scrolling='no' src='item_{item}/index.html'></iframe>", file=html)
+        print("</body></html>", file=html)  
+
+def _save_item_html(item_dir, item_id, item_name, data):
+    # replace gs:// with /root/samples/
+
+    # an html for each sample. Main html has a selector to get us id of this?
+    if not os.path.exists(item_dir):
+        os.makedirs(item_dir)
+
+    with open(f'{item_dir}/index.html', 'w') as html:
+        print(f"<html><head><title>{item_name}</title></head><body style='font-family: sans-serif; font-size: 1.4em; font-weight: bold; text-align: center; max-width:1024px; width: 100%; margin: auto;'>",
+            file=html)
+        print("<link rel='icon' href='data:;base64,iVBORw0KGgo='>", file=html)
+        total_length = data['total_length']
+        total_tokens = data['total_tokens']
+        alignment = data['alignment']
+        lyrics = data["info"]["lyrics"]
+        wav, sr = data['wav'], data['sr']
+        genre, artist = data["info"]["genre"], data["info"]["artist"]
+
+        # Strip unused columns
+        if alignment is not None:
+            assert alignment.shape == (total_length, total_tokens)
+            assert len(lyrics) == total_tokens, f'Total_tokens: {total_tokens}, Lyrics Len: {len(lyrics)}. Lyrics: {lyrics}'
+            max_attn_at_token = np.max(alignment, axis=0)
+            assert len(max_attn_at_token) == total_tokens
+            for token in reversed(range(total_tokens)):
+                if max_attn_at_token[token] > 0:
+                    break
+            alignment = alignment[:,:token+1]
+            lyrics = lyrics[:token+1]
+            total_tokens = token+1
+
+            # Small alignment image
+            im = Image.fromarray(np.uint8(alignment * 255)).resize((512, 1024)).transpose(Image.ROTATE_90)
+            img_src = f'align.png'
+            im.save(f'{item_dir}/{img_src}')
+            print(f"<img id='{img_src}' src='{img_src}' \>", file=html)
+
+            # Smaller alignment json for animation
+            total_alignment_length = total_length // 16
+            alignment = Image.fromarray(np.uint8(alignment * 255)).resize((total_tokens, total_alignment_length))
+            alignment = alignment.filter(ImageFilter.GaussianBlur(radius=1.5))
+            alignment = np.asarray(alignment).tolist()
+            align_src = f'align.json'
+            with open(f'{item_dir}/{align_src}', 'w') as f:
+                json.dump(alignment, f)
+
+        # Audio
+        wav_src = f'audio.wav'
+        soundfile.write(f'{item_dir}/{wav_src}', wav, samplerate=sr, format='wav')
+        print(f"<audio id='{wav_src}' src='{wav_src}' style='width: 100%;' controls></audio>", file=html)
+
+
+        # Labels and Lyrics
+        print(f"<pre style='white-space: pre-wrap;'>", end="", file=html)
+        print(f"<div>Artist {artist}, Genre {genre}</div>", file=html)
+        lyrics = [c for c in lyrics]  # already characters actually
+        lyrics = [''] + lyrics[:-1]  # input lyrics are shifted by 1
+        for i, c in enumerate(lyrics):
+            print(f"<span id='{item_id}/{i}'>{c}</span>", end="", file=html)
+        print(f"</pre>", file=html)
+        with open(f'{item_dir}/lyrics.json', 'w') as f:
+            json.dump(lyrics, f)
+
+        if alignment is not None:
+            # JS for alignment animation
+            print("""<script>
+            async function fetchAsync (url) {
+                let response = await fetch(url);
+                let data = await response.json();
+                return data;
+            }
+    
+            var audio = document.getElementById('""" + f'{wav_src}' + """');
+            audio.onplay = function () {
+                track = '""" + f'{item_id}' + """'
+                fetchAsync('""" + f'{align_src}' + """')
+                .then(data => animateLyrics(data, track, this))
+                .catch(reason => console.log(reason.message))
+            }; 
+    
+            function animateLyrics(data, track, audio) {
+                var animate = setInterval(function () {
+                    var time = Math.floor(audio.currentTime*""" + f'{total_alignment_length}' + """/audio.duration);
+                    if (!(time == 0 || time == """ + f'{total_alignment_length}' + """)) {
+                        console.log(time);
+                        changeColor(data, track, audio, time);
+                    }
+                    if (audio.paused) {
+                        clearInterval(animate);
+                    }
+                }, 50);
+            }
+    
+            function changeColor(data, track, audio, time) {
+                colors = data[time]
+                for (i = 0; i < colors.length; i++){
+                    character = document.getElementById(track + '/' + i.toString());
+                    color = Math.max(230 - 10*colors[i], 0).toString();
+                    character.style.color = 'rgb(255,' + color + ',' + color + ')';
+                }
+            }
+            </script>""", file=html)
+        print("</body></html>", file=html)

jukebox/train.py

@@ -0,0 +1,345 @@
+"""
+Ability to train vq-vae and prior
+First try for random inputs
+Then from maestros
+"""
+import sys
+import fire
+import warnings
+import numpy as np
+import torch as t
+import jukebox.utils.dist_adapter as dist
+from torch.nn.parallel import DistributedDataParallel
+
+from jukebox.hparams import setup_hparams
+from jukebox.make_models import make_vqvae, make_prior, restore_opt, save_checkpoint
+from jukebox.utils.logger import init_logging
+from jukebox.utils.audio_utils import audio_preprocess, audio_postprocess
+from jukebox.utils.torch_utils import zero_grad, count_parameters
+from jukebox.utils.dist_utils import print_once, allreduce, allgather
+from jukebox.utils.ema import CPUEMA, FusedEMA, EMA
+from jukebox.utils.fp16 import FP16FusedAdam, FusedAdam, LossScalar, clipped_grad_scale, backward
+from jukebox.data.data_processor import DataProcessor
+
+def prepare_aud(x, hps):
+    x = audio_postprocess(x.detach().contiguous(), hps)
+    return allgather(x)
+
+def log_aud(logger, tag, x, hps):
+    logger.add_audios(tag, prepare_aud(x, hps), hps.sr, max_len=hps.max_len, max_log=hps.max_log)
+    logger.flush()
+
+def log_labels(logger, labeller, tag, y, hps):
+    y = y.cpu().numpy()
+    txt = ''
+    for item in range(y.shape[0]):
+        description = labeller.describe_label(y[item])
+        artist, genre, lyrics = description['artist'], description['genre'], description['lyrics']
+        txt += f'{item} artist:{artist}, genre:{genre}, lyrics:{lyrics}\n'
+    logger.add_text(tag, txt)
+    logger.flush()
+
+def get_ddp(model, hps):
+    rank = dist.get_rank()
+    local_rank = rank % 8
+    ddp = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False, bucket_cap_mb=hps.bucket)
+    return ddp
+
+def get_ema(model, hps):
+    mu = hps.mu or (1. - (hps.bs * hps.ngpus/8.)/1000)
+    ema = None
+    if hps.ema and hps.train:
+        if hps.cpu_ema:
+            if dist.get_rank() == 0:
+                print("Using CPU EMA")
+            ema = CPUEMA(model.parameters(), mu=mu, freq=hps.cpu_ema_freq)
+        elif hps.ema_fused:
+            ema = FusedEMA(model.parameters(), mu=mu)
+        else:
+            ema = EMA(model.parameters(), mu=mu)
+    return ema
+
+def get_lr_scheduler(opt, hps):
+    def lr_lambda(step):
+        if hps.lr_use_linear_decay:
+            lr_scale = hps.lr_scale * min(1.0, step / hps.lr_warmup)
+            decay = max(0.0, 1.0 - max(0.0, step - hps.lr_start_linear_decay) / hps.lr_decay)
+            if decay == 0.0:
+                if dist.get_rank() == 0:
+                    print("Reached end of training")
+            return lr_scale * decay
+        else:
+            return hps.lr_scale * (hps.lr_gamma ** (step // hps.lr_decay)) * min(1.0, step / hps.lr_warmup)
+
+    shd = t.optim.lr_scheduler.LambdaLR(opt, lr_lambda)
+
+    return shd
+
+def get_optimizer(model, hps):
+    # Optimizer
+    betas = (hps.beta1, hps.beta2)
+    if hps.fp16_opt:
+        opt = FP16FusedAdam(model.parameters(), lr=hps.lr, weight_decay=hps.weight_decay, betas=betas, eps=hps.eps)
+    else:
+        opt = FusedAdam(model.parameters(), lr=hps.lr, weight_decay=hps.weight_decay, betas=betas, eps=hps.eps)
+
+    # lr scheduler
+    shd = get_lr_scheduler(opt, hps)
+
+    restore_path = hps.restore_prior if hps.prior else hps.restore_vqvae
+    restore_opt(opt, shd, restore_path)
+
+    # fp16 dynamic loss scaler
+    scalar = None
+    if hps.fp16:
+        rank = dist.get_rank()
+        local_rank = rank % 8
+        scalar = LossScalar(hps.fp16_loss_scale, scale_factor=2 ** (1./hps.fp16_scale_window))
+        if local_rank == 0: print(scalar.__dict__)
+
+    zero_grad(model)
+    return opt, shd, scalar
+
+def log_inputs(orig_model, logger, x_in, y, x_out, hps, tag="train"):
+    print(f"Logging {tag} inputs/ouputs")
+    log_aud(logger, f'{tag}_x_in', x_in, hps)
+    log_aud(logger, f'{tag}_x_out', x_out, hps)
+    bs = x_in.shape[0]
+    if hps.prior:
+        if hps.labels:
+            log_labels(logger, orig_model.labeller, f'{tag}_y_in', allgather(y.cuda()), hps)
+    else:
+        zs_in = orig_model.encode(x_in, start_level=0, bs_chunks=bs)
+        x_ds = [orig_model.decode(zs_in[level:], start_level=level, bs_chunks=bs) for level in range(0, hps.levels)]
+        for i in range(len(x_ds)):
+            log_aud(logger, f'{tag}_x_ds_start_{i}', x_ds[i], hps)
+    logger.flush()
+
+def sample_prior(orig_model, ema, logger, x_in, y, hps):
+    if ema is not None: ema.swap()
+    orig_model.eval()
+
+    x_in = x_in[:hps.bs_sample]
+    bs = x_in.shape[0]
+    zs_in = orig_model.encode(x_in, start_level=0, bs_chunks=bs)
+    assert len(zs_in) == hps.levels
+    x_ds = [orig_model.decode(zs_in[level:], start_level=level, bs_chunks=bs) for level in range(0, hps.levels)]
+
+    if not hps.labels:
+        y = None
+    elif hps.level == (hps.levels - 1):
+        # Topmost level labels in order
+        y = y[:hps.bs_sample]  # t.ones((hps.bs_sample, 1), device=y.device, dtype=t.long) * dist.get_rank()
+    else:
+        # Other levels keep labels to match x_cond
+        y = y[:hps.bs_sample]
+
+    # Temp 1.0
+    _, *z_conds = orig_model.encode(x_in, bs_chunks=bs)
+    z = orig_model.sample(hps.bs_sample, z_conds=z_conds, y=y, fp16=False, temp=1.0)
+    x_sample = orig_model.decode([z, *z_conds], bs_chunks=bs)
+
+    log_aud(logger, 'sample_x_T1', x_sample, hps)
+    if hps.prior and hps.labels:
+        log_labels(logger, orig_model.labeller, f'sample_x_T1', allgather(y.cuda()), hps)
+
+    # Recons
+    for i in range(len(x_ds)):
+        log_aud(logger, f'x_ds_start_{i}', x_ds[i], hps)
+    orig_model.train()
+    if ema is not None: ema.swap()
+    logger.flush()
+
+def evaluate(model, orig_model, logger, metrics, data_processor, hps):
+    model.eval()
+    orig_model.eval()
+    if hps.prior:
+        _print_keys = dict(l="loss", bpd="bpd")
+    else:
+        _print_keys = dict(l="loss", rl="recons_loss", sl="spectral_loss")
+
+    with t.no_grad():
+        for i, x in logger.get_range(data_processor.test_loader):
+            if isinstance(x, (tuple, list)):
+                x, y = x
+            else:
+                y = None
+
+            x = x.to('cuda', non_blocking=True)
+            if y is not None:
+                y = y.to('cuda', non_blocking=True)
+
+            x_in = x = audio_preprocess(x, hps)
+            log_input_output = (i==0)
+
+            if hps.prior:
+                forw_kwargs = dict(y=y, fp16=hps.fp16, decode=log_input_output)
+            else:
+                forw_kwargs = dict(loss_fn=hps.loss_fn, hps=hps)
+
+            x_out, loss, _metrics = model(x, **forw_kwargs)
+
+            # Logging
+            for key, val in _metrics.items():
+                _metrics[key] = val.item()
+            _metrics["loss"] = loss = loss.item() # Make sure to call to free graph
+
+            # Average and log
+            for key, val in _metrics.items():
+                _metrics[key] = metrics.update(f"test_{key}", val, x.shape[0])
+
+            with t.no_grad():
+                if log_input_output:
+                    log_inputs(orig_model, logger, x_in, y, x_out, hps)
+
+            logger.set_postfix(**{print_key:_metrics[key] for print_key, key in _print_keys.items()})
+
+    for key, val in _metrics.items():
+        logger.add_scalar(f"test_{key}", metrics.avg(f"test_{key}"))
+
+    logger.close_range()
+    return {key: metrics.avg(f"test_{key}") for key in _metrics.keys()}
+
+def train(model, orig_model, opt, shd, scalar, ema, logger, metrics, data_processor, hps):
+    model.train()
+    orig_model.train()
+    if hps.prior:
+        _print_keys = dict(l="loss", bpd="bpd", gn="gn", g_l="gen_loss", p_l="prime_loss")
+    else:
+        _print_keys = dict(l="loss", sl="spectral_loss", rl="recons_loss", e="entropy", u="usage", uc="used_curr", gn="gn", pn="pn", dk="dk")
+
+    for i, x in logger.get_range(data_processor.train_loader):
+        if isinstance(x, (tuple, list)):
+            x, y = x
+        else:
+            y = None
+
+        x = x.to('cuda', non_blocking=True)
+        if y is not None:
+            y = y.to('cuda', non_blocking=True)
+
+        x_in = x = audio_preprocess(x, hps)
+        log_input_output = (logger.iters % hps.save_iters == 0)
+
+        if hps.prior:
+            forw_kwargs = dict(y=y, fp16=hps.fp16, decode=log_input_output)
+        else:
+            forw_kwargs = dict(loss_fn=hps.loss_fn, hps=hps)
+
+        # Forward
+        x_out, loss, _metrics = model(x, **forw_kwargs)
+
+        # Backward
+        loss, scale, grad_norm, overflow_loss, overflow_grad = backward(loss=loss, params=list(model.parameters()),
+                                                                         scalar=scalar, fp16=hps.fp16, logger=logger)
+        # Skip step if overflow
+        grad_norm = allreduce(grad_norm, op=dist.ReduceOp.MAX)
+        if overflow_loss or overflow_grad or grad_norm > hps.ignore_grad_norm > 0:
+            zero_grad(orig_model)
+            continue
+
+        # Step opt. Divide by scale to include clipping and fp16 scaling
+        logger.step()
+        opt.step(scale=clipped_grad_scale(grad_norm, hps.clip, scale))
+        zero_grad(orig_model)
+        lr = hps.lr if shd is None else shd.get_lr()[0]
+        if shd is not None: shd.step()
+        if ema is not None: ema.step()
+        next_lr = hps.lr if shd is None else shd.get_lr()[0]
+        finished_training = (next_lr == 0.0)
+
+        # Logging
+        for key, val in _metrics.items():
+            _metrics[key] = val.item()
+        _metrics["loss"] = loss = loss.item() * hps.iters_before_update # Make sure to call to free graph
+        _metrics["gn"] = grad_norm
+        _metrics["lr"] = lr
+        _metrics["lg_loss_scale"] = np.log2(scale)
+
+        # Average and log
+        for key, val in _metrics.items():
+            _metrics[key] = metrics.update(key, val, x.shape[0])
+            if logger.iters % hps.log_steps == 0:
+                logger.add_scalar(key, _metrics[key])
+
+        # Save checkpoint
+        with t.no_grad():
+            if hps.save and (logger.iters % hps.save_iters == 1 or finished_training):
+                if ema is not None: ema.swap()
+                orig_model.eval()
+                name = 'latest' if hps.prior else f'step_{logger.iters}'
+                if dist.get_rank() % 8 == 0:
+                    save_checkpoint(logger, name, orig_model, opt, dict(step=logger.iters), hps)
+                orig_model.train()
+                if ema is not None: ema.swap()
+
+        # Sample
+        with t.no_grad():
+            if (logger.iters % 12000) in list(range(1, 1 + hps.iters_before_update)) or finished_training:
+                if hps.prior:
+                    sample_prior(orig_model, ema, logger, x_in, y, hps)
+
+        # Input/Output
+        with t.no_grad():
+            if log_input_output:
+                log_inputs(orig_model, logger, x_in, y, x_out, hps)
+
+        logger.set_postfix(**{print_key:_metrics[key] for print_key, key in _print_keys.items()})
+        if finished_training:
+            dist.barrier()
+            exit()
+    logger.close_range()
+    return {key: metrics.avg(key) for key in _metrics.keys()}
+
+def run(hps="teeny", port=29500, **kwargs):
+    from jukebox.utils.dist_utils import setup_dist_from_mpi
+    rank, local_rank, device = setup_dist_from_mpi(port=port)
+    hps = setup_hparams(hps, kwargs)
+    hps.ngpus = dist.get_world_size()
+    hps.argv = " ".join(sys.argv)
+    hps.bs_sample = hps.nworkers = hps.bs
+
+    # Setup dataset
+    data_processor = DataProcessor(hps)
+
+    # Setup models
+    vqvae = make_vqvae(hps, device)
+    print_once(f"Parameters VQVAE:{count_parameters(vqvae)}")
+    if hps.prior:
+        prior = make_prior(hps, vqvae, device)
+        print_once(f"Parameters Prior:{count_parameters(prior)}")
+        model = prior
+    else:
+        model = vqvae
+
+    # Setup opt, ema and distributed_model.
+    opt, shd, scalar = get_optimizer(model, hps)
+    ema = get_ema(model, hps)
+    distributed_model = get_ddp(model, hps)
+
+    logger, metrics = init_logging(hps, local_rank, rank)
+    logger.iters = model.step
+
+    # Run training, eval, sample
+    for epoch in range(hps.curr_epoch, hps.epochs):
+        metrics.reset()
+        data_processor.set_epoch(epoch)
+        if hps.train:
+            train_metrics = train(distributed_model, model, opt, shd, scalar, ema, logger, metrics, data_processor, hps)
+            train_metrics['epoch'] = epoch
+            if rank == 0:
+                print('Train',' '.join([f'{key}: {val:0.4f}' for key,val in train_metrics.items()]))
+            dist.barrier()
+
+        if hps.test:
+            if ema: ema.swap()
+            test_metrics = evaluate(distributed_model, model, logger, metrics, data_processor, hps)
+            test_metrics['epoch'] = epoch
+            if rank == 0:
+                print('Ema',' '.join([f'{key}: {val:0.4f}' for key,val in test_metrics.items()]))
+            dist.barrier()
+            if ema: ema.swap()
+        dist.barrier()
+
+if __name__ == '__main__':
+    fire.Fire(run)

jukebox/transformer/factored_attention.py

@@ -0,0 +1,510 @@
+# Factored attention
+import math
+import numpy as np
+import torch as t
+import torch.nn as nn
+import torch.nn.functional as F
+from jukebox.transformer.ops import Conv1D
+from jukebox.utils.checkpoint import checkpoint
+
+def repeat(x, n, dim):
+    if dim == -1:
+        dim = len(x.shape) - 1
+    return x.view(int(np.prod(x.shape[:dim+1])), 1, int(np.prod(x.shape[dim+1:]))).repeat(1,n,1).view(*x.shape[:dim], n * x.shape[dim], *x.shape[dim+1:])
+
+def get_mask(mask, q_l, kv_l, blocks, spread, device, sample, sample_t):
+    # returns a mask of shape 1 x 1 x q_l x kv_l or None if masking is not needed.
+    if mask is None or q_l == 1:
+        return None
+    offset = sample_t - q_l if sample else max(kv_l - q_l, 0)
+    if mask == 'autoregressive':
+        # Masked dense
+        mask = t.ones(q_l, kv_l, device=device).tril(offset)
+    elif mask == 'summary':
+        # Masked summary
+        mask = t.nn.functional.pad(t.ones(q_l, q_l, device=device).tril().view(q_l, blocks, q_l // blocks)[:,:-1,-kv_l//blocks:],(0,0,1,0),value=1).contiguous().view(q_l, kv_l)
+    elif mask == 'prime':
+        mask = t.ones(q_l, kv_l, device=device).tril(offset)
+    return mask.view(1,1,q_l,kv_l)
+
+class FactoredAttention(nn.Module):
+    def __init__(self, n_in, n_ctx, n_state, n_head,
+                 attn_dropout=0.0, resid_dropout=0.0,
+                 scale=True, mask=False,
+                 zero_out=False, init_scale=1.0,
+                 checkpoint_attn=0,
+                 attn_func=0, blocks=None, spread=None,
+                 encoder_dims=None, prime_len=None):
+        super().__init__()
+        self.n_in = n_in
+        self.n_ctx = n_ctx # NOTE: n_ctx could be different within operations. This is complete n_ctx
+        self.n_state = n_state
+        assert n_state % n_head == 0
+        self.n_head = n_head
+        self.scale = scale
+        self.mask = mask
+        if attn_func == 6:
+            self.c_attn = Conv1D(n_in, n_state, init_scale=init_scale)
+            self.c_enc_kv = Conv1D(n_in, n_state * 2, init_scale=init_scale)
+        else:
+            self.c_attn = Conv1D(n_in, n_state * 3, init_scale=init_scale)
+        self.c_proj = Conv1D(n_state, n_in, zero_out, init_scale=init_scale)
+        self.attn_dropout = nn.Dropout(attn_dropout) if attn_dropout > 0.0 else lambda x: x
+        self.resid_dropout = nn.Dropout(resid_dropout) if resid_dropout > 0.0 else lambda x: x
+
+        # Sequence of length l is factored as [blocks, l // blocks]
+        self.attn_func = attn_func
+        self.qkv, self.attn, self.attn_mask = {
+            0: (self.factored_qkv, self.dense_attn, 'autoregressive'),              # Attend to all positions
+            1: (self.factored_qkv, self.block_attn, 'autoregressive'),              # Attend to your block
+            2: (self.factored_qkv, self.transpose_block_attn, 'autoregressive'),    # Attend to transpose block
+            3: (self.factored_qkv, self.prev_block_attn, None),                     # Attend to previous block
+            4: (self.factored_qkv, self.summary_attn, 'summary'),                   # Attend to last position of each block
+            5: (self.factored_qkv, self.summary_spread_attn, 'summary'),
+            6: (self.decode_qkv, self.decode_attn, None),
+            7: (self.prime_qkv, self.prime_attn, 'prime')
+        }[attn_func] # Attend to last k position of each block
+
+        self.blocks = blocks
+        self.spread = spread
+        if blocks is not None:
+            assert n_ctx % blocks == 0
+            self.block_ctx = n_ctx // blocks
+        self.checkpoint_attn = checkpoint_attn # 0: None, 1: Attn after heads split, 2: Attn
+
+        self.sample_t = 0
+        self.cache = {}
+        self.encoder_dims = encoder_dims
+        self.prime_len = prime_len
+        self.record_attn = False
+        self.w = None
+
+    def _attn(self, q, k, v, sample):
+        scale = 1. / math.sqrt(math.sqrt(self.n_state // self.n_head))
+        if self.training:
+            w = t.matmul(q * scale, k * scale)
+        else:
+            w = t.matmul(q, k)
+            w.mul_(scale*scale)
+        wtype = w.dtype
+        w = w.float()
+        if self.mask:
+            # Generate appropriate mask to mask out all positions before current
+            # Might take up lot of memory for dense, so can cache it
+            mask = get_mask(self.attn_mask, q.size(-2), k.size(-1), self.blocks, self.spread, w.device, sample, self.sample_t)
+            if mask is not None:
+                #print(mask)
+                w = w * mask + -1e9 * (1 - mask)
+            w = F.softmax(w, dim=-1).type(wtype)
+        else:
+            w = F.softmax(w, dim=-1).type(wtype)
+        if self.record_attn:
+            self.w = w #.float().cpu().numpy()
+            if self.attn_func == 7:
+                # only keep music queries and lyrics keys/values
+                self.w = self.w[:,:,self.prime_len:,:self.prime_len]
+        w = self.attn_dropout(w)
+        a = t.matmul(w, v)
+        return a
+
+    def merge_heads(self, x):
+        x = x.permute(0, 2, 1, 3).contiguous()
+        new_x_shape = (*x.size()[:-2], x.size(-2) * x.size(-1))
+        return x.view(*new_x_shape)  # in Tensorflow implem: fct merge_states
+
+    def split_heads(self, x, k=False):
+        new_x_shape = (*x.size()[:-1], self.n_head, x.size(-1) // self.n_head)
+        x = x.view(*new_x_shape)  # in Tensorflow implem: fct split_states
+        if k:
+            return x.permute(0, 2, 3, 1)
+        else:
+            return x.permute(0, 2, 1, 3)
+
+    def dense_attn(self, query, key, value, sample):
+        query = self.split_heads(query)
+        key = self.split_heads(key, k=True)
+        value = self.split_heads(value)
+        if self.checkpoint_attn == 1 and not sample:
+            a = checkpoint(lambda q,k,v,s=sample: self._attn(q,k,v,s), (query, key, value),
+                       (), True)
+        else:
+            a = self._attn(query,key,value,sample)
+        a = self.merge_heads(a)
+        return a
+
+    def block_attn(self, q, k, v, sample):
+        blocks, block_ctx = self.blocks, self.block_ctx # block_ctx is l // blocks for complete l ie l = n_ctx. Sampling has less l
+        bs, l, d = v.shape # For sample, q_l = 1, k_l = v_l = sample_t
+        if sample:
+            assert l == self._suff_cache_len(), f"{l} != {self._suff_cache_len()}"
+            return self.dense_attn(q, k, v, sample).view(bs, 1, d)
+        else:
+            ql = q.shape[1]
+            q = q.view(bs * ql // block_ctx, block_ctx, d)
+            if ql < l:
+                l = ql
+                k = k[:, -l:].contiguous()
+                v = v[:, -l:].contiguous()
+            k = k.view(bs * l // block_ctx, block_ctx, d)
+            v = v.view(bs * l // block_ctx, block_ctx, d)
+            return self.dense_attn(q, k, v, sample).view(bs, l, d)
+
+    def transpose_block_attn(self, q, k, v, sample):
+        blocks, block_ctx = self.blocks, self.block_ctx # block_ctx is l // blocks for complete l ie l = n_ctx. Sampling has less l
+        bs, l, d = v.shape # For sample, q_l = 1, k_l = v_l = sample_t
+        if sample:
+            block_l = (l - 1) % block_ctx
+            k = k[:,block_l::block_ctx,:]
+            v = v[:,block_l::block_ctx,:]
+            return self.dense_attn(q, k, v, sample).view(bs, 1, d)
+        else:
+            ql = q.shape[1]
+            q = q.view(bs, ql // block_ctx, block_ctx, d).transpose(1,2).contiguous().view(bs * block_ctx, ql // block_ctx, d)
+            k = k.view(bs,  l // block_ctx, block_ctx, d).transpose(1,2).contiguous().view(bs * block_ctx,  l // block_ctx, d)
+            v = v.view(bs,  l // block_ctx, block_ctx, d).transpose(1,2).contiguous().view(bs * block_ctx,  l // block_ctx, d)
+            return self.dense_attn(q, k, v, sample).view(bs, block_ctx, ql // block_ctx, d).transpose(1,2).contiguous().view(bs, ql, d)
+
+    def prev_block_attn(self, q, k, v, sample):
+        blocks, block_ctx = self.blocks, self.block_ctx # block_ctx is l // blocks for complete l ie l = n_ctx. Sampling has less l
+        bs, l, d = v.shape # For sample, q_l = 1, k_l = v_l = sample_t
+        if sample:
+            assert l == self._suff_cache_len(), f"{l} != {self._suff_cache_len()}"
+            block = (l - 1) // block_ctx
+            prev_l = (block - 1) * block_ctx
+            if block > 0:
+                assert prev_l == 0
+                k = k[:, prev_l:prev_l + block_ctx, :]
+                v = v[:, prev_l:prev_l + block_ctx, :]
+            else:
+                k = t.zeros(bs, block_ctx, d, device=q.device, dtype=q.dtype)
+                v = t.zeros(bs, block_ctx, d, device=q.device, dtype=q.dtype)
+            return self.dense_attn(q, k, v, sample).view(bs, 1, d)
+        else:
+            ql = q.shape[1]
+            q = q.view(bs * ql // block_ctx, block_ctx, d)
+            k = t.nn.functional.pad(k.view(bs, l // block_ctx, block_ctx, d)[:, :-1, :, :], (0,0,0,0,1,0)).view(bs * l // block_ctx, block_ctx, d)
+            v = t.nn.functional.pad(v.view(bs, l // block_ctx, block_ctx, d)[:, :-1, :, :], (0,0,0,0,1,0)).view(bs * l // block_ctx, block_ctx, d)
+            if ql < l:
+                qb = ql // block_ctx
+                kb =  l // block_ctx
+                l = ql
+                k = k.view(bs, kb, block_ctx, d)[:, -qb:].contiguous().view(bs * qb, block_ctx, d)
+                v = v.view(bs, kb, block_ctx, d)[:, -qb:].contiguous().view(bs * qb, block_ctx, d)
+            return self.dense_attn(q, k, v, sample).view(bs, l, d)
+
+    def summary_attn(self, q, k, v, sample):
+        blocks, block_ctx = self.blocks, self.block_ctx # block_ctx is l // blocks for complete l ie l = n_ctx. Sampling has less l
+        bs, l, d = v.shape # For sample, q_l = 1, k_l = v_l = sample_t
+        if sample:
+            k = t.nn.functional.pad(k[:, block_ctx-1:blocks*block_ctx-1:block_ctx, :],(0,0,1,0))
+            v = t.nn.functional.pad(v[:, block_ctx-1:blocks*block_ctx-1:block_ctx, :],(0,0,1,0))
+            return self.dense_attn(q, k, v, sample).view(bs, 1, d)
+        else:
+            k = t.nn.functional.pad(k.view(bs, blocks, l // blocks, d)[:, :-1, -1, :],(0,0,1,0)) # bs, blocks, d
+            v = t.nn.functional.pad(v.view(bs, blocks, l // blocks, d)[:, :-1, -1, :],(0,0,1,0)) # bs, blocks, d
+            return self.dense_attn(q, k, v, sample).view(bs, l, d)
+
+    def summary_spread_attn(self, q, k, v, sample):
+        blocks, block_ctx, spread = self.blocks, self.block_ctx, self.spread # block_ctx is l // blocks for complete l ie l = n_ctx. Sampling has less l
+        bs, l, d = v.shape # For sample, q_l = 1, k_l = v_l = sample_t
+        if sample:
+            assert False, "Not yet implemented"
+            # k = t.nn.functional.pad(k,(0,0,block_ctx,(-l)%block_ctx)).view(bs, -1, block_ctx, d)[:,:-1,-spread:,:].contiguous().view(bs, -1, d)
+            # v = t.nn.functional.pad(v,(0,0,block_ctx,(-l)%block_ctx)).view(bs, -1, block_ctx, d)[:,:-1,-spread:,:].contiguous().view(bs, -1, d)
+            # return self.dense_attn(q, k, v, sample).view(bs, 1, d)
+        else:
+            k = t.nn.functional.pad(k.view(bs, blocks, l // blocks, d)[:, :-1, -spread:, :],(0,0,0,0,1,0)).contiguous().view(bs, blocks * spread, d)  # bs, blocks * spread, d
+            v = t.nn.functional.pad(v.view(bs, blocks, l // blocks, d)[:, :-1, -spread:, :],(0,0,0,0,1,0)).contiguous().view(bs, blocks * spread, d)  # bs, blocks * spread, d
+            return self.dense_attn(q, k, v, sample).view(bs, l, d)
+
+    def prime_attn(self, q, k, v, sample):
+        prime_len = self._prime_len
+        k = k[:, :prime_len]
+        v = v[:, :prime_len]
+        return self.dense_attn(q, k, v, sample)
+
+    def decode_attn(self, q, k, v, sample):
+        assert k.shape[1] == v.shape[1] == self.encoder_dims, f'k: {k.shape}, v: {v.shape}, enc_dims: {self.encoder_dims}'
+        return self.dense_attn(q, k, v, sample)
+
+    def factored_qkv(self, x, encoder_kv=None, sample=False):
+        curr_ctx = x.shape[1]
+        assert encoder_kv is None
+        query, key, value = x.chunk(3, dim=2)
+        if sample:
+            self.sample_t += curr_ctx
+            key, value = self._append_cache(key, value)
+            l_cache = self._suff_cache_len()
+            if self._cache_len() > l_cache:
+                self._slice_cache(-l_cache)
+            if curr_ctx > 1:
+                if self.attn_func != 0:
+                    query = self._pad_to_block_ctx(query, query=True)
+                    key = self._pad_to_block_ctx(key)
+                    value = self._pad_to_block_ctx(value)
+                    assert key.shape[1] % self.block_ctx == 0
+                    assert query.shape[1] % self.block_ctx == 0
+                assert key.shape[1] == value.shape[1]
+                assert query.shape[1] <= key.shape[1]
+                sample = False
+            else:
+                key = self.cache['key']
+                value = self.cache['value']
+        return query, key, value, sample
+
+    def prime_qkv(self, x, encoder_kv=None, sample=False):
+        curr_ctx = x.shape[1]
+        assert encoder_kv is None
+        query, key, value = x.chunk(3, dim=2)
+        if sample:
+            if self._cache_len() < self._prime_len:
+                self._append_cache(key, value)
+            if self._cache_len() > self._prime_len:
+                self._slice_cache(0, self._prime_len)
+            key, value = self.cache['key'], self.cache['value']
+            self.sample_t += curr_ctx
+            assert key.shape[1] == value.shape[1] == self._suff_cache_len(), f'k: {key.shape}, v: {value.shape}, prime_dims: {self._suff_cache_len()}'
+        else:
+            assert key.shape[1] == value.shape[1] == self.n_ctx, f'k: {key.shape}, v: {value.shape}, prime_dims: {self.n_ctx}'
+        assert key.shape[0] == value.shape[0] == query.shape[0], f'k: {key.shape}, v: {value.shape}, q: {query.shape}'
+        assert key.shape[2] == value.shape[2] == query.shape[2], f'k: {key.shape}, v: {value.shape}, q: {query.shape}'
+        return query, key, value, sample
+
+    def decode_qkv(self, x, encoder_kv=None, sample=False):
+        curr_ctx = x.shape[1]
+        assert encoder_kv is not None
+        query = x
+        if sample:
+            if self.sample_t == 0:
+                self.cache['key'], self.cache['value'] = self.c_enc_kv(encoder_kv.type_as(x)).chunk(2, dim=2)
+            key, value = self.cache['key'], self.cache['value']
+            self.sample_t += curr_ctx
+        else:
+            key, value = self.c_enc_kv(encoder_kv.type_as(x)).chunk(2, dim=2)
+        assert key.shape[0] == value.shape[0] == query.shape[0], f'k: {key.shape}, v: {value.shape}, q: {query.shape}'
+        assert key.shape[1] == value.shape[1] == self.encoder_dims, f'k: {key.shape}, v: {value.shape}, enc_dims: {self.encoder_dims}'
+        assert key.shape[2] == value.shape[2] == query.shape[2], f'k: {key.shape}, v: {value.shape}, q: {query.shape}'
+        return query, key, value, sample
+
+    def forward(self, x, encoder_kv=None, sample=False):
+        curr_ctx = x.shape[1]
+        x = self.c_attn(x)
+        query, key, value, sample = self.qkv(x, encoder_kv=encoder_kv, sample=sample)
+        if self.checkpoint_attn == 2 and not sample:
+            a = checkpoint(lambda q,k,v,s=sample: self.attn(q,k,v,s), (query, key, value), (), True)
+        else:
+            a = self.attn(query,key,value,sample)
+        if a.shape[1] != curr_ctx:
+            offset = self._offset(curr_ctx)
+            a = a[:,offset:offset + curr_ctx,:].contiguous()
+        a = self.c_proj(a)
+        return self.resid_dropout(a)
+
+    @property
+    def _prime_len(self):
+        prime_len = self.prime_len
+        assert prime_len is not None
+        prime_blocks = (prime_len // self.blocks) + 1
+        return prime_blocks * self.blocks
+
+    def _offset(self, curr_ctx):
+        if self.attn_func == 0:
+            return 0
+        return (self.sample_t - curr_ctx) % self.block_ctx
+
+    def _pad_to_block_ctx(self, x, query=False):
+        l = x.shape[1]
+        offset = self._offset(l) if query else 0
+        n_blocks = (l + offset + self.block_ctx - 1) // self.block_ctx
+        pad = n_blocks * self.block_ctx - l - offset
+        if pad == 0 and offset == 0:
+            return x
+        else:
+            return F.pad(x, (0, 0, offset, pad))
+
+    def _cache_len(self):
+        return 0 if 'key' not in self.cache else self.cache['key'].shape[1]
+
+    def _suff_cache_len(self):
+        """
+        Precondition:
+            key and value are appended with the current context and
+            self.sample_t reflects the 1-indexed sample location in the
+            context.
+        """
+        if self.attn_func == 0:
+            return self.sample_t
+        elif self.attn_func == 1:
+            return (self.sample_t - 1) % self.block_ctx + 1
+        elif self.attn_func == 2:
+            return self.sample_t
+        elif self.attn_func == 3:
+            if self.sample_t <= self.block_ctx:
+                return self.sample_t
+            else:
+                curr_block = (self.sample_t - 1) % self.block_ctx + 1
+                prev_block = self.block_ctx
+                return curr_block + prev_block
+        elif self.attn_func == 6:
+            return self.encoder_dims
+        elif self.attn_func == 7:
+            return min(self.sample_t, self._prime_len)
+        else:
+            raise NotImplementedError()
+
+    def _slice_cache(self, start, end=None):
+        self.cache['key'] = self.cache['key'][:, start:end]
+        self.cache['value'] = self.cache['value'][:, start:end]
+
+    def _append_cache(self, key, value):
+        if 'key' not in self.cache:
+            self.cache['key'] = key
+            self.cache['value'] = value
+        else:
+            old_key, old_value = key, value
+            key = t.cat([self.cache['key'], key], dim=1)
+            value = t.cat([self.cache['value'], value], dim=1)
+            del self.cache['key']
+            del self.cache['value']
+            del old_key
+            del old_value
+            self.cache['key'] = key
+            self.cache['value'] = value
+        return self.cache['key'], self.cache['value']
+
+    def del_cache(self):
+        self.sample_t = 0
+        if 'key' in self.cache:
+            del self.cache['key']
+        if 'value' in self.cache:
+            del self.cache['value']
+        self.cache = {}
+
+    def check(self):
+        blocks = self.blocks or 1
+        spread = self.spread or 1
+        bs, l, d = (4, self.n_ctx, self.n_in)
+        x = t.randn(bs, l, d).cuda()
+        x.requires_grad = True
+        x_out = self.forward(x) # bs, l, d
+        loss = x_out.mean(dim = -1) # bs, l
+        pos = 60
+        grad = t.autograd.grad(loss[2, pos], x)[0]
+
+        assert grad.shape == (bs, l, d)
+        assert (grad[:2] == 0).all()
+        assert (grad[3:] == 0).all()
+        assert (grad[2, (pos + 1):] == 0).all()
+        pos_grad = (t.sum(grad[2] ** 2, dim=-1) > 0).nonzero().view(-1).cpu()
+
+        block_pos = pos - (pos % (l // blocks))
+        exp_pos_grad = {0: t.arange(pos),
+                        1: t.arange(block_pos, pos),
+                        2: t.arange(pos % (l // blocks), pos, l // blocks),
+                        3: t.arange(block_pos - l // blocks, block_pos),
+                        4: t.arange(l // blocks - 1, pos, l // blocks),
+                        5: ((t.arange(pos) % (l // blocks) >= (l // blocks - spread)) & (t.arange(pos) < block_pos)).nonzero().view(-1)}[self.attn_func]
+        exp_pos_grad = t.cat([exp_pos_grad, t.tensor([pos])], dim=-1)
+
+        assert (len(pos_grad) == len(exp_pos_grad)) and (pos_grad == exp_pos_grad).all(), \
+            f"Expected pos grad {exp_pos_grad} got {pos_grad} for attn_func {self.attn_func} pos {pos} l {l} blocks {blocks}"
+
+    def check_cache(self, n_samples, sample_t, fp16):
+        assert self.sample_t == sample_t, f"{self.sample_t} != {sample_t}"
+        if sample_t == 0:
+            assert self.cache == {}
+        else:
+            dtype = {True: t.float16, False: t.float32}[fp16]
+            l_cache = self._suff_cache_len()
+            assert self.cache['key'].shape == (n_samples, l_cache, self.n_state)
+            assert self.cache['value'].shape == (n_samples, l_cache, self.n_state)
+            assert self.cache['key'].dtype == dtype, f"Expected {dtype}, got {self.cache['key'].dtype}"
+            assert self.cache['value'].dtype == dtype, f"Expected {dtype}, got {self.cache['value'].dtype}"
+
+    def check_sample(self):
+        t.manual_seed(42)
+        bs, l, d = (4, self.n_ctx, self.n_in)
+        prime = 5
+        x = t.randn(bs, l, d).cuda()
+        xs = t.chunk(x, l, dim=1)
+        assert self.sample_t == 0
+        assert self.cache == {}
+
+        with t.no_grad():
+            enc_l = self.encoder_dims
+            encoder_kv = None
+            if self.attn_func == 6:
+                encoder_kv = t.randn(bs, enc_l, d).cuda()
+
+            # Normal path
+            x_out_normal = self.forward(x, encoder_kv=encoder_kv)
+
+            # Sampling path
+            x_out_sample = t.cat([self.forward(xs[i], encoder_kv=encoder_kv, sample=True) for i in range(l)],dim=1)
+        max_err = t.max(t.abs(x_out_sample - x_out_normal))
+        assert max_err < 1e-8, f"Max sampling err is {max_err} {[i for i in range(l) if t.max(t.abs(x_out_sample - x_out_normal)[:,i,:]) > 1e-8]}"
+
+        with t.no_grad():
+            x_out_normal = x_out_normal[:,:prime,:]
+            # Prime sampling path
+            self.del_cache()
+            x_out_sample = self.forward(x[:,:prime,:].contiguous(), encoder_kv=encoder_kv, sample=True)
+            self.check_cache(bs, prime, False)
+
+        max_err = t.max(t.abs(x_out_sample - x_out_normal))
+        assert max_err < 1e-8, f"Max prime sampling err is {max_err} {[i for i in range(prime) if t.max(t.abs(x_out_sample - x_out_normal)[:,i,:]) > 1e-8]}"
+
+    def check_chunks(self, chunk_size):
+        t.manual_seed(42)
+        bs, l, d = (4, self.n_ctx, self.n_in)
+        enc_l = self.encoder_dims
+        assert l % chunk_size == 0
+        n_chunks = l // chunk_size
+        with t.no_grad():
+            encoder_kv = None
+            x = t.randn(bs, l, d).cuda()
+            if self.attn_func == 6:
+                encoder_kv = t.randn(bs, enc_l, d).cuda()
+
+            self.del_cache()
+            y_forw = self.forward(x, encoder_kv=encoder_kv, sample=False)
+            self.del_cache()
+            y_forw_sample = self.forward(x, encoder_kv=encoder_kv, sample=True)
+            max_err = t.max(t.abs(y_forw - y_forw_sample))
+            assert max_err <= 1e-6, f"Max err is {max_err} {[i for i in range(l) if t.max(t.abs(y_forw - y_forw_sample)[:, i, :]) > 1e-6]}"
+
+            self.del_cache()
+            x_chunks = t.chunk(x, n_chunks, dim=1)
+            y_chunks = []
+            total_len = 0
+            for x_chunk in x_chunks:
+                y_chunk = self.forward(x_chunk.contiguous(), encoder_kv=encoder_kv, sample=True)
+                total_len += x_chunk.shape[1]
+                self.check_cache(bs, total_len, False)
+                y_chunks.append(y_chunk)
+            y_forw_in_chunks = t.cat(y_chunks, dim=1)
+
+            max_err = t.max(t.abs(y_forw - y_forw_in_chunks))
+            assert max_err <= 1e-6, f"Max err is {max_err} {[i for i in range(l) if t.max(t.abs(y_forw - y_forw_in_chunks)[:, i, :]) > 1e-6]}"
+
+
+if __name__ == '__main__':
+    from jukebox.utils.dist_utils import setup_dist_from_mpi
+    setup_dist_from_mpi(port=29600)
+    n_in = 16
+    n_state = n_in * 2
+    n_ctx = 6144
+    n_head = 4
+    n_depth = 12
+    blocks = 64
+    chunk_size = 8
+    for attn_func in [0, 1, 2, 3, 6, 7]:
+        encoder_dims = {0: 0, 1: 0, 2: 0, 3: 0, 6: 64, 7: 0}[attn_func]
+        prime_len = {0: 0, 1: 0, 2: 0, 3: 0, 6: 0, 7: 384}[attn_func]
+        attn = FactoredAttention(n_in, n_ctx + prime_len, n_state, n_head, mask=True,
+                                 attn_func=attn_func, blocks=blocks,
+                                 encoder_dims=encoder_dims, prime_len=prime_len)
+        attn.training = False
+        attn.check_sample()
+        attn.check_chunks(chunk_size)
+        print(f"Checked attn_func: {attn_func}")

jukebox/transformer/ops.py

@@ -0,0 +1,142 @@
+import math
+import numpy as np
+import torch as t
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import FusedLayerNorm if we have apex, otherwise use regular LayerNorm
+try:
+    from apex.normalization import FusedLayerNorm
+    print("Using apex FusedLayerNorm")
+except ImportError:
+    from torch.nn import LayerNorm as FusedLayerNorm
+
+class LayerNorm(FusedLayerNorm):
+    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
+        super().__init__(normalized_shape, eps=eps, elementwise_affine=elementwise_affine)
+        self.width = np.prod(normalized_shape)
+        self.max_numel = 65535*self.width
+
+    def forward(self, input):
+        if input.numel() > self.max_numel:
+            return F.layer_norm(input.float(), self.normalized_shape, self.weight, self.bias, self.eps).type_as(input)
+        else:
+            return super(LayerNorm, self).forward(input.float()).type_as(input)
+
+def gelu(x):
+    return 0.5 * x * (1 + t.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * t.pow(x, 3))))
+
+
+def swish(x):
+    return x * t.sigmoid(x)
+
+@t.jit.script
+def quick_gelu(x):
+    return x * t.sigmoid(1.702 * x)
+
+@t.jit.script
+def quick_gelu_bwd(x, grad_output):
+    sig = t.sigmoid(1.702 * x)
+    return grad_output * sig * (1.702 * x * (1 - sig) + 1.)
+
+class QuickGelu(t.autograd.Function):
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        return quick_gelu(x)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return quick_gelu_bwd(ctx.saved_tensors[0], grad_output)
+
+def memory_efficient_quick_gelu(x):
+    return QuickGelu.apply(x)
+
+ACT_FNS = {
+    'relu': t.nn.functional.relu,
+    'swish': swish,
+    'gelu': gelu,
+    'quick_gelu': memory_efficient_quick_gelu #quick_gelu
+}
+
+def _move_to_gpu_and_convert_conv_weights_to_fp16(l):
+    l.cuda()
+    if isinstance(l, Conv1D):
+        l.w.data = l.w.data.half()
+
+def _convert_conv_weights_to_fp32(l):
+    if isinstance(l, Conv1D):
+        l.w.data = l.w.data.float()
+
+def _convert_conv_weights_to_fp16(l):
+    if isinstance(l, Conv1D):
+        l.w.data = l.w.data.half()
+
+def _convert_embedding_weights_to_fp16(l):
+    if isinstance(l, t.nn.Embedding):
+        l.weight.data = l.weight.data.half()
+
+def _convert_embedding_weights_to_fp32(l):
+    if isinstance(l, t.nn.Embedding):
+        l.weight.data = l.weight.data.float()
+
+class Conv1D(nn.Module):
+    def __init__(self, n_in, n_out, zero_out=False, init_scale=1.0):
+        super(Conv1D, self).__init__()
+        self.n_in = n_in
+        self.n_out = n_out
+        if zero_out:
+            w = t.zeros(n_in, n_out)
+        else:
+            w = t.empty(n_in, n_out)
+            nn.init.normal_(w, std=0.02 * init_scale)
+        b = t.zeros(n_out)
+        self.w = nn.Parameter(w)
+        self.b = nn.Parameter(b)
+
+    def forward(self, x):
+        size_out = (*x.size()[:-1], self.n_out)
+        x = t.addmm(self.b.type_as(x), x.view(-1, x.size(-1)), self.w.type_as(x)) # If x if float then float else half
+        x = x.view(*size_out)
+        return x
+
+# For large contexts, mask's can take up memory, so you can make a single saved mask for all layers
+class Mask(nn.Module):
+    def __init__(self, n_ctx):
+        super().__init__()
+        self.register_buffer('b', t.tril(t.ones(n_ctx, n_ctx)).view(1, 1, n_ctx, n_ctx))
+
+    def forward(self, w):
+        w = w * self.b + -1e9 * (1 - self.b)  # For fp16 do w = w.float().masked_fill(self.b, float('-inf')
+        return w
+
+def filter_logits(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):
+    """ Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+        Args:
+            logits: logits distribution shape (vocabulary size)
+            top_k >0: keep only top k tokens with highest probability (top-k filtering).
+            top_p >0.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+    """
+    #assert logits.dim() == 2  # batch size 1 for now - could be updated for more but the code would be less clear
+    logits = logits.clone()
+    top_k = min(top_k, logits.size(-1))  # Safety check
+    assert (top_k == 0) or (top_p == 0.0)
+    if top_k > 0:
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < t.topk(logits, top_k, dim=-1)[0][..., -1:]
+        logits[indices_to_remove] = filter_value
+
+    if top_p > 0.0:
+        sorted_logits, sorted_indices = t.sort(logits, descending=True, dim=-1)
+        cumulative_probs = t.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold
+        sorted_indices_to_remove = cumulative_probs > top_p
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        #indices_to_remove = sorted_indices[sorted_indices_to_remove]
+        indices_to_remove = t.zeros_like(logits, dtype=t.uint8).scatter_(dim=-1, index=sorted_indices, src=sorted_indices_to_remove)
+        logits[indices_to_remove] = filter_value
+    return logits

jukebox/transformer/transformer.py

@@ -0,0 +1,239 @@
+import functools
+import numpy as np
+import torch as t
+import torch.nn as nn
+import jukebox.utils.dist_adapter as dist
+
+from jukebox.transformer.ops import Conv1D, ACT_FNS, LayerNorm
+from jukebox.transformer.factored_attention import FactoredAttention
+from jukebox.utils.checkpoint import checkpoint
+
+def _convert_mlp_traced(l):
+    if isinstance(l, ResAttnBlock):
+        l.mlp = t.jit.trace(l.mlp, t.randn(1, 1, l.n_in).cuda())
+
+def _convert_mlp_traced_fp16(l):
+    if isinstance(l, ResAttnBlock):
+        l.mlp = t.jit.trace(l.mlp, t.randn(1, 1, l.n_in).cuda().half())
+
+class MLP(nn.Module):
+    def __init__(self, n_in, n_state, resid_dropout=0.0, afn='quick_gelu', zero_out=False, init_scale=1.0):
+        super().__init__()
+        self.c_fc = Conv1D(n_in, n_state, init_scale=init_scale)
+        self.c_proj = Conv1D(n_state, n_in, zero_out, init_scale=init_scale)
+        self.act = ACT_FNS[afn]
+        self.resid_dropout = nn.Dropout(resid_dropout) if resid_dropout > 0.0 else lambda x: x
+
+    def forward(self, x):
+        m = self.act(self.c_fc(x))
+        m = self.c_proj(m)
+        return self.resid_dropout(m)
+
+class ResAttnBlock(nn.Module):
+    def __init__(self, n_in, n_ctx, n_head,
+                 attn_dropout=0.0, resid_dropout=0.0,
+                 afn='quick_gelu', scale=True, mask=False,
+                 zero_out=False, init_scale=1.0, res_scale=1.0,
+                 m_attn = 0.25, m_mlp = 1.,
+                 checkpoint_attn = 0, checkpoint_mlp = 0,
+                 attn_func=0, blocks=None, spread=None,
+                 encoder_dims=None, prime_len=None):
+        super().__init__()
+        self.attn = FactoredAttention(n_in=n_in, n_ctx=n_ctx, n_state=int(m_attn * n_in), n_head=n_head,
+                                      attn_dropout=attn_dropout, resid_dropout=resid_dropout,
+                                      scale=scale, mask=mask,
+                                      zero_out=zero_out, init_scale=init_scale,
+                                      checkpoint_attn=checkpoint_attn,
+                                      attn_func=attn_func, blocks=blocks, spread=spread,
+                                      encoder_dims=encoder_dims, prime_len=prime_len)
+        self.ln_0 = LayerNorm(n_in)
+        self.mlp = MLP(n_in=n_in, n_state=int(m_mlp * n_in),
+                       resid_dropout=resid_dropout,
+                       afn=afn,
+                       zero_out=zero_out, init_scale=init_scale)
+        self.ln_1 = LayerNorm(n_in)
+        self.res_scale = res_scale
+
+        self.checkpoint_attn = checkpoint_attn
+        self.checkpoint_mlp = checkpoint_mlp
+        self.n_in = n_in
+        self.attn_func = attn_func
+
+    def forward(self, x, encoder_kv, sample=False):
+        if sample:
+            a = self.attn(self.ln_0(x), encoder_kv, sample)
+            m = self.mlp(self.ln_1(x + a))
+        else:
+            if self.attn_func == 6:
+                assert encoder_kv is not None
+                a = checkpoint(lambda _x,_enc_kv,_s=sample: self.attn(self.ln_0(_x),_enc_kv,_s),
+                               (x,encoder_kv),
+                               (*self.attn.parameters(), *self.ln_0.parameters()),
+                               self.checkpoint_attn == 3)  # 2 recomputes after the projections, and 1 recomputes after head splitting.
+            else:
+                assert encoder_kv is None
+                a = checkpoint(lambda _x,_enc_kv=None,_s=sample: self.attn(self.ln_0(_x),_enc_kv,_s),
+                               (x,),
+                               (*self.attn.parameters(), *self.ln_0.parameters()),
+                               self.checkpoint_attn == 3)  # 2 recomputes after the projections, and 1 recomputes after head splitting.
+            m = checkpoint(lambda _x: self.mlp(self.ln_1(_x)), (x + a,),
+                           (*self.mlp.parameters(), *self.ln_1.parameters()),
+                           self.checkpoint_mlp == 1)
+        if self.res_scale == 1.0:
+            h = x + a + m
+        else:
+            h = x + self.res_scale * (a + m)
+        return h
+
+class Transformer(nn.Module):
+    def __init__(self, n_in, n_ctx, n_head, n_depth,
+                 attn_dropout=0.0, resid_dropout=0.0,
+                 afn='quick_gelu', scale=True, mask=False,
+                 zero_out=False, init_scale=1.0, res_scale=False,
+                 m_attn=0.25, m_mlp=1.,
+                 checkpoint_attn=0, checkpoint_mlp=0, checkpoint_res=0,
+                 attn_order=0, blocks=None, spread=None,
+                 encoder_dims=None, prime_len=None):
+        super().__init__()
+        self.n_in = n_in
+        self.n_ctx = n_ctx
+        self.encoder_dims = encoder_dims
+        self.blocks = blocks
+        if blocks is not None:
+            assert n_ctx % blocks == 0
+            self.block_ctx = n_ctx // blocks
+        self.prime_len = prime_len
+        self.n_head = n_head
+
+        res_scale = 1.0 / n_depth if res_scale else 1.0
+
+        # Orders of attn_func
+        attn_func = {0: lambda d: 0,                    # Complete dense attn
+                     1: lambda d: [1,2][d%2],           # Alternate row and column attn
+                     2: lambda d: [1,2,3][d % 3],       # Alternate row, column and previous row attn
+                     3: lambda d: [1,4][d % 2],         # Alternate row and last column
+                     4: lambda d: [1,5][d % 2],         # Alternate row and last k columns
+                     5: lambda d: [1,4,1,1][d % 4],      # Alternate row, last column, row, row
+                     6: lambda d: [1,2,3,6][d % 4],
+                     7: lambda d: [*[1,2,3]*5,6][d%16],
+                     8: lambda d: [1,2,3,1,2,3,1,2,3,6][d%10], # Used by separated_enc_dec model with lyrics
+                     9: lambda d: [1,2,3,0][d % 4],
+                     10: lambda d: [*[1,2,3,1,2,3,1,2,3],*[1,2,3,1,2,3,1,2,3,6]*7][d%79], # Used by large separated_enc_dec model with lyrics
+                     11: lambda d: [6,6,0][d%3] if d%16 == 15 else [1,2,3][d%3],
+                     12: lambda d: [7,7,0][d%3] if d%16 == 15 else [1,2,3][d%3], # Used by single_enc_dec model with lyrics
+                     }[attn_order]
+
+        attn_cycle = {0:1, 1:2, 2:3, 3:2, 4:2, 5:4, 6:4, 7:16, 8:10, 9:4, 10:79, 11:16, 12:16}[attn_order]
+        #assert n_depth % attn_cycle == 0, f'Depth {n_depth} not a multiple of cycle {attn_cycle} for attn_order {attn_order}'
+
+        attn_block = lambda d: ResAttnBlock(n_in=n_in, n_ctx=n_ctx, n_head=n_head,
+                                  attn_dropout=attn_dropout, resid_dropout=resid_dropout,
+                                  afn=afn, scale=scale, mask=mask,
+                                  zero_out=zero_out if attn_func(d) !=6 else True,
+                                  init_scale=init_scale, res_scale=res_scale,
+                                  m_attn=m_attn, m_mlp=m_mlp,
+                                  checkpoint_attn=checkpoint_attn, checkpoint_mlp=checkpoint_mlp,
+                                  attn_func=attn_func(d), blocks=blocks, spread=spread,
+                                  encoder_dims=encoder_dims, prime_len=prime_len)
+
+        self.checkpoint_res = checkpoint_res
+        self._attn_mods = nn.ModuleList()
+        for d in range(n_depth):
+            self._attn_mods.append(attn_block(d))
+        self.ws = []
+
+
+    def set_record_attn(self, record_attn):
+        """
+        Arguments:
+            record_attn (bool or set): Makes forward prop dump self-attention
+                softmaxes to self.ws. Either a set of layer indices indicating
+                which layers to store, or a boolean value indicating whether to
+                dump all.
+        """
+        def _should_record_attn(layer_idx):
+            if isinstance(record_attn, bool):
+                return record_attn
+            return layer_idx in record_attn
+        for i, l in enumerate(self._attn_mods):
+            l.attn.record_attn = _should_record_attn(i)
+        if record_attn:
+            assert self.ws == []
+            for l in self._attn_mods:
+                assert l.attn.w == None
+        else:
+            self.ws = []
+            for l in self._attn_mods:
+                l.attn.w = None
+
+    def forward(self, x, encoder_kv=None, sample=False, fp16=False, fp16_out=False):
+        if fp16:
+            x = x.half()
+
+        # Blocks
+        for i,l in enumerate(self._attn_mods):
+            if self.checkpoint_res == 1 and not sample:
+                if l.attn_func == 6:
+                    assert encoder_kv is not None
+                    f = functools.partial(l, sample=sample)
+                    x = checkpoint(f, (x, encoder_kv), l.parameters(), True)
+                else:
+                    f = functools.partial(l, encoder_kv=None, sample=sample)
+                    x = checkpoint(f, (x,), l.parameters(), True)
+            else:
+                if l.attn_func == 6:
+                    x = l(x, encoder_kv=encoder_kv, sample=sample)
+                else:
+                    x = l(x, encoder_kv=None, sample=sample)
+            if l.attn.record_attn:
+                self.ws.append(l.attn.w)
+        if not fp16_out:
+            x = x.float()
+        return x
+
+    def check_cache(self, n_samples, sample_t, fp16):
+        for l in self._attn_mods:
+            l.attn.check_cache(n_samples, sample_t, fp16)
+
+    def del_cache(self):
+        for l in self._attn_mods:
+            l.attn.del_cache()
+
+    def check_sample(self):
+        bs, l, s, d = (4, self.n_ctx, self.encoder_dims, self.n_in)
+        prime = 5
+        with t.no_grad():
+            encoder_kv = t.randn(bs, s, d).cuda()
+            x = t.randn(bs, l, d).cuda()
+            y_forw = self.forward(x, encoder_kv=encoder_kv, sample=True)
+
+            self.del_cache()
+            x_chunks = t.chunk(x, 4, dim=1)
+            y_chunks = []
+            n = 0
+            for x_chunk in x_chunks:
+                self.check_cache(bs, n, False)
+                y_chunk = self.forward(x_chunk, encoder_kv=encoder_kv, sample=True)
+                y_chunks.append(y_chunk)
+                n += x_chunk.shape[1]
+            self.check_cache(bs, n, False)
+            y_forw_in_chunks = t.cat(y_chunks, dim=1)
+
+            max_err = t.max(t.abs(y_forw - y_forw_in_chunks))
+            assert max_err <= 1e-6, f"Max err is {max_err} {[i for i in range(l) if t.max(t.abs(y_forw - y_forw_in_chunks)[:, i, :]) > 1e-6]}"
+
+
+if __name__ == '__main__':
+    from jukebox.utils.dist_utils import setup_dist_from_mpi
+    setup_dist_from_mpi(port=29600)
+    n_in = 16
+    n_ctx = 192
+    n_head = 4
+    n_depth = 12
+    blocks = 16
+    for attn_order in [0,2,6]:
+        encoder_dims = {0: 0, 2: 0, 6: 64}[attn_order]
+        prior = Transformer(n_in, n_ctx, n_head, n_depth, mask=True, attn_order=attn_order, encoder_dims=encoder_dims, blocks=blocks).cuda()
+        prior.training = False
+        prior.check_sample()
+        print(f"Checked attn_order: {attn_order}")

jukebox/utils/audio_utils.py

@@ -0,0 +1,148 @@
+import numpy as np
+import torch as t
+import jukebox.utils.dist_adapter as dist
+import soundfile
+import librosa
+from jukebox.utils.dist_utils import print_once
+
+class DefaultSTFTValues:
+    def __init__(self, hps):
+        self.sr = hps.sr
+        self.n_fft = 2048
+        self.hop_length = 256
+        self.window_size = 6 * self.hop_length
+
+class STFTValues:
+    def __init__(self, hps, n_fft, hop_length, window_size):
+        self.sr = hps.sr
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.window_size = window_size
+
+def calculate_bandwidth(dataset, hps, duration=600):
+    hps = DefaultSTFTValues(hps)
+    n_samples = int(dataset.sr * duration)
+    l1, total, total_sq, n_seen, idx = 0.0, 0.0, 0.0, 0.0, dist.get_rank()
+    spec_norm_total, spec_nelem = 0.0, 0.0
+    while n_seen < n_samples:
+        x = dataset[idx]
+        if isinstance(x, (tuple, list)):
+            x, y = x
+        samples = x.astype(np.float64)
+        stft = librosa.core.stft(np.mean(samples, axis=1), hps.n_fft, hop_length=hps.hop_length, win_length=hps.window_size)
+        spec = np.absolute(stft)
+        spec_norm_total += np.linalg.norm(spec)
+        spec_nelem += 1
+        n_seen += int(np.prod(samples.shape))
+        l1 += np.sum(np.abs(samples))
+        total += np.sum(samples)
+        total_sq += np.sum(samples ** 2)
+        idx += max(16, dist.get_world_size())
+
+    if dist.is_available():
+        from jukebox.utils.dist_utils import allreduce
+        n_seen = allreduce(n_seen)
+        total = allreduce(total)
+        total_sq = allreduce(total_sq)
+        l1 = allreduce(l1)
+        spec_nelem = allreduce(spec_nelem)
+        spec_norm_total = allreduce(spec_norm_total)
+
+    mean = total / n_seen
+    bandwidth = dict(l2 = total_sq / n_seen - mean ** 2,
+                     l1 = l1 / n_seen,
+                     spec = spec_norm_total / spec_nelem)
+    print_once(bandwidth)
+    return bandwidth
+
+def audio_preprocess(x, hps):
+    # Extra layer in case we want to experiment with different preprocessing
+    # For two channel, blend randomly into mono (standard is .5 left, .5 right)
+
+    # x: NTC
+    x = x.float()
+    if x.shape[-1]==2:
+        if hps.aug_blend:
+            mix=t.rand((x.shape[0],1), device=x.device) #np.random.rand()
+        else:
+            mix = 0.5
+        x=(mix*x[:,:,0]+(1-mix)*x[:,:,1])
+    elif x.shape[-1]==1:
+        x=x[:,:,0]
+    else:
+        assert False, f'Expected channels {hps.channels}. Got unknown {x.shape[-1]} channels'
+
+    # x: NT -> NTC
+    x = x.unsqueeze(2)
+    return x
+
+def audio_postprocess(x, hps):
+    return x
+
+def stft(sig, hps):
+    return t.stft(sig, hps.n_fft, hps.hop_length, win_length=hps.window_size, window=t.hann_window(hps.window_size, device=sig.device))
+
+def spec(x, hps):
+    return t.norm(stft(x, hps), p=2, dim=-1)
+
+def norm(x):
+    return (x.view(x.shape[0], -1) ** 2).sum(dim=-1).sqrt()
+
+def squeeze(x):
+    if len(x.shape) == 3:
+        assert x.shape[-1] in [1,2]
+        x = t.mean(x, -1)
+    if len(x.shape) != 2:
+        raise ValueError(f'Unknown input shape {x.shape}')
+    return x
+
+def spectral_loss(x_in, x_out, hps):
+    hps = DefaultSTFTValues(hps)
+    spec_in = spec(squeeze(x_in.float()), hps)
+    spec_out = spec(squeeze(x_out.float()), hps)
+    return norm(spec_in - spec_out)
+
+def multispectral_loss(x_in, x_out, hps):
+    losses = []
+    assert len(hps.multispec_loss_n_fft) == len(hps.multispec_loss_hop_length) == len(hps.multispec_loss_window_size)
+    args = [hps.multispec_loss_n_fft,
+            hps.multispec_loss_hop_length,
+            hps.multispec_loss_window_size]
+    for n_fft, hop_length, window_size in zip(*args):
+        hps = STFTValues(hps, n_fft, hop_length, window_size)
+        spec_in = spec(squeeze(x_in.float()), hps)
+        spec_out = spec(squeeze(x_out.float()), hps)
+        losses.append(norm(spec_in - spec_out))
+    return sum(losses) / len(losses)
+
+def spectral_convergence(x_in, x_out, hps, epsilon=2e-3):
+    hps = DefaultSTFTValues(hps)
+    spec_in = spec(squeeze(x_in.float()), hps)
+    spec_out = spec(squeeze(x_out.float()), hps)
+
+    gt_norm = norm(spec_in)
+    residual_norm = norm(spec_in - spec_out)
+    mask = (gt_norm > epsilon).float()
+    return (residual_norm * mask) / t.clamp(gt_norm, min=epsilon)
+
+def log_magnitude_loss(x_in, x_out, hps, epsilon=1e-4):
+    hps = DefaultSTFTValues(hps)
+    spec_in = t.log(spec(squeeze(x_in.float()), hps) + epsilon)
+    spec_out = t.log(spec(squeeze(x_out.float()), hps) + epsilon)
+    return t.mean(t.abs(spec_in - spec_out))
+
+def load_audio(file, sr, offset, duration, mono=False):
+    # Librosa loads more filetypes than soundfile
+    x, _ = librosa.load(file, sr=sr, mono=mono, offset=offset/sr, duration=duration/sr)
+    if len(x.shape) == 1:
+        x = x.reshape((1, -1))
+    return x    
+
+
+def save_wav(fname, aud, sr):
+    # clip before saving?
+    aud = t.clamp(aud, -1, 1).cpu().numpy()
+    for i in list(range(aud.shape[0])):
+        soundfile.write(f'{fname}/item_{i}.wav', aud[i], samplerate=sr, format='wav')
+
+

jukebox/utils/checkpoint.py

@@ -0,0 +1,32 @@
+# Simple gradient checkpointing. Works with distributed data parallel
+import torch as t
+
+def checkpoint(func, inputs, params, flag):
+    if flag:
+        args = inputs + tuple(params)
+        return CheckpointFunction.apply(func, len(inputs), *args)
+    else:
+        return func(*inputs)
+
+class CheckpointFunction(t.autograd.Function):
+    @staticmethod
+    def forward(ctx, run_function, length, *args):
+        ctx.run_function = run_function
+        ctx.input_tensors = list(args[:length])
+        ctx.input_params = list(args[length:])
+        with t.no_grad():
+            output_tensors = ctx.run_function(*ctx.input_tensors)
+        return output_tensors
+
+    @staticmethod
+    def backward(ctx, *output_grads):
+        for i in range(len(ctx.input_tensors)):
+            temp = ctx.input_tensors[i]
+            ctx.input_tensors[i] = temp.detach()
+            ctx.input_tensors[i].requires_grad = temp.requires_grad
+        with t.enable_grad():
+            output_tensors = ctx.run_function(*ctx.input_tensors)
+        input_grads = t.autograd.grad(output_tensors, ctx.input_tensors + ctx.input_params, output_grads, allow_unused=True)
+        del ctx.input_tensors
+        del output_tensors
+        return (None, None) + input_grads

jukebox/utils/dist_adapter.py

@@ -0,0 +1,86 @@
+import torch.distributed as dist
+from enum import Enum
+
+class ReduceOp(Enum):
+    SUM = 0,
+    PRODUCT = 1,
+    MIN = 2,
+    MAX = 3
+
+    def ToDistOp(self):
+        return {
+            self.SUM: dist.ReduceOp.SUM,
+            self.PRODUCT: dist.ReduceOp.PRODUCT,
+            self.MIN: dist.ReduceOp.MIN,
+            self.MAX: dist.ReduceOp.MAX
+        }[self]
+
+def is_available():
+    return dist.is_available()
+
+def get_rank():
+    if is_available():
+        return _get_rank()
+    else:
+        return 0
+
+def get_world_size():
+    if is_available():
+        return _get_world_size()
+    else:
+        return 1
+
+def barrier():
+    if is_available():
+        return _barrier()
+    #else: do nothing
+
+def all_gather(tensor_list, tensor):
+    if is_available():
+        return _all_gather(tensor_list, tensor)
+    else:
+        tensor_list[0] = tensor
+
+def all_reduce(tensor, op=ReduceOp.SUM):
+    if is_available():
+        return _all_reduce(tensor, op)
+    #else: do nothing
+
+def reduce(tensor, dst, op=ReduceOp.SUM):
+    if is_available():
+        return _reduce(tensor, dst, op)
+    #else: do nothing
+
+def broadcast(tensor, src):
+    if is_available():
+        return _broadcast(tensor, src)
+    #else: do nothing
+
+def init_process_group(backend, init_method):
+    if is_available():
+        return _init_process_group(backend, init_method)
+    #else: do nothing
+
+def _get_rank():
+    return dist.get_rank()
+
+def _barrier():
+    return dist.barrier()
+
+def _get_world_size():
+    return dist.get_world_size()
+
+def _all_gather(tensor_list, tensor):
+    return dist.all_gather(tensor_list, tensor)
+
+def _all_reduce(tensor, op):
+    return dist.all_reduce(tensor, op.ToDistOp())
+
+def _reduce(tensor, dst, op):
+    return dist.reduce(tensor, dst, op.ToDistOp())
+
+def _broadcast(tensor, src):
+    return dist.broadcast(tensor, src)
+
+def _init_process_group(backend, init_method):
+    return dist.init_process_group(backend, init_method)

jukebox/utils/dist_utils.py

@@ -0,0 +1,101 @@
+import os
+from time import sleep
+import torch
+import jukebox.utils.dist_adapter as dist
+
+def print_once(msg):
+    if (not dist.is_available()) or dist.get_rank()==0:
+        print(msg)
+
+def print_all(msg):
+    if (not dist.is_available()):
+        print(msg)
+    elif dist.get_rank()%8==0:
+        print(f'{dist.get_rank()//8}: {msg}')
+
+def allgather(x):
+    xs = [torch.empty_like(x) for _ in range(dist.get_world_size())]
+    dist.all_gather(xs, x)
+    xs = torch.cat(xs, dim=0)
+    return xs
+
+def allreduce(x, op=dist.ReduceOp.SUM):
+    x = torch.tensor(x).float().cuda()
+    dist.all_reduce(x, op=op)
+    return x.item()
+
+def allgather_lists(xs):
+    bs = len(xs)
+    total_bs = dist.get_world_size()*len(xs)
+    lengths = torch.tensor([len(x) for x in xs], dtype=t.long, device='cuda')
+    lengths = allgather(lengths)
+    assert lengths.shape == (total_bs,)
+    max_length = torch.max(lengths).item()
+
+    xs = torch.tensor([[*x, *[0]*(max_length - len(x))] for x in xs], device='cuda')
+    assert xs.shape == (bs, max_length), f'Expected {(bs, max_length)}, got {xs.shape}'
+    xs = allgather(xs)
+    assert xs.shape == (total_bs,max_length), f'Expected {(total_bs, max_length)}, got {xs.shape}'
+
+    return [xs[i][:lengths[i]].cpu().numpy().tolist() for i in range(total_bs)]
+
+def setup_dist_from_mpi(
+    master_addr="127.0.0.1", backend="nccl", port=29500, n_attempts=5, verbose=False
+):
+    if dist.is_available():
+        return _setup_dist_from_mpi(master_addr, backend, port, n_attempts, verbose)
+    else:
+        use_cuda = torch.cuda.is_available()
+        print(f'Using cuda {use_cuda}')
+
+        mpi_rank = 0
+        local_rank = 0
+
+        device = torch.device("cuda", local_rank) if use_cuda else torch.device("cpu")
+        torch.cuda.set_device(local_rank)
+
+        return mpi_rank, local_rank, device
+
+def _setup_dist_from_mpi(master_addr, backend, port, n_attempts, verbose):
+    from mpi4py import MPI  # This must be imported in order to get e   rrors from all ranks to show up
+
+    mpi_rank = MPI.COMM_WORLD.Get_rank()
+    mpi_size = MPI.COMM_WORLD.Get_size()
+
+
+    os.environ["RANK"] = str(mpi_rank)
+    os.environ["WORLD_SIZE"] = str(mpi_size)
+    os.environ["MASTER_ADDR"] = master_addr
+    os.environ["MASTER_PORT"] = str(port)
+    os.environ["NCCL_LL_THRESHOLD"] = "0"
+    os.environ["NCCL_NSOCKS_PERTHREAD"] = "2"
+    os.environ["NCCL_SOCKET_NTHREADS"] = "8"
+
+    # Pin this rank to a specific GPU on the node
+    local_rank = mpi_rank % 8
+    if torch.cuda.is_available():
+        torch.cuda.set_device(local_rank)
+
+    if verbose:
+        print(f"Connecting to master_addr: {master_addr}")
+
+    # There is a race condition when initializing NCCL with a large number of ranks (e.g 500 ranks)
+    # We guard against the failure and then retry
+    for attempt_idx in range(n_attempts):
+        try:
+            dist.init_process_group(backend=backend, init_method=f"env://")
+            assert dist.get_rank() == mpi_rank
+
+            use_cuda = torch.cuda.is_available()
+            print(f'Using cuda {use_cuda}')
+            local_rank = mpi_rank % 8
+            device = torch.device("cuda", local_rank) if use_cuda else torch.device("cpu")
+            torch.cuda.set_device(local_rank)
+
+            return mpi_rank, local_rank, device
+        except RuntimeError as e:
+            print(f"Caught error during NCCL init (attempt {attempt_idx} of {n_attempts}): {e}")
+            sleep(1 + (0.01 * mpi_rank))  # Sleep to avoid thundering herd
+            pass
+
+    raise RuntimeError("Failed to initialize NCCL")

jukebox/utils/ema.py

@@ -0,0 +1,94 @@
+import torch
+from torch._utils import _flatten_dense_tensors
+import numpy as np
+
+# EMA always in float, as accumulation needs lots of bits
+class EMA:
+    def __init__(self, params, mu=0.999):
+        self.mu = mu
+        self.state = [(p, self.get_model_state(p)) for p in params if p.requires_grad]
+
+    def get_model_state(self, p):
+        return p.data.float().detach().clone()
+
+    def step(self):
+        for p, state in self.state:
+            state.mul_(self.mu).add_(1 - self.mu, p.data.float())
+
+    def swap(self):
+        # swap ema and model params
+        for p, state in self.state:
+            other_state = self.get_model_state(p)
+            p.data.copy_(state.type_as(p.data))
+            state.copy_(other_state)
+
+
+class CPUEMA:
+    def __init__(self, params, mu=0.999, freq=1):
+        self.mu = mu**freq
+        self.state = [(p, self.get_model_state(p)) for p in params if p.requires_grad]
+        self.freq = freq
+        self.steps = 0
+
+    def get_model_state(self, p):
+        with torch.no_grad():
+            state = p.data.float().detach().cpu().numpy()
+        return state
+
+    def step(self):
+        with torch.no_grad():
+            self.steps += 1
+            if self.steps % self.freq == 0:
+                for i in range(len(self.state)):
+                    p, state = self.state[i]
+                    state = torch.from_numpy(state).cuda()
+                    state.mul_(self.mu).add_(1 - self.mu, p.data.float())
+                    self.state[i] = (p, state.cpu().numpy())
+
+    def swap(self):
+        with torch.no_grad():
+            # swap ema and model params
+            for p, state in self.state:
+                other_state = self.get_model_state(p)
+                p.data.copy_(torch.from_numpy(state).type_as(p.data))
+                np.copyto(state, other_state)
+
+class FusedEMA:
+    def __init__(self, params, mu=0.999):
+        self.mu = mu
+        params = list(params)
+        self.params = {}
+        self.params['fp16'] = [p for p in params if p.requires_grad and p.data.dtype == torch.float16]
+        self.params['fp32'] = [p for p in params if p.requires_grad and p.data.dtype != torch.float16]
+        self.groups = [group for group in self.params.keys() if len(self.params[group]) > 0]
+        self.state = {}
+        for group in self.groups:
+            self.state[group] = self.get_model_state(group)
+
+    def get_model_state(self, group):
+        params = self.params[group]
+        return _flatten_dense_tensors([p.data.float() for p in params])
+        # if self.fp16:
+        #     return _flatten_dense_tensors([p.data.half() for p in self.param_group if p.dtype])
+        # else:
+        #     return _flatten_dense_tensors([p.data for p in self.param_group])
+
+    def step(self):
+        for group in self.groups:
+            self.state[group].mul_(self.mu).add_(1 - self.mu, self.get_model_state(group))
+
+    def swap(self):
+        # swap ema and model params
+        for group in self.groups:
+            other_state = self.get_model_state(group)
+            state = self.state[group]
+            params = self.params[group]
+            offset = 0
+            for p in params:
+                numel = p.data.numel()
+                p.data = state.narrow(0, offset, numel).view_as(p.data).type_as(p.data)
+                offset += numel
+
+            self.state[group] = other_state
+
+

jukebox/utils/fp16.py

@@ -0,0 +1,303 @@
+# Utils for fp16 training.
+import importlib
+import math
+import numpy as np
+import torch
+import jukebox.utils.dist_adapter as dist
+from torch.optim import Optimizer
+from torch._utils import _flatten_dense_tensors
+
+from jukebox.utils.dist_utils import allreduce
+
+def adam_step(p: torch.Tensor, out_p: torch.Tensor, exp_avg: torch.Tensor, exp_avg_sq: torch.Tensor, grad: torch.Tensor,
+              lr: float, beta1: float, beta2: float, eps: float, scale: float, step: int, eps_mode: int, bias_correction: int, weight_decay: float):
+    assert bias_correction == 1
+    assert eps_mode == 1
+
+    grad = grad.float()
+    grad.div_(scale)
+
+    # Decay the first and second moment running average coefficient
+    exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+    exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+    denom = exp_avg_sq.sqrt().add_(eps)
+
+    bias_correction1 = 1 - beta1 ** step
+    bias_correction2 = 1 - beta2 ** step
+    step_size = lr * math.sqrt(bias_correction2) / bias_correction1
+
+    p.add_(exp_avg/denom + weight_decay*p.float(), alpha=-step_size)
+
+# Import fused_adam if we have apex, otherwise use regular adam
+try:
+    fused_adam_cuda = importlib.import_module("fused_adam_cuda")
+    fused_adam_step = fused_adam_cuda.adam
+    print("Using apex fused_adam_cuda")
+except ModuleNotFoundError:
+    fused_adam_step = adam_step
+
+def backward(loss, params, scalar, fp16, logger):
+    # Perform backward
+    if not fp16:
+        scale = 1.0
+        loss.backward()
+        gn = grad_norm(params, scale)
+        return loss, scale, gn, False, False
+    else:
+        scale = scalar.get_scale()
+        loss = (loss.float())*scale
+        overflow_loss = check_overflow(loss.item())
+        overflow_loss = allreduce(int(overflow_loss), op=dist.ReduceOp.MAX) > 0
+        if not overflow_loss:
+            loss.backward()
+            gn = grad_norm(params, scale)
+            overflow_grad = check_overflow(gn)
+            overflow_grad = allreduce(int(overflow_grad), op=dist.ReduceOp.MAX) > 0
+            scalar.update_scale(overflow_grad)
+        else:
+            gn = 0.0
+            overflow_grad = True
+        loss = (loss.detach().float()) / scale # Should delete computation graph for overflow
+        if logger.rank == 0:
+            if loss > 12.: print(f"\nWarning. Loss is {loss}")
+            if overflow_loss: print(f"\nOverflow in forward. Loss {loss}, lgscale {np.log2(scale)}. Skipping batch completely (no backward, scale update)")
+            elif overflow_grad: print(f"\nOverflow in backward. Loss {loss}, grad norm {gn}, lgscale {np.log2(scale)}, new lgscale {np.log2(scalar.get_scale())}")
+        return loss, scale, gn, overflow_loss, overflow_grad
+
+# Automatic loss scaling
+class LossScalar(object):
+    def __init__(self,
+                 loss_scale,
+                 init_scale=2. ** 16,
+                 scale_factor=2. ** (1. / 1000),
+                 scale_window=1):
+        if loss_scale == None:
+            # Use dynamic loss scaling
+            self.dynamic = True
+            self.loss_scale = init_scale
+        else:
+            self.dynamic = False
+            self.loss_scale = loss_scale
+        self.max_loss_scale = 2.**24
+        self.scale_factor = scale_factor
+        self.scale_window  = scale_window
+        self.unskipped = 0
+        self.overflow = False
+
+    def get_scale(self):
+        return self.loss_scale
+
+    def update_scale(self, overflow):
+        if overflow and self.dynamic:
+            self.loss_scale /= 2.
+            self.unskipped = 0
+        else:
+            self.unskipped += 1
+
+        if self.unskipped == self.scale_window and self.dynamic:
+            self.loss_scale = min(self.max_loss_scale, self.loss_scale * self.scale_factor)
+            self.unskipped = 0
+
+def check_overflow(val):
+    return (val == float('inf')) or (val == -float('inf')) or (val != val)
+
+def grad_norm(params, scale, flat=False):
+    params = list(params)
+    if flat:
+        # Faster but more memory
+        fp16_grads = [p.grad for p in params if p.grad is not None and p.data.dtype == torch.float16]
+        fp16_norm = 0.0 if len(fp16_grads) == 0 else float(_flatten_dense_tensors(fp16_grads).norm(p=2, dtype=torch.float32))
+        fp32_grads = [p.grad for p in params if p.grad is not None and p.data.dtype != torch.float16]
+        fp32_norm = 0.0 if len(fp32_grads) == 0 else float(_flatten_dense_tensors(fp32_grads).norm(p=2))
+        grad_norm = (fp16_norm**2 + fp32_norm**2)**0.5
+    else:
+        # Slightly slower but less memory
+        grad_norm = 0.0
+        for p in params:
+            if p.grad is not None:
+                grad_norm += p.grad.norm(p=2, dtype=torch.float32)**2
+        grad_norm = float(grad_norm**0.5)
+    return grad_norm / scale
+
+def clipped_grad_scale(grad_norm, max_grad_norm, scale):
+    clip = grad_norm / max_grad_norm
+    if clip > 1:
+        scale = clip * scale
+    return scale
+
+class FP16FusedAdam(Optimizer):
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        bias_correction=True,
+        betas=(0.9, 0.999),
+        eps=1e-8,
+        eps_inside_sqrt=False,
+        weight_decay=0.0,
+        amsgrad=False,
+    ):
+        if amsgrad:
+            raise RuntimeError("FusedAdam does not support the AMSGrad variant.")
+        defaults = dict(
+            lr=lr, bias_correction=bias_correction, betas=betas, eps=eps, weight_decay=weight_decay
+        )
+        super(FP16FusedAdam, self).__init__(params, defaults)
+        self.eps_mode = 0 if eps_inside_sqrt else 1
+        self.FLOAT16_MAX = 65504.0
+        self.init_state()
+
+    def init_state(self):
+        for group in self.param_groups:
+            for p in group["params"]:
+                assert p.requires_grad == True
+                state = self.state[p]
+                if len(state) == 0:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros_like(p.data)
+                    if p.data.dtype == torch.float16:
+                        state["scale_exp_avg"] = 1.0
+                        state["scale_exp_avg_sq"] = 1.0
+
+    def step(self, closure=None, scale=1.0):
+        """Performs a single optimization step. Scales gradients down by scale
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+            scale (float, optional): factor to divide gradient tensor values
+                by before applying to weights. (default: 1)
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            bias_correction = 1 if group["bias_correction"] else 0
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+
+                state = self.state[p]
+
+                if p.data.dtype == torch.float16:
+                    exp_avg, exp_avg_sq = (
+                        state["exp_avg"].float() * state["scale_exp_avg"],
+                        state["exp_avg_sq"].float() * state["scale_exp_avg_sq"],
+                    )
+                else:
+                    exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+
+                out_p = torch.tensor([], dtype=torch.float)
+                fused_adam_step(
+                    p.data,
+                    out_p,
+                    exp_avg,
+                    exp_avg_sq,
+                    grad,
+                    group["lr"],
+                    beta1,
+                    beta2,
+                    group["eps"],
+                    scale,
+                    state["step"],
+                    self.eps_mode,
+                    bias_correction,
+                    group["weight_decay"],
+                )
+
+                if p.data.dtype == torch.float16:
+                    state["scale_exp_avg"] = (
+                        1e-8 + float(torch.norm(exp_avg, float("inf"))) / self.FLOAT16_MAX
+                    )
+                    state["scale_exp_avg_sq"] = (
+                        1e-8 + float(torch.norm(exp_avg_sq, float("inf"))) / self.FLOAT16_MAX
+                    )
+                    state["exp_avg"] = (exp_avg / state["scale_exp_avg"]).half()
+                    state["exp_avg_sq"] = (exp_avg_sq / state["scale_exp_avg_sq"]).half()
+
+        return loss
+
+
+class FusedAdam(Optimizer):
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        bias_correction=True,
+        betas=(0.9, 0.999),
+        eps=1e-8,
+        eps_inside_sqrt=False,
+        weight_decay=0.0,
+        amsgrad=False,
+    ):
+        if amsgrad:
+            raise RuntimeError("FusedAdam does not support the AMSGrad variant.")
+        defaults = dict(
+            lr=lr, bias_correction=bias_correction, betas=betas, eps=eps, weight_decay=weight_decay
+        )
+        super(FusedAdam, self).__init__(params, defaults)
+        self.eps_mode = 0 if eps_inside_sqrt else 1
+
+    def step(self, closure=None, scale=1.0):
+        """Performs a single optimization step. Scales gradients down by scale
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+            scale (float, optional): factor to divide gradient tensor values
+                by before applying to weights. (default: 1)
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            bias_correction = 1 if group["bias_correction"] else 0
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p.data).float()
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros_like(p.data).float()
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+
+                out_p = torch.tensor([], dtype=torch.float)
+                fused_adam_step(
+                    p.data,
+                    out_p,
+                    exp_avg,
+                    exp_avg_sq,
+                    grad,
+                    group["lr"],
+                    beta1,
+                    beta2,
+                    group["eps"],
+                    scale,
+                    state["step"],
+                    self.eps_mode,
+                    bias_correction,
+                    group["weight_decay"],
+                )
+
+        return loss
+

jukebox/utils/io.py

@@ -0,0 +1,136 @@
+import numpy as np
+import av
+import torch as t
+import jukebox.utils.dist_adapter as dist
+
+def get_duration_sec(file, cache=False):
+    try:
+        with open(file + '.dur', 'r') as f:
+            duration = float(f.readline().strip('\n'))
+        return duration
+    except:
+        container = av.open(file)
+        audio = container.streams.get(audio=0)[0]
+        duration = audio.duration * float(audio.time_base)
+        if cache:
+            with open(file + '.dur', 'w') as f:
+                f.write(str(duration) + '\n')
+        return duration
+
+def load_audio(file, sr, offset, duration, resample=True, approx=False, time_base='samples', check_duration=True):
+    if time_base == 'sec':
+        offset = offset * sr
+        duration = duration * sr
+    # Loads at target sr, stereo channels, seeks from offset, and stops after duration
+    container = av.open(file)
+    audio = container.streams.get(audio=0)[0] # Only first audio stream
+    audio_duration = audio.duration * float(audio.time_base)
+    if approx:
+        if offset + duration > audio_duration*sr:
+            # Move back one window. Cap at audio_duration
+            offset = np.min(audio_duration*sr - duration, offset - duration)
+    else:
+        if check_duration:
+            assert offset + duration <= audio_duration*sr, f'End {offset + duration} beyond duration {audio_duration*sr}'
+    if resample:
+        resampler = av.AudioResampler(format='fltp',layout='stereo', rate=sr)
+    else:
+        assert sr == audio.sample_rate
+    offset = int(offset / sr / float(audio.time_base)) #int(offset / float(audio.time_base)) # Use units of time_base for seeking
+    duration = int(duration) #duration = int(duration * sr) # Use units of time_out ie 1/sr for returning
+    sig = np.zeros((2, duration), dtype=np.float32)
+    container.seek(offset, stream=audio)
+    total_read = 0
+    for frame in container.decode(audio=0): # Only first audio stream
+        if resample:
+            frame.pts = None
+            frame = resampler.resample(frame)
+        frame = frame.to_ndarray(format='fltp') # Convert to floats and not int16
+        read = frame.shape[-1]
+        if total_read + read > duration:
+            read = duration - total_read
+        sig[:, total_read:total_read + read] = frame[:, :read]
+        total_read += read
+        if total_read == duration:
+            break
+    assert total_read <= duration, f'Expected {duration} frames, got {total_read}'
+    return sig, sr
+
+def test_simple_loader():
+    import librosa
+    from tqdm import tqdm
+
+    collate_fn = lambda batch: t.stack([t.from_numpy(b) for b in batch], dim=0)
+
+    def get_batch(file, loader):
+        y1, sr = loader(file, sr=44100, offset=0.0, duration=6.0, time_base='sec')
+        y2, sr = loader(file, sr=44100, offset=20.0, duration=6.0, time_base='sec')
+        return [y1, y2]
+
+    def load(file, loader):
+        batch = get_batch(file, loader)  # np
+        x = collate_fn(batch)  # torch cpu
+        x = x.to('cuda', non_blocking=True)  # torch gpu
+        return x
+
+    files = librosa.util.find_files('/root/data/', ['mp3', 'm4a', 'opus'])
+    print(files[:10])
+    loader = load_audio
+    print("Loader", loader.__name__)
+    x = t.randn(2, 2).cuda()
+    x = load(files[0], loader)
+    for i,file in enumerate(tqdm(files)):
+        x = load(file, loader)
+        if i == 100:
+            break
+
+def test_dataset_loader():
+    from tqdm import tqdm
+    from torch.utils.data import DataLoader
+    from torch.utils.data.distributed import DistributedSampler
+    from jukebox.utils.audio_utils import audio_preprocess, audio_postprocess
+    from jukebox.hparams import setup_hparams
+    from jukebox.data.files_dataset import FilesAudioDataset
+    hps = setup_hparams("teeny", {})
+    hps.sr = 22050  # 44100
+    hps.hop_length = 512
+    hps.labels = False
+    hps.channels = 2
+    hps.aug_shift = False
+    hps.bs = 2
+    hps.nworkers = 2 # Getting 20 it/s with 2 workers, 10 it/s with 1 worker
+    print(hps)
+    dataset = hps.dataset
+    root = hps.root
+    from tensorboardX import SummaryWriter
+    sr = {22050: '22k', 44100: '44k', 48000: '48k'}[hps.sr]
+    writer = SummaryWriter(f'{root}/{dataset}/logs/{sr}/logs')
+    dataset = FilesAudioDataset(hps)
+    print("Length of dataset", len(dataset))
+
+    # Torch Loader
+    collate_fn = lambda batch: t.stack([t.from_numpy(b) for b in batch], 0)
+    sampler = DistributedSampler(dataset)
+    train_loader = DataLoader(dataset, batch_size=hps.bs, num_workers=hps.nworkers, pin_memory=False, sampler=sampler,
+                              drop_last=True, collate_fn=collate_fn)
+
+    dist.barrier()
+    sampler.set_epoch(0)
+    for i, x in enumerate(tqdm(train_loader)):
+        x = x.to('cuda', non_blocking=True)
+        for j, aud in enumerate(x):
+            writer.add_audio('in_' + str(i*hps.bs + j), aud, 1, hps.sr)
+        print("Wrote in")
+        x = audio_preprocess(x, hps)
+        x = audio_postprocess(x, hps)
+        for j, aud in enumerate(x):
+            writer.add_audio('out_' + str(i*hps.bs + j), aud, 1, hps.sr)
+        print("Wrote out")
+        dist.barrier()
+        break
+
+if __name__ == '__main__':
+    from jukebox.utils.dist_utils import setup_dist_from_mpi
+    setup_dist_from_mpi(port=29500)
+    test_dataset_loader()
+

jukebox/utils/logger.py

@@ -0,0 +1,147 @@
+import torch as t
+import jukebox.utils.dist_adapter as dist
+from tqdm import tqdm
+from datetime import date
+import os
+import sys
+
+def def_tqdm(x):
+    return tqdm(x, leave=True, file=sys.stdout, bar_format="{n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]")
+
+def get_range(x):
+    if dist.get_rank() == 0:
+        return def_tqdm(x)
+    else:
+        return x
+
+def init_logging(hps, local_rank, rank):
+    logdir = f"{hps.local_logdir}/{hps.name}"
+    if local_rank == 0:
+        if not os.path.exists(logdir):
+            os.makedirs(logdir)
+        with open(logdir + 'argv.txt', 'w') as f:
+            f.write(hps.argv + '\n')
+        print("Logging to", logdir)
+    logger = Logger(logdir, rank)
+    metrics = Metrics()
+    logger.add_text('hps', str(hps))
+    return logger, metrics
+
+def get_name(hps):
+    name = ""
+    for key, value in hps.items():
+        name += f"{key}_{value}_"
+    return name
+
+def average_metrics(_metrics):
+    metrics = {}
+    for _metric in _metrics:
+        for key, val in _metric.items():
+            if key not in metrics:
+                metrics[key] = []
+            metrics[key].append(val)
+    return {key: sum(vals)/len(vals) for key, vals in metrics.items()}
+
+class Metrics:
+    def __init__(self):
+        self.sum = {}
+        self.n = {}
+
+    def update(self, tag, val, batch):
+        # v is average value over batch
+        # store total value and total batch, returns dist average
+        sum = t.tensor(val * batch).float().cuda()
+        n = t.tensor(batch).float().cuda()
+        dist.all_reduce(sum)
+        dist.all_reduce(n)
+        sum = sum.item()
+        n = n.item()
+        self.sum[tag] = self.sum.get(tag, 0.0) + sum
+        self.n[tag] = self.n.get(tag, 0.0) + n
+        return sum / n
+
+    def avg(self, tag):
+        if tag in self.sum:
+            return self.sum[tag] / self.n[tag]
+        else:
+            return 0.0
+
+    def reset(self):
+        self.sum = {}
+        self.n = {}
+
+class Logger:
+    def __init__(self, logdir, rank):
+        if rank == 0:
+            from tensorboardX import SummaryWriter
+            self.sw = SummaryWriter(f"{logdir}/logs")
+        self.iters = 0
+        self.rank = rank
+        self.works = []
+        self.logdir = logdir
+
+    def step(self):
+        self.iters += 1
+
+    def flush(self):
+        if self.rank == 0:
+            self.sw.flush()
+
+    def add_text(self, tag, text):
+        if self.rank == 0:
+            self.sw.add_text(tag, text, self.iters)
+
+    def add_audios(self, tag, auds, sample_rate=22050, max_len=None, max_log=8):
+        if self.rank == 0:
+            for i in range(min(len(auds), max_log)):
+                if max_len:
+                    self.sw.add_audio(f"{i}/{tag}", auds[i][:max_len * sample_rate], self.iters, sample_rate)
+                else:
+                    self.sw.add_audio(f"{i}/{tag}", auds[i], self.iters, sample_rate)
+
+    def add_audio(self, tag, aud, sample_rate=22050):
+        if self.rank == 0:
+            self.sw.add_audio(tag, aud, self.iters, sample_rate)
+
+    def add_images(self, tag, img, dataformats="NHWC"):
+        if self.rank == 0:
+            self.sw.add_images(tag, img, self.iters, dataformats=dataformats)
+
+    def add_image(self, tag, img):
+        if self.rank == 0:
+            self.sw.add_image(tag, img, self.iters)
+
+    def add_scalar(self, tag, val):
+        if self.rank == 0:
+            self.sw.add_scalar(tag, val, self.iters)
+
+    def get_range(self, loader):
+        if self.rank == 0:
+            self.trange = def_tqdm(loader)
+        else:
+            self.trange = loader
+        return enumerate(self.trange)
+
+    def close_range(self):
+        if self.rank == 0:
+            self.trange.close()
+
+    def set_postfix(self, *args, **kwargs):
+        if self.rank == 0:
+            self.trange.set_postfix(*args, **kwargs)
+
+    # For logging summaries of varies graph ops
+    def add_reduce_scalar(self, tag, layer, val):
+        if self.iters % 100 == 0:
+            with t.no_grad():
+                val = val.float().norm()/float(val.numel())
+            work = dist.reduce(val, 0, async_op=True)
+            self.works.append((tag, layer, val, work))
+
+    def finish_reduce(self):
+        for tag, layer, val, work in self.works:
+            work.wait()
+            if self.rank == 0:
+                val = val.item()/dist.get_world_size()
+                self.lw[layer].add_scalar(tag, val, self.iters)
+        self.works = []

jukebox/utils/remote_utils.py

@@ -0,0 +1,42 @@
+import sys
+import subprocess
+
+def download(remote_path, local_path, async_download=False):
+    args = ['wget', '-O', local_path, remote_path]
+    print("Running ", " ".join(args))
+    if async_download:
+        subprocess.Popen(args)
+    else:
+        subprocess.call(args)
+
+# GCE
+def gs_download(gs_path, local_path, async_download=False):
+    args = ['gsutil',
+            '-o', 'GSUtil:parallel_thread_count=1',
+            '-o', 'GSUtil:sliced_object_download_max_components=8',
+            'cp', gs_path, local_path]
+    if async_download:
+        subprocess.Popen(args)
+    else:
+        subprocess.call(args)
+
+
+def gs_upload(local_path, gs_path, async_upload=False):
+    # NOTE: Download and upload have differ -o flags.
+    # We also use -n to prevent clobbering checkpoints by mistake
+    assert not local_path.startswith("gs://")
+    assert gs_path.startswith("gs://")
+    args = ['gsutil',
+            '-o', 'GSUtil:parallel_composite_upload_threshold=150M',
+            'cp', '-n', local_path, gs_path]
+    if async_upload:
+        subprocess.Popen(args)
+    else:
+        subprocess.call(args)
+
+def ls(regex):
+    outputs = subprocess.check_output(['gsutil', 'ls', regex]).decode(sys.stdout.encoding)
+    outputs = outputs.split('\n')
+    outputs = [output for output in outputs if output is not '']
+    return outputs
+

jukebox/utils/sample_utils.py

@@ -0,0 +1,22 @@
+import torch as t
+
+def split_batch(obj, n_samples, split_size):
+    n_passes = (n_samples + split_size - 1) // split_size
+    if isinstance(obj, t.Tensor):
+        return t.split(obj, split_size, dim=0)
+    elif isinstance(obj, list):
+        return list(zip(*[t.split(item, split_size, dim=0) for item in obj]))
+    elif obj is None:
+        return [None] * n_passes
+    else:
+        raise TypeError('Unknown input type')
+
+# Break total_length into hops/windows of size n_ctx separated by hop_length
+def get_starts(total_length, n_ctx, hop_length):
+    starts = []
+    for start in range(0, total_length - n_ctx + hop_length, hop_length):
+        if start + n_ctx >= total_length:
+            # Last hop could be smaller, we make it n_ctx to maximise context
+            start = total_length - n_ctx
+        starts.append(start)
+    return starts

jukebox/utils/torch_utils.py

@@ -0,0 +1,32 @@
+import gc
+import torch as t
+
+def freeze_model(model):
+    model.eval()
+    for params in model.parameters():
+        params.requires_grad = False
+
+
+def unfreeze_model(model):
+    model.train()
+    for params in model.parameters():
+        params.requires_grad = True
+
+def zero_grad(model):
+    for p in model.parameters():
+        if p.requires_grad and p.grad is not None:
+            p.grad = None
+
+def empty_cache():
+    gc.collect()
+    t.cuda.empty_cache()
+
+def assert_shape(x, exp_shape):
+    assert x.shape == exp_shape, f"Expected {exp_shape} got {x.shape}"
+
+def count_parameters(model):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+def count_state(model):
+    return sum(s.numel() for s in model.state_dict().values())
+

jukebox/vqvae/bottleneck.py

@@ -0,0 +1,248 @@
+import numpy as np
+import torch as t
+import torch.nn as nn
+import torch.nn.functional as F
+import jukebox.utils.dist_adapter as dist
+
+class BottleneckBlock(nn.Module):
+    def __init__(self, k_bins, emb_width, mu):
+        super().__init__()
+        self.k_bins = k_bins
+        self.emb_width = emb_width
+        self.mu = mu
+        self.reset_k()
+        self.threshold = 1.0
+
+    def reset_k(self):
+        self.init = False
+        self.k_sum = None
+        self.k_elem = None
+        self.register_buffer('k', t.zeros(self.k_bins, self.emb_width).cuda())
+
+    def _tile(self, x):
+        d, ew = x.shape
+        if d < self.k_bins:
+            n_repeats = (self.k_bins + d - 1) // d
+            std = 0.01 / np.sqrt(ew)
+            x = x.repeat(n_repeats, 1)
+            x = x + t.randn_like(x) * std
+        return x
+
+    def init_k(self, x):
+        mu, emb_width, k_bins = self.mu, self.emb_width, self.k_bins
+        self.init = True
+        # init k_w using random vectors from x
+        y = self._tile(x)
+        _k_rand = y[t.randperm(y.shape[0])][:k_bins]
+        dist.broadcast(_k_rand, 0)
+        self.k = _k_rand
+        assert self.k.shape == (k_bins, emb_width)
+        self.k_sum = self.k
+        self.k_elem = t.ones(k_bins, device=self.k.device)
+
+    def restore_k(self, num_tokens=None, threshold=1.0):
+        mu, emb_width, k_bins = self.mu, self.emb_width, self.k_bins
+        self.init = True
+        assert self.k.shape == (k_bins, emb_width)
+        self.k_sum = self.k.clone()
+        self.k_elem = t.ones(k_bins, device=self.k.device)
+        if num_tokens is not None:
+            expected_usage = num_tokens / k_bins
+            self.k_elem.data.mul_(expected_usage)
+            self.k_sum.data.mul_(expected_usage)
+        self.threshold = threshold
+
+    def update_k(self, x, x_l):
+        mu, emb_width, k_bins = self.mu, self.emb_width, self.k_bins
+        with t.no_grad():
+            # Calculate new centres
+            x_l_onehot = t.zeros(k_bins, x.shape[0], device=x.device)  # k_bins, N * L
+            x_l_onehot.scatter_(0, x_l.view(1, x.shape[0]), 1)
+
+            _k_sum = t.matmul(x_l_onehot, x)  # k_bins, w
+            _k_elem = x_l_onehot.sum(dim=-1)  # k_bins
+            y = self._tile(x)
+            _k_rand = y[t.randperm(y.shape[0])][:k_bins]
+
+            dist.broadcast(_k_rand, 0)
+            dist.all_reduce(_k_sum)
+            dist.all_reduce(_k_elem)
+
+            # Update centres
+            old_k = self.k
+            self.k_sum = mu * self.k_sum + (1. - mu) * _k_sum  # w, k_bins
+            self.k_elem = mu * self.k_elem + (1. - mu) * _k_elem  # k_bins
+            usage = (self.k_elem.view(k_bins, 1) >= self.threshold).float()
+            self.k = usage * (self.k_sum.view(k_bins, emb_width) / self.k_elem.view(k_bins, 1)) \
+                     + (1 - usage) * _k_rand
+            _k_prob = _k_elem / t.sum(_k_elem)  # x_l_onehot.mean(dim=-1)  # prob of each bin
+            entropy = -t.sum(_k_prob * t.log(_k_prob + 1e-8))  # entropy ie how diverse
+            used_curr = (_k_elem >= self.threshold).sum()
+            usage = t.sum(usage)
+            dk = t.norm(self.k - old_k) / np.sqrt(np.prod(old_k.shape))
+        return dict(entropy=entropy,
+                    used_curr=used_curr,
+                    usage=usage,
+                    dk=dk)
+
+    def preprocess(self, x):
+        # NCT -> NTC -> [NT, C]
+        x = x.permute(0, 2, 1).contiguous()
+        x = x.view(-1, x.shape[-1])  # x_en = (N * L, w), k_j = (w, k_bins)
+
+        if x.shape[-1] == self.emb_width:
+            prenorm = t.norm(x - t.mean(x)) / np.sqrt(np.prod(x.shape))
+        elif x.shape[-1] == 2 * self.emb_width:
+            x1, x2 = x[...,:self.emb_width], x[...,self.emb_width:]
+            prenorm = (t.norm(x1 - t.mean(x1)) / np.sqrt(np.prod(x1.shape))) + (t.norm(x2 - t.mean(x2)) / np.sqrt(np.prod(x2.shape)))
+
+            # Normalise
+            x = x1 + x2
+        else:
+            assert False, f"Expected {x.shape[-1]} to be (1 or 2) * {self.emb_width}"
+        return x, prenorm
+
+    def postprocess(self, x_l, x_d, x_shape):
+        # [NT, C] -> NTC -> NCT
+        N, T = x_shape
+        x_d = x_d.view(N, T, -1).permute(0, 2, 1).contiguous()
+        x_l = x_l.view(N, T)
+        return x_l, x_d
+
+    def quantise(self, x):
+        # Calculate latent code x_l
+        k_w = self.k.t()
+        distance = t.sum(x ** 2, dim=-1, keepdim=True) - 2 * t.matmul(x, k_w) + t.sum(k_w ** 2, dim=0,
+                                                                                            keepdim=True)  # (N * L, b)
+        min_distance, x_l = t.min(distance, dim=-1)
+        fit = t.mean(min_distance)
+        return x_l, fit
+
+    def dequantise(self, x_l):
+        x = F.embedding(x_l, self.k)
+        return x
+
+    def encode(self, x):
+        N, width, T = x.shape
+
+        # Preprocess.
+        x, prenorm = self.preprocess(x)
+
+        # Quantise
+        x_l, fit = self.quantise(x)
+
+        # Postprocess.
+        x_l = x_l.view(N, T)
+        return x_l
+
+    def decode(self, x_l):
+        N, T = x_l.shape
+        width = self.emb_width
+
+        # Dequantise
+        x_d = self.dequantise(x_l)
+
+        # Postprocess
+        x_d = x_d.view(N, T, width).permute(0, 2, 1).contiguous()
+        return x_d
+
+    def forward(self, x, update_k=True):
+        N, width, T = x.shape
+
+        # Preprocess
+        x, prenorm = self.preprocess(x)
+
+        # Init k if not inited
+        if update_k and not self.init:
+            self.init_k(x)
+
+        # Quantise and dequantise through bottleneck
+        x_l, fit = self.quantise(x)
+        x_d = self.dequantise(x_l)
+
+        # Update embeddings
+        if update_k:
+            update_metrics = self.update_k(x, x_l)
+        else:
+            update_metrics = {}
+
+        # Loss
+        commit_loss = t.norm(x_d.detach() - x) ** 2 / np.prod(x.shape)
+
+        # Passthrough
+        x_d = x + (x_d - x).detach()
+
+        # Postprocess
+        x_l, x_d = self.postprocess(x_l, x_d, (N,T))
+        return x_l, x_d, commit_loss, dict(fit=fit,
+                                           pn=prenorm,
+                                           **update_metrics)
+
+
+class Bottleneck(nn.Module):
+    def __init__(self, l_bins, emb_width, mu, levels):
+        super().__init__()
+        self.levels = levels
+        level_block = lambda level: BottleneckBlock(l_bins, emb_width, mu)
+        self.level_blocks = nn.ModuleList()
+        for level in range(self.levels):
+            self.level_blocks.append(level_block(level))
+
+    def encode(self, xs):
+        zs = [level_block.encode(x) for (level_block, x) in zip(self.level_blocks, xs)]
+        return zs
+
+    def decode(self, zs, start_level=0, end_level=None):
+        if end_level is None:
+            end_level = self.levels
+        xs_quantised = [level_block.decode(z) for (level_block, z) in zip(self.level_blocks[start_level:end_level], zs)]
+        return xs_quantised
+
+    def forward(self, xs):
+        zs, xs_quantised, commit_losses, metrics = [], [], [], []
+        for level in range(self.levels):
+            level_block = self.level_blocks[level]
+            x = xs[level]
+            z, x_quantised, commit_loss, metric = level_block(x, update_k=self.training)
+            zs.append(z)
+            if not self.training:
+                # Be extra paranoid and make sure the encoder weights can't
+                # change from straight-through estimator
+                x_quantised = x_quantised.detach()
+            xs_quantised.append(x_quantised)
+            commit_losses.append(commit_loss)
+            if self.training:
+                metrics.append(metric)
+        return zs, xs_quantised, commit_losses, metrics
+
+class NoBottleneckBlock(nn.Module):
+    def restore_k(self):
+        pass
+
+class NoBottleneck(nn.Module):
+    def __init__(self, levels):
+        super().__init__()
+        self.level_blocks = nn.ModuleList()
+        self.levels = levels
+        for level in range(levels):
+            self.level_blocks.append(NoBottleneckBlock())
+
+    def encode(self, xs):
+        return xs
+
+    def decode(self, zs, start_level=0, end_level=None):
+        if end_level is None:
+            end_level = self.levels
+        return zs
+
+    def forward(self, xs):
+        zero = t.zeros(()).cuda()
+        commit_losses = [zero for _ in range(self.levels)]
+        metrics = [dict(entropy=zero, usage=zero, used_curr=zero, pn=zero, dk=zero) for _ in range(self.levels)]
+        return xs, xs, commit_losses, metrics
+
+if __name__ == '__main__':
+    from jukebox.utils.dist_utils import setup_dist_from_mpi
+    rank, local_rank, device = setup_dist_from_mpi(port=29600)
+    bottleneck = Bottleneck(256, 64, 0.99, 2).to(device)
+    bottleneck.check()

jukebox/vqvae/encdec.py

@@ -0,0 +1,131 @@
+import torch as t
+import torch.nn as nn
+from jukebox.vqvae.resnet import Resnet, Resnet1D
+from jukebox.utils.torch_utils import assert_shape
+
+class EncoderConvBlock(nn.Module):
+    def __init__(self, input_emb_width, output_emb_width, down_t,
+                 stride_t, width, depth, m_conv,
+                 dilation_growth_rate=1, dilation_cycle=None, zero_out=False,
+                 res_scale=False):
+        super().__init__()
+        blocks = []
+        filter_t, pad_t = stride_t * 2, stride_t // 2
+        if down_t > 0:
+            for i in range(down_t):
+                block = nn.Sequential(
+                    nn.Conv1d(input_emb_width if i == 0 else width, width, filter_t, stride_t, pad_t),
+                    Resnet1D(width, depth, m_conv, dilation_growth_rate, dilation_cycle, zero_out, res_scale),
+                )
+                blocks.append(block)
+            block = nn.Conv1d(width, output_emb_width, 3, 1, 1)
+            blocks.append(block)
+        self.model = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.model(x)
+
+class DecoderConvBock(nn.Module):
+    def __init__(self, input_emb_width, output_emb_width, down_t,
+                 stride_t, width, depth, m_conv, dilation_growth_rate=1, dilation_cycle=None, zero_out=False, res_scale=False, reverse_decoder_dilation=False, checkpoint_res=False):
+        super().__init__()
+        blocks = []
+        if down_t > 0:
+            filter_t, pad_t = stride_t * 2, stride_t // 2
+            block = nn.Conv1d(output_emb_width, width, 3, 1, 1)
+            blocks.append(block)
+            for i in range(down_t):
+                block = nn.Sequential(
+                    Resnet1D(width, depth, m_conv, dilation_growth_rate, dilation_cycle, zero_out=zero_out, res_scale=res_scale, reverse_dilation=reverse_decoder_dilation, checkpoint_res=checkpoint_res),
+                    nn.ConvTranspose1d(width, input_emb_width if i == (down_t - 1) else width, filter_t, stride_t, pad_t)
+                )
+                blocks.append(block)
+        self.model = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.model(x)
+
+class Encoder(nn.Module):
+    def __init__(self, input_emb_width, output_emb_width, levels, downs_t,
+                 strides_t, **block_kwargs):
+        super().__init__()
+        self.input_emb_width = input_emb_width
+        self.output_emb_width = output_emb_width
+        self.levels = levels
+        self.downs_t = downs_t
+        self.strides_t = strides_t
+
+        block_kwargs_copy = dict(**block_kwargs)
+        if 'reverse_decoder_dilation' in block_kwargs_copy:
+            del block_kwargs_copy['reverse_decoder_dilation']
+        level_block = lambda level, down_t, stride_t: EncoderConvBlock(input_emb_width if level == 0 else output_emb_width,
+                                                           output_emb_width,
+                                                           down_t, stride_t,
+                                                           **block_kwargs_copy)
+        self.level_blocks = nn.ModuleList()
+        iterator = zip(list(range(self.levels)), downs_t, strides_t)
+        for level, down_t, stride_t in iterator:
+            self.level_blocks.append(level_block(level, down_t, stride_t))
+
+    def forward(self, x):
+        N, T = x.shape[0], x.shape[-1]
+        emb = self.input_emb_width
+        assert_shape(x, (N, emb, T))
+        xs = []
+
+        # 64, 32, ...
+        iterator = zip(list(range(self.levels)), self.downs_t, self.strides_t)
+        for level, down_t, stride_t in iterator:
+            level_block = self.level_blocks[level]
+            x = level_block(x)
+            emb, T = self.output_emb_width, T // (stride_t ** down_t)
+            assert_shape(x, (N, emb, T))
+            xs.append(x)
+
+        return xs
+
+class Decoder(nn.Module):
+    def __init__(self, input_emb_width, output_emb_width, levels, downs_t,
+                 strides_t, **block_kwargs):
+        super().__init__()
+        self.input_emb_width = input_emb_width
+        self.output_emb_width = output_emb_width
+        self.levels = levels
+
+        self.downs_t = downs_t
+
+        self.strides_t = strides_t
+
+        level_block = lambda level, down_t, stride_t: DecoderConvBock(output_emb_width,
+                                                          output_emb_width,
+                                                          down_t, stride_t,
+                                                          **block_kwargs)
+        self.level_blocks = nn.ModuleList()
+        iterator = zip(list(range(self.levels)), downs_t, strides_t)
+        for level, down_t, stride_t in iterator:
+            self.level_blocks.append(level_block(level, down_t, stride_t))
+
+        self.out = nn.Conv1d(output_emb_width, input_emb_width, 3, 1, 1)
+
+    def forward(self, xs, all_levels=True):
+        if all_levels:
+            assert len(xs) == self.levels
+        else:
+            assert len(xs) == 1
+        x = xs[-1]
+        N, T = x.shape[0], x.shape[-1]
+        emb = self.output_emb_width
+        assert_shape(x, (N, emb, T))
+
+        # 32, 64 ...
+        iterator = reversed(list(zip(list(range(self.levels)), self.downs_t, self.strides_t)))
+        for level, down_t, stride_t in iterator:
+            level_block = self.level_blocks[level]
+            x = level_block(x)
+            emb, T = self.output_emb_width, T * (stride_t ** down_t)
+            assert_shape(x, (N, emb, T))
+            if level != 0 and all_levels:
+                x = x + xs[level - 1]
+
+        x = self.out(x)
+        return x

jukebox/vqvae/resnet.py

@@ -0,0 +1,75 @@
+import math
+import torch.nn as nn
+import jukebox.utils.dist_adapter as dist
+from jukebox.utils.checkpoint import checkpoint
+
+class ResConvBlock(nn.Module):
+    def __init__(self, n_in, n_state):
+        super().__init__()
+        self.model = nn.Sequential(
+            nn.ReLU(),
+            nn.Conv2d(n_in, n_state, 3, 1, 1),
+            nn.ReLU(),
+            nn.Conv2d(n_state, n_in, 1, 1, 0),
+        )
+
+    def forward(self, x):
+        return x + self.model(x)
+
+class Resnet(nn.Module):
+    def __init__(self, n_in, n_depth, m_conv=1.0):
+        super().__init__()
+        self.model = nn.Sequential(*[ResConvBlock(n_in, int(m_conv * n_in)) for _ in range(n_depth)])
+
+    def forward(self, x):
+        return self.model(x)
+
+class ResConv1DBlock(nn.Module):
+    def __init__(self, n_in, n_state, dilation=1, zero_out=False, res_scale=1.0):
+        super().__init__()
+        padding = dilation
+        self.model = nn.Sequential(
+            nn.ReLU(),
+            nn.Conv1d(n_in, n_state, 3, 1, padding, dilation),
+            nn.ReLU(),
+            nn.Conv1d(n_state, n_in, 1, 1, 0),
+        )
+        if zero_out:
+            out = self.model[-1]
+            nn.init.zeros_(out.weight)
+            nn.init.zeros_(out.bias)
+        self.res_scale = res_scale
+
+    def forward(self, x):
+        return x + self.res_scale * self.model(x)
+
+class Resnet1D(nn.Module):
+    def __init__(self, n_in, n_depth, m_conv=1.0, dilation_growth_rate=1, dilation_cycle=None, zero_out=False, res_scale=False, reverse_dilation=False, checkpoint_res=False):
+        super().__init__()
+        def _get_depth(depth):
+            if dilation_cycle is None:
+                return depth
+            else:
+                return depth % dilation_cycle
+        blocks = [ResConv1DBlock(n_in, int(m_conv * n_in),
+                                 dilation=dilation_growth_rate ** _get_depth(depth),
+                                 zero_out=zero_out,
+                                 res_scale=1.0 if not res_scale else 1.0 / math.sqrt(n_depth))
+                  for depth in range(n_depth)]
+        if reverse_dilation:
+            blocks = blocks[::-1]
+        self.checkpoint_res = checkpoint_res
+        if self.checkpoint_res == 1:
+            if dist.get_rank() == 0:
+                print("Checkpointing convs")
+            self.blocks = nn.ModuleList(blocks)
+        else:
+            self.model = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        if self.checkpoint_res == 1:
+            for block in self.blocks:
+                x = checkpoint(block, (x, ), block.parameters(), True)
+            return x
+        else:
+            return self.model(x)

jukebox/vqvae/vqvae.py

@@ -0,0 +1,228 @@
+import numpy as np
+import torch as t
+import torch.nn as nn
+
+from jukebox.vqvae.encdec import Encoder, Decoder, assert_shape
+from jukebox.vqvae.bottleneck import NoBottleneck, Bottleneck
+from jukebox.utils.logger import average_metrics
+from jukebox.utils.audio_utils import spectral_convergence, spectral_loss, multispectral_loss, audio_postprocess
+
+def dont_update(params):
+    for param in params:
+        param.requires_grad = False
+
+def update(params):
+    for param in params:
+        param.requires_grad = True
+
+def calculate_strides(strides, downs):
+    return [stride ** down for stride, down in zip(strides, downs)]
+
+def _loss_fn(loss_fn, x_target, x_pred, hps):
+    if loss_fn == 'l1':
+        return t.mean(t.abs(x_pred - x_target)) / hps.bandwidth['l1']
+    elif loss_fn == 'l2':
+        return t.mean((x_pred - x_target) ** 2) / hps.bandwidth['l2']
+    elif loss_fn == 'linf':
+        residual = ((x_pred - x_target) ** 2).reshape(x_target.shape[0], -1)
+        values, _ = t.topk(residual, hps.linf_k, dim=1)
+        return t.mean(values) / hps.bandwidth['l2']
+    elif loss_fn == 'lmix':
+        loss = 0.0
+        if hps.lmix_l1:
+            loss += hps.lmix_l1 * _loss_fn('l1', x_target, x_pred, hps)
+        if hps.lmix_l2:
+            loss += hps.lmix_l2 * _loss_fn('l2', x_target, x_pred, hps)
+        if hps.lmix_linf:
+            loss += hps.lmix_linf * _loss_fn('linf', x_target, x_pred, hps)
+        return loss
+    else:
+        assert False, f"Unknown loss_fn {loss_fn}"
+
+class VQVAE(nn.Module):
+    def __init__(self, input_shape, levels, downs_t, strides_t,
+                 emb_width, l_bins, mu, commit, spectral, multispectral,
+                 multipliers=None, use_bottleneck=True, **block_kwargs):
+        super().__init__()
+
+        self.sample_length = input_shape[0]
+        x_shape, x_channels = input_shape[:-1], input_shape[-1]
+        self.x_shape = x_shape
+
+        self.downsamples = calculate_strides(strides_t, downs_t)
+        self.hop_lengths = np.cumprod(self.downsamples)
+        self.z_shapes = z_shapes = [(x_shape[0] // self.hop_lengths[level],) for level in range(levels)]
+        self.levels = levels
+
+        if multipliers is None:
+            self.multipliers = [1] * levels
+        else:
+            assert len(multipliers) == levels, "Invalid number of multipliers"
+            self.multipliers = multipliers
+        def _block_kwargs(level):
+            this_block_kwargs = dict(block_kwargs)
+            this_block_kwargs["width"] *= self.multipliers[level]
+            this_block_kwargs["depth"] *= self.multipliers[level]
+            return this_block_kwargs
+
+        encoder = lambda level: Encoder(x_channels, emb_width, level + 1,
+                                        downs_t[:level+1], strides_t[:level+1], **_block_kwargs(level))
+        decoder = lambda level: Decoder(x_channels, emb_width, level + 1,
+                                        downs_t[:level+1], strides_t[:level+1], **_block_kwargs(level))
+        self.encoders = nn.ModuleList()
+        self.decoders = nn.ModuleList()
+        for level in range(levels):
+            self.encoders.append(encoder(level))
+            self.decoders.append(decoder(level))
+
+        if use_bottleneck:
+            self.bottleneck = Bottleneck(l_bins, emb_width, mu, levels)
+        else:
+            self.bottleneck = NoBottleneck(levels)
+
+        self.downs_t = downs_t
+        self.strides_t = strides_t
+        self.l_bins = l_bins
+        self.commit = commit
+        self.spectral = spectral
+        self.multispectral = multispectral
+
+    def preprocess(self, x):
+        # x: NTC [-1,1] -> NCT [-1,1]
+        assert len(x.shape) == 3
+        x = x.permute(0,2,1).float()
+        return x
+
+    def postprocess(self, x):
+        # x: NTC [-1,1] <- NCT [-1,1]
+        x = x.permute(0,2,1)
+        return x
+
+    def _decode(self, zs, start_level=0, end_level=None):
+        # Decode
+        if end_level is None:
+            end_level = self.levels
+        assert len(zs) == end_level - start_level
+        xs_quantised = self.bottleneck.decode(zs, start_level=start_level, end_level=end_level)
+        assert len(xs_quantised) == end_level - start_level
+
+        # Use only lowest level
+        decoder, x_quantised = self.decoders[start_level], xs_quantised[0:1]
+        x_out = decoder(x_quantised, all_levels=False)
+        x_out = self.postprocess(x_out)
+        return x_out
+
+    def decode(self, zs, start_level=0, end_level=None, bs_chunks=1):
+        z_chunks = [t.chunk(z, bs_chunks, dim=0) for z in zs]
+        x_outs = []
+        for i in range(bs_chunks):
+            zs_i = [z_chunk[i] for z_chunk in z_chunks]
+            x_out = self._decode(zs_i, start_level=start_level, end_level=end_level)
+            x_outs.append(x_out)
+        return t.cat(x_outs, dim=0)
+
+    def _encode(self, x, start_level=0, end_level=None):
+        # Encode
+        if end_level is None:
+            end_level = self.levels
+        x_in = self.preprocess(x)
+        xs = []
+        for level in range(self.levels):
+            encoder = self.encoders[level]
+            x_out = encoder(x_in)
+            xs.append(x_out[-1])
+        zs = self.bottleneck.encode(xs)
+        return zs[start_level:end_level]
+
+    def encode(self, x, start_level=0, end_level=None, bs_chunks=1):
+        x_chunks = t.chunk(x, bs_chunks, dim=0)
+        zs_list = []
+        for x_i in x_chunks:
+            zs_i = self._encode(x_i, start_level=start_level, end_level=end_level)
+            zs_list.append(zs_i)
+        zs = [t.cat(zs_level_list, dim=0) for zs_level_list in zip(*zs_list)]
+        return zs
+
+    def sample(self, n_samples):
+        zs = [t.randint(0, self.l_bins, size=(n_samples, *z_shape), device='cuda') for z_shape in self.z_shapes]
+        return self.decode(zs)
+
+    def forward(self, x, hps, loss_fn='l1'):
+        metrics = {}
+
+        N = x.shape[0]
+
+        # Encode/Decode
+        x_in = self.preprocess(x)
+        xs = []
+        for level in range(self.levels):
+            encoder = self.encoders[level]
+            x_out = encoder(x_in)
+            xs.append(x_out[-1])
+
+        zs, xs_quantised, commit_losses, quantiser_metrics = self.bottleneck(xs)
+        x_outs = []
+        for level in range(self.levels):
+            decoder = self.decoders[level]
+            x_out = decoder(xs_quantised[level:level+1], all_levels=False)
+            assert_shape(x_out, x_in.shape)
+            x_outs.append(x_out)
+
+        # Loss
+        def _spectral_loss(x_target, x_out, hps):
+            if hps.use_nonrelative_specloss:
+                sl = spectral_loss(x_target, x_out, hps) / hps.bandwidth['spec']
+            else:
+                sl = spectral_convergence(x_target, x_out, hps)
+            sl = t.mean(sl)
+            return sl
+
+        def _multispectral_loss(x_target, x_out, hps):
+            sl = multispectral_loss(x_target, x_out, hps) / hps.bandwidth['spec']
+            sl = t.mean(sl)
+            return sl
+
+        recons_loss = t.zeros(()).to(x.device)
+        spec_loss = t.zeros(()).to(x.device)
+        multispec_loss = t.zeros(()).to(x.device)
+        x_target = audio_postprocess(x.float(), hps)
+
+        for level in reversed(range(self.levels)):
+            x_out = self.postprocess(x_outs[level])
+            x_out = audio_postprocess(x_out, hps)
+            this_recons_loss = _loss_fn(loss_fn, x_target, x_out, hps)
+            this_spec_loss = _spectral_loss(x_target, x_out, hps)
+            this_multispec_loss = _multispectral_loss(x_target, x_out, hps)
+            metrics[f'recons_loss_l{level + 1}'] = this_recons_loss
+            metrics[f'spectral_loss_l{level + 1}'] = this_spec_loss
+            metrics[f'multispectral_loss_l{level + 1}'] = this_multispec_loss
+            recons_loss += this_recons_loss
+            spec_loss += this_spec_loss
+            multispec_loss += this_multispec_loss
+
+        commit_loss = sum(commit_losses)
+        loss = recons_loss + self.spectral * spec_loss + self.multispectral * multispec_loss + self.commit * commit_loss
+
+        with t.no_grad():
+            sc = t.mean(spectral_convergence(x_target, x_out, hps))
+            l2_loss = _loss_fn("l2", x_target, x_out, hps)
+            l1_loss = _loss_fn("l1", x_target, x_out, hps)
+            linf_loss = _loss_fn("linf", x_target, x_out, hps)
+
+        quantiser_metrics = average_metrics(quantiser_metrics)
+
+        metrics.update(dict(
+            recons_loss=recons_loss,
+            spectral_loss=spec_loss,
+            multispectral_loss=multispec_loss,
+            spectral_convergence=sc,
+            l2_loss=l2_loss,
+            l1_loss=l1_loss,
+            linf_loss=linf_loss,
+            commit_loss=commit_loss,
+            **quantiser_metrics))
+
+        for key, val in metrics.items():
+            metrics[key] = val.detach()
+
+        return x_out, loss, metrics

requirements.txt

@@ -1,5 +1,6 @@
 git+https://github.com/TEAMuP-dev/pyharp.git@v0.3.0
 # model-specific deps below:
+setuptools<82  # madmom imports pkg_resources at runtime; removed in setuptools 82+
 torch
 torchaudio
 transformers
@@ -7,6 +8,17 @@ miditoolkit
 questionary
 soundfile
 mpi4py
-sheetsage @ git+https://github.com/tanchihpin0517/PiCoGen-sheetsage.git
 beat_this @ https://github.com/CPJKU/beat_this/archive/main.zip
-madmom @ git+https://github.com/CPJKU/madmom.git@0551aa8f48d71a367d92b5d3a347a0cf7cd97cc9
+madmom @ git+https://github.com/CPJKU/madmom.git@0551aa8f48d71a367d92b5d3a347a0cf7cd97cc9
+# sheetsage/ and jukebox/ are vendored directly via this repo since jukebox's own 
+# packaging is broken upstream. Here are their runtime deps:
+numpy<2
+scipy
+pretty-midi
+validators
+pillow
+fire
+unidecode
+numba
+librosa
+resampy

sheetsage/__init__.py

@@ -0,0 +1,17 @@
+import pathlib
+from os import environ as os_env
+
+LIB_DIR = pathlib.Path(__file__).resolve().parent
+
+if "SHEETSAGE_CACHE_DIR" in os_env:
+    CACHE_DIR = pathlib.Path(os_env["SHEETSAGE_CACHE_DIR"])
+else:
+    CACHE_DIR = pathlib.Path(pathlib.Path.home(), ".sheetsage")
+CACHE_DIR = CACHE_DIR.resolve()
+
+
+# NOTE: This changes the test discovery pattern from "test*.py" (default) to "*test.py".
+def load_tests(loader, standard_tests, pattern):
+    package_tests = loader.discover(start_dir=LIB_DIR, pattern="*test.py")
+    standard_tests.addTests(package_tests)
+    return standard_tests

sheetsage/align.py

@@ -0,0 +1,29 @@
+import numpy as np
+from scipy.interpolate import interp1d
+
+
+def _extrapolating_linear_interp1d(a, b, safe=True):
+    if safe:
+        if isinstance(a, np.ndarray):
+            a = a.tolist()
+        if isinstance(b, np.ndarray):
+            b = b.tolist()
+        if a != sorted(a):
+            raise ValueError()
+        if b != sorted(b):
+            raise ValueError()
+        if len(a) != len(b):
+            raise ValueError()
+        if len(np.unique(a)) != len(a):
+            raise ValueError()
+        if len(np.unique(b)) != len(b):
+            raise ValueError()
+    return interp1d(a, b, kind="linear", fill_value="extrapolate")
+
+
+def create_beat_to_time_fn(beats, times, safe=True):
+    return _extrapolating_linear_interp1d(beats, times, safe=safe)
+
+
+def create_time_to_beat_fn(beats, times, safe=True):
+    return _extrapolating_linear_interp1d(times, beats, safe=safe)

sheetsage/assets.py

@@ -0,0 +1,165 @@
+import json
+import logging
+import pathlib
+import urllib.request
+
+from . import CACHE_DIR, LIB_DIR
+from .utils import compute_checksum
+
+_DEFAULT_CHUNK_SIZE = 4096
+_ASSETS = None
+
+
+def _init_assets():
+    global _ASSETS
+    if _ASSETS is not None:
+        raise Exception("Should only run this once")
+
+    _ASSETS = {}
+    asset_paths = set()
+    for json_path in sorted(pathlib.Path(LIB_DIR, "assets").rglob("*.json")):
+        with open(json_path, "r") as f:
+            d = json.load(f)
+        for tag, asset in d.items():
+            if "checksum" not in asset:
+                raise AssertionError("Missing checksum")
+            try:
+                asset["path"] = pathlib.PurePosixPath(asset["path"].strip())
+            except:
+                raise AssertionError("Invalid path")
+            if asset["path"] in asset_paths:
+                raise AssertionError("Duplicate path")
+            asset_paths.add(asset["path"])
+            asset["path_abs"] = pathlib.Path(CACHE_DIR, asset["path"])
+        _ASSETS.update(d)
+
+
+_init_assets()
+
+
+def get_asset_tags():
+    return set(_ASSETS.keys())
+
+
+def _download(url, dest_path, chunk_size=_DEFAULT_CHUNK_SIZE):
+    with open(dest_path, "wb") as f:
+        r = urllib.request.urlopen(url)
+        while True:
+            chunk = r.read(chunk_size)
+            if not chunk:
+                break
+            f.write(chunk)
+
+
+def retrieve_asset(tag, delete_wrong=False, chunk_size=_DEFAULT_CHUNK_SIZE, log=True):
+    """Attempts to acquire and/or verify existence of a tagged asset in the cache.
+
+    Returns
+    -------
+    str
+       Absolute file path for asset, if verified.
+
+    Raises
+    ------
+    :class:`ValueError`
+       Invalid asset tag.
+    :class:`Exception`
+       Asset could not be verified.
+    """
+    # Retrieve asset
+    if tag not in _ASSETS:
+        raise ValueError()
+    asset = _ASSETS[tag]
+    path = asset["path_abs"]
+    checksum = asset["checksum"]
+    if log:
+        logging.info(f"Verifying asset: {tag}")
+        logging.info(f"Asset location: {path}")
+
+    # Create parent directory
+    if not path.parent.is_dir():
+        if log:
+            logging.info(f"Creating parent: {path.parent}")
+        path.parent.mkdir(parents=True)
+
+    def verify():
+        assert path.is_file()
+        if checksum is not None:
+            if len(checksum) == 32:
+                algorithm = "md5"
+            elif len(checksum) == 40:
+                algorithm = "sha1"
+            elif len(checksum) == 64:
+                algorithm = "sha256"
+            else:
+                raise AssertionError("Unknown checksum algorithm")
+            computed = compute_checksum(
+                path, algorithm=algorithm, chunk_size=chunk_size
+            )
+            if computed != checksum:
+                raise Exception(f"File {path} has wrong checksum.")
+
+    # Delete incorrect files
+    already_verified = False
+    if delete_wrong and path.is_file():
+        try:
+            verify()
+            already_verified = True
+        except Exception:
+            logging.warning(f"Deleting file with bad checksum: {path}")
+            path.unlink()
+
+    # Attempt to download
+    if not path.is_file():
+        url = asset.get("url")
+        if url is None:
+            raise Exception("File is missing and cannot be downloaded")
+        if log:
+            logging.info(f"Downloading from: {url}")
+        try:
+            _download(url, path)
+        except Exception as e:
+            if path.is_file():
+                path.unlink()
+            raise Exception(f"Download failed: {e}")
+    assert path.is_file()
+
+    # Ensure file integrity
+    if not already_verified:
+        verify()
+    if log:
+        logging.info(f"Verified!")
+
+    return path
+
+
+if __name__ == "__main__":
+    import multiprocessing
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser()
+
+    parser.add_argument("startswith", nargs="?")
+    parser.add_argument("--delete_wrong", action="store_true", dest="delete_wrong")
+    parser.add_argument("--num_parallel", "-n", type=int)
+
+    parser.set_defaults(startswith=None, num_parallel=1, delete_wrong=False)
+
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO)
+
+    tags = sorted(list(get_asset_tags()))
+    if args.startswith is not None:
+        tags = [t for t in tags if t.startswith(args.startswith.strip().upper())]
+
+    def task(t):
+        logging.info("-" * 80)
+        try:
+            retrieve_asset(t, delete_wrong=args.delete_wrong)
+        except Exception as e:
+            logging.error(e)
+            raise e
+
+    with multiprocessing.Pool(args.num_parallel) as p:
+        p.map(task, tags)

sheetsage/assets/hooktheory.json

@@ -0,0 +1,42 @@
+{
+  "HOOKTHEORY": {
+    "path": "hooktheory/Hooktheory.json.gz",
+    "url": "https://github.com/chrisdonahue/sheetsage-data/raw/refs/heads/main/hooktheory/Hooktheory.json.gz",
+    "checksum": "917b7cd58f5f4e07d6c36acf7bfad958c99ee05472dab3555399141094698e0c"
+  },
+  "HOOKTHEORY_TRAIN_SEGMENTS": {
+    "path": "hooktheory/Hooktheory_Train_Segments.json",
+    "url": "https://github.com/chrisdonahue/sheetsage-data/raw/refs/heads/main/hooktheory/Hooktheory_Train_Segments.json",
+    "checksum": "f2601eb544f2e5028ffad54d3827912865578fb9ad96e6768b35e4714d5c7207"
+  },
+  "HOOKTHEORY_TRAIN_MIDI": {
+    "path": "hooktheory/Hooktheory_Train_MIDI.tar.gz",
+    "url": "https://github.com/chrisdonahue/sheetsage-data/raw/refs/heads/main/hooktheory/Hooktheory_Train_MIDI.tar.gz",
+    "checksum": "a2345e13564c81740c087b79731b47e8323c4ceb7b85d40a1a11582af58145cb"
+  },
+  "HOOKTHEORY_VALID_SEGMENTS": {
+    "path": "hooktheory/Hooktheory_Valid_Segments.json",
+    "url": "https://github.com/chrisdonahue/sheetsage-data/raw/refs/heads/main/hooktheory/Hooktheory_Valid_Segments.json",
+    "checksum": "12526962f77c2eb41cd117c8effa678b39c2b350384a7b048b327aff287b0c48"
+  },
+  "HOOKTHEORY_VALID_MIDI": {
+    "path": "hooktheory/Hooktheory_Valid_MIDI.tar.gz",
+    "url": "https://github.com/chrisdonahue/sheetsage-data/raw/refs/heads/main/hooktheory/Hooktheory_Valid_MIDI.tar.gz",
+    "checksum": "e369fd4a3072c7524e3cafe506bbdad6e908de969d0f1ba7abf08bf5148989fe"
+  },
+  "HOOKTHEORY_TEST_SEGMENTS": {
+    "path": "hooktheory/Hooktheory_Test_Segments.json",
+    "url": "https://github.com/chrisdonahue/sheetsage-data/raw/refs/heads/main/hooktheory/Hooktheory_Test_Segments.json",
+    "checksum": "72be80045d4d28842352383e605e8712d50b3437a07b15faa541ee9d17283d5a"
+  },
+  "HOOKTHEORY_TEST_MIDI": {
+    "path": "hooktheory/Hooktheory_Test_MIDI.tar.gz",
+    "url": "https://github.com/chrisdonahue/sheetsage-data/raw/refs/heads/main/hooktheory/Hooktheory_Test_MIDI.tar.gz",
+    "checksum": "3baebe9d4e19a5006d0f24bc7f0c92a4f66039ab376be86bf7b37a136d4fb6c8"
+  },
+  "HOOKTHEORY_RAW": {
+    "path": "hooktheory/Hooktheory_Raw.json.gz",
+    "url": "https://github.com/chrisdonahue/sheetsage-data/raw/refs/heads/main/hooktheory/Hooktheory_Raw.json.gz",
+    "checksum": "716af2979f060400c302ab098dd45d9f8c5fe4d4b3b1fe61c478dd9bdf041634"
+  }
+}