initial deploy

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.pyc
+.DS_Store
+STAGE-repo/

README.md

@@ -1,14 +1,12 @@
 ---
 title: STAGE
-emoji: 🔥
+emoji: 🏟️
 colorFrom: blue
-colorTo: yellow
+colorTo: purple
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.12'
+sdk_version: 5.28.0
+python_version: '3.11'
 app_file: app.py
 pinned: false
 short_description: Single Stem Accompaniment Generation
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

__init__.py

@@ -0,0 +1,2 @@
+def main() -> None:
+    print("Hello from stage!")

app.py

@@ -0,0 +1,218 @@
+import sys
+sys.stdout.reconfigure(line_buffering=True)
+
+try:
+    import spaces
+except ImportError:
+    class spaces:
+        class GPU:
+            def __init__(self, func=None, duration=60):
+                self.func = func
+
+            def __call__(self, *args, **kwargs):
+                if self.func is not None:
+                    return self.func(*args, **kwargs)
+                func = args[0]
+                return func
+
+import tempfile
+import threading
+from pathlib import Path
+import torch
+import torchaudio
+import soundfile as sf
+import gradio as gr
+from huggingface_hub import hf_hub_download
+from safetensors import torch as sft
+
+import hyperparameters as hp
+from conditioning.condition_type import ConditionType
+from conditioning.conditioning_method import ConditioningMethod
+from conditioning.prompt_processor import InterleavedContextPromptProcessor
+from conditioning.t5embedder import T5EmbedderGPU
+from models.lightning_musicgen import LightningMusicgen
+from pyharp import ModelCard, build_endpoint
+
+REPO_ID = "teamup-tech/STAGE-checkpoints"
+SAMPLE_RATE = 32_000
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+drums_model: LightningMusicgen | None = None
+bass_model: LightningMusicgen | None = None
+models_on_device = False
+model_loading = True
+model_error: str | None = None
+
+
+def build_params(encodec_weights: str, lm_weights: str) -> hp.MusicgenParams:
+    """Assemble MusicgenParams with explicit weight paths."""
+    return hp.MusicgenParams(
+        encodec_params=hp.EncodecParams(
+            sample_rate=32_000,
+            seanet_params=hp.SeaNetParams(128, 64, (8, 5, 4, 4), False, True),
+            quantizer_params=hp.QuantizerParams(128, 4, 2048),
+            sum_loss_mulitiplier=0,
+            weights=encodec_weights,
+        ),
+        prompt_processor_params=hp.PromptProcessorParams(
+            keep_only_valid_steps=True,
+            model_class=InterleavedContextPromptProcessor,
+            context_dropout=0.1,
+        ),
+        conditioning_params=hp.ConditioningParams(
+            embedder_types={ConditionType.DESCRIPTION: T5EmbedderGPU},
+            conditioning_methods={
+                ConditionType.DESCRIPTION: ConditioningMethod.CROSS_ATTENTION
+            },
+            conditioning_dropout=0.5,
+        ),
+        lm_params=hp.PretrainedSmallLmParams(sep_token=2049, weights=lm_weights),
+    )
+
+
+def load_checkpoint(params: hp.MusicgenParams, ckp_path: str) -> LightningMusicgen:
+    """Instantiate model from params and load fine-tuned checkpoint onto CPU."""
+    model: LightningMusicgen = params.instantiate()
+    sft.load_model(model, ckp_path)
+    return model.cpu().eval()
+
+
+def load_models():
+    """Download all weights and build both models on CPU. Background thread only — no CUDA."""
+    global drums_model, bass_model, model_loading, model_error
+    try:
+        print("Downloading shared weights...")
+        encodec_path = hf_hub_download(REPO_ID, "encodec_32khz.pt")
+        lm_path = hf_hub_download(REPO_ID, "lm-small-weights.pt")
+        params = build_params(encodec_path, lm_path)
+
+        print("Building drums model (CPU)...")
+        drums_ckp = hf_hub_download(REPO_ID, "stage-drums.safetensors")
+        drums_model = load_checkpoint(params, drums_ckp)
+        print("Drums model ready.")
+
+        print("Building bass model (CPU)...")
+        bass_ckp = hf_hub_download(REPO_ID, "stage-bass.safetensors")
+        bass_model = load_checkpoint(params, bass_ckp)
+        print("Bass model ready.")
+    except Exception as e:
+        model_error = str(e)
+        print(f"Load error: {e}")
+    finally:
+        model_loading = False
+
+
+threading.Thread(target=load_models, daemon=True).start()
+
+
+model_card = ModelCard(
+    name="STAGE",
+    description=(
+        "Single Stem Accompaniment Generation. Provide an audio mix or click track "
+        "and STAGE generates a coherent drums or bass stem to accompany it."
+    ),
+    author="Giorgio Strano, Vansh Chaudhary, Derek Tran",
+    tags=["music-generation", "accompaniment", "stems"],
+)
+
+
+@spaces.GPU(duration=120)
+@torch.inference_mode()
+def process_fn(
+    input_audio_path: str,
+    instrument: str,
+    gen_seconds: int,
+    description: str,
+) -> str:
+    """Load context audio, run STAGE autoregressive generation, return the generated stem."""
+    global drums_model, bass_model, models_on_device
+
+    if model_loading:
+        raise gr.Error("Model is still loading — please wait a moment and try again.")
+    if model_error:
+        raise gr.Error(f"Model failed to load: {model_error}")
+
+    model = drums_model if instrument == "Drums" else bass_model
+    if model is None:
+        raise gr.Error(f"Model for {instrument} not available.")
+
+    if not models_on_device:
+        drums_model = drums_model.to(DEVICE)  # type: ignore[union-attr]
+        bass_model = bass_model.to(DEVICE)    # type: ignore[union-attr]
+        models_on_device = True
+
+    # load and prepare context audio
+    audio_np, orig_sr = sf.read(input_audio_path, always_2d=True)
+    context = torch.from_numpy(audio_np.T).float()
+    context = torchaudio.functional.resample(context, orig_sr, SAMPLE_RATE)
+    if context.shape[0] > 1:
+        context = context.mean(dim=0, keepdim=True)
+    context = context.reshape(1, 1, -1).to(DEVICE)
+
+    desc = [description.strip()] if description and description.strip() else [None]
+
+    out = model.generate(
+        n_samples=1,
+        gen_seconds=gen_seconds,
+        prompt=None,
+        context=context,
+        style=None,
+        beat=None,
+        description=desc,
+    )
+
+    # out: (1, 1, T) or (1, 2, T) — save as stereo wav
+    audio_out = out.squeeze(0).cpu().float()
+    if audio_out.shape[0] == 1:
+        audio_out = audio_out.repeat(2, 1)
+
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+        out_path = f.name
+    torchaudio.save(out_path, audio_out, sample_rate=SAMPLE_RATE)
+    return out_path
+
+
+with gr.Blocks() as demo:
+    input_components = [
+        gr.Audio(
+            type="filepath",
+            label="Context Track",
+        ).harp_required(True),
+        gr.Radio(
+            choices=["Drums", "Bass"],
+            value="Drums",
+            label="Instrument",
+            info="Which accompaniment stem to generate.",
+        ),
+        gr.Slider(
+            minimum=5,
+            maximum=20,
+            step=1,
+            value=10,
+            label="Length (seconds)",
+            info="Duration of the generated stem (default: 10 sec, per paper).",
+        ),
+        gr.Textbox(
+            value="",
+            label="Style Description",
+            info="Optional — describe the mood or style (e.g. 'lo-fi chill groove').",
+            placeholder="lo-fi chill groove with soft kick...",
+        ),
+    ]
+    output_components = [
+        gr.Audio(
+            type="filepath",
+            label="Generated Stem",
+        ).set_info(
+            "The generated accompaniment stem. Mix it back with your context track."
+        ),
+    ]
+
+    build_endpoint(
+        model_card=model_card,
+        input_components=input_components,
+        output_components=output_components,
+        process_fn=process_fn,
+    )
+
+demo.queue().launch(pwa=True)

conditioning/__init__.py

@@ -0,0 +1,6 @@
+from typing import Union
+from conditioning.beat_embedder import DirectSinusoidalBeatEmbedder, SinusoidalBeatEmbedder
+from conditioning.t5embedder import T5EmbedderCPU, T5EmbedderGPU
+
+ConcreteEmbedder = Union[T5EmbedderCPU, T5EmbedderGPU, SinusoidalBeatEmbedder,
+                         DirectSinusoidalBeatEmbedder]

conditioning/beat_embedder.py

@@ -0,0 +1,249 @@
+from typing import List, Sequence
+from torch import Tensor
+import numpy as np
+import torch
+from torch import nn
+from dataclasses import dataclass
+
+from utils.audio import load_audio, make_variable_frequency_sinewave
+from utils.plotting import plot_waveforms
+from conditioning.embedded_condition import EmbeddedCondition
+from conditioning.embedder import Embedder, LinearProjectionEmbedder
+import config as cfg
+
+
+@dataclass
+class Beat:
+    """
+        beats / downbeats: 
+            Tensor of Int or Long type, containing indices of samples where
+            beats / downbeats occur
+        seq_len: int representing the total length of the audio in samples
+    """
+    beats: Tensor
+    downbeats: Tensor
+    seq_len: int
+
+
+class DirectSinusoidalBeatEmbedder(Embedder):
+
+    def __init__(self, embedding_dim: int):
+        super().__init__(input_dim=2, embedding_dim=embedding_dim)
+        self.encodec_framerate: int = 50
+        self.sample_rate: int = 32_000
+
+    def forward(self,
+                x: Sequence[Beat],
+                duplicate_for_cfg: bool = False) -> EmbeddedCondition:
+        """
+            Returns: A tensor of shape [B, S, H], where B is batch, H is 
+            embedding dim, and S is the length of the sequence
+        """
+
+        assert len(set([t.seq_len for t in x])) == 1
+
+        emb_list: List[Tensor] = [None] * len(x)  # type: ignore
+
+        ratio = self.encodec_framerate / self.sample_rate
+        for i, track in enumerate(x):
+
+            # convert indices from explicit to latent space
+            beats = torch.floor(track.beats.float() * ratio).long()
+            downbeats = torch.floor(track.downbeats.float() * ratio).long()
+            seq_len = round(track.seq_len * ratio)
+
+            # make a variable-frequency sinewave lined with the beats
+            beat_emb = make_variable_frequency_sinewave(seq_len, beats)
+            downbeat_emb = make_variable_frequency_sinewave(seq_len, downbeats)
+
+            emb_list[i] = torch.stack((beat_emb, downbeat_emb), dim=-1)
+
+        embeds: Tensor = torch.stack(emb_list)
+
+        # create an embedding that's just the beat sinewave in every dimension
+        embeds = torch.cat(
+            (embeds[..., :1].repeat(1, 1, self.embedding_dim // 2),
+             embeds[..., 1:].repeat(1, 1, self.embedding_dim // 2)),
+            dim=-1)
+        mask = torch.ones(embeds.shape[:-1],
+                          device=embeds.device,
+                          dtype=torch.bool)
+
+        # concatenate an empty condition
+        if duplicate_for_cfg:
+            embeds = torch.cat((embeds, torch.zeros_like(embeds)), dim=0)
+            mask = torch.cat((mask, torch.zeros_like(mask)), dim=0)
+
+        return EmbeddedCondition(embeds, mask)
+
+    def null_condition(self, batch_size: int) -> EmbeddedCondition:
+        return EmbeddedCondition(
+            torch.zeros(batch_size,
+                        1,
+                        self.embedding_dim,
+                        dtype=torch.float32,
+                        device=self.output_proj.weight.device),
+            torch.zeros(batch_size,
+                        1,
+                        dtype=torch.bool,
+                        device=self.output_proj.weight.device))
+
+
+class SinusoidalBeatEmbedder(LinearProjectionEmbedder):
+
+    def __init__(self, embedding_dim: int):
+        super().__init__(input_dim=2, embedding_dim=embedding_dim)
+        self.encodec_framerate: int = 50
+        self.sample_rate: int = 32_000
+
+    def forward(self,
+                x: Sequence[Beat],
+                duplicate_for_cfg: bool = False) -> EmbeddedCondition:
+        """
+            Returns: A tensor of shape [B, S, H], where B is batch, H is 
+            embedding dim, and S is the length of the sequence
+        """
+
+        assert len(set([t.seq_len for t in x])) == 1
+
+        emb_list: List[Tensor] = [None] * len(x)  # type: ignore
+
+        for i, track in enumerate(x):
+
+            beats = (track.beats / self.sample_rate *
+                     self.encodec_framerate).round().int()
+            downbeats = (track.downbeats / self.sample_rate *
+                         self.encodec_framerate).round().int()
+
+            seq_len = round(track.seq_len / self.sample_rate *
+                            self.encodec_framerate)
+
+            beat_emb = make_variable_frequency_sinewave(seq_len, beats)
+            downbeat_emb = make_variable_frequency_sinewave(seq_len, downbeats)
+
+            emb_list[i] = torch.stack((beat_emb, downbeat_emb), dim=-1)
+
+        embeds: Tensor = torch.stack(emb_list)
+
+        embeds = self.output_proj(embeds.to(self.output_proj.weight))
+        mask = torch.ones(embeds.shape[:-1],
+                          device=embeds.device,
+                          dtype=torch.bool)
+
+        if duplicate_for_cfg:
+            embeds = torch.cat((embeds, torch.zeros_like(embeds)), dim=0)
+            mask = torch.cat((mask, torch.zeros_like(mask)), dim=0)
+
+        return EmbeddedCondition(embeds, mask)
+
+    def null_condition(self, batch_size: int) -> EmbeddedCondition:
+        return EmbeddedCondition(
+            torch.zeros(batch_size,
+                        1,
+                        self.embedding_dim,
+                        dtype=torch.float32,
+                        device=self.output_proj.weight.device),
+            torch.zeros(batch_size,
+                        1,
+                        dtype=torch.bool,
+                        device=self.output_proj.weight.device))
+
+
+class SinusoidalBeatEmbedderMLP(Embedder):
+
+    def __init__(self, embedding_dim: int):
+        super().__init__(input_dim=2, embedding_dim=embedding_dim)
+        # self.embedding_dim: int = embedding_dim
+        self.encodec_framerate: int = 50
+        self.sample_rate: int = 32_000
+
+        self.output_proj: nn.Sequential = nn.Sequential(
+            nn.Linear(self.input_dim, 512), nn.ReLU(),
+            nn.Linear(512, self.embedding_dim))
+
+    def forward(self,
+                x: Sequence[Beat],
+                duplicate_for_cfg: bool = False) -> EmbeddedCondition:
+        """
+            Returns: A tensor of shape [B, S, H], where B is batch, H is 
+            embedding dim, and S is the length of the sequence
+        """
+
+        assert len(set([t.seq_len for t in x])) == 1
+
+        emb_list: List[Tensor] = [None] * len(x)  # type: ignore
+
+        for i, track in enumerate(x):
+
+            beats = (track.beats / self.sample_rate *
+                     self.encodec_framerate).round().int()
+            downbeats = (track.downbeats / self.sample_rate *
+                         self.encodec_framerate).round().int()
+
+            seq_len = round(track.seq_len / self.sample_rate *
+                            self.encodec_framerate)
+
+            beat_emb = make_variable_frequency_sinewave(seq_len, beats)
+            downbeat_emb = make_variable_frequency_sinewave(seq_len, downbeats)
+
+            emb_list[i] = torch.stack((beat_emb, downbeat_emb), dim=-1)
+
+        embeds: Tensor = torch.stack(emb_list)
+
+        embeds = self.output_proj(embeds.to(self.output_proj[-1].weight))
+        mask = torch.ones(embeds.shape[:-1],
+                          device=embeds.device,
+                          dtype=torch.bool)
+
+        if duplicate_for_cfg:
+            embeds = torch.cat((embeds, torch.zeros_like(embeds)), dim=0)
+            mask = torch.cat((mask, torch.zeros_like(mask)), dim=0)
+
+        return EmbeddedCondition(embeds, mask)
+
+    def null_condition(self, batch_size: int) -> EmbeddedCondition:
+        return EmbeddedCondition(
+            torch.zeros(batch_size,
+                        1,
+                        self.embedding_dim,
+                        dtype=torch.float32,
+                        device=self.output_proj[-1].weight.device),
+            torch.zeros(batch_size,
+                        1,
+                        dtype=torch.bool,
+                        device=self.output_proj[-1].weight.device))
+
+
+if __name__ == "__main__":
+    song_path = cfg.moises_path(
+    ) / "0d528a19-cb0f-4421-b250-444f9343e51c/mixed.wav"
+    song = load_audio(song_path)
+    n_samples = song.shape[-1]
+    loaded = np.load(
+        "/home/tkol/dev/datasets/moisesdb/moisesdb_v0.1/0d528a19-cb0f-4421-b250-444f9343e51c/beatthis.npz"
+    )
+    beats = torch.tensor([int(b * 32_000) for b in loaded["beats"]],
+                         dtype=torch.int32)
+    downbeats = torch.tensor([int(b * 32_000) for b in loaded["downbeats"]],
+                             dtype=torch.int32)
+    assert n_samples > beats.max()
+
+    b = Beat(beats=beats, downbeats=downbeats, seq_len=n_samples)
+    embedder = SinusoidalBeatEmbedder(1024)
+    emb = embedder([b])
+
+    directembedder = DirectSinusoidalBeatEmbedder(1024)
+    directemb = directembedder([b])
+    # from matplotlib import pyplot as plt
+    # plt.style.use(cfg.ROOT / "tkol.mplstyle")
+    # plt.figure(figsize=(100, 10))
+    # plt.scatter(x=torch.arange(emb.shape[-1]) / 50 * 32_000, y=emb)
+    # plt.show()
+
+    # plot_waveforms(
+    #     song,
+    #     emb,
+    #     # savepath=cfg.ROOT / "plots/beatontrack.png",
+    #     figsize=(50, 3),
+    #     dpi=100,
+    # )

conditioning/condition_dispatcher.py

@@ -0,0 +1,157 @@
+from typing import Callable, List, Sequence, Dict, Type
+from torch import nn, Tensor
+import torch
+
+from conditioning.condition_type import ConditionType
+from conditioning.conditioning_method import ConditioningMethod
+from conditioning.embedded_condition import EmbeddedCondition
+
+
+class ConditionDispatcher(nn.Module):
+
+    def __init__(self, cond_to_method: Dict[ConditionType, ConditioningMethod],
+                 embedding_dim: int, condition_dropout: float):
+        super().__init__()
+        self.embedding_dim: int = embedding_dim
+        self.condition_dropout: float = condition_dropout
+
+        # what fusing method to use for each conditioning type
+        self.cond_to_method: Dict[ConditionType,
+                                  ConditioningMethod] = cond_to_method
+
+        # what condition types is each method applied to
+        self.method_to_conds: Dict[ConditioningMethod, List[ConditionType]] = {}
+        for c, m in self.cond_to_method.items():
+            self.method_to_conds[m] = self.method_to_conds.get(m, []) + [c]
+
+        # instantiate fuser classes
+        self.method_to_fuser: nn.ModuleDict = nn.ModuleDict({
+            m.value:
+                METHOD_TO_FUSER_TYPE[m](len(self.method_to_conds[m]),
+                                        self.embedding_dim)
+            for m in self.method_to_conds.keys()
+        })
+
+    def dropout(self, x: EmbeddedCondition):
+        if self.training and self.condition_dropout > 0:
+            drop = torch.rand(x.data.shape[0]) < self.condition_dropout
+            # x.data[drop] = torch.zeros_like(x.data[drop])  #.detach()
+            x.data[drop] = x.data[drop] * 0
+            if x.mask is not None:
+                x.mask[drop] = torch.zeros_like(x.mask[drop])  #.detach()
+        return x
+
+    def forward(
+        self, embedded_conditions: Dict[ConditionType, EmbeddedCondition]
+    ) -> Dict[ConditioningMethod, EmbeddedCondition]:
+        assert embedded_conditions.keys() == self.cond_to_method.keys()
+
+        # figure out where does each conditioning go, between:
+        #  - summed to the input of the lm (after embedding)
+        #  - prepended to the input of the lm (after embedding)
+        #  - in cross-attention of the lm
+        method_to_embeddings: Dict[ConditioningMethod,
+                                   List[EmbeddedCondition]] = {}
+
+        # for condition_type, emb in embedded_conditions.items():
+        for condition_type in sorted(embedded_conditions.keys()):
+            emb = embedded_conditions[condition_type]
+
+            # apply dropout to each condition
+            emb = self.dropout(emb)
+
+            method: ConditioningMethod = self.cond_to_method[condition_type]
+            method_to_embeddings[method] = method_to_embeddings.get(method,
+                                                                    []) + [emb]
+
+        # has to return a dict that associates one embedded condition to each
+        # conditioning method
+        method_to_fused_embedding: Dict[ConditioningMethod,
+                                        EmbeddedCondition] = {}
+
+        # if multiple conditions are dispatched to the same method, we have to
+        # fuse them
+        for m, emb_list in method_to_embeddings.items():
+            fused_embedding: EmbeddedCondition = self.method_to_fuser[m.value](
+                emb_list)
+            method_to_fused_embedding[m] = fused_embedding
+
+        return method_to_fused_embedding
+
+
+class CrossAttentionFuser(nn.Module):
+
+    def __init__(self, n_conditions: int, embedding_dim: int):
+        super().__init__()
+        self.n_conditions: int = n_conditions
+        self.embedding_dim: int = embedding_dim
+
+        # use concatenation with segment embedding to merge conditions
+        if self.n_conditions > 1:
+            self.segment_embedding = nn.Embedding(self.n_conditions - 1,
+                                                  self.embedding_dim)
+
+    def forward(self, conds: Sequence[EmbeddedCondition]) -> EmbeddedCondition:
+        if len(conds) == 0:
+            raise RuntimeError("Received a list of 0 length. There is a bug.")
+        if len(conds) == 1:
+            return conds[0]
+
+        # apply segment embedding
+        if len(conds) == 2:
+            assert self.n_conditions == 2
+            assert conds[0].mask is not None and conds[1].mask is not None
+            cond1_emb = conds[1].data
+            cond1_ids = torch.zeros(cond1_emb.shape[0],
+                                    cond1_emb.shape[1],
+                                    dtype=torch.long,
+                                    device=conds[1].data.device)
+            cond1_segembed = self.segment_embedding(cond1_ids)
+            cond1_emb = cond1_emb + cond1_segembed
+            total_cond = torch.cat((conds[0].data, cond1_emb), dim=1)
+            total_mask = torch.cat((conds[0].mask, conds[1].mask), dim=1)
+            return EmbeddedCondition(total_cond, total_mask)
+
+        # multiple conditions in cross-atteniton are not implemented
+        raise NotImplementedError()
+
+
+class PrependFuser(nn.Module):
+
+    def __init__(self, n_conditions: int, embedding_dim: int):
+        super().__init__()
+        self.n_conditions: int = n_conditions
+        self.embedding_dim: int = embedding_dim
+
+    def forward(self, conds: Sequence[EmbeddedCondition]):
+        if len(conds) == 0:
+            raise RuntimeError("Received a list of 0 length. There is a bug.")
+        if len(conds) == 1:
+            return conds[0]
+
+        # multiple conditions prepended to input are not implemented
+        raise NotImplementedError()
+
+
+class SumFuser(nn.Module):
+
+    def __init__(self, n_conditions: int, embedding_dim: int):
+        super().__init__()
+        self.n_conditions: int = n_conditions
+        self.embedding_dim: int = embedding_dim
+
+    def forward(self, conds: Sequence[EmbeddedCondition]):
+        if len(conds) == 0:
+            raise RuntimeError("Received a list of 0 length. There is a bug.")
+        if len(conds) == 1:
+            return conds[0]
+
+        # multiple conditions summed to input are not implemented
+        raise NotImplementedError()
+
+
+METHOD_TO_FUSER_TYPE: Dict[ConditioningMethod, Type[nn.Module]] = {
+    ConditioningMethod.CROSS_ATTENTION: CrossAttentionFuser,
+    ConditioningMethod.INPUT_PREPEND: PrependFuser,
+    ConditioningMethod.INPUT_SUM: SumFuser,
+}

conditioning/condition_provider.py

@@ -0,0 +1,91 @@
+from typing import Dict, Optional, Type, Any, Union
+from torch import Tensor, nn
+import torch
+
+from conditioning.condition_type import ConditionType
+from conditioning.embedded_condition import EmbeddedCondition
+from conditioning.t5embedder import T5EmbedderCPU
+from conditioning import ConcreteEmbedder
+
+
+class ConditionProvider(nn.Module):
+
+    def __init__(self, embedding_dim: int,
+                 embedder_types: Dict[ConditionType, Type[ConcreteEmbedder]]):
+        super().__init__()
+        self.embedding_dim: int = embedding_dim
+        self.embedders: nn.ModuleDict = nn.ModuleDict()
+
+        # instantiate embedders for all condition types
+        for condition, embedder_type in embedder_types.items():
+            self.embedders[condition.value] = embedder_type(
+                embedding_dim=self.embedding_dim)
+
+    # def duplicate_conditions_for_cfg(self, conditions: Dict[str, Any]):
+    #     for cond_name, cond_data in conditions.items():
+    #         if isinstance(cond_data, list) and isinstance(cond_data[0], str):
+    #             conditions[cond_name] += [""] * len(cond_data)
+    #         elif isinstance(cond_data, Tensor):
+    #             conditions[cond_name] = torch.cat(
+    #                 (cond_data, torch.zeros_like(cond_data)), dim=0)
+    #         else:
+    #             raise RuntimeError(
+    #                 f"I don't know what an emtpy condition for type: {type}")
+
+    # embed condition into tensors with the corresponding embedders
+    def process_conditions(
+        self,
+        conditions: Dict[str, Any],
+        duplicate_for_cfg: bool = False,
+        batch_size: Optional[int] = None,
+    ) -> Dict[ConditionType, EmbeddedCondition]:
+
+        processed_conditions: Dict[ConditionType, EmbeddedCondition] = {}
+
+        # for each condition in the dict, try to find the corresponding embedder
+        # for cond_name, cond_data in conditions.items():
+        #     condtype: ConditionType = ConditionType(cond_name)
+        #     if condtype.value not in self.embedders:
+        #         raise RuntimeError(f"I don't have an embedder for a condition "
+        #                            f"named {cond_name}")
+        #     embedded: EmbeddedCondition = self.embedders[condtype.value](
+        #         cond_data, duplicate_for_cfg=duplicate_for_cfg)
+        #     processed_conditions[condtype] = embedded
+
+        # for each of my embedders, find the corresponding condition in the dict
+        for cond_name, embedder in self.embedders.items():
+            condtype: ConditionType = ConditionType(cond_name)
+            if (condtype.value in conditions and
+                    conditions[condtype.value] is not None):
+                embedded: EmbeddedCondition = embedder(
+                    conditions[condtype.value],
+                    duplicate_for_cfg=duplicate_for_cfg)
+                processed_conditions[condtype] = embedded
+            else:
+                if batch_size is None:
+                    raise RuntimeError(
+                        f"Condition {cond_name} was not provided in batch. "
+                        f"I need the batch size to generate a null condition")
+                embedded: EmbeddedCondition = embedder.null_condition(
+                    batch_size + (batch_size * duplicate_for_cfg))
+                processed_conditions[condtype] = embedded
+
+        return processed_conditions
+
+
+if __name__ == "__main__":
+    import torch
+
+    conditions = {
+        "description": "The quick brown fox jumps over the lazy dog.",
+        # "context": torch.rand(4, 32_000),
+        # "style": torch.rand(4, 32_000),
+    }
+
+    provider = ConditionProvider(embedding_dim=2048,
+                                 embedder_types={
+                                     ConditionType.DESCRIPTION: T5EmbedderCPU,
+                                 })
+
+    with torch.no_grad():
+        processed_conditions = provider.process_conditions(conditions)

conditioning/condition_type.py

@@ -0,0 +1,26 @@
+from enum import Enum
+from functools import total_ordering
+
+
+class ConditionType(Enum):
+    DESCRIPTION = "description"
+    CONTEXT = "context"
+    STYLE = "style"
+    BEAT = "beat"
+
+    def __lt__(self, other):
+
+        order = ["DESCRIPTION", "BEAT", "CONTEXT", "STYLE"]
+        if not isinstance(other, ConditionType):
+            raise NotImplementedError()
+        return order.index(self.name) < order.index(other.name)
+
+
+if __name__ == "__main__":
+    c = ConditionType("description")
+    print(f'{c=}')
+    print(f'{c.value=}')
+
+    c2 = ConditionType("beat")
+
+    print(c < c2)

conditioning/conditioning_method.py

@@ -0,0 +1,14 @@
+from enum import Enum
+
+
+class ConditioningMethod(Enum):
+
+    INPUT_PREPEND = "input_prepend"
+    INPUT_SUM = "input_sum"
+    CROSS_ATTENTION = "cross_attention"
+
+
+if __name__ == "__main__":
+    c = ConditioningMethod.INPUT_PREPEND
+    print(f'{c=}')
+    print(f'{c.value=}')

conditioning/embedded_condition.py

@@ -0,0 +1,13 @@
+from torch import Tensor
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass(repr=False)
+class EmbeddedCondition:
+    data: Tensor
+    mask: Optional[Tensor]
+
+    def __repr__(self) -> str:
+        return (f"data: {list(self.data.shape)}; mask: "
+                f"{list(self.mask.shape) if self.mask is not None else 'None'}")

conditioning/embedder.py

@@ -0,0 +1,50 @@
+"""
+    Abstract base classes for conditioning embedders.
+    An embedder takes some conditioning data and embeds/encodes it into vectors
+    that are then fed to the LM according to the condition strategy specified 
+    in the config.
+    
+    Embedding vectors should have the same latent dimensionality of the 
+    transformer, so every embedder takes as input to the constructor a 
+    parameter `embedding_dim`, that should be the same as the hidden dim of 
+    the transformer.
+"""
+
+from torch import nn
+from abc import ABC, abstractmethod
+
+from conditioning.embedded_condition import EmbeddedCondition
+
+
+class Embedder(ABC, nn.Module):
+
+    def __init__(self, input_dim: int, embedding_dim: int):
+        super().__init__()
+        self.input_dim: int = input_dim
+        self.embedding_dim: int = embedding_dim
+
+    @abstractmethod
+    def forward(self, x, duplicate_for_cfg: bool) -> EmbeddedCondition:
+        ...
+
+    @abstractmethod
+    def null_condition(self, batch_size: int) -> EmbeddedCondition:
+        ...
+
+
+class LinearProjectionEmbedder(ABC, nn.Module):
+
+    def __init__(self, input_dim: int, embedding_dim: int):
+        super().__init__()
+        self.input_dim: int = input_dim
+        self.embedding_dim: int = embedding_dim
+        self.output_proj: nn.Linear = nn.Linear(self.input_dim,
+                                                self.embedding_dim)
+
+    @abstractmethod
+    def forward(self, x, duplicate_for_cfg: bool) -> EmbeddedCondition:
+        ...
+
+    @abstractmethod
+    def null_condition(self, batch_size: int) -> EmbeddedCondition:
+        ...

conditioning/prompt_processor.py

@@ -0,0 +1,625 @@
+"""
+    Processing the prompt of the language model.
+    With "prompt", we always refer to the sequence that is fed as input to 
+    the transformer model.
+
+    This sequence might contain conditioning data depending on the architecture
+    of the model.
+
+    Every prompt processor should handle at least:
+    - Encoding of the input tokens with encodec
+    - Interleaving the EnCodec tokens according to the proper pattern
+
+    All classes should extend the abstract PromptProcessor class and implement
+    their custom preprocessing strategy
+"""
+from torch import Tensor, nn
+import torch
+from typing import Callable, Dict, Any, List, Optional, Tuple
+from abc import ABC, abstractmethod
+from config import ConfigurationError
+from models.modules.codebooks_patterns import DelayedPatternProvider, Pattern
+from models.encodec import EncodecModel
+from utils.audio import pad_stack
+
+
+class PromptProcessor(ABC, nn.Module):
+    uses_sep_token: bool = False
+    """
+        Abstract base class for any prompt processor.
+    """
+
+    def __init__(self, encodec_model: EncodecModel, special_token: int,
+                 keep_only_valid_steps: bool, **kwargs):
+        super().__init__()
+
+        self.__dict__["encodec_model"] = encodec_model
+        # self.encodec_model: EncodecModel = encodec_model
+        self.special_token: int = special_token
+        self.keep_only_valid_steps = keep_only_valid_steps
+        self.pattern_provider = DelayedPatternProvider(
+            self.encodec_model.num_codebooks)
+
+    def interleave(
+            self, prompt: Tensor,
+            keep_only_valid_steps: bool) -> Tuple[Tensor, Tensor, Pattern]:
+        # interleave codes
+        prompt = prompt.contiguous()
+        pattern: Pattern = self.pattern_provider.get_pattern(prompt.shape[-1])
+        prompt, indices, mask = pattern.build_pattern_sequence(
+            prompt,
+            self.special_token,
+            keep_only_valid_steps=keep_only_valid_steps)
+
+        return prompt, mask, pattern
+
+    def deinterleave_logits(self, logits: Tensor,
+                            pattern: Pattern) -> Tuple[Tensor, Tensor]:
+        logits = logits.permute(0, 3, 1, 2)
+        logits, _, logits_mask = pattern.revert_pattern_logits(
+            logits,
+            float("nan"),
+            keep_only_valid_steps=self.keep_only_valid_steps)
+        logits = logits.permute(0, 2, 3, 1)
+        logits_mask = logits_mask[None, :, :].expand(logits.shape[0], -1, -1)
+        return logits, logits_mask
+
+    def encode_and_pad_list(self, tensors: List[Tensor],
+                            pad_value: int) -> Tensor:
+        encoded = [self.encode(t.view(1, 1, -1)).squeeze(0) for t in tensors]
+
+        max_len = max(t.shape[-1] for t in encoded)
+        padded_tensors = [
+            torch.nn.functional.pad(t, (max_len - t.shape[-1], 0),
+                                    value=pad_value) for t in encoded
+        ]
+        stack = torch.stack(padded_tensors, dim=0)
+        return stack
+
+    def encode(self, x: Tensor) -> Tensor:
+        if x.shape[-1] < 7:
+            return torch.empty(x.shape[1],
+                               self.encodec_model.num_codebooks,
+                               0,
+                               dtype=torch.long,
+                               device=x.device)
+        self.encodec_model.eval()
+        with torch.no_grad():
+            return self.encodec_model.encode(x)
+
+    @abstractmethod
+    def preprocess(
+        self, batch: Dict[str, Any]
+    ) -> Tuple[Tensor, Tensor, Tensor, Callable[[Tensor], Tuple[Tensor,
+                                                                Tensor]]]:
+        ...
+
+
+class ContextPromptProcessor(PromptProcessor):
+    """
+        Base class for a prompt processor that encodes target and
+        context with encodec.
+    """
+
+    def __init__(self, encodec_model: EncodecModel, special_token: int,
+                 keep_only_valid_steps: bool, context_dropout: float):
+        super().__init__(encodec_model, special_token, keep_only_valid_steps)
+        self.context_dropout: float = context_dropout
+
+    def encode_target_and_context_sequential(
+            self, target: Tensor, context: Tensor) -> Tuple[Tensor, Tensor]:
+        target_codes = self.encode(target)
+        context_codes = self.encode(context)
+        return target_codes, context_codes
+
+    def encode_target_and_context_parallel(
+            self, target: Tensor, context: Tensor) -> Tuple[Tensor, Tensor]:
+
+        batch_size = target.shape[0]
+        assert context.shape[0] == batch_size
+
+        # concatenate along batch target and context to encode in parallel
+        encodec_input = torch.cat((target, context), dim=0)
+
+        # encode with encodec
+        encodec_output: Tensor = self.encode(encodec_input)
+
+        target_codes: Tensor = encodec_output[:batch_size, ...]
+        context_codes: Tensor = encodec_output[-batch_size:, ...]
+        return target_codes, context_codes
+
+    def encode_target_and_context(
+            self, target: Tensor,
+            context: Tensor | List[Tensor]) -> Tuple[Tensor, Tensor]:
+
+        if isinstance(context, list):
+            context_codes: Tensor = self.encode_and_pad_list(
+                context, self.special_token)
+            target_codes: Tensor = self.encode(target)
+            return target_codes, context_codes
+
+        assert isinstance(context, Tensor)
+
+        target_len = target.shape[-1]
+        context_len = context.shape[-1]
+
+        if target_len == context_len:
+            return self.encode_target_and_context_parallel(target, context)
+        else:
+            return self.encode_target_and_context_sequential(target, context)
+
+    @abstractmethod
+    def prepare_for_generation(
+        self, prompt: Optional[Tensor], context: Optional[Tensor],
+        gen_sequence: Tensor, use_cfg: bool,
+        context_dropout_mask: Optional[Tensor]
+    ) -> Tuple[Tensor, Tensor, int, Callable[[Tensor], Tuple[Tensor, Tensor]]]:
+        """
+            Take prompt and context, prepare them, apply cfg, and insert them
+            in dummy sequence full of -1 for generation. Return: 
+            - gen_sequence
+            - gen_mask 
+            - start_offset 
+            - closure to decode sequence once filled with gen tokens
+        """
+
+
+class DefaultPromptProcessor(PromptProcessor):
+    """
+        A prompt processor that only encodes the target with encodec, without
+        conditioning on context.
+    """
+
+    def preprocess(self, batch: Dict[str, Any]):
+        target = batch["target"]
+
+        # encode target
+        target_codes: Tensor = self.encode(target)
+
+        # interleave
+        prompt, _, pattern = self.interleave(target_codes,
+                                             self.keep_only_valid_steps)
+
+        def decode_logits_fn(logits: Tensor) -> Tuple[Tensor, Tensor]:
+            return self.deinterleave_logits(logits, pattern)
+
+        mask = (prompt != 2048)
+
+        return prompt, mask, target_codes, decode_logits_fn
+
+    def prepare_for_generation(
+        self,
+        prompt: Optional[Tensor],
+        context: Optional[Tensor],
+        gen_sequence: Tensor,
+        use_cfg: bool,
+        context_dropout_mask: Optional[Tensor],
+    ) -> Tuple[Tensor, Tensor, int, Callable[[Tensor], Tuple[Tensor, Tensor]]]:
+
+        if context is not None or context_dropout_mask is not None:
+            raise ConfigurationError(
+                "This model does not support context conditioning in the prompt"
+            )
+        if prompt is None:
+            prompt_codes = torch.empty(gen_sequence.shape[0],
+                                       gen_sequence.shape[1],
+                                       0,
+                                       dtype=torch.long,
+                                       device=gen_sequence.device)
+        else:
+            prompt_codes = self.encode(prompt)
+
+        # TODO: make sure that duplicating for cfg is needed here
+        if use_cfg:
+            gen_sequence = torch.cat((gen_sequence, gen_sequence), dim=0)
+            prompt_codes = torch.cat((prompt_codes, prompt_codes), dim=0)
+
+        # insert encoded prompt into sequence
+        gen_sequence[..., :prompt_codes.shape[-1]] = prompt_codes
+
+        # interleave sequence (keep_only_valid_steps always false for inference)
+        gen_sequence, gen_mask, pattern = self.interleave(
+            gen_sequence, keep_only_valid_steps=False)
+
+        # compute first offset of the sequence to generate
+        start_offset: int = pattern.get_first_step_with_timesteps(  # type: ignore
+            prompt_codes.shape[-1])
+
+        mask: Tensor = (gen_sequence != 2048)
+
+        def decode_sequence_closure(sequence) -> Tuple[Tensor, Tensor]:
+            out_codes, _, out_mask = pattern.revert_pattern_sequence(
+                sequence, special_token=-1, keep_only_valid_steps=False)
+
+            return out_codes, out_mask
+
+        return gen_sequence, mask, start_offset, decode_sequence_closure
+
+
+class StraightContextPromptProcessor(ContextPromptProcessor):
+    """
+        Corresponding implementation to Bart and Azir models.
+        
+        Interleaves only the target. Concatenates straight context to 
+        interleaved target.
+    """
+
+    def preprocess(self, batch: Dict[str, Any]):
+        target = batch["target"]
+        context = batch["context"]
+
+        # encode target and context
+        target_codes, context_codes = self.encode_target_and_context(
+            target, context)
+
+        # apply dropout to context
+        if self.training and self.context_dropout > 0:
+            drop = torch.rand(context_codes.shape[0]) < self.context_dropout
+            context_codes[drop] = torch.full_like(context_codes[drop],
+                                                  self.special_token)
+
+        # interleave target only
+        prompt, mask, pattern = self.interleave(target_codes,
+                                                self.keep_only_valid_steps)
+
+        # concatenate context and prompt
+        prompt = torch.cat((context_codes, prompt), dim=-1)
+
+        def decode_logits_fn(logits: Tensor) -> Tuple[Tensor, Tensor]:
+            # remove context
+            n_context_codes = context_codes.shape[-1]
+            logits = logits[:, :, n_context_codes:, :]
+
+            # deinterleave
+            return self.deinterleave_logits(logits, pattern)
+
+        raise NotImplementedError("mask implementation missing")
+
+        return prompt, target_codes, decode_logits_fn
+
+    def prepare_for_generation(
+        self,
+        prompt: Optional[Tensor],
+        context: Optional[Tensor],
+        gen_sequence: Tensor,
+        use_cfg: bool,
+        context_dropout_mask: Optional[Tensor],
+    ) -> Tuple[Tensor, Tensor, int, Callable[[Tensor], Tuple[Tensor, Tensor]]]:
+        raise NotImplementedError("Missing context dropout mask implmentation")
+        # dummy prompt and context
+        prompt_codes = torch.empty(gen_sequence.shape[0],
+                                   gen_sequence.shape[1],
+                                   0,
+                                   dtype=torch.long,
+                                   device=gen_sequence.device)
+        context_codes = torch.empty(gen_sequence.shape[0],
+                                    gen_sequence.shape[1],
+                                    0,
+                                    dtype=torch.long,
+                                    device=gen_sequence.device)
+
+        # encode prompt and context
+        if prompt is not None and context is not None:
+            prompt_codes, context_codes = self.encode_target_and_context(
+                prompt, context)
+        elif prompt is not None:
+            prompt_codes = self.encode(prompt)
+        elif context is not None:
+            context_codes = self.encode(context)
+
+        # duplicate prompt, gen_sequence and context if using cfg
+        if use_cfg:
+            gen_sequence = torch.cat((gen_sequence, gen_sequence), dim=0)
+            prompt_codes = torch.cat((prompt_codes, prompt_codes), dim=0)
+            context_codes = torch.cat(
+                (context_codes,
+                 torch.full_like(context_codes, self.special_token)),
+                dim=0)
+
+        # insert encoded prompt into sequence
+        gen_sequence[..., :prompt_codes.shape[-1]] = prompt_codes
+
+        # interleave sequence
+        gen_sequence, gen_mask, pattern = self.interleave(
+            gen_sequence, keep_only_valid_steps=False)
+        start_offset: int = pattern.get_first_step_with_timesteps(
+            prompt_codes.shape[-1]) + context_codes.shape[-1]  # type: ignore
+
+        # prepend non-interleaved context codes to the sequence
+        gen_sequence = torch.cat((context_codes, gen_sequence), dim=-1)
+        gen_mask = torch.cat((torch.zeros(context_codes.shape[1:],
+                                          dtype=torch.bool,
+                                          device=gen_mask.device), gen_mask),
+                             dim=-1)
+
+        def decode_sequence_closure(sequence):
+            # remove context
+            sequence = sequence[..., context_codes.shape[-1]:]
+            # de-interleave
+            sequence = sequence.contiguous()
+            out_codes, _, out_mask = pattern.revert_pattern_sequence(
+                sequence, special_token=-1, keep_only_valid_steps=False)
+            return out_codes, out_mask
+
+        raise NotImplementedError("mask implementation missing")
+
+        return gen_sequence, gen_mask, start_offset, decode_sequence_closure
+
+
+class InterleavedContextPromptProcessor(ContextPromptProcessor):
+    uses_sep_token: bool = True
+    """
+        Concatenates context and target, separating them with a single timestep
+        of a "stop tensor", made of n_q special tokens.
+
+        The sequence (context, stop_token, target) is then interleaved.
+    """
+
+    def preprocess(self, batch: Dict[str, Any]):
+        target = batch["target"]
+        context = batch["context"]
+
+        # encode target and context
+        target_codes, context_codes = self.encode_target_and_context(
+            target, context)
+
+        # apply dropout to context
+        if self.training and self.context_dropout > 0:
+            drop = torch.rand(context_codes.shape[0]) < self.context_dropout
+            context_codes[drop] = torch.full_like(context_codes[drop],
+                                                  self.special_token)
+
+        # concatenate context and target with stop token
+        stop_tensor = torch.full(
+            (target_codes.shape[0], target_codes.shape[1], 1),
+            self.special_token + 1,
+            dtype=target_codes.dtype,
+            device=target_codes.device)
+        prompt = torch.cat((context_codes, stop_tensor, target_codes), dim=-1)
+
+        # interleave prompt
+        prompt, prompt_mask, pattern = self.interleave(
+            prompt, self.keep_only_valid_steps)
+
+        # all tokens behind the stop tokens are conditioning
+        cond_tokens_mask = torch.zeros_like(prompt, dtype=torch.bool)
+        n_q = prompt.shape[1]
+        context_len = context_codes.shape[-1]
+        for i in range(n_q):
+            cond_tokens_mask[:, i, range(context_len + i + 1)] = True
+        special_tokens_mask = prompt == self.special_token
+        invalid_cond_mask = special_tokens_mask * cond_tokens_mask
+        final_mask = (~invalid_cond_mask) * prompt_mask.repeat(
+            invalid_cond_mask.shape[0], 1, 1)
+
+        def decode_logits_closure(logits) -> Tuple[Tensor, Tensor]:
+            # remove interleaving pattern
+            deinterleaved_logits, logits_mask = self.deinterleave_logits(
+                logits, pattern)
+
+            # remove context
+            n_target_codes = target_codes.shape[-1]
+            deinterleaved_logits = deinterleaved_logits[:, :,
+                                                        -n_target_codes:, :]
+            logits_mask = logits_mask[:, :, -n_target_codes:]
+
+            return deinterleaved_logits, logits_mask
+
+        return prompt, final_mask, target_codes, decode_logits_closure
+
+    def prepare_for_generation(
+        self, prompt: Optional[Tensor], context: Optional[Tensor],
+        gen_sequence: Tensor, use_cfg: bool,
+        context_dropout_mask: Optional[Tensor]
+    ) -> Tuple[Tensor, Tensor, int, Callable[[Tensor], Tuple[Tensor, Tensor]]]:
+        """
+            Take prompt and context, prepare them, apply cfg, and insert them
+            in dummy sequence full of -1 for generation. Return: 
+            - gen_sequence
+            - gen_mask 
+            - start_offset 
+            - closure to decode sequence once filled with gen tokens
+        """
+        # encode prompt and context
+        if prompt is not None and context is not None:
+            prompt_codes, context_codes = self.encode_target_and_context(
+                prompt, context)
+            # context_mask = torch.ones_like(context_codes, dtype=torch.bool)
+        else:
+            if prompt is None:
+                prompt_codes = torch.empty(gen_sequence.shape[0],
+                                           gen_sequence.shape[1],
+                                           0,
+                                           dtype=torch.long,
+                                           device=gen_sequence.device)
+            else:
+                prompt_codes = self.encode(prompt)
+            if context is None:
+                context_codes = torch.full(
+                    (gen_sequence.shape[0], gen_sequence.shape[1], 0),
+                    self.special_token,
+                    dtype=torch.long,
+                    device=gen_sequence.device)
+                # context_mask = torch.zeros_like(context_codes, dtype=torch.bool)
+            else:
+                if isinstance(context, list):
+                    context_codes = self.encode_and_pad_list(
+                        context, self.special_token)
+                else:
+                    assert isinstance(context, Tensor)
+                    context_codes = self.encode(context)
+                # context_mask = torch.ones_like(context_codes, dtype=torch.bool)
+
+            if context_dropout_mask is not None:
+                if context_dropout_mask.shape != torch.Size(
+                    (context_codes.shape[0],)):
+                    raise ValueError(f"Context dropout mask should be of shape "
+                                     f"[{context_codes.shape[0]} but it is "
+                                     f"{context_dropout_mask.shape}]")
+                for i in range(len(context_dropout_mask)):
+                    if not context_dropout_mask[i]:
+                        context_codes[i] = torch.full_like(
+                            context_codes[i], self.special_token)
+                # context_codes = torch.where(
+                #     context_dropout_mask, context_codes,
+                #     torch.full_like(context_codes, self.special_token))
+
+        # duplicate prompt, gen_sequence and context if using cfg
+        if use_cfg:
+            gen_sequence = torch.cat((gen_sequence, gen_sequence), dim=0)
+            prompt_codes = torch.cat((prompt_codes, prompt_codes), dim=0)
+            context_codes = torch.cat(
+                (context_codes,
+                 torch.full_like(context_codes, self.special_token)),
+                dim=0)
+            # context_mask = torch.cat(
+            #     (context_mask, torch.zeros_like(context_mask,
+            #                                     dtype=torch.bool)),
+            #     dim=0)
+
+        # insert encoded prompt into sequence
+        gen_sequence[..., :prompt_codes.shape[-1]] = prompt_codes
+
+        # prepend context and stop token to gen_sequence
+        stop_tensor = torch.ones(
+            (gen_sequence.shape[0], gen_sequence.shape[1], 1),
+            dtype=gen_sequence.dtype,
+            device=gen_sequence.device) * self.special_token + 1
+        # stop_mask = torch.ones_like(stop_tensor, dtype=torch.bool)
+        # conditioning_mask = torch.cat(
+        #     (context_mask, stop_mask,
+        #      torch.ones_like(gen_sequence, dtype=torch.bool)),
+        #     dim=-1)
+        gen_sequence = torch.cat((context_codes, stop_tensor, gen_sequence),
+                                 dim=-1)
+
+        # interleave sequence (keep_only_valid_steps always false for inference)
+        gen_sequence, gen_mask, pattern = self.interleave(
+            gen_sequence, keep_only_valid_steps=False)
+
+        # interleaved_conditioning_mask, mask_mask, _ = self.interleave(
+        #     conditioning_mask, keep_only_valid_steps=False)
+
+        # assert gen_sequence.shape == interleaved_conditioning_mask.shape
+
+        # interleaved_conditioning_mask[~gen_mask] = 0
+        # interleaved_conditioning_mask[~mask_mask] = 0
+
+        # compute first offset of the sequence to generate
+        start_offset: int = pattern.get_first_step_with_timesteps(  # type: ignore
+            prompt_codes.shape[-1] + context_codes.shape[-1] + 1)
+
+        # set gen_mask of stop token to false
+        n_q = gen_sequence.shape[1]
+        # gen_mask[range(n_q),
+        #          range(start_offset - 1, start_offset - 1 + n_q)] = False
+
+        # all tokens behind the stop token are conditioning
+        cond_tokens_mask = torch.zeros_like(gen_sequence, dtype=torch.bool)
+        for i in range(n_q):
+            cond_tokens_mask[:, i, range(start_offset + i - 1)] = True
+        special_tokens_mask = gen_sequence == self.special_token
+        invalid_cond_mask = special_tokens_mask * cond_tokens_mask
+
+        final_mask = (~invalid_cond_mask) * gen_mask.repeat(
+            invalid_cond_mask.shape[0], 1, 1)
+
+        def decode_sequence_closure(sequence):
+            # deinterleave
+            sequence = sequence.contiguous()
+            out_codes, _, out_mask = pattern.revert_pattern_sequence(
+                sequence, special_token=-1, keep_only_valid_steps=False)
+            # remove context and stop token
+            out_codes = out_codes[..., (context_codes.shape[-1] + 1):]
+            return out_codes, out_mask
+
+        backwards_compatible_mode = False
+        if backwards_compatible_mode:
+            return gen_sequence, gen_mask, start_offset, decode_sequence_closure
+
+        return gen_sequence, final_mask, start_offset, decode_sequence_closure
+
+
+def parallel_vs_sequential_test():
+
+    from time import time
+    import hyperparameters as hp
+    from models.encodec import EncodecModel
+    device = torch.device("cuda")
+
+    encodec_params = hp.pretrained_encodec_meta_32khz_params
+    encodec = EncodecModel.from_params(encodec_params).eval().to(device)
+
+    processor = StraightContextPromptProcessor(encodec_model=encodec,
+                                               special_token=2048,
+                                               keep_only_valid_steps=False,
+                                               context_dropout=0.9).eval()
+
+    # quick benchmark for parallel vs sequential encoding performance
+    batch_parallel = {
+        "target": torch.rand(4, 1, 320_000).to(device),
+        "context": torch.rand(4, 1, 320_000).to(device),
+    }
+    t0 = time()
+    out1 = processor.preprocess(batch_parallel)
+    t1 = time()
+    print(f"parallel encoding: {round((t1 - t0) * 1000, 2)} ms")
+
+    batch_sequential = {
+        "target": torch.rand(4, 1, 320_000).to(device),
+        "context": torch.rand(4, 1, 310_000).to(device),
+    }
+    t0 = time()
+    out2 = processor.preprocess(batch_sequential)
+    t1 = time()
+    print(f"sequential encoding: {round((t1 - t0) * 1000, 2)} ms")
+
+    batch_parallel = {
+        "target": torch.rand(4, 1, 320_000).to(device),
+        "context": torch.rand(4, 1, 320_000).to(device),
+    }
+    t0 = time()
+    out3 = processor.preprocess(batch_parallel)
+    t1 = time()
+    print(f"parallel encoding: {round((t1 - t0) * 1000, 2)} ms")
+
+    batch_sequential = {
+        "target": torch.rand(4, 1, 320_000).to(device),
+        "context": torch.rand(4, 1, 310_000).to(device),
+    }
+    t0 = time()
+    out4 = processor.preprocess(batch_sequential)
+    t1 = time()
+    print(f"sequential encoding: {round((t1 - t0) * 1000, 2)} ms")
+
+    batch_parallel = {
+        "target": torch.rand(4, 1, 320_000).to(device),
+        "context": torch.rand(4, 1, 320_000).to(device),
+    }
+    t0 = time()
+    out5 = processor.preprocess(batch_parallel)
+    t1 = time()
+    print(f"parallel encoding: {round((t1 - t0) * 1000, 2)} ms")
+
+    batch_sequential = {
+        "target": torch.rand(4, 1, 320_000).to(device),
+        "context": torch.rand(4, 1, 310_000).to(device),
+    }
+    t0 = time()
+    out6 = processor.preprocess(batch_sequential)
+    t1 = time()
+    print(f"sequential encoding: {round((t1 - t0) * 1000, 2)} ms")
+
+
+if __name__ == "__main__":
+    import hyperparameters as hp
+    from models.encodec import EncodecModel
+    device = torch.device("cuda")
+
+    encodec_params = hp.pretrained_encodec_meta_32khz_params
+    encodec = EncodecModel.from_params(encodec_params).eval().to(device)
+
+    processor = InterleavedContextPromptProcessor(encodec_model=encodec,
+                                                  special_token=2048,
+                                                  keep_only_valid_steps=False,
+                                                  context_dropout=0.9).eval()

conditioning/t5embedder.py

@@ -0,0 +1,211 @@
+import logging
+from transformers import T5Tokenizer, T5EncoderModel
+from typing import List, Optional, Dict, Sequence, Tuple
+import random
+import torch
+from torch import Tensor
+
+from conditioning.embedder import Embedder, LinearProjectionEmbedder
+from conditioning.embedded_condition import EmbeddedCondition
+
+
+class T5Embedder(LinearProjectionEmbedder):
+    """T5-based TextConditioner.
+
+    Args:
+        name (str): Name of the T5 model.
+        output_dim (int): Output dim of the conditioner.
+        finetune (bool): Whether to fine-tune T5 at train time.
+        word_dropout (float, optional): Word dropout probability.
+        normalize_text (bool, optional): Whether to apply text normalization.
+    """
+    MODELS = [
+        "t5-small", "t5-base", "t5-large", "t5-3b", "t5-11b",
+        "google/flan-t5-small", "google/flan-t5-base", "google/flan-t5-large",
+        "google/flan-t5-xl", "google/flan-t5-xxl"
+    ]
+    MODELS_DIMS = {
+        "t5-small": 512,
+        "t5-base": 768,
+        "t5-large": 1024,
+        "t5-3b": 1024,
+        "t5-11b": 1024,
+        "google/flan-t5-small": 512,
+        "google/flan-t5-base": 768,
+        "google/flan-t5-large": 1024,
+        "google/flan-t5-3b": 1024,
+        "google/flan-t5-11b": 1024,
+    }
+
+    def __init__(self,
+                 embedding_dim: int,
+                 t5_on_cpu: bool,
+                 name: str = "t5-base",
+                 finetune: bool = False,
+                 word_dropout: float = 0.3):
+        assert name in self.MODELS, f"Unrecognized t5 model name (should in {self.MODELS})"
+        super().__init__(self.MODELS_DIMS[name], embedding_dim)
+        self.name = name
+        self.finetune = finetune
+        self.word_dropout = word_dropout
+        self.t5_on_cpu: bool = t5_on_cpu
+        if self.t5_on_cpu and finetune:
+            raise ValueError("Can't finetune t5 if it's locked on cpu")
+
+        # Let's disable logging temporarily because T5 will vomit some errors otherwise.
+        # thanks https://gist.github.com/simon-weber/7853144
+        previous_level = logging.root.manager.disable
+        logging.disable(logging.ERROR)
+        self.t5_tokenizer = T5Tokenizer.from_pretrained(
+            name,
+            clean_up_tokenization_spaces=False,
+        )
+
+        t5 = T5EncoderModel.from_pretrained(name).train(mode=finetune)
+
+        if self.t5_on_cpu:
+            t5 = t5.cpu()
+            self.__dict__["t5"] = t5
+        else:
+            self.t5 = t5
+
+        if not self.finetune:
+            for p in self.t5.parameters():
+                p.requires_grad = False
+
+        # if not self.finetune:
+        #     if self.t5_on_cpu:
+        #         t5 = t5.cpu()
+        #     self.__dict__["t5"] = t5.eval()
+        #     for p in self.t5.parameters():
+        #         p.requires_grad = False
+        # else:
+        #     self.t5 = t5
+
+        # with warnings.catch_warnings():
+
+        #     warnings.simplefilter("ignore")
+        #     try:
+        #         self.t5_tokenizer = T5Tokenizer.from_pretrained(name)
+        #         t5 = T5EncoderModel.from_pretrained(name).train(mode=finetune)
+        #     finally:
+        #         logging.disable(previous_level)
+        # if finetune:
+        #     self.t5 = t5
+        # else:
+        #     # this makes sure that the t5 models is not part
+        #     # of the saved checkpoint
+        #     self.__dict__['t5'] = t5
+
+        # self.normalize_text = normalize_text
+        # if normalize_text:
+        #     self.text_normalizer = WhiteSpaceTokenizer(1,
+        #                                                lemma=True,
+        #                                                stopwords=True)
+
+    # def to(self, *args, **kwargs):
+    #     return super().to("cpu")
+
+    # def cuda(self, *args, **kwargs):
+    #     return self.to("cpu")
+
+    def tokenize(self, x: Sequence[Optional[str]]) -> Dict[str, torch.Tensor]:
+        # if current sample doesn't have a certain attribute, replace with empty string
+        entries: List[str] = [xi if xi is not None else "" for xi in x]
+        # if self.normalize_text:
+        #     _, _, entries = self.text_normalizer(  # type: ignore
+        #         entries, return_text=True)
+        if self.word_dropout > 0. and self.training:
+            new_entries = []
+            for entry in entries:
+                words = [
+                    word for word in entry.split(" ")
+                    if random.random() >= self.word_dropout
+                ]
+                new_entries.append(" ".join(words))
+            entries = new_entries
+
+        empty_idx = torch.LongTensor(
+            [i for i, xi in enumerate(entries) if xi == ""])
+
+        inputs = self.t5_tokenizer(entries, return_tensors='pt',
+                                   padding=True).to(
+                                       next(iter(self.parameters())).device)
+        mask = inputs['attention_mask']
+        mask[empty_idx, :] = 0  # zero-out index where the input is non-existant
+        return inputs
+
+    def forward(self,
+                descriptions: List[str],
+                duplicate_for_cfg: bool = False) -> EmbeddedCondition:
+
+        if duplicate_for_cfg:
+            descriptions = descriptions + ([""] * len(descriptions))
+
+        tokenized = self.tokenize(descriptions)
+        embedding, mask = self.embed(tokenized)
+        mask = mask.bool()
+        return EmbeddedCondition(embedding, mask)
+
+    def embed(self, inputs: Dict[str, torch.Tensor]) -> Tuple[Tensor, Tensor]:
+        if not self.finetune:
+            self.t5.eval()
+        mask = inputs['attention_mask']
+        with torch.set_grad_enabled(self.finetune):
+            if self.t5_on_cpu:
+                inputs = {k: v.to("cpu") for k, v in inputs.items()}
+            with torch.autocast(device_type="cuda", enabled=False):
+                embeds = self.t5(**inputs).last_hidden_state
+        embeds = self.output_proj(embeds.to(self.output_proj.weight))
+        embeds = (embeds * mask.unsqueeze(-1).to(embeds))
+        return embeds, mask
+
+    def null_condition(self, batch_size: int):
+        return EmbeddedCondition(
+            torch.zeros(batch_size,
+                        1,
+                        self.embedding_dim,
+                        dtype=torch.float32,
+                        device=self.output_proj.weight.device),
+            torch.ones(batch_size,
+                       1,
+                       dtype=torch.bool,
+                       device=self.output_proj.weight.device))
+
+
+class T5EmbedderCPU(T5Embedder):
+
+    def __init__(self, embedding_dim: int):
+        super().__init__(embedding_dim,
+                         t5_on_cpu=True,
+                         name="t5-base",
+                         finetune=False,
+                         word_dropout=0.3)
+
+
+class T5EmbedderGPU(T5Embedder):
+
+    def __init__(self, embedding_dim: int):
+        super().__init__(embedding_dim,
+                         t5_on_cpu=False,
+                         name="t5-base",
+                         finetune=False,
+                         word_dropout=0.3)
+
+
+if __name__ == "__main__":
+    from utils.inspection import print_params
+    from time import time
+
+    t5 = T5EmbedderCPU(1024).eval().cpu()
+    # print_params(t5, 1, False)
+
+    descriptions = ["the quick"]
+
+    embedded_desc = t5(descriptions)
+    # print(embedded_desc)
+    # print(embedded_desc.data)
+    # print(embedded_desc.mask)
+
+    # print(embedded_desc.data[..., -1, ...])
+    # print(embedded_desc.data[0, -1, ...])

config.py

@@ -0,0 +1,77 @@
+from pathlib import Path
+import os
+from typing import Optional
+import os
+
+ENTITY: str = ""
+PROJECT: str = ""
+
+ROOT = Path(__file__).parent
+LOCAL_WEIGHTS_DIR: Path = ROOT / "weights"
+LOCAL_CKP_DIR: Path = ROOT / "checkpoints"
+AUDIO_DIR: Path = ROOT / "audio"
+EVAL_DIR: Path = ROOT / "eval"
+EXP_DIR: Path = ROOT / "experiments"
+CKP_DIR = LOCAL_CKP_DIR
+
+
+class ConfigurationError(ValueError):
+    ...
+
+
+def first_existing(*paths: Path | str) -> Optional[Path]:
+    for path in paths:
+        if Path(path).exists():
+            return Path(path)
+
+
+def moises_path() -> Path:
+    path = first_existing(Path.home() / "dev/dataset/moisesdb/moisesdb_v0.1",
+                          Path.home() / "dev/datasets/moisesdb/moisesdb_v0.1",
+                          Path.home() / "lag-data/lag-moisesdb",
+                          ROOT / "datasets/moisesdb")
+    if path is None:
+        raise RuntimeError("I can't find moisesdb")
+    return path
+
+
+def mus_path() -> Path:
+    path = first_existing(
+        Path.home() / "dev/dataset/moisesdb/musdb",
+        Path.home() / "datasets/moisesdb/musdb",
+        Path.home() / "lag-data/musdb",
+        ROOT / "datasets/musdb",
+    )
+    if path is None:
+        raise RuntimeError("I can't find musdb")
+    return path
+
+
+def mixdata_path() -> Path:
+    path = first_existing(
+        Path.home() / "dev/datasets/moisesdb",
+        Path.home() / "datasets/moisesdb",
+        Path.home() / "lag-data",
+        Path.home() / "datasets/lag-data",
+    )
+    if path is None:
+        raise RuntimeError("I can't find the mixed dataset path")
+    return path
+
+
+def weights_dir() -> Path:
+    return LOCAL_WEIGHTS_DIR
+
+
+def output_dir() -> Path:
+    return CKP_DIR
+
+
+def shutdown():
+    print(f"Shutting down myself 💀")
+    os.system("sudo shutdown now")
+
+
+import torch
+
+torch.set_float32_matmul_precision("medium")

data/auto_labelling.py

@@ -0,0 +1,147 @@
+#%% Imports
+from pathlib import Path
+import subprocess as sp
+import essentia.standard as es
+import config as cfg
+import numpy as np
+import librosa
+from multipledispatch import dispatch
+from torch import Tensor
+import torchaudio
+
+from data.labels import GENRE_LABELS, MOOD_THEME_CLASSES, INSTRUMENT_CLASSES
+from utils import audio as audio_utils
+import essentia
+
+#%% Download models
+if False:
+    sp.call([
+        "curl",
+        "https://essentia.upf.edu/models/classification-heads/genre_discogs400/genre_discogs400-discogs-effnet-1.pb",
+        "--output", "genre_discogs400-discogs-effnet-1.pb"
+    ])
+    sp.call([
+        "curl",
+        "https://essentia.upf.edu/models/feature-extractors/discogs-effnet/discogs-effnet-bs64-1.pb",
+        "--output", "discogs-effnet-bs64-1.pb"
+    ])
+    sp.call([
+        "curl",
+        "https://essentia.upf.edu/models/classification-heads/mtg_jamendo_moodtheme/mtg_jamendo_moodtheme-discogs-effnet-1.pb",
+        "--output", "mtg_jamendo_moodtheme-discogs-effnet-1.pb"
+    ])
+    sp.call([
+        "curl",
+        "https://essentia.upf.edu/models/classification-heads/mtg_jamendo_instrument/mtg_jamendo_instrument-discogs-effnet-1.pb",
+        "--output", "mtg_jamendo_instrument-discogs-effnet-1.pb"
+    ])
+
+
+def filter_predictions(predictions, class_list, threshold=0.1):
+    predictions_mean = np.mean(predictions, axis=0)
+    sorted_indices = np.argsort(predictions_mean)[::-1]
+    filtered_indices = [
+        i for i in sorted_indices if predictions_mean[i] > threshold
+    ]
+    filtered_labels = [class_list[i] for i in filtered_indices]
+    filtered_values = [predictions_mean[i] for i in filtered_indices]
+    return filtered_labels, filtered_values
+
+
+def make_comma_separated_unique(tags):
+    seen_tags = set()
+    result = []
+    for tag in ', '.join(tags).split(', '):
+        if tag not in seen_tags:
+            result.append(tag)
+            seen_tags.add(tag)
+    return ', '.join(result)
+
+
+@dispatch(Path)
+def get_audio_features(audio_filename: Path):  # type: ignore
+    audio = audio_utils.load_audio(audio_filename, 16_000, False).squeeze()
+    # audio = es.MonoLoader(filename=str(audio_filename),
+    #                       sampleRate=16000,
+    #                       resampleQuality=4)()
+
+    return get_audio_features(audio, 16_000)
+
+
+@dispatch(Tensor, int, Path)
+def get_audio_features(audio: Tensor, sr: int,
+                       models_dir: Path):  # type: ignore
+
+    essentia.log.infoActive = False
+    audio = audio_utils.to_mono(audio)
+    audio = torchaudio.functional.resample(audio, sr, 16_000).squeeze()
+    audio = audio.numpy()
+
+    embedding_model = es.TensorflowPredictEffnetDiscogs(
+        graphFilename=str(models_dir / "discogs-effnet-bs64-1.pb"),
+        output="PartitionedCall:1")
+    embeddings = embedding_model(audio)
+
+    result_dict = {}
+
+    # Predicting genres
+    genre_model = es.TensorflowPredict2D(
+        graphFilename=str(models_dir / "genre_discogs400-discogs-effnet-1.pb"),
+        input="serving_default_model_Placeholder",
+        output="PartitionedCall:0")
+    predictions = genre_model(embeddings)
+
+    filtered_labels, _ = filter_predictions(predictions, GENRE_LABELS)
+    filtered_labels = ', '.join(filtered_labels).replace("---",
+                                                         ", ").split(', ')
+    result_dict['genres'] = make_comma_separated_unique(filtered_labels)
+
+    # Predicting mood/theme
+    mood_model = es.TensorflowPredict2D(
+        graphFilename=str(models_dir /
+                          "mtg_jamendo_moodtheme-discogs-effnet-1.pb"))
+    predictions = mood_model(embeddings)
+    filtered_labels, _ = filter_predictions(predictions,
+                                            MOOD_THEME_CLASSES,
+                                            threshold=0.05)
+    result_dict['moods'] = make_comma_separated_unique(filtered_labels)
+
+    bpm, key = get_bpm_key(audio, sr)
+
+    result_dict["bpm"] = bpm
+    result_dict["key"] = key
+
+    # Predicting instruments
+    # instrument_model = es.TensorflowPredict2D(
+    #     graphFilename="mtg_jamendo_instrument-discogs-effnet-1.pb")
+    # predictions = instrument_model(embeddings)
+    # filtered_labels, _ = filter_predictions(predictions, INSTRUMENT_CLASSES)
+    # result_dict['instruments'] = filtered_labels
+
+    return result_dict
+
+
+@dispatch(Path)
+def get_bpm_key(audio_filename):
+    y, sr = librosa.load(str(audio_filename))
+    get_bpm_key(y, sr)
+
+
+@dispatch(np.ndarray, int)
+def get_bpm_key(audio: np.ndarray, sr: int):
+    tempo, _ = librosa.beat.beat_track(y=audio, sr=sr)
+    tempo = round(tempo[0])
+    chroma = librosa.feature.chroma_stft(y=audio, sr=sr)
+    key = np.argmax(np.sum(chroma, axis=1))
+    key = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B'][key]
+    length = librosa.get_duration(y=audio, sr=sr)
+
+    return tempo, key
+
+
+#%% Test on demo audio
+if __name__ == "__main__":
+    audio_filename = cfg.AUDIO_DIR / "cake.wav"
+    features = get_audio_features(audio_filename)
+
+# %

data/dataset_mixed.py

@@ -0,0 +1,79 @@
+import torch
+import torch.utils.data
+from pathlib import Path
+
+from data.stemmed_dataset import StemmedDataset
+import config as cfg
+
+
+class MixDataset(torch.utils.data.Dataset):
+
+    def __init__(self, root_dir: Path, *args, **kwargs):
+
+        for dirname in [
+                "moisesdb_v0.1", "lag-moisesdb", "lag_moisesdb", "moisesdb",
+                "moises"
+        ]:
+            if (root_dir / dirname).exists():
+                self.root_dir_moises: Path = root_dir / dirname
+                break
+        else:
+            raise FileNotFoundError(
+                f"Couldn't find subdirectory for moisesdb under {root_dir}")
+
+        self.root_dir_mus: Path = root_dir / "musdb"
+
+        self.moisesdb_dataset: StemmedDataset = StemmedDataset(
+            self.root_dir_moises,
+            *args,
+            **kwargs,
+        )
+        self.musdb_dataset: StemmedDataset = StemmedDataset(
+            self.root_dir_mus,
+            *args,
+            **kwargs,
+        )
+
+        # self.n_samples = len(self.moisesdb_dataset) + len(self.musdb_dataset)
+        self.n_samples_moises = len(self.moisesdb_dataset)
+        self.n_samples_mus = len(self.musdb_dataset)
+
+    def __len__(self):
+        return self.n_samples_moises + self.n_samples_mus
+
+    def __getitem__(self, x):
+        if x < self.n_samples_moises:
+            return self.moisesdb_dataset[x]
+        else:
+            return self.musdb_dataset[x - self.n_samples_moises]
+
+
+if __name__ == "__main__":
+    from data.stem import Stem
+
+    moises_root = cfg.moises_path()
+    mus_root = cfg.mus_path()
+
+    stems = {
+        Stem.DRUMS, Stem.GUITAR, Stem.BASS, Stem.PIANO, Stem.KEYBOARD,
+        Stem.STRINGS
+    }
+
+    d = MixDataset(
+        cfg.mixdata_path(),
+        stems,
+        target_stem=Stem.DRUMS,
+        single_stem=True,
+        add_click=False,
+        bpm_in_caption=False,
+        sync_chunks=False,
+        train=True,
+        sample_rate=32_000,
+        chunk_size_samples=32_000 * 10,
+        speed_transform_p=1,
+        pitch_transform_p=1,
+        stereo=False,
+        n_samples_per_epoch=2000,
+    )
+
+    sample = next(iter(d))

data/labels.py

@@ -0,0 +1,424 @@
+# @title metadata (labels) for essentia
+
+GENRE_LABELS = [
+    "Blues---Boogie Woogie",
+    "Blues---Chicago Blues",
+    "Blues---Country Blues",
+    "Blues---Delta Blues",
+    "Blues---Electric Blues",
+    "Blues---Harmonica Blues",
+    "Blues---Jump Blues",
+    "Blues---Louisiana Blues",
+    "Blues---Modern Electric Blues",
+    "Blues---Piano Blues",
+    "Blues---Rhythm & Blues",
+    "Blues---Texas Blues",
+    "Brass & Military---Brass Band",
+    "Brass & Military---Marches",
+    "Brass & Military---Military",
+    "Children's---Educational",
+    "Children's---Nursery Rhymes",
+    "Children's---Story",
+    "Classical---Baroque",
+    "Classical---Choral",
+    "Classical---Classical",
+    "Classical---Contemporary",
+    "Classical---Impressionist",
+    "Classical---Medieval",
+    "Classical---Modern",
+    "Classical---Neo-Classical",
+    "Classical---Neo-Romantic",
+    "Classical---Opera",
+    "Classical---Post-Modern",
+    "Classical---Renaissance",
+    "Classical---Romantic",
+    "Electronic---Abstract",
+    "Electronic---Acid",
+    "Electronic---Acid House",
+    "Electronic---Acid Jazz",
+    "Electronic---Ambient",
+    "Electronic---Bassline",
+    "Electronic---Beatdown",
+    "Electronic---Berlin-School",
+    "Electronic---Big Beat",
+    "Electronic---Bleep",
+    "Electronic---Breakbeat",
+    "Electronic---Breakcore",
+    "Electronic---Breaks",
+    "Electronic---Broken Beat",
+    "Electronic---Chillwave",
+    "Electronic---Chiptune",
+    "Electronic---Dance-pop",
+    "Electronic---Dark Ambient",
+    "Electronic---Darkwave",
+    "Electronic---Deep House",
+    "Electronic---Deep Techno",
+    "Electronic---Disco",
+    "Electronic---Disco Polo",
+    "Electronic---Donk",
+    "Electronic---Downtempo",
+    "Electronic---Drone",
+    "Electronic---Drum n Bass",
+    "Electronic---Dub",
+    "Electronic---Dub Techno",
+    "Electronic---Dubstep",
+    "Electronic---Dungeon Synth",
+    "Electronic---EBM",
+    "Electronic---Electro",
+    "Electronic---Electro House",
+    "Electronic---Electroclash",
+    "Electronic---Euro House",
+    "Electronic---Euro-Disco",
+    "Electronic---Eurobeat",
+    "Electronic---Eurodance",
+    "Electronic---Experimental",
+    "Electronic---Freestyle",
+    "Electronic---Future Jazz",
+    "Electronic---Gabber",
+    "Electronic---Garage House",
+    "Electronic---Ghetto",
+    "Electronic---Ghetto House",
+    "Electronic---Glitch",
+    "Electronic---Goa Trance",
+    "Electronic---Grime",
+    "Electronic---Halftime",
+    "Electronic---Hands Up",
+    "Electronic---Happy Hardcore",
+    "Electronic---Hard House",
+    "Electronic---Hard Techno",
+    "Electronic---Hard Trance",
+    "Electronic---Hardcore",
+    "Electronic---Hardstyle",
+    "Electronic---Hi NRG",
+    "Electronic---Hip Hop",
+    "Electronic---Hip-House",
+    "Electronic---House",
+    "Electronic---IDM",
+    "Electronic---Illbient",
+    "Electronic---Industrial",
+    "Electronic---Italo House",
+    "Electronic---Italo-Disco",
+    "Electronic---Italodance",
+    "Electronic---Jazzdance",
+    "Electronic---Juke",
+    "Electronic---Jumpstyle",
+    "Electronic---Jungle",
+    "Electronic---Latin",
+    "Electronic---Leftfield",
+    "Electronic---Makina",
+    "Electronic---Minimal",
+    "Electronic---Minimal Techno",
+    "Electronic---Modern Classical",
+    "Electronic---Musique Concrète",
+    "Electronic---Neofolk",
+    "Electronic---New Age",
+    "Electronic---New Beat",
+    "Electronic---New Wave",
+    "Electronic---Noise",
+    "Electronic---Nu-Disco",
+    "Electronic---Power Electronics",
+    "Electronic---Progressive Breaks",
+    "Electronic---Progressive House",
+    "Electronic---Progressive Trance",
+    "Electronic---Psy-Trance",
+    "Electronic---Rhythmic Noise",
+    "Electronic---Schranz",
+    "Electronic---Sound Collage",
+    "Electronic---Speed Garage",
+    "Electronic---Speedcore",
+    "Electronic---Synth-pop",
+    "Electronic---Synthwave",
+    "Electronic---Tech House",
+    "Electronic---Tech Trance",
+    "Electronic---Techno",
+    "Electronic---Trance",
+    "Electronic---Tribal",
+    "Electronic---Tribal House",
+    "Electronic---Trip Hop",
+    "Electronic---Tropical House",
+    "Electronic---UK Garage",
+    "Electronic---Vaporwave",
+    "Folk, World, & Country---African",
+    "Folk, World, & Country---Bluegrass",
+    "Folk, World, & Country---Cajun",
+    "Folk, World, & Country---Canzone Napoletana",
+    "Folk, World, & Country---Catalan Music",
+    "Folk, World, & Country---Celtic",
+    "Folk, World, & Country---Country",
+    "Folk, World, & Country---Fado",
+    "Folk, World, & Country---Flamenco",
+    "Folk, World, & Country---Folk",
+    "Folk, World, & Country---Gospel",
+    "Folk, World, & Country---Highlife",
+    "Folk, World, & Country---Hillbilly",
+    "Folk, World, & Country---Hindustani",
+    "Folk, World, & Country---Honky Tonk",
+    "Folk, World, & Country---Indian Classical",
+    "Folk, World, & Country---Laïkó",
+    "Folk, World, & Country---Nordic",
+    "Folk, World, & Country---Pacific",
+    "Folk, World, & Country---Polka",
+    "Folk, World, & Country---Raï",
+    "Folk, World, & Country---Romani",
+    "Folk, World, & Country---Soukous",
+    "Folk, World, & Country---Séga",
+    "Folk, World, & Country---Volksmusik",
+    "Folk, World, & Country---Zouk",
+    "Folk, World, & Country---Éntekhno",
+    "Funk / Soul---Afrobeat",
+    "Funk / Soul---Boogie",
+    "Funk / Soul---Contemporary R&B",
+    "Funk / Soul---Disco",
+    "Funk / Soul---Free Funk",
+    "Funk / Soul---Funk",
+    "Funk / Soul---Gospel",
+    "Funk / Soul---Neo Soul",
+    "Funk / Soul---New Jack Swing",
+    "Funk / Soul---P.Funk",
+    "Funk / Soul---Psychedelic",
+    "Funk / Soul---Rhythm & Blues",
+    "Funk / Soul---Soul",
+    "Funk / Soul---Swingbeat",
+    "Funk / Soul---UK Street Soul",
+    "Hip Hop---Bass Music",
+    "Hip Hop---Boom Bap",
+    "Hip Hop---Bounce",
+    "Hip Hop---Britcore",
+    "Hip Hop---Cloud Rap",
+    "Hip Hop---Conscious",
+    "Hip Hop---Crunk",
+    "Hip Hop---Cut-up/DJ",
+    "Hip Hop---DJ Battle Tool",
+    "Hip Hop---Electro",
+    "Hip Hop---G-Funk",
+    "Hip Hop---Gangsta",
+    "Hip Hop---Grime",
+    "Hip Hop---Hardcore Hip-Hop",
+    "Hip Hop---Horrorcore",
+    "Hip Hop---Instrumental",
+    "Hip Hop---Jazzy Hip-Hop",
+    "Hip Hop---Miami Bass",
+    "Hip Hop---Pop Rap",
+    "Hip Hop---Ragga HipHop",
+    "Hip Hop---RnB/Swing",
+    "Hip Hop---Screw",
+    "Hip Hop---Thug Rap",
+    "Hip Hop---Trap",
+    "Hip Hop---Trip Hop",
+    "Hip Hop---Turntablism",
+    "Jazz---Afro-Cuban Jazz",
+    "Jazz---Afrobeat",
+    "Jazz---Avant-garde Jazz",
+    "Jazz---Big Band",
+    "Jazz---Bop",
+    "Jazz---Bossa Nova",
+    "Jazz---Contemporary Jazz",
+    "Jazz---Cool Jazz",
+    "Jazz---Dixieland",
+    "Jazz---Easy Listening",
+    "Jazz---Free Improvisation",
+    "Jazz---Free Jazz",
+    "Jazz---Fusion",
+    "Jazz---Gypsy Jazz",
+    "Jazz---Hard Bop",
+    "Jazz---Jazz-Funk",
+    "Jazz---Jazz-Rock",
+    "Jazz---Latin Jazz",
+    "Jazz---Modal",
+    "Jazz---Post Bop",
+    "Jazz---Ragtime",
+    "Jazz---Smooth Jazz",
+    "Jazz---Soul-Jazz",
+    "Jazz---Space-Age",
+    "Jazz---Swing",
+    "Latin---Afro-Cuban",
+    "Latin---Baião",
+    "Latin---Batucada",
+    "Latin---Beguine",
+    "Latin---Bolero",
+    "Latin---Boogaloo",
+    "Latin---Bossanova",
+    "Latin---Cha-Cha",
+    "Latin---Charanga",
+    "Latin---Compas",
+    "Latin---Cubano",
+    "Latin---Cumbia",
+    "Latin---Descarga",
+    "Latin---Forró",
+    "Latin---Guaguancó",
+    "Latin---Guajira",
+    "Latin---Guaracha",
+    "Latin---MPB",
+    "Latin---Mambo",
+    "Latin---Mariachi",
+    "Latin---Merengue",
+    "Latin---Norteño",
+    "Latin---Nueva Cancion",
+    "Latin---Pachanga",
+    "Latin---Porro",
+    "Latin---Ranchera",
+    "Latin---Reggaeton",
+    "Latin---Rumba",
+    "Latin---Salsa",
+    "Latin---Samba",
+    "Latin---Son",
+    "Latin---Son Montuno",
+    "Latin---Tango",
+    "Latin---Tejano",
+    "Latin---Vallenato",
+    "Non-Music---Audiobook",
+    "Non-Music---Comedy",
+    "Non-Music---Dialogue",
+    "Non-Music---Education",
+    "Non-Music---Field Recording",
+    "Non-Music---Interview",
+    "Non-Music---Monolog",
+    "Non-Music---Poetry",
+    "Non-Music---Political",
+    "Non-Music---Promotional",
+    "Non-Music---Radioplay",
+    "Non-Music---Religious",
+    "Non-Music---Spoken Word",
+    "Pop---Ballad",
+    "Pop---Bollywood",
+    "Pop---Bubblegum",
+    "Pop---Chanson",
+    "Pop---City Pop",
+    "Pop---Europop",
+    "Pop---Indie Pop",
+    "Pop---J-pop",
+    "Pop---K-pop",
+    "Pop---Kayōkyoku",
+    "Pop---Light Music",
+    "Pop---Music Hall",
+    "Pop---Novelty",
+    "Pop---Parody",
+    "Pop---Schlager",
+    "Pop---Vocal",
+    "Reggae---Calypso",
+    "Reggae---Dancehall",
+    "Reggae---Dub",
+    "Reggae---Lovers Rock",
+    "Reggae---Ragga",
+    "Reggae---Reggae",
+    "Reggae---Reggae-Pop",
+    "Reggae---Rocksteady",
+    "Reggae---Roots Reggae",
+    "Reggae---Ska",
+    "Reggae---Soca",
+    "Rock---AOR",
+    "Rock---Acid Rock",
+    "Rock---Acoustic",
+    "Rock---Alternative Rock",
+    "Rock---Arena Rock",
+    "Rock---Art Rock",
+    "Rock---Atmospheric Black Metal",
+    "Rock---Avantgarde",
+    "Rock---Beat",
+    "Rock---Black Metal",
+    "Rock---Blues Rock",
+    "Rock---Brit Pop",
+    "Rock---Classic Rock",
+    "Rock---Coldwave",
+    "Rock---Country Rock",
+    "Rock---Crust",
+    "Rock---Death Metal",
+    "Rock---Deathcore",
+    "Rock---Deathrock",
+    "Rock---Depressive Black Metal",
+    "Rock---Doo Wop",
+    "Rock---Doom Metal",
+    "Rock---Dream Pop",
+    "Rock---Emo",
+    "Rock---Ethereal",
+    "Rock---Experimental",
+    "Rock---Folk Metal",
+    "Rock---Folk Rock",
+    "Rock---Funeral Doom Metal",
+    "Rock---Funk Metal",
+    "Rock---Garage Rock",
+    "Rock---Glam",
+    "Rock---Goregrind",
+    "Rock---Goth Rock",
+    "Rock---Gothic Metal",
+    "Rock---Grindcore",
+    "Rock---Grunge",
+    "Rock---Hard Rock",
+    "Rock---Hardcore",
+    "Rock---Heavy Metal",
+    "Rock---Indie Rock",
+    "Rock---Industrial",
+    "Rock---Krautrock",
+    "Rock---Lo-Fi",
+    "Rock---Lounge",
+    "Rock---Math Rock",
+    "Rock---Melodic Death Metal",
+    "Rock---Melodic Hardcore",
+    "Rock---Metalcore",
+    "Rock---Mod",
+    "Rock---Neofolk",
+    "Rock---New Wave",
+    "Rock---No Wave",
+    "Rock---Noise",
+    "Rock---Noisecore",
+    "Rock---Nu Metal",
+    "Rock---Oi",
+    "Rock---Parody",
+    "Rock---Pop Punk",
+    "Rock---Pop Rock",
+    "Rock---Pornogrind",
+    "Rock---Post Rock",
+    "Rock---Post-Hardcore",
+    "Rock---Post-Metal",
+    "Rock---Post-Punk",
+    "Rock---Power Metal",
+    "Rock---Power Pop",
+    "Rock---Power Violence",
+    "Rock---Prog Rock",
+    "Rock---Progressive Metal",
+    "Rock---Psychedelic Rock",
+    "Rock---Psychobilly",
+    "Rock---Pub Rock",
+    "Rock---Punk",
+    "Rock---Rock & Roll",
+    "Rock---Rockabilly",
+    "Rock---Shoegaze",
+    "Rock---Ska",
+    "Rock---Sludge Metal",
+    "Rock---Soft Rock",
+    "Rock---Southern Rock",
+    "Rock---Space Rock",
+    "Rock---Speed Metal",
+    "Rock---Stoner Rock",
+    "Rock---Surf",
+    "Rock---Symphonic Rock",
+    "Rock---Technical Death Metal",
+    "Rock---Thrash",
+    "Rock---Twist",
+    "Rock---Viking Metal",
+    "Rock---Yé-Yé",
+    "Stage & Screen---Musical",
+    "Stage & Screen---Score",
+    "Stage & Screen---Soundtrack",
+    "Stage & Screen---Theme",
+]
+MOOD_THEME_CLASSES = [
+    "action", "adventure", "advertising", "background", "ballad", "calm",
+    "children", "christmas", "commercial", "cool", "corporate", "dark", "deep",
+    "documentary", "drama", "dramatic", "dream", "emotional", "energetic",
+    "epic", "fast", "film", "fun", "funny", "game", "groovy", "happy", "heavy",
+    "holiday", "hopeful", "inspiring", "love", "meditative", "melancholic",
+    "melodic", "motivational", "movie", "nature", "party", "positive",
+    "powerful", "relaxing", "retro", "romantic", "sad", "sexy", "slow", "soft",
+    "soundscape", "space", "sport", "summer", "trailer", "travel", "upbeat",
+    "uplifting"
+]
+INSTRUMENT_CLASSES = [
+    "accordion", "acousticbassguitar", "acousticguitar", "bass", "beat", "bell",
+    "bongo", "brass", "cello", "clarinet", "classicalguitar", "computer",
+    "doublebass", "drummachine", "drums", "electricguitar", "electricpiano",
+    "flute", "guitar", "harmonica", "harp", "horn", "keyboard", "oboe",
+    "orchestra", "organ", "pad", "percussion", "piano", "pipeorgan", "rhodes",
+    "sampler", "saxophone", "strings", "synthesizer", "trombone", "trumpet",
+    "viola", "violin", "voice"
+]

data/stem.py

@@ -0,0 +1,62 @@
+from enum import IntEnum
+
+
+class Stem(IntEnum):
+    DRUMS = 0
+    BASS = 1
+    GUITAR = 2
+    KEYBOARD = 3
+    PIANO = 4
+    STRINGS = 5
+    OTHER = 6
+    VOCALS = 7
+    ANY = 8
+
+    @staticmethod
+    def fromstring(s: str):
+        match s:
+            case "drums":
+                return Stem.DRUMS
+            case "drums_mixed":
+                return Stem.DRUMS
+            case "bass":
+                return Stem.BASS
+            case "guitar":
+                return Stem.GUITAR
+            case "other_keys":
+                return Stem.KEYBOARD
+            case "keyboard":
+                return Stem.KEYBOARD
+            case "piano":
+                return Stem.PIANO
+            case "bowed_strings":
+                return Stem.STRINGS
+            case "strings":
+                return Stem.STRINGS
+            case "other":
+                return Stem.OTHER
+            case "vocals":
+                return Stem.VOCALS
+            case _:
+                raise ValueError("unknown stem name")
+
+    def getname(self):
+        match self:
+            case Stem.DRUMS:
+                return "drums_mixed"
+            case Stem.BASS:
+                return "bass"
+            case Stem.GUITAR:
+                return "guitar"
+            case Stem.KEYBOARD:
+                return "other_keys"
+            case Stem.PIANO:
+                return "piano"
+            case Stem.STRINGS:
+                return "bowed_strings"
+            case Stem.OTHER:
+                return "other"
+            case Stem.VOCALS:
+                return "vocals"
+            case Stem.ANY:
+                raise ValueError("Stem ANY has no name")

data/stemmed_datamodule.py

@@ -0,0 +1,190 @@
+import lightning as L
+from pathlib import Path
+from typing import Set, Optional
+from config import ConfigurationError
+from data.stem import Stem
+from data.stemmed_dataset import StemmedDataset
+from data.dataset_mixed import MixDataset
+import hyperparameters as hp
+import torch.utils.data
+from torch import Tensor
+import random
+
+
+class StemmedDatamodule(L.LightningDataModule):
+
+    def __init__(self, params: hp.StemmedDatasetParams):
+        super().__init__()
+        # self.stems: Set[Stem] = params.stems
+        # self.single_stem: bool = params.single_stem
+        # self.root_dir: Path = Path(params.root_dir)
+        # self.clip_length_in_seconds: int = params.clip_length_in_seconds
+        # self.sample_rate: int = params.sample_rate
+        # self.batch_size_train: int = params.batch_size_train
+        # self.batch_size_test: int = params.batch_size_test
+        # self.num_workers: int = params.num_workers
+        # self.speed_transform_p: float = params.speed_transform_p
+        # self.pitch_transform_p: float = params.pitch_transform_p
+        # self.n_samples_per_epoch: int = params.n_samples_per_epoch
+        # self.target_stem: Stem = params.target_stem
+        # self.add_click: bool = params.add_click
+        # self.sync_chunks: bool = params.sync_chunks
+        # self.bpm_in_caption: bool = params.bpm_in_caption
+        self.params = params
+
+        if isinstance(params, hp.MixDatasetParams):
+            self.dataset_class = MixDataset
+        else:
+            self.dataset_class = StemmedDataset
+
+        self.setup(None)
+
+        self.lengths = {
+            "train": len(self.train_dataloader()),
+            "valid": len(self.val_dataloader())
+        }
+
+    def _collate_fn(self, batch):
+        if self.params.min_context_seconds > self.params.clip_length_in_seconds:
+            raise ConfigurationError(
+                "Context has to be smaller than clip length")
+        if (self.params.min_context_seconds ==
+                self.params.clip_length_in_seconds):
+            inputs = {
+                k:
+                    torch.stack([s[k] for s in batch]) if isinstance(
+                        batch[0][k], Tensor) else [s[k] for s in batch]
+                for k in batch[0].keys()
+                # if k != "name"
+            }
+        else:
+            inputs = {
+                k:
+                    torch.stack([s[k] for s in batch]) if
+                    (isinstance(batch[0][k], Tensor) and
+                     k != "context") else [s[k] for s in batch]
+                for k in batch[0].keys()
+                # if k != "name"
+            }
+
+        # inputs = {
+        #     "target": torch.stack([s["target"] for s in batch]),
+        #     "context": torch.stack([s["context"] for s in batch]),
+        #     "description": [s["description"] for s in batch],
+        #     "style": torch.stack([s["style"] for s in batch])
+        # }
+        return inputs
+
+    def setup(self, stage: Optional[str]):
+        self.train_dataset = self.dataset_class(
+            Path(self.params.root_dir),
+            self.params.stems,
+            train=True,
+            target_stem=self.params.target_stem,
+            single_stem=self.params.single_stem,
+            min_context_seconds=self.params.min_context_seconds,
+            use_style_conditioning=self.params.use_style_conditioning,
+            use_beat_conditioning=self.params.use_beat_conditioning,
+            add_click=self.params.add_click,
+            sync_chunks=self.params.sync_chunks,
+            bpm_in_caption=self.params.bpm_in_caption,
+            sample_rate=self.params.sample_rate,
+            type_of_context=self.params.type_of_context,
+            chunk_size_samples=self.params.clip_length_in_seconds *
+            self.params.sample_rate,
+            speed_transform_p=self.params.speed_transform_p,
+            pitch_transform_p=self.params.pitch_transform_p,
+            n_samples_per_epoch=self.params.n_samples_per_epoch,
+            stereo=False,
+            max_genres_in_description=3,
+            max_moods_in_description=3,
+        )
+
+        self.val_dataset = self.dataset_class(
+            Path(self.params.root_dir),
+            self.params.stems,
+            train=False,
+            target_stem=self.params.target_stem,
+            single_stem=self.params.single_stem,
+            min_context_seconds=self.params.min_context_seconds,
+            use_style_conditioning=self.params.use_style_conditioning,
+            use_beat_conditioning=self.params.use_beat_conditioning,
+            add_click=self.params.add_click,
+            sync_chunks=self.params.sync_chunks,
+            bpm_in_caption=self.params.bpm_in_caption,
+            sample_rate=self.params.sample_rate,
+            type_of_context=self.params.type_of_context,
+            chunk_size_samples=self.params.clip_length_in_seconds *
+            self.params.sample_rate,
+            speed_transform_p=self.params.speed_transform_p,
+            pitch_transform_p=self.params.pitch_transform_p,
+            n_samples_per_epoch=None,
+            stereo=False,
+            max_genres_in_description=3,
+            max_moods_in_description=3,
+        )
+
+    def train_dataloader(self):
+        return torch.utils.data.DataLoader(
+            dataset=self.train_dataset,
+            batch_size=self.params.batch_size_train,
+            shuffle=True,
+            num_workers=self.params.num_workers,
+            collate_fn=self._collate_fn,
+            pin_memory=True,
+            worker_init_fn=lambda id: random.seed(id),
+            # prefetch_factor=1,
+        )
+
+    def val_dataloader(self):
+        return torch.utils.data.DataLoader(
+            self.val_dataset,
+            self.params.batch_size_test,
+            shuffle=False,
+            num_workers=self.params.num_workers,
+            collate_fn=self._collate_fn,
+            pin_memory=True,
+            worker_init_fn=lambda id: random.seed(id),
+            # prefetch_factor=1,
+        )
+
+
+# def main():
+if __name__ == "__main__":
+
+    from tqdm import tqdm
+    import config as cfg
+
+    # root_dir = cfg.mixdata_path()
+    root_dir = cfg.moises_path()
+    stems = {
+        Stem.DRUMS, Stem.GUITAR, Stem.BASS, Stem.PIANO, Stem.KEYBOARD,
+        Stem.STRINGS, Stem.OTHER
+    }
+    dataset_params = hp.MixDatasetParams(
+        root_dir=root_dir,
+        stems=stems,
+        single_stem=True,
+        min_context_seconds=5,
+        use_style_conditioning=True,
+        use_beat_conditioning=True,
+        target_stem=Stem.DRUMS,
+        add_click=False,
+        sync_chunks=False,
+        bpm_in_caption=False,
+        batch_size_train=4,
+        batch_size_test=4,
+        num_workers=8,
+        clip_length_in_seconds=10,
+        sample_rate=32_000,
+        speed_transform_p=1,
+        pitch_transform_p=0.5,
+        n_samples_per_epoch=2000,
+    )
+    d = dataset_params.instantiate()
+
+    vd = d.val_dataloader()
+    vbatch = next(iter(vd))
+
+    td = d.train_dataloader()
+    tbatch = next(iter(td))

data/stemmed_dataset.py

@@ -0,0 +1,801 @@
+import random
+from typing import Dict, Iterable, List, Optional, Sequence, Set, Tuple
+import numpy as np
+import torch.utils.data
+from pathlib import Path
+import torch
+import itertools
+
+from conditioning.beat_embedder import Beat
+from data.stem import Stem
+from utils import audio as audio_utils
+import json
+from torch import Tensor
+import torchaudio
+# from collections.abc import Sized
+import librosa
+
+# import pyrubberband as pyrb
+# import pylibrb
+
+N_VALID_SAMPLES = 24
+# 014f37 is removed
+# EXPECTED_N_SONGS = 240
+# EXPECTED_N_SONGS = 150
+
+
+class StemmedDataset(torch.utils.data.Dataset):
+
+    def __init__(self,
+                 root_dir: Path,
+                 stems: Set[Stem],
+                 target_stem: Stem,
+                 single_stem: bool,
+                 min_context_seconds: int,
+                 use_style_conditioning: bool,
+                 use_beat_conditioning: bool,
+                 type_of_context: str,
+                 bpm_in_caption: bool,
+                 add_click: bool,
+                 sync_chunks: bool,
+                 train: bool,
+                 sample_rate: int,
+                 chunk_size_samples: int,
+                 speed_transform_p: float,
+                 pitch_transform_p: float,
+                 stereo: bool = False,
+                 max_genres_in_description: int = 3,
+                 max_moods_in_description: int = 3,
+                 n_samples_per_epoch: Optional[int] = None,
+                 verbose: bool = False):
+        self.root_dir: Path = root_dir
+        self.stems: Set[Stem] = stems
+        self.single_stem: bool = single_stem
+        self.min_context_seconds: int = min_context_seconds
+        self.target_stem: Stem = target_stem
+        self.use_style_conditioning: bool = use_style_conditioning
+        self.use_beat_conditioning: bool = use_beat_conditioning
+        if self.use_style_conditioning and not self.single_stem:
+            raise ValueError("You can only use style conditioning if "
+                             "the target is a single stem")
+
+        if self.target_stem != Stem.ANY:
+            assert self.single_stem
+
+        self.add_click: bool = add_click
+        self.bpm_in_caption: bool = bpm_in_caption
+        self.sync_chunks: bool = sync_chunks
+        self.stem_names: Set[str] = {s.getname() for s in self.stems}
+        self.train: bool = train
+        self.sample_rate: int = sample_rate
+        self.chunk_size_samples: int = chunk_size_samples
+        self.speed_transform_p: float = speed_transform_p
+        self.pitch_transform_p: float = pitch_transform_p
+        self.stereo: bool = stereo
+        self.max_genres_in_description: int = max_genres_in_description
+        self.max_moods_in_description: int = max_moods_in_description
+        self.verbose: bool = verbose
+
+        self.type_of_context: str = type_of_context
+        assert self.type_of_context in ["stems", "beats", "stems or beats"]
+
+        if self.add_click or self.sync_chunks:
+            if not (self.root_dir / "sync.json").exists():
+                raise FileNotFoundError(
+                    "If you want click or sync, I need a 'sync.json' file in "
+                    "the top-level dir of the dataset")
+
+            with open(self.root_dir / "sync.json", "r") as f:
+                self.syncdata: Dict[str, List[int]] = json.load(f)
+
+        # load all song directories
+        self.song_names: List[str] = sorted(
+            [p.name for p in self.root_dir.iterdir() if p.is_dir()])
+        # assert len(self.song_names) == EXPECTED_N_SONGS
+
+        # if has a single target stem filter out songs that don't have that stem
+        if self.target_stem != Stem.ANY:
+            toremove: List[str] = []
+            for song_name in self.song_names:
+                if not (self.root_dir / song_name /
+                        self.target_stem.getname()).exists():
+                    toremove.append(song_name)
+            self.song_names = [s for s in self.song_names if s not in toremove]
+
+        # train/valid split
+        if self.train:
+            self.song_names = self.song_names[:-N_VALID_SAMPLES]
+        else:
+            self.song_names = self.song_names[-N_VALID_SAMPLES:]
+
+        # if self.verbose:
+        print(f"Loaded {len(self.song_names)} for "
+              f"{'train' if self.train else 'valid'} dataset.")
+
+        # WTF? TODO: remove
+        # if self.target_stem != Stem.ANY:
+        #     # if only interested in a stem, remove songs without it
+        #     for songname in self.song_names:
+        #         songdir = self.root_dir / songname
+        #         if not (songdir / self.target_stem.getname()).exists():
+        #             self.song_names.remove(songname)
+
+        # create iterator to run n_sample times
+        self.n_samples: int = n_samples_per_epoch or len(self.song_names)
+
+        self.song_iterator: List[str] = list(
+            itertools.islice(itertools.cycle(iter(self.song_names)),
+                             self.n_samples))
+
+    def __len__(self):
+        return self.n_samples
+
+    def save_sample(self, sample: Dict, path: Path):
+        audio_utils.save_audio(sample["wav"], path / "input.wav",
+                               self.sample_rate)
+        audio_utils.save_audio(sample["conditioning"].wav, path / "cond.wav",
+                               self.sample_rate)
+        audio_utils.save_audio(sample["conditioning"].wav + sample["wav"],
+                               path / "mix.wav", self.sample_rate)
+
+    def get_description(self,
+                        features: Dict[str, str | int],
+                        instruments: Sequence[Stem],
+                        speed_factor: Optional[float] = None) -> str:
+
+        genres: List[str] = str(features["genres"]).split(",")
+        moods: List[str] = str(features["moods"]).split(",")
+        description = ""
+
+        # Genre
+        if len(genres) > 0:
+            # if more than max number of genres, choose first few
+            # if len(genres) > self.max_genres:
+            #     genres = random.sample(genres, self.max_genres)
+            genres = [
+                s.strip() for s in genres[:self.max_genres_in_description]
+            ]
+            description += (f"Genre{'s' if len(genres) > 1 else ''}: "
+                            f"{', '.join(genres)}. ")
+
+        # Mood
+        if len(moods) > 0:
+            # if more than max number of moods, choose first few
+            # if len(moods) > self.max_moods:
+            #     moods = random.sample(moods, self.max_moods)
+            moods = [s.strip() for s in moods[:self.max_moods_in_description]]
+            description += (f"Mood{'s' if len(moods) > 1 else ''}: "
+                            f"{', '.join(moods)}. ")
+
+        # Instruments
+        if not (self.single_stem and self.target_stem != Stem.ANY):
+            instrument_names: List[str] = [s.name.lower() for s in instruments]
+            random.shuffle(instrument_names)
+            description += f"Instruments: {', '.join(instrument_names)}."
+
+        # BPM
+        if self.bpm_in_caption:
+            bpm = features["bpm"]
+            if speed_factor:
+                bpm = round(int(bpm) / speed_factor)
+            description += f" Bpm: {bpm}."
+
+        # Key
+        # key = features["key"]
+        # description += f"Key: {key}."
+
+        return description
+
+    def _transform_chunk(self, t: Tensor, speed_factor: float,
+                         pitch_factor: int, target_size: int):
+        if self.stereo:
+            raise NotImplementedError(
+                "No augmentations for stereo audio implemented")
+
+        stretched = audio_utils.stretch_with_timeout(
+            t,
+            self.sample_rate,
+            speed_factor,
+            pitch_factor,
+            2,
+        )
+
+        if stretched.shape[-1] < target_size:
+            stretched = torch.nn.functional.pad(
+                stretched, (0, target_size - stretched.shape[-1]), "constant",
+                0)
+        elif stretched.shape[-1] > target_size:
+            stretched = stretched[..., :target_size]
+
+        return stretched
+
+    def load_stems(
+        self,
+        song_path: Path,
+        song_stems: Iterable[Stem],
+        start_offset: int,
+        n_frames: int,
+    ) -> Dict[Stem, Tensor]:
+
+        # load audio chunks for each stem
+        stem_tensors: Dict[Stem, Tensor] = {}
+        for stem in song_stems:
+            stemdir = song_path / stem.getname()
+            stem_tensor: Tensor = torch.zeros(2 if self.stereo else 1,
+                                              n_frames,
+                                              dtype=torch.float32)
+
+            # for each track of a stem
+            for trackpath in stemdir.iterdir():
+                # load wav
+                chunk = audio_utils.load_audio_chunk(trackpath,
+                                                     start_offset,
+                                                     n_frames,
+                                                     stereo=self.stereo)
+                stem_tensor += chunk
+
+            # if not silent, include it in dictionary of stems
+            if not audio_utils.is_silent(stem_tensor, threshold=0.01):
+                stem_tensors[stem] = stem_tensor
+
+        return stem_tensors
+
+    # def choose_conditioning(
+    #         self,
+    #         stems: Sequence[Stem]) -> Tuple[Sequence[Stem], Sequence[Stem]]:
+    #     n_stems: int = len(stems)
+    #     if n_stems == 1:
+    #         # if only one stem, use it as input with no conditioning
+    #         return stems[:], []
+    #     n_conditioning_stems: int = random.randint(1, n_stems - 1)
+    #     conditioning_stems: Sequence[Stem] = random.sample(
+    #         stems, n_conditioning_stems)
+    #     input_stems = list(filter(lambda x: x not in conditioning_stems, stems))
+    #     return input_stems, conditioning_stems
+
+    def choose_input_and_conditioning(
+            self,
+            stems: Sequence[Stem]) -> Tuple[Sequence[Stem], Sequence[Stem]]:
+        n_stems: int = len(stems)
+
+        if n_stems == 1:
+            raise RuntimeError("This song has only one stem")
+            return stems[:], []
+
+        # choose a random number of context stems,
+        # leaving at least 1 for the input
+        if self.target_stem != Stem.ANY:
+            assert self.target_stem in stems
+            possible_conditioning_stems = [
+                s for s in stems if s != self.target_stem
+            ]
+            n_conditioning_stems: int = random.randint(
+                1, len(possible_conditioning_stems))
+            conditioning_stems: Sequence[Stem] = random.sample(
+                possible_conditioning_stems, n_conditioning_stems)
+        else:
+            n_conditioning_stems: int = random.randint(1, n_stems - 1)
+            conditioning_stems: Sequence[Stem] = random.sample(
+                stems, n_conditioning_stems)
+
+        # choose a random number of the remaining stems as input
+        n_input_stems: int = 1 if self.single_stem else random.randint(
+            1, n_stems - n_conditioning_stems)
+
+        # choose input stem
+        if self.target_stem != Stem.ANY:
+            # if target stem != ANY, that HAS to be the input
+            assert self.target_stem in stems
+            input_stems: Sequence[Stem] = [self.target_stem]
+        else:
+            possible_input_stems: Sequence[Stem] = list(
+                filter(lambda x: x not in conditioning_stems, stems))
+            input_stems: Sequence[Stem] = random.sample(possible_input_stems,
+                                                        n_input_stems)
+
+        return input_stems, conditioning_stems
+
+    def add_click_to_track(self, wav: Tensor, wav_sr: int,
+                           click_frames: Sequence[int],
+                           start_offset: int) -> Tensor:
+
+        for idx, frame in enumerate(click_frames):
+            if frame >= start_offset:
+                first_relevant_index = idx
+                break
+        else:
+            return wav
+
+        shifted_click_frames: List[int] = [
+            f - start_offset for f in click_frames[first_relevant_index:]
+        ]
+
+        click_track: Tensor = audio_utils.create_click(list(wav.shape), wav_sr,
+                                                       shifted_click_frames)
+        assert click_track.shape == wav.shape
+
+        mix: Tensor = wav + click_track * 0.5
+        return mix
+
+    def mix_input_and_conditioning(self, stem_tensors: Dict[Stem, Tensor],
+                                   input_stems: Sequence[Stem],
+                                   condition_stems: Sequence[Stem]):
+
+        assert len(input_stems) > 0
+        input_tensor = torch.stack([stem_tensors[s] for s in input_stems
+                                   ]).sum(dim=-0)
+
+        if len(condition_stems) > 0:
+            condition_tensor = torch.stack(
+                [stem_tensors[s] for s in condition_stems]).sum(dim=-0)
+        else:
+            condition_tensor = None
+
+        return input_tensor, condition_tensor
+
+    def find_good_chunk(
+            self, song_name: str, n_frames_to_take: int, song_n_frames: int,
+            song_path: Path,
+            song_stems: Iterable[Stem]) -> Tuple[Dict[Stem, Tensor], int]:
+        found_good_chunk: bool = False
+        attempts: int = 0
+        while not found_good_chunk:
+            attempts += 1
+            if attempts > 10 and (attempts - 1) % 10 == 0:
+                print(
+                    f"Tried to find some non-silent chunk of song {song_name} "
+                    f"for {attempts} times but it's so hard please master "
+                    "I am tired let me rest")
+
+            # choose random chunk
+            if self.sync_chunks:
+                choices: List[int] = self.syncdata[song_name]
+                choices = list(
+                    filter(lambda x: (x + n_frames_to_take) < song_n_frames,
+                           choices))
+                start_offset: int = random.choice(choices)
+
+            else:
+                start_offset: int = random.randint(
+                    0, song_n_frames - n_frames_to_take)
+
+            # load song stems, filter out silent ones
+            stem_tensors: Dict[Stem, Tensor] = self.load_stems(
+                song_path,
+                song_stems,
+                start_offset,
+                n_frames_to_take,
+            )
+            nonsilent_stems: List[Stem] = list(stem_tensors.keys())
+
+            if self.target_stem != Stem.ANY:
+                if self.target_stem in nonsilent_stems and (
+                        len(nonsilent_stems)
+                        > (1 if len(list(song_stems)) > 1 else 0)):
+                    found_good_chunk = True
+            else:
+                if len(nonsilent_stems) > 0:
+                    found_good_chunk = True
+
+        return stem_tensors, start_offset  # type: ignore
+
+    def __getitem__(self, idx: int) -> Dict[str, Tensor | str]:
+        # output = {
+        #     "name": "",
+        #     "target": torch.rand(1, 320_000),
+        #     "description": "",
+        #     "context": torch.rand(1, 320_000),
+        #     "style": torch.rand(1, 320_000),
+        #     "beat": Beat(beats, downbeats, seq_len)
+        # }
+        # return output
+
+        song_name = self.song_iterator[idx]
+        song_path: Path = self.root_dir / song_name
+
+        # toss a coin to decide whether to augment data
+        apply_speed_transform: bool = random.random() < self.speed_transform_p
+        apply_pitch_transform: bool = random.random() < self.pitch_transform_p
+
+        # choose random augmentation factors if needed
+        speed_factor = round(random.random() * 0.4 +
+                             0.80, 2) if apply_speed_transform else 1
+        pitch_factor = (random.randint(-4, 4) if apply_pitch_transform else 0)
+
+        # get song features
+        with open(song_path / "features.json", "r") as f:
+            features: Dict[str, str | int] = json.load(f)
+        song_n_frames: int = int(features["num_frames"])
+        song_sr: int = int(features["sample_rate"])
+        n_frames_to_take_orig: int = int(self.chunk_size_samples /
+                                         self.sample_rate * song_sr)
+        n_frames_to_take = int(n_frames_to_take_orig / speed_factor)
+
+        # get all song stems, including possibly silent ones
+        song_stems: List[Stem] = []
+        for stem in self.stems:
+            if (song_path / stem.getname()).exists():
+                song_stems.append(stem)
+
+        if self.target_stem != Stem.ANY and self.target_stem not in song_stems:
+            raise RuntimeError(f"Target stem is {self.target_stem} but song "
+                               f"{song_name} has no interesting stems. "
+                               "Maybe remove it from the dataset?")
+
+        if len(song_stems) == 0:
+            raise RuntimeError(f"Song {song_name} has no interesting stems. "
+                               "Maybe remove it from the dataset?")
+        if len(song_stems) == 1:
+            raise RuntimeError(f"Song {song_name} has only one stem, which is"
+                               f"{song_stems[0]}. What to do?")
+
+        # find a good chunk of the song
+        stem_tensors: Dict[Stem, Tensor]
+        start_offset: int
+        stem_tensors, start_offset = self.find_good_chunk(
+            song_name, n_frames_to_take, song_n_frames, song_path, song_stems)
+        nonsilent_stems: List[Stem] = list(stem_tensors.keys())
+
+        # split stems between input and conditioning
+        input_stems, condition_stems = self.choose_input_and_conditioning(
+            nonsilent_stems)
+
+        # mix input and condition tensors
+        input_tensor, condition_tensor = self.mix_input_and_conditioning(
+            stem_tensors, input_stems, condition_stems)
+
+        # if applying style conditioning, find a good style conditioning chunk
+        style_tensor: Optional[Tensor] = None
+        if self.use_style_conditioning:
+            assert self.single_stem
+            assert len(input_stems) == 1
+            inputstem = input_stems[0]
+            style_tensor = self.find_good_chunk(song_name,
+                                                n_frames_to_take_orig,
+                                                song_n_frames, song_path,
+                                                [inputstem])[0][inputstem]
+
+        # if using beat conditioning, compute beats data for current chunk
+        beats_conditioning: Optional[Beat] = None
+        if self.use_beat_conditioning:
+            beatfile = song_path / "beatthis.npz"
+            if not beatfile.exists():
+                raise FileNotFoundError(
+                    f"Couldn't find beat annotations for song {song_path}")
+            loaded = np.load(beatfile)
+            beats_sec = torch.from_numpy(loaded["beats"])
+            downbeats_sec = torch.from_numpy(loaded["downbeats"])
+
+            # if using a speed augmentation, reposition beats
+            beats_frames: Tensor = (beats_sec * speed_factor *
+                                    self.sample_rate).round().long()
+            downbeats_frames: Tensor = (downbeats_sec * speed_factor *
+                                        self.sample_rate).round().long()
+            start_offset = round(start_offset / song_sr * self.sample_rate *
+                                 speed_factor)
+            min_max_beat: Tensor = torch.tensor(
+                [start_offset, start_offset + self.chunk_size_samples])
+
+            beats_start_idx, beats_end_idx = torch.searchsorted(beats_frames,
+                                                                min_max_beat,
+                                                                right=False)
+            beats_cut = beats_frames[beats_start_idx:beats_end_idx]
+
+            downbeats_start_idx, downbeats_end_idx = torch.searchsorted(
+                downbeats_frames, min_max_beat, right=False)
+            downbeats_cut = downbeats_frames[
+                downbeats_start_idx:downbeats_end_idx]
+
+            assert (downbeats_end_idx + 1 >= len(downbeats_frames) or
+                    downbeats_frames[downbeats_end_idx + 1]
+                    >= start_offset + self.chunk_size_samples)
+            assert (beats_end_idx + 1 >= len(beats_frames) or
+                    beats_frames[beats_end_idx + 1]
+                    >= start_offset + self.chunk_size_samples)
+
+            beats_cut -= start_offset
+            downbeats_cut -= start_offset
+
+            beats_conditioning = Beat(beats_cut, downbeats_cut,
+                                      self.chunk_size_samples)
+
+            match self.type_of_context:
+                case "beats":
+                    beats_as_context = True
+                case "stems":
+                    beats_as_context = False
+                case "stems or beats":
+                    beats_as_context = random.random() < 0.5
+
+            beats_time = beats_cut.numpy() / self.sample_rate
+            downbeats_time = downbeats_cut.numpy() / self.sample_rate
+            clicks = librosa.clicks(times=beats_time,
+                                    sr=self.sample_rate,
+                                    click_freq=1000,
+                                    length=self.chunk_size_samples)
+            downbeat_clicks = librosa.clicks(times=downbeats_time,
+                                             sr=self.sample_rate,
+                                             click_freq=1000,
+                                             length=self.chunk_size_samples)
+            # Combine clicks (downbeats are stronger)
+            audio_cliks = clicks + downbeat_clicks
+            audio_cliks = np.clip(audio_cliks, -1.0, 1.0)
+            # mono audio
+            if audio_cliks.ndim > 1:
+                audio_cliks = audio_cliks.mean(axis=0)
+            context_beats = torch.tensor(audio_cliks)
+            context_beats = torch.unsqueeze(context_beats, 0)
+            if beats_as_context:
+                condition_tensor = context_beats
+
+        # add click
+        if self.add_click:
+            raise NotImplementedError(
+                "Add click is not implemented for new dataset with style")
+            click_frames: List[int] = self.syncdata[song_name]
+            input_tensor = self.add_click_to_track(input_tensor, song_sr,
+                                                   click_frames, start_offset)
+            if condition_tensor is not None:
+                condition_tensor = self.add_click_to_track(
+                    condition_tensor, song_sr, click_frames, start_offset)
+
+        # get description of song input
+        description: str = self.get_description(features, input_stems,
+                                                speed_factor)
+
+        # resample input and conditioning to desired sample rate
+        input_tensor = torchaudio.functional.resample(input_tensor, song_sr,
+                                                      self.sample_rate)
+        if condition_tensor is not None and not beats_as_context:
+            condition_tensor = torchaudio.functional.resample(
+                condition_tensor, song_sr, self.sample_rate)
+
+        if style_tensor is not None:
+            style_tensor = torchaudio.functional.resample(
+                style_tensor, song_sr, self.sample_rate)
+
+        # data augmentation to input and conditioning
+        if apply_speed_transform or apply_pitch_transform:
+            input_tensor = self._transform_chunk(input_tensor, speed_factor,
+                                                 pitch_factor,
+                                                 self.chunk_size_samples)
+            if condition_tensor is not None and not beats_as_context:
+                condition_tensor = self._transform_chunk(
+                    condition_tensor, speed_factor, pitch_factor,
+                    self.chunk_size_samples)
+
+        if condition_tensor is not None:
+            # cut conditioning to a random length
+            min_context_samples = self.min_context_seconds * self.sample_rate
+            if min_context_samples < self.chunk_size_samples:
+                if torch.rand(1).item() > 0.95:
+                    index = self.chunk_size_samples
+                else:
+                    index = torch.randint(min_context_samples,
+                                          self.chunk_size_samples + 1,
+                                          (1,)).item()
+                condition_tensor = condition_tensor[..., :index]
+
+        output = {
+            "name": song_name,
+            "target": input_tensor,
+            "description": description,
+            "context": condition_tensor,
+        }
+
+        if self.use_beat_conditioning:
+            output["beat_seconds"] = beats_time
+
+        if style_tensor is not None:
+            output["style"] = style_tensor
+
+        if beats_conditioning is not None:
+            output["beat"] = beats_conditioning
+
+        return output
+
+
+# def generate_sync_data(
+#         output_path: Path) -> Dict[str, Dict[str, float | List[int]]]:
+#     db = MoisesDB(data_path=str(cfg.DATA_DIR / "moisesdb"),
+#                   sample_rate=32_000)
+#     if not output_path.exists():
+#         raise FileNotFoundError("output path doesn't seem to exist.")
+#     sync_path = output_path / "sync.json"
+#     if sync_path.exists():
+#         raise FileExistsError()
+#     data = {}
+#     errors = 0
+#     for song in tqdm(db, total=len(db)): # type: ignore
+#         try:
+#             songid = song.id
+#             sr = song.sr
+#             audio = librosa.to_mono(song.audio)
+#             # utils.save_audio(audio)
+#             tempo, beats = librosa.beat.beat_track(y=audio,
+#                                                    sr=sr,
+#                                                    units="samples")
+#             beats = beats.tolist()
+#             data[songid] = {
+#                 "tempo": tempo,
+#                 "beats": beats,
+#             }
+#         except:
+#             errors += 1
+#     with open(sync_path, "w") as fp:
+#         json.dump(data, fp)
+#     print(f"Saved sync data of {len(db) - errors} songs, with {errors} errors.")
+#     return data
+
+
+def prepare_data(root_dir: Path, save_mixed_drums: bool, save_mix: bool,
+                 extract_features: bool, track_bpm: bool):
+    from tqdm import tqdm
+    import torchaudio
+    from torch import Tensor
+    from lag.data.auto_labelling import get_audio_features
+
+    subdirs: List[Path] = sorted([p for p in root_dir.iterdir() if p.is_dir()],
+                                 key=lambda x: x.name)
+    # assert len(subdirs) == EXPECTED_N_SONGS
+
+    if track_bpm:
+        syncdata: Dict[str, List[int]] = {}
+
+    for song in tqdm(subdirs):
+
+        for stemdir in (p for p in song.iterdir() if p.is_dir()):
+            if len(list(stemdir.iterdir())) == 0:
+                raise FileNotFoundError(f"Song {song} contains no stems. WTF")
+
+        # mix drums
+        if (song / "drums").exists():
+            drums_sample_rates: Set[int] = set()
+            drums_audios: List[Tensor] = []
+            for drum_stem in (song / "drums").iterdir():
+                audio, sr = torchaudio.load(str(drum_stem))
+                drums_sample_rates.add(sr)
+                drums_audios.append(audio.permute(1, 0))
+
+            if len(drums_sample_rates) != 1:
+                raise ValueError(f"song {song} contains drums stems of "
+                                 "different sample rates")
+            drums_sr = drums_sample_rates.pop()
+
+            # mix drums
+            drums_tensor = torch.nn.utils.rnn.pad_sequence(
+                drums_audios, batch_first=True,
+                padding_value=0).permute(0, 2, 1).sum(dim=0)
+
+            assert drums_tensor.ndim == 2
+
+            if save_mixed_drums:
+                target_dir = song / "drums_mixed"
+                target_dir.mkdir(exist_ok=True)
+                torchaudio.save(target_dir / "drums.wav", drums_tensor,
+                                drums_sr)
+
+        # load mixed song
+        stem_subdirs = [p for p in song.iterdir() if p.is_dir()]
+
+        # for each stem
+        stem_tracks: List[Tensor] = []
+        sample_rates: Set[int] = set()
+        for stem_subdir in stem_subdirs:
+            if stem_subdir.name == "drums":
+                continue
+
+            # for each track of that stem
+            for audio_path in stem_subdir.iterdir():
+                audio, sr = torchaudio.load(str(audio_path))
+                assert audio.ndim == 2
+                stem_tracks.append(audio.permute(1, 0))
+
+                sample_rates.add(sr)
+                if len(sample_rates) > 1:
+                    raise ValueError(f"song {song} contains stems of "
+                                     "different sample rates")
+        sr = sample_rates.pop()
+
+        # pad shorter tracks
+        stems_tensor: Tensor = torch.nn.utils.rnn.pad_sequence(
+            stem_tracks,
+            batch_first=True,
+            padding_value=0.,
+        ).permute(0, 2, 1)
+        assert stems_tensor.ndim == 3
+
+        # mix song
+        mixed_tensor = stems_tensor.sum(dim=0)
+        num_frames = mixed_tensor.shape[-1]
+        assert mixed_tensor.ndim == 2
+
+        if save_mix:
+            mix_out_path = song / "mixed.wav"
+            audio_utils.save_audio(mixed_tensor, mix_out_path, sr)
+
+        # neural classification to get metadata
+        if extract_features:
+            features = get_audio_features(mixed_tensor, sr, cfg.weights_dir())
+
+            features["sample_rate"] = sr
+            features["num_frames"] = num_frames
+            out_file = song / "features.json"
+            with open(out_file, "w") as f:
+                json.dump(features, f)
+
+        # bpm tracking
+        if track_bpm:
+            song_numpy = audio_utils.to_mono(mixed_tensor).squeeze().numpy()
+
+            # audio = librosa.to_mono(song.audio)
+            # utils.save_audio(audio)
+            try:
+                tempo, beats = librosa.beat.beat_track(y=song_numpy,
+                                                       sr=sr,
+                                                       units="samples")
+            except Exception as e:
+                print(f"Error tracking beats of song {song.name}")
+                raise e
+
+            beats = beats.tolist()
+
+            syncdata[song.name] = beats  # type: ignore
+
+    if track_bpm:
+        sync_path: Path = root_dir / "sync.json"
+        with open(sync_path, "w") as f:
+            json.dump(syncdata, f)  # type: ignore
+
+
+# if __name__ == "__main__":
+#     # from lag import config as cfg
+#     from tqdm import tqdm
+
+#     # root_dir = Path("/home/tkol/dev/datasets") / "moisesdb" / "moisesdb_v0.1"
+#     # root_dir = Path("/home/tkol/dev/datasets") / "moisesdb" / "musdb"
+#     root_dir = cfg.moises_path()
+
+#     '''prepare_data(root_dir,
+#                  save_mixed_drums=False,
+#                  save_mix=False,
+#                  extract_features=True,
+#                  track_bpm=False)'''
+
+#     stems = {
+#         Stem.DRUMS, Stem.GUITAR, Stem.BASS, Stem.PIANO, Stem.KEYBOARD,
+#         Stem.STRINGS
+#     }
+
+#     dataset = StemmedDataset(
+#         root_dir,
+#         stems,
+#         target_stem=Stem.DRUMS,
+#         single_stem=True,
+#         min_context_seconds=5,
+#         use_style_conditioning=True,
+#         use_beat_conditioning=True,
+#         add_click=False,
+#         bpm_in_caption=False,
+#         sync_chunks=False,
+#         train=False,
+#         sample_rate=32_000,
+#         chunk_size_samples=32_000 * 10,
+#         speed_transform_p=1,
+#         pitch_transform_p=1,
+#         stereo=False,
+#         n_samples_per_epoch=None,
+#     )
+
+#     dataset_iterator = iter(dataset)
+#     for i in tqdm(range(10)):
+#         sample = next(dataset_iterator)
+
+#         target: Tensor = sample["target"]  # type: ignore
+#         context: Tensor = sample["context"] if sample["context"] is not None else sample["context"]
+
+#         audio_utils.save_audio(target, cfg.AUDIO_DIR / "temp" / f"target{i}.wav")
+#         audio_utils.save_audio(context, cfg.AUDIO_DIR / "temp" / f"context{i}.wav")
+#         mix = target + torch.nn.functional.pad(
+#             context, (0, target.shape[-1] - context.shape[-1]))
+#         audio_utils.save_audio(mix, cfg.AUDIO_DIR / "temp" / f"mix{i}.wav")

hyperparameters.py

@@ -0,0 +1,244 @@
+from dataclasses import dataclass, field
+from typing import Dict, Tuple, Optional, List, Type, Set
+from pydoc import locate
+from frozendict import frozendict
+from copy import deepcopy
+from pathlib import Path
+
+import config as cfg
+from conditioning.condition_type import ConditionType
+from conditioning.conditioning_method import ConditioningMethod
+from conditioning import ConcreteEmbedder
+# from conditioning.prompt_processor import InterleavedContextPromptProcessor
+# from conditioning.t5embedder import T5EmbedderGPU
+from data.stem import Stem
+from utils.logging import to_loggable
+
+import typing
+if typing.TYPE_CHECKING:
+    from conditioning.prompt_processor import PromptProcessor
+
+
+class Loggable:
+
+    def to_dict(self) -> Dict:
+        return to_loggable(self.__dict__)  # type: ignore
+
+
+# generic parameters for a model that can be instantiated
+# from a hyperparameter class
+@dataclass(unsafe_hash=True, kw_only=True)
+class ModelParams(Loggable):
+    model_class: str
+
+    # def to_dict(self) -> Dict:
+    #     return to_loggable(self.__dict__)  # type: ignore
+
+    def instantiate(self):
+        klass = locate(self.model_class)
+        return klass(self)  # type: ignore
+
+
+# --- ENCODEC ---
+@dataclass(unsafe_hash=True)
+class QuantizerParams:
+    dimension: int
+    n_q: int
+    bins: int
+
+    q_dropout: bool = False
+    decay: float = 0.99
+    kmeans_init: bool = True
+    kmeans_iters: int = 50
+    threshold_ema_dead_code: int = 2
+    orthogonal_reg_weight: float = 0.0
+    orthogonal_reg_active_codes_only: bool = False
+    orthogonal_reg_max_codes: int | None = None
+
+
+@dataclass(unsafe_hash=True)
+class SeaNetParams:
+    dimension: int
+    n_filters: int
+    ratios: Tuple[int, int, int, int]
+    causal: bool
+    true_skip: bool
+
+    channels: int = 1
+    n_residual_layers: int = 1
+    activation: str = "ELU"
+    activation_params: frozendict = field(default_factory=frozendict)
+    norm: str = "weight_norm"
+    norm_params: frozendict = field(default_factory=frozendict)
+    kernel_size: int = 7
+    last_kernel_size: int = 7
+    residual_kernel_size: int = 3
+    dilation_base: int = 2
+    pad_mode: str = "reflect"
+    compress: int = 2
+    lstm: int = 2
+
+
+@dataclass(unsafe_hash=True, kw_only=True)
+class EncodecParams(ModelParams):
+    sample_rate: int
+    seanet_params: SeaNetParams
+    quantizer_params: QuantizerParams
+    sum_loss_mulitiplier: int
+    weights: Optional[str] = None
+
+    model_class: str = "stage.models.lightning_encodec.LightningEncodec"
+
+    def to_dict(self) -> Dict:
+        d = deepcopy(self.__dict__)
+        for key, value in self.seanet_params.__dict__.items():
+            d["seanet_" + key] = value
+        for key, value in self.quantizer_params.__dict__.items():
+            d["qt_" + key] = value
+        del d["quantizer_params"]
+        del d["seanet_params"]
+        return d
+
+
+# --- CONDITIONING ---
+@dataclass(unsafe_hash=True)
+class PromptProcessorParams(Loggable):
+    keep_only_valid_steps: bool
+    model_class: Type["PromptProcessor"]
+    context_dropout: Optional[float] = None
+
+
+@dataclass(unsafe_hash=True)
+class ConditioningParams(Loggable):
+    embedder_types: Dict[ConditionType, Type[ConcreteEmbedder]]
+    conditioning_methods: Dict[ConditionType, ConditioningMethod]
+    conditioning_dropout: float
+
+
+# --- LM ---
+@dataclass(unsafe_hash=True, kw_only=True)
+class LmParams(ModelParams):
+    dim: int
+    n_layers: int
+    n_heads: int
+    card: int = 2048
+    padding_token: Optional[int] = 2048
+    sep_token: Optional[int] = None
+    cross_attend: bool = True
+    weights: Optional[str] = None
+    model_class: str = "stage.models.musicgen_lm.MusicgenLm"
+
+
+@dataclass(unsafe_hash=True, kw_only=True)
+class PretrainedSmallLmParams(LmParams):
+    dim: int = 1024
+    n_layers: int = 24
+    n_heads: int = 16
+    weights: Optional[str] = str(cfg.weights_dir() / "lm-small-weights.pt")
+
+
+@dataclass(unsafe_hash=True, kw_only=True)
+class PretrainedLargeLmParams(LmParams):
+    dim: int = 2048
+    n_layers: int = 48
+    n_heads: int = 32
+    weights: Optional[str] = str(cfg.weights_dir() / "lm-large-weights.pt")
+
+
+@dataclass(unsafe_hash=True, kw_only=True)
+class PretrainedMelodyLmParams(LmParams):
+    dim: int = 1536
+    n_layers: int = 48
+    n_heads: int = 24
+    cross_attend: bool = False
+    weights: Optional[str] = str(cfg.weights_dir() / "lm-melody-weights.pt")
+
+
+@dataclass(unsafe_hash=True, kw_only=True)
+class FioraSmallLmParams(PretrainedSmallLmParams):
+    sep_token: Optional[int] = 2049
+
+
+# --- LORA ---
+@dataclass(unsafe_hash=True)
+class LoraParams:
+    r: int
+    alpha: int
+    dropout: float
+    layers: List[str]
+
+
+# --- MUSICGEN ---
+@dataclass(unsafe_hash=True, kw_only=True)
+class MusicgenParams(ModelParams):
+    encodec_params: EncodecParams
+    prompt_processor_params: PromptProcessorParams
+    conditioning_params: ConditioningParams
+    lm_params: LmParams
+    lora_params: Optional[LoraParams] = None
+    model_class: str = "stage.models.lightning_musicgen.LightningMusicgen"
+
+
+# ------ DATA -------
+
+
+@dataclass(unsafe_hash=True, kw_only=True)
+class DatasetParams:
+    datamodule_class: str
+    clip_length_in_seconds: int
+    sample_rate: int
+
+    def to_dict(self) -> Dict:
+        return self.__dict__
+
+    def instantiate(self):
+        klass = locate(self.datamodule_class)
+        return klass(self)  # type: ignore
+
+
+@dataclass(kw_only=True)
+class StemmedDatasetParams(DatasetParams):
+    root_dir: Path
+    stems: Set[Stem] = field(
+        default_factory=lambda: {
+            Stem.DRUMS,
+            Stem.BASS,
+            Stem.GUITAR,
+            Stem.KEYBOARD,
+            Stem.PIANO,
+            Stem.STRINGS,
+            Stem.OTHER,
+        })
+    single_stem: bool
+    target_stem: Stem
+    min_context_seconds: int
+    use_style_conditioning: bool
+    use_beat_conditioning: bool
+    type_of_context: str
+    add_click: bool
+    sync_chunks: bool
+    bpm_in_caption: bool
+    batch_size_train: int
+    batch_size_test: int
+    num_workers: int
+    clip_length_in_seconds: int
+    sample_rate: int
+    speed_transform_p: float
+    pitch_transform_p: float
+    n_samples_per_epoch: int
+    datamodule_class: str = "stage.data.stemmed_datamodule.StemmedDatamodule"
+
+
+@dataclass(kw_only=True)
+class MixDatasetParams(StemmedDatasetParams):
+    ...
+
+
+# --------- CONFIGURATIONS ---------
+pretrained_encodec_meta_32khz_params: EncodecParams = EncodecParams(
+    sample_rate=32_000,
+    seanet_params=SeaNetParams(128, 64, (8, 5, 4, 4), False, True),
+    quantizer_params=QuantizerParams(128, 4, 2048),
+    sum_loss_mulitiplier=0,
+    weights=str(cfg.weights_dir() / "encodec_32khz.pt"),
+)

inference.py

@@ -0,0 +1,54 @@
+import torch
+import lightning as L
+
+import config as cfg
+from loader import load_model
+from utils.audio import load_audio, save_audio
+"""
+REQUIREMENTS:
+
+weights/
+ - encodec_32khz.pt
+ - lm-small-weights.pt in
+
+checkpoints/
+ - stage-drums-ckp1.pt
+ - stage-bass-ckp1.pt
+"""
+
+#%% Load model
+INSTRUMENT = "drums"
+checkpoint_path = cfg.CKP_DIR / f"stage-{INSTRUMENT}.safetensors"
+model = load_model(checkpoint_path)
+
+#%% Load conditioning, generate and save
+
+# load context audio, description
+SAMPLE = "sample2"
+SEED = 42
+
+# load description if present
+desc_path = cfg.AUDIO_DIR / INSTRUMENT / f"{SAMPLE}-desc.txt"
+desc = desc_path.read_text().strip() if desc_path.exists() else None
+
+# load audio context
+wav = load_audio(cfg.AUDIO_DIR / INSTRUMENT / f"{SAMPLE}.wav").to(model.device)
+
+# generate
+L.seed_everything(SEED)
+out = model.generate(n_samples=1,
+                     gen_seconds=10,
+                     prompt=None,
+                     context=wav,
+                     style=None,
+                     beat=None,
+                     description=[desc],
+                     prog_bar=True)
+
+# save output and mix
+save_audio(out, cfg.AUDIO_DIR / "gen" / f"{SAMPLE}_{INSTRUMENT}_{SEED}.wav")
+padded_wav = torch.nn.functional.pad(wav,
+                                     tuple((0, out.shape[-1] - wav.shape[-1])),
+                                     value=0)
+mix = out + padded_wav
+save_audio(mix, cfg.AUDIO_DIR / "gen" / f"{SAMPLE}_{INSTRUMENT}_{SEED}_mix.wav")

loader.py

@@ -0,0 +1,42 @@
+from pathlib import Path
+from tabnanny import check
+from typing import Optional
+import torch
+from safetensors import torch as sft
+
+from conditioning.condition_type import ConditionType
+from conditioning.conditioning_method import ConditioningMethod
+from conditioning.prompt_processor import InterleavedContextPromptProcessor
+from conditioning.t5embedder import T5EmbedderGPU
+from models.lightning_musicgen import LightningMusicgen
+import hyperparameters as hp
+
+
+def load_model(checkpoint_path: Path,
+               device: Optional[str] = None) -> LightningMusicgen:
+
+    stage_params = hp.MusicgenParams(
+        encodec_params=hp.pretrained_encodec_meta_32khz_params,
+        prompt_processor_params=hp.PromptProcessorParams(
+            keep_only_valid_steps=True,
+            model_class=InterleavedContextPromptProcessor,
+            context_dropout=0.1),
+        conditioning_params=hp.ConditioningParams(
+            embedder_types={
+                ConditionType.DESCRIPTION: T5EmbedderGPU,
+            },
+            conditioning_methods={
+                ConditionType.DESCRIPTION: ConditioningMethod.CROSS_ATTENTION,
+            },
+            conditioning_dropout=0.5),
+        lm_params=hp.PretrainedSmallLmParams(sep_token=2049))
+
+    model: LightningMusicgen = stage_params.instantiate()
+    # state_dict = torch.load(checkpoint_path)
+    # model.load_state_dict(state_dict)
+    sft.load_model(model, checkpoint_path)
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(device).eval()
+
+    return model

models/__init__.py

@@ -0,0 +1 @@
+

models/encodec.py

@@ -0,0 +1,436 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""Compression models or wrapper around existing models.
+Also defines the main interface that a model must follow to be usable as an audio tokenizer.
+"""
+
+import logging
+import contextlib
+import os
+import math
+from pathlib import Path
+import typing as tp
+import numpy as np
+import torch
+from torch import nn
+# from torchmetrics.functional.audio import scale_invariant_signal_distortion_ratio
+
+import config as cfg
+from models.quantization.base import QuantizedResult
+from models import quantization as qt
+import hyperparameters as hp
+# from modules import SEANetDecoder, SEANetEncoder
+
+# audiocraft imports
+with contextlib.redirect_stderr(open(os.devnull, "w")):
+    from models.modules import SEANetEncoder, SEANetDecoder
+
+logger = logging.getLogger()
+
+
+class EncodecModel(nn.Module):
+    """Encodec model operating on the raw waveform.
+
+    Args:
+        encoder (nn.Module): Encoder network.
+        decoder (nn.Module): Decoder network.
+        quantizer (qt.BaseQuantizer): Quantizer network.
+        frame_rate (int): Frame rate for the latent representation.
+        sample_rate (int): Audio sample rate.
+        channels (int): Number of audio channels.
+        causal (bool): Whether to use a causal version of the model.
+        renormalize (bool): Whether to renormalize the audio before running the model.
+    """
+    # we need assignment to override the property in the abstract class,
+    # I couldn't find a better way...
+    frame_rate: float = 0  # type: ignore
+    sample_rate: int = 0
+    channels: int = 0
+
+    def __init__(self,
+                 encoder: SEANetEncoder,
+                 decoder: SEANetDecoder,
+                 quantizer: qt.ResidualVectorQuantizer,
+                 sample_rate: int,
+                 channels: int,
+                 causal: bool = True,
+                 renormalize: bool = False):
+        super().__init__()
+        self.encoder: SEANetEncoder = encoder
+        self.decoder: SEANetDecoder = decoder
+        self.quantizer: qt.ResidualVectorQuantizer = quantizer
+        # self.frame_rate = frame_rate
+        self.sample_rate = sample_rate
+        self.channels = channels
+        self.renormalize = renormalize
+        self.causal = causal
+        self.frame_rate: int = int(
+            math.ceil(self.sample_rate /
+                      np.prod(self.encoder.ratios)))  # type: ignore
+
+        if self.causal:
+            # we force disabling here to avoid handling linear overlap of segments
+            # as supported in original EnCodec codebase.
+            assert not self.renormalize, 'Causal model does not support renormalize'
+
+    @property
+    def total_codebooks(self):
+        """Total number of quantizer codebooks available."""
+        return self.quantizer.total_codebooks
+
+    @property
+    def num_codebooks(self):
+        """Active number of codebooks used by the quantizer."""
+        return self.quantizer.num_codebooks
+
+    def set_num_codebooks(self, n: int):
+        """Set the active number of codebooks used by the quantizer."""
+        self.quantizer.set_num_codebooks(n)
+
+    @property
+    def cardinality(self):
+        """Cardinality of each codebook."""
+        return self.quantizer.bins
+
+    def preprocess(
+            self, x: torch.Tensor
+    ) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        scale: tp.Optional[torch.Tensor]
+        if self.renormalize:
+            mono = x.mean(dim=1, keepdim=True)
+            volume = mono.pow(2).mean(dim=2, keepdim=True).sqrt()
+            scale = 1e-8 + volume
+            x = x / scale
+            scale = scale.view(-1, 1)
+        else:
+            scale = None
+        return x, scale
+
+    def postprocess(self,
+                    x: torch.Tensor,
+                    scale: tp.Optional[torch.Tensor] = None) -> torch.Tensor:
+        if scale is not None:
+            assert self.renormalize
+            x = x * scale.view(-1, 1, 1)
+        return x
+
+    def forward_with_sum_loss(
+            self,
+            x: torch.Tensor,
+            sum_loss_multiplier: float = 1.) -> qt.QuantizedResult:
+        """ Forward pass enforcing additivity in the latent space:
+            Q(x1) + Q(x2) = Q(x1 + x2)
+        """
+        raise NotImplementedError()
+
+        # if we are using sum_loss, batch size needs to be even
+        if x.shape[0] % 2 != 0:
+            raise ValueError("Batch size needs to be even to use sum loss. "
+                             f"Received {x.shape}")
+        # y = x1 + x2
+        y = x.clone()
+        half_bs: int = x.shape[0] // 2
+        y = y[:half_bs] + y[half_bs:]
+
+        # encode y: E(y)
+        y, y_scale = self.preprocess(y)
+        encoded_y = self.encoder(y)
+
+        # encode x: E(x)
+        length = x.shape[-1]
+        x, scale = self.preprocess(x)
+        encoded_x = self.encoder(x)
+
+        # quantize x = x1, x2: get Q(x) = Q(x1), Q(x2)
+        x_quantized: QuantizedResult = self.quantizer(encoded_x,
+                                                      self.frame_rate)
+
+        # quantize y = x1 + x2: get Q(y) = Q(x1 + x2)
+        y_quantized = self.quantizer(encoded_y, self.frame_rate)  # type: ignore
+
+        # sum_of_quantized_layers = Q(x1) + Q(x2)
+        sum_of_quantized_layers = (
+            x_quantized.quantized_layers[:half_bs] +  # type: ignore
+            x_quantized.quantized_layers[half_bs:])  # type: ignore
+
+        # quantization_of_sum = Q(x1 + x2) = Q(y)
+        quantization_of_sum_layers = y_quantized.quantized_layers
+
+        # compute sum_loss with L2.  prediction: Q(x1) + Q(x2) target: Q(x1+x2)
+        sum_loss = nn.functional.mse_loss(sum_of_quantized_layers,
+                                          quantization_of_sum_layers)
+        sum_loss = sum_loss * sum_loss_multiplier
+        x_quantized.sum_loss = sum_loss
+
+        # decode Q(x) to get D(x) = D(x1), D(x2)
+        decoded_x = self.decoder(x_quantized.x)
+
+        # remove extra padding added by the encoder and decoder
+        assert decoded_x.shape[-1] >= length, (decoded_x.shape[-1], length)
+        decoded_x = decoded_x[..., :length]
+        # put in x_quantized.x the decoded version to return it
+        x_quantized.x = self.postprocess(decoded_x, scale)
+
+        # compute informative losses (metrics)
+        with torch.no_grad():
+
+            # only take last layer (sum of all layers)
+            sum_of_quantized = sum_of_quantized_layers[:, -1, ...]
+            quantization_of_sum = quantization_of_sum_layers[:, -1, ...]
+
+            # decode Q(y) to get D(y)
+            decoded_y = self.decoder(y_quantized.x)
+            assert decoded_y.shape[-1] >= length, (decoded_y.shape[-1], length)
+            decoded_y = decoded_y[..., :length]
+
+            # decode Q(x1) + Q(x2) to get D(Q(x1) + Q(x2))
+            decoded_sum_of_quantized = self.decoder(sum_of_quantized)
+            assert decoded_sum_of_quantized.shape[-1] >= length, (
+                decoded_sum_of_quantized.shape[-1], length)
+            decoded_sum_of_quantized = decoded_sum_of_quantized[..., :length]
+
+            # sum_of_decoded_quantized = D(Q(x1)) + D(Q(x2))
+            sum_of_decoded_quantized = decoded_x[:half_bs] + decoded_x[half_bs:]
+
+            # variable names recap:
+            # sum_of_quantized = Q(x1) + Q(x2)
+            # quantization_of_sum = Q(x1 + x2)
+            # decoded_sum_of_quantized = D(Q(x1) + Q(x2))
+            # sum_of_decoded_quantized = D(Q(x1)) + D(Q(x2))
+
+            # cosine similarity between Q(x1) + Q(x2) and Q(x1 + x2)
+            cos_sim = torch.nn.functional.cosine_similarity(sum_of_quantized,
+                                                            quantization_of_sum,
+                                                            dim=2).mean()
+
+            # recon-2
+            # sisdr1:       prediction: D(Q(x1)) + D(Q(x2))     target: D(y)
+            sisdr1 = scale_invariant_signal_distortion_ratio(
+                sum_of_decoded_quantized, decoded_y).mean()
+
+            # comp-1
+            # sisdr2:       prediction: D(Q(x1) + Q(x2))        target: D(y)
+            sisdr2 = scale_invariant_signal_distortion_ratio(
+                decoded_sum_of_quantized, decoded_y).mean()
+
+            # recon-1
+            # sisdr3:       prediction: D(y)                    target: y
+            sisdr3 = scale_invariant_signal_distortion_ratio(decoded_y,
+                                                             y).mean()
+
+            # comp-2
+            # sisdr4:       prediction: D(Q(x1) + Q(x2))        target: y
+            sisdr4 = scale_invariant_signal_distortion_ratio(
+                decoded_sum_of_quantized, y).mean()
+
+        metrics: tp.Dict[str, float] = {
+            "quantizer/cos1": cos_sim.item(),
+            "quantizer/recon2": sisdr1.item(),
+            "quantizer/comp1": sisdr2.item(),
+            "quantizer/recon1": sisdr3.item(),
+            "quantizer/comp2": sisdr4.item(),
+        }
+        x_quantized.metrics = metrics
+
+        return x_quantized
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        sum_loss_amount: float = 0.,
+    ) -> qt.QuantizedResult:
+
+        assert x.dim() == 3
+        if sum_loss_amount > 0:
+            return self.forward_with_sum_loss(x, sum_loss_amount)
+
+        length = x.shape[-1]
+        x, scale = self.preprocess(x)
+        emb = self.encoder(x)
+        q_res: QuantizedResult = self.quantizer(emb, self.frame_rate)
+        out = self.decoder(q_res.x)
+
+        # remove extra padding added by the encoder and decoder
+        assert out.shape[-1] >= length, (out.shape[-1], length)
+        out = out[..., :length]
+        q_res.x = self.postprocess(out, scale)
+
+        return q_res
+
+    def quantize_embedding(self, emb: torch.Tensor) -> QuantizedResult:
+        return self.quantizer(emb, self.frame_rate)
+
+    def quantize(self, x: torch.Tensor) -> QuantizedResult:
+        """ Pass x through encoder and quantizer and return QuantizedResult"""
+        assert x.dim() == 3
+        emb = self.encoder(x)
+        quantized: QuantizedResult = self.quantizer(emb, self.frame_rate)
+        return quantized
+
+    def dequantize(self, x: torch.Tensor | QuantizedResult) -> torch.Tensor:
+        if isinstance(x, QuantizedResult):
+            x = x.x
+        decoded: torch.Tensor = self.decoder(x)
+        return decoded
+
+    def encode(
+        self,
+        x: torch.Tensor,
+        return_scales: bool = False
+        # ) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]] | torch.Tensor:
+    ) -> torch.Tensor:
+        """Encode the given input tensor to quantized representation along with scale parameter.
+
+        Args:
+            x (torch.Tensor): Float tensor of shape [B, C, T]
+
+        Returns:
+            codes, scale (tuple of torch.Tensor, torch.Tensor): Tuple composed of:
+                codes a float tensor of shape [B, K, T] with K the number of codebooks used and T the timestep.
+                scale a float tensor containing the scale for audio renormalizealization.
+        """
+        assert x.dim() == 3
+        x, scale = self.preprocess(x)
+        emb = self.encoder(x)
+        codes = self.quantizer.encode(emb)
+        # if return_scales:
+        #     return codes, scale
+        return codes
+
+    def decode(self,
+               codes: torch.Tensor,
+               scale: tp.Optional[torch.Tensor] = None):
+        """Decode the given codes to a reconstructed representation, using the scale to perform
+        audio denormalization if needed.
+
+        Args:
+            codes (torch.Tensor): Int tensor of shape [B, K, T]
+            scale (torch.Tensor, optional): Float tensor containing the scale value.
+
+        Returns:
+            out (torch.Tensor): Float tensor of shape [B, C, T], the reconstructed audio.
+        """
+        emb = self.decode_latent(codes)
+        out = self.decoder(emb)
+        out = self.postprocess(out, scale)
+        # out contains extra padding added by the encoder and decoder
+        return out
+
+    def decode_latent(self, codes: torch.Tensor):
+        """Decode from the discrete codes to continuous latent space."""
+        return self.quantizer.decode(codes)
+
+    @staticmethod
+    def from_pretrained(name: str):
+        if name == "facebook/encodec_32khz":
+            model = torch.load(cfg.weights_dir() / "encodec_32khz.pt")
+        else:
+            raise NotImplementedError()
+
+        return model
+
+    @staticmethod
+    def from_params(params: hp.EncodecParams):
+
+        seanet_params = params.seanet_params
+        qt_params = params.quantizer_params
+        encoder: SEANetEncoder = SEANetEncoder(**seanet_params.__dict__)
+        decoder: SEANetDecoder = SEANetDecoder(**seanet_params.__dict__)
+        quantizer: qt.ResidualVectorQuantizer = qt.ResidualVectorQuantizer(
+            **qt_params.__dict__)
+
+        model = EncodecModel(encoder,
+                             decoder,
+                             quantizer,
+                             params.sample_rate,
+                             channels=1,
+                             causal=params.seanet_params.causal,
+                             renormalize=False)
+
+        if params.weights is not None:
+            weights = torch.load(Path(params.weights),
+                                 weights_only=True,
+                                 map_location="cpu")
+            model.load_state_dict(weights)
+
+        return model
+
+    # @staticmethod
+    # def _get_model(target_bandwidths: tp.List[float],
+    #                sample_rate: int = 24_000,
+    #                channels: int = 1,
+    #                causal: bool = True,
+    #                model_norm: str = 'weight_norm',
+    #                audio_normalize: bool = False,
+    #                segment: tp.Optional[float] = None,
+    #                name: str = 'unset'):
+    #     encoder = m.SEANetEncoder(channels=channels,
+    #                               norm=model_norm,
+    #                               causal=causal)
+    #     decoder = m.SEANetDecoder(channels=channels,
+    #                               norm=model_norm,
+    #                               causal=causal)
+    #     n_q = int(1000 * target_bandwidths[-1] //
+    #               (math.ceil(sample_rate / encoder.hop_length) * 10))  # = 32
+    #     quantizer = qt.ResidualVectorQuantizer(
+    #         dimension=encoder.dimension,
+    #         n_q=n_q,
+    #         bins=1024,
+    #     )
+    #     model = EncodecModel(
+    #         encoder,
+    #         decoder,
+    #         quantizer,
+    #         # target_bandwidths,
+    #         sample_rate,
+    #         channels,
+    #         causal=True,
+    #         renormalize=audio_normalize,
+    #         # segment=segment,
+    #         # name=name,
+    #     )
+    #     return model
+
+    # @staticmethod
+    # def _get_pretrained(checkpoint_name: str,
+    #                     repository: tp.Optional[Path] = None):
+    #     if repository is not None:
+    #         if not repository.is_dir():
+    #             raise ValueError(f"{repository} must exist and be a directory.")
+    #         file = repository / checkpoint_name
+    #         checksum = file.stem.split('-')[1]
+    #         _check_checksum(file, checksum)
+    #         return torch.load(file)
+    #     else:
+    #         url = _get_checkpoint_url(cfg.ROOT_URL, checkpoint_name)
+    #         return torch.hub.load_state_dict_from_url(
+    #             url, map_location=cfg.device, check_hash=True)  # type:ignore
+
+    # @staticmethod
+    # def encodec_model_24khz(pretrained: bool = True,
+    #                         repository: tp.Optional[Path] = None):
+    #     """Return the pretrained causal 24khz model.
+    #     """
+    #     if repository:
+    #         assert pretrained
+    #     target_bandwidths = [1.5, 3., 6, 12., 24.]
+    #     checkpoint_name = 'encodec_24khz-d7cc33bc.th'
+    #     sample_rate = 24_000
+    #     channels = 1
+    #     model = EncodecModel._get_model(
+    #         target_bandwidths,
+    #         sample_rate,
+    #         channels,
+    #         causal=True,
+    #         model_norm='weight_norm',
+    #         audio_normalize=False,
+    #         name='encodec_24khz' if pretrained else 'unset')
+    #     if pretrained:
+    #         state_dict = EncodecModel._get_pretrained(checkpoint_name,
+    #                                                   repository)
+    #         model.load_state_dict(state_dict)
+    #     model.eval()
+    #     return model

models/lightning_musicgen.py

@@ -0,0 +1,622 @@
+import itertools
+from pathlib import Path
+import lightning as L
+from typing import Dict, Any, List, Mapping, Type, Optional, Sequence
+from torch import Tensor
+import torch
+from dataclasses import dataclass
+from lightning.pytorch import utilities as lightning_utils
+from tqdm import tqdm
+import loralib as lora
+
+import hyperparameters as hp
+from conditioning.beat_embedder import Beat, SinusoidalBeatEmbedder
+from conditioning.condition_dispatcher import ConditionDispatcher
+from conditioning.condition_provider import ConditionProvider
+from conditioning.condition_type import ConditionType
+from conditioning.conditioning_method import ConditioningMethod
+from conditioning.embedded_condition import EmbeddedCondition
+from conditioning.t5embedder import T5Embedder, T5EmbedderCPU, T5EmbedderGPU
+from models.encodec import EncodecModel
+from conditioning.prompt_processor import (
+    DefaultPromptProcessor, InterleavedContextPromptProcessor, PromptProcessor,
+    StraightContextPromptProcessor)
+from models.loss import compute_cross_entropy
+from models.musicgen_lm import MusicgenLm
+from utils.audio import load_audio, save_audio
+from utils.inspection import sanity_check
+from utils.sample import eval_decorator, sample_top_k
+import config as cfg
+
+
+class LightningMusicgen(L.LightningModule):
+
+    def __init__(self, params: hp.MusicgenParams):
+        super().__init__()
+        self.params: hp.MusicgenParams = params
+
+        # instantiate EnCodec
+        self.encodec_model: EncodecModel = EncodecModel.from_params(
+            self.params.encodec_params)
+        self.sample_rate: int = self.encodec_model.sample_rate
+        self.n_q: int = self.encodec_model.num_codebooks
+        self.special_token: int = 2048
+        assert self.n_q == 4
+
+        # freeze encodec
+        for p in self.encodec_model.parameters():
+            p.requires_grad = False
+
+        # instantiate prompt processor
+        self.prompt_processor: PromptProcessor = self.params.prompt_processor_params.model_class(
+            self.encodec_model,
+            self.special_token,
+            keep_only_valid_steps=params.prompt_processor_params.
+            keep_only_valid_steps,
+            context_dropout=params.prompt_processor_params.context_dropout)
+
+        # check consistency in prompt processor and lm params
+        if self.prompt_processor.uses_sep_token:
+            if self.params.lm_params.sep_token is None:
+                raise cfg.ConfigurationError(
+                    "This prompt processor requires an LM with "
+                    "support for a separator token.")
+
+        # instantiate lm
+        self.lm: MusicgenLm = MusicgenLm(self.params.lm_params)
+
+        # check consistency in conditioning parameters
+        if (self.params.conditioning_params.embedder_types.keys()
+                != self.params.conditioning_params.conditioning_methods.keys()):
+            t1 = set(
+                p.value
+                for p in self.params.conditioning_params.embedder_types.keys())
+            t2 = set(
+                p.value for p in
+                self.params.conditioning_params.conditioning_methods.keys())
+            raise ValueError(
+                "Embeddings produced by the condition provider don't match the "
+                f"conditioning methods given in params. "
+                f"processed conditions: {t1} "
+                f"conditioning methods: {t2}")
+
+        # instantiate condition provider (embedder)
+        self.condition_provider = ConditionProvider(
+            self.params.lm_params.dim,
+            embedder_types=self.params.conditioning_params.embedder_types,
+        )
+
+        if self.params.lm_params.weights is not None:
+            w = torch.load(Path(self.params.lm_params.weights),
+                           map_location=None,
+                           weights_only=True)
+            newstatedict = {
+                "condition_provider.embedders.description.output_proj.weight":
+                    w['conditioner.output_proj.weight'],
+                "condition_provider.embedders.description.output_proj.bias":
+                    w['conditioner.output_proj.bias']
+            }
+            self.load_state_dict(newstatedict, strict=False)
+
+        # instantiate condition dispatcher (contains fusers)
+        self.condition_dispatcher: ConditionDispatcher = ConditionDispatcher(
+            self.params.conditioning_params.conditioning_methods,
+            self.params.lm_params.dim,
+            self.params.conditioning_params.conditioning_dropout)
+
+        # inject lora if needed
+        if self.params.lora_params is not None:
+            self._inject_lora()
+
+        self.save_hyperparameters()
+
+    def _inject_lora(self):
+        assert self.params.lora_params is not None
+
+        # freeze the decoder model
+        # for p in self.lm.decoder.parameters():
+        #     p.requires_grad = False
+
+        layers: torch.nn.ModuleList = self.lm.decoder.attn_layers.layers
+        # inject lora in every attention layer
+        for att_idx in range(self.params.lm_params.n_layers):
+            for sublayer_idx in range(2):
+                layeridx = att_idx * 3 + sublayer_idx
+                sublayer: torch.nn.ModuleList = layers[layeridx]  # type: ignore
+                # for all layer types we want to swap
+                for layername in self.params.lora_params.layers:
+                    source_layer = sublayer[1].__getattr__(f"to_{layername}")
+                    new_layer = lora.Linear(source_layer.in_features,
+                                            source_layer.out_features,
+                                            self.params.lora_params.r,
+                                            self.params.lora_params.alpha,
+                                            self.params.lora_params.dropout,
+                                            bias=False)
+                    with torch.no_grad():
+                        new_layer.weight.data.copy_(source_layer.weight.data)
+                    if layername == "q":
+                        sublayer[1].to_q = new_layer
+                    elif layername == "k":
+                        sublayer[1].to_k = new_layer
+                    elif layername == "v":
+                        sublayer[1].to_v = new_layer
+                    elif layername == "out":
+                        sublayer[1].to_out = new_layer
+                    else:
+                        raise RuntimeError(f"unknown layer name {layername}")
+                    # sublayer[1].__setattr__(f"to_{layername}", new_layer)
+
+        lora.mark_only_lora_as_trainable(self)
+
+    def configure_optimizers(self):  # type: ignore
+        opt = torch.optim.AdamW(self.parameters(),
+                                lr=1e-5,
+                                betas=(0.9, 0.95),
+                                weight_decay=0.001)
+        return opt
+        # n_warmup_steps: int = 1000
+
+        # projections = (
+        #     p for n, p in self.named_parameters() if "output_proj" in n)
+        # embeddings = (p for n, p in self.named_parameters() if "token_emb" in n)
+        # warmup_params = itertools.chain(projections, embeddings)
+        # others = (p for n, p in self.named_parameters()
+        #           if "output_proj" not in n and "token_emb" not in n)
+
+        # multigroup_optim = torch.optim.AdamW(
+        #     ({
+        #         "params": warmup_params
+        #     }, {
+        #         "params": others
+        #     }),
+        #     lr=2e-5,
+        #     betas=(0.9, 0.95),
+        #     weight_decay=0.1,
+        # )
+
+        # lambda_encoder = lambda x: (1 + 1.5 * (1 - (x / n_warmup_steps))
+        #                            ) if x < n_warmup_steps else 1.
+        # lambda_decoder = lambda x: 0. if x < n_warmup_steps else 1.
+
+        # multigroup_scheduler = torch.optim.lr_scheduler.LambdaLR(
+        #     multigroup_optim,
+        #     lr_lambda=[lambda_encoder, lambda_decoder],
+        # )
+
+        # scheduler_config = {
+        #     "scheduler": multigroup_scheduler,
+        #     "interval": "step"
+        # }
+
+        # return {"optimizer": multigroup_optim, "lr_scheduler": scheduler_config}
+
+    # def on_train_batch_end(self, outputs: Tensor | Mapping[str, Any] | None,
+    #                        batch: Any, batch_idx: int) -> None:
+    # def on_before_optimizer_step(self, optimizer):
+    #     decoder_grads = lightning_utils.grad_norm(self.lm.decoder, 2)
+    #     self.log_dict(decoder_grads)
+
+    def training_step(self, batch, batch_idx) -> Tensor:
+        self.train(True)
+        loss = self.run_step(batch)
+        self.log(
+            "train/loss",
+            loss,
+            # prog_bar=True,
+            batch_size=len(batch["target"]),
+            # sync_dist=True,
+            # on_step=True,
+        )
+        self.log("global_step", self.global_step, prog_bar=True, logger=False)
+        if self._trainer is not None and self.lr_schedulers() is not None:
+            self.log(
+                "train/new_params_lr",
+                self.lr_schedulers().get_last_lr()[0],  # type: ignore
+                prog_bar=True)
+            self.log(
+                "train/old_params_lr",
+                self.lr_schedulers().get_last_lr()[1],  # type: ignore
+                prog_bar=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx) -> Tensor:
+        self.train(False)
+        with torch.no_grad():
+            loss = self.run_step(batch)
+        self.log(
+            "val/loss",
+            loss,
+            prog_bar=True,
+            batch_size=len(batch["target"]),
+            sync_dist=True,
+        )
+        # sanity_check(self, interrupt=True)
+        return loss
+
+    def run_step(self, batch: Dict[str, Any]) -> Tensor:
+        """
+            Expects batch to be a dictionary like: 
+            {
+                "target": Tensor,
+                "context": Tensor,     - optional
+                "style": Tensor,       - optional
+                "description": string, - optional
+            }
+        """
+
+        # call prompt pre-processor
+        (prompt, prompt_mask, target,
+         decode_logits_fn) = self.prompt_processor.preprocess(batch)
+
+        attention_mask = prompt_mask.sum(dim=-2) > 0
+
+        # embed/encode conditioning data
+        processed_conditions: Dict[ConditionType, EmbeddedCondition] = (
+            self.condition_provider.process_conditions(batch))
+
+        # dispatch eatch conditioning to the proper method, fusing if necessary
+        method_to_cond: Dict[ConditioningMethod,
+                             EmbeddedCondition] = self.condition_dispatcher(
+                                 processed_conditions)
+
+        # call language model
+        logits = self.lm(
+            x=prompt,
+            attention_mask=attention_mask,
+            cross_attention_input=method_to_cond.get(
+                ConditioningMethod.CROSS_ATTENTION),
+            prepend_embeds=method_to_cond.get(ConditioningMethod.INPUT_PREPEND),
+            sum_embeds=method_to_cond.get(ConditioningMethod.INPUT_SUM),
+        )
+
+        # de-interleave logits and postprocess prompt
+        logits, logits_mask = decode_logits_fn(logits)
+
+        # compute cross entropy
+        cross_entropy_loss, _ = compute_cross_entropy(logits, target,
+                                                      logits_mask)
+
+        return cross_entropy_loss
+
+    def sample_next_token(
+            self, current_sequence: Tensor, attention_mask: Tensor,
+            method_to_cond: Dict[ConditioningMethod,
+                                 EmbeddedCondition]) -> Tensor:
+
+        if not current_sequence.isfinite().all():
+            if current_sequence.isnan().any():
+                print(f"Before forward pass some logits are nan")
+            else:
+                print(f"Before forward pass some logits are not finite")
+
+        # call language model
+        logits: Tensor = self.lm(
+            x=current_sequence,
+            attention_mask=attention_mask,
+            cross_attention_input=method_to_cond.get(
+                ConditioningMethod.CROSS_ATTENTION),
+            prepend_embeds=method_to_cond.get(ConditioningMethod.INPUT_PREPEND),
+            sum_embeds=method_to_cond.get(ConditioningMethod.INPUT_SUM),
+        )
+
+        if not logits.isfinite().all():
+            if logits.isnan().any():
+                print(f"After forward pass some logits are nan")
+            else:
+                print(f"After forward pass some logits are not finite")
+
+        # classifier-free guidance
+        cond_logits, uncond_logits = logits.split(
+            current_sequence.shape[0] // 2,
+            dim=0,
+        )
+        logits = uncond_logits + (cond_logits - uncond_logits) * 3.0
+
+        # get logits for last token
+        logits = logits.permute(0, 1, 3, 2)  # B, K, card, T
+        logits = logits[..., -1]  # B, K, card,
+
+        # apply softmax
+        probs = torch.softmax(logits, dim=-1)
+
+        # sample
+        next_token = sample_top_k(probs, k=250)
+        return next_token
+
+    def predict_step(self, batch):
+        batch["prog_bar"] = False
+        return self.generate(**batch)
+
+    # @torch.inference_mode()
+    @eval_decorator
+    @torch.no_grad()
+    def generate(self,
+                 n_samples: int,
+                 gen_seconds: float | int,
+                 prompt: Optional[Tensor],
+                 context: Optional[Tensor | List[Tensor]],
+                 style: Optional[Tensor],
+                 beat: Optional[List[Beat]],
+                 description: Optional[List[str]],
+                 context_dropout_mask: Optional[Tensor] = None,
+                 prog_bar: bool = False) -> Tensor:
+        """Run autoregressive generation
+
+        Args:
+            n_samples (int): number of samples to generate (batch size). All other input parameters should match this.
+            gen_seconds (float | int): total length of generation in seconds, including prompt if present.
+            prompt (Optional[Tensor]): a piece of input to continue.
+            context (Optional[Tensor  |  List[Tensor]]): a musical context to generate an accompaniment for.
+            style (Optional[Tensor]): a piece of music to use as stylistic reference.
+            beat (Optional[List[Beat]]): a beat object to follow
+            description (Optional[List[str]]): a list of descriptions to use as conditioning
+            context_dropout_mask (Optional[Tensor], optional): Defaults to None.
+            prog_bar (bool, optional): whether to display a progress bar. Defaults to False.
+
+        Raises:
+            ValueError: _description_
+
+        Returns:
+            Tensor: _description_
+        """
+
+        n_gen_frames = int(self.encodec_model.frame_rate * gen_seconds)
+
+        # generate empty sequence
+        gen_sequence = torch.full((n_samples, self.n_q, n_gen_frames),
+                                  -1,
+                                  dtype=torch.long,
+                                  device=self.device)
+
+        # pre-process prompt to feed to the lm
+        (gen_sequence, gen_mask, start_offset,
+         decode_sequence_fn) = self.prompt_processor.prepare_for_generation(
+             prompt,
+             context,
+             gen_sequence,
+             use_cfg=True,
+             context_dropout_mask=context_dropout_mask)
+
+        # attention mask: in timesteps in which ALL residual layers are invalid,
+        # set attention mask to False.
+        attention_mask = gen_mask.sum(dim=-2) > 0
+
+        # from now on we only need the first part of the gen_mask, the second was cfg
+        gen_mask = gen_mask[:gen_mask.shape[0] // 2]
+
+        # embed/encode conditioning data
+        conditions = {
+            "description": description,
+            # "context": context,
+            "style": style,
+            "beat": beat,
+        }
+
+        # check for compatibilty of conditions
+        for c_name, c_value in conditions.items():
+            if c_value is None:
+                continue
+            condtype: ConditionType = ConditionType(c_name)
+            if condtype.value not in self.condition_provider.embedders:
+                raise ValueError(
+                    f"This version of the model does not support conditioning "
+                    f"with {c_name}. You should pass None.")
+
+        processed_conditions: Dict[ConditionType, EmbeddedCondition] = (
+            self.condition_provider.process_conditions(conditions,
+                                                       duplicate_for_cfg=True,
+                                                       batch_size=n_samples))
+
+        # dispatch eatch conditioning to the proper method, fusing if necessary
+        method_to_cond: Dict[ConditioningMethod,
+                             EmbeddedCondition] = self.condition_dispatcher(
+                                 processed_conditions)
+        # autoregression
+
+        iterator = range(start_offset, gen_sequence.shape[-1])
+        if prog_bar:
+            iterator = tqdm(iterator, desc="generating autoregressively...")
+
+        for offset in iterator:
+            current_sequence = gen_sequence[..., :offset]
+            current_mask = attention_mask[..., :offset]
+            next_token = self.sample_next_token(current_sequence, current_mask,
+                                                method_to_cond)
+            valid_mask = gen_mask[
+                ...,  # TODO: I can't figure out if this is correct or if it matters at all anyways
+                offset:offset + 1].expand(n_samples, -1, -1)
+            next_token[~valid_mask] = self.special_token
+            gen_sequence[:n_samples, :, offset:offset + 1] = torch.where(
+                gen_sequence[:n_samples, :, offset:offset + 1] == -1,
+                next_token,
+                gen_sequence[:n_samples, :, offset:offset + 1],
+            )
+            gen_sequence[n_samples:, :, offset:offset + 1] = torch.where(
+                gen_sequence[n_samples:, :, offset:offset + 1] == -1,
+                next_token,
+                gen_sequence[n_samples:, :, offset:offset + 1],
+            )
+
+            if prog_bar and torch.cuda.is_available():
+                torch.cuda.synchronize()
+
+        assert not (gen_sequence == -1).any()
+        gen_sequence = gen_sequence[:gen_sequence.shape[0] // 2]
+        # assert (gen_sequence == torch.where(
+        #     gen_mask[None, ...].expand(n_samples, -1, -1),
+        #     gen_sequence,
+        #     self.special_token,
+        # )).all()
+
+        out_codes, out_mask = decode_sequence_fn(gen_sequence)
+
+        self.encodec_model.eval()
+        with torch.no_grad():
+            out_audio = self.encodec_model.decode(out_codes)
+        return out_audio
+
+    @staticmethod
+    def load_from_checkpoint_replacing_paths(ckp_path: Path | str):
+        ckp_path = Path(ckp_path)
+
+        def swap_parent(filepath: Path | str, new_parent: Path) -> Path | str:
+            if isinstance(filepath, Path):
+                typeout = Path
+            elif isinstance(filepath, str):
+                typeout = str
+            else:
+                raise RuntimeError("expected Path or str")
+
+            filepath = Path(filepath)
+            return typeout(new_parent / filepath.name)
+
+        ckp = torch.load(ckp_path, map_location="cpu")
+        params = ckp["hyper_parameters"]["params"]
+        params.encodec_params.weights = swap_parent(
+            params.encodec_params.weights, cfg.weights_dir())
+        params.lm_params.weights = swap_parent(params.lm_params.weights,
+                                               cfg.weights_dir())
+        model = LightningMusicgen(params)
+        model.load_state_dict(ckp["state_dict"])
+        return model
+
+
+if __name__ == "__main__":
+    from time import time
+
+    device = torch.device("cuda")
+
+    # musicgen params
+    model_params = hp.MusicgenParams(
+        encodec_params=hp.pretrained_encodec_meta_32khz_params,
+        prompt_processor_params=hp.PromptProcessorParams(
+            model_class=InterleavedContextPromptProcessor,
+            keep_only_valid_steps=True,
+            context_dropout=0.5,
+        ),
+        conditioning_params=hp.ConditioningParams(
+            embedder_types={
+                ConditionType.DESCRIPTION: T5EmbedderGPU,
+                ConditionType.BEAT: SinusoidalBeatEmbedder,
+            },
+            conditioning_methods={
+                ConditionType.DESCRIPTION: ConditioningMethod.CROSS_ATTENTION,
+                ConditionType.BEAT: ConditioningMethod.INPUT_PREPEND,
+            },
+            conditioning_dropout=0.5,
+        ),
+        lm_params=hp.PretrainedSmallLmParams(sep_token=2049),
+    )
+
+    model = model_params.instantiate().to(device)
+
+    context = [
+        torch.rand(1, 1, 200_000).to(device),
+        # torch.rand(1, 1, 1234).to(device)
+    ]
+
+    # TEST TRAINING/VALIDATION STEP
+    n_tries = 5
+    for _ in range(n_tries):
+        batch = {
+            "target": torch.rand(1, 1, 320_000).to(device),
+            "context": context,
+            "beat": [
+                Beat(beats=(torch.arange(18) * 16_000).long(),
+                     downbeats=(torch.arange(0, 18, 4) * 16_000).long(),
+                     seq_len=320_000)
+            ],
+            "description": [""],
+        }
+        t0 = time()
+        loss = model.run_step(batch)
+        torch.cuda.synchronize()
+        t1 = time()
+        print(f"training step in {t1 - t0} seconds")
+
+    # print("MODEL COMPILED")
+    # model = torch.compile(model, fullgraph=True,
+    #                       backend="eager")  # type: ignore
+    # n_tries = 5
+    # for _ in range(n_tries):
+    #     batch = {
+    #         "target": torch.rand(2, 1, 1_000).to(device),
+    #         "context": torch.rand(2, 1, 1_000).to(device) * 2,
+    #         "style": torch.rand(2, 1, 1_000).to(device),
+    #         "description": ["", ""],
+    #     }
+    #     t0 = time()
+    #     loss = model.run_step(batch)
+    #     t1 = time()
+    #     print(f"training step in {t1 - t0} seconds")
+
+    # TEST INFERENCE
+    L.seed_everything(42)
+    model.eval()
+
+    # audio1 = load_audio(cfg.AUDIO_DIR / "42cpu.wav").to(device)
+    # audio2 = load_audio(cfg.AUDIO_DIR / "42gpuT5cpu.wav").to(device)
+
+    # PROMPT
+    # prompt = audio1
+    prompt = None
+
+    # CONTEXT
+    # context = torch.rand(2, 1, 320_000).to(device)
+    # context = None
+    # context = torch.cat((audio1, audio2), dim=0)
+    # context = audio1.reshape(1, 1, -1)
+    context = load_audio(cfg.EXP_DIR / "experiment_1" /
+                         "context.wav").to(device).reshape(1, 1, -1)
+    # context[0, ...] = 0
+
+    # STYLE
+    # style = torch.cat((audio1, audio2), dim=0)
+    # style = torch.rand(2, 1, 320_000).to(device)
+    style = None
+
+    # DESCRIPTION
+    description = [
+        ""
+        # "lo-fi chill beat with drums, keyboard and bass playing in a relaxed mood",
+        # "lo-fi chill beat with drums, keyboard and bass playing in a relaxed mood"
+    ]
+    # description = None
+
+    # with torch.autocast(device_type="cuda"):
+    # t0 = time()
+    # gen_audio = model.generate(
+    #     n_samples=len(description),
+    #     gen_seconds=10,
+    #     prompt=prompt,
+    #     context=context,
+    #     style=style,
+    #     description=description,
+    #     prog_bar=True,
+    # )
+    # torch.cuda.synchronize()
+    # t1 = time()
+    # print(f"inference completed in {t1 - t0} seconds")
+
+    # # for i in range(gen_audio.shape[0]):
+    # save_audio(gen_audio, cfg.AUDIO_DIR / f"temp.wav")
+
+    # args = {
+    #     "n_samples": 1,
+    #     "gen_seconds": 10,
+    #     "prompt": prompt,
+    #     "context": context,
+    #     "style": style,
+    #     "description": description,
+    # }
+    # i = iter((args,))
+
+    # trainer = L.Trainer(precision="32", enable_progress_bar=False)
+
+    # t0 = time()
+    # gen_audio = trainer.predict(model, i)
+    # torch.cuda.synchronize()
+    # t1 = time()
+    # print(f"predict completed in {t1 - t0} seconds")
+
+    # save_audio(gen_audio[0], cfg.AUDIO_DIR / f"temp.wav")  # type: ignore

models/loss.py

@@ -0,0 +1,48 @@
+import torch
+from torch import Tensor
+from typing import Tuple, List
+from torch.nn import functional as F
+
+
+def compute_cross_entropy(
+    logits: Tensor,
+    targets: Tensor,
+    mask: Tensor,
+) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+    """Compute cross entropy between multi-codebook targets and model's logits.
+    The cross entropy is computed per codebook to provide codebook-level cross entropy.
+    Valid timesteps for each of the codebook are pulled from the mask, where invalid
+    timesteps are set to 0.
+
+    Args:
+        logits (torch.Tensor): Model's logits of shape [B, K, T, card].
+        targets (torch.Tensor): Target codes, of shape [B, K, T].
+        mask (torch.Tensor): Mask for valid target codes, of shape [B, K, T].
+    Returns:
+        ce (torch.Tensor): Cross entropy averaged over the codebooks
+        ce_per_codebook (list of torch.Tensor): Cross entropy per codebook (detached).
+    """
+    B, K, T = targets.shape
+    assert logits.shape[:-1] == targets.shape
+    assert mask.shape == targets.shape
+    ce = torch.zeros([], device=targets.device)
+    ce_per_codebook: List[Tensor] = []
+    for k in range(K):
+        logits_k = (logits[:, k, ...].contiguous().view(-1, logits.size(-1))
+                   )  # [B x T, card]
+        targets_k = targets[:, k, ...].contiguous().view(-1)  # [B x T]
+        mask_k = mask[:, k, ...].contiguous().view(-1)  # [B x T]
+        ce_targets = targets_k[mask_k]
+        ce_logits = logits_k[mask_k]
+
+        # if the codebook is masked out, the loss is 0
+        if mask_k.sum() == 0:
+            q_ce = torch.tensor(0.0, device=targets.device)
+        else:
+            q_ce = F.cross_entropy(ce_logits, ce_targets)
+
+        ce += q_ce
+        ce_per_codebook.append(q_ce.detach())
+    # average cross entropy across codebooks
+    ce = ce / K
+    return ce, ce_per_codebook

models/modules/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""Torch modules."""
+
+# flake8: noqa
+from .conv import (
+    pad1d,
+    unpad1d,
+    pad_for_conv1d,
+    NormConv1d,
+    NormConvTranspose1d,
+    NormConv2d,
+    NormConvTranspose2d,
+    SConv1d,
+    SConvTranspose1d,
+)
+from .lstm import SLSTM
+from .seanet import SEANetEncoder, SEANetDecoder
+from .codebooks_patterns import DelayedPatternProvider
+# from .transformer import StreamingTransformerEncoder

models/modules/codebooks_patterns.py

@@ -0,0 +1,548 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from collections import namedtuple
+from dataclasses import dataclass
+from functools import lru_cache
+import logging
+import typing as tp
+
+from abc import ABC, abstractmethod
+import torch
+
+LayoutCoord = namedtuple('LayoutCoord', ['t', 'q'])  # (timestep, codebook index)
+PatternLayout = tp.List[tp.List[LayoutCoord]]  # Sequence of coordinates
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Pattern:
+    """Base implementation of a pattern over a sequence with multiple codebooks.
+
+    The codebook pattern consists in a layout, defining for each sequence step
+    the list of coordinates of each codebook timestep in the resulting interleaved sequence.
+    The first item of the pattern is always an empty list in order to properly insert a special token
+    to start with. For convenience, we also keep track of ``n_q`` the number of codebooks used for the pattern
+    and ``timesteps`` the number of timesteps corresponding to the original sequence.
+
+    The pattern provides convenient methods to build and revert interleaved sequences from it:
+    ``build_pattern_sequence`` maps a given a dense input tensor of multi-codebook sequence from [B, K, T]
+        to the interleaved sequence of shape [B, K, S] applying the pattern, with B being the batch size,
+        K being the number of codebooks, T the number of original timesteps and S the number of sequence steps
+        for the output sequence. The unfilled positions are replaced with a special token and the built sequence
+        is returned along with a mask indicating valid tokens.
+    ``revert_pattern_sequence`` maps back an interleaved sequence of shape [B, K, S] to the original alignment
+        of codebooks across timesteps to an output tensor of shape [B, K, T], using again a special token and a mask
+        to fill and specify invalid positions if needed.
+    See the dedicated methods for more details.
+    """
+    # Pattern layout, for each sequence step, we have a list of coordinates
+    # corresponding to the original codebook timestep and position.
+    # The first list is always an empty list in order to properly insert
+    # a special token to start with.
+    layout: PatternLayout
+    timesteps: int
+    n_q: int
+
+    def __post_init__(self):
+        assert len(self.layout) > 0
+        self._validate_layout()
+        self._build_reverted_sequence_scatter_indexes = lru_cache(100)(self._build_reverted_sequence_scatter_indexes)
+        self._build_pattern_sequence_scatter_indexes = lru_cache(100)(self._build_pattern_sequence_scatter_indexes)
+        logger.info("New pattern, time steps: %d, sequence steps: %d", self.timesteps, len(self.layout))
+
+    def _validate_layout(self):
+        """Runs checks on the layout to ensure a valid pattern is defined.
+        A pattern is considered invalid if:
+            - Multiple timesteps for a same codebook are defined in the same sequence step
+            - The timesteps for a given codebook are not in ascending order as we advance in the sequence
+              (this would mean that we have future timesteps before past timesteps).
+        """
+        q_timesteps = {q: 0 for q in range(self.n_q)}
+        for s, seq_coords in enumerate(self.layout):
+            if len(seq_coords) > 0:
+                qs = set()
+                for coord in seq_coords:
+                    qs.add(coord.q)
+                    last_q_timestep = q_timesteps[coord.q]
+                    assert coord.t >= last_q_timestep, \
+                        f"Past timesteps are found in the sequence for codebook = {coord.q} at step {s}"
+                    q_timesteps[coord.q] = coord.t
+                # each sequence step contains at max 1 coordinate per codebook
+                assert len(qs) == len(seq_coords), \
+                    f"Multiple entries for a same codebook are found at step {s}"
+
+    @property
+    def num_sequence_steps(self):
+        return len(self.layout) - 1
+
+    @property
+    def max_delay(self):
+        max_t_in_seq_coords = 0
+        for seq_coords in self.layout[1:]:
+            for coords in seq_coords:
+                max_t_in_seq_coords = max(max_t_in_seq_coords, coords.t + 1)
+        return max_t_in_seq_coords - self.timesteps
+
+    @property
+    def valid_layout(self):
+        valid_step = len(self.layout) - self.max_delay
+        return self.layout[:valid_step]
+
+    def starts_with_special_token(self):
+        return self.layout[0] == []
+
+    def get_sequence_coords_with_timestep(self, t: int, q: tp.Optional[int] = None):
+        """Get codebook coordinates in the layout that corresponds to the specified timestep t
+        and optionally to the codebook q. Coordinates are returned as a tuple with the sequence step
+        and the actual codebook coordinates.
+        """
+        assert t <= self.timesteps, "provided timesteps is greater than the pattern's number of timesteps"
+        if q is not None:
+            assert q <= self.n_q, "provided number of codebooks is greater than the pattern's number of codebooks"
+        coords = []
+        for s, seq_codes in enumerate(self.layout):
+            for code in seq_codes:
+                if code.t == t and (q is None or code.q == q):
+                    coords.append((s, code))
+        return coords
+
+    def get_steps_with_timestep(self, t: int, q: tp.Optional[int] = None) -> tp.List[int]:
+        return [step for step, coords in self.get_sequence_coords_with_timestep(t, q)]
+
+    def get_first_step_with_timesteps(self, t: int, q: tp.Optional[int] = None) -> tp.Optional[int]:
+        steps_with_timesteps = self.get_steps_with_timestep(t, q)
+        return steps_with_timesteps[0] if len(steps_with_timesteps) > 0 else None
+
+    def _build_pattern_sequence_scatter_indexes(self, timesteps: int, n_q: int, keep_only_valid_steps: bool,
+                                                device: tp.Union[torch.device, str] = 'cpu'):
+        """Build scatter indexes corresponding to the pattern, up to the provided sequence_steps.
+
+        Args:
+            timesteps (int): Maximum number of timesteps steps to consider.
+            keep_only_valid_steps (bool): Restrict the pattern layout to match only valid steps.
+            device (torch.device or str): Device for created tensors.
+        Returns:
+            indexes (torch.Tensor): Indexes corresponding to the sequence, of shape [K, S].
+            mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes, of shape [K, S].
+        """
+        assert n_q == self.n_q, f"invalid number of codebooks for the sequence and the pattern: {n_q} != {self.n_q}"
+        assert timesteps <= self.timesteps, "invalid number of timesteps used to build the sequence from the pattern"
+        # use the proper layout based on whether we limit ourselves to valid steps only or not,
+        # note that using the valid_layout will result in a truncated sequence up to the valid steps
+        ref_layout = self.valid_layout if keep_only_valid_steps else self.layout
+        # single item indexing being super slow with pytorch vs. numpy, so we use numpy here
+        indexes = torch.zeros(n_q, len(ref_layout), dtype=torch.long).numpy()
+        mask = torch.zeros(n_q, len(ref_layout), dtype=torch.bool).numpy()
+        # fill indexes with last sequence step value that will correspond to our special token
+        # the last value is n_q * timesteps as we have flattened z and append special token as the last token
+        # which will correspond to the index: n_q * timesteps
+        indexes[:] = n_q * timesteps
+        # iterate over the pattern and fill scattered indexes and mask
+        for s, sequence_coords in enumerate(ref_layout):
+            for coords in sequence_coords:
+                if coords.t < timesteps:
+                    indexes[coords.q, s] = coords.t + coords.q * timesteps
+                    mask[coords.q, s] = 1
+        indexes = torch.from_numpy(indexes).to(device)
+        mask = torch.from_numpy(mask).to(device)
+        return indexes, mask
+
+    def build_pattern_sequence(self, z: torch.Tensor, special_token: int, keep_only_valid_steps: bool = False):
+        """Build sequence corresponding to the pattern from the input tensor z.
+        The sequence is built using up to sequence_steps if specified, and non-pattern
+        coordinates are filled with the special token.
+
+        Args:
+            z (torch.Tensor): Input tensor of multi-codebooks sequence, of shape [B, K, T].
+            special_token (int): Special token used to fill non-pattern coordinates in the new sequence.
+            keep_only_valid_steps (bool): Build a sequence from the pattern up to valid (= fully defined) steps.
+                Steps that are beyond valid steps will be replaced by the special_token in that case.
+        Returns:
+            values (torch.Tensor): Interleaved sequence matching the pattern, of shape [B, K, S] with S
+                corresponding either to the sequence_steps if provided, otherwise to the length of the pattern.
+            indexes (torch.Tensor): Indexes corresponding to the interleaved sequence, of shape [K, S].
+            mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes of shape [K, S].
+        """
+        B, K, T = z.shape
+        indexes, mask = self._build_pattern_sequence_scatter_indexes(
+            T, K, keep_only_valid_steps=keep_only_valid_steps, device=str(z.device)
+        )
+        z = z.view(B, -1)
+        # we append the special token as the last index of our flattened z tensor
+        z = torch.cat([z, torch.zeros_like(z[:, :1]) + special_token], dim=1)
+        values = z[:, indexes.view(-1)]
+        values = values.view(B, K, indexes.shape[-1])
+        return values, indexes, mask
+
+    def _build_reverted_sequence_scatter_indexes(self, sequence_steps: int, n_q: int,
+                                                 keep_only_valid_steps: bool = False,
+                                                 is_model_output: bool = False,
+                                                 device: tp.Union[torch.device, str] = 'cpu'):
+        """Builds scatter indexes required to retrieve the original multi-codebook sequence
+        from interleaving pattern.
+
+        Args:
+            sequence_steps (int): Sequence steps.
+            n_q (int): Number of codebooks.
+            keep_only_valid_steps (bool): Build a sequence from the pattern up to valid (= fully defined) steps.
+                Steps that are beyond valid steps will be replaced by the special_token in that case.
+            is_model_output (bool): Whether to keep the sequence item corresponding to initial special token or not.
+            device (torch.device or str): Device for created tensors.
+        Returns:
+            indexes (torch.Tensor): Indexes for reconstructing the output, of shape [K, T].
+            mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes of shape [K, T].
+        """
+        ref_layout = self.valid_layout if keep_only_valid_steps else self.layout
+        # TODO(jade): Do we want to further truncate to only valid timesteps here as well?
+        timesteps = self.timesteps
+        assert n_q == self.n_q, f"invalid number of codebooks for the sequence and the pattern: {n_q} != {self.n_q}"
+        assert sequence_steps <= len(ref_layout), \
+            f"sequence to revert is longer than the defined pattern: {sequence_steps} > {len(ref_layout)}"
+
+        # ensure we take the appropriate indexes to keep the model output from the first special token as well
+        if is_model_output and self.starts_with_special_token():
+            ref_layout = ref_layout[1:]
+
+        # single item indexing being super slow with pytorch vs. numpy, so we use numpy here
+        indexes = torch.zeros(n_q, timesteps, dtype=torch.long).numpy()
+        mask = torch.zeros(n_q, timesteps, dtype=torch.bool).numpy()
+        # fill indexes with last sequence step value that will correspond to our special token
+        indexes[:] = n_q * sequence_steps
+        for s, sequence_codes in enumerate(ref_layout):
+            if s < sequence_steps:
+                for code in sequence_codes:
+                    if code.t < timesteps:
+                        indexes[code.q, code.t] = s + code.q * sequence_steps
+                        mask[code.q, code.t] = 1
+        indexes = torch.from_numpy(indexes).to(device)
+        mask = torch.from_numpy(mask).to(device)
+        return indexes, mask
+
+    def revert_pattern_sequence(self, s: torch.Tensor, special_token: int, keep_only_valid_steps: bool = False):
+        """Revert a sequence built from the pattern back to the original multi-codebook sequence without interleaving.
+        The sequence is reverted using up to timesteps if specified, and non-pattern coordinates
+        are filled with the special token.
+
+        Args:
+            s (torch.Tensor): Interleaved sequence tensor obtained from the pattern, of shape [B, K, S].
+            special_token (int or float): Special token used to fill non-pattern coordinates in the new sequence.
+        Returns:
+            values (torch.Tensor): Interleaved sequence matching the pattern, of shape [B, K, T] with T
+                corresponding either to the timesteps if provided, or the total timesteps in pattern otherwise.
+            indexes (torch.Tensor): Indexes corresponding to the interleaved sequence, of shape [K, T].
+            mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes of shape [K, T].
+        """
+        B, K, S = s.shape
+        indexes, mask = self._build_reverted_sequence_scatter_indexes(
+            S, K, keep_only_valid_steps, is_model_output=False, device=str(s.device)
+        )
+        s = s.view(B, -1)
+        # we append the special token as the last index of our flattened z tensor
+        s = torch.cat([s, torch.zeros_like(s[:, :1]) + special_token], dim=1)
+        values = s[:, indexes.view(-1)]
+        values = values.view(B, K, indexes.shape[-1])
+        return values, indexes, mask
+
+    def revert_pattern_logits(self, logits: torch.Tensor, special_token: float, keep_only_valid_steps: bool = False):
+        """Revert model logits obtained on a sequence built from the pattern
+        back to a tensor matching the original sequence.
+
+        This method is similar to ``revert_pattern_sequence`` with the following specificities:
+        1. It is designed to work with the extra cardinality dimension
+        2. We return the logits for the first sequence item that matches the special_token and
+        which matching target in the original sequence is the first item of the sequence,
+        while we skip the last logits as there is no matching target
+        """
+        B, card, K, S = logits.shape
+        indexes, mask = self._build_reverted_sequence_scatter_indexes(
+            S, K, keep_only_valid_steps, is_model_output=True, device=logits.device
+        )
+        logits = logits.reshape(B, card, -1)
+        # we append the special token as the last index of our flattened z tensor
+        logits = torch.cat([logits, torch.zeros_like(logits[:, :, :1]) + special_token], dim=-1)  # [B, card, K x S]
+        values = logits[:, :, indexes.view(-1)]
+        values = values.view(B, card, K, indexes.shape[-1])
+        return values, indexes, mask
+
+
+class CodebooksPatternProvider(ABC):
+    """Abstraction around providing pattern for interleaving codebooks.
+
+    The CodebooksPatternProvider abstraction allows to implement various strategies to
+    define interleaving pattern of sequences composed of multiple codebooks. For a given
+    number of codebooks `n_q`, the pattern provider can generate a specified pattern
+    corresponding to a sequence of `T` timesteps with `n_q` parallel codebooks. This pattern
+    can be used to construct a new sequence from the original codes respecting the specified
+    pattern. The pattern is defined as a list of list of code coordinates, code coordinate
+    being a tuple with the original timestep and codebook to build the new sequence.
+    Note that all patterns must start with an empty list that is then used to insert a first
+    sequence step of special tokens in the newly generated sequence.
+
+    Args:
+        n_q (int): number of codebooks.
+        cached (bool): if True, patterns for a given length are cached. In general
+            that should be true for efficiency reason to avoid synchronization points.
+    """
+    def __init__(self, n_q: int, cached: bool = True):
+        assert n_q > 0
+        self.n_q = n_q
+        self.get_pattern = lru_cache(100)(self.get_pattern)  # type: ignore
+
+    @abstractmethod
+    def get_pattern(self, timesteps: int) -> Pattern:
+        """Builds pattern with specific interleaving between codebooks.
+
+        Args:
+            timesteps (int): Total number of timesteps.
+        """
+        raise NotImplementedError()
+
+
+class DelayedPatternProvider(CodebooksPatternProvider):
+    """Provider for delayed pattern across delayed codebooks.
+    Codebooks are delayed in the sequence and sequence steps will contain codebooks
+    from different timesteps.
+
+    Example:
+        Taking timesteps=4 and n_q=3, delays=None, the multi-codebook sequence:
+        [[1, 2, 3, 4],
+        [1, 2, 3, 4],
+        [1, 2, 3, 4]]
+        The resulting sequence obtained from the returned pattern is:
+        [[S, 1, 2, 3, 4],
+        [S, S, 1, 2, 3],
+        [S, S, S, 1, 2]]
+        (with S being a special token)
+
+    Args:
+        n_q (int): Number of codebooks.
+        delays (list of int, optional): Delay for each of the codebooks.
+            If delays not defined, each codebook is delayed by 1 compared to the previous one.
+        flatten_first (int): Flatten the first N timesteps.
+        empty_initial (int): Prepend with N empty list of coordinates.
+    """
+    def __init__(self, n_q: int, delays: tp.Optional[tp.List[int]] = None,
+                 flatten_first: int = 0, empty_initial: int = 0):
+        super().__init__(n_q)
+        if delays is None:
+            delays = list(range(n_q))
+        self.delays = delays
+        self.flatten_first = flatten_first
+        self.empty_initial = empty_initial
+        assert len(self.delays) == self.n_q
+        assert sorted(self.delays) == self.delays
+
+    def get_pattern(self, timesteps: int) -> Pattern:
+        omit_special_token = self.empty_initial < 0
+        out: PatternLayout = [] if omit_special_token else [[]]
+        max_delay = max(self.delays)
+        if self.empty_initial:
+            out += [[] for _ in range(self.empty_initial)]
+        if self.flatten_first:
+            for t in range(min(timesteps, self.flatten_first)):
+                for q in range(self.n_q):
+                    out.append([LayoutCoord(t, q)])
+        for t in range(self.flatten_first, timesteps + max_delay):
+            v = []
+            for q, delay in enumerate(self.delays):
+                t_for_q = t - delay
+                if t_for_q >= self.flatten_first:
+                    v.append(LayoutCoord(t_for_q, q))
+            out.append(v)
+        return Pattern(out, n_q=self.n_q, timesteps=timesteps)
+
+
+class ParallelPatternProvider(DelayedPatternProvider):
+    """Provider for parallel pattern across codebooks.
+    This pattern provider is a special case of the delayed pattern with actually no delay,
+    hence delays=repeat(0, n_q).
+
+    Args:
+        n_q (int): Number of codebooks.
+        empty_initial (int): Prepend with N empty list of coordinates.
+    """
+    def __init__(self, n_q: int, empty_initial: int = 0):
+        super().__init__(n_q, [0] * n_q, empty_initial=empty_initial)
+
+
+class UnrolledPatternProvider(CodebooksPatternProvider):
+    """Provider for unrolling codebooks pattern.
+    This pattern provider enables to represent the codebook flattened completely or only to some extend
+    while also specifying a given delay between the flattened codebooks representation, allowing to
+    unroll the codebooks in the sequence.
+
+    Example:
+        1. Flattening of the codebooks.
+        By default, the pattern provider will fully flatten the codebooks such as flattening=range(n_q),
+        taking n_q = 3 and timesteps = 4:
+        [[1, 2, 3, 4],
+         [1, 2, 3, 4],
+         [1, 2, 3, 4]]
+        will result into:
+        [[S, S, 1, S, S, 2, S, S, 3, S, S, 4],
+         [S, 1, S, S, 2, S, S, 3, S, S, 4, S],
+         [1, S, S, 2, S, S, 3, S, S, 4, S, S]]
+        2. Partial flattening of the codebooks. The ``flattening`` parameter allows to specify the inner step
+        for each of the codebook, allowing to define which codebook to flatten (or keep in parallel), for example
+        taking n_q = 3, timesteps = 4 and flattening = [0, 1, 1]:
+        [[1, 2, 3, 4],
+         [1, 2, 3, 4],
+         [1, 2, 3, 4]]
+        will result into:
+        [[S, 1, S, S, 2, S, S, 3, S, S, 4, S],
+         [S, 1, S, S, 2, S, S, 3, S, S, 4, S],
+         [1, S, S, 2, S, S, 3, S, S, 4, S, S]]
+        3. Flattening with delay. The ``delay`` parameter allows to further unroll the sequence of codebooks
+        allowing to specify the delay per codebook. Note that the delay between codebooks flattened to the
+        same inner timestep should be coherent. For example, taking n_q = 3, timesteps = 4, flattening = [0, 1, 1]
+        and delays = [0, 3, 3]:
+        [[1, 2, 3, 4],
+         [1, 2, 3, 4],
+         [1, 2, 3, 4]]
+        will result into:
+        [[S, S, S, 1, S, 2, S, 3, S, 4],
+         [S, S, S, 1, S, 2, S, 3, S, 4],
+         [1, 2, 3, S, 4, S, 5, S, 6, S]]
+
+    Args:
+        n_q (int): Number of codebooks.
+        flattening (list of int, optional): Flattening schema over the codebooks. If not defined,
+            the codebooks will be flattened to 1 codebook per step, meaning that the sequence will
+            have n_q extra steps for each timestep.
+        delays (list of int, optional): Delay for each of the codebooks. If not defined,
+            no delay is added and therefore will default to [0] * ``n_q``.
+            Note that two codebooks that will be flattened to the same inner step
+            should have the same delay, otherwise the pattern is considered as invalid.
+    """
+    FlattenedCodebook = namedtuple('FlattenedCodebook', ['codebooks', 'delay'])
+
+    def __init__(self, n_q: int, flattening: tp.Optional[tp.List[int]] = None,
+                 delays: tp.Optional[tp.List[int]] = None):
+        super().__init__(n_q)
+        if flattening is None:
+            flattening = list(range(n_q))
+        if delays is None:
+            delays = [0] * n_q
+        assert len(flattening) == n_q
+        assert len(delays) == n_q
+        assert sorted(flattening) == flattening
+        assert sorted(delays) == delays
+        self._flattened_codebooks = self._build_flattened_codebooks(delays, flattening)
+        self.max_delay = max(delays)
+
+    def _build_flattened_codebooks(self, delays: tp.List[int], flattening: tp.List[int]):
+        """Build a flattened codebooks representation as a dictionary of inner step
+        and the actual codebook indices corresponding to the flattened codebook. For convenience, we
+        also store the delay associated to the flattened codebook to avoid maintaining an extra mapping.
+        """
+        flattened_codebooks: dict = {}
+        for q, (inner_step, delay) in enumerate(zip(flattening, delays)):
+            if inner_step not in flattened_codebooks:
+                flat_codebook = UnrolledPatternProvider.FlattenedCodebook(codebooks=[q], delay=delay)
+            else:
+                flat_codebook = flattened_codebooks[inner_step]
+                assert flat_codebook.delay == delay, (
+                    "Delay and flattening between codebooks is inconsistent: ",
+                    "two codebooks flattened to the same position should have the same delay."
+                )
+                flat_codebook.codebooks.append(q)
+            flattened_codebooks[inner_step] = flat_codebook
+        return flattened_codebooks
+
+    @property
+    def _num_inner_steps(self):
+        """Number of inner steps to unroll between timesteps in order to flatten the codebooks.
+        """
+        return max([inner_step for inner_step in self._flattened_codebooks.keys()]) + 1
+
+    def num_virtual_steps(self, timesteps: int) -> int:
+        return timesteps * self._num_inner_steps + 1
+
+    def get_pattern(self, timesteps: int) -> Pattern:
+        """Builds pattern for delay across codebooks.
+
+        Args:
+            timesteps (int): Total number of timesteps.
+        """
+        # the PatternLayout is built as a tuple of sequence position and list of coordinates
+        # so that it can be reordered properly given the required delay between codebooks of given timesteps
+        indexed_out: list = [(-1, [])]
+        max_timesteps = timesteps + self.max_delay
+        for t in range(max_timesteps):
+            # for each timestep, we unroll the flattened codebooks,
+            # emitting the sequence step with the corresponding delay
+            for step in range(self._num_inner_steps):
+                if step in self._flattened_codebooks:
+                    # we have codebooks at this virtual step to emit
+                    step_codebooks = self._flattened_codebooks[step]
+                    t_for_q = t + step_codebooks.delay
+                    coords = [LayoutCoord(t, q) for q in step_codebooks.codebooks]
+                    if t_for_q < max_timesteps and t < max_timesteps:
+                        indexed_out.append((t_for_q, coords))
+                else:
+                    # there is no codebook in this virtual step so we emit an empty list
+                    indexed_out.append((t, []))
+        out = [coords for _, coords in sorted(indexed_out)]
+        return Pattern(out, n_q=self.n_q, timesteps=timesteps)
+
+
+class CoarseFirstPattern(CodebooksPatternProvider):
+    """First generates all the codebooks #1 (e.g. coarser), then the remaining ones,
+    potentially with delays.
+
+    ..Warning:: You must always generate the full training duration at test time, for instance,
+        30 seconds, as otherwise, the fine codebooks will start being generated in an unexpected
+        location. This is due to the non causality of the remaining codebooks with respect to
+        the first ones.
+
+    Args:
+        n_q (int): Number of codebooks.
+        delays (list of int, optional): Delay for each of the codebooks.
+            If delays not defined, each codebook is delayed by 1 compared to the previous one.
+    """
+    def __init__(self, n_q: int, delays: tp.Optional[tp.List[int]] = None):
+        super().__init__(n_q)
+        if delays is None:
+            delays = [0] * (n_q - 1)
+        self.delays = delays
+        assert len(self.delays) == self.n_q - 1
+        assert sorted(self.delays) == self.delays
+
+    def get_pattern(self, timesteps: int) -> Pattern:
+        out: PatternLayout = [[]]
+        for t in range(timesteps):
+            out.append([LayoutCoord(t, 0)])
+        max_delay = max(self.delays)
+        for t in range(timesteps + max_delay):
+            v = []
+            for q, delay in enumerate(self.delays):
+                t_for_q = t - delay
+                if t_for_q >= 0:
+                    v.append(LayoutCoord(t_for_q, q + 1))
+            out.append(v)
+        return Pattern(out, n_q=self.n_q, timesteps=timesteps)
+
+
+class MusicLMPattern(CodebooksPatternProvider):
+    """Almost MusicLM style pattern. This is equivalent to full flattening
+    but in a different order.
+
+    Args:
+        n_q (int): Number of codebooks.
+        group_by (int): Number of codebooks to group together.
+    """
+    def __init__(self, n_q: int, group_by: int = 2):
+        super().__init__(n_q)
+        self.group_by = group_by
+
+    def get_pattern(self, timesteps: int) -> Pattern:
+        out: PatternLayout = [[]]
+        for offset in range(0, self.n_q, self.group_by):
+            for t in range(timesteps):
+                for q in range(offset, offset + self.group_by):
+                    out.append([LayoutCoord(t, q)])
+        return Pattern(out, n_q=self.n_q, timesteps=timesteps)

models/modules/conv.py

@@ -0,0 +1,346 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""Convolutional layers wrappers and utilities."""
+
+import math
+import typing as tp
+import warnings
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.utils import spectral_norm
+from torch.nn.utils.parametrizations import weight_norm
+# from torch.nn.utils import weight_norm
+
+from .norm import ConvLayerNorm
+
+CONV_NORMALIZATIONS = frozenset([
+    'none', 'weight_norm', 'spectral_norm', 'time_layer_norm', 'layer_norm',
+    'time_group_norm'
+])
+
+
+# self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
+def apply_parametrization_norm(module: nn.Module,
+                               norm: str = 'none') -> nn.Module:
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'weight_norm':
+        return weight_norm(module)
+    elif norm == 'spectral_norm':
+        return spectral_norm(module)
+    else:
+        # We already check was in CONV_NORMALIZATION, so any other choice
+        # doesn't need reparametrization.
+        return module
+
+
+def get_norm_module(module: nn.Module,
+                    causal: bool = False,
+                    norm: str = 'none',
+                    **norm_kwargs) -> nn.Module:
+    """Return the proper normalization module. If causal is True, this will ensure the returned
+    module is causal, or return an error if the normalization doesn't support causal evaluation.
+    """
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'layer_norm':
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return ConvLayerNorm(module.out_channels, **norm_kwargs)
+    elif norm == 'time_group_norm':
+        if causal:
+            raise ValueError("GroupNorm doesn't support causal evaluation.")
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
+    else:
+        return nn.Identity()
+
+
+def get_extra_padding_for_conv1d(x: torch.Tensor,
+                                 kernel_size: int,
+                                 stride: int,
+                                 padding_total: int = 0) -> int:
+    """See `pad_for_conv1d`.
+    """
+    length = x.shape[-1]
+    n_frames = (length - kernel_size + padding_total) / stride + 1
+    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size -
+                                                         padding_total)
+    return ideal_length - length
+
+
+def pad_for_conv1d(x: torch.Tensor,
+                   kernel_size: int,
+                   stride: int,
+                   padding_total: int = 0):
+    """Pad for a convolution to make sure that the last window is full.
+    Extra padding is added at the end. This is required to ensure that we can rebuild
+    an output of the same length, as otherwise, even with padding, some time steps
+    might get removed.
+    For instance, with total padding = 4, kernel size = 4, stride = 2:
+        0 0 1 2 3 4 5 0 0   # (0s are padding)
+        1   2   3           # (output frames of a convolution, last 0 is never used)
+        0 0 1 2 3 4 5 0     # (output of tr. conv., but pos. 5 is going to get removed as padding)
+            1 2 3 4         # once you removed padding, we are missing one time step !
+    """
+    extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride,
+                                                 padding_total)
+    return F.pad(x, (0, extra_padding))
+
+
+def pad1d(x: torch.Tensor,
+          paddings: tp.Tuple[int, int],
+          mode: str = 'zero',
+          value: float = 0.):
+    """Tiny wrapper around F.pad, just to allow for reflect padding on small input.
+    If this is the case, we insert extra 0 padding to the right before the reflection happen.
+    """
+    length = x.shape[-1]
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left,
+                                                      padding_right)
+    if mode == 'reflect':
+        max_pad = max(padding_left, padding_right)
+        extra_pad = 0
+        if length <= max_pad:
+            extra_pad = max_pad - length + 1
+            x = F.pad(x, (0, extra_pad))
+        padded = F.pad(x, paddings, mode, value)
+        end = padded.shape[-1] - extra_pad
+        return padded[..., :end]
+    else:
+        return F.pad(x, paddings, mode, value)
+
+
+def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
+    """Remove padding from x, handling properly zero padding. Only for 1d!"""
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left,
+                                                      padding_right)
+    assert (padding_left + padding_right) <= x.shape[-1]
+    end = x.shape[-1] - padding_right
+    return x[..., padding_left:end]
+
+
+class NormConv1d(nn.Module):
+    """Wrapper around Conv1d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+
+    def __init__(self,
+                 *args,
+                 causal: bool = False,
+                 norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {},
+                 **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        #print("inputNormConv1d:", x.shape, torch.sum(x), torch.sum(torch.abs(x)), "norm:", self.norm_type, "conv:", self.conv)
+        in_x = x.clone()
+        x = self.conv(x)
+
+        #print("betweenNormConv1d:", x.shape, torch.sum(x), torch.sum(torch.abs(x)))
+        # if torch.isnan(torch.sum(x)).any():
+        #     print("got nan", x.shape, self.conv)
+        #     raise ValueError("wtf")
+        x = self.norm(x)
+        #print("outputNormConv1d:", x.shape, torch.sum(x), torch.sum(torch.abs(x)))
+        return x
+
+
+class NormConv2d(nn.Module):
+    """Wrapper around Conv2d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+
+    def __init__(self,
+                 *args,
+                 norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {},
+                 **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv,
+                                    causal=False,
+                                    norm=norm,
+                                    **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        #print("inputNormConv2d:", x.shape, torch.sum(x), torch.sum(torch.abs(x)))
+        x = self.conv(x)
+        x = self.norm(x)
+        #print("outputNormConv2d:", x.shape, torch.sum(x), torch.sum(torch.abs(x)))
+        return x
+
+
+class NormConvTranspose1d(nn.Module):
+    """Wrapper around ConvTranspose1d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+
+    def __init__(self,
+                 *args,
+                 causal: bool = False,
+                 norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {},
+                 **kwargs):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(
+            nn.ConvTranspose1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        #print("inputNormConvTranspose1d:", x.shape, torch.sum(x), torch.sum(torch.abs(x)))
+        x = self.convtr(x)
+        x = self.norm(x)
+        #print("outputNormConvTranspose1d:", x.shape, torch.sum(x), torch.sum(torch.abs(x)))
+        return x
+
+
+class NormConvTranspose2d(nn.Module):
+    """Wrapper around ConvTranspose2d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+
+    def __init__(self,
+                 *args,
+                 norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {},
+                 **kwargs):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(
+            nn.ConvTranspose2d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.convtr,
+                                    causal=False,
+                                    norm=norm,
+                                    **norm_kwargs)
+
+    def forward(self, x):
+        #print("inputNormConvTranspose2d:", x.shape, torch.sum(x), torch.sum(torch.abs(x)))
+        x = self.convtr(x)
+        x = self.norm(x)
+        #print("outputNormConvTranspose2d:", x.shape, torch.sum(x), torch.sum(torch.abs(x)))
+        return x
+
+
+class SConv1d(nn.Module):
+    """Conv1d with some builtin handling of asymmetric or causal padding
+    and normalization.
+    """
+
+    def __init__(self,
+                 in_channels: int,
+                 out_channels: int,
+                 kernel_size: int,
+                 stride: int = 1,
+                 dilation: int = 1,
+                 groups: int = 1,
+                 bias: bool = True,
+                 causal: bool = False,
+                 norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {},
+                 pad_mode: str = 'reflect'):
+        super().__init__()
+        # warn user on unusual setup between dilation and stride
+        if stride > 1 and dilation > 1:
+            warnings.warn(
+                'SConv1d has been initialized with stride > 1 and dilation > 1'
+                f' (kernel_size={kernel_size} stride={stride}, dilation={dilation}).'
+            )
+        self.conv = NormConv1d(in_channels,
+                               out_channels,
+                               kernel_size,
+                               stride,
+                               dilation=dilation,
+                               groups=groups,
+                               bias=bias,
+                               causal=causal,
+                               norm=norm,
+                               norm_kwargs=norm_kwargs)
+        self.causal = causal
+        self.pad_mode = pad_mode
+
+    def forward(self, x):
+        #print("inputSConv1d:", x.shape, torch.sum(x), torch.sum(torch.abs(x)))
+        B, C, T = x.shape
+        kernel_size = self.conv.conv.kernel_size[0]
+        stride = self.conv.conv.stride[0]
+        dilation = self.conv.conv.dilation[0]
+        padding_total = (kernel_size - 1) * dilation - (stride - 1)
+        extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride,
+                                                     padding_total)
+        if self.causal:
+            # Left padding for causal
+            x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
+        else:
+            # Asymmetric padding required for odd strides
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            x = pad1d(x, (padding_left, padding_right + extra_padding),
+                      mode=self.pad_mode)
+        x = self.conv(x)
+        #print("outputSConv1d:", x.shape, torch.sum(x), torch.sum(torch.abs(x)))
+        return x
+
+
+class SConvTranspose1d(nn.Module):
+    """ConvTranspose1d with some builtin handling of asymmetric or causal padding
+    and normalization.
+    """
+
+    def __init__(self,
+                 in_channels: int,
+                 out_channels: int,
+                 kernel_size: int,
+                 stride: int = 1,
+                 causal: bool = False,
+                 norm: str = 'none',
+                 trim_right_ratio: float = 1.,
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}):
+        super().__init__()
+        self.convtr = NormConvTranspose1d(in_channels,
+                                          out_channels,
+                                          kernel_size,
+                                          stride,
+                                          causal=causal,
+                                          norm=norm,
+                                          norm_kwargs=norm_kwargs)
+        self.causal = causal
+        self.trim_right_ratio = trim_right_ratio
+        assert self.causal or self.trim_right_ratio == 1., \
+            "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
+        assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1.
+
+    def forward(self, x):
+        #print("inputSConvTranspose1d:", x.shape, torch.sum(x), torch.sum(torch.abs(x)))
+        kernel_size = self.convtr.convtr.kernel_size[0]
+        stride = self.convtr.convtr.stride[0]
+        padding_total = kernel_size - stride
+
+        y = self.convtr(x)
+
+        # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be
+        # removed at the very end, when keeping only the right length for the output,
+        # as removing it here would require also passing the length at the matching layer
+        # in the encoder.
+        if self.causal:
+            # Trim the padding on the right according to the specified ratio
+            # if trim_right_ratio = 1.0, trim everything from right
+            padding_right = math.ceil(padding_total * self.trim_right_ratio)
+            padding_left = padding_total - padding_right
+            y = unpad1d(y, (padding_left, padding_right))
+        else:
+            # Asymmetric padding required for odd strides
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            y = unpad1d(y, (padding_left, padding_right))
+        #print("outputSConvTranspose1d:", y.shape, torch.sum(y), torch.sum(torch.abs(y)))
+        return y

models/modules/decoder.py

@@ -0,0 +1,155 @@
+from abc import ABC, abstractmethod
+from torch import nn, Tensor
+import torch
+import x_transformers as xt
+from typing import Optional
+
+import config as cfg
+
+
+class ResidualTokenEmbedding(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 num_tokens: int,
+                 n_layers: int,
+                 padding_token: Optional[int] = None):
+        super().__init__()
+        self.padding_token: Optional[int] = padding_token
+        self.emb = nn.ModuleList([
+            nn.Embedding(num_tokens, dim, padding_idx=self.padding_token)
+            for _ in range(n_layers)
+        ])
+
+        for layer in self.emb:
+            nn.init.kaiming_normal_(layer.weight)
+
+    def forward(self, x):
+        n_res_layers = x.shape[1]
+        assert n_res_layers == len(self.emb)
+        token_emb: Tensor = sum(  # type: ignore
+            [self.emb[i](x[:, i]) for i in range(n_res_layers)])
+        return token_emb
+
+
+def create_sin_embedding(
+    positions: torch.Tensor,
+    dim: int,
+    max_period: float = 10000,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """Create sinusoidal positional embedding, with shape `[B, T, C]`.
+
+    Args:
+        positions (torch.Tensor): LongTensor of positions.
+        dim (int): Dimension of the embedding.
+        max_period (float): Maximum period of the cosine/sine functions.
+        dtype (torch.dtype or str): dtype to use to generate the embedding.
+    Returns:
+        torch.Tensor: Sinusoidal positional embedding.
+    """
+    # We aim for BTC format
+    assert dim % 2 == 0
+    half_dim = dim // 2
+    positions = positions.to(dtype)
+    adim = torch.arange(half_dim, device=positions.device,
+                        dtype=dtype).view(1, 1, -1)
+    max_period_tensor = torch.full([],
+                                   max_period,
+                                   device=positions.device,
+                                   dtype=dtype)  # avoid sync point
+    phase = positions / (max_period_tensor**(adim / (half_dim - 1)))
+    return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1)
+
+
+class ResidualSinusoidalEmbedding(nn.Module):
+
+    def __init__(self, dim, theta=10000):
+        super().__init__()
+        assert dim % 2 == 0
+        self.scale = 1
+        self.theta = theta
+        self.dim = dim
+
+    def forward(self, x, pos=None, seq_start_pos=None):
+        B, K, T = x.shape
+        positions = (pos if pos is not None else torch.arange(
+            T, device=x.device).view(1, -1, 1))
+        pos_emb = create_sin_embedding(positions,
+                                       self.dim,
+                                       max_period=self.theta,
+                                       dtype=x.dtype)
+        return pos_emb * self.scale
+
+
+class ResidualOutputProj(nn.Module):
+
+    def __init__(self, input_dim: int, output_dim: int, n_layers: int,
+                 use_bias: bool):
+        super().__init__()
+        self.linears = nn.ModuleList([
+            nn.Linear(input_dim, output_dim, use_bias) for _ in range(n_layers)
+        ])
+
+    def forward(self, x):
+        return torch.stack([layer(x) for layer in self.linears], dim=1)
+
+
+class TransformerDecoder(ABC, nn.Module):
+
+    @abstractmethod
+    def forward(x, mask, context, context_mask, prepend_data, prepend_mask,
+                sum_data) -> Tensor:
+        ...
+
+
+class XTransformerDecoder(TransformerDecoder):
+
+    def __init__(
+        self,
+        num_tokens: int,
+        max_seq_len: int,
+        use_abs_pos_emb: bool,
+        scaled_sinu_pos_emb: bool,
+        dim: int,
+        depth: int,
+        heads: int,
+        attn_dim_head: int,
+        attn_flash: bool,
+        ff_no_bias: bool,
+        cross_attend: bool,
+    ):
+
+        self.decoder: xt.TransformerWrapper = xt.TransformerWrapper(
+            num_tokens=self.input_card,
+            max_seq_len=500,
+            use_abs_pos_emb=True,
+            scaled_sinu_pos_emb=True,
+            attn_layers=xt.Decoder(
+                dim=self.dim,
+                depth=self.n_layers,
+                heads=self.n_heads,
+                attn_dim_head=64,
+                attn_flash=True,
+                ff_no_bias=True,
+                cross_attend=self.cross_attend,
+            ),
+        )
+
+        self.decoder.token_emb = ResidualTokenEmbedding(  # type: ignore
+            self.dim,
+            self.input_card,
+            self.n_q,
+            self.padding_token,
+        )
+        self.decoder.pos_emb = ResidualSinusoidalEmbedding(  # type: ignore
+            dim=self.dim)
+        self.decoder.to_logits = ResidualOutputProj(self.dim, self.card,
+                                                    self.n_q, False)
+
+        # TODO: this is horrendous, gotta find a fix
+        if not self.cross_attend:
+            self.nullwav_embeds = torch.load(cfg.weights_dir() /
+                                             "nullwav_embeds.pt")[0]
+        else:
+            self.nullwav_embeds = None

models/modules/lstm.py

@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""LSTM layers module."""
+
+from torch import nn
+
+
+class SLSTM(nn.Module):
+    """
+    LSTM without worrying about the hidden state, nor the layout of the data.
+    Expects input as convolutional layout.
+    """
+    def __init__(self, dimension: int, num_layers: int = 2, skip: bool = True):
+        super().__init__()
+        self.skip = skip
+        self.lstm = nn.LSTM(dimension, dimension, num_layers)
+
+    def forward(self, x):
+        x = x.permute(2, 0, 1)
+        y, _ = self.lstm(x)
+        if self.skip:
+            y = y + x
+        y = y.permute(1, 2, 0)
+        return y

models/modules/norm.py

@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Normalization modules."""
+
+import typing as tp
+
+import einops
+import torch
+from torch import nn
+
+
+class ConvLayerNorm(nn.LayerNorm):
+    """
+    Convolution-friendly LayerNorm that moves channels to last dimensions
+    before running the normalization and moves them back to original position right after.
+    """
+    def __init__(self, normalized_shape: tp.Union[int, tp.List[int], torch.Size], **kwargs):
+        super().__init__(normalized_shape, **kwargs)
+
+    def forward(self, x):
+        x = einops.rearrange(x, 'b ... t -> b t ...')
+        x = super().forward(x)
+        x = einops.rearrange(x, 'b t ... -> b ... t')
+        return

models/modules/rope.py

@@ -0,0 +1,125 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import typing as tp
+
+from torch import nn
+import torch
+
+
+class XPos(nn.Module):
+    """Length-extrapolatable positional embedding (xPos) from [Sun et al 2022](https://arxiv.org/abs/2212.10554v1).
+    This applies an exponential decay to the RoPE rotation matrix.
+
+    Args:
+        dim (int): Embedding dimension.
+        smoothing (float): Smoothing factor applied to the decay rates.
+        base_scale (int): Base decay rate, given in terms of scaling time.
+        device (torch.device, optional): Device on which to initialize the module.
+        dtype (torch.dtype): dtype to use to generate the embedding.
+    """
+    def __init__(self, dim: int, smoothing: float = 0.4, base_scale: int = 512,
+                 device=None, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        assert dim % 2 == 0
+        assert dtype in [torch.float64, torch.float32]
+        self.dtype = dtype
+        self.base_scale = base_scale
+
+        half_dim = dim // 2
+        adim = torch.arange(half_dim, device=device, dtype=dtype)
+        decay_rates = (adim / half_dim + smoothing) / (1.0 + smoothing)
+        self.register_buffer("decay_rates", decay_rates)
+        self.decay: tp.Optional[torch.Tensor] = None
+
+    def get_decay(self, start: int, end: int):
+        """Create complex decay tensor, cache values for fast computation."""
+        if self.decay is None or end > self.decay.shape[0]:
+            assert isinstance(self.decay_rates, torch.Tensor)  # Satisfy type checker.
+            idx = torch.arange(end, device=self.decay_rates.device, dtype=self.dtype)
+            power = idx / self.base_scale
+            scale = self.decay_rates ** power.unsqueeze(-1)
+            self.decay = torch.polar(scale, torch.zeros_like(scale))
+        return self.decay[start:end]  # [T, C/2]
+
+
+class RotaryEmbedding(nn.Module):
+    """Rotary positional embedding (RoPE) from [Su et al 2022](https://arxiv.org/abs/2104.09864).
+
+    Args:
+        dim (int): Embedding dimension (twice the number of frequencies).
+        max_period (float): Maximum period of the rotation frequencies.
+        xpos (bool): Use xPos, applies an exponential decay to rotation matrix.
+        scale (float): Scale of positional embedding, set to 0 to deactivate.
+        device (torch.device, optional): Device on which to initialize the module.
+        dtype (torch.dtype): dtype to use to generate the embedding.
+    """
+    def __init__(self, dim: int, max_period: float = 10000.0, xpos: bool = False,
+                 scale: float = 1.0, device=None, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        assert dim % 2 == 0
+        self.scale = scale
+        assert dtype in [torch.float64, torch.float32]
+        self.dtype = dtype
+
+        adim = torch.arange(0, dim, 2, device=device, dtype=dtype)[: (dim // 2)]
+        frequencies = 1.0 / (max_period ** (adim / dim))
+        self.register_buffer("frequencies", frequencies)
+        self.rotation: tp.Optional[torch.Tensor] = None
+
+        self.xpos = XPos(dim, device=device, dtype=dtype) if xpos else None
+
+    def get_rotation(self, start: int, end: int):
+        """Create complex rotation tensor, cache values for fast computation."""
+        if self.rotation is None or end > self.rotation.shape[0]:
+            assert isinstance(self.frequencies, torch.Tensor)  # Satisfy type checker.
+            idx = torch.arange(end, device=self.frequencies.device, dtype=self.dtype)
+            angles = torch.outer(idx, self.frequencies)
+            self.rotation = torch.polar(torch.ones_like(angles), angles)
+        return self.rotation[start:end]
+
+    def rotate(self, x: torch.Tensor, start: int = 0, time_dim: int = 1, invert_decay: bool = False):
+        """Apply rope rotation to query or key tensor."""
+        T = x.shape[time_dim]
+        target_shape = [1] * x.dim()
+        target_shape[time_dim] = T
+        target_shape[-1] = -1
+        rotation = self.get_rotation(start, start + T).view(target_shape)
+
+        if self.xpos:
+            decay = self.xpos.get_decay(start, start + T).view(target_shape)
+        else:
+            decay = 1.0
+
+        if invert_decay:
+            decay = decay ** -1
+
+        x_complex = torch.view_as_complex(x.to(self.dtype).reshape(*x.shape[:-1], -1, 2))
+        scaled_rotation = (rotation * decay) * self.scale + (1.0 - self.scale)
+        x_out = torch.view_as_real(x_complex * scaled_rotation).view_as(x)
+
+        return x_out.type_as(x)
+
+    def rotate_qk(self, query: torch.Tensor, key: torch.Tensor, start: int = 0, time_dim: int = 1):
+        """ Apply rope rotation to both query and key tensors.
+        Supports streaming mode, in which query and key are not expected to have the same shape.
+        In streaming mode, key will be of length [P + C] with P the cached past timesteps, but
+        query will be [C] (typically C == 1).
+
+        Args:
+            query (torch.Tensor): Query to rotate.
+            key (torch.Tensor): Key to rotate.
+            start (int): Start index of the sequence for time offset.
+            time_dim (int): which dimension represent the time steps.
+        """
+        query_timesteps = query.shape[time_dim]
+        key_timesteps = key.shape[time_dim]
+        streaming_offset = key_timesteps - query_timesteps
+
+        query_out = self.rotate(query, start + streaming_offset, time_dim)
+        key_out = self.rotate(key, start, time_dim, invert_decay=True)
+
+        return query_out, key_out

models/modules/seanet.py

@@ -0,0 +1,257 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Encodec SEANet-based encoder and decoder implementation."""
+
+import typing as tp
+
+import numpy as np
+import torch.nn as nn
+
+from . import (
+    SConv1d,
+    SConvTranspose1d,
+    SLSTM
+)
+
+
+class SEANetResnetBlock(nn.Module):
+    """Residual block from SEANet model.
+    Args:
+        dim (int): Dimension of the input/output
+        kernel_sizes (list): List of kernel sizes for the convolutions.
+        dilations (list): List of dilations for the convolutions.
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3)
+        true_skip (bool): Whether to use true skip connection or a simple convolution as the skip connection.
+    """
+    def __init__(self, dim: int, kernel_sizes: tp.List[int] = [3, 1], dilations: tp.List[int] = [1, 1],
+                 activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 norm: str = 'weight_norm', norm_params: tp.Dict[str, tp.Any] = {}, causal: bool = False,
+                 pad_mode: str = 'reflect', compress: int = 2, true_skip: bool = True):
+        super().__init__()
+        assert len(kernel_sizes) == len(dilations), 'Number of kernel sizes should match number of dilations'
+        act = getattr(nn, activation)
+        hidden = dim // compress
+        block = []
+        for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)): # this is always length 2
+            in_chs = dim if i == 0 else hidden
+            out_chs = dim if i == len(kernel_sizes) - 1 else hidden
+            # print(in_chs, "_", out_chs) # 32 _ 16; 16 _ 32; 64 _ 32; 32 _ 64; etc until 256 _ 128; 128_ 256 for encode
+            block += [
+                act(**activation_params),
+                SConv1d(in_chs, out_chs, kernel_size=kernel_size, dilation=dilation,
+                        norm=norm, norm_kwargs=norm_params,
+                        causal=causal, pad_mode=pad_mode),
+            ]
+        self.block = nn.Sequential(*block)
+        self.shortcut: nn.Module
+        # true_skip is always false since the default in SEANetEncoder / SEANetDecoder does not get changed
+        if true_skip:
+            self.shortcut = nn.Identity()
+        else:
+            self.shortcut = SConv1d(dim, dim, kernel_size=1, norm=norm, norm_kwargs=norm_params,
+                                    causal=causal, pad_mode=pad_mode)
+
+    def forward(self, x):
+        return self.shortcut(x) + self.block(x) # This is simply the sum of two tensors of the same size
+
+# Only channels, norm, causal are different between 24HZ & 48HZ, everything else is default parameter
+# 24HZ -> channels = 1, norm = weight_norm, causal = True
+# 48HZ -> channels = 2, norm = time_group_norm, causal = False
+class SEANetEncoder(nn.Module):
+    """SEANet encoder.
+    Args:
+        channels (int): Audio channels.
+        dimension (int): Intermediate representation dimension.
+        n_filters (int): Base width for the model.
+        n_residual_layers (int): nb of residual layers.
+        ratios (Sequence[int]): kernel size and stride ratios. The encoder uses downsampling ratios instead of
+            upsampling ratios, hence it will use the ratios in the reverse order to the ones specified here
+            that must match the decoder order
+        activation (str): Activation function. ELU = Exponential Linear Unit
+        activation_params (dict): Parameters to provide to the activation function
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        kernel_size (int): Kernel size for the initial convolution.
+        last_kernel_size (int): Kernel size for the initial convolution.
+        residual_kernel_size (int): Kernel size for the residual layers.
+        dilation_base (int): How much to increase the dilation with each layer.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        true_skip (bool): Whether to use true skip connection or a simple
+            (streamable) convolution as the skip connection in the residual network blocks.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        lstm (int): Number of LSTM layers at the end of the encoder.
+    """
+    def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 1,
+                 ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 norm: str = 'weight_norm', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,
+                 last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,
+                 pad_mode: str = 'reflect', true_skip: bool = False, compress: int = 2, lstm: int = 2):
+        super().__init__()
+        self.channels = channels
+        self.dimension = dimension
+        self.n_filters = n_filters
+        self.ratios = list(reversed(ratios))
+        del ratios
+        self.n_residual_layers = n_residual_layers
+        self.hop_length = np.prod(self.ratios)
+
+        act = getattr(nn, activation)
+        mult = 1
+        model: tp.List[nn.Module] = [
+            SConv1d(channels, mult * n_filters, kernel_size, norm=norm, norm_kwargs=norm_params,
+                    causal=causal, pad_mode=pad_mode)
+        ]
+        # Downsample to raw audio scale
+        for ratio in self.ratios: # CHANGED from: for i, ratio in enumerate(self.ratios):
+            # Add residual layers
+            for j in range(n_residual_layers): # This is always 1, parameter never gets changed from default anywhere
+                model += [
+                    SEANetResnetBlock(mult * n_filters, kernel_sizes=[residual_kernel_size, 1],
+                                      dilations=[dilation_base ** j, 1],
+                                      norm=norm, norm_params=norm_params,
+                                      activation=activation, activation_params=activation_params,
+                                      causal=causal, pad_mode=pad_mode, compress=compress, true_skip=true_skip)]
+
+            # Add downsampling layers
+            model += [
+                act(**activation_params),
+                SConv1d(mult * n_filters, mult * n_filters * 2,
+                        kernel_size=ratio * 2, stride=ratio,
+                        norm=norm, norm_kwargs=norm_params,
+                        causal=causal, pad_mode=pad_mode),
+            ]
+            mult *= 2
+
+        if lstm:
+            model += [SLSTM(mult * n_filters, num_layers=lstm)]
+
+        model += [
+            act(**activation_params),
+            SConv1d(mult * n_filters, dimension, last_kernel_size, norm=norm, norm_kwargs=norm_params,
+                    causal=causal, pad_mode=pad_mode)
+        ]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class SEANetDecoder(nn.Module):
+    """SEANet decoder.
+    Args:
+        channels (int): Audio channels.
+        dimension (int): Intermediate representation dimension.
+        n_filters (int): Base width for the model.
+        n_residual_layers (int): nb of residual layers.
+        ratios (Sequence[int]): kernel size and stride ratios
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function
+        final_activation (str): Final activation function after all convolutions.
+        final_activation_params (dict): Parameters to provide to the activation function
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        kernel_size (int): Kernel size for the initial convolution.
+        last_kernel_size (int): Kernel size for the initial convolution.
+        residual_kernel_size (int): Kernel size for the residual layers.
+        dilation_base (int): How much to increase the dilation with each layer.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        true_skip (bool): Whether to use true skip connection or a simple
+            (streamable) convolution as the skip connection in the residual network blocks.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        lstm (int): Number of LSTM layers at the end of the encoder.
+        trim_right_ratio (float): Ratio for trimming at the right of the transposed convolution under the causal setup.
+            If equal to 1.0, it means that all the trimming is done at the right.
+    """
+    def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 1,
+                 ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 final_activation: tp.Optional[str] = None, final_activation_params: tp.Optional[dict] = None,
+                 norm: str = 'weight_norm', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,
+                 last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,
+                 pad_mode: str = 'reflect', true_skip: bool = False, compress: int = 2, lstm: int = 2,
+                 trim_right_ratio: float = 1.0):
+        super().__init__()
+        self.dimension = dimension
+        self.channels = channels
+        self.n_filters = n_filters
+        self.ratios = ratios
+        del ratios
+        self.n_residual_layers = n_residual_layers
+        self.hop_length = np.prod(self.ratios)
+
+        act = getattr(nn, activation)
+        mult = int(2 ** len(self.ratios))
+        model: tp.List[nn.Module] = [
+            SConv1d(dimension, mult * n_filters, kernel_size, norm=norm, norm_kwargs=norm_params,
+                    causal=causal, pad_mode=pad_mode)
+        ]
+
+        if lstm:
+            model += [SLSTM(mult * n_filters, num_layers=lstm)]
+
+        # Upsample to raw audio scale
+        for i, ratio in enumerate(self.ratios):
+            # Add upsampling layers
+            model += [
+                act(**activation_params),
+                SConvTranspose1d(mult * n_filters, mult * n_filters // 2,
+                                 kernel_size=ratio * 2, stride=ratio,
+                                 norm=norm, norm_kwargs=norm_params,
+                                 causal=causal, trim_right_ratio=trim_right_ratio),
+            ]
+            # Add residual layers
+            for j in range(n_residual_layers):
+                model += [
+                    SEANetResnetBlock(mult * n_filters // 2, kernel_sizes=[residual_kernel_size, 1],
+                                      dilations=[dilation_base ** j, 1],
+                                      activation=activation, activation_params=activation_params,
+                                      norm=norm, norm_params=norm_params, causal=causal,
+                                      pad_mode=pad_mode, compress=compress, true_skip=true_skip)]
+
+            mult //= 2
+
+        # Add final layers
+        model += [
+            act(**activation_params),
+            SConv1d(n_filters, channels, last_kernel_size, norm=norm, norm_kwargs=norm_params,
+                    causal=causal, pad_mode=pad_mode)
+        ]
+        # Add optional final activation to decoder (eg. tanh)
+        if final_activation is not None: # This is always None
+            final_act = getattr(nn, final_activation)
+            final_activation_params = final_activation_params or {}
+            model += [
+                final_act(**final_activation_params)
+            ]
+        self.model = nn.Sequential(*model)
+
+    def forward(self, z):
+        y = self.model(z)
+        return y
+
+
+def test():
+    import torch
+    encoder = SEANetEncoder()
+    decoder = SEANetDecoder()
+    x = torch.randn(1, 1, 24000)
+    z = encoder(x)
+    assert list(z.shape) == [1, 128, 75], z.shape
+    y = decoder(z)
+    assert y.shape == x.shape, (x.shape, y.shape)
+
+
+if __name__ == '__main__':
+    test()

models/modules/streaming.py

@@ -0,0 +1,131 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Streaming module API that should be implemented by all Streaming components,
+"""
+
+from contextlib import contextmanager
+import typing as tp
+from torch import nn
+import torch
+
+
+State = tp.Dict[str, torch.Tensor]
+
+
+class StreamingModule(nn.Module):
+    """Common API for streaming components.
+
+    Each streaming component has a streaming state, which is just a dict[str, Tensor].
+    By convention, the first dim of each tensor must be the batch size.
+    Don't use dots in the key names, as this would clash with submodules
+    (like in state_dict).
+
+    If `self._is_streaming` is True, the component should use and remember
+    the proper state inside `self._streaming_state`.
+
+    To set a streaming component in streaming state, use
+
+        with module.streaming():
+            ...
+
+    This will automatically reset the streaming state when exiting the context manager.
+    This also automatically propagates to all streaming children module.
+
+    Some module might also implement the `StreamingModule.flush` method, although
+    this one is trickier, as all parents module must be StreamingModule and implement
+    it as well for it to work properly. See `StreamingSequential` after.
+    """
+    def __init__(self) -> None:
+        super().__init__()
+        self._streaming_state: State = {}
+        self._is_streaming = False
+
+    def _apply_named_streaming(self, fn: tp.Any):
+        for name, module in self.named_modules():
+            if isinstance(module, StreamingModule):
+                fn(name, module)
+
+    def _set_streaming(self, streaming: bool):
+        def _set_streaming(name, module):
+            module._is_streaming = streaming
+        self._apply_named_streaming(_set_streaming)
+
+    @contextmanager
+    def streaming(self):
+        """Context manager to enter streaming mode. Reset streaming state on exit."""
+        self._set_streaming(True)
+        try:
+            yield
+        finally:
+            self._set_streaming(False)
+            self.reset_streaming()
+
+    def reset_streaming(self):
+        """Reset the streaming state."""
+        def _reset(name: str, module: StreamingModule):
+            module._streaming_state.clear()
+
+        self._apply_named_streaming(_reset)
+
+    def get_streaming_state(self) -> State:
+        """Return the streaming state, including that of sub-modules."""
+        state: State = {}
+
+        def _add(name: str, module: StreamingModule):
+            if name:
+                name += "."
+            for key, value in module._streaming_state.items():
+                state[name + key] = value
+
+        self._apply_named_streaming(_add)
+        return state
+
+    def set_streaming_state(self, state: State):
+        """Set the streaming state, including that of sub-modules."""
+        state = dict(state)
+
+        def _set(name: str, module: StreamingModule):
+            if name:
+                name += "."
+            module._streaming_state.clear()
+            for key, value in list(state.items()):
+                # complexity is not ideal here, but probably fine.
+                if key.startswith(name):
+                    local_key = key[len(name):]
+                    if '.' not in local_key:
+                        module._streaming_state[local_key] = value
+                        del state[key]
+
+        self._apply_named_streaming(_set)
+        assert len(state) == 0, list(state.keys())
+
+    def flush(self, x: tp.Optional[torch.Tensor] = None):
+        """Flush any remaining outputs that were waiting for completion.
+        Typically, for convolutions, this will add the final padding
+        and process the last buffer.
+
+        This should take an optional argument `x`, which will be provided
+        if a module before this one in the streaming pipeline has already
+        spitted out a flushed out buffer.
+        """
+        if x is None:
+            return None
+        else:
+            return self(x)
+
+
+class StreamingSequential(StreamingModule, nn.Sequential):
+    """A streaming compatible alternative of `nn.Sequential`.
+    """
+    def flush(self, x: tp.Optional[torch.Tensor] = None):
+        for module in self:
+            if isinstance(module, StreamingModule):
+                x = module.flush(x)
+            elif x is not None:
+                x = module(x)
+        return x

models/modules/transformer.py

@@ -0,0 +1,755 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Transformer model, with streaming support, xformer attention support
+and easy causal attention with a potentially finite receptive field.
+
+See `StreamingTransformer` for more information.
+
+Unlike regular PyTorch Transformer, we make the hard choice that batches are first.
+"""
+
+import typing as tp
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint as torch_checkpoint
+from xformers import ops
+
+from .rope import RotaryEmbedding
+from .streaming import StreamingModule
+
+_efficient_attention_backend: str = 'torch'
+
+
+def set_efficient_attention_backend(backend: str = 'torch'):
+    # Using torch by default, it seems a bit faster on older P100 GPUs (~20% faster).
+    global _efficient_attention_backend
+    assert _efficient_attention_backend in ['xformers', 'torch']
+    _efficient_attention_backend = backend
+
+
+def _get_attention_time_dimension(memory_efficient: bool) -> int:
+    if _efficient_attention_backend == 'torch' and memory_efficient:
+        return 2
+    else:
+        return 1
+
+
+def _is_profiled() -> bool:
+    # Return true if we are currently running with a xformers profiler activated.
+    try:
+        from xformers.profiler import profiler
+    except ImportError:
+        return False
+    return profiler._Profiler._CURRENT_PROFILER is not None
+
+
+def create_norm_fn(norm_type: str, dim: int, **kwargs) -> nn.Module:
+    """Create normalization module for transformer encoder layer.
+
+    Args:
+        norm_type (str): Normalization method.
+        dim (int): Dimension of the normalized layer.
+        **kwargs (dict): Additional parameters for normalization layer.
+    Returns:
+        nn.Module: Normalization module.
+    """
+    if norm_type == 'layer_norm':
+        return nn.LayerNorm(dim, eps=1e-5, **kwargs)
+    else:
+        raise ValueError(f"Unknown norm type: {norm_type}")
+
+
+def create_sin_embedding(positions: torch.Tensor, dim: int, max_period: float = 10000,
+                         dtype: torch.dtype = torch.float32) -> torch.Tensor:
+    """Create sinusoidal positional embedding, with shape `[B, T, C]`.
+
+    Args:
+        positions (torch.Tensor): LongTensor of positions.
+        dim (int): Dimension of the embedding.
+        max_period (float): Maximum period of the cosine/sine functions.
+        dtype (torch.dtype or str): dtype to use to generate the embedding.
+    Returns:
+        torch.Tensor: Sinusoidal positional embedding.
+    """
+    # We aim for BTC format
+    assert dim % 2 == 0
+    half_dim = dim // 2
+    positions = positions.to(dtype)
+    adim = torch.arange(half_dim, device=positions.device, dtype=dtype).view(1, 1, -1)
+    max_period_tensor = torch.full([], max_period, device=positions.device, dtype=dtype)  # avoid sync point
+    phase = positions / (max_period_tensor ** (adim / (half_dim - 1)))
+    return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1)
+
+
+def expand_repeated_kv(x: torch.Tensor, n_rep: int, memory_efficient: bool) -> torch.Tensor:
+    """torch.repeat_interleave(x, dim=2, repeats=n_rep) from xlformers."""
+    if n_rep == 1:
+        return x
+    if _efficient_attention_backend == 'torch' and memory_efficient:
+        bs, n_kv_heads, slen, head_dim = x.shape
+        return (
+            x[:, :, None, :, :]
+            .expand(bs, n_kv_heads, n_rep, slen, head_dim)
+            .reshape(bs, n_kv_heads * n_rep, slen, head_dim)
+        )
+    else:
+        bs, slen, n_kv_heads, head_dim = x.shape
+        return (
+            x[:, :, :, None, :]
+            .expand(bs, slen, n_kv_heads, n_rep, head_dim)
+            .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
+        )
+
+
+class LayerScale(nn.Module):
+    """Layer scale from [Touvron et al 2021] (https://arxiv.org/pdf/2103.17239.pdf).
+    This rescales diagonally the residual outputs close to 0, with a learnt scale.
+
+    Args:
+        channels (int): Number of channels.
+        init (float): Initial scale.
+        channel_last (bool): If True, expect `[*, C]` shaped tensors, otherwise, `[*, C, T]`.
+        device (torch.device or str, optional): Device on which to initialize the module.
+        dtype (torch.dtype, optional): dtype to use to initialize the module.
+    """
+    def __init__(self, channels: int, init: float = 1e-4, channel_last: bool = True,
+                 device=None, dtype=None):
+        super().__init__()
+        self.channel_last = channel_last
+        self.scale = nn.Parameter(
+            torch.full((channels,), init,
+                       requires_grad=True, device=device, dtype=dtype))
+
+    def forward(self, x: torch.Tensor):
+        if self.channel_last:
+            return self.scale * x
+        else:
+            return self.scale[:, None] * x
+
+
+class StreamingMultiheadAttention(StreamingModule):
+    """Similar to `nn.MultiheadAttention` but with support for streaming, causal evaluation.
+
+    Args:
+        embed_dim (int): Dimension to project to.
+        num_heads (int): Number of heads.
+        dropout (float): Dropout level.
+        bias (bool): Use bias in projections.
+        causal (bool): Causal mask applied automatically.
+        past_context (int, optional): Receptive field for the causal mask, infinite if None.
+        custom (bool): Use custom MHA implementation, for testing / benchmarking.
+        memory_efficient (bool): Use xformers based memory efficient attention.
+        attention_as_float32 (bool): Perform the attention as float32
+            (especially important with memory_efficient as autocast won't do this automatically).
+        rope (`RotaryEmbedding`, optional): Rope embedding to use.
+        cross_attention: Should be true when used as a cross attention.
+            All keys and values must be available at once, streaming is only for the queries.
+            Cannot be used with `causal` or `rope` (as it wouldn't make sens to
+            interpret the time steps in the keys relative to those in the queries).
+        safe_streaming (bool): Bug fix, will go away with xformers update.
+        qk_layer_norm (bool): Layer normalization applied to queries and keys before dot product.
+        kv_repeat (int): If > 1, will repeat keys and queries multiple times (need to divide num_heads).
+            This will lead to faster decoding time on A100 or other GPUs with tensorcore.
+        device (torch.device, optional): Device on which to initialize.
+        dtype (torch.dtype, optional): dtype to use.
+    """
+    def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0, bias: bool = True,
+                 causal: bool = False, past_context: tp.Optional[int] = None, custom: bool = False,
+                 memory_efficient: bool = False, attention_as_float32: bool = False,
+                 rope: tp.Optional[RotaryEmbedding] = None, cross_attention: bool = False,
+                 safe_streaming: bool = True, qk_layer_norm: bool = False, kv_repeat: int = 1,
+                 device=None, dtype=None):
+        super().__init__()
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        if past_context is not None:
+            assert causal
+
+        self.embed_dim = embed_dim
+        self.causal = causal
+        self.past_context = past_context
+        self.memory_efficient = memory_efficient
+        self.attention_as_float32 = attention_as_float32
+        self.rope = rope
+        self.cross_attention = cross_attention
+        self.safe_streaming = safe_streaming
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.kv_repeat = kv_repeat
+        if cross_attention:
+            assert not causal, "Causal cannot work with cross attention."
+            assert rope is None, "Rope cannot work with cross attention."
+
+        if memory_efficient:
+            _verify_xformers_memory_efficient_compat()
+
+        self.custom = _is_custom(custom, memory_efficient)
+        if self.custom:
+            out_dim = embed_dim
+            assert num_heads % kv_repeat == 0
+            assert not cross_attention or kv_repeat == 1
+            num_kv = num_heads // kv_repeat
+            kv_dim = (embed_dim // num_heads) * num_kv
+            out_dim += 2 * kv_dim
+            in_proj = nn.Linear(embed_dim, out_dim, bias=bias, **factory_kwargs)
+            # We try to follow the default PyTorch MHA convention, to easily compare results.
+            self.in_proj_weight = in_proj.weight
+            self.in_proj_bias = in_proj.bias
+            if bias:
+                self.in_proj_bias.data.zero_()  # Following Pytorch convention
+            self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias, **factory_kwargs)
+            if bias:
+                self.out_proj.bias.data.zero_()
+        else:
+            assert not qk_layer_norm
+            assert kv_repeat == 1
+            self.mha = nn.MultiheadAttention(
+                embed_dim, num_heads, dropout=dropout, bias=bias, batch_first=True,
+                **factory_kwargs)
+        self.qk_layer_norm = qk_layer_norm
+        if qk_layer_norm:
+            assert self.custom
+            assert kv_repeat == 1
+            ln_dim = embed_dim
+            self.q_layer_norm = nn.LayerNorm(ln_dim)
+            self.k_layer_norm = nn.LayerNorm(ln_dim)
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        if not self.custom:
+            # Support compat with regular MHA
+            keys = [n for n, _ in self.mha.named_parameters()]
+            for key in keys:
+                if prefix + key in state_dict:
+                    state_dict[prefix + "mha." + key] = state_dict.pop(prefix + key)
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
+    def _get_mask(self, current_steps: int, device: torch.device, dtype: torch.dtype):
+        # Return a causal mask, accounting for potentially stored past keys/values
+        # We actually return a bias for the attention score, as this has the same
+        # convention both in the builtin MHA in Pytorch, and Xformers functions.
+        time_dim = _get_attention_time_dimension(self.memory_efficient)
+        if self.memory_efficient:
+            from xformers.ops import LowerTriangularMask
+            if current_steps == 1:
+                # If we only have one step, then we do not need a mask.
+                return None
+            elif 'past_keys' in self._streaming_state:
+                raise RuntimeError("Not supported at the moment")
+            else:
+                # Then we can safely use a lower triangular mask
+                return LowerTriangularMask()
+        if self._streaming_state:
+            past_keys = self._streaming_state['past_keys']
+            past_steps = past_keys.shape[time_dim]
+        else:
+            past_steps = 0
+
+        queries_pos = torch.arange(
+            past_steps, current_steps + past_steps, device=device).view(-1, 1)
+        keys_pos = torch.arange(past_steps + current_steps, device=device).view(1, -1)
+        delta = queries_pos - keys_pos
+        valid = delta >= 0
+        if self.past_context is not None:
+            valid &= (delta <= self.past_context)
+        return torch.where(
+            valid,
+            torch.zeros([], device=device, dtype=dtype),
+            torch.full([], float('-inf'), device=device, dtype=dtype))
+
+    def _complete_kv(self, k, v):
+        time_dim = _get_attention_time_dimension(self.memory_efficient)
+        if self.cross_attention:
+            # With cross attention we assume all keys and values
+            # are already available, and streaming is with respect
+            # to the queries only.
+            return k, v
+        # Complete the key/value pair using the streaming state.
+        if self._streaming_state:
+            pk = self._streaming_state['past_keys']
+            nk = torch.cat([pk, k], dim=time_dim)
+            if v is k:
+                nv = nk
+            else:
+                pv = self._streaming_state['past_values']
+                nv = torch.cat([pv, v], dim=time_dim)
+        else:
+            nk = k
+            nv = v
+
+        assert nk.shape[time_dim] == nv.shape[time_dim]
+        offset = 0
+        if self.past_context is not None:
+            offset = max(0, nk.shape[time_dim] - self.past_context)
+        if self._is_streaming:
+            self._streaming_state['past_keys'] = nk[:, offset:]
+            if v is not k:
+                self._streaming_state['past_values'] = nv[:, offset:]
+            if 'offset' in self._streaming_state:
+                self._streaming_state['offset'] += offset
+            else:
+                self._streaming_state['offset'] = torch.tensor(0)
+        return nk, nv
+
+    def _apply_rope(self, query: torch.Tensor, key: torch.Tensor):
+        time_dim = _get_attention_time_dimension(self.memory_efficient)
+        # Apply rope embeddings to query and key tensors.
+        assert self.rope is not None
+        if 'past_keys' in self._streaming_state:
+            past_keys_offset = self._streaming_state['past_keys'].shape[1]
+        else:
+            past_keys_offset = 0
+        if 'offset' in self._streaming_state:
+            past_context_offset = int(self._streaming_state['offset'].item())
+        else:
+            past_context_offset = 0
+        streaming_offset = past_context_offset + past_keys_offset
+        return self.rope.rotate_qk(query, key, start=streaming_offset, time_dim=time_dim)
+
+    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
+                key_padding_mask=None, need_weights=False, attn_mask=None,
+                average_attn_weights=True, is_causal=False):
+        assert not is_causal, ("New param added in torch 2.0.1 not supported, "
+                               "use the causal args in the constructor.")
+
+        time_dim = _get_attention_time_dimension(self.memory_efficient)
+        if time_dim == 2:
+            layout = "b h t d"
+        else:
+            layout = "b t h d"
+        dtype = query.dtype
+        if self._is_streaming:
+            assert self.causal or self.cross_attention, \
+                "Streaming only available for causal or cross attention"
+
+        custom_attn_mask = attn_mask is not None
+
+        if self.causal:
+            assert attn_mask is None
+            # At the moment we specialize only for the self-attention case.
+            assert query.shape[1] == key.shape[1], "Causal only for same length query / key / value"
+            assert value.shape[1] == key.shape[1], "Causal only for same length query / key / value"
+            attn_mask = self._get_mask(query.shape[1], query.device, query.dtype)
+
+        if self.custom:
+            # custom implementation
+            assert need_weights is False
+            assert key_padding_mask is None
+            if self.cross_attention:
+                # Different queries, keys, values, we have to spit manually the weights
+                # before applying the linear.
+                dim = self.in_proj_weight.shape[0] // 3
+                if self.in_proj_bias is None:
+                    bias_q, bias_k, bias_v = None, None, None
+                else:
+                    bias_q = self.in_proj_bias[:dim]
+                    bias_k = self.in_proj_bias[dim: 2 * dim]
+                    bias_v = self.in_proj_bias[2 * dim:]
+                q = nn.functional.linear(query, self.in_proj_weight[:dim], bias_q)
+                # todo: when streaming, we could actually save k, v and check the shape actually match.
+                k = nn.functional.linear(key, self.in_proj_weight[dim: 2 * dim], bias_k)
+                v = nn.functional.linear(value, self.in_proj_weight[2 * dim:], bias_v)
+                if self.qk_layer_norm is True:
+                    q = self.q_layer_norm(q)
+                    k = self.k_layer_norm(k)
+                q, k, v = [rearrange(x, f"b t (h d) -> {layout}", h=self.num_heads) for x in [q, k, v]]
+            else:
+                if not _is_profiled():
+                    # profiling breaks that propertysomehow.
+                    assert query is key, "specialized implementation"
+                    assert value is key, "specialized implementation"
+                projected = nn.functional.linear(query, self.in_proj_weight, self.in_proj_bias)
+                if self.kv_repeat == 1:
+                    if time_dim == 2:
+                        bound_layout = "b h p t d"
+                    else:
+                        bound_layout = "b t p h d"
+                    packed = rearrange(projected, f"b t (p h d) -> {bound_layout}", p=3, h=self.num_heads)
+                    q, k, v = ops.unbind(packed, dim=2)
+                else:
+                    embed_dim = self.embed_dim
+                    per_head_dim = (embed_dim // self.num_heads)
+                    kv_heads = self.num_heads // self.kv_repeat
+                    q = projected[:, :, :embed_dim]
+                    start = embed_dim
+                    end = start + per_head_dim * kv_heads
+                    k = projected[:, :, start: end]
+                    v = projected[:, :, end:]
+                    q = rearrange(q, f"b t (h d) -> {layout}", h=self.num_heads)
+                    k = rearrange(k, f"b t (h d) -> {layout}", h=kv_heads)
+                    v = rearrange(v, f"b t (h d) -> {layout}", h=kv_heads)
+
+                if self.qk_layer_norm is True:
+                    assert self.kv_repeat == 1
+                    q, k = [rearrange(x, f"{layout} -> b t (h d)") for x in [q, k]]
+                    q = self.q_layer_norm(q)
+                    k = self.k_layer_norm(k)
+                    q, k = [rearrange(x, f"b t (h d) -> {layout}", h=self.num_heads) for x in [q, k]]
+                if self.rope:
+                    q, k = self._apply_rope(q, k)
+                k, v = self._complete_kv(k, v)
+                if self.kv_repeat > 1:
+                    k = expand_repeated_kv(k, self.kv_repeat, self.memory_efficient)
+                    v = expand_repeated_kv(v, self.kv_repeat, self.memory_efficient)
+            if self.attention_as_float32:
+                q, k, v = [x.float() for x in [q, k, v]]
+            if self.memory_efficient:
+                if custom_attn_mask:
+                    # When using a custom attn mask:
+                    # Move to query's device, repeat for each sample, remove align8 padding
+                    seq_len = query.shape[1]
+                    attn_mask = attn_mask.to(q.dtype)
+                    attn_mask = attn_mask.repeat((q.shape[0], 1, 1, 1))
+                    attn_mask = attn_mask[..., :seq_len, :seq_len]
+
+                p = self.dropout if self.training else 0
+                if _efficient_attention_backend == 'torch':
+                    x = torch.nn.functional.scaled_dot_product_attention(
+                        q, k, v, is_causal=attn_mask is not None, dropout_p=p)
+                else:
+                    x = ops.memory_efficient_attention(q, k, v, attn_mask, p=p)
+            else:
+                # We include the dot product as float32, for consistency
+                # with the other implementations that include that step
+                # as part of the attention. Note that when using `autocast`,
+                # the einsums would be done as bfloat16, but the softmax
+                # would be done as bfloat16, so `attention_as_float32` will
+                # extend a bit the range of operations done in float32,
+                # although this should make no difference.
+                q = q / q.shape[-1] ** 0.5
+                key_layout = layout.replace('t', 'k')
+                query_layout = layout
+                if self._is_streaming and self.safe_streaming and q.device.type == 'cuda':
+                    with torch.autocast(device_type=q.device.type, dtype=torch.float32):
+                        pre_w = torch.einsum(f"{query_layout},{key_layout}-> b h t k", q, k)
+                else:
+                    pre_w = torch.einsum(f"{query_layout},{key_layout}-> b h t k", q, k)
+                if attn_mask is not None:
+                    pre_w = pre_w + attn_mask
+                w = torch.softmax(pre_w, dim=-1)
+                w = F.dropout(w, self.dropout, training=self.training).to(v)
+                # Key and value have the same format.
+                x = torch.einsum(f"b h t k, {key_layout} -> {layout}", w, v)
+            x = x.to(dtype)
+            x = rearrange(x, f"{layout} -> b t (h d)", h=self.num_heads)
+            x = self.out_proj(x)
+        else:
+            key, value = self._complete_kv(key, value)
+            if self.attention_as_float32:
+                query, key, value = [x.float() for x in [query, key, value]]
+            x, _ = self.mha(
+                query, key, value, key_padding_mask,
+                need_weights, attn_mask, average_attn_weights)
+            x = x.to(dtype)
+
+        return x, None
+
+
+class StreamingTransformerLayer(nn.TransformerEncoderLayer):
+    """TransformerLayer with Streaming / Causal support.
+    This also integrates cross_attention, when passing `cross_attention=True`,
+    rather than having two separate classes like in PyTorch.
+
+    Args:
+        d_model (int): Dimension of the data.
+        num_heads (int): Number of heads.
+        dim_feedforward (int): Intermediate dimension of FF module.
+        dropout (float): Dropout both for MHA and FF.
+        bias_ff (bool): Use bias for FF.
+        bias_attn (bool): Use bias for MHA.
+        causal (bool): Causal mask applied automatically.
+        past_context (int, optional): Receptive field for the causal mask, infinite if None.
+        custom (bool): Use custom MHA implementation, for testing / benchmarking.
+        memory_efficient (bool): Use xformers based memory efficient attention.
+        attention_as_float32 (bool): Perform the attention as float32
+            (especially important with memory_efficient as autocast won't do this automatically).
+        qk_layer_norm (bool): Layer normalization applied to queries and keys before dot product in attention.
+        qk_layer_norm_cross (bool): Same for the cross attention.
+        cross_attention (bool): If True, expect to get secondary input for cross-attention.
+            Cross attention will use the default MHA, as it typically won't require
+            special treatment.
+        layer_scale (float, optional): If not None, LayerScale will be used with
+            the given value as initial scale.
+        rope (`RotaryEmbedding`, optional): Rope embedding to use.
+        attention_dropout (float, optional): If not None, separate the value of the dimension dropout
+            in FFN and of the attention dropout.
+        kv_repeat (int): If > 1, will repeat keys and queries multiple times (need to divide num_heads).
+            This will lead to faster decoding time on A100 or other GPUs with tensorcore.
+        device (torch.device, optional): Device on which to initialize.
+        dtype (torch.dtype, optional): dtype to use.
+        **kwargs: See `nn.TransformerEncoderLayer`.
+    """
+    def __init__(self, d_model: int, num_heads: int, dim_feedforward: int = 2048, dropout: float = 0.1,
+                 bias_ff: bool = True, bias_attn: bool = True, causal: bool = False,
+                 past_context: tp.Optional[int] = None, custom: bool = False,
+                 memory_efficient: bool = False, attention_as_float32: bool = False,
+                 qk_layer_norm: bool = False, qk_layer_norm_cross: bool = False,
+                 cross_attention: bool = False, layer_scale: tp.Optional[float] = None,
+                 rope: tp.Optional[RotaryEmbedding] = None, attention_dropout: tp.Optional[float] = None,
+                 kv_repeat: int = 1, norm: str = 'layer_norm', device=None, dtype=None, **kwargs):
+        super().__init__(d_model, num_heads, dim_feedforward, dropout,
+                         device=device, dtype=dtype, batch_first=True, **kwargs)
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        # Redefine self_attn to our streaming multi-head attention
+        attn_kwargs: tp.Dict[str, tp.Any] = {
+            'embed_dim': d_model,
+            'num_heads': num_heads,
+            'dropout': dropout if attention_dropout is None else attention_dropout,
+            'bias': bias_attn,
+            'custom': custom,
+            'memory_efficient': memory_efficient,
+            'attention_as_float32': attention_as_float32,
+        }
+        self.self_attn: StreamingMultiheadAttention = StreamingMultiheadAttention(
+            causal=causal, past_context=past_context, rope=rope, qk_layer_norm=qk_layer_norm,
+            kv_repeat=kv_repeat, **attn_kwargs, **factory_kwargs)  # type: ignore
+        # Redefine feedforward layers to expose bias parameter
+        self.linear1 = nn.Linear(d_model, dim_feedforward, bias=bias_ff, **factory_kwargs)
+        self.linear2 = nn.Linear(dim_feedforward, d_model, bias=bias_ff, **factory_kwargs)
+
+        self.layer_scale_1: nn.Module
+        self.layer_scale_2: nn.Module
+        if layer_scale is None:
+            self.layer_scale_1 = nn.Identity()
+            self.layer_scale_2 = nn.Identity()
+        else:
+            self.layer_scale_1 = LayerScale(d_model, layer_scale, **factory_kwargs)
+            self.layer_scale_2 = LayerScale(d_model, layer_scale, **factory_kwargs)
+
+        self.cross_attention: tp.Optional[nn.Module] = None
+        if cross_attention:
+            self.cross_attention = StreamingMultiheadAttention(
+                cross_attention=True, qk_layer_norm=qk_layer_norm_cross,
+                **attn_kwargs, **factory_kwargs)
+            # Norm and dropout
+            self.dropout_cross = nn.Dropout(dropout)
+            # eps value matching that used in PyTorch reference implementation.
+            self.norm_cross = nn.LayerNorm(d_model, eps=1e-5, **factory_kwargs)
+            self.layer_scale_cross: nn.Module
+            if layer_scale is None:
+                self.layer_scale_cross = nn.Identity()
+            else:
+                self.layer_scale_cross = LayerScale(d_model, layer_scale, **factory_kwargs)
+        self.norm1 = create_norm_fn(norm, d_model, **factory_kwargs)  # type: ignore
+        self.norm2 = create_norm_fn(norm, d_model, **factory_kwargs)  # type: ignore
+
+    def _cross_attention_block(self, src: torch.Tensor,
+                               cross_attention_src: torch.Tensor) -> torch.Tensor:
+        assert self.cross_attention is not None
+        # queries are from src, keys and values from cross_attention_src.
+        x = self.cross_attention(
+            src, cross_attention_src, cross_attention_src, need_weights=False)[0]
+        return self.dropout_cross(x)  # type: ignore
+
+    def forward(self, src: torch.Tensor, src_mask: tp.Optional[torch.Tensor] = None,  # type: ignore
+                src_key_padding_mask: tp.Optional[torch.Tensor] = None,
+                cross_attention_src: tp.Optional[torch.Tensor] = None):
+        if self.cross_attention is None:
+            assert cross_attention_src is None
+        else:
+            assert cross_attention_src is not None
+        x = src
+        if self.norm_first:
+            x = x + self.layer_scale_1(
+                self._sa_block(self.norm1(x), src_mask, src_key_padding_mask))
+            if cross_attention_src is not None:
+                x = x + self.layer_scale_cross(
+                    self._cross_attention_block(
+                        self.norm_cross(x), cross_attention_src))
+            x = x + self.layer_scale_2(self._ff_block(self.norm2(x)))
+        else:
+            x = self.norm1(x + self.layer_scale_1(
+                self._sa_block(x, src_mask, src_key_padding_mask)))
+            if cross_attention_src is not None:
+                x = self.norm_cross(
+                    x + self.layer_scale_cross(
+                        self._cross_attention_block(src, cross_attention_src)))
+            x = self.norm2(x + self.layer_scale_2(self._ff_block(x)))
+        return x
+
+
+class StreamingTransformer(StreamingModule):
+    """Transformer with Streaming / Causal support.
+
+    Args:
+        d_model (int): Dimension of the data.
+        num_heads (int): Number of heads.
+        dim_feedforward (int): Intermediate dimension of FF module.
+        dropout (float): Dropout both for MHA and FF.
+        bias_ff (bool): Use bias for FF.
+        bias_attn (bool): Use bias for MHA.
+        causal (bool): Causal mask applied automatically.
+        past_context (int, optional): Receptive field for the causal mask, infinite if None.
+        custom (bool): Use custom MHA implementation, for testing / benchmarking.
+        memory_efficient (bool): Use xformers based memory efficient attention.
+        attention_as_float32 (bool): Perform the attention as float32
+            (especially important with memory_efficient as autocast won't do this automatically).
+        cross_attention (bool): If True, expect to get secondary input for cross-attention.
+        layer_scale (float, optional): If not None, LayerScale will be used
+            with the given value as initial scale.
+        positional_embedding (str): Positional embedding strategy (sin, rope, or sin_rope).
+        max_period (float): Maximum period of the time embedding.
+        positional_scale (float): Scale of positional embedding, set to 0 to deactivate.
+        xpos (bool): Apply xpos exponential decay to positional embedding (rope only).
+        lr (float, optional): learning rate override through the `make_optim_group` API.
+        weight_decay (float, optional): Weight_decay override through the `make_optim_group` API.
+        layer_class: (subclass of `StreamingTransformerLayer): class to use
+            to initialize the layers, allowing further customization outside of AudioCraft.
+        checkpointing (str): Checkpointing strategy to reduce memory usage.
+            No checkpointing if set to 'none'. Per layer checkpointing using PyTorch
+            if set to 'torch' (entire layer checkpointed, i.e. linears are evaluated twice,
+            minimal memory usage, but maximal runtime). Finally, `xformers_default` provide
+            a policy for opting-out some operations of the checkpointing like
+            linear layers and attention, providing a middle ground between speed and memory.
+        device (torch.device, optional): Device on which to initialize.
+        dtype (torch.dtype, optional): dtype to use.
+        **kwargs: See `nn.TransformerEncoderLayer`.
+    """
+    def __init__(self, d_model: int, num_heads: int, num_layers: int, dim_feedforward: int = 2048,
+                 dropout: float = 0.1, bias_ff: bool = True, bias_attn: bool = True,
+                 causal: bool = False, past_context: tp.Optional[int] = None,
+                 custom: bool = False, memory_efficient: bool = False, attention_as_float32: bool = False,
+                 cross_attention: bool = False, layer_scale: tp.Optional[float] = None,
+                 positional_embedding: str = 'sin', max_period: float = 10_000, positional_scale: float = 1.,
+                 xpos: bool = False, lr: tp.Optional[float] = None, weight_decay: tp.Optional[float] = None,
+                 layer_class: tp.Type[StreamingTransformerLayer] = StreamingTransformerLayer,
+                 checkpointing: str = 'none', device=None, dtype=None, **kwargs):
+        super().__init__()
+        assert d_model % num_heads == 0
+
+        self.positional_embedding = positional_embedding
+        self.max_period = max_period
+        self.positional_scale = positional_scale
+        self.weight_decay = weight_decay
+        self.lr = lr
+
+        assert positional_embedding in ['sin', 'rope', 'sin_rope']
+        self.rope: tp.Optional[RotaryEmbedding] = None
+        if self.positional_embedding in ['rope', 'sin_rope']:
+            assert _is_custom(custom, memory_efficient)
+            self.rope = RotaryEmbedding(d_model // num_heads, max_period=max_period,
+                                        xpos=xpos, scale=positional_scale, device=device)
+
+        self.checkpointing = checkpointing
+
+        assert checkpointing in ['none', 'torch', 'xformers_default', 'xformers_mm']
+        if self.checkpointing.startswith('xformers'):
+            _verify_xformers_internal_compat()
+
+        self.layers = nn.ModuleList()
+        for idx in range(num_layers):
+            self.layers.append(
+                layer_class(
+                    d_model=d_model, num_heads=num_heads, dim_feedforward=dim_feedforward,
+                    dropout=dropout, bias_ff=bias_ff, bias_attn=bias_attn,
+                    causal=causal, past_context=past_context, custom=custom,
+                    memory_efficient=memory_efficient, attention_as_float32=attention_as_float32,
+                    cross_attention=cross_attention, layer_scale=layer_scale, rope=self.rope,
+                    device=device, dtype=dtype, **kwargs))
+
+        if self.checkpointing != 'none':
+            for layer in self.layers:
+                # see audiocraft/optim/fsdp.py, magic signal to indicate this requires fixing the
+                # backward hook inside of FSDP...
+                layer._magma_checkpointed = True  # type: ignore
+
+    def _apply_layer(self, layer, *args, **kwargs):
+        method = self.checkpointing
+        if method == 'none':
+            return layer(*args, **kwargs)
+        elif method == 'torch':
+            return torch_checkpoint(layer, *args, use_reentrant=False, **kwargs)
+        elif method.startswith('xformers'):
+            from xformers.checkpoint_fairinternal import checkpoint, _get_default_policy
+            if method == 'xformers_default':
+                # those operations will be saved, and not recomputed.
+                # According to Francisco we can get smarter policies but this is a good start.
+                allow_list = [
+                    "xformers.efficient_attention_forward_cutlass.default",
+                    "xformers_flash.flash_fwd.default",
+                    "aten.addmm.default",
+                    "aten.mm.default",
+                ]
+            elif method == 'xformers_mm':
+                # those operations will be saved, and not recomputed.
+                # According to Francisco we can get smarter policies but this is a good start.
+                allow_list = [
+                    "aten.addmm.default",
+                    "aten.mm.default",
+                ]
+            else:
+                raise ValueError(f"xformers checkpointing xformers policy {method} is not known.")
+            policy_fn = _get_default_policy(allow_list)
+            return checkpoint(layer, *args, policy_fn=policy_fn, **kwargs)
+        else:
+            raise ValueError(f"Checkpointing method {method} is unknown.")
+
+    def forward(self, x: torch.Tensor, *args, **kwargs):
+        B, T, C = x.shape
+
+        if 'offsets' in self._streaming_state:
+            offsets = self._streaming_state['offsets']
+        else:
+            offsets = torch.zeros(B, dtype=torch.long, device=x.device)
+
+        if self.positional_embedding in ['sin', 'sin_rope']:
+            positions = torch.arange(T, device=x.device).view(1, -1, 1)
+            positions = positions + offsets.view(-1, 1, 1)
+            pos_emb = create_sin_embedding(positions, C, max_period=self.max_period, dtype=x.dtype)
+            x = x + self.positional_scale * pos_emb
+
+        for layer in self.layers:
+            x = self._apply_layer(layer, x, *args, **kwargs)
+
+        if self._is_streaming:
+            self._streaming_state['offsets'] = offsets + T
+
+        return x
+
+    def make_optim_group(self):
+        group = {"params": list(self.parameters())}
+        if self.lr is not None:
+            group["lr"] = self.lr
+        if self.weight_decay is not None:
+            group["weight_decay"] = self.weight_decay
+        return group
+
+
+# special attention related function
+
+def _verify_xformers_memory_efficient_compat():
+    try:
+        from xformers.ops import memory_efficient_attention, LowerTriangularMask  # noqa
+    except ImportError:
+        raise ImportError(
+            "xformers is not installed. Please install it and try again.\n"
+            "To install on AWS and Azure, run \n"
+            "FORCE_CUDA=1 TORCH_CUDA_ARCH_LIST='8.0'\\\n"
+            "pip install -U git+https://git@github.com/fairinternal/xformers.git#egg=xformers\n"
+            "To install on FAIR Cluster, run \n"
+            "FORCE_CUDA=1 TORCH_CUDA_ARCH_LIST='6.0;7.0'\\\n"
+            "pip install -U git+https://git@github.com/fairinternal/xformers.git#egg=xformers\n")
+
+
+def _verify_xformers_internal_compat():
+    try:
+        from xformers.checkpoint_fairinternal import checkpoint, _get_default_policy  # noqa
+    except ImportError:
+        raise ImportError(
+            "Francisco's fairinternal xformers is not installed. Please install it and try again.\n"
+            "To install on AWS and Azure, run \n"
+            "FORCE_CUDA=1 TORCH_CUDA_ARCH_LIST='8.0'\\\n"
+            "pip install -U git+https://git@github.com/fairinternal/xformers.git#egg=xformers\n"
+            "To install on FAIR Cluster, run \n"
+            "FORCE_CUDA=1 TORCH_CUDA_ARCH_LIST='6.0;7.0'\\\n"
+            "pip install -U git+https://git@github.com/fairinternal/xformers.git#egg=xformers\n")
+
+
+def _is_custom(custom: bool, memory_efficient: bool):
+    return custom or memory_efficient

models/musicgen_lm.py

@@ -0,0 +1,367 @@
+from pathlib import Path
+from torch import nn, Tensor
+import torch
+from typing import Optional, Tuple, List, Dict
+import x_transformers as xt
+
+from conditioning.condition_type import ConditionType
+from conditioning.conditioning_method import ConditioningMethod
+from conditioning.embedded_condition import EmbeddedCondition
+import hyperparameters as hp
+from utils.inspection import printshape
+import config as cfg
+
+
+class ResidualTokenEmbedding(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 num_tokens: int,
+                 n_layers: int,
+                 padding_token: Optional[int] = None):
+        super().__init__()
+        self.padding_token: Optional[int] = padding_token
+        self.emb = nn.ModuleList([
+            nn.Embedding(num_tokens, dim, padding_idx=self.padding_token)
+            for _ in range(n_layers)
+        ])
+
+        for layer in self.emb:
+            nn.init.kaiming_normal_(layer.weight)
+
+    def forward(self, x):
+        n_res_layers = x.shape[1]
+        assert n_res_layers == len(self.emb)
+        token_emb: Tensor = sum(  # type: ignore
+            [self.emb[i](x[:, i]) for i in range(n_res_layers)])
+        return token_emb
+
+
+def create_sin_embedding(
+    positions: torch.Tensor,
+    dim: int,
+    max_period: float = 10000,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """Create sinusoidal positional embedding, with shape `[B, T, C]`.
+
+    Args:
+        positions (torch.Tensor): LongTensor of positions.
+        dim (int): Dimension of the embedding.
+        max_period (float): Maximum period of the cosine/sine functions.
+        dtype (torch.dtype or str): dtype to use to generate the embedding.
+    Returns:
+        torch.Tensor: Sinusoidal positional embedding.
+    """
+    # We aim for BTC format
+    assert dim % 2 == 0
+    half_dim = dim // 2
+    positions = positions.to(dtype)
+    adim = torch.arange(half_dim, device=positions.device,
+                        dtype=dtype).view(1, 1, -1)
+    max_period_tensor = torch.full([],
+                                   max_period,
+                                   device=positions.device,
+                                   dtype=dtype)  # avoid sync point
+    phase = positions / (max_period_tensor**(adim / (half_dim - 1)))
+    return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1)
+
+
+class ResidualSinusoidalEmbedding(nn.Module):
+
+    def __init__(self, dim, theta=10000):
+        super().__init__()
+        assert dim % 2 == 0
+        self.scale = 1
+        self.theta = theta
+        self.dim = dim
+
+    def forward(self, x, pos=None, seq_start_pos=None):
+        B, K, T = x.shape
+        positions = (pos if pos is not None else torch.arange(
+            T, device=x.device).view(1, -1, 1))
+        pos_emb = create_sin_embedding(positions,
+                                       self.dim,
+                                       max_period=self.theta,
+                                       dtype=x.dtype)
+        return pos_emb * self.scale
+
+
+class ResidualOutputProj(nn.Module):
+
+    def __init__(self, input_dim: int, output_dim: int, n_layers: int,
+                 use_bias: bool):
+        super().__init__()
+        self.linears = nn.ModuleList([
+            nn.Linear(input_dim, output_dim, use_bias) for _ in range(n_layers)
+        ])
+
+    def forward(self, x):
+        return torch.stack([layer(x) for layer in self.linears], dim=1)
+
+
+class DropoutModule(nn.Module):
+    """Base module for all dropout modules."""
+
+    def __init__(self, seed: int = 1234):
+        super().__init__()
+        self.rng = torch.Generator()
+        self.rng.manual_seed(seed)
+
+
+class ConditioningDropout(DropoutModule):
+
+    def __init__(self, p: float, seed: int = 1234):
+        super().__init__(seed=seed)
+        self.p = p
+
+    def forward(self, x: Tensor, mask: Tensor) -> Tuple[Tensor, Tensor]:
+        if not self.training:
+            return x, mask
+        drop = torch.rand(1, generator=self.rng).item() < self.p
+        if not drop:
+            return x, mask
+        return torch.zeros_like(x), torch.zeros_like(mask)
+
+
+class ClassifierFreeGuidanceDropout(DropoutModule):
+    """Classifier Free Guidance dropout.
+    All attributes are dropped with the same probability.
+
+    Args:
+        p (float): Probability to apply condition dropout during training.
+        seed (int): Random seed.
+    """
+
+    def __init__(self, p: float, seed: int = 1234):
+        super().__init__(seed=seed)
+        self.p = p
+
+    def forward(self, samples: List) -> List[str | None]:
+        if not self.training:
+            return samples
+        drop = torch.rand(1, generator=self.rng).item() < self.p
+        if not drop:
+            return samples
+        return [None for _ in range(len(samples))]
+
+    def __repr__(self):
+        return f"ClassifierFreeGuidanceDropout(p={self.p})"
+
+
+class MusicgenLm(nn.Module):
+
+    def __init__(self, params: hp.LmParams):
+        super().__init__()
+
+        self.dim: int = params.dim
+        self.n_layers: int = params.n_layers
+        self.n_heads: int = params.n_heads
+        self.card: int = params.card
+        self.cross_attend: bool = params.cross_attend
+        self.padding_token: Optional[int] = params.padding_token
+        self.sep_token: Optional[int] = params.sep_token
+
+        # if using an extra padding token, input card is card+1
+        self.input_card: int = self.card + (
+            1 if self.padding_token is not None and
+            self.padding_token >= self.card else 0)
+
+        if self.sep_token is not None:
+            self.input_card += 1
+
+        self.cfg_coef: float = 3.0
+        self.n_q = 4
+
+        # DECODER
+        self.decoder: xt.TransformerWrapper = xt.TransformerWrapper(
+            num_tokens=self.input_card,
+            max_seq_len=500,
+            use_abs_pos_emb=True,
+            scaled_sinu_pos_emb=True,
+            attn_layers=xt.Decoder(
+                dim=self.dim,
+                depth=self.n_layers,
+                heads=self.n_heads,
+                attn_dim_head=64,
+                attn_flash=True,
+                ff_no_bias=True,
+                cross_attend=self.cross_attend,
+            ),
+        )
+        self.decoder.token_emb = ResidualTokenEmbedding(  # type: ignore
+            self.dim,
+            self.input_card,
+            self.n_q,
+            self.padding_token,
+        )
+        self.decoder.pos_emb = ResidualSinusoidalEmbedding(  # type: ignore
+            dim=self.dim)
+        self.decoder.to_logits = ResidualOutputProj(self.dim, self.card,
+                                                    self.n_q, False)
+
+        # TODO: this is horrendous, gotta find a fix
+        if not self.cross_attend:
+            self.nullwav_embeds = torch.load(cfg.weights_dir() /
+                                             "nullwav_embeds.pt")[0]
+        else:
+            self.nullwav_embeds = None
+
+        if params.weights is not None:
+
+            ckpt_state_dict: Dict = torch.load(Path(params.weights),
+                                               map_location=None,
+                                               weights_only=True)
+            for key in ("conditioner.output_proj.weight",
+                        "conditioner.output_proj.bias"):
+                if key in ckpt_state_dict:
+                    ckpt_state_dict.pop(key)
+
+            if self.sep_token is not None:
+                for k in range(self.n_q):
+                    param_name = f"decoder.token_emb.emb.{k}.weight"
+                    param = ckpt_state_dict[param_name]
+                    assert self.sep_token == param.shape[0]
+                    sep_token_emb = torch.empty_like(param[:1])
+                    nn.init.kaiming_normal_(sep_token_emb)
+                    param = torch.cat((param, sep_token_emb))
+                    ckpt_state_dict[param_name] = param
+
+            self.load_state_dict(ckpt_state_dict, strict=True)
+
+            # zero-out padding token in embedding
+            # TODO: figure out how to handle doing or not this shit
+            if self.cross_attend:
+                if self.padding_token is not None:
+                    with torch.no_grad():
+                        for k in range(self.n_q):
+                            self.decoder.token_emb.emb[  # type: ignore
+                                k].weight[
+                                    self.padding_token] = torch.zeros_like(
+                                        self.decoder.token_emb.
+                                        emb[k].  # type: ignore
+                                        weight[0])
+
+    def forward(
+        self,
+        x: Tensor,
+        attention_mask: Tensor,
+        cross_attention_input: Optional[EmbeddedCondition] = None,
+        prepend_embeds: Optional[EmbeddedCondition] = None,
+        sum_embeds: Optional[EmbeddedCondition] = None,
+    ) -> Tensor:
+
+        # unpack cross_attention
+        if cross_attention_input is not None:
+            cross_attention_data = cross_attention_input.data
+            cross_attention_mask = cross_attention_input.mask
+
+            # for samples in the batch in which xatt-conditioning is all padding,
+            # place dummy zero-valued conditioning vector, with mask set to True
+            if cross_attention_mask is not None:
+                masked_rows_coords = cross_attention_mask.sum(dim=-1) == 0
+                cross_attention_data[masked_rows_coords] = torch.zeros_like(
+                    cross_attention_data[masked_rows_coords])
+                # cross_attention_mask[masked_rows_coords, ..., 0] = True
+                cross_attention_mask[masked_rows_coords] = torch.ones_like(
+                    cross_attention_mask[masked_rows_coords])
+        else:
+            if self.cross_attend:
+                cross_attention_data = torch.zeros(x.shape[0],
+                                                   1,
+                                                   self.dim,
+                                                   dtype=torch.float32,
+                                                   device=x.device)
+                cross_attention_mask = torch.ones(x.shape[0],
+                                                  1,
+                                                  dtype=torch.bool,
+                                                  device=x.device)
+            else:
+                cross_attention_data = None
+                cross_attention_mask = None
+
+        # unpack prepend embeddings
+        prepend_embeds_data, prepend_embeds_mask = None, None
+        if prepend_embeds is not None:
+            prepend_embeds_data = prepend_embeds.data
+            prepend_embeds_mask = prepend_embeds.mask
+
+            if self.nullwav_embeds is not None:
+                raise RuntimeError(
+                    "CAREFUL! You are using the musicgen-melody hack!")
+                nwe = self.nullwav_embeds.repeat(prepend_embeds_data.shape[0],
+                                                 1, 1).to(prepend_embeds_data)
+                prepend_embeds_data = torch.cat((nwe, prepend_embeds_data),
+                                                dim=1)
+                if prepend_embeds_mask is not None:
+                    nwm = torch.ones(nwe.shape[0],
+                                     nwe.shape[1],
+                                     device=nwe.device,
+                                     dtype=torch.bool)
+                    prepend_embeds_mask = torch.cat((nwm, prepend_embeds_mask),
+                                                    dim=1)
+
+        # unpack sum embeddings
+        sum_embeds_data, sum_embeds_mask = None, None
+        if sum_embeds is not None:
+            sum_embeds_data = sum_embeds.data
+            sum_embeds_mask = sum_embeds.mask
+            if sum_embeds_mask is not None:
+                sum_embeds_data[~sum_embeds_mask] = 0
+
+        assert attention_mask.shape == torch.Size([x.shape[0], x.shape[-1]])
+
+        # create manually positional embedding indices
+        positions = torch.zeros_like(attention_mask, dtype=torch.int64)
+        B, S = positions.shape
+        first_valid_indices = torch.argmax(attention_mask.long(), dim=-1)
+        sequence = torch.arange(S, device=positions.device).unsqueeze(0).expand(
+            B, S)
+        mask = torch.arange(S, device=positions.device).unsqueeze(
+            0) >= first_valid_indices.unsqueeze(1)
+        positions[mask] = (sequence - first_valid_indices.unsqueeze(1))[mask]
+        positions = positions.unsqueeze(-1)
+
+        # create mask
+        mask = torch.ones((x.shape[0], x.shape[-1]),
+                          dtype=torch.bool,
+                          device=x.device)
+
+        # positional embeddings are manually applied to prepend data
+        # if not self.cross_attend and prepend_embeds_data is not None:
+        if prepend_embeds_data is not None:
+            positions += prepend_embeds_data.shape[-2]
+            prepend_embeds_data = prepend_embeds_data + self.decoder.pos_emb(
+                prepend_embeds_data[..., 0].unsqueeze(1))
+
+        # forward through decoder
+        logits = self.decoder(
+            x,
+            # mask=mask,
+            mask=attention_mask,
+            pos=positions,
+            context=cross_attention_data,
+            context_mask=cross_attention_mask,
+            prepend_embeds=prepend_embeds_data,
+            prepend_mask=prepend_embeds_mask,
+            sum_embeds=sum_embeds_data,
+        )
+        # if something is prepended to the input, remove it from the logits
+        if prepend_embeds_data is not None:
+            logits = logits[:, :, prepend_embeds_data.shape[-2]:, :]
+
+        return logits
+
+
+if __name__ == "__main__":
+    # params = hp.PretrainedSmallLmParams()
+
+    # model = params.instantiate()
+
+    # model.eval()
+    res = ResidualSinusoidalEmbedding(1024)
+
+    x = torch.rand(1, 1024, 3)
+    emb = res(x)
+    print(f'{emb.shape=}')
+    print(f'{emb=}')

models/quantization/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""RVQ."""
+# flake8: noqa
+from .vq import ResidualVectorQuantizer
+from .base import BaseQuantizer, DummyQuantizer, QuantizedResult

models/quantization/base.py

@@ -0,0 +1,104 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""
+Base class for all quantizers.
+"""
+
+from dataclasses import dataclass, field
+import typing as tp
+
+import torch
+from torch import nn
+
+
+@dataclass
+class QuantizedResult:
+    quantized_layers: torch.Tensor  # q1, q1+q2, q1+q2+q3 ...
+    x: torch.Tensor  # final quantized vector: Σ(q_i) == quantized_layers[-1]
+    codes: torch.Tensor
+    bandwidth: torch.Tensor  # bandwidth in kb/s used, per batch item.
+    penalty: tp.Optional[torch.Tensor] = None
+    sum_loss: tp.Optional[torch.Tensor] = None
+    metrics: dict = field(default_factory=dict)
+
+
+class BaseQuantizer(nn.Module):
+    """Base class for quantizers.
+    """
+
+    def forward(self, x: torch.Tensor, frame_rate: int) -> QuantizedResult:
+        """
+        Given input tensor x, returns first the quantized (or approximately quantized)
+        representation along with quantized codes, bandwidth, and any penalty term for the loss.
+        Finally, this returns a dict of metrics to update logging etc.
+        Frame rate must be passed so that the bandwidth is properly computed.
+        """
+        raise NotImplementedError()
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified sample rate at the given bandwidth."""
+        raise NotImplementedError()
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation."""
+        raise NotImplementedError()
+
+    @property
+    def total_codebooks(self):
+        """Total number of codebooks."""
+        raise NotImplementedError()
+
+    @property
+    def num_codebooks(self):
+        """Number of active codebooks."""
+        raise NotImplementedError()
+
+    def set_num_codebooks(self, n: int):
+        """Set the number of active codebooks."""
+        raise NotImplementedError()
+
+
+class DummyQuantizer(BaseQuantizer):
+    """Fake quantizer that actually does not perform any quantization.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor, frame_rate: int):
+        q = x.unsqueeze(1)
+        return QuantizedResult(
+            x, q,
+            torch.tensor(q.numel() * 32 * frame_rate / 1000 / len(x)).to(x))
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified sample rate at the given bandwidth.
+        In the case of the DummyQuantizer, the codes are actually identical
+        to the input and resulting quantized representation as no quantization is done.
+        """
+        return x.unsqueeze(1)
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation.
+        In the case of the DummyQuantizer, the codes are actually identical
+        to the input and resulting quantized representation as no quantization is done.
+        """
+        return codes.squeeze(1)
+
+    @property
+    def total_codebooks(self):
+        """Total number of codebooks."""
+        return 1
+
+    @property
+    def num_codebooks(self):
+        """Total number of codebooks."""
+        return self.total_codebooks
+
+    def set_num_codebooks(self, n: int):
+        """Set the number of active codebooks."""
+        raise AttributeError(
+            "Cannot override the number of codebooks for the dummy quantizer")

models/quantization/core_vq.py

@@ -0,0 +1,420 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import typing as tp
+
+from einops import rearrange, repeat
+# import flashy
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+
+
+def exists(val: tp.Optional[tp.Any]) -> bool:
+    return val is not None
+
+
+def default(val: tp.Any, d: tp.Any) -> tp.Any:
+    return val if exists(val) else d
+
+
+def l2norm(t):
+    return F.normalize(t, p=2, dim=-1)
+
+
+def ema_inplace(moving_avg, new, decay: float):
+    moving_avg.data.mul_(decay).add_(new.detach(), alpha=(1 - decay))
+
+
+def laplace_smoothing(x, n_categories: int, epsilon: float = 1e-5):
+    return (x + epsilon) / (x.sum() + n_categories * epsilon)
+
+
+def uniform_init(*shape: int):
+    t = torch.empty(shape)
+    nn.init.kaiming_uniform_(t)
+    return t
+
+
+def sample_vectors(samples, num: int):
+    num_samples, device = samples.shape[0], samples.device
+
+    if num_samples >= num:
+        indices = torch.randperm(num_samples, device=device)[:num]
+    else:
+        indices = torch.randint(0, num_samples, (num,), device=device)
+
+    return samples[indices]
+
+
+def kmeans(samples, num_clusters: int, num_iters: int = 10):
+    dim, dtype = samples.shape[-1], samples.dtype
+
+    means = sample_vectors(samples, num_clusters)
+
+    for _ in range(num_iters):
+        diffs = rearrange(samples, "n d -> n () d") - rearrange(
+            means, "c d -> () c d")
+        dists = -(diffs**2).sum(dim=-1)
+
+        buckets = dists.max(dim=-1).indices
+        bins = torch.bincount(buckets, minlength=num_clusters)
+        zero_mask = bins == 0
+        bins_min_clamped = bins.masked_fill(zero_mask, 1)
+
+        new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)
+        new_means.scatter_add_(0, repeat(buckets, "n -> n d", d=dim), samples)
+        new_means = new_means / bins_min_clamped[..., None]
+
+        means = torch.where(zero_mask[..., None], means, new_means)
+
+    return means, bins
+
+
+def orthogonal_loss_fn(t):
+    # eq (2) from https://arxiv.org/abs/2112.00384
+    n = t.shape[0]
+    normed_codes = l2norm(t)
+    identity = torch.eye(n, device=t.device)
+    cosine_sim = einsum("i d, j d -> i j", normed_codes, normed_codes)
+    return ((cosine_sim - identity)**2).sum() / (n**2)
+
+
+class EuclideanCodebook(nn.Module):
+    """Codebook with Euclidean distance.
+
+    Args:
+        dim (int): Dimension.
+        codebook_size (int): Codebook size.
+        kmeans_init (bool): Whether to use k-means to initialize the codebooks.
+            If set to true, run the k-means algorithm on the first training batch and use
+            the learned centroids as initialization.
+        kmeans_iters (int): Number of iterations used for k-means algorithm at initialization.
+        decay (float): Decay for exponential moving average over the codebooks.
+        epsilon (float): Epsilon value for numerical stability.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        codebook_size: int,
+        kmeans_init: int = False,
+        kmeans_iters: int = 10,
+        decay: float = 0.8,
+        epsilon: float = 1e-5,
+        threshold_ema_dead_code: int = 2,
+    ):
+        super().__init__()
+        self.decay = decay
+        init_fn: tp.Union[
+            tp.Callable[..., torch.Tensor],
+            tp.Any] = uniform_init if not kmeans_init else torch.zeros
+        embed = init_fn(codebook_size, dim)
+
+        self.codebook_size = codebook_size
+
+        self.kmeans_iters = kmeans_iters
+        self.epsilon = epsilon
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+
+        self.register_buffer("inited", torch.Tensor([not kmeans_init]))
+        self.register_buffer("cluster_size", torch.zeros(codebook_size))
+        self.register_buffer("embed", embed)
+        self.register_buffer("embed_avg", embed.clone())
+
+    # @torch.jit.ignore
+    def init_embed_(self, data):
+        if self.inited:
+            return
+
+        embed, cluster_size = kmeans(data, self.codebook_size,
+                                     self.kmeans_iters)
+        self.embed.data.copy_(embed)
+        self.embed_avg.data.copy_(embed.clone())
+        self.cluster_size.data.copy_(cluster_size)
+        self.inited.data.copy_(torch.Tensor([True]))
+        # Make sure all buffers across workers are in sync after initialization
+        # flashy.distrib.broadcast_tensors(self.buffers())
+
+    def replace_(self, samples, mask):
+        modified_codebook = torch.where(
+            mask[..., None], sample_vectors(samples, self.codebook_size),
+            self.embed)
+        self.embed.data.copy_(modified_codebook)
+
+    def expire_codes_(self, batch_samples):
+        if self.threshold_ema_dead_code == 0:
+            return
+
+        expired_codes = self.cluster_size < self.threshold_ema_dead_code
+        if not torch.any(expired_codes):
+            return
+
+        batch_samples = rearrange(batch_samples, "... d -> (...) d")
+        self.replace_(batch_samples, mask=expired_codes)
+        # flashy.distrib.broadcast_tensors(self.buffers())
+
+    def preprocess(self, x):
+        x = rearrange(x, "... d -> (...) d")
+        return x
+
+    def quantize(self, x):
+        embed = self.embed.t()
+        dist = -(x.pow(2).sum(1, keepdim=True) - 2 * x @ embed +
+                 embed.pow(2).sum(0, keepdim=True))
+        embed_ind = dist.max(dim=-1).indices
+        return embed_ind
+
+    def postprocess_emb(self, embed_ind, shape):
+        return embed_ind.view(*shape[:-1])
+
+    def dequantize(self, embed_ind):
+        quantize = F.embedding(embed_ind, self.embed)
+        return quantize
+
+    def encode(self, x):
+        shape = x.shape
+        # pre-process
+        x = self.preprocess(x)
+        # quantize
+        embed_ind = self.quantize(x)
+        # post-process
+        embed_ind = self.postprocess_emb(embed_ind, shape)
+        return embed_ind
+
+    def decode(self, embed_ind):
+        quantize = self.dequantize(embed_ind)
+        return quantize
+
+    def forward(self, x):
+        shape, dtype = x.shape, x.dtype
+        x = self.preprocess(x)
+        self.init_embed_(x)
+
+        embed_ind = self.quantize(x)
+        embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
+        embed_ind = self.postprocess_emb(embed_ind, shape)
+        quantize = self.dequantize(embed_ind)
+
+        if self.training:
+            # We do the expiry of code at that point as buffers are in sync
+            # and all the workers will take the same decision.
+            self.expire_codes_(x)
+            ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)
+            embed_sum = x.t() @ embed_onehot
+            ema_inplace(self.embed_avg, embed_sum.t().detach(), self.decay)
+            cluster_size = (laplace_smoothing(
+                self.cluster_size, self.codebook_size, self.epsilon) *
+                            self.cluster_size.sum())
+            embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
+            self.embed.data.copy_(embed_normalized)
+
+        return quantize, embed_ind
+
+
+class VectorQuantization(nn.Module):
+    """Vector quantization implementation.
+    Currently supports only euclidean distance.
+
+    Args:
+        dim (int): Dimension
+        codebook_size (int): Codebook size
+        codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim.
+        decay (float): Decay for exponential moving average over the codebooks.
+        epsilon (float): Epsilon value for numerical stability.
+        kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
+        kmeans_iters (int): Number of iterations used for kmeans initialization.
+        threshold_ema_dead_code (int):
+        channels_last (bool): Channels are the last dimension in the input tensors.
+        commitment_weight (float): Weight for commitment loss.
+        orthogonal_reg_weight (float): Orthogonal regularization weights.
+        orthogonal_reg_active_codes_only (bool): Apply orthogonal regularization only on active codes.
+        orthogonal_reg_max_codes (optional int): Maximum number of codes to consider
+            for orthogonal regularization.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        codebook_size: int,
+        codebook_dim: tp.Optional[int] = None,
+        decay: float = 0.8,
+        epsilon: float = 1e-5,
+        kmeans_init: bool = False,
+        kmeans_iters: int = 10,
+        threshold_ema_dead_code: int = 2,
+        channels_last: bool = False,
+        commitment_weight: float = 1.,
+        orthogonal_reg_weight: float = 0.0,
+        orthogonal_reg_active_codes_only: bool = False,
+        orthogonal_reg_max_codes: tp.Optional[int] = None,
+    ):
+        super().__init__()
+        _codebook_dim: int = default(codebook_dim, dim)
+
+        requires_projection = _codebook_dim != dim
+        self.project_in = (nn.Linear(dim, _codebook_dim)
+                           if requires_projection else nn.Identity())
+        self.project_out = (nn.Linear(_codebook_dim, dim)
+                            if requires_projection else nn.Identity())
+
+        self.epsilon = epsilon
+        self.commitment_weight = commitment_weight
+
+        self.orthogonal_reg_weight = orthogonal_reg_weight
+        self.orthogonal_reg_active_codes_only = orthogonal_reg_active_codes_only
+        self.orthogonal_reg_max_codes = orthogonal_reg_max_codes
+
+        self._codebook = EuclideanCodebook(
+            dim=_codebook_dim,
+            codebook_size=codebook_size,
+            kmeans_init=kmeans_init,
+            kmeans_iters=kmeans_iters,
+            decay=decay,
+            epsilon=epsilon,
+            threshold_ema_dead_code=threshold_ema_dead_code)
+        self.codebook_size = codebook_size
+
+        self.channels_last = channels_last
+
+    @property
+    def codebook(self):
+        return self._codebook.embed
+
+    @property
+    def inited(self):
+        return self._codebook.inited
+
+    def _preprocess(self, x):
+        if not self.channels_last:
+            x = rearrange(x, "b d n -> b n d")
+        return x
+
+    def _postprocess(self, quantize):
+        if not self.channels_last:
+            quantize = rearrange(quantize, "b n d -> b d n")
+        return quantize
+
+    def encode(self, x):
+        x = self._preprocess(x)
+        x = self.project_in(x)
+        embed_in = self._codebook.encode(x)
+        return embed_in
+
+    def decode(self, embed_ind):
+        quantize = self._codebook.decode(embed_ind)
+        quantize = self.project_out(quantize)
+        quantize = self._postprocess(quantize)
+        return quantize
+
+    def forward(self, x):
+        device = x.device
+        x = self._preprocess(x)
+
+        x = self.project_in(x)
+        quantize, embed_ind = self._codebook(x)
+
+        if self.training:
+            quantize = x + (quantize - x).detach()
+
+        loss = torch.tensor([0.0], device=device, requires_grad=self.training)
+
+        if self.training:
+            if self.commitment_weight > 0:
+                commit_loss = F.mse_loss(quantize.detach(), x)
+                loss = loss + commit_loss * self.commitment_weight
+
+            if self.orthogonal_reg_weight > 0:
+                codebook = self.codebook
+
+                if self.orthogonal_reg_active_codes_only:
+                    # only calculate orthogonal loss for the activated codes for this batch
+                    unique_code_ids = torch.unique(embed_ind)
+                    codebook = codebook[unique_code_ids]
+
+                num_codes = codebook.shape[0]
+                if exists(self.orthogonal_reg_max_codes
+                         ) and num_codes > self.orthogonal_reg_max_codes:
+                    rand_ids = torch.randperm(
+                        num_codes,
+                        device=device)[:self.orthogonal_reg_max_codes]
+                    codebook = codebook[rand_ids]
+
+                orthogonal_reg_loss = orthogonal_loss_fn(codebook)
+                loss = loss + orthogonal_reg_loss * self.orthogonal_reg_weight
+
+        quantize = self.project_out(quantize)
+        quantize = self._postprocess(quantize)
+
+        return quantize, embed_ind, loss
+
+
+class ResidualVectorQuantization(nn.Module):
+    """Residual vector quantization implementation.
+
+    Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf
+    """
+
+    def __init__(self, *, num_quantizers, **kwargs):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [VectorQuantization(**kwargs) for _ in range(num_quantizers)])
+
+    def forward(self, x, n_q: tp.Optional[int] = None):
+        # quantized_out = 0.0
+        quantized_out = torch.tensor(0.0)
+
+        residual = x
+
+        all_losses = []
+        all_indices = []
+        quantized_list = []
+
+        n_q = n_q or len(self.layers)
+
+        for i, layer in enumerate(self.layers[:n_q]):
+            quantized, indices, loss = layer(residual)
+            residual = residual - quantized
+            quantized_out = quantized_out + quantized
+
+            # to train sum_loss, also return quantization layers
+            quantized_list.append(quantized_out)
+
+            all_indices.append(indices)
+            all_losses.append(loss)
+
+        out_losses, out_indices = map(torch.stack, (all_losses, all_indices))
+        out_layers = torch.stack(quantized_list, dim=1)
+        assert out_layers[:, -1, ...].equal(quantized_out)
+        return out_layers, quantized_out, out_indices, out_losses
+
+    def encode(self,
+               x: torch.Tensor,
+               n_q: tp.Optional[int] = None) -> torch.Tensor:
+        residual = x
+        all_indices = []
+        n_q = n_q or len(self.layers)
+        for layer in self.layers[:n_q]:
+            indices = layer.encode(residual)
+            quantized = layer.decode(indices)
+            residual = residual - quantized
+            all_indices.append(indices)
+        out_indices = torch.stack(all_indices)
+        return out_indices
+
+    def decode(self, q_indices: torch.Tensor) -> torch.Tensor:
+        quantized_out = torch.tensor(0.0, device=q_indices.device)
+        for i, indices in enumerate(q_indices):
+            layer = self.layers[i]
+            quantized = layer.decode(indices)
+            quantized_out = quantized_out + quantized
+        return quantized_out

models/quantization/vq.py

@@ -0,0 +1,121 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+import typing as tp
+
+import torch
+
+from .base import BaseQuantizer, QuantizedResult
+from .core_vq import ResidualVectorQuantization
+
+
+class ResidualVectorQuantizer(BaseQuantizer):
+    """Residual Vector Quantizer.
+
+    Args:
+        dimension (int): Dimension of the codebooks.
+        n_q (int): Number of residual vector quantizers used.
+        q_dropout (bool): Random quantizer drop out at train time.
+        bins (int): Codebook size.
+        decay (float): Decay for exponential moving average over the codebooks.
+        kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
+        kmeans_iters (int): Number of iterations used for kmeans initialization.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+        orthogonal_reg_weight (float): Orthogonal regularization weights.
+        orthogonal_reg_active_codes_only (bool): Apply orthogonal regularization only on active codes.
+        orthogonal_reg_max_codes (optional int): Maximum number of codes to consider.
+            for orthogonal regularization.
+    """
+
+    def __init__(
+        self,
+        dimension: int = 256,
+        n_q: int = 8,
+        q_dropout: bool = False,
+        bins: int = 1024,
+        decay: float = 0.99,
+        kmeans_init: bool = True,
+        kmeans_iters: int = 10,
+        threshold_ema_dead_code: int = 2,
+        orthogonal_reg_weight: float = 0.0,
+        orthogonal_reg_active_codes_only: bool = False,
+        orthogonal_reg_max_codes: tp.Optional[int] = None,
+    ):
+        super().__init__()
+        self.max_n_q = n_q
+        self.n_q = n_q
+        self.q_dropout = q_dropout
+        self.dimension = dimension
+        self.bins = bins
+        self.decay = decay
+        self.kmeans_init = kmeans_init
+        self.kmeans_iters = kmeans_iters
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+        self.orthogonal_reg_weight = orthogonal_reg_weight
+        self.orthogonal_reg_active_codes_only = orthogonal_reg_active_codes_only
+        self.orthogonal_reg_max_codes = orthogonal_reg_max_codes
+        self.vq = ResidualVectorQuantization(
+            dim=self.dimension,
+            codebook_size=self.bins,
+            num_quantizers=self.n_q,
+            decay=self.decay,
+            kmeans_init=self.kmeans_init,
+            kmeans_iters=self.kmeans_iters,
+            threshold_ema_dead_code=self.threshold_ema_dead_code,
+            orthogonal_reg_weight=self.orthogonal_reg_weight,
+            orthogonal_reg_active_codes_only=self.
+            orthogonal_reg_active_codes_only,
+            orthogonal_reg_max_codes=self.orthogonal_reg_max_codes,
+            channels_last=False)
+
+    def forward(self, x: torch.Tensor, frame_rate: int):
+        n_q = self.n_q
+        if self.training and self.q_dropout:
+            n_q = int(torch.randint(1, self.n_q + 1, (1,)).item())
+        bw_per_q = math.log2(self.bins) * frame_rate / 1000  #kbps
+        # to support sum_loss, we also return quantized layers
+        quantized_layers, quantized, codes, commit_loss = self.vq(x, n_q=n_q)
+        codes = codes.transpose(0, 1)
+        # codes is [B, K, T], with T frames, K nb of codebooks.
+        bw = torch.tensor(n_q * bw_per_q).to(x)
+        return QuantizedResult(quantized_layers,
+                               quantized,
+                               codes,
+                               bw,
+                               penalty=torch.mean(commit_loss))
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified frame rate at the given bandwidth.
+        The RVQ encode method sets the appropriate number of quantizer to use
+        and returns indices for each quantizer.
+        """
+        n_q = self.n_q
+        codes = self.vq.encode(x, n_q=n_q)
+        codes = codes.transpose(0, 1)
+        # codes is [B, K, T], with T frames, K nb of codebooks.
+        return codes
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation."""
+        # codes is [B, K, T], with T frames, K nb of codebooks, vq.decode expects [K, B, T].
+        codes = codes.transpose(0, 1)
+        quantized = self.vq.decode(codes)
+        return quantized
+
+    @property
+    def total_codebooks(self):
+        return self.max_n_q
+
+    @property
+    def num_codebooks(self):
+        return self.n_q
+
+    def set_num_codebooks(self, n: int):
+        assert n > 0 and n <= self.max_n_q
+        self.n_q = n

requirements.txt

@@ -0,0 +1,15 @@
+git+https://github.com/TEAMuP-dev/pyharp.git@v0.3.0
+# model-specific deps:
+torch
+torchaudio
+transformers
+sentencepiece
+safetensors
+lightning
+x-transformers==1.26.0
+frozendict
+loralib
+soundfile
+numpy
+matplotlib
+einops

train_stage_drums.py

@@ -0,0 +1,87 @@
+"""
+Train STAGE for DRUMS generation, with either a mixture or a metronome track as context
+"""
+import torch
+import lightning as L
+
+import hyperparameters as hp
+from conditioning.condition_type import ConditionType
+from conditioning.conditioning_method import ConditioningMethod
+from conditioning.prompt_processor import InterleavedContextPromptProcessor
+from conditioning.t5embedder import T5EmbedderGPU
+from data.stem import Stem
+from training.train import train
+import config as cfg
+
+RUN_NAME = "stage-drums"
+
+# distributed strategy - set to DDP to train on multi-gpu machines
+STRATEGY = None
+
+
+def launch_train():
+    lm_params = hp.PretrainedSmallLmParams(sep_token=2049)
+
+    conditioning_params = hp.ConditioningParams(
+        embedder_types={
+            ConditionType.DESCRIPTION: T5EmbedderGPU,
+        },
+        conditioning_methods={
+            ConditionType.DESCRIPTION: ConditioningMethod.CROSS_ATTENTION,
+        },
+        conditioning_dropout=0.5)
+
+    prompt_params = hp.PromptProcessorParams(
+        keep_only_valid_steps=True,
+        model_class=InterleavedContextPromptProcessor,
+        context_dropout=0.1)
+
+    encodec_params = hp.pretrained_encodec_meta_32khz_params
+
+    model_params = hp.MusicgenParams(encodec_params=encodec_params,
+                                     prompt_processor_params=prompt_params,
+                                     conditioning_params=conditioning_params,
+                                     lm_params=lm_params)
+    batch_size_train = 2
+    max_steps = 100_000
+    max_time = "00:24:00:00"
+
+    n_samples_per_epoch = 10_000
+
+    accumulate_grad_batches = 4
+
+    dataset_params = hp.StemmedDatasetParams(
+        clip_length_in_seconds=10,
+        sample_rate=32_000,
+        root_dir=cfg.moises_path(),
+        single_stem=True,
+        target_stem=Stem.DRUMS,
+        min_context_seconds=5,
+        use_style_conditioning=True,
+        use_beat_conditioning=True,
+        type_of_context="stems or beats",
+        add_click=False,
+        sync_chunks=False,
+        bpm_in_caption=False,
+        batch_size_train=batch_size_train,
+        batch_size_test=12,
+        num_workers=11,
+        speed_transform_p=0.5,
+        pitch_transform_p=0.5,
+        n_samples_per_epoch=n_samples_per_epoch)
+
+    train(
+        model_params=model_params,
+        dataset_params=dataset_params,
+        max_time=max_time,
+        max_steps=max_steps,
+        accumulate_grad_batches=accumulate_grad_batches,
+        run_name=RUN_NAME,
+        distributed_strategy=STRATEGY,
+        log=True,
+        kill_on_end=False,
+    )
+
+
+if __name__ == "__main__":
+    launch_train()

training/callback.py

@@ -0,0 +1,333 @@
+from concurrent.futures import ThreadPoolExecutor
+import pickle
+import librosa
+import lightning as L
+from lightning.pytorch.utilities.seed import isolate_rng
+from lightning.pytorch.callbacks import Callback
+from torch import Tensor
+from typing import Dict, Any
+from pathlib import Path
+import random
+import wandb
+import torch
+from lightning.pytorch.utilities import rank_zero_only
+
+from conditioning.beat_embedder import Beat
+from utils import audio
+import config as cfg
+# from cocola.contrastive_model import CoCola
+# from cocola.feature_extraction import CoColaFeatureExtractor
+# from cocola import constants as cocola_constants
+# from utils.logging import upload_to_s3
+
+# Custom progress bar to refresh only twice per epoch
+# class SlowProgressBar(TQDMProgressBar):
+
+#     def init_train_tqdm(self):
+#         bar = super().init_train_tqdm()
+#         # Refresh rate: update the bar only twice per epoch
+#         bar.refresh_rate = self.total_train_batches // 2
+#         return bar
+
+
+class SaveDemoOnValidationCallback(Callback):
+
+    @rank_zero_only
+    def __init__(self, save_path: Path, save_model: bool, n_demos: int):
+        self.save_path: Path = save_path
+        self.save_model: bool = save_model
+        self.save_path.mkdir(parents=True, exist_ok=True)
+        self.n_demos: int = n_demos
+        self.executor = ThreadPoolExecutor(max_workers=3)
+
+    # @rank_zero_only
+    # def on_validation_end(self, trainer: L.Trainer,
+    #                       pl_module: L.LightningModule) -> None:
+    #     if self.save_model:
+    #         if not (self.save_path / "params.pkl").exists():
+    #             with open(self.save_path / "params.pkl", "wb") as f:
+    #                 pickle.dump(pl_module.params, f)
+
+    #         checkpoint_path = self.save_path / f"step={trainer.global_step}.ckpt"
+    #         trainer.save_checkpoint(checkpoint_path, weights_only=False)
+
+    #         # upload checkpoint to S3
+    #         run_name = self.save_path.name
+    #         future = self.executor.submit(
+    #             upload_to_s3,
+    #             checkpoint_path,
+    #             "lag-modular",
+    #             f"checkpoints/{run_name}/{checkpoint_path.name}",
+    #         )
+
+    #         def log_upload_status(fut):
+    #             try:
+    #                 success = fut.result()
+    #                 status = "successful" if success else "failed"
+    #                 print(f"Checkpoint upload on S3 {status}")
+    #             except Exception as e:
+    #                 print(f"Error during upload of {checkpoint_path}: {e}")
+
+    #         future.add_done_callback(log_upload_status)
+
+    @rank_zero_only
+    def on_validation_epoch_start(self, trainer: L.Trainer,
+                                  pl_module: L.LightningModule) -> None:
+
+        total_batches = trainer.num_val_batches[0]
+        assert isinstance(total_batches, int)
+        self.random_batch_idx = random.randint(0, total_batches - 1)
+
+    def on_train_end(self, trainer, pl_module):
+        print("Shutting down uploader")
+        self.executor.shutdown(wait=True)
+        print("Uploader shut down")
+
+    @rank_zero_only
+    def on_validation_batch_end(self,
+                                trainer: L.Trainer,
+                                pl_module: L.LightningModule,
+                                outputs: Any,
+                                batch: Dict[str, Any],
+                                batch_idx: int,
+                                dataloader_idx: int = 0) -> None:
+        if batch_idx == self.random_batch_idx:
+            self.generate_and_save_demo(trainer, pl_module, batch)
+
+    def generate_and_save_demo(
+        self,
+        trainer: L.Trainer,
+        pl_module: L.LightningModule,
+        batch: Dict[str, Any],
+    ):
+        logger = trainer.logger.experiment if trainer.logger is not None else None  # type: ignore
+
+        save_path: Path = self.save_path / "demos" / f"step={trainer.global_step}"
+        save_path.mkdir(parents=True, exist_ok=True)
+
+        # batch_size = min(len(batch["target"]), self.n_demos)
+        batch_size = len(batch["target"])
+        if self.n_demos < len(batch["target"]):
+            chosen_samples = random.sample(list(range(batch_size)),
+                                           self.n_demos)
+        else:
+            chosen_samples = list(range(batch_size))
+        n_samples = len(chosen_samples)
+
+        with isolate_rng(include_cuda=True):
+            # extract conditioning data from batch
+            target = batch.get("target")
+            assert target is not None
+            context = batch.get("context")
+            context_dropout_mask = None
+            if context is not None:
+                if isinstance(context, list):
+                    context = [
+                        c for i, c in enumerate(context) if i in chosen_samples
+                    ]
+                else:
+                    assert isinstance(context, Tensor)
+                    context = context[chosen_samples]
+                context_dropout_mask = torch.full((len(context),),
+                                                  True,
+                                                  dtype=torch.bool)
+                # context_dropout_mask[:n_samples // 2] = False # todo uncomment me
+            style = batch.get("style")
+            if style is not None:
+                style = style[chosen_samples]
+            description = batch.get("description")
+            if description is not None:
+                description = [
+                    v for i, v in enumerate(description) if i in chosen_samples
+                ]
+            beat = batch.get("beat")
+            if beat is not None:
+                beat = [v for i, v in enumerate(beat) if i in chosen_samples]
+
+            # generate
+            L.seed_everything(42)
+            with torch.autocast(device_type="cuda"):
+                gen_audio = pl_module.generate(
+                    n_samples=n_samples,
+                    gen_seconds=10,
+                    prompt=None,
+                    context=context,
+                    context_dropout_mask=context_dropout_mask,
+                    style=style,
+                    beat=beat,
+                    description=description,
+                    prog_bar=cfg.running_locally())
+
+            # compute cocola score between generated audio and context/style
+            '''if style is not None or context is not None:
+                cocola = CoCola(
+                    embedding_mode=cocola_constants.EmbeddingMode.BOTH)
+                cocola.load_state_dict(
+                    torch.load(cfg.weights_dir() / "cocola-weights.pt",
+                               weights_only=True))
+                # cocola = CoCola.load_from_checkpoint(cfg.weights_dir() /
+                #                                      "cocola.pt",
+                #                                      map_location="cpu").eval()
+                cocola.eval()
+                cocola.set_embedding_mode(cocola_constants.EmbeddingMode.BOTH)
+                feature_extractor = CoColaFeatureExtractor()
+                gen_features = feature_extractor(gen_audio.cpu())
+                if style is not None:
+                    assert len(gen_audio) == len(style)
+                    style_features = feature_extractor(style.cpu())
+                    style_score = cocola.score(gen_features, style_features)
+                if context is not None and isinstance(context, Tensor):
+                    assert len(gen_audio) == len(context)
+                    context_features = feature_extractor(context.cpu())
+                    context_score = cocola.score(gen_features, context_features)'''
+
+        columns = ([
+            s for s in ("description", "context", "beat") if s in batch
+        ] + ["generated"])
+        if "context" in columns:
+            columns.append("mix")
+            # mixed_audio = gen_audio + context
+        if "beat" in columns and "context" in columns:
+            columns.append("context with beat")
+        # if "style" in columns:
+        #     columns.append("style cocola score")
+        # if "context" in columns and isinstance(context, Tensor):
+        #     columns.append("context cocola score")
+
+        data = [[] for _ in range(n_samples)]
+
+        # for idx in range(n_samples):
+        for i, idx in enumerate(chosen_samples):
+            gen_filename: Path = save_path / f"demo{i}_gen.wav"
+            audio.save_audio(gen_audio[i], gen_filename)
+
+            if description is not None:
+                desc_filename: Path = save_path / f"demo{i}_description.txt"
+                desc_filename.write_text(description[i])
+            '''if style is not None:
+                style_filename: Path = save_path / f"demo{i}_style.wav"
+                audio.save_audio(style[i], style_filename)'''
+            if context is not None and context_dropout_mask[i]:
+                context_filename: Path = save_path / f"demo{i}_context.wav"
+                mix_filename: Path = save_path / f"demo{i}_mix.wav"
+                audio.save_audio(context[i], context_filename)
+                mix = torch.nn.functional.pad(
+                    context[i],
+                    (0, gen_audio[i].shape[-1] - context[i].shape[-1]),
+                    value=0) + gen_audio[i]
+                audio.save_audio(mix, mix_filename)
+            if "target" in batch:
+                target_filename: Path = save_path / f"demo{i}_target.wav"
+                audio.save_audio(target[idx], target_filename)
+
+            if beat is not None:
+                beat_filename: Path = save_path / f"demo{i}_beat.wav"
+                b: Beat = beat[i]
+                beats = (b.beats / pl_module.sample_rate).cpu().numpy()
+                downbeats = (b.downbeats / pl_module.sample_rate).cpu().numpy()
+                clicks = librosa.clicks(times=beats,
+                                        sr=pl_module.sample_rate,
+                                        click_freq=1000,
+                                        length=gen_audio[i].shape[-1])
+                high_clicks = librosa.clicks(times=downbeats,
+                                             sr=pl_module.sample_rate,
+                                             click_freq=2000,
+                                             length=gen_audio[i].shape[-1])
+                gen_with_beat = gen_audio[i].cpu() + (clicks *
+                                                      0.5) + (high_clicks * 0.5)
+                audio.save_audio(gen_with_beat, beat_filename)
+                if context is not None and context_dropout_mask[i]:
+                    context_with_beats_filename: Path = (
+                        save_path / f"demo{i}_context_with_beats.wav")
+                    context_with_beats = torch.nn.functional.pad(
+                        context[i].cpu(),
+                        (0, len(clicks) - context[i].shape[-1])) + (
+                            clicks * 0.5) + (high_clicks * 0.5)
+                    audio.save_audio(context_with_beats,
+                                     context_with_beats_filename)
+
+            if logger is not None:
+                # if pl_module.logger is not None:
+                row = []
+                if description is not None:
+                    row.append(description[i])
+                '''if style is not None:
+                    row.append(
+                        wandb.Audio(str(style_filename),
+                                    sample_rate=pl_module.sample_rate))'''
+                if context is not None:
+                    assert context_dropout_mask is not None
+                    if context_dropout_mask[i]:
+                        row.append(
+                            wandb.Audio(str(context_filename),
+                                        sample_rate=pl_module.sample_rate))
+                    else:
+                        row.append(None)
+
+                if beat is not None:
+                    row.append(
+                        wandb.Audio(str(beat_filename),
+                                    sample_rate=pl_module.sample_rate))
+
+                row.append(
+                    wandb.Audio(str(gen_filename),
+                                sample_rate=pl_module.sample_rate))
+                if context is not None:
+                    assert context_dropout_mask is not None
+                    if context_dropout_mask[i]:
+                        row.append(
+                            wandb.Audio(str(mix_filename),
+                                        sample_rate=pl_module.sample_rate))
+                    else:
+                        row.append(None)
+                if context is not None and beat is not None:
+                    assert context_dropout_mask is not None
+                    if context_dropout_mask[i]:
+                        row.append(
+                            wandb.Audio(str(context_with_beats_filename),
+                                        sample_rate=pl_module.sample_rate))
+                    else:
+                        row.append(None)
+                # if style is not None:
+                #     row.append(style_score[i])
+                # if context is not None and isinstance(context, Tensor):
+                #     row.append(context_score[i])
+
+                data[i] = row
+
+        # if pl_module.logger is not None:
+        if logger is not None:
+            table = wandb.Table(columns=columns, data=data)
+            logger.log({f"demos/step:{trainer.global_step}": table})
+            # if style is not None:
+            #     logger.log({"valid/style_cocola": style_score.mean()})
+            # if context is not None and isinstance(context, Tensor):
+            #     logger.log({"valid/context_cocola": context_score.mean()})
+            # logger.log_table(  # type: ignore
+            #     key=f"demos_step{trainer.global_step}",
+            #     columns=columns,
+            #     data=data)
+
+
+class PrintLossesCallback(L.Callback):
+
+    def __init__(self):
+        super().__init__()
+        self.train_losses = []
+        self.val_losses = []
+
+    def on_train_epoch_end(self, trainer, pl_module):
+        # Compute and print the average training loss for the epoch
+        train_loss = trainer.callback_metrics.get("train_loss")
+        if train_loss is None:
+            return
+        print(f"Epoch {trainer.current_epoch +1} / "
+              f"train loss: {train_loss:.4f}")
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        # Compute and print the average validation loss for the epoch
+        val_loss = trainer.callback_metrics.get("val_loss")
+        if val_loss is None:
+            return
+        print(f"Epoch {trainer.current_epoch + 1} / "
+              f"validation loss: {val_loss:.4f}")

training/train.py

@@ -0,0 +1,260 @@
+from pathlib import Path
+from time import sleep
+import traceback
+from typing import List, Optional
+import lightning as L
+from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch.callbacks import OnExceptionCheckpoint
+import wandb
+from torch import Tensor
+import torch
+
+import hyperparameters as hp
+from conditioning.clap_embedder import LinearClapEmbedder
+from conditioning.condition_type import ConditionType
+from conditioning.conditioning_method import ConditioningMethod
+from conditioning.prompt_processor import InterleavedContextPromptProcessor, StraightContextPromptProcessor
+from conditioning.t5embedder import T5EmbedderCPU, T5EmbedderGPU
+import config as cfg
+from data.stem import Stem
+from models.lightning_musicgen import LightningMusicgen
+from training.callback import PrintLossesCallback, SaveDemoOnValidationCallback
+from lightning.pytorch.callbacks import TQDMProgressBar, ModelCheckpoint
+from utils.inspection import print_params
+from utils.logging import get_or_create_run_id
+
+
+def train(model_params: hp.ModelParams,
+          dataset_params: hp.DatasetParams,
+          max_steps: Optional[int] = None,
+          max_time: Optional[str] = None,
+          max_epochs: Optional[int] = None,
+          log: Optional[bool] = None,
+          accumulate_grad_batches: int = 1,
+          validate_every_n_steps: Optional[int] = None,
+          distributed_strategy: Optional[str] = None,
+          seed: Optional[int] = None,
+          run_name: Optional[str] = None,
+          kill_on_end: bool = False,
+          n_demos_per_epoch: int = 6,
+          devices: int = -1):
+
+    if log is None:
+        log = run_name is not None
+    if log and run_name is None:
+        raise ValueError("Need a run name to log on wandb")
+
+    # init model and datamodule
+    full_run_params = {
+        "model_params": model_params.to_dict(),
+        "data_params": dataset_params.to_dict()
+    }
+    model = model_params.instantiate()
+    datamodule = dataset_params.instantiate()
+    n_train_batches, n_valid_batches = (datamodule.lengths["train"],
+                                        datamodule.lengths["valid"])
+
+    # setup callbacks and logger
+    logger: WandbLogger | bool = False
+    resume_from_checkpoint: Optional[Path] = None
+    if log:
+        assert run_name is not None
+        output_dir = cfg.output_dir() / run_name
+        if output_dir.exists():
+            print(
+                f"Output directory for a run named {run_name} exists. Resuming training..."
+            )
+            resume_from_checkpoint = output_dir / "last.ckpt"
+        wandb_run_id = get_or_create_run_id(output_dir)
+        logger = WandbLogger(
+            entity=cfg.ENTITY,
+            project=cfg.PROJECT,
+            name=run_name,
+            id=wandb_run_id,
+            resume="allow",
+            config=full_run_params,
+            # settings=wandb.Settings(start_method="fork"),
+        )
+    callbacks: List[L.Callback] = []
+    if not cfg.running_locally():
+        # printlossescallback = PrintLossesCallback()
+        # callbacks.append(printlossescallback)
+        progbar: L.Callback = TQDMProgressBar(refresh_rate=n_train_batches // 2)
+        callbacks.append(progbar)
+    if run_name is not None:
+        output_dir = cfg.output_dir() / run_name
+        savedemocallback = SaveDemoOnValidationCallback(
+            output_dir, save_model=False, n_demos=n_demos_per_epoch)
+        interruptcallback = OnExceptionCheckpoint(output_dir,
+                                                  filename="interrupted")
+        modelcheckpoint = ModelCheckpoint(dirpath=output_dir,
+                                          save_last=True,
+                                          every_n_epochs=1,
+                                          save_top_k=-1)
+        callbacks += [
+            interruptcallback,
+            savedemocallback,
+            modelcheckpoint,
+        ]
+
+    # init trainer
+    trainer = L.Trainer(
+        enable_model_summary=True,
+        accelerator="auto",
+        max_steps=max_steps or -1,
+        max_epochs=max_epochs,
+        max_time=max_time,
+        devices=devices,
+        strategy=distributed_strategy or "auto",
+        gradient_clip_val=1.0,
+        accumulate_grad_batches=accumulate_grad_batches,
+        gradient_clip_algorithm="value",
+        precision="16-mixed",
+        callbacks=callbacks,
+        logger=logger,
+        log_every_n_steps=10,
+        val_check_interval=validate_every_n_steps or 1.0,
+        # check_val_every_n_epoch=None,
+        limit_train_batches=n_train_batches,
+        limit_val_batches=n_valid_batches,
+        num_sanity_val_steps=-1,
+        enable_progress_bar=True,
+    )
+
+    # set seed
+    if seed is not None:
+        L.seed_everything(seed)
+
+    # run training
+    if kill_on_end:
+        try:
+            trainer.fit(model, datamodule=datamodule)
+            print("Training is finished. Killing myself in five minutes.")
+            try:
+                wandb.finish()
+                sleep(300)
+                cfg.shutdown()
+            except KeyboardInterrupt:
+                print("You saved me! I'll never forget that.")
+                return
+            except Exception:
+                cfg.shutdown()
+
+        except KeyboardInterrupt:
+            print("Received keyboard interrupt. Stopping training "
+                  "without shutting down...")
+            wandb.finish()
+
+        except Exception as e:
+            (cfg.output_dir() / "exception.txt").write_text(
+                f"Exception: {str(e)}\n\n "
+                f"Stacktrace: {traceback.format_exc()}\n")
+            print(f"training broke with exception {e}")
+            print(f"Killing myself in five minutes")
+            try:
+                wandb.finish()
+                sleep(300)
+                cfg.shutdown()
+            except KeyboardInterrupt:
+                print("You saved me! I'll never forget that.")
+                return
+            except Exception:
+                cfg.shutdown()
+
+    else:
+        trainer.fit(model, datamodule=datamodule)
+        if run_name is not None:
+            wandb.finish()
+
+    return
+
+
+if __name__ == "__main__":
+    from time import time
+
+    encodec_params = hp.pretrained_encodec_meta_32khz_params
+    lm_params = hp.FioraSmallLmParams()
+    prompt_processor_params = hp.PromptProcessorParams(
+        keep_only_valid_steps=True,
+        model_class=InterleavedContextPromptProcessor,
+        context_dropout=0.5)
+    conditioning_params = hp.ConditioningParams(
+        embedder_types={
+            ConditionType.DESCRIPTION: T5EmbedderGPU,
+            ConditionType.STYLE: LinearClapEmbedder
+        },
+        conditioning_methods={
+            ConditionType.DESCRIPTION: ConditioningMethod.CROSS_ATTENTION,
+            ConditionType.STYLE: ConditioningMethod.INPUT_SUM,
+        },
+        conditioning_dropout=0.5)
+
+    model_params: hp.MusicgenParams = hp.MusicgenParams(
+        encodec_params=encodec_params,
+        lm_params=lm_params,
+        prompt_processor_params=prompt_processor_params,
+        conditioning_params=conditioning_params)
+
+    dataset_params = hp.MixDatasetParams(clip_length_in_seconds=10,
+                                         sample_rate=32_000,
+                                         root_dir=cfg.mixdata_path(),
+                                         single_stem=True,
+                                         target_stem=Stem.DRUMS,
+                                         min_context_seconds=5,
+                                         use_style_conditioning=True,
+                                         use_beat_conditioning=False,
+                                         type_of_context="stems",
+                                         add_click=False,
+                                         sync_chunks=False,
+                                         bpm_in_caption=False,
+                                         batch_size_train=2,
+                                         batch_size_test=12,
+                                         num_workers=8,
+                                         speed_transform_p=0.5,
+                                         pitch_transform_p=0.5,
+                                         n_samples_per_epoch=2000)
+
+    device = "cuda"
+    model: LightningMusicgen = model_params.instantiate().to(device)
+    datamodule = dataset_params.instantiate()
+
+    # validation step test
+    # model.eval()
+    # vd = iter(datamodule.val_dataloader())
+    # for i in range(2):
+    #     batch = next(vd)
+    #     batch = {
+    #         k: v.to(device) if isinstance(v, Tensor) else v
+    #         for k, v in batch.items()
+    #     }
+    #     t0 = time()
+    #     with torch.autocast(device_type="cuda"):
+    #         val_loss = model.validation_step(batch, i)
+    #     t1 = time()
+    #     print(f"val step in {t1 - t0} seconds")
+
+    # training step test
+    # td = iter(datamodule.train_dataloader())
+    # model.train()
+    # for i in range(10):
+    #     batch = next(td)
+    #     batch = {
+    #         k: v.to(device) if isinstance(v, Tensor) else v
+    #         for k, v in batch.items()
+    #     }
+    #     t0 = time()
+    #     with torch.autocast(device_type="cuda"):
+    #         train_loss = model.training_step(batch, i)
+    #     t1 = time()
+    #     print(f"training step in {t1 - t0} seconds")
+
+    trainer = L.Trainer(accelerator="auto",
+                        precision="16-mixed",
+                        enable_model_summary=True,
+                        logger=None,
+                        enable_checkpointing=False,
+                        num_sanity_val_steps=2,
+                        limit_train_batches=50,
+                        limit_val_batches=2,
+                        max_epochs=10)
+    trainer.fit(model, datamodule=datamodule)

utils/audio.py

@@ -0,0 +1,293 @@
+from functools import partial
+import random
+import time
+from torch import Tensor
+from pathlib import Path
+import torch
+import torchaudio
+import soundfile as sf
+from typing import Dict, List, Optional, Sequence, Tuple, Union
+import concurrent.futures
+import numpy as np
+
+from data.stem import Stem
+
+
+def pad_stack(tensors: List[Tensor],
+              pad_value: int | float,
+              padding_dim: int = -1,
+              pad_start: bool = False,
+              stack_dim: int = 0):
+    dtypes = {t.dtype for t in tensors}
+    if len(dtypes) > 1:
+        raise ValueError("Input tensors have different types")
+    ndims = {t.ndim for t in tensors}
+    if len(ndims) > 1:
+        raise ValueError("Input tensors have different number of dimensions")
+    if len(tensors) == 0:
+        raise ValueError("Input list cannot be empty")
+    ndim = len(tensors[0].shape)
+    shapes = [t.shape for t in tensors]
+    for i in range(1, len(shapes)):
+        shape_i = list(tensors[i].shape)
+        shape_ii = list(tensors[i - 1].shape)
+        shape_i.pop(padding_dim)
+        shape_ii.pop(padding_dim)
+        if shape_i != shape_ii:
+            raise ValueError(
+                f"Shape of tensors at indices ({i-1}, {i}) don't match")
+
+    max_len = max(t.shape[padding_dim] for t in tensors)
+    if padding_dim < 0:
+        padding_dim = ndim + padding_dim
+
+    def get_pad_code(n):
+        pre = [0] * (2 * (ndim - padding_dim - 1))
+        mid = [n, 0] if pad_start else [0, n]
+        post = [0] * (2 * (padding_dim))
+        pad_code = pre + mid + post
+        return pad_code
+
+    padded_tensors = [
+        torch.nn.functional.pad(t,
+                                get_pad_code(max_len - t.shape[padding_dim]),
+                                value=pad_value) for t in tensors
+    ]
+    return torch.stack(padded_tensors, dim=stack_dim)
+
+
+def to_stereo(audio: Tensor) -> Tensor:
+    if audio.dim() == 1:
+        return audio.repeat(2, 1)
+    if audio.dim() >= 2 and audio.shape[-2] == 2:
+        return audio
+    return torch.cat((audio, audio), dim=-2)
+
+
+def to_mono(audio: Tensor) -> Tensor:
+    if audio.ndim == 1:
+        return audio
+    if audio.shape[-2] == 1:
+        return audio
+    return audio.mean(dim=-2, keepdim=True)
+
+
+def save_audio(audio: Tensor, path: Path, sample_rate: int = 32_000):
+    audio = audio.float()
+    if audio.shape[-1] == 1:
+        audio.squeeze(-1)
+    l = audio.shape[-1]
+    audio = audio.reshape(-1, l)
+    # audio = mono_to_stereo(audio.detach().cpu())
+    audio = to_stereo(audio.detach().cpu())
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if audio.dim() != 2:
+        print(f'{audio.shape=}')
+
+    torchaudio.save(str(path), audio, sample_rate=sample_rate)  # type: ignore
+
+
+def load_audio(path: Path,
+               sample_rate: int = 32_000,
+               stereo: bool = False) -> Tensor:
+    # soundfile instead of torchaudio.load to avoid torchcodec dependency
+    audio_np, orig_sr = sf.read(str(path), always_2d=True)
+    audio = torch.from_numpy(audio_np.T).float()
+    audio = torchaudio.functional.resample(audio, orig_sr, sample_rate)
+    if not stereo:
+        audio = to_mono(audio)
+    return audio.reshape(1, 1, -1)  # type: ignore
+
+
+def load_audio_chunk(audio_path: Path, start_offset: int, num_frames: int,
+                     stereo: bool) -> Tensor:
+
+    # info = torchaudio.info(str(audio_path))
+    # length = info.num_frames
+    # file_sample_rate = info.sample_rate
+
+    # assert file_sample_rate == sample_rate
+    try:
+        wav, sr = torchaudio.load(str(audio_path),
+                                  frame_offset=start_offset,
+                                  num_frames=num_frames,
+                                  backend="soundfile")
+    except:
+        wav = torch.zeros(2, num_frames)
+
+    # if start_offset + num_frames >= length:
+    #     wav = torch.zeros(2, num_frames)
+
+    # wav = torchaudio.functional.resample(wav, sr, sample_rate)
+
+    if wav.shape[-1] < num_frames:
+        wav = torch.nn.functional.pad(wav,
+                                      pad=(0, num_frames - wav.shape[-1]),
+                                      mode="constant",
+                                      value=0)
+
+    if not stereo:
+        # wav: Tensor = stereo_to_mono(wav).reshape(1, -1)
+        wav: Tensor = to_mono(wav).reshape(1, -1)
+
+    return wav
+
+
+def is_silent(audio: Tensor, threshold: float = 1e-2):
+    return audio.max().item() < threshold
+
+
+def create_click(shape: Sequence[int],
+                 sr: int,
+                 beats: Sequence[int],
+                 click_freq: int = 440,
+                 click_length: int = 200) -> Tensor:
+    click_track: Tensor = torch.zeros(shape)
+    sinewave: Tensor = create_sine_wave(click_freq, sr, click_length)
+    for beat in beats:
+        if beat >= click_track.shape[-1]:
+            break
+        for offset in range(200):
+            idx = beat + offset
+            if idx < click_track.shape[-1]:
+                click_track[..., idx] = sinewave[offset]
+    return click_track
+
+
+def create_sine_wave(freq: float, sr: int, length: int) -> Tensor:
+    cycle_len = int(sr // freq)
+    cycle = torch.linspace(start=0, end=2 * torch.pi, steps=cycle_len)
+    cycle = cycle.repeat(length // cycle_len + 1)
+    cycle = cycle[:length]
+    wave = cycle.sin()
+    return wave
+
+
+def stretch(audio: Tensor, sample_rate: int, speed_factor: float,
+            pitch_factor: int) -> Tensor:
+    import pylibrb  # optional dep, not needed for inference
+    stretcher = pylibrb.RubberBandStretcher(
+        sample_rate=sample_rate,
+        channels=1,
+        options=pylibrb.Option.PROCESS_OFFLINE | pylibrb.Option.ENGINE_FASTER,
+        initial_time_ratio=speed_factor,
+        initial_pitch_scale=pow(2, pitch_factor / 12))
+    stretcher.set_max_process_size(audio.shape[-1])
+    audio_in = audio.reshape(1, -1).numpy()
+    stretcher.study(audio_in, final=True)
+    stretcher.process(audio_in, final=True)
+    audio_out = torch.from_numpy(stretcher.retrieve_available()).reshape(
+        1, -1).float()
+    return audio_out
+
+
+def stretch_with_timeout(audio: Tensor, sample_rate: int, speed_factor: float,
+                         pitch_factor: int, timeout_seconds: float):
+    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+        future = executor.submit(stretch,
+                                 audio,
+                                 sample_rate=sample_rate,
+                                 speed_factor=speed_factor,
+                                 pitch_factor=pitch_factor)
+        try:
+            return future.result(timeout=timeout_seconds)
+        except concurrent.futures.TimeoutError:
+            print("Timeout occurred in stretching audio.")
+            return audio
+        except Exception as e:
+            print(f"Exception occurred in stretching audio: {e}.")
+            return audio
+
+
+def make_variable_frequency_sinewave(t_end: int,
+                                     peak_indices: Tensor) -> Tensor:
+    if not isinstance(peak_indices, Tensor):
+        peak_indices = torch.tensor(peak_indices, dtype=torch.int32)
+
+    device = peak_indices.device
+
+    if peak_indices.shape[-1] < 2:
+        return torch.zeros((t_end,), device=device)
+
+    # Define the fine-grained time array
+    # time = torch.linspace(t_start, t_end, 100_000)
+    time = torch.arange(t_end, device=device)
+
+    # Calculate frequencies for each interval
+    intervals = torch.diff(peak_indices)  # Time intervals between beats
+    frequencies = 1 / intervals  # Frequencies for each interval
+
+    # Find segment indices for each time point
+    segment_indices = torch.searchsorted(peak_indices, time, right=True) - 1
+    segment_indices = torch.clamp(segment_indices, 0, len(frequencies) - 1)
+
+    # Compute sinewave for all time points
+    phase_shift = torch.pi / 2
+    relative_time = time - peak_indices[segment_indices]
+    wave = torch.sin(2 * torch.pi * frequencies[segment_indices] *
+                     relative_time + phase_shift)
+
+    # Extend before the first peak
+    before_mask = time < peak_indices[0]
+    freq_before = 1 / (peak_indices[1] - peak_indices[0])
+    wave[before_mask] = torch.sin(2 * torch.pi * freq_before *
+                                  (time[before_mask] - peak_indices[0]) +
+                                  phase_shift)
+
+    # Extend after the last peak
+    after_mask = time >= peak_indices[-1]
+    freq_after = 1 / (peak_indices[-1] - peak_indices[-2])
+    wave[after_mask] = torch.sin(2 * torch.pi * freq_after *
+                                 (time[after_mask] - peak_indices[-1]) +
+                                 phase_shift)
+
+    return wave
+
+
+def normalize(audio: Tensor, new_min, new_max):
+    if len(audio.shape) == 2 and audio.shape[0] == 2:
+        audio = audio.mean(dim=-1)
+    audio = to_mono(audio)
+    # Calculate the min and max of the original array
+    old_min = audio.min()
+    old_max = audio.max()
+
+    # Apply the normalization formula
+    normalized_arr = (audio - old_min) / (old_max - old_min) * (
+        new_max - new_min) + new_min
+    return normalized_arr
+
+
+def play(waveform: torch.Tensor, sr: int):
+    import IPython.display
+    import sounddevice as sd
+    waveform = to_stereo(waveform)
+    waveform_np = waveform.cpu().float().detach().numpy()
+    if is_interactive():
+        IPython.display.display(IPython.display.Audio(waveform_np, rate=sr))
+    else:
+        sd.play(waveform_np.T, sr)
+        sd.wait()
+
+
+def is_interactive():
+    import sys
+    return "ipykernel" in sys.modules
+
+
+def inject_clicks(audio_tensor: Tensor, beat_positions: Tensor,
+                  sample_rate: int):
+    """
+    Add short clicks at `beat_positions` (sample indices) in `audio_tensor`.
+    If audio is multi-channel, we'll apply the same clicks to each channel.
+    """
+    import librosa
+    beat_positions_seconds = beat_positions / sample_rate
+    click_track: Tensor = torch.tensor(
+        librosa.clicks(times=beat_positions_seconds.cpu().detach().numpy(),
+                       hop_length=1,
+                       length=audio_tensor.shape[-1],
+                       sr=sample_rate))
+
+    click_track = click_track.broadcast_to(audio_tensor.shape)
+    return click_track + audio_tensor.detach().cpu()

utils/aws.py

@@ -0,0 +1,119 @@
+"""
+This is the script that runs on the "controller" machine: the lightweight 
+machine that requests the big boi instance(s) from EC2 and configures the 
+training job. In this example, I want the controller machine to be my home
+computer. This means that I can run this script from my computer and start a 
+training job on the remote big boi instance, without an intermediary machine
+from aws.
+
+This script can run in two modes, local and remote, depending on the value of
+`run_locally`.
+
+VERY IMPORTANT NOTE: running this script locally is not at all the same as 
+running a normal training locally (i.e: running locally the file `train.py`):
+Running this script in local mode still creates a sagemaker job and does 
+everything in the same exact way as it would to run on a remote instance, 
+EXCEPT that the big boi instance in EC2 is replaced by your local machine.
+
+So the script creates a docker container, downloads it, creates the virtual 
+environment, and runs the training job in the same exact way as it would on the
+remote instance. Running this training job in local mode successfuly should 
+GUARANTEE that it will run successfully on a remote instance.
+"""
+
+import sagemaker.pytorch
+import sagemaker.local
+import sagemaker
+from datetime import datetime
+
+
+def launch_aws_job(
+    experiment_name: str,
+    instance_type: str,
+    max_training_hours: int,
+    entry_point: str,
+    input_mode: str = "File",
+    data_dir: str = "",
+):
+    # True: run training job locally for testing purposes (needs docker installed)
+    # False: run training job on an AWS instance
+    run_locally: bool = False
+
+    # define names for project, experiment and run
+    project_name = "latent-accompaniment-generation"
+    # experiment_name = "bart-drums-mixdata"
+
+    checkpoint_ec2_path = "/opt/ml/checkpoints"
+    # checkpoint_s3_path = "s3://latent-accompaniment-generation/checkpoints/ash-drums"
+
+    # NOTE: job name has to be unique
+    now = datetime.now()
+    datestring = now.strftime('%Y-%m-%d-%H-%M')
+    job_name = experiment_name + "-" + datestring
+
+    # define session and role
+    sagemaker_session = (sagemaker.local.LocalSession()
+                         if run_locally else sagemaker.Session())
+    role = ("arn:aws:iam::076456026604:role/service-role/"
+            "AmazonSageMaker-ExecutionRole-20240407T195108")
+
+    # define input and output directories
+    if run_locally:
+        data_path: str = "file://./data/mtg-jamendo-low/"
+        output_path: str = ("file://./checkpoints/"
+                            f"{project_name}/{experiment_name}/{job_name}")
+    else:
+        # NOTE: whatever directory is specified in here will be downloaded to the
+        # training machine in its entirety (if using "File" input mode, which is
+        # the default). So make sure that the s3 directory specified here only
+        # contains data that is needed for the training.
+        # WRONG: s3://audio-data-bucket/data/       (.../MNIST, .../moisesdb)
+        # OK:    s3://audio-data-bucket/data/MNIST
+        data_path: str = "s3://audio-data-bucket/data/lag-data/" + data_dir
+        weights_path: str = "s3://latent-accompaniment-generation/weights/"
+        output_path = f"s3://{project_name}/{experiment_name}"
+        # ckp_path = checkpoint_s3_path
+
+    # instance_type: str = "ml.g4dn.xlarge"
+    # instance_type: str = "ml.m5.xlarge"
+    # instance_type: str = "ml.g5.2xlarge" # 1x nvidia A10
+    # instance_type: str = "ml.g5.12xlarge"  # 4x nvidia A10
+    # instance_type: str = "ml.p4d.24xlarge"  # 8x nvidia A100 40gb 💀
+    # instance_type: str = "ml.p4de.24xlarge"  # 8x nvidia A100 80gb 💀
+    # instance_type: str = "ml.trn1.2xlarge"
+
+    # max_training_time = 3 * 24 * 60 * 60  # 3 days
+    # max_training_time = 16 * 60 * 60  # 16 hours
+    max_training_time = max_training_hours * 60 * 60
+
+    estimator = sagemaker.pytorch.PyTorch(
+        # entry_point="scripts/train_bart.py",
+        entry_point=entry_point,
+        role=role,
+        max_run=max_training_time,
+        instance_count=1,
+        framework_version="2.1.0",
+        py_version="py310",
+        dependencies=["requirements.txt"],
+        source_dir="src",
+        output_path=output_path,
+        instance_type="local" if run_locally else instance_type,
+        local_code=run_locally,
+        # input_mode="FastFile",
+        # input_mode="File",
+        input_mode=input_mode,
+        checkpoint_local_path=checkpoint_ec2_path,
+        # checkpoint_s3_uri=checkpoint_s3_path,
+        # distribution={"pytorchddp": {
+        #     "enabled": "true"
+        # }},
+    )
+
+    estimator.fit(
+        {
+            "data": data_path,
+            # "ckp": ckp_path,
+            "weights": weights_path,
+        },
+        job_name=job_name,
+    )